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