1// RUN: mlir-opt -linalg-bufferize  -canonicalize -cse -split-input-file %s | FileCheck %s
2
3#map0 = affine_map<(d0) -> (d0)>
4
5// In-depth checking of a basic case, this is testing
6// - memref.buffer_cast / memref.tensor_load materializations are properly inserted
7// - payload is correctly carried over
8// - affine maps are correctly carried over
9// Later tests will not check all these details.
10
11// CHECK: #map = affine_map<(d0) -> (d0)>
12// CHECK-LABEL:   func @basic(
13// CHECK-SAME:                %[[TENSOR:.*]]: tensor<4xf32>) -> tensor<4xf32> {
14// CHECK:           %[[MEMREF:.*]] = memref.buffer_cast %[[TENSOR]] : memref<4xf32>
15// CHECK:           %[[RESULT_MEMREF:.*]] = memref.alloc() : memref<4xf32>
16// CHECK:           linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]}
17// CHECK-SAME:      ins(%[[MEMREF]] : memref<4xf32>)
18// CHECK-SAME:      outs(%[[RESULT_MEMREF]] : memref<4xf32>) {
19// CHECK:           ^bb0(%[[RESULT1:.*]]: f32, %[[UNUSED:.*]]: f32):
20// CHECK:             %[[DIM1:.*]] = math.exp %[[RESULT1]] : f32
21// CHECK:             linalg.yield %[[DIM1]] : f32
22// CHECK:           }
23// CHECK:           %[[RESULT:.*]] = memref.tensor_load %[[RESULT_MEMREF]] : memref<4xf32>
24// CHECK:           return %[[RESULT]] : tensor<4xf32>
25func @basic(%arg0: tensor<4xf32>) -> tensor<4xf32> {
26    %0 = linalg.generic {
27      indexing_maps = [#map0, #map0],
28      iterator_types = ["parallel"]
29    } ins(%arg0 : tensor<4xf32>)
30      outs(%arg0 : tensor<4xf32>) {
31      ^bb0(%gen_arg1: f32, %out: f32):
32        %tmp1 = math.exp %gen_arg1 : f32
33        linalg.yield %tmp1 : f32
34    } -> tensor<4xf32>
35    return %0 : tensor<4xf32>
36}
37
38
39// -----
40
41#map0 = affine_map<(d0) -> (d0)>
42
43// Same as above but with linalg.init_tensor op.
44
45// CHECK: #map = affine_map<(d0) -> (d0)>
46// CHECK-LABEL: func @init_tensor(
47// CHECK-SAME:      %[[IN:.*]]: tensor<?xf32>, %[[SIZE:.*]]: index)
48// CHECK:         %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
49// CHECK:         %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
50// CHECK:         linalg.generic
51// CHECK-SAME:    ins(%[[MEMREF]] : memref<?xf32>)
52// CHECK-SAME:    outs(%[[OUT_BUF]] : memref<?xf32>) {
53func @init_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> {
54  %init = linalg.init_tensor [%size] : tensor<?xf32>
55  %0 = linalg.generic {
56    indexing_maps = [#map0, #map0],
57    iterator_types = ["parallel"]
58  } ins(%in : tensor<?xf32>)
59    outs(%init : tensor<?xf32>) {
60    ^bb0(%gen_arg1: f32, %out: f32):
61      %tmp1 = math.exp %gen_arg1 : f32
62      linalg.yield %tmp1 : f32
63  } -> tensor<?xf32>
64  return %0 : tensor<?xf32>
65}
66
67
68// -----
69
70#map0 = affine_map<(d0) -> (d0)>
71
72// CHECK-LABEL:   func @multiple_results
73// CHECK:           %[[RESULT0:.*]] = memref.alloc() : memref<4xf32>
74// CHECK:           %[[RESULT1:.*]] = memref.alloc() : memref<4xf32>
75// CHECK:           linalg.generic
76// CHECK-SAME:      ins(%{{.*}} : memref<4xf32>)
77// CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<4xf32>, memref<4xf32>)
78// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
79func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
80    %0, %1 = linalg.generic {
81      indexing_maps = [#map0, #map0, #map0],
82      iterator_types = ["parallel"]
83    } ins(%arg0 : tensor<4xf32>)
84      outs (%arg0, %arg0 : tensor<4xf32>, tensor<4xf32>) {
85      ^bb0(%gen_arg1: f32, %out1: f32, %out2: f32):
86        %tmp1 = math.exp %gen_arg1 : f32
87        linalg.yield %tmp1, %tmp1 : f32, f32
88    } -> tensor<4xf32>, tensor<4xf32>
89    return %0, %1 : tensor<4xf32>, tensor<4xf32>
90}
91
92// -----
93
94#map0 = affine_map<(d0) -> (d0)>
95
96// CHECK-LABEL:   func @multiple_results_indexed
97// CHECK:           %[[RESULT0:.*]] = memref.alloc() : memref<4xi32>
98// CHECK:           %[[RESULT1:.*]] = memref.alloc() : memref<4xi32>
99// CHECK:           linalg.indexed_generic
100// CHECK-SAME:      ins(%{{.*}} : memref<4xi32>)
101// CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<4xi32>, memref<4xi32>)
102// CHECK-NEXT: ^bb0(%{{.*}}: index, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32):
103func @multiple_results_indexed(%arg0: tensor<4xi32>)
104        -> (tensor<4xi32>, tensor<4xi32>) {
105    %0, %1 = linalg.indexed_generic {
106      indexing_maps = [#map0, #map0, #map0],
107      iterator_types = ["parallel"]
108    } ins(%arg0 : tensor<4xi32>)
109      outs (%arg0, %arg0 : tensor<4xi32>, tensor<4xi32>) {
110      ^bb0(%i: index, %gen_arg1: i32, %out1: i32, %out2: i32):
111        %i_i32 = index_cast %i : index to i32
112        %tmp1 = addi %gen_arg1, %i_i32 : i32
113        linalg.yield %tmp1, %tmp1 : i32, i32
114    } -> tensor<4xi32>, tensor<4xi32>
115    return %0, %1 : tensor<4xi32>, tensor<4xi32>
116}
117
118// -----
119
120#map_2d = affine_map<(d0, d1) -> (d0, d1)>
121
122// Check that the allocs properly consider the different shapes of the output
123// operands. The permuted indexing maps translate to different output shapes.
124
125// CHECK-LABEL:   func @dynamic_results(
126// CHECK-SAME:                          %[[ARG:.*]]: tensor<?x?xf32>
127// CHECK:           %[[C0:.*]] = constant 0 : index
128// CHECK:           %[[C1:.*]] = constant 1 : index
129// CHECK:           %[[MEMREF_ARG:.*]] = memref.buffer_cast %[[ARG]] : memref<?x?xf32>
130// CHECK:           %[[DIM0:.*]] = memref.dim %[[ARG]], %[[C0]] : tensor<?x?xf32>
131// CHECK:           %[[DIM1:.*]] = memref.dim %[[ARG]], %[[C1]] : tensor<?x?xf32>
132// CHECK:           %[[RESULT0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
133// CHECK:           %[[RESULT1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
134// CHECK:           linalg.generic
135// CHECK-SAME:      ins(%[[MEMREF_ARG]] : memref<?x?xf32>)
136// CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<?x?xf32>, memref<?x?xf32>)
137func @dynamic_results(%arg0: tensor<?x?xf32>)
138         -> (tensor<?x?xf32>, tensor<?x?xf32>) {
139    %0, %1 = linalg.generic {
140      indexing_maps = [#map_2d, #map_2d, #map_2d],
141      iterator_types = ["parallel", "parallel"]
142    } ins(%arg0 : tensor<?x?xf32>)
143      outs (%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) {
144      ^bb0(%gen_arg1: f32, %out1: f32, %out2: f32):
145        %tmp1 = math.exp %gen_arg1 : f32
146        linalg.yield %tmp1, %tmp1 : f32, f32
147    } -> tensor<?x?xf32>, tensor<?x?xf32>
148    return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32>
149}
150
151// -----
152
153#accesses = [
154  affine_map<(i, j, k) -> (j, i, k)>,
155  affine_map<(i, j, k) -> (i, j)>
156]
157
158#trait = {
159  indexing_maps = #accesses,
160  iterator_types = ["parallel", "parallel", "reduction"]
161}
162
163// Check the bufferization of init tensors.
164
165// CHECK-LABEL:   func @generic_with_init_tensor(
166// CHECK-SAME:                                   %[[ARG0_TENSOR:.*]]: tensor<2x3x4xvector<3x4xi4>>,
167// CHECK-SAME:                                   %[[ARG1_TENSOR:.*]]: tensor<3x2xf32>) -> tensor<3x2xf32> {
168// CHECK:           %[[ARG0_MEMREF:.*]] = memref.buffer_cast %[[ARG0_TENSOR]] : memref<2x3x4xvector<3x4xi4>>
169// CHECK:           %[[ARG1_MEMREF:.*]] = memref.buffer_cast %[[ARG1_TENSOR]] : memref<3x2xf32>
170// CHECK:           %[[INIT_BUFFER:.*]] = memref.alloc() : memref<3x2xf32>
171// CHECK:           linalg.copy(%[[ARG1_MEMREF]], %[[INIT_BUFFER]]) : memref<3x2xf32>, memref<3x2xf32>
172// CHECK:           linalg.generic
173// CHECK-SAME:      ins(%[[ARG0_MEMREF]] : memref<2x3x4xvector<3x4xi4>>)
174// CHECK-SAME:      outs(%[[INIT_BUFFER]] : memref<3x2xf32>) {
175func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>,
176  %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) {
177
178  %0 = linalg.generic #trait
179    ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>)
180   outs(%arg1 : tensor<3x2xf32>) {
181    ^bb(%v0: vector<3x4xi4>, %v1: f32) :
182      linalg.yield %v1 : f32
183  } -> tensor<3x2xf32>
184
185  return %0 : tensor<3x2xf32>
186}
187
188// -----
189
190// CHECK-DAG: #[[$MAP0:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
191// CHECK-DAG: #[[$MAP1:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1 * 2)>
192
193func private @make_index() -> index
194
195// CHECK-LABEL: func @bufferize_subtensor(
196//  CHECK-SAME:   %[[T:[0-9a-z]*]]: tensor<?x?xf32>
197func @bufferize_subtensor(%t : tensor<?x?xf32>) -> (tensor<2x3xf32>, tensor<2x?xf32>) {
198  //      CHECK: %[[IDX:.*]] = call @make_index() : () -> index
199  %i0 = call @make_index() : () -> index
200
201  //      CHECK: %[[M0:.*]] = memref.buffer_cast %[[T]] : memref<?x?xf32>
202  // CHECK-NEXT: %[[A0:.*]] = memref.alloc() : memref<2x3xf32>
203  // CHECK-NEXT: %[[SM0:.*]] = memref.subview %[[M0]][0, 0] [2, 3] [1, 1]
204  // CHECK-SAME:   memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]>
205  // CHECK-NEXT: linalg.copy(%[[SM0]], %[[A0]]) : memref<2x3xf32, #[[$MAP0]]>, memref<2x3xf32>
206  // CHECK-NEXT: %[[RT0:.*]] = memref.tensor_load %[[A0]] : memref<2x3xf32>
207  %st0 = subtensor %t[0, 0][2, 3][1, 1] : tensor<?x?xf32> to tensor<2x3xf32>
208
209  //      CHECK: %[[M1:.*]] = memref.buffer_cast %[[T]] : memref<?x?xf32>
210  // CHECK-NEXT: %[[A1:.*]] = memref.alloc(%[[IDX]]) : memref<2x?xf32>
211  // CHECK-NEXT: %[[SM1:.*]] = memref.subview %[[M1]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2]
212  // CHECK-SAME:   memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]>
213  // CHECK-NEXT: linalg.copy(%[[SM1]], %[[A1]]) : memref<2x?xf32, #[[$MAP1]]>, memref<2x?xf32>
214  // CHECK-NEXT: %[[RT1:.*]] = memref.tensor_load %[[A1]] : memref<2x?xf32>
215  %st1 = subtensor %t[0, %i0][2, %i0][1, 2] : tensor<?x?xf32> to tensor<2x?xf32>
216
217  // CHECK-NEXT: return %[[RT0]], %[[RT1]]
218  return %st0, %st1 : tensor<2x3xf32>, tensor<2x?xf32>
219}
220
221// -----
222
223// CHECK-DAG: #[[$MAP0:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
224// CHECK-DAG: #[[$MAP1:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1 * 2)>
225
226func private @make_index() -> index
227
228// CHECK-LABEL: func @bufferize_subtensor_insert(
229//  CHECK-SAME:   %[[T:[0-9a-z]*]]: tensor<?x?xf32>
230//  CHECK-SAME:   %[[ST0:[0-9a-z]*]]: tensor<2x3xf32>
231//  CHECK-SAME:   %[[ST1:[0-9a-z]*]]: tensor<2x?xf32>
232func @bufferize_subtensor_insert(%t : tensor<?x?xf32>, %st0 : tensor<2x3xf32>, %st1 : tensor<2x?xf32>) ->
233    (tensor<?x?xf32>, tensor<?x?xf32>) {
234  %c0 = constant 0 : index
235  %c1 = constant 1 : index
236  // CHECK-NEXT: %[[C0:.*]] = constant 0 : index
237  // CHECK-NEXT: %[[C1:.*]] = constant 1 : index
238  %i0 = call @make_index() : () -> index
239  // CHECK: %[[IDX:.*]] = call @make_index() : () -> index
240
241
242  // CHECK-DAG: %[[M0:.*]] = memref.buffer_cast %[[T]] : memref<?x?xf32>
243  // CHECK-DAG: %[[SM0:.*]] = memref.buffer_cast %[[ST0]] : memref<2x3xf32>
244  // CHECK-NEXT: %[[DIM0:.*]] = memref.dim %[[T]], %[[C0]] : tensor<?x?xf32>
245  // CHECK-NEXT: %[[DIM1:.*]] = memref.dim %[[T]], %[[C1]] : tensor<?x?xf32>
246  // CHECK-NEXT: %[[M0_COPY:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
247  // CHECK-NEXT: linalg.copy(%[[M0]], %[[M0_COPY]]) : memref<?x?xf32>, memref<?x?xf32>
248  // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[M0_COPY]][0, 0] [2, 3] [1, 1]
249  // CHECK-SAME:   memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]>
250  // CHECK-NEXT: linalg.copy(%[[SM0]], %[[SUBVIEW0]]) : memref<2x3xf32>, memref<2x3xf32, #[[$MAP0]]>
251  // CHECK-NEXT: %[[RT0:.*]] = memref.tensor_load %[[M0_COPY]] : memref<?x?xf32>
252  %t0 = subtensor_insert %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor<?x?xf32>
253
254  //  CHECK-DAG: %[[M1:.*]] = memref.buffer_cast %[[T]] : memref<?x?xf32>
255  //  CHECK-DAG: %[[SM1:.*]] = memref.buffer_cast %[[ST1]] : memref<2x?xf32>
256  // CHECK-NEXT: %[[M1_COPY:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
257  // CHECK-NEXT: linalg.copy(%[[M1]], %[[M1_COPY]]) : memref<?x?xf32>, memref<?x?xf32>
258  // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[M1_COPY]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2]
259  // CHECK-SAME:   memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]>
260  // CHECK-NEXT: linalg.copy(%[[SM1]], %[[SUBVIEW1]]) : memref<2x?xf32>, memref<2x?xf32, #[[$MAP1]]>
261  // CHECK-NEXT: %[[RT1:.*]] = memref.tensor_load %[[M1_COPY]] : memref<?x?xf32>
262  %t1 = subtensor_insert %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor<?x?xf32>
263
264  //     CHECK: return %[[RT0]], %[[RT1]]
265  return %t0, %t1: tensor<?x?xf32>, tensor<?x?xf32>
266}
267
268// -----
269
270// CHECK-LABEL: func @bufferize_fill(
271// CHECK-SAME:    %[[IN:.*]]: tensor<?xf32>
272func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
273  %c0 = constant 0.0 : f32
274  // CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
275  // CHECK: linalg.fill(%[[MEMREF]], %cst) : memref<?xf32>, f32
276  // CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[MEMREF]] : memref<?xf32>
277  // CHECK: return %[[TENSOR]]
278  %0 = linalg.fill(%arg0, %c0) : tensor<?xf32>, f32 -> tensor<?xf32>
279  return %0 : tensor<?xf32>
280}
281