mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

   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>
   8 func.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>
  20 func.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>
  32 func.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>
  43 func.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>
  54 func.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>
  66 func.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>
  77 func.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.divui %[[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>
  92 func.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>
 103 func.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>
 114 func.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>
 125 func.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.divui %[[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>
 141 func.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>
 154 func.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>
 167 func.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>
 179 func.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.divui %[[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>
 196 func.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.divui %[[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>
 213 func.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.divui %[[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>
 229 func.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.divui %[[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>
 246 func.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.divui %[[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>
 263 func.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>
 274 func.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.divui %[[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>
 291 func.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.divui %[[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>
 308 func.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>
 318 func.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.divui %[[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>
 335 func.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>
 347 func.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.divui %[[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>
 363 func.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.divui %[[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>
 379 func.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>
 390 func.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
 405 func.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>
 416 func.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>)
 425 func.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
 438 func.func @slice(%arg0: tensor<6xf32>) ->() {
 439   // CHECK: [[SLICE:%.+]] = tensor.extract_slice %arg0[2] [1] [1]
 440   %0 = "tosa.slice"(%arg0) {start = array<i64: 2>, size = array<i64: 1>} : (tensor<6xf32>)  -> (tensor<1xf32>)
 441   return
 442 }
 443
 444 // -----
 445
 446 // CHECK-LABEL: @slice_dyn
 447 func.func @slice_dyn(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
 448   // CHECK: %[[C0:.+]] = arith.constant 0 : index
 449   // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
 450   // CHECK: %[[C2:.+]] = arith.constant 2 : index
 451   // CHECK: %[[SUB:.+]] = arith.subi %[[DIM]], %[[C2]]
 452   // CHECK: tensor.extract_slice %arg0[2] [%[[SUB]]] [1]
 453   %0 = "tosa.slice"(%arg0) {start = array<i64: 2>, size = array<i64: -1>} : (tensor<?xf32>)  -> (tensor<?xf32>)
 454   return %0 : tensor<?xf32>
 455 }
 456
 457 // -----
 458
 459 // CHECK-LABEL: @pad_float
 460 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 461 func.func @pad_float(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
 462   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
 463   // TODO: Output contains multiple "arith.constant 1 : index".
 464   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
 465   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
 466   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
 467   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
 468   // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
 469   // CHECK: tensor.pad %[[ARG0]] low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
 470   // CHECK:   tensor.yield [[CST]]
 471   // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
 472   %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xf32>, tensor<2x2xi32>)  -> (tensor<4x9xf32>)
 473   return %1 : tensor<4x9xf32>
 474 }
 475
 476 func.func @pad_int(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
 477   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
 478   // CHECK: [[CST:%.+]] = arith.constant 0 : i32
 479   // CHECK: tensor.pad
 480   // CHECK:   tensor.yield [[CST]]
 481   %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xi32>, tensor<2x2xi32>)  -> (tensor<4x9xi32>)
 482   return %1 : tensor<4x9xi32>
 483 }
 484
 485 func.func @pad_quant(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
 486   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
 487   // CHECK: [[CST:%.+]] = arith.constant 42 : i32
 488   // CHECK: tensor.pad
 489   // CHECK:   tensor.yield [[CST]]
 490   %1 = "tosa.pad"(%arg0, %0) {quantization_info = #tosa.pad_quant<input_zp = 42>} : (tensor<1x2xi32>, tensor<2x2xi32>)  -> (tensor<4x9xi32>)
 491   return %1 : tensor<4x9xi32>
 492 }
 493
 494 // -----
 495
 496 func.func @pad_float_explicit(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
 497   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
 498   // TODO: Output contains multiple "arith.constant 1 : index".
 499   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
 500   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
 501   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
 502   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
 503   // CHECK-DAG: [[CST:%.+]] = arith.constant 4.200000e+01 : f32
 504   // CHECK: tensor.pad %[[ARG0]] low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
 505   // CHECK:   tensor.yield [[CST]]
 506   // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
 507   %1 = arith.constant dense<42.0> : tensor<f32>
 508   %2 = "tosa.pad"(%arg0, %0, %1)  : (tensor<1x2xf32>, tensor<2x2xi32>, tensor<f32>)  -> (tensor<4x9xf32>)
 509   return %2 : tensor<4x9xf32>
 510 }
 511
 512 // -----
 513
 514 func.func @pad_dyn_input(%arg0 : tensor<?x2xf32>) -> (tensor<?x9xf32>) {
 515   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
 516   // TODO: Output contains multiple "arith.constant 1 : index".
 517   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
 518   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
 519   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
 520   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
 521   // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
 522   // CHECK: tensor.pad %[[ARG0]] low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
 523   // CHECK:   tensor.yield [[CST]]
 524   // CHECK: } : tensor<?x2xf32> to tensor<?x9xf32>
 525   %1 = "tosa.pad"(%arg0, %0)  : (tensor<?x2xf32>, tensor<2x2xi32>)  -> (tensor<?x9xf32>)
 526   return %1 : tensor<?x9xf32>
 527 }
 528
 529 func.func @pad_dyn_padding(%arg0 : tensor<1x2xf32>) -> (tensor<?x9xf32>) {
 530   %0 = arith.constant dense<[[-1, 2], [3, 4]]> : tensor<2x2xi32>
 531   // TODO: Output contains multiple "arith.constant 1 : index".
 532   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
 533   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
 534   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
 535   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
 536   // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
 537   // CHECK: tensor.pad %[[ARG0]] low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
 538   // CHECK:   tensor.yield [[CST]]
 539   // CHECK: } : tensor<1x2xf32> to tensor<?x9xf32>
 540   %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xf32>, tensor<2x2xi32>)  -> (tensor<?x9xf32>)
 541   return %1 : tensor<?x9xf32>
 542 }
 543
 544 // -----
 545
 546 // CHECK-LABEL: @concat
 547 // CHECK-SAME: %[[ARG0:.+]]: tensor<5x1xf32>
 548 // CHECK-SAME: %[[ARG1:.+]]: tensor<6x1xf32>
 549 func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
 550   // CHECK-DAG: [[INIT:%.+]] = tensor.empty() : tensor<11x1xf32>
 551   // CHECK-DAG: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
 552   // CHECK-DAG: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG1]] into [[INSERT0]][5, 0] [6, 1] [1, 1]
 553   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i32} : (tensor<5x1xf32>, tensor<6x1xf32>)  -> (tensor<11x1xf32>)
 554
 555   // CHECK-DAG: [[INIT:%.+]] = tensor.empty() : tensor<5x2xf32>
 556   // CHECK-DAG: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
 557   // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG0]] into [[INSERT0]][0, 1] [5, 1] [1, 1]
 558   %1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i32} : (tensor<5x1xf32>, tensor<5x1xf32>)  -> (tensor<5x2xf32>)
 559   return
 560 }
 561
 562 // -----
 563
 564 // CHECK-LABEL: @concat_non_axis_dyn
 565 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 566 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]
 567 func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -> () {
 568   // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0
 569   // CHECK-DAG: %[[IDX1:.+]] = arith.constant 1
 570   // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX1]]
 571   // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[DIM0]]) : tensor<11x?xf32>
 572   // CHECK-DAG: %[[IDX1_1:.+]] = arith.constant 1 : index
 573   // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[IDX1_1]]
 574   // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [5, %[[DIM1]]] [1, 1]
 575   // CHECK-DAG: %[[IDX1_2:.+]] = arith.constant 1 : index
 576   // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG1]], %[[IDX1_2]] : tensor<6x?xf32>
 577   // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][5, 0] [6, %[[DIM2]]] [1, 1]
 578   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i32} : (tensor<5x?xf32>, tensor<6x?xf32>)  -> (tensor<11x?xf32>)
 579   return
 580 }
 581
 582 // -----
 583
 584 // CHECK-LABEL: @concat_axis_dyn
 585 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 586 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
 587 func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> () {
 588   // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0 : index
 589   // CHECK-DAG: %[[IDX0:.+]] = arith.constant 0 : index
 590   // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX0]] : tensor<?x3xf32>
 591   // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG1]], %[[AXIS]] : tensor<?x3xf32>
 592   // CHECK-DAG: %[[SUM:.+]] = arith.addi %[[DIM0]], %[[DIM1]] : index
 593   // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[SUM]]) : tensor<?x3xf32>
 594   // CHECK-DAG: %[[IDX0_1:.+]] = arith.constant 0 : index
 595   // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[IDX0_1]] : tensor<?x3xf32>
 596   // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM2]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
 597   // CHECK-DAG: %[[IDX0_2:.+]] = arith.constant 0 : index
 598   // CHECK-DAG: %[[DIM3:.+]] = tensor.dim %[[ARG1]], %[[IDX0_2]] : tensor<?x3xf32>
 599   // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[DIM0]], 0] [%[[DIM3]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
 600
 601   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i32} : (tensor<?x3xf32>, tensor<?x3xf32>)  -> (tensor<?x3xf32>)
 602   return
 603 }
 604
 605 // -----
 606
 607 // CHECK-LABEL: @concat_axis_dyn_mixed
 608 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 609 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
 610 // CHECK-SAME:  %[[ARG2:[0-9a-zA-Z_]*]]:
 611 func.func @concat_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> () {
 612   // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
 613   // CHECK-DAG: %[[C0_0:.+]] = arith.constant 0 : index
 614   // CHECK-DAG: %[[OFFSET0:.+]] = tensor.dim %[[ARG0]], %[[C0_0]] : tensor<?x1xf32>
 615   // CHECK-DAG: %[[DIM1_0:.+]] = tensor.dim %[[ARG1]], %[[C0]] : tensor<?x1xf32>
 616   // CHECK-DAG: %[[OFFSET1:.+]] = arith.addi %[[OFFSET0]], %[[DIM1_0]] : index
 617   // CHECK-DAG: %[[DIM2_2:.+]] = tensor.dim %[[ARG2]], %[[C0]] : tensor<?x1xf32>
 618   // CHECK-DAG: %[[OFFSET2:.+]] = arith.addi %[[OFFSET1]], %[[DIM2_2]] : index
 619   // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x1xf32>
 620   // CHECK-DAG: %[[C0_3:.+]] = arith.constant 0 : index
 621   // CHECK-DAG: %[[DIM_4:.+]] = tensor.dim %[[ARG0]], %[[C0_3]] : tensor<?x1xf32>
 622   // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM_4]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
 623   // CHECK-DAG: %[[C0_4:.+]] = arith.constant 0 : index
 624   // CHECK-DAG: %[[DIM_6:.+]] = tensor.dim %[[ARG1]], %[[C0_4]] : tensor<?x1xf32>
 625   // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[OFFSET0]], 0] [%[[DIM_6]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
 626   // CHECK-DAG: %[[C0_8:.+]] = arith.constant 0 : index
 627   // CHECK-DAG: %[[DIM_9:.+]] = tensor.dim %[[ARG2]], %[[C0_8]] : tensor<?x1xf32>
 628   // CHECK-DAG: %[[INSERT3:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][%[[OFFSET1]], 0] [%[[DIM_9]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
 629
 630   // CHECK: return
 631
 632   %0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 0 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x1xf32>
 633   return
 634 }
 635
 636 // -----
 637
 638 // CHECK-LABEL: @concat_non_axis_dyn_mixed
 639 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 640 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
 641 // CHECK-SAME:  %[[ARG2:[0-9a-zA-Z_]*]]:
 642 func.func @concat_non_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> () {
 643   // CHECK-DAG: %[[UNUSED0:.+]] = arith.constant 0 : index
 644   // CHECK-DAG: %[[UNUSED1:.+]] = tensor.dim %[[ARG0]], %[[UNUSED0]] : tensor<?x1xf32>
 645
 646   // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x3xf32>
 647   // CHECK-DAG: %[[C0_0:.+]] = arith.constant 0 : index
 648   // CHECK-DAG: %[[DIM0_0:.+]] = tensor.dim %[[ARG0]], %[[C0_0]] : tensor<?x1xf32>
 649   // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM0_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
 650   // CHECK-DAG: %[[C0_1:.+]] = arith.constant 0 : index
 651   // CHECK-DAG: %[[DIM1_0:.+]] = tensor.dim %[[ARG1]], %[[C0_1]] : tensor<?x1xf32>
 652   // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][0, 1] [%[[DIM1_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
 653   // CHECK-DAG: %[[C0_2:.+]] = arith.constant 0 : index
 654   // CHECK-DAG: %[[DIM2_0:.+]] = tensor.dim %[[ARG2]], %[[C0_2]] : tensor<?x1xf32>
 655   // CHECK-DAG: %[[INSERT2:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][0, 2] [%[[DIM2_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
 656   // CHECK: return
 657
 658   %0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 1 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x3xf32>
 659   return
 660 }