1// RUN: mlir-opt %s -tensor-bufferize -cse | FileCheck %s
2
3// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
4// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 20 + s0 + d1)>
5// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 140 + d1 * 20 + d2 * 5 + d3 + s0)>
6// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0 + 1)>
7// CHECK-DAG: #[[$MAP4:.*]] = affine_map<() -> (1)>
8// CHECK-DAG: #[[$MAP5:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1)>
9
10// CHECK-LABEL:   func @dim(
11// CHECK-SAME:              %[[TENSOR:.*]]: tensor<f32>,
12// CHECK-SAME:              %[[INDEX:.*]]: index) -> index {
13// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<f32>
14// CHECK:           %[[EXTENT:.*]] = memref.dim %[[MEMREF]], %[[INDEX]] : memref<f32>
15// CHECK:           return %[[EXTENT]] : index
16func.func @dim(%arg0: tensor<f32>, %arg1: index) -> index {
17  %0 = tensor.dim %arg0, %arg1 : tensor<f32>
18  return %0 : index
19}
20
21// CHECK-LABEL: func @rank(
22// CHECK-SAME:    %[[TENSOR:.*]]: tensor<*xf32>) -> index {
23// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]]
24// CHECK:           %[[EXTENT:.*]] = memref.rank %[[MEMREF]] : memref<*xf32>
25func.func @rank(%arg0: tensor<*xf32>) -> index {
26  %0 = tensor.rank %arg0 : tensor<*xf32>
27  return %0 : index
28}
29
30// CHECK-LABEL:   func @tensor.cast(
31// CHECK-SAME:                      %[[TENSOR:.*]]: tensor<?xindex>) -> tensor<2xindex> {
32// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]]
33// CHECK:           %[[CASTED:.*]] = memref.cast %[[MEMREF]] : memref<?xindex> to memref<2xindex>
34// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[CASTED]]
35// CHECK:           return %[[RET]] : tensor<2xindex>
36func.func @tensor.cast(%arg0: tensor<?xindex>) -> tensor<2xindex> {
37  %0 = tensor.cast %arg0 : tensor<?xindex> to tensor<2xindex>
38  return %0 : tensor<2xindex>
39}
40
41// CHECK-LABEL:   func @tensor.cast_from_unranked(
42// CHECK-SAME:                                    %[[TENSOR:.*]]: tensor<*xf32>) -> tensor<2xf32> {
43// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<*xf32>
44// CHECK:           %[[CASTED_MEMREF:.*]] = memref.cast %[[MEMREF]] : memref<*xf32> to memref<2xf32>
45// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[CASTED_MEMREF]] : memref<2xf32>
46// CHECK:           return %[[RET]] : tensor<2xf32>
47func.func @tensor.cast_from_unranked(%arg0: tensor<*xf32>) -> tensor<2xf32> {
48  %0 = tensor.cast %arg0 : tensor<*xf32> to tensor<2xf32>
49  return %0 : tensor<2xf32>
50}
51
52// CHECK-LABEL:   func @tensor.cast_to_unranked(
53// CHECK-SAME:                                  %[[TENSOR:.*]]: tensor<2xf32>) -> tensor<*xf32> {
54// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<2xf32>
55// CHECK:           %[[CASTED_MEMREF:.*]] = memref.cast %[[MEMREF]] : memref<2xf32> to memref<*xf32>
56// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[CASTED_MEMREF]] : memref<*xf32>
57// CHECK:           return %[[RET]] : tensor<*xf32>
58func.func @tensor.cast_to_unranked(%arg0: tensor<2xf32>) -> tensor<*xf32> {
59  %0 = tensor.cast %arg0 : tensor<2xf32> to tensor<*xf32>
60  return %0 : tensor<*xf32>
61}
62
63// CHECK-LABEL:   func @tensor.extract(
64// CHECK-SAME:                  %[[TENSOR:.*]]: tensor<?xf32>,
65// CHECK-SAME:                  %[[IDX:.*]]: index) -> f32 {
66// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<?xf32>
67// CHECK:           %[[RET:.*]] = memref.load %[[MEMREF]][%[[IDX]]] : memref<?xf32>
68// CHECK:           return %[[RET]] : f32
69// CHECK:         }
70func.func @tensor.extract(%arg0: tensor<?xf32>, %arg1: index) -> f32 {
71  %0 = tensor.extract %arg0[%arg1] : tensor<?xf32>
72  return %0 : f32
73}
74
75// CHECK-LABEL:   func @tensor.from_elements_0d(
76// CHECK-SAME:        %[[ELEM0:.*]]: index) -> tensor<index> {
77// CHECK:           %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<index>
78// CHECK:           store %[[ELEM0]], %[[MEMREF]]
79// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
80// CHECK:           return %[[RET]] : tensor<index>
81func.func @tensor.from_elements_0d(%arg0: index) -> tensor<index> {
82  %0 = tensor.from_elements %arg0 : tensor<index>
83  return %0 : tensor<index>
84}
85
86// CHECK-LABEL:   func @tensor.from_elements_1d(
87// CHECK-SAME:                               %[[ELEM0:.*]]: index,
88// CHECK-SAME:                               %[[ELEM1:.*]]: index) -> tensor<2xindex> {
89// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
90// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
91// CHECK-DAG:       %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<2xindex>
92// CHECK:           store %[[ELEM0]], %[[MEMREF]][%[[C0]]]
93// CHECK:           store %[[ELEM1]], %[[MEMREF]][%[[C1]]]
94// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
95// CHECK:           return %[[RET]] : tensor<2xindex>
96func.func @tensor.from_elements_1d(%arg0: index, %arg1: index) -> tensor<2xindex> {
97  %0 = tensor.from_elements %arg0, %arg1 : tensor<2xindex>
98  return %0 : tensor<2xindex>
99}
100
101// CHECK-LABEL: func @tensor.from_elements_2d(
102// CHECK-SAME:      %[[ELEM0:.*]]: index, %[[ELEM1:.*]]: index)
103// CHECK-SAME:      -> tensor<3x2xindex> {
104// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
105// CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
106// CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
107// CHECK-DAG:     %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<3x2xindex>
108// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C0]]]
109// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C0]], %[[C1]]]
110// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C1]], %[[C0]]]
111// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C1]]]
112// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C2]], %[[C0]]]
113// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C2]], %[[C1]]]
114// CHECK:         %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
115// CHECK:         return %[[RET]] : tensor<3x2xindex>
116func.func @tensor.from_elements_2d(%arg0: index, %arg1: index) -> tensor<3x2xindex> {
117  %0 = tensor.from_elements %arg0, %arg1, %arg0, %arg1, %arg0, %arg1
118         : tensor<3x2xindex>
119  return %0 : tensor<3x2xindex>
120}
121
122// CHECK-LABEL: func @tensor.from_elements_3d(
123//  CHECK-SAME:     %[[F0:.*]]: f32
124
125// CHECK-DAG: %[[F1:.*]] = arith.constant 1.0{{0+}}e+00
126// CHECK-DAG: %[[F2:.*]] = arith.constant 2.0
127// CHECK-DAG: %[[F3:.*]] = arith.constant 3.0
128// CHECK-DAG: %[[F4:.*]] = arith.constant 4.0
129// CHECK-DAG: %[[F5:.*]] = arith.constant 5.0
130// CHECK-DAG: %[[F6:.*]] = arith.constant 6.0
131// CHECK-DAG: %[[F7:.*]] = arith.constant 7.0
132// CHECK-DAG: %[[F8:.*]] = arith.constant 8.0
133// CHECK-DAG: %[[F9:.*]] = arith.constant 9.0
134// CHECK-DAG: %[[F10:.*]] = arith.constant 1.0{{0+}}e+01
135// CHECK-DAG: %[[F11:.*]] = arith.constant 1.1{{0+}}e+01
136
137// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
138// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
139// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
140
141// CHECK-DAG: %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<3x2x2xf32>
142
143// CHECK: store %[[F0]], %[[MEMREF]][%[[C0]], %[[C0]], %[[C0]]]
144// CHECK: store %[[F1]], %[[MEMREF]][%[[C0]], %[[C0]], %[[C1]]]
145// CHECK: store %[[F2]], %[[MEMREF]][%[[C0]], %[[C1]], %[[C0]]]
146// CHECK: store %[[F3]], %[[MEMREF]][%[[C0]], %[[C1]], %[[C1]]]
147// CHECK: store %[[F4]], %[[MEMREF]][%[[C1]], %[[C0]], %[[C0]]]
148// CHECK: store %[[F5]], %[[MEMREF]][%[[C1]], %[[C0]], %[[C1]]]
149// CHECK: store %[[F6]], %[[MEMREF]][%[[C1]], %[[C1]], %[[C0]]]
150// CHECK: store %[[F7]], %[[MEMREF]][%[[C1]], %[[C1]], %[[C1]]]
151// CHECK: store %[[F8]], %[[MEMREF]][%[[C2]], %[[C0]], %[[C0]]]
152// CHECK: store %[[F9]], %[[MEMREF]][%[[C2]], %[[C0]], %[[C1]]]
153// CHECK: store %[[F10]], %[[MEMREF]][%[[C2]], %[[C1]], %[[C0]]]
154// CHECK: store %[[F11]], %[[MEMREF]][%[[C2]], %[[C1]], %[[C1]]]
155
156// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
157// CHECK: return %[[RET]] : tensor<3x2x2xf32>
158func.func @tensor.from_elements_3d(%f0 : f32) -> tensor<3x2x2xf32> {
159  %f1 = arith.constant 1.0 : f32
160  %f2 = arith.constant 2.0 : f32
161  %f3 = arith.constant 3.0 : f32
162  %f4 = arith.constant 4.0 : f32
163  %f5 = arith.constant 5.0 : f32
164  %f6 = arith.constant 6.0 : f32
165  %f7 = arith.constant 7.0 : f32
166  %f8 = arith.constant 8.0 : f32
167  %f9 = arith.constant 9.0 : f32
168  %f10 = arith.constant 10.0 : f32
169  %f11 = arith.constant 11.0 : f32
170  %0 = tensor.from_elements %f0,%f1,%f2,%f3,%f4,%f5,%f6,%f7,%f8,%f9,%f10,%f11
171         : tensor<3x2x2xf32>
172  return %0 : tensor<3x2x2xf32>
173}
174
175// CHECK-LABEL:   func @tensor.generate(
176// CHECK-SAME:                                       %[[ARG:.*]]: tensor<*xf32>,
177// CHECK-SAME:                                       %[[DYNAMIC_EXTENT:.*]]: index) -> tensor<?xindex> {
178// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
179// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
180// CHECK-DAG:       %[[CASTED:.*]] = bufferization.to_memref %[[ARG]] : memref<*xf32>
181// CHECK-DAG:       %[[MEMREF:.*]] = memref.alloc(%[[DYNAMIC_EXTENT]]) {{.*}} : memref<?xindex>
182// CHECK:           scf.parallel (%[[I:.*]]) = (%[[C0]]) to (%[[DYNAMIC_EXTENT]]) step (%[[C1]]) {
183// CHECK:             %[[ELEM:.*]] = memref.dim %[[CASTED]], %[[I]] : memref<*xf32>
184// CHECK:             store %[[ELEM]], %[[MEMREF]][%[[I]]] : memref<?xindex>
185// CHECK:             scf.yield
186// CHECK:           }
187// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<?xindex>
188// CHECK:           return %[[RET]] : tensor<?xindex>
189// CHECK:         }
190func.func @tensor.generate(%arg: tensor<*xf32>, %dynamic_extent: index) -> tensor<?xindex> {
191  %result = tensor.generate %dynamic_extent {
192  ^bb0(%i : index):
193    %elem = tensor.dim %arg, %i : tensor<*xf32>
194    tensor.yield %elem : index
195  } : tensor<?xindex>
196  return %result : tensor<?xindex>
197}
198
199// Additional test that checks the logic for intermixed static and dynamic
200// extents.
201//
202// CHECK-LABEL:   func @tensor.generate_static_and_dynamic(
203// CHECK-SAME:        %[[DYNAMIC_EXTENT:.*]]: index) -> tensor<16x?xindex> {
204// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
205// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
206// CHECK-DAG:       %[[C16:.*]] = arith.constant 16 : index
207// CHECK-DAG:       %[[MEMREF:.*]] = memref.alloc(%[[DYNAMIC_EXTENT]]) {{.*}} : memref<16x?xindex>
208// CHECK:           scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]]) to (%[[C16]], %[[DYNAMIC_EXTENT]]) step (%[[C1]], %[[C1]]) {
209// CHECK:             %[[VAL_7:.*]] = arith.addi %[[I]], %[[J]] : index
210// CHECK:             store %[[VAL_7]], %[[MEMREF]][%[[I]], %[[J]]] : memref<16x?xindex>
211// CHECK:             scf.yield
212// CHECK:           }
213// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<16x?xindex>
214// CHECK:           return %[[RET]] : tensor<16x?xindex>
215// CHECK:         }
216func.func @tensor.generate_static_and_dynamic(%arg0: index) -> tensor<16x?xindex> {
217  %result = tensor.generate %arg0 {
218  ^bb0(%i: index, %j: index):
219    %sum = arith.addi %i, %j : index
220    tensor.yield %sum : index
221  } : tensor<16x?xindex>
222  return %result : tensor<16x?xindex>
223}
224
225// CHECK-LABEL: func @tensor.generate_unknown_ops_in_body
226func.func @tensor.generate_unknown_ops_in_body(%arg0: index) -> tensor<?xindex> {
227  // CHECK-NOT: tensor.generate
228  %tensor = tensor.generate %arg0 {
229  ^bb0(%iv: index):
230    // CHECK: test.source
231    %0 = "test.source"() : () -> index
232    tensor.yield %0 : index
233  } : tensor<?xindex>
234  return %tensor : tensor<?xindex>
235}
236
237// CHECK-LABEL: func @tensor.extract_slice(
238//  CHECK-SAME:     %[[t1:.*]]: tensor<?x?xf32>, %[[idx1:.*]]: index, %[[idx2:.*]]: index
239func.func @tensor.extract_slice(
240    %t1: tensor<?x?xf32>, %idx1: index, %idx2: index) -> tensor<?x10xf32> {
241  // CHECK: %[[m:.*]] = bufferization.to_memref %[[t1]] : memref<?x?xf32>
242  // CHECK: %[[r:.*]] = memref.subview %[[m]][5, %[[idx2]]] [%[[idx1]], 10] [1, 1] : memref<?x?xf32> to memref<?x10xf32, #[[$MAP0]]>
243  %0 = tensor.extract_slice %t1[5, %idx2][%idx1, 10][1, 1]
244      : tensor<?x?xf32> to tensor<?x10xf32>
245  // CHECK: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
246  // CHECK: return %[[r_tensor]]
247  return %0 : tensor<?x10xf32>
248}
249
250// CHECK-LABEL: func @tensor.extract_slice_rank_reducing(
251//  CHECK-SAME:     %[[t1:.*]]: tensor<?x10x?xf32>, %[[idx1:.*]]: index,
252//  CHECK-SAME:     %[[idx2:.*]]: index
253func.func @tensor.extract_slice_rank_reducing(
254    %t1: tensor<?x10x?xf32>, %idx1: index, %idx2: index) -> tensor<?x15xf32> {
255  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10x?xf32>
256  // CHECK: %[[r:.*]] = memref.subview %[[m1]][5, %[[idx1]], 10] [%[[idx2]], 1, 15] [1, 1, 1] : memref<?x10x?xf32> to memref<?x15xf32, #[[$MAP0]]>
257  %0 = tensor.extract_slice %t1[5, %idx1, 10][%idx2, 1, 15][1, 1, 1]
258      : tensor<?x10x?xf32> to tensor<?x15xf32>
259  // CHECK: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
260  // CHECK: return %[[r_tensor]]
261  return %0 : tensor<?x15xf32>
262}
263
264// CHECK-LABEL: func @tensor.insert_slice(
265//  CHECK-SAME:     %[[t1:.*]]: tensor<?x?xf32>, %[[t2:.*]]: tensor<?x10xf32>,
266//  CHECK-SAME:     %[[idx1:.*]]: index, %[[idx2:.*]]: index
267func.func @tensor.insert_slice(%t1: tensor<?x?xf32>, %t2: tensor<?x10xf32>,
268                          %idx1: index, %idx2: index) -> tensor<?x?xf32> {
269  // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
270  // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
271  // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x?xf32>
272  // CHECK-DAG: %[[m2:.*]] = bufferization.to_memref %[[t2]] : memref<?x10xf32>
273  // CHECK-DAG: %[[dim0:.*]] = memref.dim %[[m1]], %[[c0]]
274  // CHECK-DAG: %[[dim1:.*]] = memref.dim %[[m1]], %[[c1]]
275  //     CHECK: %[[alloc:.*]] = memref.alloc(%[[dim0]], %[[dim1]])
276  //     CHECK: memref.copy %[[m1]], %[[alloc]]
277  //     CHECK: %[[subview:.*]] = memref.subview %[[alloc]][%[[idx1]], 5] [%[[idx2]], 10] [1, 1]
278  //     CHECK: memref.copy %[[m2]], %[[subview]]
279  %0 = tensor.insert_slice %t2 into %t1[%idx1, 5][%idx2, 10][1, 1]
280      : tensor<?x10xf32> into tensor<?x?xf32>
281
282  //     CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
283  //     CHECK: return %[[r]]
284  return %0 : tensor<?x?xf32>
285}
286
287// CHECK-LABEL: func @tensor.insert(
288//  CHECK-SAME:     %[[t1:.*]]: tensor<5xf32>, %[[idx1:.*]]: index,
289//  CHECK-SAME:     %[[f:.*]]: f32
290func.func @tensor.insert(%t1: tensor<5xf32>, %idx1: index, %f: f32) -> tensor<5xf32> {
291  // CHECK-DAG: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
292  // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<5xf32>
293  // CHECK: memref.copy %[[m1]], %[[alloc]]
294  // CHECK: memref.store %[[f]], %[[alloc]][%[[idx1]]]
295  %0 = tensor.insert %f into %t1[%idx1] : tensor<5xf32>
296
297  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
298  // CHECK: return %[[r]]
299  return %0 : tensor<5xf32>
300}
301
302// CHECK-LABEL: func @tensor.expand_shape(
303//  CHECK-SAME:     %[[t1:.*]]: tensor<?x10xf32>
304func.func @tensor.expand_shape(%t1: tensor<?x10xf32>) -> tensor<2x?x10xf32> {
305  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10xf32>
306  // CHECK: %[[expanded:.*]] = memref.expand_shape %[[m1]] [
307  // CHECK-SAME: [0, 1], [2]] : memref<?x10xf32> into memref<2x?x10xf32>
308  %0 = tensor.expand_shape %t1 [[0, 1], [2]]
309      : tensor<?x10xf32> into tensor<2x?x10xf32>
310
311  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[expanded]]
312  // CHECK: return %[[r]]
313  return %0 : tensor<2x?x10xf32>
314}
315
316// CHECK-LABEL: func @tensor.expand_shape_of_slice(
317//  CHECK-SAME:     %[[t1:.*]]: tensor<?x20xf32>
318func.func @tensor.expand_shape_of_slice(
319    %t1: tensor<?x20xf32>, %o1: index, %s1: index) -> tensor<?x7x2x5xf32> {
320  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x20xf32>
321  // CHECK: %[[subview:.*]] = memref.subview %[[m1]][%{{.*}}, 5] [%{{.*}}, 10] [1, 1] : memref<?x20xf32> to memref<?x10xf32, #[[$MAP1]]>
322  %0 = tensor.extract_slice %t1[%o1, 5][%s1, 10][1, 1] :
323      tensor<?x20xf32> to tensor<?x10xf32>
324  // CHECK: %[[expanded:.*]] = memref.expand_shape %[[subview]] [
325  // CHECK-SAME: [0, 1], [2, 3]] : memref<?x10xf32, #[[$MAP1]]> into memref<?x7x2x5xf32, #[[$MAP2]]>
326  %1 = tensor.expand_shape %0 [[0, 1], [2, 3]] :
327      tensor<?x10xf32> into tensor<?x7x2x5xf32>
328  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[expanded]]
329  // CHECK: return %[[r]]
330  return %1 : tensor<?x7x2x5xf32>
331}
332
333// CHECK-LABEL: func @tensor.expand_shape_of_slice2(
334//  CHECK-SAME:     %[[t1:.*]]: tensor<1x2xf32>
335func.func @tensor.expand_shape_of_slice2(%t1: tensor<1x2xf32>) -> tensor<1xf32> {
336  // CHECK: memref.subview {{.*}} : memref<1x2xf32> to memref<1x1xf32, #[[$MAP5]]>
337  %0 = tensor.extract_slice %t1[0, 0][1, 1][1, 1] : tensor<1x2xf32> to tensor<1x1xf32>
338  // CHECK: memref.collapse_shape %{{.*}} [
339  // CHECK-SAME: [0, 1]] : memref<1x1xf32, #[[$MAP5]]> into memref<1xf32>
340  %1 = tensor.collapse_shape %0 [[0, 1]] : tensor<1x1xf32> into tensor<1xf32>
341  return %1 : tensor<1xf32>
342}
343
344// CHECK-LABEL: func @tensor.collapse_shape(
345//  CHECK-SAME:     %[[t1:.*]]: tensor<2x?x?xf32>
346func.func @tensor.collapse_shape(%t1: tensor<2x?x?xf32>) -> tensor<?x?xf32> {
347  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<2x?x?xf32>
348  // CHECK: %[[collapsed:.*]] = memref.collapse_shape %[[m1]] [
349  // CHECK-SAME: [0, 1], [2]] : memref<2x?x?xf32> into memref<?x?xf32>
350  %0 = tensor.collapse_shape %t1 [[0, 1], [2]]
351      : tensor<2x?x?xf32> into tensor<?x?xf32>
352
353  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[collapsed]]
354  // CHECK: return %[[r]]
355  return %0 : tensor<?x?xf32>
356}
357
358// CHECK-LABEL: func @tensor.collapse_shape_to_scalar(
359//  CHECK-SAME:     %[[t1:.*]]: tensor<1x1x1xf32>
360func.func @tensor.collapse_shape_to_scalar(%t1: tensor<1x1x1xf32>) -> tensor<f32> {
361  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<1x1x1xf32>
362  // CHECK: %[[collapsed:.*]] = memref.collapse_shape %[[m1]] [] : memref<1x1x1xf32> into memref<f32>
363  %0 = tensor.collapse_shape %t1 []
364      : tensor<1x1x1xf32> into tensor<f32>
365
366  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[collapsed]]
367  // CHECK: return %[[r]]
368  return %0 : tensor<f32>
369}
370
371// CHECK-LABEL: func @tensor.collapse_shape_of_slice(
372func.func @tensor.collapse_shape_of_slice(%arg0: tensor<2xi32>) -> tensor<i32> {
373  // CHECK: memref.subview %{{.*}}[1] [1] [1] : memref<2xi32> to memref<1xi32, #[[$MAP3]]>
374  %0 = tensor.extract_slice %arg0[1] [1] [1] : tensor<2xi32> to tensor<1xi32>
375  // CHECK: memref.collapse_shape %{{.*}} [] : memref<1xi32, #[[$MAP3]]> into memref<i32, #[[$MAP4]]>
376  %1 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
377  return %1 : tensor<i32>
378}
379
380// CHECK-LABEL: func @tensor.collapse_shape_of_slice2(
381func.func @tensor.collapse_shape_of_slice2(
382    %arg0: tensor<?x?x?x?xi64>, %o1: index, %o2: index, %o3: index, %o4: index)
383    -> tensor<87x63648xi64> {
384  // CHECK: %[[subview:.*]] = memref.subview %{{.*}} : memref<?x?x?x?xi64> to memref<87x78x68x12xi64, #{{.*}}>
385  %0 = tensor.extract_slice %arg0[%o1, %o2, %o3, %o4] [87, 78, 68, 12] [1, 1, 1, 1] : tensor<?x?x?x?xi64> to tensor<87x78x68x12xi64>
386
387  // This memref may not be collapsible, so the buffer must be copied to get rid
388  // of the layout map.
389  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<87x78x68x12xi64>
390  // CHECK: memref.copy %[[subview]], %[[alloc]]
391  // CHECK: memref.collapse_shape %[[alloc]] [
392  // CHECK-SAME: [0], [1, 2, 3]] : memref<87x78x68x12xi64> into memref<87x63648xi64>
393  %1 = tensor.collapse_shape %0 [[0], [1, 2, 3]] : tensor<87x78x68x12xi64> into tensor<87x63648xi64>
394  return %1 : tensor<87x63648xi64>
395}
396