xref: /llvm-project/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir (revision 956c0707d9098499a2682297b71f46b0a562eed9)
1// RUN: mlir-opt --split-input-file --tosa-to-tensor %s -o -| FileCheck %s
2
3// -----
4
5// CHECK-LABEL: test_reshape_0d_same_s2s_explicit
6// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<f32>
7// CHECK: return %[[ARG_0]] : tensor<f32>
8func.func @test_reshape_0d_same_s2s_explicit(%arg0: tensor<f32>) -> tensor<f32> {
9  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64>} : (tensor<f32>) -> tensor<f32>
10  return %0 : tensor<f32>
11}
12
13// -----
14
15// CHECK-LABEL: test_reshape_0d_up_s2d_auto
16// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<f32>
17// CHECK: %[[VAL_0:.*]] = tensor.expand_shape %[[ARG_0]] [] output_shape [1] : tensor<f32> into tensor<1xf32>
18// CHECK: %[[VAL_1:.*]] = tensor.cast %[[VAL_0]] : tensor<1xf32> to tensor<?xf32>
19// CHECK: return %[[VAL_1]] : tensor<?xf32>
20func.func @test_reshape_0d_up_s2d_auto(%arg0: tensor<f32>) -> tensor<?xf32> {
21  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1>} : (tensor<f32>) -> tensor<?xf32>
22  return %0 : tensor<?xf32>
23}
24
25// -----
26
27// CHECK-LABEL: test_reshape_0d_up_s2d_explicit
28// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<f32>
29// CHECK: %[[VAL_0:.*]] = tensor.expand_shape %[[ARG_0]] [] output_shape [1] : tensor<f32> into tensor<1xf32>
30// CHECK: %[[VAL_1:.*]] = tensor.cast %[[VAL_0]] : tensor<1xf32> to tensor<?xf32>
31// CHECK: return %[[VAL_1]] : tensor<?xf32>
32func.func @test_reshape_0d_up_s2d_explicit(%arg0: tensor<f32>) -> tensor<?xf32> {
33  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 1>} : (tensor<f32>) -> tensor<?xf32>
34  return %0 : tensor<?xf32>
35}
36
37// -----
38
39// CHECK-LABEL: test_reshape_0d_up_s2s_auto
40// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<f32>
41// CHECK: %[[VAL_0:.*]] = tensor.expand_shape %[[ARG_0]] [] output_shape [1] : tensor<f32> into tensor<1xf32>
42// CHECK: return %[[VAL_0]] : tensor<1xf32>
43func.func @test_reshape_0d_up_s2s_auto(%arg0: tensor<f32>) -> tensor<1xf32> {
44  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1>} : (tensor<f32>) -> tensor<1xf32>
45  return %0 : tensor<1xf32>
46}
47
48// -----
49
50// CHECK-LABEL: test_reshape_0d_up_s2s_explicit
51// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<f32>
52// CHECK: %[[VAL_0:.*]] = tensor.expand_shape %[[ARG_0]] [] output_shape [1] : tensor<f32> into tensor<1xf32>
53// CHECK: return %[[VAL_0]] : tensor<1xf32>
54func.func @test_reshape_0d_up_s2s_explicit(%arg0: tensor<f32>) -> tensor<1xf32> {
55  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 1>} : (tensor<f32>) -> tensor<1xf32>
56  return %0 : tensor<1xf32>
57}
58
59// -----
60
61// CHECK-LABEL: test_reshape_1d_down_d2s_explicit
62// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?xf32>
63// CHECK: %[[VAL_0:.*]] = tensor.cast %[[ARG_0]] : tensor<?xf32> to tensor<1xf32>
64// CHECK: %[[VAL_1:.*]] = tensor.collapse_shape %[[VAL_0]] [] : tensor<1xf32> into tensor<f32>
65// CHECK: return %[[VAL_1]] : tensor<f32>
66func.func @test_reshape_1d_down_d2s_explicit(%arg0: tensor<?xf32>) -> tensor<f32> {
67  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64>} : (tensor<?xf32>) -> tensor<f32>
68  return %0 : tensor<f32>
69}
70
71// -----
72
73// CHECK-LABEL: test_reshape_1d_down_s2s_explicit
74// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<1xf32>
75// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] [] : tensor<1xf32> into tensor<f32>
76// CHECK: return %[[VAL_0]] : tensor<f32>
77func.func @test_reshape_1d_down_s2s_explicit(%arg0: tensor<1xf32>) -> tensor<f32> {
78  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64>} : (tensor<1xf32>) -> tensor<f32>
79  return %0 : tensor<f32>
80}
81
82// -----
83
84// CHECK-LABEL: test_reshape_1d_up_d2d_auto
85// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?xf32>
86// CHECK: %[[C0:.*]] = arith.constant 0 : index
87// CHECK: %[[DIM:.*]] = tensor.dim %arg0, %[[C0]] : tensor<?xf32>
88// CHECK: %[[C2:.*]] = arith.constant 2 : index
89// CHECK: %[[VAL_0:.*]] = arith.divsi %[[DIM]], %[[C2]] : index
90// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[ARG_0]] {{\[\[}}0, 1]] output_shape [2, %[[VAL_0]]] : tensor<?xf32> into tensor<2x?xf32>
91// CHECK: return %[[EXPANDED]] : tensor<2x?xf32>
92func.func @test_reshape_1d_up_d2d_auto(%arg0: tensor<?xf32>) -> tensor<2x?xf32> {
93  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, -1>} : (tensor<?xf32>) -> tensor<2x?xf32>
94  return %0 : tensor<2x?xf32>
95}
96
97// -----
98
99// CHECK-LABEL: test_reshape_1d_up_s2s_explicit
100// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<6xf32>
101// CHECK: %[[VAL_0:.*]] = tensor.expand_shape %[[ARG_0]] {{\[\[}}0, 1]] output_shape [2, 3] : tensor<6xf32> into tensor<2x3xf32>
102// CHECK: return %[[VAL_0]] : tensor<2x3xf32>
103func.func @test_reshape_1d_up_s2s_explicit(%arg0: tensor<6xf32>) -> tensor<2x3xf32> {
104  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3>} : (tensor<6xf32>) -> tensor<2x3xf32>
105  return %0 : tensor<2x3xf32>
106}
107
108// -----
109
110// CHECK-LABEL: test_reshape_2d_down_d2d_auto
111// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<2x?xf32>
112// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1]] : tensor<2x?xf32> into tensor<?xf32>
113// CHECK: return %[[VAL_0]] : tensor<?xf32>
114func.func @test_reshape_2d_down_d2d_auto(%arg0: tensor<2x?xf32>) -> tensor<?xf32> {
115  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1>} : (tensor<2x?xf32>) -> tensor<?xf32>
116  return %0 : tensor<?xf32>
117}
118
119// -----
120
121// CHECK-LABEL: test_reshape_2d_down_s2s_explicit
122// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<2x3xf32>
123// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1]] : tensor<2x3xf32> into tensor<6xf32>
124// CHECK: return %[[VAL_0]] : tensor<6xf32>
125func.func @test_reshape_2d_down_s2s_explicit(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {
126  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 6>} : (tensor<2x3xf32>) -> tensor<6xf32>
127  return %0 : tensor<6xf32>
128}
129
130// -----
131
132// CHECK-LABEL: test_reshape_2d_same_d2d_auto
133// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x2xf32>
134// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1]] : tensor<?x2xf32> into tensor<?xf32>
135// CHECK: %[[C0:.*]] = arith.constant 0 : index
136// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
137// CHECK: %[[C2:.*]] = arith.constant 2 : index
138// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C2]] : index
139// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1]] output_shape [2, %[[DIV]]] : tensor<?xf32> into tensor<2x?xf32>
140// CHECK: return %[[EXPANDED]] : tensor<2x?xf32>
141func.func @test_reshape_2d_same_d2d_auto(%arg0: tensor<?x2xf32>) -> tensor<2x?xf32> {
142  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, -1>} : (tensor<?x2xf32>) -> tensor<2x?xf32>
143  return %0 : tensor<2x?xf32>
144}
145
146// -----
147
148// CHECK-LABEL: test_reshape_2d_same_s2d_auto
149// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<2x4xf32>
150// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1]] : tensor<2x4xf32> into tensor<8xf32>
151// CHECK: %[[VAL_1:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1]] output_shape [4, 2] : tensor<8xf32> into tensor<4x2xf32>
152// CHECK: %[[VAL_2:.*]] = tensor.cast %[[VAL_1]] : tensor<4x2xf32> to tensor<?x2xf32>
153// CHECK: return %[[VAL_2]] : tensor<?x2xf32>
154func.func @test_reshape_2d_same_s2d_auto(%arg0: tensor<2x4xf32>) -> tensor<?x2xf32> {
155  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1, 2>} : (tensor<2x4xf32>) -> tensor<?x2xf32>
156  return %0 : tensor<?x2xf32>
157}
158
159// -----
160
161// CHECK-LABEL: test_reshape_2d_same_s2d_explicit
162// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<2x4xf32>
163// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1]] : tensor<2x4xf32> into tensor<8xf32>
164// CHECK: %[[VAL_1:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1]] output_shape [4, 2] : tensor<8xf32> into tensor<4x2xf32>
165// CHECK: %[[VAL_2:.*]] = tensor.cast %[[VAL_1]] : tensor<4x2xf32> to tensor<?x2xf32>
166// CHECK: return %[[VAL_2]] : tensor<?x2xf32>
167func.func @test_reshape_2d_same_s2d_explicit(%arg0: tensor<2x4xf32>) -> tensor<?x2xf32> {
168  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 4, 2>} : (tensor<2x4xf32>) -> tensor<?x2xf32>
169  return %0 : tensor<?x2xf32>
170}
171
172// -----
173
174// CHECK-LABEL: test_reshape_2d_same_s2s_explicit
175// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<3x2xf32>
176// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1]] : tensor<3x2xf32> into tensor<6xf32>
177// CHECK: %[[VAL_1:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1]] output_shape [2, 3] : tensor<6xf32> into tensor<2x3xf32>
178// CHECK: return %[[VAL_1]] : tensor<2x3xf32>
179func.func @test_reshape_2d_same_s2s_explicit(%arg0: tensor<3x2xf32>) -> tensor<2x3xf32> {
180  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3>} : (tensor<3x2xf32>) -> tensor<2x3xf32>
181  return %0 : tensor<2x3xf32>
182}
183
184// -----
185
186// CHECK-LABEL: test_reshape_3d_same_d2d_auto_empty
187// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<3x2x?xf32>
188// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<3x2x?xf32> into tensor<?xf32>
189// CHECK: %[[C0:.*]] = arith.constant 0 : index
190// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
191// CHECK: %[[C0_0:.*]] = arith.constant 0 : index
192// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C0_0]] : index
193// CHECK: %[[VAL_1:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [0, 3, %[[DIV]]] : tensor<?xf32> into tensor<0x3x?xf32>
194// CHECK: %[[VAL_2:.*]] = tensor.cast %[[VAL_1]] : tensor<0x3x?xf32> to tensor<?x?x?xf32>
195// CHECK: return %[[VAL_2]] : tensor<?x?x?xf32>
196func.func @test_reshape_3d_same_d2d_auto_empty(%arg0: tensor<3x2x?xf32>) -> tensor<?x?x?xf32> {
197  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 0, 3, -1>} : (tensor<3x2x?xf32>) -> tensor<?x?x?xf32>
198  return %0 : tensor<?x?x?xf32>
199}
200
201// -----
202
203// CHECK-LABEL: test_reshape_3d_same_d2d_auto
204// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<2x?x?xf32>
205// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<2x?x?xf32> into tensor<?xf32>
206// CHECK: %[[C0:.*]] = arith.constant 0 : index
207// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
208// CHECK: %[[C8:.*]] = arith.constant 8 : index
209// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C8]] : index
210// CHECK: %[[VAL_1:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [2, %[[DIV]], 4] : tensor<?xf32> into tensor<2x?x4xf32>
211// CHECK: %[[VAL_2:.*]] = tensor.cast %[[VAL_1]] : tensor<2x?x4xf32> to tensor<?x?x?xf32>
212// CHECK: return %[[VAL_2]] : tensor<?x?x?xf32>
213func.func @test_reshape_3d_same_d2d_auto(%arg0: tensor<2x?x?xf32>) -> tensor<?x?x?xf32> {
214  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, -1, 4>} : (tensor<2x?x?xf32>) -> tensor<?x?x?xf32>
215  return %0 : tensor<?x?x?xf32>
216}
217
218// -----
219
220// CHECK-LABEL: test_reshape_3d_same_d2d_auto_identity
221// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x3x4xf32>
222// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<?x3x4xf32> into tensor<?xf32>
223// CHECK: %[[C0:.*]] = arith.constant 0 : index
224// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
225// CHECK: %[[C6:.*]] = arith.constant 6 : index
226// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C6]] : index
227// CHECK: %[[VAL_1:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [2, 3, %[[DIV]]] : tensor<?xf32> into tensor<2x3x?xf32>
228// CHECK: return %[[VAL_1]] : tensor<2x3x?xf32>
229func.func @test_reshape_3d_same_d2d_auto_identity(%arg0: tensor<?x3x4xf32>) -> tensor<2x3x?xf32> {
230  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3, -1>} : (tensor<?x3x4xf32>) -> tensor<2x3x?xf32>
231  return %0 : tensor<2x3x?xf32>
232}
233
234// -----
235
236// CHECK-LABEL: test_reshape_3d_same_d2d_explicit_empty
237// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<3x2x?xf32>
238// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<3x2x?xf32> into tensor<?xf32>
239// CHECK: %[[C0:.*]] = arith.constant 0 : index
240// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
241// CHECK: %[[C6:.*]] = arith.constant 6 : index
242// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C6]] : index
243// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [%[[DIV]], 3, 2] : tensor<?xf32> into tensor<?x3x2xf32>
244// CHECK: %[[VAL_2:.*]] = tensor.cast %[[EXPANDED]] : tensor<?x3x2xf32> to tensor<?x?x?xf32>
245// CHECK: return %[[VAL_2]] : tensor<?x?x?xf32>
246func.func @test_reshape_3d_same_d2d_explicit_empty(%arg0: tensor<3x2x?xf32>) -> tensor<?x?x?xf32> {
247  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 0, 3, 2>} : (tensor<3x2x?xf32>) -> tensor<?x?x?xf32>
248  return %0 : tensor<?x?x?xf32>
249}
250
251// -----
252
253// CHECK-LABEL: test_reshape_3d_same_d2d_explicit
254// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
255// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<?x?x?xf32> into tensor<?xf32>
256// CHECK: %[[C0:.*]] = arith.constant 0 : index
257// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
258// CHECK: %[[C12:.*]] = arith.constant 12 : index
259// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C12]] : index
260// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [%[[DIV]], 3, 4] : tensor<?xf32> into tensor<?x3x4xf32>
261// CHECK: %[[VAL_2:.*]] = tensor.cast %[[EXPANDED]] : tensor<?x3x4xf32> to tensor<?x?x?xf32>
262// CHECK: return %[[VAL_2]] : tensor<?x?x?xf32>
263func.func @test_reshape_3d_same_d2d_explicit(%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
264  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3, 4>} : (tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
265  return %0 : tensor<?x?x?xf32>
266}
267
268// -----
269
270// CHECK-LABEL: test_reshape_3d_same_d2d_explicit_identity
271// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x3x4xf32>
272// CHECK: %[[VAL_0:.*]] = tensor.cast %[[ARG_0]] : tensor<?x3x4xf32> to tensor<2x3x?xf32>
273// CHECK: return %[[VAL_0]] : tensor<2x3x?xf32>
274func.func @test_reshape_3d_same_d2d_explicit_identity(%arg0: tensor<?x3x4xf32>) -> tensor<2x3x?xf32> {
275  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3, 4>} : (tensor<?x3x4xf32>) -> tensor<2x3x?xf32>
276  return %0 : tensor<2x3x?xf32>
277}
278
279// -----
280
281// CHECK-LABEL: test_reshape_3d_same_d2s_auto
282// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
283// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<?x?x?xf32> into tensor<?xf32>
284// CHECK: %[[C0:.*]] = arith.constant 0 : index
285// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
286// CHECK: %[[C8:.*]] = arith.constant 8 : index
287// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C8]] : index
288// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [2, %[[DIV]], 4] : tensor<?xf32> into tensor<2x?x4xf32>
289// CHECK: %[[VAL_2:.*]] = tensor.cast %[[EXPANDED]] : tensor<2x?x4xf32> to tensor<2x3x4xf32>
290// CHECK: return %[[VAL_2]] : tensor<2x3x4xf32>
291func.func @test_reshape_3d_same_d2s_auto(%arg0: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
292  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, -1, 4>} : (tensor<?x?x?xf32>) -> tensor<2x3x4xf32>
293  return %0 : tensor<2x3x4xf32>
294}
295
296// -----
297
298// CHECK-LABEL: test_reshape_3d_same_d2s_explicit
299// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
300// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<?x?x?xf32> into tensor<?xf32>
301// CHECK: %[[C0:.*]] = arith.constant 0 : index
302// CHECK: %[[DIM:.*]] = tensor.dim %[[VAL_0]], %[[C0]] : tensor<?xf32>
303// CHECK: %[[C12:.*]] = arith.constant 12 : index
304// CHECK: %[[DIV:.*]] = arith.divsi %[[DIM]], %[[C12]] : index
305// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[VAL_0]] {{\[\[}}0, 1, 2]] output_shape [%[[DIV]], 3, 4] : tensor<?xf32> into tensor<?x3x4xf32>
306// CHECK: %[[VAL_2:.*]] = tensor.cast %[[EXPANDED]] : tensor<?x3x4xf32> to tensor<2x3x4xf32>
307// CHECK: return %[[VAL_2]] : tensor<2x3x4xf32>
308func.func @test_reshape_3d_same_d2s_explicit(%arg0: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
309  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3, 4>} : (tensor<?x?x?xf32>) -> tensor<2x3x4xf32>
310  return %0 : tensor<2x3x4xf32>
311}
312
313// -----
314
315// CHECK-LABEL: test_reshape_3d_same_s2s_explicit_identity
316// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<2x3x4xf32>
317// CHECK: return %[[ARG_0]] : tensor<2x3x4xf32>
318func.func @test_reshape_3d_same_s2s_explicit_identity(%arg0: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
319  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3, 4>} : (tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
320  return %0 : tensor<2x3x4xf32>
321}
322
323// -----
324
325// CHECK-LABEL: test_reshape_3d_up_d2s_explicit
326// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
327// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2]] : tensor<?x?x?xf32> into tensor<?xf32>
328// CHECK: %[[C0:.*]] = arith.constant 0 : index
329// CHECK: %[[DIM:.*]] = tensor.dim %[[COLLAPSED]], %[[C0]] : tensor<?xf32>
330// CHECK: %[[C6:.*]] = arith.constant 6 : index
331// CHECK: %[[VAL_0:.*]] = arith.divsi %[[DIM]], %[[C6]] : index
332// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSED]] {{\[\[}}0, 1, 2, 3]] output_shape [%[[VAL_0]], 3, 2, 1] : tensor<?xf32> into tensor<?x3x2x1xf32>
333// CHECK: %[[CAST:.*]] = tensor.cast %[[EXPANDED]] : tensor<?x3x2x1xf32> to tensor<1x3x2x1xf32>
334// CHECK: return %[[CAST]] : tensor<1x3x2x1xf32>
335func.func @test_reshape_3d_up_d2s_explicit(%input: tensor<?x?x?xf32>) -> tensor<1x3x2x1xf32> {
336  %0 = tosa.reshape %input {new_shape = array<i64: 1, 3, 2, 1>} : (tensor<?x?x?xf32>) -> tensor<1x3x2x1xf32>
337  return %0 : tensor<1x3x2x1xf32>
338}
339
340// -----
341
342// CHECK-LABEL: test_reshape_4d_down_d2s_explicit
343// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x?x?x?xf32>
344// CHECK: %[[VAL_0:.*]] = tensor.cast %[[ARG_0]] : tensor<?x?x?x?xf32> to tensor<1x1x1x1xf32>
345// CHECK: %[[VAL_1:.*]] = tensor.collapse_shape %[[VAL_0]] [] : tensor<1x1x1x1xf32> into tensor<f32>
346// CHECK: return %[[VAL_1]] : tensor<f32>
347func.func @test_reshape_4d_down_d2s_explicit(%arg0: tensor<?x?x?x?xf32>) -> tensor<f32> {
348  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64>} : (tensor<?x?x?x?xf32>) -> tensor<f32>
349  return %0 : tensor<f32>
350}
351
352// -----
353
354// CHECK-LABEL: test_reshape_5d_down_d2d_auto
355// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?x?x?x2x3xf32>
356// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2, 3, 4]] : tensor<?x?x?x2x3xf32> into tensor<?xf32>
357// CHECK: %[[C0:.*]] = arith.constant 0 : index
358// CHECK: %[[DIM:.*]] = tensor.dim %[[COLLAPSED]], %[[C0]] : tensor<?xf32>
359// CHECK: %[[C6:.*]] = arith.constant 6 : index
360// CHECK: %[[VAL_0:.*]] = arith.divsi %[[DIM]], %[[C6]] : index
361// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSED]] {{\[\[}}0, 1, 2]] output_shape [%[[VAL_0]], 2, 3] : tensor<?xf32> into tensor<?x2x3xf32>
362// CHECK: return %[[EXPANDED]] : tensor<?x2x3xf32>
363func.func @test_reshape_5d_down_d2d_auto(%arg0: tensor<?x?x?x2x3xf32>) -> tensor<?x2x3xf32> {
364  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1, 2, 3>} : (tensor<?x?x?x2x3xf32>) -> tensor<?x2x3xf32>
365  return %0 : tensor<?x2x3xf32>
366}
367
368// -----
369
370// CHECK-LABEL: test_reshape_6d_down_d2d_auto
371// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<1x2x?x5x7x11xf32>
372// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2, 3, 4, 5]] : tensor<1x2x?x5x7x11xf32> into tensor<?xf32>
373// CHECK: %[[C0:.*]] = arith.constant 0 : index
374// CHECK: %[[DIM:.*]] = tensor.dim %[[COLLAPSED]], %[[C0]] : tensor<?xf32>
375// CHECK: %[[C385:.*]] = arith.constant 385 : index
376// CHECK: %[[VAL_0:.*]] = arith.divsi %[[DIM]], %[[C385]] : index
377// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSED]] {{\[\[}}0, 1, 2]] output_shape [%[[VAL_0]], 5, 77] : tensor<?xf32> into tensor<?x5x77xf32>
378// CHECK: return %[[EXPANDED]] : tensor<?x5x77xf32>
379func.func @test_reshape_6d_down_d2d_auto(%arg0: tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32> {
380  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1, 5, 77>} : (tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32>
381  return %0 : tensor<?x5x77xf32>
382}
383
384// -----
385
386// CHECK-LABEL: test_reshape_6d_down_s2s_auto
387// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<1x2x3x5x7x11xf32>
388// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2], [3], [4, 5]] : tensor<1x2x3x5x7x11xf32> into tensor<6x5x77xf32>
389// CHECK: return %[[VAL_0]] : tensor<6x5x77xf32>
390func.func @test_reshape_6d_down_s2s_auto(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32> {
391  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 6, 5, -1>} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
392  return %0 : tensor<6x5x77xf32>
393}
394
395// -----
396
397// This test would previously fail on GCC with certain compiler flags.
398// The GCC issue would cause invalid IR after tosa-to-tensor, so this test
399// locks down that the code goes through tosa-to-tensor and verifies.
400//
401// See https://github.com/llvm/llvm-project/pull/91521 for a full description.
402
403// CHECK-LABEL: reshape_bug_fix
404// CHECK: tensor.expand_shape
405func.func @reshape_bug_fix(%arg0: tensor<?xf32>) -> tensor<1x1x1x?xf32> {
406  %0 = tosa.reshape %arg0 {new_shape = array<i64: 1, 1, 1, -1>} : (tensor<?xf32>) -> tensor<1x1x1x?xf32>
407  return %0 : tensor<1x1x1x?xf32>
408}
409
410// -----
411
412// CHECK-LABEL: test_reshape_6d_down_s2s_explicit
413// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<1x2x3x5x7x11xf32>
414// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2], [3], [4, 5]] : tensor<1x2x3x5x7x11xf32> into tensor<6x5x77xf32>
415// CHECK: return %[[VAL_0]] : tensor<6x5x77xf32>
416func.func @test_reshape_6d_down_s2s_explicit(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32> {
417  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 6, 5, 77>} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
418  return %0 : tensor<6x5x77xf32>
419}
420
421// -----
422
423// CHECK-LABEL: @test_reshape_samerank_unsigned
424//  CHECK-SAME: (%[[ARG0:.*]]: tensor<3x2xui8>)
425func.func @test_reshape_samerank_unsigned(%arg0: tensor<3x2xui8>) -> tensor<2x3xui8> {
426  // CHECK-NEXT: %[[CAST1:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : tensor<3x2xui8> to tensor<3x2xi8>
427  // CHECK-NEXT: %[[RESHAPE1:.*]] = tensor.collapse_shape %[[CAST1]] {{\[}}[0, 1]] : tensor<3x2xi8> into tensor<6xi8>
428  // CHECK-NEXT: %[[RESHAPE2:.*]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1]] output_shape {{\[}}2, 3] : tensor<6xi8> into tensor<2x3xi8>
429  // CHECK-NEXT: %[[CAST2:.*]] = builtin.unrealized_conversion_cast %[[RESHAPE2]] : tensor<2x3xi8> to tensor<2x3xui8
430  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3>} : (tensor<3x2xui8>) -> tensor<2x3xui8>
431  // CHECK-NEXT: return %[[CAST2]]
432  return %0 : tensor<2x3xui8>
433}
434
435// -----
436
437// CHECK-LABEL: func @slice
438func.func @slice(%arg0: tensor<6xf32>) ->() {
439  // CHECK: [[SLICE:%.+]] = tensor.extract_slice %arg0[2] [1] [1]
440  %0 = tosa.const_shape  {value = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1>
441  %1 = tosa.const_shape  {value = dense<1> : tensor<1xindex>} : () -> !tosa.shape<1>
442  %2 = tosa.slice %arg0, %0, %1 : (tensor<6xf32>, !tosa.shape<1>, !tosa.shape<1>) -> tensor<1xf32>
443  return
444}
445
446// -----
447
448// CHECK-LABEL: @slice_dyn
449func.func @slice_dyn(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
450  // CHECK: %[[C0:.+]] = arith.constant 0 : index
451  // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
452  // CHECK: %[[C2:.+]] = arith.constant 2 : index
453  // CHECK: %[[SUB:.+]] = arith.subi %[[DIM]], %[[C2]]
454  // CHECK: tensor.extract_slice %arg0[2] [%[[SUB]]] [1]
455  %0 = tosa.const_shape  {value = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1>
456  %1 = tosa.const_shape  {value = dense<-1> : tensor<1xindex>} : () -> !tosa.shape<1>
457  %2 = tosa.slice %arg0, %0, %1 : (tensor<?xf32>, !tosa.shape<1>, !tosa.shape<1>) -> tensor<?xf32>
458  return %2 : tensor<?xf32>
459}
460
461// -----
462
463// CHECK-LABEL: @pad_float
464// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
465func.func @pad_float(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
466  %0 = tosa.const_shape {value = dense<[1, 2, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
467  // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
468  // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
469  // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
470  // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
471  // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
472  // CHECK: tensor.pad %[[ARG0]] low{{\[}}[[INDEX1]], [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
473  // CHECK:   tensor.yield [[CST]]
474  // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
475  %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xf32>, !tosa.shape<4>)  -> (tensor<4x9xf32>)
476  return %1 : tensor<4x9xf32>
477}
478// -----
479
480func.func @pad_int(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
481  %0 = tosa.const_shape {value = dense<[1, 2, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
482  // CHECK: [[CST:%.+]] = arith.constant 0 : i32
483  // CHECK: tensor.pad
484  // CHECK:   tensor.yield [[CST]]
485  %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xi32>, !tosa.shape<4>)  -> (tensor<4x9xi32>)
486  return %1 : tensor<4x9xi32>
487}
488// -----
489
490func.func @pad_quant(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
491  %0 = tosa.const_shape {value = dense<[1, 2, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
492  // CHECK: [[CST:%.+]] = arith.constant 42 : i32
493  // CHECK: tensor.pad
494  // CHECK:   tensor.yield [[CST]]
495  %1 = "tosa.pad"(%arg0, %0) {quantization_info = #tosa.pad_quant<input_zp = 42>} : (tensor<1x2xi32>, !tosa.shape<4>)  -> (tensor<4x9xi32>)
496  return %1 : tensor<4x9xi32>
497}
498
499// -----
500
501func.func @pad_float_explicit(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
502  %0 = tosa.const_shape {value = dense<[1, 2, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
503  // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
504  // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
505  // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
506  // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
507  // CHECK-DAG: [[CST:%.+]] = arith.constant 4.200000e+01 : f32
508  // CHECK: tensor.pad %[[ARG0]] low{{\[}}[[INDEX1]], [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
509  // CHECK:   tensor.yield [[CST]]
510  // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
511  %1 = arith.constant dense<42.0> : tensor<f32>
512  %2 = "tosa.pad"(%arg0, %0, %1)  : (tensor<1x2xf32>, !tosa.shape<4>, tensor<f32>)  -> (tensor<4x9xf32>)
513  return %2 : tensor<4x9xf32>
514}
515
516// -----
517
518func.func @pad_dyn_input(%arg0 : tensor<?x2xf32>) -> (tensor<?x9xf32>) {
519  %0 = tosa.const_shape {value = dense<[1, 2, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
520  // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
521  // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
522  // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
523  // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
524  // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
525  // CHECK: tensor.pad %[[ARG0]] low{{\[}}[[INDEX1]], [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
526  // CHECK:   tensor.yield [[CST]]
527  // CHECK: } : tensor<?x2xf32> to tensor<?x9xf32>
528  %1 = "tosa.pad"(%arg0, %0)  : (tensor<?x2xf32>, !tosa.shape<4>)  -> (tensor<?x9xf32>)
529  return %1 : tensor<?x9xf32>
530}
531// -----
532
533func.func @pad_dyn_padding(%arg0 : tensor<1x2xf32>) -> (tensor<?x9xf32>) {
534  %0 = tosa.const_shape {value = dense<[-1, 2, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
535  // CHECK-DAG: [[INDEX1:%.+]] = arith.constant -1 : index
536  // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
537  // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
538  // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
539  // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
540  // CHECK: tensor.pad %[[ARG0]] low{{\[}}[[INDEX1]], [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
541  // CHECK:   tensor.yield [[CST]]
542  // CHECK: } : tensor<1x2xf32> to tensor<?x9xf32>
543  %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xf32>, !tosa.shape<4>)  -> (tensor<?x9xf32>)
544  return %1 : tensor<?x9xf32>
545}
546
547// -----
548
549// CHECK-LABEL: @concat
550// CHECK-SAME: %[[ARG0:.+]]: tensor<5x1xf32>
551// CHECK-SAME: %[[ARG1:.+]]: tensor<6x1xf32>
552func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
553  // CHECK-DAG: [[INIT:%.+]] = tensor.empty() : tensor<11x1xf32>
554  // CHECK-DAG: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
555  // CHECK-DAG: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG1]] into [[INSERT0]][5, 0] [6, 1] [1, 1]
556  %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i32} : (tensor<5x1xf32>, tensor<6x1xf32>)  -> (tensor<11x1xf32>)
557
558  // CHECK-DAG: [[INIT:%.+]] = tensor.empty() : tensor<5x2xf32>
559  // CHECK-DAG: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
560  // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG0]] into [[INSERT0]][0, 1] [5, 1] [1, 1]
561  %1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i32} : (tensor<5x1xf32>, tensor<5x1xf32>)  -> (tensor<5x2xf32>)
562  return
563}
564
565// -----
566
567// CHECK-LABEL: @concat_non_axis_dyn
568// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
569// CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]
570func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -> () {
571  // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0
572  // CHECK-DAG: %[[IDX1:.+]] = arith.constant 1
573  // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX1]]
574  // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[DIM0]]) : tensor<11x?xf32>
575  // CHECK-DAG: %[[IDX1_1:.+]] = arith.constant 1 : index
576  // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[IDX1_1]]
577  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [5, %[[DIM1]]] [1, 1]
578  // CHECK-DAG: %[[IDX1_2:.+]] = arith.constant 1 : index
579  // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG1]], %[[IDX1_2]] : tensor<6x?xf32>
580  // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][5, 0] [6, %[[DIM2]]] [1, 1]
581  %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i32} : (tensor<5x?xf32>, tensor<6x?xf32>)  -> (tensor<11x?xf32>)
582  return
583}
584
585// -----
586
587// CHECK-LABEL: @concat_axis_dyn
588// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
589// CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
590func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> () {
591  // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0 : index
592  // CHECK-DAG: %[[IDX0:.+]] = arith.constant 0 : index
593  // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX0]] : tensor<?x3xf32>
594  // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG1]], %[[AXIS]] : tensor<?x3xf32>
595  // CHECK-DAG: %[[SUM:.+]] = arith.addi %[[DIM0]], %[[DIM1]] : index
596  // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[SUM]]) : tensor<?x3xf32>
597  // CHECK-DAG: %[[IDX0_1:.+]] = arith.constant 0 : index
598  // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[IDX0_1]] : tensor<?x3xf32>
599  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM2]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
600  // CHECK-DAG: %[[IDX0_2:.+]] = arith.constant 0 : index
601  // CHECK-DAG: %[[DIM3:.+]] = tensor.dim %[[ARG1]], %[[IDX0_2]] : tensor<?x3xf32>
602  // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[DIM0]], 0] [%[[DIM3]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
603
604  %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i32} : (tensor<?x3xf32>, tensor<?x3xf32>)  -> (tensor<?x3xf32>)
605  return
606}
607
608// -----
609
610// CHECK-LABEL: @concat_axis_dyn_mixed
611// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
612// CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
613// CHECK-SAME:  %[[ARG2:[0-9a-zA-Z_]*]]:
614func.func @concat_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> () {
615  // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
616  // CHECK-DAG: %[[C0_0:.+]] = arith.constant 0 : index
617  // CHECK-DAG: %[[OFFSET0:.+]] = tensor.dim %[[ARG0]], %[[C0_0]] : tensor<?x1xf32>
618  // CHECK-DAG: %[[DIM1_0:.+]] = tensor.dim %[[ARG1]], %[[C0]] : tensor<?x1xf32>
619  // CHECK-DAG: %[[OFFSET1:.+]] = arith.addi %[[OFFSET0]], %[[DIM1_0]] : index
620  // CHECK-DAG: %[[DIM2_2:.+]] = tensor.dim %[[ARG2]], %[[C0]] : tensor<?x1xf32>
621  // CHECK-DAG: %[[OFFSET2:.+]] = arith.addi %[[OFFSET1]], %[[DIM2_2]] : index
622  // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x1xf32>
623  // CHECK-DAG: %[[C0_3:.+]] = arith.constant 0 : index
624  // CHECK-DAG: %[[DIM_4:.+]] = tensor.dim %[[ARG0]], %[[C0_3]] : tensor<?x1xf32>
625  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM_4]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
626  // CHECK-DAG: %[[C0_4:.+]] = arith.constant 0 : index
627  // CHECK-DAG: %[[DIM_6:.+]] = tensor.dim %[[ARG1]], %[[C0_4]] : tensor<?x1xf32>
628  // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[OFFSET0]], 0] [%[[DIM_6]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
629  // CHECK-DAG: %[[C0_8:.+]] = arith.constant 0 : index
630  // CHECK-DAG: %[[DIM_9:.+]] = tensor.dim %[[ARG2]], %[[C0_8]] : tensor<?x1xf32>
631  // CHECK-DAG: %[[INSERT3:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][%[[OFFSET1]], 0] [%[[DIM_9]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
632
633  // CHECK: return
634
635  %0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 0 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x1xf32>
636  return
637}
638
639// -----
640
641// CHECK-LABEL: @concat_non_axis_dyn_mixed
642// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
643// CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
644// CHECK-SAME:  %[[ARG2:[0-9a-zA-Z_]*]]:
645func.func @concat_non_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> () {
646  // CHECK-DAG: %[[UNUSED0:.+]] = arith.constant 0 : index
647  // CHECK-DAG: %[[UNUSED1:.+]] = tensor.dim %[[ARG0]], %[[UNUSED0]] : tensor<?x1xf32>
648
649  // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x3xf32>
650  // CHECK-DAG: %[[C0_0:.+]] = arith.constant 0 : index
651  // CHECK-DAG: %[[DIM0_0:.+]] = tensor.dim %[[ARG0]], %[[C0_0]] : tensor<?x1xf32>
652  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM0_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
653  // CHECK-DAG: %[[C0_1:.+]] = arith.constant 0 : index
654  // CHECK-DAG: %[[DIM1_0:.+]] = tensor.dim %[[ARG1]], %[[C0_1]] : tensor<?x1xf32>
655  // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][0, 1] [%[[DIM1_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
656  // CHECK-DAG: %[[C0_2:.+]] = arith.constant 0 : index
657  // CHECK-DAG: %[[DIM2_0:.+]] = tensor.dim %[[ARG2]], %[[C0_2]] : tensor<?x1xf32>
658  // CHECK-DAG: %[[INSERT2:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][0, 2] [%[[DIM2_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
659  // CHECK: return
660
661  %0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 1 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x3xf32>
662  return
663}
664