1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
|
// RUN: mlir-opt %s --test-transform-dialect-interpreter -canonicalize -cse -split-input-file | FileCheck %s
// Offset per thread:
// CHECK-DAG: affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 10))>
// Per thread tile size.
// CHECK-DAG: affine_map<(d0)[s0] -> (-(d0 * (s0 ceildiv 10)) + s0, s0 ceildiv 10)>
// CHECK-DAG: affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 20))>
// CHECK-DAG: affine_map<(d0)[s0] -> (-(d0 * (s0 ceildiv 20)) + s0, s0 ceildiv 20)>
module {
// CHECK-LABEL: matmul(
// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<?x?xf32>
func.func @matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: scf.forall ({{.*}}) in (10, 20) shared_outs(%[[C_BLK:.*]] = %[[C]]) -> (tensor<?x?xf32>) {
// CHECK: %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[tB:.*]] = tensor.extract_slice %[[B]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[tC:.*]] = tensor.extract_slice %[[C_BLK]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[RES:.*]] = linalg.matmul
// CHECK-SAME: ins(%[[tA]], %[[tB]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[tC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: scf.forall.in_parallel {
// CHECK-NEXT: tensor.parallel_insert_slice %[[RES]] into %[[C_BLK]]{{.*}} :
// CHECK-SAME: tensor<?x?xf32> into tensor<?x?xf32>
// CHECK-NEXT: }
// CHECK-NEXT: } {mapping = [#gpu.thread<y>, #gpu.thread<x>]}
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 num_threads [10, 20] (mapping = [ #gpu.thread<y>, #gpu.thread<x> ] )
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
}
// -----
// In this test case, matmul dims and tile size are dynamic.
// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
// CHECK-DAG: #[[$map4:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>
// CHECK-LABEL: matmul_tile_size_dynamic_dynamic(
// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<?x?xf32>
func.func @matmul_tile_size_dynamic_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[tile_size_1:.*]] = "test.dummy"()
// CHECK-DAG: %[[tile_size_2:.*]] = "test.dummy"()
// CHECK-DAG: %[[M:.+]] = tensor.dim %[[A]], %[[c0]] :
// CHECK-DAG: %[[N:.+]] = tensor.dim %[[B]], %c1 :
// CHECK-DAG: %[[NT0:.+]] = affine.apply #[[$map0]]()[%[[M]], %[[tile_size_1]]]
// CHECK-DAG: %[[NT1:.+]] = affine.apply #[[$map0]]()[%[[N]], %[[tile_size_2]]]
// CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]])
// CHECK: tensor.extract_slice %[[A]]
// CHECK: tensor.extract_slice %[[B]]
// CHECK: tensor.extract_slice %[[C_BLK]]
// CHECK: linalg.matmul
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice
%tile_size_1 = "test.dummy"() : () -> (index)
%tile_size_2 = "test.dummy"() : () -> (index)
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%sz = transform.structured.match ops{["test.dummy"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 tile_sizes *(%sz : !transform.any_op)
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// -----
// Tests that dimension 0 can eliminate affine.min/max, dimension 1 cannot.
// CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * -15 + 300, 15)>
// CHECK-DAG: #[[$map1:.+]] = affine_map<(d0) -> (0, d0)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0) -> (d0 * 10)>
// CHECK-DAG: #[[$map3:.+]] = affine_map<(d0) -> (d0 * 15)>
// CHECK-LABEL: matmul_static(
// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor
// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor
// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor
func.func @matmul_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> {
// CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) in (10, 21) shared_outs(%[[C_BLK:.*]] = %[[C]])
// CHECK: %[[TSMIN:.+]] = affine.min #[[$map0]](%[[IV1]])
// CHECK: %[[TS:.+]] = affine.max #[[$map1]](%[[TSMIN]])
// CHECK-NOT: affine.min
// CHECK-NOT: affine.max
// CHECK: %[[LB0:.+]] = affine.apply #[[$map2]](%[[IV0]])
// CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]])
// CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] :
// CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] :
// CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] :
// CHECK: linalg.matmul
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice
%0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>)
outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>)
return %0 : tensor<100x300xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 num_threads [10, 21]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// -----
// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0] -> (s0 ceildiv 10)>
// CHECK-DAG: #[[$map1:.+]] = affine_map<()[s0] -> (s0 ceildiv 20)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0)[s0] -> (d0 * -10 + s0, 10)>
// CHECK-DAG: #[[$map4:.+]] = affine_map<(d0)[s0] -> (d0 * -20 + s0, 20)>
// CHECK-DAG: #[[$map5:.+]] = affine_map<(d0) -> (d0 * 10)>
// CHECK-DAG: #[[$map6:.+]] = affine_map<(d0) -> (d0 * 20)>
// CHECK-LABEL: matmul_tile_size_dynamic(
// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<?x?xf32>
func.func @matmul_tile_size_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: %[[M:.+]] = tensor.dim %[[A]], %c0 :
// CHECK: %[[N:.+]] = tensor.dim %[[B]], %c1 :
// CHECK: %[[NT0:.+]] = affine.apply #map()[%[[M]]]
// CHECK: %[[NT1:.+]] = affine.apply #map1()[%[[N]]]
// CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]])
// CHECK: %[[TS0:.+]] = affine.min #[[$map2]](%[[IV0]])[%[[M]]]
// CHECK: %[[TS1:.+]] = affine.min #[[$map4]](%[[IV1]])[%[[N]]]
// CHECK: %[[LB0:.+]] = affine.apply #[[$map5]](%[[IV0]])
// CHECK: %[[LB1:.+]] = affine.apply #[[$map6]](%[[IV1]])
// CHECK: tensor.extract_slice %[[A]]
// CHECK: tensor.extract_slice %[[B]]
// CHECK: tensor.extract_slice %[[C_BLK]]
// CHECK: linalg.matmul
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 tile_sizes [10, 20]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// -----
// Tests that dimension 0 can eliminate affine.min/max, dimension 1 cannot.
// CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * -21 + 300, 21)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0) -> (d0 * 10)>
// CHECK-DAG: #[[$map3:.+]] = affine_map<(d0) -> (d0 * 21)>
// CHECK-LABEL: matmul_tile_size_static(
// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor
// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor
// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor
func.func @matmul_tile_size_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> {
// CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) in (10, 15) shared_outs(%[[C_BLK:.*]] = %[[C]])
// CHECK: %[[TS:.+]] = affine.min #[[$map0]](%[[IV1]])
// CHECK-NOT: affine.max
// CHECK-NOT: affine.min
// CHECK: %[[LB0:.+]] = affine.apply #[[$map2]](%[[IV0]])
// CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]])
// CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] :
// CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] :
// CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] :
// CHECK: linalg.matmul
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice
%0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>)
outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>)
return %0 : tensor<100x300xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 tile_sizes [10, 21]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// -----
module {
func.func @extract_source(%A: tensor<4xf32>, %B: tensor<16xf32>) -> tensor<4xf32> {
%B1 = tensor.extract_slice %B[10] [4] [1] : tensor<16xf32> to tensor<4xf32>
%result = linalg.generic {indexing_maps = [
affine_map<(d0) -> (d0)>,affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]}
ins(%A : tensor<4xf32>) outs(%B1 : tensor<4xf32>) {
^bb0(%arg3: f32, %arg4: f32): // no predecessors
%2 = arith.addf %arg3, %arg3 : f32
linalg.yield %2 : f32
} -> tensor<4xf32>
return %result : tensor<4xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 num_threads [2] ( mapping = [#gpu.thread<x>])
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
}
// CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 2)>
// CHECK-LABEL: extract_source(
// CHECK: scf.forall (%[[ARG:.*]]) in (2) shared_outs(%{{.*}} = %{{.*}}) -> (tensor<4xf32>) {
// CHECK: %[[OFF:.*]] = affine.apply #[[$map0]](%[[ARG]])
// CHECK: scf.forall.in_parallel {
// CHECK: tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%[[OFF]]] [2] [1] : tensor<2xf32> into tensor<4xf32>
// -----
// In this test case, matmul dims and tile size are dynamic.
// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
// CHECK-DAG: #[[$map1:.+]] = affine_map<()[s0] -> (s0 ceildiv 20)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
// CHECK-DAG: #[[$map3:.+]] = affine_map<(d0)[s0] -> (d0 * -20 + s0, 20)>
// CHECK-DAG: #[[$map4:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>
// CHECK-DAG: #[[$map5:.+]] = affine_map<(d0) -> (d0 * 20)>
// CHECK-LABEL: matmul_tile_size_dynamic_dynamic(
// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<?x?xf32>
func.func @matmul_tile_size_dynamic_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[tile_size:.*]] = "test.dummy"()
// CHECK-DAG: %[[M:.+]] = tensor.dim %[[A]], %[[c0]] :
// CHECK-DAG: %[[N:.+]] = tensor.dim %[[B]], %c1 :
// CHECK-DAG: %[[NT0:.+]] = affine.apply #[[$map0]]()[%[[M]], %[[tile_size]]]
// CHECK-DAG: %[[NT1:.+]] = affine.apply #[[$map1]]()[%[[N]]]
// CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]])
// CHECK: tensor.extract_slice %[[A]]
// CHECK: tensor.extract_slice %[[B]]
// CHECK: tensor.extract_slice %[[C_BLK]]
// CHECK: linalg.matmul
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice
%tile_size = "test.dummy"() : () -> (index)
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%sz = transform.structured.match ops{["test.dummy"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1:2 = transform.structured.tile_to_forall_op %0 tile_sizes [%sz : !transform.any_op, 20]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// -----
// CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * -15 + 100, 15)>
// CHECK-DAG: #[[$map1:.+]] = affine_map<(d0) -> (0, d0)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0) -> (d0 * 15)>
// CHECK-DAG: #[[$map3:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: tile_output_multi_1d_static(
// CHECK-SAME: %[[IN1:[0-9a-z]+]]: tensor<100xf32>
// CHECK-SAME: %[[IN2:[0-9a-z]+]]: tensor<100xf32>
// CHECK-SAME: %[[ORGOUT1:[0-9a-z]+]]: tensor<100xf32>
// CHECK-SAME: %[[ORGOUT2:[0-9a-z]+]]: tensor<100xf32>
func.func @tile_output_multi_1d_static(%IN1: tensor<100xf32>, %IN2: tensor<100xf32>,
%OUT1: tensor<100xf32>, %OUT2: tensor<100xf32>)
-> (tensor<100xf32>, tensor<100xf32>) {
// CHECK: scf.forall (%[[IV0:.+]]) in (7) shared_outs(%[[OUT1:[0-9a-z]+]] = %[[ORGOUT1]], %[[OUT2:[0-9a-z]+]] = %[[ORGOUT2]])
// CHECK: %[[TSMIN:.+]] = affine.min #[[$map0]](%[[IV0]])
// CHECK: %[[TS:.+]] = affine.max #[[$map1]](%[[TSMIN]])
// CHECK-NOT: affine.min
// CHECK-NOT: affine.max
// CHECK: %[[LB:.+]] = affine.apply #[[$map2]](%[[IV0]])
// CHECK: %[[tIN1:.+]] = tensor.extract_slice %[[IN1]][%[[LB]]] [%[[TS]]] [1] :
// CHECK: %[[tIN2:.+]] = tensor.extract_slice %[[IN2]][%[[LB]]] [%[[TS]]] [1] :
// CHECK: %[[tOUT1:.+]] = tensor.extract_slice %[[OUT1]][%[[LB]]] [%[[TS]]] [1] :
// CHECK: %[[tOUT2:.+]] = tensor.extract_slice %[[OUT2]][%[[LB]]] [%[[TS]]] [1] :
// CHECK: %[[RES1:[0-9]+]]:[[RES2:[0-9]+]] = linalg.generic
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#0 into %[[OUT1]][%[[LB]]] [%[[TS]]] [1] :
// CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#1 into %[[OUT2]][%[[LB]]] [%[[TS]]] [1] :
%res1, %res2 = linalg.generic
{
indexing_maps = [affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]
} ins(%IN1, %IN2 : tensor<100xf32>, tensor<100xf32>)
outs(%OUT1, %OUT2 : tensor<100xf32>, tensor<100xf32>)
{
^bb0(%a1: f32, %a2: f32, %a3: f32, %a4: f32):
%1 = arith.addf %a1, %a3 : f32
%2 = arith.addf %a2, %a4 : f32
linalg.yield %1, %2 : f32,f32
} -> (tensor<100xf32>, tensor<100xf32>)
return %res1, %res2 : tensor<100xf32>, tensor<100xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%forall, %tiled_generic = transform.structured.tile_to_forall_op %0 num_threads [7]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// -----
// CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 75)>
// CHECK-DAG: #[[$map1:.+]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0, d1) -> (d1, d0)
// CHECK-DAG: #[[$map3:.+]] = affine_map<(d0, d1) -> (d0)>
// CHECK-DAG: #[[$map4:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: tile_output_multi_1d2d_static(
// CHECK-SAME: %[[IN1:[0-9a-z]+]]: tensor<100xf32>
// CHECK-SAME: %[[IN2:[0-9a-z]+]]: tensor<100x300xf32>
// CHECK-SAME: %[[IN3:[0-9a-z]+]]: tensor<300xf32>
// CHECK-SAME: %[[ORGOUT1:[0-9a-z]+]]: tensor<300x100xf32>
// CHECK-SAME: %[[ORGOUT2:[0-9a-z]+]]: tensor<300xf32>
func.func @tile_output_multi_1d2d_static(%IN1: tensor<100xf32>, %IN2: tensor<100x300xf32>, %IN3: tensor<300xf32>,
%OUT1: tensor<300x100xf32>, %OUT2: tensor<300xf32>)
-> (tensor<300x100xf32>, tensor<300xf32>) {
// CHECK: scf.forall (%[[IV0:.+]]) in (4) shared_outs(%[[OUT1:[0-9a-z]+]] = %[[ORGOUT1]], %[[OUT2:[0-9a-z]+]] = %[[ORGOUT2]])
// CHECK: %[[LB:.+]] = affine.apply #[[$map0]](%[[IV0]])
// CHECK: %[[tIN1:.+]] = tensor.extract_slice %[[IN2]][0, %[[LB]]] [100, 75]
// CHECK: %[[tIN2:.+]] = tensor.extract_slice %[[IN3]][%[[LB]]] [75]
// CHECK: %[[tOUT1:.+]] = tensor.extract_slice %[[OUT1]][%[[LB]], 0] [75, 100]
// CHECK: %[[tOUT2:.+]] = tensor.extract_slice %[[OUT2]][%[[LB]]] [75]
// CHECK: %[[RES1:[0-9]+]]:[[RES2:[0-9]+]] = linalg.generic
// CHECK: scf.forall.in_parallel
// CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#0 into %[[OUT1]][%[[LB]], 0] [75, 100]
// CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#1 into %[[OUT2]][%[[LB]]] [75]
%res2, %res3 = linalg.generic {
indexing_maps = [affine_map<(d0,d1) -> (d1)>,
affine_map<(d0,d1) -> (d1,d0)>,
affine_map<(d0,d1) -> (d0)>,
affine_map<(d0,d1) -> (d0,d1)>,
affine_map<(d0,d1) -> (d0)>
],
iterator_types = ["parallel", "parallel"]
} ins(%IN1, %IN2, %IN3 : tensor<100xf32>, tensor<100x300xf32>, tensor<300xf32>)
outs(%OUT1, %OUT2: tensor<300x100xf32>, tensor<300xf32>) {
^bb0(%i1: f32, %i2: f32, %i3: f32, %o1: f32, %o2: f32):
%1 = arith.addf %i1, %o1 : f32
%2 = arith.addf %i2, %1 : f32
%3 = arith.addf %i3, %2 : f32
linalg.yield %3, %i3 : f32, f32
} -> (tensor<300x100xf32>, tensor<300xf32>)
return %res2, %res3 : tensor<300x100xf32>, tensor<300xf32>
}
transform.sequence failures(propagate) {
^bb1(%IN_MAT2: !transform.any_op):
%0 = transform.structured.match ops{["linalg.generic"]} in %IN_MAT2 : (!transform.any_op) -> !transform.any_op
%forall, %tiled_generic = transform.structured.tile_to_forall_op %0 num_threads [4]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
|
