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
|
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-linalg-transform-patterns=test-patterns -split-input-file | FileCheck %s
func.func @dot(%x: memref<?xf32, strided<[1], offset: ?>>,
%y: memref<?xf32, strided<[1], offset: ?>>,
%v: memref<f32>) {
linalg.dot ins(%x, %y: memref<?xf32, strided<[1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>)
outs(%v: memref<f32>)
return
}
transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.dot"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1, %loop = transform.structured.tile %0 [8000] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
// CHECK-LABEL: func @dot
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c8000:.*]] = arith.constant 8000 : index
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c8000]] {
// CHECK: linalg.dot
// -----
func.func @matvec(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%x: memref<?xf32, strided<[1], offset: ?>>,
%y: memref<?xf32, strided<[1], offset: ?>>) {
linalg.matvec
ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>)
outs(%y: memref<?xf32, strided<[1], offset: ?>>)
return
}
transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1, %loops:2 = transform.structured.tile %0 [5, 6] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
}
// CHECK-LABEL: func @matvec
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
// CHECK: scf.for {{.*}} step %[[c5]]
// CHECK: scf.for {{.*}} step %[[c6]]
// CHECK: linalg.matvec
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?xf32, strided<[1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?xf32, strided<[1], offset: ?>>)
// -----
func.func @matmul(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C: memref<?x?xf32, strided<[?, 1], offset: ?>>)
return
}
transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
%2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
%3, %loops_3:3 = transform.structured.tile %2 [20, 30, 40] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
%4, %loops_4:3 = transform.structured.tile %3 [2, 3, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
}
// CHECK-LABEL: func @matmul
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[c20:.*]] = arith.constant 20 : index
// CHECK-DAG: %[[c30:.*]] = arith.constant 30 : index
// CHECK-DAG: %[[c40:.*]] = arith.constant 40 : index
// CHECK-DAG: %[[c200:.*]] = arith.constant 200 : index
// CHECK-DAG: %[[c300:.*]] = arith.constant 300 : index
// CHECK-DAG: %[[c400:.*]] = arith.constant 400 : index
// CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index
// CHECK-DAG: %[[c3000:.*]] = arith.constant 3000 : index
// CHECK-DAG: %[[c4000:.*]] = arith.constant 4000 : index
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2000]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3000]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4000]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c200]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c300]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c400]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4]] {
// CHECK: linalg.matmul
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>)
// -----
// Map corresponding to a 2D memory access where the stride along the last dim is known to be 1.
// CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CHECK-DAG: #[[$nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)>
// CHECK-DAG: #[[$km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)>
#matmul_accesses = [
affine_map<(m, n, k) -> (m, k)>,
affine_map<(m, n, k) -> (k, n)>,
affine_map<(m, n, k) -> (m, n)>
]
#generic_matmul_trait = {
args_in = 2,
args_out = 1,
indexing_maps = #matmul_accesses,
library_call = "linalg_matmul",
iterator_types = ["parallel", "parallel", "reduction"]
}
func.func @permute_generic(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
linalg.generic #generic_matmul_trait
ins(%A, %B : memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>) {
^bb(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b: f32
%e = arith.addf %c, %d: f32
linalg.yield %e: f32
}
return
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.structured.interchange %0 iterator_interchange = [1, 2, 0] : (!transform.any_op) -> !transform.any_op
}
// CHECK-LABEL: func @permute_generic
// CHECK: linalg.generic {
// CHECK-SAME: indexing_maps = [#[[$kn]], #[[$nm]], #[[$km]]],
// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"],
// CHECK-SAME: library_call = "linalg_matmul"}
// CHECK: memref<?x?xf32, strided<[?, 1], offset: ?>>,
// CHECK-SAME: memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK-SAME: memref<?x?xf32, strided<[?, 1], offset: ?>>
// -----
func.func @matvec_perm(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%x: memref<?xf32, strided<[1], offset: ?>>,
%y: memref<?xf32, strided<[1], offset: ?>>) {
linalg.matvec ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>)
outs(%y: memref<?xf32, strided<[1], offset: ?>>)
return
}
transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1, %loops:2 = transform.structured.tile %0 [5, 6] {interchange = [1, 0]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
}
// CHECK-LABEL: func @matvec_perm
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c6]]
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c5]]
// CHECK: linalg.matvec
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?xf32, strided<[1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?xf32, strided<[1], offset: ?>>)
// -----
func.func @matmul_perm(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>)
return
}
transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] {interchange = [1, 2, 0]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
%2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] {interchange = [1, 0, 2]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
%3, %loops_3:3 = transform.structured.tile %2 [20, 30, 40] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
}
// CHECK-LABEL: func @matmul_perm
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c20:.*]] = arith.constant 20 : index
// CHECK-DAG: %[[c30:.*]] = arith.constant 30 : index
// CHECK-DAG: %[[c40:.*]] = arith.constant 40 : index
// CHECK-DAG: %[[c200:.*]] = arith.constant 200 : index
// CHECK-DAG: %[[c300:.*]] = arith.constant 300 : index
// CHECK-DAG: %[[c400:.*]] = arith.constant 400 : index
// CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index
// CHECK-DAG: %[[c3000:.*]] = arith.constant 3000 : index
// CHECK-DAG: %[[c4000:.*]] = arith.constant 4000 : index
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3000]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4000]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2000]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c300]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c200]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c400]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
// CHECK: linalg.matmul
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>)
|
