mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

   1 // RUN: mlir-opt --split-input-file --tosa-to-linalg %s -verify-diagnostics -o -| FileCheck %s
   2
   3 // CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
   4
   5 // CHECK-LABEL: @test_abs
   6 func @test_abs(%arg0: tensor<f32>) -> tensor<f32> {
   7   // CHECK: [[INIT:%.+]] = linalg.init_tensor [] : tensor<f32>
   8   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = []} ins(%arg0 : tensor<f32>) outs([[INIT]] : tensor<f32>) {
   9   // CHECK: ^bb0(%arg1: f32, %arg2: f32):
  10   // CHECK:   [[ELEMENT:%.+]] = math.abs %arg1
  11   // CHECK:   linalg.yield [[ELEMENT]] : f32
  12   // CHECK: } -> tensor<f32>
  13
  14   %0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
  15
  16   // CHECK: return [[GENERIC]]
  17   return %0 : tensor<f32>
  18 }
  19
  20 // -----
  21
  22 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
  23
  24 // CHECK-LABEL: @test_abs
  25 func @test_abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {
  26   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32>
  27   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) {
  28   // CHECK: ^bb0(%arg1: f32, %arg2: f32):
  29   // CHECK:   [[ELEMENT:%.+]] = math.abs %arg1
  30   // CHECK:   linalg.yield [[ELEMENT]] : f32
  31   // CHECK: } -> tensor<2xf32>
  32   %0 = "tosa.abs"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
  33
  34   // CHECK: return [[GENERIC]]
  35   return %0 : tensor<2xf32>
  36 }
  37
  38 // -----
  39
  40 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
  41
  42 // CHECK-LABEL: @test_abs
  43 func @test_abs(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
  44   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3] : tensor<2x3xf32>
  45   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<2x3xf32>) outs([[INIT]] : tensor<2x3xf32>) {
  46   // CHECK: ^bb0(%arg1: f32, %arg2: f32):
  47   // CHECK:   [[ELEMENT:%.+]] = math.abs %arg1
  48   // CHECK:   linalg.yield [[ELEMENT]] : f32
  49   // CHECK: } -> tensor<2x3xf32>
  50   %0 = "tosa.abs"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
  51
  52   // CHECK: return [[GENERIC]]
  53   return %0 : tensor<2x3xf32>
  54 }
  55
  56 // -----
  57
  58 // CHECK-LABEL: @test_abs
  59 func @test_abs(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  60   // CHECK: %[[C0:.+]] = arith.constant 0
  61   // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
  62   // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM]]]
  63   // CHECK: linalg.generic
  64   // CHECK: math.abs
  65   %0 = "tosa.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
  66   return %0 : tensor<?xf32>
  67 }
  68
  69 // -----
  70
  71 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
  72
  73 // CHECK-LABEL: @test_abs_dyn
  74 func @test_abs_dyn(%arg0: tensor<2x?xf32>) -> tensor<2x?xf32> {
  75   // CHECK: %[[C1:.+]] = arith.constant 1
  76   // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C1]]
  77   // CHECK: %[[INIT:.+]] = linalg.init_tensor [2, %[[DIM]]]
  78   // CHECK: linalg.generic
  79   // CHECK: math.abs
  80   %0 = "tosa.abs"(%arg0) : (tensor<2x?xf32>) -> tensor<2x?xf32>
  81   return %0 : tensor<2x?xf32>
  82 }
  83 // -----
  84
  85
  86 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> ()>
  87 // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
  88
  89 // CHECK-LABEL: @test_broadcast
  90 func @test_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  91   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32>
  92   // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0
  93   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins([[RESHAPE]], %arg1 : tensor<f32>, tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) {
  94   // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
  95   // CHECK:   [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
  96   // CHECK:   linalg.yield [[ELEMENT]] : f32
  97   // CHECK: } -> tensor<2xf32>
  98   %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<2xf32>) -> tensor<2xf32>
  99   return %0 : tensor<2xf32>
 100 }
 101
 102 // -----
 103
 104 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 105 // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> ()>
 106
 107 // CHECK-LABEL: @test_broadcast_swapped_args
 108 func @test_broadcast_swapped_args(%arg0: tensor<2xf32>, %arg1: tensor<1xf32>) -> tensor<2xf32> {
 109   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32>
 110   // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg1
 111   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[RESHAPE]] : tensor<2xf32>, tensor<f32>) outs([[INIT]] : tensor<2xf32>) {
 112   // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
 113   // CHECK:   [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
 114   // CHECK:   linalg.yield [[ELEMENT]] : f32
 115   // CHECK: } -> tensor<2xf32>
 116   %0 = "tosa.add"(%arg0, %arg1) : (tensor<2xf32>, tensor<1xf32>) -> tensor<2xf32>
 117   return %0 : tensor<2xf32>
 118 }
 119
 120 // -----
 121
 122 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 123 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
 124 // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
 125
 126 // CHECK-LABEL: @test_multibroadcast
 127 func @test_multibroadcast(%arg0: tensor<1x3xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x3xf32> {
 128   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3] : tensor<2x3xf32>
 129   // CHECK: [[RESHAPE1:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
 130   // CHECK: [[RESHAPE2:%.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1]]
 131   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins([[RESHAPE1]], [[RESHAPE2]] : tensor<3xf32>, tensor<2xf32>) outs([[INIT]] : tensor<2x3xf32>) {
 132   // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
 133   // CHECK:   [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
 134   // CHECK:   linalg.yield [[ELEMENT]] : f32
 135   // CHECK: } -> tensor<2x3xf32>
 136   %0 = "tosa.add"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32>
 137   return %0 : tensor<2x3xf32>
 138 }
 139
 140 // -----
 141
 142 // CHECK-LABEL: @test_simple_f32
 143 func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
 144   // CHECK: linalg.generic
 145   // CHECK: tanh
 146   %0 = "tosa.tanh"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
 147
 148   // CHECK: linalg.generic
 149   // CHECK: math.abs
 150   %1 = "tosa.abs"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
 151
 152   // CHECK: linalg.generic
 153   // CHECK: arith.addf
 154   %2 = "tosa.add"(%0, %0) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 155
 156   // CHECK: linalg.generic
 157   // CHECK: arith.subf
 158   %3 = "tosa.sub"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 159
 160   // CHECK: linalg.generic
 161   // CHECK: arith.mulf
 162   %4 = "tosa.mul"(%0, %1) {shift = 0 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 163
 164   // CHECK: linalg.generic
 165   // CHECK: arith.negf
 166   %5 = "tosa.negate"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 167
 168   // CHECK: linalg.generic
 169   // CHECK: pow
 170   %6 = "tosa.pow"(%1, %2) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 171
 172   // CHECK: linalg.generic
 173   // CHECK: rsqrt
 174   %7 = "tosa.rsqrt"(%1) : (tensor<1xf32>) -> tensor<1xf32>
 175
 176   // CHECK: linalg.generic
 177   // CHECK: log
 178   %8 = "tosa.log"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
 179
 180   // CHECK: linalg.generic
 181   // CHECK: exp
 182   %9 = "tosa.exp"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
 183
 184   // CHECK: linalg.generic
 185   // CHECK: arith.cmpf
 186   %10 = "tosa.greater"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>
 187
 188   // CHECK: linalg.generic
 189   // CHECK: arith.cmpf
 190   %11 = "tosa.greater_equal"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>
 191
 192   // CHECK: linalg.generic
 193   // CHECK: arith.cmpf
 194   %12 = "tosa.equal"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>
 195
 196   // CHECK: linalg.generic
 197   // CHECK: select
 198   %13 = "tosa.select"(%10, %0, %1) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 199
 200   // CHECK: linalg.generic
 201   // CHECK: arith.cmpf
 202   // CHECK: select
 203   %14 = "tosa.maximum"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 204
 205   // CHECK: linalg.generic
 206   // CHECK: arith.cmpf
 207   // CHECK: select
 208   %15 = "tosa.minimum"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
 209
 210   // CHECK: linalg.generic
 211   // CHECK: ceil
 212   %16 = "tosa.ceil"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 213
 214   // CHECK: linalg.generic
 215   // CHECK: floor
 216   %17 = "tosa.floor"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 217
 218   // CHECK: linalg.generic
 219   // CHECK: arith.cmpf
 220   // CHECK: select
 221   %18 = "tosa.clamp"(%0) {min_int = 1 : i64, max_int = 5 : i64, min_fp = 1.0 : f32, max_fp = 5.0 : f32} : (tensor<1xf32>) -> tensor<1xf32>
 222
 223   // CHECK: linalg.generic
 224   // CHECK: arith.cmpf
 225   // CHECK: select
 226   %19 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xf32>) -> tensor<1xf32>
 227
 228   // CHECK: linalg.generic
 229   // CHECK: arith.negf
 230   // CHECK: exp
 231   // CHECK: arith.addf
 232   // CHECK: arith.divf
 233   %20 = "tosa.sigmoid"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 234
 235   // CHECK: linalg.generic
 236   // CHECK: arith.constant 0.000000e+00
 237   // CHECK: arith.constant 5.000000e-01
 238   // CHECK: arith.constant -2.14748365E+9
 239   // CHECK: arith.constant 2.14748365E+9
 240   // CHECK: arith.addf
 241   // CHECK: arith.subf
 242   // CHECK: arith.cmpf olt
 243   // CHECK: select
 244   // CHECK: arith.cmpf olt
 245   // CHECK: select
 246   // CHECK: arith.cmpf olt
 247   // CHECK: select
 248   // CHECK: arith.fptosi
 249   %21 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi32>
 250
 251   // CHECK: linalg.generic
 252   // CHECK: arith.constant 0
 253   // CHECK: arith.cmpf
 254   %22 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi1>
 255
 256   // CHECK: linalg.generic
 257   // CHECK: arith.truncf
 258   %23 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf16>
 259
 260   // CHECK: linalg.generic
 261   // CHECK: arith.divf
 262   %24 = "tosa.reciprocal"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 263
 264   return
 265 }
 266
 267 // -----
 268
 269 // CHECK-LABEL: @test_simple_f16
 270 func @test_simple_f16(%arg0: tensor<1xf16>) -> () {
 271
 272   // CHECK: linalg.generic
 273   // CHECK: arith.extf
 274   %0 = "tosa.cast"(%arg0) : (tensor<1xf16>) -> tensor<1xf32>
 275
 276   return
 277 }
 278
 279 // -----
 280
 281 // CHECK-LABEL: @test_simple_i16
 282 func @test_simple_i16(%arg0: tensor<1xi16>) -> () {
 283   // CHECK: linalg.generic
 284   // CHECK: arith.extsi
 285   // CHECK: arith.extsi
 286   // CHECK: arith.muli
 287   %0 = "tosa.mul"(%arg0, %arg0) {shift = 0 : i32} : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi32>
 288
 289   return
 290 }
 291
 292 // -----
 293
 294 // CHECK-LABEL: @test_simple_ui8
 295 func @test_simple_ui8(%arg0: tensor<1xui8>) -> () {
 296   // CHECK: arith.uitofp
 297   %0 = "tosa.cast"(%arg0) : (tensor<1xui8>) -> tensor<1xf32>
 298   return
 299 }
 300
 301 // -----
 302
 303 // CHECK-LABEL: @test_simple_i32
 304 func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
 305   // CHECK: linalg.generic
 306   // CHECK: arith.addi
 307   %0 = "tosa.add"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 308
 309   // CHECK: linalg.generic
 310   // CHECK: arith.subi
 311   %1 = "tosa.sub"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 312
 313   // CHECK: linalg.generic
 314   // CHECK: arith.muli
 315   %2 = "tosa.mul"(%arg0, %arg0) {shift = 0 : i32} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 316
 317   // CHECK: linalg.generic
 318   // CHECK: arith.constant 2
 319   // CHECK: apply_scale
 320   %3 = "tosa.mul"(%arg0, %arg0) {shift = 2 : i32} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 321
 322   // CHECK: linalg.generic
 323   // CHECK: arith.divsi
 324   %4 = "tosa.div"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 325
 326   // CHECK: linalg.generic
 327   // CHECK: [[ZERO:%.+]] = arith.constant 0
 328   // CHECK: arith.subi [[ZERO]], %arg1
 329   %5 = "tosa.negate"(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
 330
 331   // CHECK: linalg.generic
 332   // CHECK: and
 333   %6 = "tosa.bitwise_and"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 334
 335   // CHECK: linalg.generic
 336   // CHECK: or
 337   %7 = "tosa.bitwise_or"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 338
 339   // CHECK: linalg.generic
 340   // CHECK: arith.xori
 341   %8 = "tosa.bitwise_xor"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 342
 343   // CHECK: linalg.generic
 344   // CHECK: arith.shli
 345   %9 = "tosa.logical_left_shift"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 346
 347   // CHECK: linalg.generic
 348   // CHECK: arith.shrui
 349   %10 = "tosa.logical_right_shift"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 350
 351   // CHECK: linalg.generic
 352   // CHECK: arith.shrsi
 353   %11 = "tosa.arithmetic_right_shift"(%arg0, %arg0) {round = 0 : i1} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 354
 355   // CHECK: linalg.generic
 356   // CHECK: arith.constant 1
 357   // CHECK: arith.constant 0
 358   // CHECK: arith.constant true
 359   // CHECK: arith.cmpi
 360   // CHECK: arith.subi
 361   // CHECK: arith.shrsi
 362   // CHECK: arith.trunci
 363   // CHECK: and
 364   // CHECK: and
 365   // CHECK: arith.extui
 366   // CHECK: arith.addi
 367   %12 = "tosa.arithmetic_right_shift"(%arg0, %arg0) {round = 1 : i1} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 368
 369   // CHECK: scf.while
 370   // CHECK: arith.cmpi ne
 371   // CHECK: scf.condition
 372   // CHECK: arith.shrui
 373   // CHECK: arith.subi
 374   // CHECK: scf.yield
 375   %13 = "tosa.clz"(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
 376
 377   // CHECK: linalg.generic
 378   // CHECK: arith.cmpi
 379   %14 = "tosa.greater"(%0, %1) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
 380
 381   // CHECK: linalg.generic
 382   // CHECK: arith.cmpi
 383   %15 = "tosa.greater_equal"(%0, %1) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
 384
 385   // CHECK: linalg.generic
 386   // CHECK: select
 387   %16 = "tosa.select"(%14, %0, %1) : (tensor<1xi1>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 388
 389   // CHECK: linalg.generic
 390   // CHECK: arith.cmpi
 391   // CHECK: select
 392   %17 = "tosa.maximum"(%0, %1) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 393
 394   // CHECK: linalg.generic
 395   // CHECK: arith.cmpi
 396   // CHECK: select
 397   %18 = "tosa.minimum"(%0, %1) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
 398
 399   // CHECK: linalg.generic
 400   // CHECK: arith.cmpi
 401   // CHECK: select
 402   %19 = "tosa.clamp"(%0) {min_int = 1 : i64, max_int = 5 : i64, min_fp = 1.0 : f32, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 403
 404   // CHECK: linalg.generic
 405   // CHECK: arith.cmpi
 406   // CHECK: select
 407   %20 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 408
 409   // CHECK: linalg.generic
 410   // CHECK: arith.constant -32768
 411   // CHECK: arith.constant 32767
 412   // CHECK: arith.cmpi slt
 413   // CHECK: select
 414   // CHECK: arith.cmpi slt
 415   // CHECK: select
 416   // CHECK: arith.trunci
 417   %21 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi16>
 418
 419   // CHECK: linalg.generic
 420   // CHECK: arith.extsi
 421   %22 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi64>
 422
 423   // CHECK: linalg.generic
 424   // CHECK: arith.constant 0
 425   // CHECK: arith.cmpi
 426   %23 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi1>
 427
 428   // CHECK: linalg.generic
 429   // CHECK: arith.sitofp
 430   %24 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xf32>
 431
 432   // CHECK: linalg.generic
 433   // CHECK: arith.constant 0
 434   // CHECK: arith.cmpi sgt
 435   // CHECK: arith.subi
 436   // CHECK: select
 437   %25 = "tosa.abs"(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
 438
 439   return
 440 }
 441
 442 // -----
 443
 444 // CHECK-LABEL: @test_simple_ui8
 445 func @test_simple_ui8(%arg0: tensor<1xi8>) -> () {
 446
 447   // CHECK: linalg.generic
 448   // CHECK: sitofp
 449   %0 = "tosa.cast"(%arg0) : (tensor<1xi8>) -> tensor<1xf32>
 450
 451   return
 452 }
 453
 454 // -----
 455
 456 // CHECK-LABEL: @test_i8
 457 func @test_i8(%arg0: tensor<1xi8>) -> () {
 458   // CHECK: linalg.generic
 459   // CHECK-DAG: %[[C127:.+]] = arith.constant -127
 460   // CHECK-DAG: %[[C126:.+]] = arith.constant 126
 461   // CHECK-DAG: %[[CMP1:.+]] = arith.cmpi slt, %arg1, %[[C127]]
 462   // CHECK-DAG: %[[SEL1:.+]] = select %[[CMP1]], %[[C127]]
 463   // CHECK-DAG: %[[CMP2:.+]] = arith.cmpi slt, %[[C126]], %arg1
 464   // CHECK: %[[SEL2:.+]] = select %[[CMP2]], %[[C126]], %[[SEL1]]
 465   %0 = "tosa.clamp"(%arg0) {min_int = -127 : i64, max_int = 126 : i64, min_fp = 0.0 : f32, max_fp = 0.0 : f32} : (tensor<1xi8>) -> tensor<1xi8>
 466
 467   // CHECK: linalg.generic
 468   // CHECK-DAG: %[[C128:.+]] = arith.constant -128
 469   // CHECK-DAG: %[[C127:.+]] = arith.constant 127
 470   // CHECK-DAG: %[[CMP1:.+]] = arith.cmpi slt, %arg1, %[[C128]]
 471   // CHECK-DAG: %[[SEL1:.+]] = select %[[CMP1]], %[[C128]]
 472   // CHECK-DAG: %[[CMP2:.+]] = arith.cmpi slt, %[[C127]], %arg1
 473   // CHECK: %[[SEL2:.+]] = select %[[CMP2]], %[[C127]], %[[SEL1]]
 474   %1 = "tosa.clamp"(%arg0) {min_int = -130 : i64, max_int = 130 : i64, min_fp = 0.0 : f32, max_fp = 0.0 : f32} : (tensor<1xi8>) -> tensor<1xi8>
 475
 476   return
 477 }
 478
 479 // -----
 480
 481 // CHECK-LABEL: @test_bool
 482 func @test_bool(%arg0: tensor<1xi1>, %arg1: tensor<1xi1>) -> () {
 483   // CHECK: linalg.generic
 484   // CHECK: and
 485   %0 = "tosa.logical_and"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
 486
 487   // CHECK: linalg.generic
 488   // CHECK: or
 489   %1 = "tosa.logical_or"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
 490
 491   // CHECK: linalg.generic
 492   // CHECK: arith.xori
 493   %2 = "tosa.logical_xor"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
 494
 495   // CHECK: linalg.generic
 496   // CHECK: arith.constant true
 497   // CHECK: arith.xori
 498   %3 = "tosa.logical_not"(%arg0) : (tensor<1xi1>) -> tensor<1xi1>
 499
 500   return
 501 }
 502
 503 // -----
 504
 505 // CHECK-LABEL: @test_negate_quantized
 506 func @test_negate_quantized(%arg0: tensor<1xi8>) -> () {
 507   // CHECK: linalg.generic
 508   // CHECK: [[ZERO:%.+]] = arith.constant 0
 509   // CHECK: [[EXT:%.+]] = arith.extsi %arg1 : i8 to i16
 510   // CHECK: [[SUB:%.+]] = arith.subi [[ZERO]], [[EXT]]
 511   // CHECK: [[MIN:%.+]] = arith.constant -128
 512   // CHECK: [[MAX:%.+]] = arith.constant 127
 513   // CHECK: [[PRED1:%.+]] = arith.cmpi slt, [[SUB]], [[MIN]]
 514   // CHECK: [[LBOUND:%.+]] = select [[PRED1]], [[MIN]], [[SUB]]
 515   // CHECK: [[PRED2:%.+]] = arith.cmpi slt, [[MAX]], [[SUB]]
 516   // CHECK: [[UBOUND:%.+]] = select [[PRED2]], [[MAX]], [[LBOUND]]
 517   // CHECK: [[TRUNC:%.+]] = arith.trunci [[UBOUND]]
 518   // CHECK: linalg.yield [[TRUNC]]
 519   %0 = "tosa.negate"(%arg0) {quantization_info = { input_zp = 0 : i32, output_zp = 0 : i32}} : (tensor<1xi8>) -> tensor<1xi8>
 520
 521   // CHECK: linalg.generic
 522   // CHECK: [[EXT:%.+]] = arith.extsi %arg1 : i8 to i16
 523   %1 = "tosa.negate"(%arg0) {quantization_info = { input_zp = 32639 : i32, output_zp = 0 : i32}} : (tensor<1xi8>) -> tensor<1xi8>
 524
 525   // CHECK: linalg.generic
 526   // CHECK: [[EXT:%.+]] = arith.extsi %arg1 : i8 to i32
 527   %2 = "tosa.negate"(%arg0) {quantization_info = { input_zp = 32640 : i32, output_zp = 0 : i32}} : (tensor<1xi8>) -> tensor<1xi8>
 528
 529   return
 530 }
 531
 532 // -----
 533
 534 // CHECK-LABEL: @test_reshape_downrank
 535 func @test_reshape_downrank(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {
 536   // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
 537   %0 = "tosa.reshape"(%arg0) {new_shape = [6]} : (tensor<2x3xf32>) -> tensor<6xf32>
 538   // CHECK: return [[RESHAPE]]
 539   return %0 : tensor<6xf32>
 540 }
 541
 542 // -----
 543
 544 // CHECK-LABEL: @test_reshape_downrank_dyn
 545 func @test_reshape_downrank_dyn(%arg0: tensor<2x?xf32>) -> tensor<?xf32> {
 546   // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
 547   %0 = "tosa.reshape"(%arg0) {new_shape = [-1]} : (tensor<2x?xf32>) -> tensor<?xf32>
 548   // CHECK: return [[RESHAPE]]
 549   return %0 : tensor<?xf32>
 550 }
 551
 552 // -----
 553
 554 // CHECK-LABEL: @test_reshape_uprank
 555 func @test_reshape_uprank(%arg0: tensor<6xf32>) -> tensor<2x3xf32> {
 556   // CHECK: [[RESHAPE:%.+]] = tensor.expand_shape %arg0 {{\[}}[0, 1]]
 557   %0 = "tosa.reshape"(%arg0) {new_shape = [2, 3]} : (tensor<6xf32>) -> tensor<2x3xf32>
 558   // CHECK: return [[RESHAPE]]
 559   return %0 : tensor<2x3xf32>
 560 }
 561
 562 // -----
 563
 564 // CHECK-LABEL: @test_reshape_uprank_dyn
 565 func @test_reshape_uprank_dyn(%arg0: tensor<?xf32>) -> tensor<2x?xf32> {
 566   // CHECK: [[RESHAPE:%.+]] = tensor.expand_shape %arg0 {{\[}}[0, 1]]
 567   %0 = "tosa.reshape"(%arg0) {new_shape = [2, -1]} : (tensor<?xf32>) -> tensor<2x?xf32>
 568   // CHECK: return [[RESHAPE]]
 569   return %0 : tensor<2x?xf32>
 570 }
 571
 572 // -----
 573
 574 // CHECK-LABEL: @test_reshape_samerank
 575 func @test_reshape_samerank(%arg0: tensor<3x2xf32>) -> tensor<2x3xf32> {
 576   // CHECK-SAME: (%[[ARG0:.*]]: tensor<3x2xf32>)
 577   // CHECK-NEXT: %[[RESHAPE1:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]]
 578   // CHECK-NEXT: %[[RESHAPE2:.*]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1]]
 579   %0 = "tosa.reshape"(%arg0) {new_shape = [2, 3]} : (tensor<3x2xf32>) -> tensor<2x3xf32>
 580   // CHECK-NEXT: return %[[RESHAPE2]]
 581   return %0 : tensor<2x3xf32>
 582 }
 583
 584 // -----
 585
 586 // CHECK-LABEL: @test_reshape_samerank_dyn
 587 func @test_reshape_samerank_dyn(%arg0: tensor<?x2xf32>) -> tensor<2x?xf32> {
 588   // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x2xf32>)
 589   // CHECK-NEXT: %[[RESHAPE1:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]]
 590   // CHECK-NEXT: %[[RESHAPE2:.*]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1]]
 591   %0 = "tosa.reshape"(%arg0) {new_shape = [2, -1]} : (tensor<?x2xf32>) -> tensor<2x?xf32>
 592   // CHECK-NEXT: return %[[RESHAPE2]]
 593   return %0 : tensor<2x?xf32>
 594 }
 595
 596 // -----
 597
 598 // CHECK-LABEL: @test_reshape_downrank_6D
 599 func @test_reshape_downrank_6D(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32> {
 600   // CHECK: tensor.collapse_shape %arg0 {{\[}}[0, 1, 2], [3], [4, 5]]
 601   %0 = "tosa.reshape"(%arg0) {new_shape = [6, 5, 77]} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
 602   return %0 : tensor<6x5x77xf32>
 603 }
 604
 605 // -----
 606
 607 // CHECK-LABEL: @test_reshape_downrank_6D_dyn
 608 func @test_reshape_downrank_6D_dyn(%arg0: tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32> {
 609   // CHECK: tensor.collapse_shape %arg0 {{\[}}[0, 1, 2, 3, 4, 5]]
 610   // CHECK: tensor.expand_shape %0 {{\[}}[0, 1, 2]]
 611   %0 = "tosa.reshape"(%arg0) {new_shape = [-1, 5, 77]} : (tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32>
 612   return %0 : tensor<?x5x77xf32>
 613 }
 614
 615 // -----
 616
 617 // CHECK-LABEL: @test_identity
 618 func @test_identity(%arg0: tensor<1xf32>, %arg1: tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>) {
 619   %0 = "tosa.identity"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
 620   %1 = "tosa.identity"(%arg1) : (tensor<1xi32>) -> tensor<1xi32>
 621
 622   // CHECK: return %arg0, %arg1
 623   return %0, %1 : tensor<1xf32>, tensor<1xi32>
 624 }
 625
 626 // -----
 627
 628 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
 629 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 630
 631 // CHECK-LABEL: @test_transpose
 632 // CHECK-SAME: ([[ARG0:%.+]]: tensor<1x2x3xi32>)
 633 func @test_transpose(%arg0: tensor<1x2x3xi32>) -> () {
 634   %0 = arith.constant dense<[1, 2, 0]> : tensor<3xi32>
 635   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3, 1]
 636   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel"]} ins([[ARG0]] : tensor<1x2x3xi32>) outs([[OUT:%.+]] : tensor<2x3x1xi32>)
 637   // CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32)
 638   // CHECK:   linalg.yield [[ARG1]]
 639   // CHECK: }
 640   %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<2x3x1xi32>)
 641   return
 642 }
 643
 644 // -----
 645
 646 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)>
 647 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 648
 649 // CHECK-LABEL: @test_transpose_dyn
 650 // CHECK-SAME: (%[[ARG0:.+]]: tensor<1x?x3x4xi32>)
 651 func @test_transpose_dyn(%arg0: tensor<1x?x3x4xi32>) -> () {
 652   %0 = arith.constant dense<[1, 3, 0, 2]> : tensor<4xi32>
 653   // CHECK: %[[C1:.+]] = arith.constant 1
 654   // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C1]]
 655   // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM]], 4, 1, 3]
 656   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[ARG0]] : tensor<1x?x3x4xi32>) outs([[OUT:%.+]] : tensor<?x4x1x3xi32>)
 657   // CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32)
 658   // CHECK:   linalg.yield [[ARG1]]
 659   // CHECK: }
 660   %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x?x3x4xi32>, tensor<4xi32>) -> (tensor<?x4x1x3xi32>)
 661   return
 662 }
 663
 664 // -----
 665
 666 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1, d0)>
 667 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 668
 669 // CHECK-LABEL: @test_transpose_dyn
 670 // CHECK-SAME: (%[[ARG0:.+]]: tensor<?x?xf32>)
 671 func @test_transpose_dyn_multiple(%arg0: tensor<?x?xf32>) -> () {
 672   %0 = arith.constant dense<[1, 0]> : tensor<2xi32>
 673   // CHECK: %[[C0:.+]] = arith.constant 0
 674   // CHECK: %[[DIM0:.+]] = tensor.dim %arg0, %[[C0]]
 675   // CHECK: %[[C1:.+]] = arith.constant 1
 676   // CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]]
 677   // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM1]], %[[DIM0]]]
 678   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%[[ARG0]] : tensor<?x?xf32>) outs([[OUT:%.+]] : tensor<?x?xf32>)
 679   // CHECK: ^bb0([[ARG1:%.+]]: f32, [[ARG2:%.+]]: f32)
 680   // CHECK:   linalg.yield [[ARG1]]
 681   // CHECK: }
 682   %1 = "tosa.transpose"(%arg0, %0) : (tensor<?x?xf32>, tensor<2xi32>) -> (tensor<?x?xf32>)
 683   return
 684 }
 685
 686 // -----
 687
 688 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 689 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
 690 // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
 691
 692 // CHECK-LABEL: @reduce_float
 693 // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xf32>
 694 func @reduce_float(%arg0: tensor<5x4xf32>) -> () {
 695   // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
 696   // CHECK: [[CST0:%.+]] = arith.constant 0.0
 697   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
 698   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<4xf32>)
 699   // CHECK: ^bb0(%arg1: f32, %arg2: f32)
 700   // CHECK:   [[RES:%.+]] = arith.addf %arg1, %arg2 : f32
 701   // CHECK:   linalg.yield [[RES]] : f32
 702   // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xf32> into tensor<1x4xf32>
 703   %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
 704
 705   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5]
 706   // CHECK: [[CST0:%.+]] = arith.constant 0.0
 707   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
 708   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5xf32>)
 709   // CHECK: ^bb0(%arg1: f32, %arg2: f32)
 710   // CHECK:   [[RES:%.+]] = arith.addf %arg1, %arg2 : f32
 711   // CHECK:   linalg.yield [[RES]] : f32
 712   // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
 713   %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5x1xf32>
 714
 715   // CHECK: arith.constant 1.0
 716   // CHECK: linalg.fill
 717   // CHECK: linalg.generic
 718   // CHECK: arith.mulf
 719   %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
 720
 721   // CHECK: arith.constant 3.40282347E+38 : f32
 722   // CHECK: linalg.fill
 723   // CHECK: linalg.generic
 724   // CHECK: arith.cmpf olt
 725   // CHECK: select
 726   %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
 727
 728   // CHECK: arith.constant -3.40282347E+38 : f32
 729   // CHECK: linalg.fill
 730   // CHECK: linalg.generic
 731   // CHECK: arith.cmpf ogt
 732   // CHECK: select
 733   %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
 734   return
 735 }
 736
 737 // -----
 738
 739 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 740 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
 741 // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
 742
 743 // CHECK-LABEL: @reduce_int
 744 // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi32>
 745 func @reduce_int(%arg0: tensor<5x4xi32>) -> () {
 746   // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
 747   // CHECK: [[CST0:%.+]] = arith.constant 0
 748   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
 749   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<4xi32>)
 750   // CHECK: ^bb0(%arg1: i32, %arg2: i32)
 751   // CHECK:   [[RES:%.+]] = arith.addi %arg1, %arg2 : i32
 752   // CHECK:   linalg.yield [[RES]] : i32
 753   // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi32> into tensor<1x4xi32>
 754   %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
 755
 756   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5]
 757   // CHECK: [[CST0:%.+]] = arith.constant 0
 758   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
 759   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5xi32>)
 760   // CHECK: ^bb0(%arg1: i32, %arg2: i32)
 761   // CHECK:   [[RES:%.+]] = arith.addi %arg1, %arg2 : i32
 762   // CHECK:   linalg.yield [[RES]] : i32
 763   // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xi32> into tensor<5x1xi32>
 764   %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x1xi32>
 765
 766   // CHECK: arith.constant 1
 767   // CHECK: linalg.fill
 768   // CHECK: linalg.generic
 769   // CHECK: arith.muli
 770   %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
 771
 772   // CHECK: arith.constant 2147483647 : i32
 773   // CHECK: linalg.fill
 774   // CHECK: linalg.generic
 775   // CHECK: arith.cmpi slt
 776   // CHECK: select
 777   %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
 778
 779   // CHECK: arith.constant -2147483648 : i32
 780   // CHECK: linalg.fill
 781   // CHECK: linalg.generic
 782   // CHECK: arith.cmpi sgt
 783   // CHECK: select
 784   %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
 785   return
 786 }
 787
 788 // -----
 789
 790 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 791 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
 792
 793 // CHECK-LABEL: @reduce_bool
 794 // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi1>
 795 func @reduce_bool(%arg0: tensor<5x4xi1>) -> () {
 796   // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
 797   // CHECK: [[CST0:%.+]] = arith.constant true
 798   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
 799   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi1>) outs([[FILL]] : tensor<4xi1>)
 800   // CHECK: ^bb0(%arg1: i1, %arg2: i1)
 801   // CHECK:   [[RES:%.+]] = arith.andi %arg1, %arg2 : i1
 802   // CHECK:   linalg.yield [[RES]] : i1
 803   // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi1> into tensor<1x4xi1>
 804   %0 = "tosa.reduce_all"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1>
 805
 806   // CHECK: arith.constant false
 807   // CHECK: linalg.fill
 808   // CHECK: linalg.generic
 809   // CHECK: or
 810   %1 = "tosa.reduce_any"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1>
 811
 812   return
 813 }
 814
 815 // -----
 816
 817 // CHECK-LABEL: @concat
 818 func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
 819   // CHECK: [[AXIS:%.+]] = arith.constant 0
 820   // CHECK: [[STRIDE:%.+]]   = arith.constant 1
 821   // CHECK: [[OFFSET:%.+]] = arith.constant 0 : index
 822   // CHECK: [[IDX0:%.+]] = arith.constant 0 : index
 823   // CHECK: [[IDX1:%.+]] = arith.constant 1 : index
 824   // CHECK: [[INIT:%.+]] = linalg.init_tensor [11, 1]
 825   // CHECK: [[CST:%.+]] = arith.constant 0.0
 826   // CHECK: [[FILL:%.+]] = linalg.fill([[CST]], [[INIT]])
 827   // CHECK: [[INSERT0:%.+]] = tensor.insert_slice %arg0 into [[FILL]][0, 0] [5, 1] [1, 1]
 828   // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %arg1 into [[INSERT0]][5, 0] [6, 1] [1, 1]
 829   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x1xf32>, tensor<6x1xf32>)  -> (tensor<11x1xf32>)
 830
 831   // CHECK: [[AXIS:%.+]] = arith.constant 1
 832   // CHECK: [[STRIDE:%.+]]   = arith.constant 1
 833   // CHECK: [[OFFSET:%.+]] = arith.constant 0 : index
 834   // CHECK: [[IDX0:%.+]] = arith.constant 0 : index
 835   // CHECK: [[IDX1:%.+]] = arith.constant 1 : index
 836   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 2]
 837   // CHECK: [[CST:%.+]] = arith.constant 0.0
 838   // CHECK: [[FILL:%.+]] = linalg.fill([[CST]], [[INIT]])
 839   // CHECK: [[INSERT0:%.+]] = tensor.insert_slice %arg0 into [[FILL]][0, 0] [5, 1] [1, 1]
 840   // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %arg0 into [[INSERT0]][0, 1] [5, 1] [1, 1]
 841   %1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i64} : (tensor<5x1xf32>, tensor<5x1xf32>)  -> (tensor<5x2xf32>)
 842   return
 843 }
 844
 845 // -----
 846
 847 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 848
 849 // CHECK-LABEL: @rescale_i8
 850 func @rescale_i8(%arg0 : tensor<2xi8>) -> () {
 851   // CHECK: [[C0:%.+]] = arith.constant 19689
 852   // CHECK: [[C1:%.+]] = arith.constant 15
 853   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2]
 854   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xi8>) outs([[INIT]] : tensor<2xi8>)
 855   // CHECK: ^bb0([[IN:%.+]]: i8, [[UNUSED:%.+]]: i8):
 856   // CHECK: [[C17:%.+]] = arith.constant 17
 857   // CHECK: [[C22:%.+]] = arith.constant 22
 858   // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
 859   // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]]
 860   // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false}
 861   // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]]
 862   // CHECK-DAG: [[CMIN:%.+]] = arith.constant -128
 863   // CHECK-DAG: [[CMAX:%.+]] = arith.constant 127
 864   // CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]]
 865   // CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]]
 866   // CHECK-DAG: [[LOWER:%.+]] = select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]]
 867   // CHECK-DAG: [[BOUNDED:%.+]] = select [[MAXLT]], [[CMAX]], [[LOWER]]
 868   // CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]]
 869   // CHECK-DAG: linalg.yield [[TRUNC]]
 870   %0 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<2xi8>)  -> (tensor<2xi8>)
 871
 872   // CHECK: [[C0:%.+]] = arith.constant 19689
 873   // CHECK: [[C1:%.+]] = arith.constant 15
 874   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2]
 875   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xi8>) outs([[INIT]] : tensor<2xui8>)
 876   // CHECK: ^bb0([[IN:%.+]]: i8, [[UNUSED:%.+]]: ui8):
 877   // CHECK: [[C17:%.+]] = arith.constant 17
 878   // CHECK: [[C22:%.+]] = arith.constant 22
 879   // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
 880   // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]]
 881   // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false}
 882   // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]]
 883   // CHECK-DAG: [[CMIN:%.+]] = arith.constant 0
 884   // CHECK-DAG: [[CMAX:%.+]] = arith.constant 255
 885   // CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]]
 886   // CHECK-DAG: [[LOWER:%.+]] = select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]]
 887   // CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]]
 888   // CHECK-DAG: [[BOUNDED:%.+]] = select [[MAXLT]], [[CMAX]], [[LOWER]]
 889   // CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]]
 890   // CHECK-DAG: [[CAST:%.+]] = builtin.unrealized_conversion_cast [[TRUNC]] : i8 to ui8
 891   // CHECK: linalg.yield [[CAST]]
 892   %1 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<2xi8>)  -> (tensor<2xui8>)
 893
 894   // CHECK: return
 895   return
 896 }
 897
 898 // -----
 899
 900 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 901
 902 // CHECK-LABEL: @rescale_ui8
 903 func @rescale_ui8(%arg0 : tensor<2xui8>) -> () {
 904   // CHECK: [[C0:%.+]] = arith.constant 19689
 905   // CHECK: [[C1:%.+]] = arith.constant 15
 906   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2]
 907   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xui8>) outs([[INIT]] : tensor<2xi8>)
 908   // CHECK: ^bb0([[IN:%.+]]: ui8, [[UNUSED:%.+]]: i8):
 909   // CHECK: [[C17:%.+]] = arith.constant 17
 910   // CHECK: [[C22:%.+]] = arith.constant 22
 911   // CHECK-DAG: [[CAST:%.+]] = builtin.unrealized_conversion_cast [[IN]] : ui8 to i8
 912   // CHECK-DAG: [[IN32:%.+]] = arith.extui [[CAST]]
 913   // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]]
 914   // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false}
 915   // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]]
 916   // CHECK-DAG: [[CMIN:%.+]] = arith.constant -128
 917   // CHECK-DAG: [[CMAX:%.+]] = arith.constant 127
 918   // CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]]
 919   // CHECK-DAG: [[LOWER:%.+]] = select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]]
 920   // CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]]
 921   // CHECK-DAG: [[BOUNDED:%.+]] = select [[MAXLT]], [[CMAX]], [[LOWER]]
 922   // CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]]
 923   // CHECK: linalg.yield [[TRUNC]]
 924   %0 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<2xui8>)  -> (tensor<2xi8>)
 925
 926   return
 927 }
 928
 929 // -----
 930
 931 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 932
 933 // CHECK-LABEL: @rescale_per_channel
 934 func @rescale_per_channel(%arg0 : tensor<3xi8>) -> (tensor<3xi8>) {
 935   // CHECK: [[MULTIPLIERS:%.+]] = arith.constant dense<[42, 43, 0]>
 936   // CHECK: [[SHIFTS:%.+]] = arith.constant dense<[14, 15, 0]>
 937   // CHECK: [[INIT:%.+]] = linalg.init_tensor [3]
 938   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[MULTIPLIERS]], [[SHIFTS]] : tensor<3xi8>, tensor<3xi32>, tensor<3xi8>) outs([[INIT]] : tensor<3xi8>)
 939   // CHECK: ^bb0([[IN:%.+]]: i8, [[MULTIPLIER:%.+]]: i32, [[SHIFT:%.+]]: i8, [[UNUSED:%.+]]: i8):
 940   // CHECK: [[C243:%.+]] = arith.constant 243
 941   // CHECK: [[C252:%.+]] = arith.constant 252
 942
 943   // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
 944   // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C243]]
 945   // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[MULTIPLIER]], [[SHIFT]]) {double_round = false}
 946   // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C252]]
 947   // CHECK-DAG: [[CMIN:%.+]] = arith.constant -128
 948   // CHECK-DAG: [[CMAX:%.+]] = arith.constant 127
 949   // CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]]
 950   // CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]]
 951   // CHECK-DAG: [[LOWER:%.+]] = select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]]
 952   // CHECK-DAG: [[BOUNDED:%.+]] = select [[MAXLT]], [[CMAX]], [[LOWER]]
 953   // CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]]
 954   // CHECK-DAG: linalg.yield [[TRUNC]]
 955   %0 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [42 : i32, 43 : i32, 44 : i32], shift = [14 : i32, 15 : i32, 64 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<3xi8>)  -> (tensor<3xi8>)
 956
 957   // CHECK: return [[GENERIC]]
 958   return %0 : tensor<3xi8>
 959 }
 960
 961 // -----
 962
 963 // CHECK-LABEL: @rescaleDoubleRound
 964 func @rescaleDoubleRound(%arg0 : tensor<2xi8>) -> (tensor<2xi8>) {
 965   // CHECK: linalg.generic
 966   // CHECK: "tosa.apply_scale"
 967   // CHECK-SAME:  {double_round = true}
 968   %0 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [19689 : i32], shift = [33 : i32], scale32 = true, double_round = true, per_channel = false} : (tensor<2xi8>)  -> (tensor<2xi8>)
 969   return %0 : tensor<2xi8>
 970 }
 971
 972 // CHECK-LABEL: @rescaleUnnecessaryDoubleRound
 973 func @rescaleUnnecessaryDoubleRound(%arg0 : tensor<2xi8>) -> (tensor<2xi8>) {
 974   // CHECK: linalg.generic
 975   // CHECK: "tosa.apply_scale"
 976   // CHECK-SAME:  {double_round = false}
 977   %0 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = true, double_round = true, per_channel = false} : (tensor<2xi8>)  -> (tensor<2xi8>)
 978   return %0 : tensor<2xi8>
 979 }
 980
 981 // -----
 982
 983 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 984
 985 // CHECK-LABEL: @reverse
 986 func @reverse(%arg0: tensor<5x4xi32>) -> () {
 987   // CHECK: %[[C0:.+]] = arith.constant 0
 988   // CHECK: %[[RDIM:.+]] = tensor.dim %arg0, %[[C0]]
 989   // CHECK: %[[INIT:.+]] = linalg.init_tensor [5, 4]
 990   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<5x4xi32>)
 991   // CHECK-DAG:   %[[I0:.+]] = linalg.index 0
 992   // CHECK-DAG:   %[[I1:.+]] = linalg.index 1
 993   // CHECK-DAG:   %[[SUB1:.+]] = arith.constant 1
 994   // CHECK-DAG:   %[[RDIM_MINUS_C1:.+]] = arith.subi %[[RDIM]], %[[SUB1]]
 995   // CHECK-DAG:   %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I0]]
 996   // CHECK-DAG:   %[[EXTRACT:.+]] = tensor.extract %arg0[%[[READ_DIM]], %[[I1]]] : tensor<5x4xi32>
 997   // CHECK:   linalg.yield %[[EXTRACT]]
 998   %0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
 999
