1// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
2
3// CHECK-LABEL: @vectorize_matmul
4// CHECK-SAME: %[[A:.*]]: tensor<24x12xf32>
5// CHECK-SAME: %[[B:.*]]: tensor<12x25xf32>
6// CHECK-SAME: %[[C:.*]]: tensor<24x25xf32>
7func.func @vectorize_matmul(%arg0: tensor<24x12xf32>,
8                            %arg1: tensor<12x25xf32>,
9                            %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
10  // CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
11  // CHECK: %[[vB:.+]] = vector.transfer_read %[[B]]
12  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
13  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
14  // CHECK: vector.transfer_write %[[vR]], %[[C]]
15  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
16  func.return %0 : tensor<24x25xf32>
17}
18
19transform.with_pdl_patterns {
20^bb0(%arg0: !pdl.operation):
21  pdl.pattern @pdl_target : benefit(1) {
22    %args = operands
23    %results = types
24    %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
25    // TODO: we don't want this, but it is the required terminator for pdl.pattern
26    rewrite %0 with "transform.dialect"
27  }
28
29  transform.sequence %arg0 {
30  ^bb1(%arg1: !pdl.operation):
31    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
32    %1 = get_closest_isolated_parent %0
33    %2 = transform.structured.vectorize %1
34  }
35}
36
37// -----
38
39#map0 = affine_map<()[s0] -> (-s0 + 12, 7)>
40#map1 = affine_map<()[s0] -> (-s0 + 7)>
41
42// CHECK-LABEL: @vectorize_keep_pad
43// CHECK-SAME: %[[C:[a-zA-Z0-9_]+]]: tensor<24x25xf32>
44func.func @vectorize_keep_pad(
45    %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>,
46    %arg2: tensor<24x25xf32>, %arg3: index, %arg4: index,
47    %arg5: index) -> tensor<24x25xf32> {
48  %c0 = arith.constant 0 : index
49  %cst = arith.constant 0.000000e+00 : f32
50  %0 = affine.min #map0()[%arg5]
51  %1 = tensor.extract_slice %arg0[%arg3, %arg5] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
52  %2 = tensor.extract_slice %arg1[%arg5, %arg4] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
53  %3 = tensor.extract_slice %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
54  %4 = affine.apply #map1()[%0]
55  // CHECK: %[[pA:.*]] = tensor.pad
56  %5 = tensor.pad %1 nofold low[%c0, %c0] high[%c0, %4] {
57  ^bb0(%arg6: index, %arg7: index):
58    tensor.yield %cst : f32
59  } : tensor<4x?xf32> to tensor<4x7xf32>
60  %6 = affine.apply #map1()[%0]
61  // CHECK: %[[pB:.*]] = tensor.pad
62  %7 = tensor.pad %2 nofold low[%c0, %c0] high[%6, %c0] {
63  ^bb0(%arg6: index, %arg7: index):
64    tensor.yield %cst : f32
65  } : tensor<?x5xf32> to tensor<7x5xf32>
66  // CHECK: %[[vA:.+]] = vector.transfer_read %[[pA]]
67  // CHECK: %[[vB:.+]] = vector.transfer_read %[[pB]]
68  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
69  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
70  // CHECK: vector.transfer_write %[[vR]], %[[C]]
71  %8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
72  %9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
73  return %9 : tensor<24x25xf32>
74}
75
76transform.with_pdl_patterns {
77^bb0(%arg0: !pdl.operation):
78  transform.sequence %arg0 {
79  ^bb1(%arg1: !pdl.operation):
80    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
81    %1 = get_closest_isolated_parent %0
82    %2 = transform.structured.vectorize %1
83  }
84}
85
86// -----
87
88#map0 = affine_map<()[s0] -> (-s0 + 12, 7)>
89#map1 = affine_map<()[s0] -> (-s0 + 7)>
90
91// CHECK-LABEL: @vectorize_pad
92// CHECK-SAME: %[[A:.+]]: tensor<24x12xf32>
93// CHECK-SAME: %[[B:.+]]: tensor<12x25xf32>
94// CHECK-SAME: %[[C:.+]]: tensor<24x25xf32>
95func.func @vectorize_pad(
96    %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>,
97    %arg2: tensor<24x25xf32>, %arg3: index, %arg4: index,
98    %arg5: index) -> tensor<24x25xf32> {
99  %c0 = arith.constant 0 : index
100  %cst = arith.constant 0.000000e+00 : f32
101  %0 = affine.min #map0()[%arg5]
102  // CHECK: %[[sA:.+]] = tensor.extract_slice %[[A]]
103  // CHECK: %[[sB:.+]] = tensor.extract_slice %[[B]]
104  %1 = tensor.extract_slice %arg0[%arg3, %arg5] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
105  %2 = tensor.extract_slice %arg1[%arg5, %arg4] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
106  %3 = tensor.extract_slice %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
107  // CHECK: %[[vA:.+]] = vector.transfer_read %[[sA]]
108  %4 = affine.apply #map1()[%0]
109  %5 = tensor.pad %1 nofold low[%c0, %c0] high[%c0, %4] {
110  ^bb0(%arg6: index, %arg7: index):
111    tensor.yield %cst : f32
112  } : tensor<4x?xf32> to tensor<4x7xf32>
113  %6 = affine.apply #map1()[%0]
114  // CHECK: %[[vB:.+]] = vector.transfer_read %[[sB]]
115  %7 = tensor.pad %2 nofold low[%c0, %c0] high[%6, %c0] {
116  ^bb0(%arg6: index, %arg7: index):
117    tensor.yield %cst : f32
118  } : tensor<?x5xf32> to tensor<7x5xf32>
119  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
120  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
121  // CHECK: vector.transfer_write %[[vR]], %[[C]]
122  %8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
123  %9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
124  return %9 : tensor<24x25xf32>
125}
126
127transform.with_pdl_patterns {
128^bb0(%arg0: !pdl.operation):
129  transform.sequence %arg0 {
130  ^bb1(%arg1: !pdl.operation):
131    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
132    %1 = get_closest_isolated_parent %0
133    %2 = transform.structured.vectorize %1 {vectorize_padding = true}
134  }
135}
136
137// -----
138
139func.func @vectorize(%arg0: tensor<24x12xf32>,
140                     %arg1: tensor<12x25xf32>,
141                     %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
142  // expected-note @below {{non-isolated target}}
143  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
144  func.return %0 : tensor<24x25xf32>
145}
146
147transform.with_pdl_patterns {
148^bb0(%arg0: !pdl.operation):
149  transform.sequence %arg0 {
150  ^bb1(%arg1: !pdl.operation):
151    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
152    // expected-error @below {{op requires isolated-from-above targets}}
153    %2 = transform.structured.vectorize %0
154  }
155}
156