1000   // CHECK: %[[C1:.+]] = arith.constant 1
1001   // CHECK: %[[RDIM:.+]] = tensor.dim %arg0, %[[C1]]
1002   // CHECK: %[[INIT:.+]] = linalg.init_tensor [5, 4]
1003   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<5x4xi32>)
1004   // CHECK-DAG:   %[[I0:.+]] = linalg.index 0
1005   // CHECK-DAG:   %[[I1:.+]] = linalg.index 1
1006   // CHECK-DAG:   %[[SUB1:.+]] = arith.constant 1
1007   // CHECK-DAG:   %[[RDIM_MINUS_C1:.+]] = arith.subi %[[RDIM]], %[[SUB1]]
1008   // CHECK-DAG:   %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I1]]
1009   // CHECK-DAG:   %[[EXTRACT:.+]] = tensor.extract %arg0[%[[I0]], %[[READ_DIM]]] : tensor<5x4xi32>
1010   // CHECK:   linalg.yield %[[EXTRACT]]
1011   %1 = "tosa.reverse"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
1012   return
1013 }
1014
1015 // -----
1016
1017 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
1018
1019 // CHECK-LABEL: @reverse_dyn
1020 func @reverse_dyn(%arg0: tensor<?xi32>) -> () {
1021   // CHECK: %[[C0_1:.+]] = arith.constant 0
1022   // CHECK: %[[D0_1:.+]] = tensor.dim %arg0, %[[C0_1]]
1023   // CHECK: %[[C0_2:.+]] = arith.constant 0
1024   // CHECK: %[[D0_2:.+]] = tensor.dim %arg0, %[[C0_2]]
1025   // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[D0_1]]]
1026   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel"]} outs(%[[INIT]] : tensor<?xi32>)
1027   // CHECK-DAG:   %[[I0:.+]] = linalg.index 0
1028   // CHECK-DAG:   %[[SUB1:.+]] = arith.constant 1
1029   // CHECK-DAG:   %[[RDIM_MINUS_C1:.+]] = arith.subi %[[D0_2]], %[[SUB1]]
1030   // CHECK-DAG:   %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I0]]
1031   // CHECK-DAG:   %[[EXTRACT:.+]] = tensor.extract %arg0[%[[READ_DIM]]] : tensor<?xi32>
1032   // CHECK:   linalg.yield %[[EXTRACT]]
1033   %0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<?xi32>) -> tensor<?xi32>
1034   return
1035 }
1036
1037 // -----
1038
1039 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
1040 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
1041
1042 // CHECK-LABEL: @tile
1043 func @tile(%arg0 : tensor<2x3xi8>) -> () {
1044   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 2, 1, 3]
1045   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<2x2x1x3xi8>)
1046   // CHECK:   linalg.yield %arg1 : i8
1047   // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1, 2], [3]]
1048   %0 = "tosa.tile"(%arg0) {multiples = [2, 1]} : (tensor<2x3xi8>)  -> (tensor<4x3xi8>)
1049
1050   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 2, 2, 3]
1051   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<1x2x2x3xi8>)
1052   // CHECK:   linalg.yield %arg1 : i8
1053   // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
1054   %1 = "tosa.tile"(%arg0) {multiples = [1, 2]} : (tensor<2x3xi8>)  -> (tensor<2x6xi8>)
1055
1056   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 2, 7, 3]
1057   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<5x2x7x3xi8>)
1058   // CHECK:   linalg.yield %arg1 : i8
1059   // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
1060   %2 = "tosa.tile"(%arg0) {multiples = [5, 7]} : (tensor<2x3xi8>)  -> (tensor<10x21xi8>)
1061
1062   return
1063 }
1064
1065 // -----
1066
1067 func @pad_float(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
1068   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
1069   // TODO: Output contains multiple "arith.constant 1 : index".
1070   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
1071   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
1072   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
1073   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
1074   // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
1075   // CHECK: linalg.pad_tensor %arg0 low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
1076   // CHECK: ^bb0(%arg1: index, %arg2: index):  // no predecessors
1077   // CHECK:   linalg.yield [[CST]]
1078   // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
1079   %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xf32>, tensor<2x2xi32>)  -> (tensor<4x9xf32>)
1080   return %1 : tensor<4x9xf32>
1081 }
1082
1083 func @pad_int(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
1084   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
1085   // CHECK: [[CST:%.+]] = arith.constant 0 : i32
1086   // CHECK: linalg.pad_tensor
1087   // CHECK:   linalg.yield [[CST]]
1088   %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xi32>, tensor<2x2xi32>)  -> (tensor<4x9xi32>)
1089   return %1 : tensor<4x9xi32>
1090 }
1091
1092 func @pad_quant(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
1093   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
1094   // CHECK: [[CST:%.+]] = arith.constant 42 : i32
1095   // CHECK: linalg.pad_tensor
1096   // CHECK:   linalg.yield [[CST]]
1097   %1 = "tosa.pad"(%arg0, %0) { quantization_info = { input_zp = 42 : i32}} : (tensor<1x2xi32>, tensor<2x2xi32>)  -> (tensor<4x9xi32>)
1098   return %1 : tensor<4x9xi32>
1099 }
1100
1101 // -----
1102
1103 func @pad_float_explicit(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
1104   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
1105   // TODO: Output contains multiple "arith.constant 1 : index".
1106   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
1107   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
1108   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
1109   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
1110   // CHECK-DAG: [[CST:%.+]] = arith.constant 4.200000e+01 : f32
1111   // CHECK: linalg.pad_tensor %arg0 low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
1112   // CHECK: ^bb0(%arg1: index, %arg2: index):  // no predecessors
1113   // CHECK:   linalg.yield [[CST]]
1114   // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
1115   %1 = arith.constant dense<42.0> : tensor<f32>
1116   %2 = "tosa.pad"(%arg0, %0, %1)  : (tensor<1x2xf32>, tensor<2x2xi32>, tensor<f32>)  -> (tensor<4x9xf32>)
1117   return %2 : tensor<4x9xf32>
1118 }
1119
1120 // -----
1121
1122 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
1123 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
1124 // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
1125 // CHECK: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
1126 // CHECK: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
1127
1128 func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
1129   // CHECK: [[IDX_INIT:%.+]] = linalg.init_tensor [2]
1130   // CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
1131   // CHECK: [[IDX_FILL:%.+]] = linalg.fill([[IDX_MIN]], [[IDX_INIT]])
1132   // CHECK: [[VAL_INIT:%.+]] = linalg.init_tensor [2]
1133   // CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648
1134   // CHECK: [[VAL_FILL:%.+]] = linalg.fill([[VAL_MIN]], [[VAL_INIT]])
1135   // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%arg0 : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<2xi32>, tensor<2xi32>)
1136   // CHECK:   [[IDX:%.+]] = linalg.index 0
1137   // CHECK:   [[CAST:%.+]] = arith.index_cast [[IDX]]
1138   // CHECK:   [[CMP:%.+]] = arith.cmpi sgt, %arg2, %arg4
1139   // CHECK:   [[SELECT_VAL:%.+]] = select [[CMP]], %arg2, %arg4
1140   // CHECK:   [[SELECT_IDX:%.+]] = select [[CMP]], [[CAST]], %arg3
1141   // CHECK:   linalg.yield [[SELECT_IDX]], [[SELECT_VAL]]
1142   %0 = "tosa.argmax"(%arg0) { axis = 0 : i64} : (tensor<3x2xi32>)  -> (tensor<2xi32>)
1143
1144   // CHECK: [[IDX_INIT:%.+]] = linalg.init_tensor [3]
1145   // CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
1146   // CHECK: [[IDX_FILL:%.+]] = linalg.fill([[IDX_MIN]], [[IDX_INIT]])
1147   // CHECK: [[VAL_INIT:%.+]] = linalg.init_tensor [3]
1148   // CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648
1149   // CHECK: [[VAL_FILL:%.+]] = linalg.fill([[VAL_MIN]], [[VAL_INIT]])
1150   // CHECK: linalg.generic {indexing_maps = [#map0, #map2, #map2], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
1151   // CHECK:   [[IDX:%.+]] = linalg.index 1
1152   // CHECK:   [[CAST:%.+]] = arith.index_cast [[IDX]]
1153   // CHECK:   [[CMP:%.+]] = arith.cmpi sgt, %arg2, %arg4
1154   // CHECK:   [[SELECT_VAL:%.+]] = select [[CMP]], %arg2, %arg4
1155   // CHECK:   [[SELECT_IDX:%.+]] = select [[CMP]], [[CAST]], %arg3
1156   // CHECK:   linalg.yield [[SELECT_IDX]], [[SELECT_VAL]]
1157   %1 = "tosa.argmax"(%arg0) { axis = 1 : i64} : (tensor<3x2xi32>)  -> (tensor<3xi32>)
1158
1159   // CHECK: arith.constant -3.40282347E+38 : f32
1160   // CHECK: linalg.index
1161   // CHECK: arith.index_cast
1162   // CHECK: arith.cmpf ogt
1163   // CHECK: select
1164   // CHECK: select
1165   // CHECK: linalg.yield
1166   %2 = "tosa.argmax"(%arg1) { axis = 0 : i64} : (tensor<6xf32>)  -> (tensor<i32>)
1167
1168   return
1169 }
1170
1171 // -----
1172
1173 // CHECK-LABEL: @gather_float
1174 func @gather_float(%arg0: tensor<2x3x2xf32>, %arg1: tensor<2x3xi32>) -> () {
1175   // CHECK: %[[INIT:.+]] = linalg.init_tensor [2, 3, 2]
1176   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xf32>)
1177   // CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: f32)
1178   // CHECK:   %[[IDX0:.+]] = linalg.index 0
1179   // CHECK:   %[[CAST:.+]] = arith.index_cast %[[ARG0]]
1180   // CHECK:   %[[IDX2:.+]] = linalg.index 2
1181   // CHECK:   %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xf32>
1182   // CHECK:   linalg.yield %[[EXTRACT]]
1183   %0 = "tosa.gather"(%arg0, %arg1)  : (tensor<2x3x2xf32>, tensor<2x3xi32>)  -> (tensor<2x3x2xf32>)
1184   return
1185 }
1186
1187 // CHECK-LABEL: @gather_int
1188 func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () {
1189   // CHECK: %[[INIT:.+]] = linalg.init_tensor [2, 3, 2]
1190   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xi32>)
1191   // CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: i32)
1192   // CHECK:   %[[IDX0:.+]] = linalg.index 0
1193   // CHECK:   %[[CAST:.+]] = arith.index_cast %[[ARG0]]
1194   // CHECK:   %[[IDX2:.+]] = linalg.index 2
1195   // CHECK:   %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xi32>
1196   // CHECK:   linalg.yield %[[EXTRACT]]
1197   %0 = "tosa.gather"(%arg0, %arg1)  : (tensor<2x3x2xi32>, tensor<2x3xi32>)  -> (tensor<2x3x2xi32>)
1198   return
1199 }
1200
1201 // -----
1202
1203 // CHECK-LABEL: @table8
1204 func @table8(%arg0: tensor<6xi8>, %arg1: tensor<512xi8>) -> () {
1205   // CHECK: %[[INIT:.+]] = linalg.init_tensor [6]
1206   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi8>) outs(%[[INIT]] : tensor<6xi8>)
1207   // CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8)
1208   // CHECK:   %[[CAST:.+]] = arith.index_cast %[[ARG_IN]]
1209   // CHECK:   %[[OFFSET:.+]] = arith.constant 128
1210   // CHECK:   %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]]
1211   // CHECK:   %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]]
1212   // CHECK:   linalg.yield %[[EXTRACT]]
1213   %0 = "tosa.table"(%arg0, %arg1)  : (tensor<6xi8>, tensor<512xi8>)  -> (tensor<6xi8>)
1214   return
1215 }
1216
1217 // -----
1218
1219 // CHECK-LABEL: @table16
1220 func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () {
1221   // CHECK: %[[INIT:.+]] = linalg.init_tensor [6]
1222   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi16>) outs(%[[INIT]] : tensor<6xi32>)
1223   // CHECK: ^bb0(%arg2: i16, %arg3: i32)
1224   // CHECK: %[[EXT_IN:.+]] = arith.extsi %arg2
1225   // CHECK: %[[C32768:.+]] = arith.constant 32768
1226   // CHECK: %[[C7:.+]] = arith.constant 7
1227   // CHECK: %[[C1:.+]] = arith.constant 1
1228   // CHECK: %[[C127:.+]] = arith.constant 127
1229   // CHECK: %[[INADD:.+]] = arith.addi %[[EXT_IN]], %[[C32768]]
1230   // CHECK: %[[IDX:.+]] = arith.shrui %[[INADD]], %[[C7]]
1231   // CHECK: %[[FRACTION:.+]] = arith.andi %[[INADD]], %[[C127]]
1232   // CHECK: %[[IDXPLUS1:.+]] = arith.addi %[[IDX]], %[[C1]]
1233   // CHECK: %[[IDX_CAST:.+]] = arith.index_cast %[[IDX]]
1234   // CHECK: %[[IDXPLUS1_CAST:.+]] = arith.index_cast %[[IDXPLUS1]]
1235   // CHECK: %[[BASE:.+]] = tensor.extract %arg1[%[[IDX_CAST]]]
1236   // CHECK: %[[NEXT:.+]] = tensor.extract %arg1[%[[IDXPLUS1_CAST]]]
1237   // CHECK: %[[BASE_EXT:.+]] = arith.extsi %[[BASE]]
1238   // CHECK: %[[NEXT_EXT:.+]] = arith.extsi %[[NEXT]]
1239   // CHECK: %[[BASE_MUL:.+]] = arith.shli %[[BASE_EXT]], %[[C7]]
1240   // CHECK: %[[DIFF:.+]] = arith.subi %[[NEXT_EXT]], %[[BASE_EXT]]
1241   // CHECK: %[[DIFF_MUL:.+]] = arith.muli %[[DIFF]], %[[FRACTION]]
1242   // CHECK: %[[RESULT:.+]] = arith.addi %[[BASE_MUL]], %[[DIFF_MUL]]
1243   // CHECK: linalg.yield %[[RESULT]]
1244   %0 = "tosa.table"(%arg0, %arg1)  : (tensor<6xi16>, tensor<513xi16>)  -> (tensor<6xi32>)
1245   return
1246 }
1247
1248 // -----
1249
1250 // CHECK-LABEL: @resize_nearest
1251 func @resize_nearest(%input: tensor<1x2x2x1xf32>) -> () {
1252   // CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1]
1253   // CHECK: %[[GENERIC:.+]] = linalg.generic
1254   // CHECK: %[[IDX0:.+]] = linalg.index 0
1255   // CHECK: %[[IDX1:.+]] = linalg.index 1
1256   // CHECK: %[[IDX2:.+]] = linalg.index 2
1257   // CHECK: %[[IDX3:.+]] = linalg.index 3
1258   // CHECK-DAG: %[[XYMIN:.+]] = arith.constant 0
1259   // CHECK-DAG: %[[YMAX:.+]] = arith.constant 1
1260   // CHECK-DAG: %[[XMAX:.+]] = arith.constant 1
1261   // CHECK-DAG: %[[Y:.+]] = arith.index_cast %[[IDX1]]
1262   // CHECK-DAG: %[[X:.+]] = arith.index_cast %[[IDX2]]
1263   // CHECK-DAG: %[[STRIDEY:.+]] = arith.constant 5.000000e-01
1264   // CHECK-DAG: %[[STRIDEX:.+]] = arith.constant 5.000000e-01
1265   // CHECK-DAG: %[[OFFSETY:.+]] = arith.constant 1.000000e-01
1266   // CHECK-DAG: %[[OFFSETX:.+]] = arith.constant 2.000000e-01
1267   // CHECK-DAG: %[[VAL4:.+]] = arith.uitofp %[[Y]]
1268   // CHECK-DAG: %[[VAL5:.+]] = arith.uitofp %[[X]]
1269   // CHECK-DAG: %[[VAL6:.+]] = arith.mulf %[[VAL4]], %[[STRIDEY]]
1270   // CHECK-DAG: %[[VAL7:.+]] = arith.mulf %[[VAL5]], %[[STRIDEX]]
1271   // CHECK-DAG: %[[VAL8:.+]] = arith.addf %[[VAL6]], %[[OFFSETY]]
1272   // CHECK-DAG: %[[VAL9:.+]] = arith.addf %[[VAL7]], %[[OFFSETX]]
1273
1274   // Find the remainder and integer component of the target index.
1275
1276   // CHECK-DAG: %[[VAL10:.+]] = math.floor %[[VAL8]]
1277   // CHECK-DAG: %[[VAL11:.+]] = math.floor %[[VAL9]]
1278   // CHECK-DAG: %[[VAL12:.+]] = arith.subf %[[VAL8]], %[[VAL10]]
1279   // CHECK-DAG: %[[VAL13:.+]] = arith.subf %[[VAL9]], %[[VAL11]]
1280   // CHECK-DAG: %[[VAL14:.+]] = arith.fptosi %[[VAL10]]
1281   // CHECK-DAG: %[[VAL15:.+]] = arith.fptosi %[[VAL11]]
1282
1283   // Round to the nearest index.
1284
1285   // CHECK-DAG: %[[ROUND:.+]] = arith.constant 5.000000e-01
1286   // CHECK-DAG: %[[VAL16:.+]] = arith.cmpf oge, %[[VAL12]], %[[ROUND]]
1287   // CHECK-DAG: %[[VAL17:.+]] = arith.cmpf oge, %[[VAL13]], %[[ROUND]]
1288   // CHECK-DAG: %[[ZERO:.+]] = arith.constant 0
1289   // CHECK-DAG: %[[ONE:.+]] = arith.constant 1
1290   // CHECK-DAG: %[[VAL18:.+]] = select %[[VAL16]], %[[ONE]], %[[ZERO]]
1291   // CHECK-DAG: %[[VAL19:.+]] = select %[[VAL17]], %[[ONE]], %[[ZERO]]
1292   // CHECK-DAG: %[[VAL20:.+]] = arith.addi %[[VAL14]], %[[VAL18]]
1293   // CHECK-DAG: %[[VAL21:.+]] = arith.addi %[[VAL15]], %[[VAL19]]
1294
1295   // This section applies bound checking to be within the input image.
1296
1297   // CHECK-DAG: %[[VAL22:.+]] = arith.cmpi slt, %[[VAL20]], %[[XYMIN]]
1298   // CHECK-DAG: %[[VAL23:.+]] = select %[[VAL22]], %[[XYMIN]], %[[VAL20]]
1299   // CHECK-DAG: %[[VAL24:.+]] = arith.cmpi slt, %[[YMAX]], %[[VAL20]]
1300   // CHECK-DAG: %[[VAL25:.+]] = select %[[VAL24]], %[[YMAX]], %[[VAL23]]
1301   // CHECK-DAG: %[[VAL26:.+]] = arith.cmpi slt, %[[VAL21]], %[[XYMIN]]
1302   // CHECK-DAG: %[[VAL27:.+]] = select %[[VAL26]], %[[XYMIN]], %[[VAL21]]
1303   // CHECK-DAG: %[[VAL28:.+]] = arith.cmpi slt, %[[XMAX]], %[[VAL21]]
1304   // CHECK-DAG: %[[VAL29:.+]] = select %[[VAL28]], %[[XMAX]], %[[VAL27]]
1305
1306   // Extract the nearest value using the computed indices.
1307
1308   // CHECK-DAG: %[[IDY:.+]] = arith.index_cast %[[VAL25]]
1309   // CHECK-DAG: %[[IDX:.+]] = arith.index_cast %[[VAL29]]
1310   // CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[IDY]], %[[IDX]], %[[IDX3]]]
1311   // CHECK: linalg.yield %[[EXTRACT]]
1312   %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [0, 0], offset = [0, 0], stride_fp = [0.5 : f32, 0.5 : f32], offset_fp = [0.1 : f32, 0.2 : f32], shift = 0 : i32, mode = "NEAREST_NEIGHBOR" } : (tensor<1x2x2x1xf32>)  -> (tensor<1x4x4x1xf32>)
1313
1314   return
1315 }
1316
1317 // -----
1318
1319 // CHECK-LABEL: @resize_bilinear
1320 func @resize_bilinear(%input: tensor<1x2x2x1xf32>) -> () {
1321   // CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1]
1322   // CHECK: %[[GENERIC:.+]] = linalg.generic
1323   // CHECK: %[[IDX0:.+]] = linalg.index 0
1324   // CHECK: %[[IDX1:.+]] = linalg.index 1
1325   // CHECK: %[[IDX2:.+]] = linalg.index 2
1326   // CHECK: %[[IDX3:.+]] = linalg.index 3
1327   // CHECK: %[[XYMIN:.+]] = arith.constant 0
1328   // CHECK: %[[YMAX:.+]] = arith.constant 1
1329   // CHECK: %[[XMAX:.+]] = arith.constant 1
1330
1331   // CHECK: %[[VAL10:.+]] = math.floor %[[VAL8:.+]]
1332   // CHECK: %[[VAL11:.+]] = math.floor %[[VAL9:.+]]
1333
1334   // CHECK: %[[DY:.+]] = arith.subf %[[VAL8:.+]], %[[VAL10]]
1335   // CHECK: %[[DX:.+]] = arith.subf %[[VAL9:.+]], %[[VAL11]]
1336
1337   // CHECK: %[[Y0:.+]] = arith.fptosi %[[VAL10]]
1338   // CHECK: %[[X0:.+]] = arith.fptosi %[[VAL11]]
1339
1340   // Compute the left, right, and top indices for the bilinear interpolation.
1341
1342   // CHECK: %[[ONE:.+]] = arith.constant 1
1343   // CHECK: %[[Y1:.+]] = arith.addi %[[Y0]], %[[ONE]]
1344   // CHECK: %[[X1:.+]] = arith.addi %[[X0]], %[[ONE]]
1345
1346   // Bound check each dimension.
1347
1348   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[Y0]], %[[XYMIN]]
1349   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[Y0]]
1350   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[YMAX]], %[[Y0]]
1351   // CHECK: %[[YLO:.+]] = select %[[PRED]], %[[YMAX]], %[[BOUND]]
1352
1353   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[Y1]], %[[XYMIN]]
1354   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[Y1]]
1355   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[YMAX]], %[[Y1]]
1356   // CHECK: %[[YHI:.+]] = select %[[PRED]], %[[YMAX]], %[[BOUND]]
1357
1358   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[X0]], %[[XYMIN]]
1359   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[X0]]
1360   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[XMAX]], %[[X0]]
1361   // CHECK: %[[XLO:.+]] = select %[[PRED]], %[[XMAX]], %[[BOUND]]
1362
1363   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[X1]], %[[XYMIN]]
1364   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[X1]]
1365   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[XMAX]], %[[X1]]
1366   // CHECK: %[[XHI:.+]] = select %[[PRED]], %[[XMAX]], %[[BOUND]]
1367
1368   // Extract each corner of the bilinear interpolation.
1369
1370   // CHECK: %[[YLOI:.+]] = arith.index_cast %[[YLO]]
1371   // CHECK: %[[YHII:.+]] = arith.index_cast %[[YHI]]
1372   // CHECK: %[[XLOI:.+]] = arith.index_cast %[[XLO]]
1373   // CHECK: %[[XHII:.+]] = arith.index_cast %[[XHI]]
1374
1375   // CHECK: %[[LOLO:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YLOI]], %[[XLOI]], %[[IDX3]]]
1376   // CHECK: %[[LOHI:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YLOI]], %[[XHII]], %[[IDX3]]]
1377   // CHECK: %[[HILO:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YHII]], %[[XLOI]], %[[IDX3]]]
1378   // CHECK: %[[HIHI:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YHII]], %[[XHII]], %[[IDX3]]]
1379
1380   // Compute the bilinear interpolation.
1381
1382   // CHECK: %[[ONE:.+]] = arith.constant 1.000000e+00
1383   // CHECK: %[[NDX:.+]] = arith.subf %[[ONE]], %[[DX]]
1384   // CHECK: %[[WLOLO:.+]] = arith.mulf %[[LOLO]], %[[NDX]]
1385   // CHECK: %[[WLOHI:.+]] = arith.mulf %[[LOHI]], %[[DX]]
1386   // CHECK: %[[LO:.+]] = arith.addf %[[WLOLO]], %[[WLOHI]]
1387   // CHECK: %[[WHILO:.+]] = arith.mulf %[[HILO]], %[[NDX]]
1388   // CHECK: %[[WHIHI:.+]] = arith.mulf %[[HIHI]], %[[DX]]
1389   // CHECK: %[[HI:.+]] = arith.addf %[[WHILO]], %[[WHIHI]]
1390   // CHECK: %[[NDY:.+]] = arith.subf %[[ONE]], %[[DY]]
1391   // CHECK: %[[WLO:.+]] = arith.mulf %[[LO]], %[[NDY]]
1392   // CHECK: %[[WHI:.+]] = arith.mulf %[[HI]], %[[DY]]
1393   // CHECK: %[[RESULT:.+]] = arith.addf %[[WLO]], %[[WHI]]
1394   // CHECK: linalg.yield %[[RESULT]]
1395   %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [0, 0], offset = [0, 0], stride_fp = [0.5 : f32, 0.5 : f32], offset_fp = [0.1 : f32, 0.2 : f32], shift = 0 : i32, mode = "BILINEAR" } : (tensor<1x2x2x1xf32>)  -> (tensor<1x4x4x1xf32>)
1396   return
1397 }
1398
1399 // -----
1400
1401 // CHECK-LABEL: @resize_nearest_int
1402 func @resize_nearest_int(%input: tensor<1x2x2x1xi32>) -> () {
1403   // CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1]
1404   // CHECK: %[[GENERIC:.+]] = linalg.generic
1405   // CHECK: %[[IDX0:.+]] = linalg.index 0
1406   // CHECK: %[[IDX1:.+]] = linalg.index 1
1407   // CHECK: %[[IDX2:.+]] = linalg.index 2
1408   // CHECK: %[[IDX3:.+]] = linalg.index 3
1409   // CHECK-DAG: %[[XYMIN:.+]] = arith.constant 0
1410   // CHECK-DAG: %[[YMAX:.+]] = arith.constant 1
1411   // CHECK-DAG: %[[XMAX:.+]] = arith.constant 1
1412   // CHECK-DAG: %[[Y:.+]] = arith.index_cast %[[IDX1]]
1413   // CHECK-DAG: %[[X:.+]] = arith.index_cast %[[IDX2]]
1414   // CHECK-DAG: %[[STRIDEY:.+]] = arith.constant 128
1415   // CHECK-DAG: %[[STRIDEX:.+]] = arith.constant 128
1416   // CHECK-DAG: %[[OFFSETY:.+]] = arith.constant 1
1417   // CHECK-DAG: %[[OFFSETX:.+]] = arith.constant 2
1418   // CHECK-DAG: %[[EIGHT:.+]] = arith.constant 8
1419   // CHECK-DAG: %[[VAL4:.+]] = arith.muli %[[Y]], %[[STRIDEY]]
1420   // CHECK-DAG: %[[VAL5:.+]] = arith.muli %[[X]], %[[STRIDEX]]
1421   // CHECK-DAG: %[[VAL6:.+]] = arith.addi %[[VAL4]], %[[OFFSETY]]
1422   // CHECK-DAG: %[[VAL7:.+]] = arith.addi %[[VAL5]], %[[OFFSETX]]
1423
1424   // Find the remainder and integer component of the target index.
1425
1426
1427   // CHECK-DAG: %[[VAL8:.+]] = arith.shrsi %[[VAL6]], %[[EIGHT]]
1428   // CHECK-DAG: %[[VAL9:.+]] = arith.shrsi %[[VAL7]], %[[EIGHT]]
1429   // CHECK-DAG: %[[VAL10:.+]] = arith.shli %[[VAL8]], %[[EIGHT]]
1430   // CHECK-DAG: %[[VAL11:.+]] = arith.shli %[[VAL9]], %[[EIGHT]]
1431   // CHECK-DAG: %[[VAL12:.+]] = arith.subi %[[VAL6]], %[[VAL10]]
1432   // CHECK-DAG: %[[VAL13:.+]] = arith.subi %[[VAL7]], %[[VAL11]]
1433
1434   // Round to the nearest index.
1435
1436   // CHECK-DAG: %[[ROUND:.+]] = arith.constant 128
1437   // CHECK-DAG: %[[VAL16:.+]] = arith.cmpi sge, %[[VAL12]], %[[ROUND]]
1438   // CHECK-DAG: %[[VAL17:.+]] = arith.cmpi sge, %[[VAL13]], %[[ROUND]]
1439   // CHECK-DAG: %[[ZERO:.+]] = arith.constant 0
1440   // CHECK-DAG: %[[ONE:.+]] = arith.constant 1
1441   // CHECK-DAG: %[[VAL18:.+]] = select %[[VAL16]], %[[ONE]], %[[ZERO]]
1442   // CHECK-DAG: %[[VAL19:.+]] = select %[[VAL17]], %[[ONE]], %[[ZERO]]
1443   // CHECK-DAG: %[[VAL20:.+]] = arith.addi %[[VAL8]], %[[VAL18]]
1444   // CHECK-DAG: %[[VAL21:.+]] = arith.addi %[[VAL9]], %[[VAL19]]
1445
1446   // This section applies bound checking to be within the input image.
1447
1448   // CHECK-DAG: %[[VAL22:.+]] = arith.cmpi slt, %[[VAL20]], %[[XYMIN]]
1449   // CHECK-DAG: %[[VAL23:.+]] = select %[[VAL22]], %[[XYMIN]], %[[VAL20]]
1450   // CHECK-DAG: %[[VAL24:.+]] = arith.cmpi slt, %[[YMAX]], %[[VAL20]]
1451   // CHECK-DAG: %[[VAL25:.+]] = select %[[VAL24]], %[[YMAX]], %[[VAL23]]
1452   // CHECK-DAG: %[[VAL26:.+]] = arith.cmpi slt, %[[VAL21]], %[[XYMIN]]
1453   // CHECK-DAG: %[[VAL27:.+]] = select %[[VAL26]], %[[XYMIN]], %[[VAL21]]
1454   // CHECK-DAG: %[[VAL28:.+]] = arith.cmpi slt, %[[XMAX]], %[[VAL21]]
1455   // CHECK-DAG: %[[VAL29:.+]] = select %[[VAL28]], %[[XMAX]], %[[VAL27]]
1456
1457   // Extract the nearest value using the computed indices.
1458
1459   // CHECK-DAG: %[[IDY:.+]] = arith.index_cast %[[VAL25]]
1460   // CHECK-DAG: %[[IDX:.+]] = arith.index_cast %[[VAL29]]
1461   // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[IDY]], %[[IDX]], %[[IDX3]]]
1462   // CHECK: linalg.yield %[[EXTRACT]]
1463   %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [128, 128], offset = [1, 2], stride_fp = [0. : f32, 0. : f32], offset_fp = [0. : f32, 0. : f32], shift = 8 : i32, mode = "NEAREST_NEIGHBOR" } : (tensor<1x2x2x1xi32>)  -> (tensor<1x4x4x1xi32>)
1464   return
1465 }
1466
1467 // -----
1468
1469 // CHECK-LABEL: @resize_bilinear_int
1470 func @resize_bilinear_int(%input: tensor<1x2x2x1xi8>) -> () {
1471   // CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1]
1472   // CHECK: %[[GENERIC:.+]] = linalg.generic
1473
1474   // CHECK: %[[IDX0:.+]] = linalg.index 0
1475   // CHECK: %[[IDX3:.+]] = linalg.index 3
1476
1477   // CHECK: %[[XYMIN:.+]] = arith.constant 0
1478   // CHECK: %[[YMAX:.+]] = arith.constant 1
1479   // CHECK: %[[XMAX:.+]] = arith.constant 1
1480
1481   // CHECK: %[[Y0:.+]] = arith.shrsi
1482   // CHECK: %[[X0:.+]] = arith.shrsi
1483   // CHECK: %[[ROUNDY:.+]] = arith.shli %[[Y0]]
1484   // CHECK: %[[ROUNDX:.+]] = arith.shli %[[X0]]
1485   // CHECK: %[[DY:.+]] = arith.subi %10, %[[ROUNDY]]
1486   // CHECK: %[[DX:.+]] = arith.subi %11, %[[ROUNDX]]
1487
1488   // Compute the left, right, and top indices for the bilinear interpolation.
1489
1490   // CHECK: %[[ONE:.+]] = arith.constant 1
1491   // CHECK: %[[Y1:.+]] = arith.addi %[[Y0]], %[[ONE]]
1492   // CHECK: %[[X1:.+]] = arith.addi %[[X0]], %[[ONE]]
1493
1494   // Bound check each dimension.
1495
1496   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[Y0]], %[[XYMIN]]
1497   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[Y0]]
1498   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[YMAX]], %[[Y0]]
1499   // CHECK: %[[YLO:.+]] = select %[[PRED]], %[[YMAX]], %[[BOUND]]
1500
1501   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[Y1]], %[[XYMIN]]
1502   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[Y1]]
1503   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[YMAX]], %[[Y1]]
1504   // CHECK: %[[YHI:.+]] = select %[[PRED]], %[[YMAX]], %[[BOUND]]
1505
1506   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[X0]], %[[XYMIN]]
1507   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[X0]]
1508   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[XMAX]], %[[X0]]
1509   // CHECK: %[[XLO:.+]] = select %[[PRED]], %[[XMAX]], %[[BOUND]]
1510
1511   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[X1]], %[[XYMIN]]
1512   // CHECK: %[[BOUND:.+]] = select %[[PRED]], %[[XYMIN]], %[[X1]]
1513   // CHECK: %[[PRED:.+]] = arith.cmpi slt, %[[XMAX]], %[[X1]]
1514   // CHECK: %[[XHI:.+]] = select %[[PRED]], %[[XMAX]], %[[BOUND]]
1515
1516   // Extract each corner of the bilinear interpolation.
1517
1518   // CHECK: %[[YLOI:.+]] = arith.index_cast %[[YLO]]
1519   // CHECK: %[[YHII:.+]] = arith.index_cast %[[YHI]]
1520   // CHECK: %[[XLOI:.+]] = arith.index_cast %[[XLO]]
1521   // CHECK: %[[XHII:.+]] = arith.index_cast %[[XHI]]
1522
1523   // CHECK: %[[LOLO:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YLOI]], %[[XLOI]], %[[IDX3]]]
1524   // CHECK: %[[LOHI:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YLOI]], %[[XHII]], %[[IDX3]]]
1525   // CHECK: %[[HILO:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YHII]], %[[XLOI]], %[[IDX3]]]
1526   // CHECK: %[[HIHI:.+]] = tensor.extract %arg0[%[[IDX0]], %[[YHII]], %[[XHII]], %[[IDX3]]]
1527
1528   // CHECK: %[[XLOLO:.+]] = arith.extsi %[[LOLO]]
1529   // CHECK: %[[XLOHI:.+]] = arith.extsi %[[LOHI]]
1530   // CHECK: %[[XHILO:.+]] = arith.extsi %[[HILO]]
1531   // CHECK: %[[XHIHI:.+]] = arith.extsi %[[HIHI]]
1532
1533   // Compute the bilinear interpolation.
1534
1535   // CHECK: %[[SCALE:.+]] = arith.constant 256
1536   // CHECK: %[[NDX:.+]] = arith.subi %[[SCALE]], %[[DX]]
1537   // CHECK: %[[WLOLO:.+]] = arith.muli %[[XLOLO]], %[[NDX]]
1538   // CHECK: %[[WLOHI:.+]] = arith.muli %[[XLOHI]], %[[DX]]
1539   // CHECK: %[[LO:.+]] = arith.addi %[[WLOLO]], %[[WLOHI]]
1540   // CHECK: %[[WHILO:.+]] = arith.muli %[[XHILO]], %[[NDX]]
1541   // CHECK: %[[WHIHI:.+]] = arith.muli %[[XHIHI]], %[[DX]]
1542   // CHECK: %[[HI:.+]] = arith.addi %[[WHILO]], %[[WHIHI]]
1543   // CHECK: %[[NDY:.+]] = arith.subi %[[SCALE]], %[[DY]]
1544   // CHECK: %[[WLO:.+]] = arith.muli %[[LO]], %[[NDY]]
1545   // CHECK: %[[WHI:.+]] = arith.muli %[[HI]], %[[DY]]
1546   // CHECK: %[[RESULT:.+]] = arith.addi %[[WLO]], %[[WHI]]
1547   // CHECK: linalg.yield %[[RESULT]]
1548   %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [128, 128], offset = [1, 2], stride_fp = [0. : f32, 0. : f32], offset_fp = [0. : f32, 0. : f32], shift = 8 : i32, mode = "BILINEAR" } : (tensor<1x2x2x1xi8>)  -> (tensor<1x4x4x1xi32>)
1549   return
1550 }