mlir/test/IR/core-ops.mlir

   1 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
   2 // Verify the printed output can be parsed.
   3 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
   4 // Verify the generic form can be parsed.
   5 // RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
   6
   7 // CHECK: #map0 = affine_map<(d0) -> (d0 + 1)>
   8
   9 // CHECK: #map1 = affine_map<()[s0] -> (s0 + 1)>
  10
  11 // CHECK-DAG: #[[$BASE_MAP0:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>
  12 // CHECK-DAG: #[[$BASE_MAP3:map[0-9]+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
  13
  14 // CHECK-LABEL: func @func_with_ops
  15 // CHECK-SAME: %[[ARG:.*]]: f32
  16 func @func_with_ops(f32) {
  17 ^bb0(%a : f32):
  18   // CHECK: %[[T:.*]] = "getTensor"() : () -> tensor<4x4x?xf32>
  19   %t = "getTensor"() : () -> tensor<4x4x?xf32>
  20
  21   // CHECK: %[[C2:.*]] = arith.constant 2 : index
  22   // CHECK-NEXT: %{{.*}} = tensor.dim %[[T]], %[[C2]] : tensor<4x4x?xf32>
  23   %c2 = arith.constant 2 : index
  24   %t2 = "tensor.dim"(%t, %c2) : (tensor<4x4x?xf32>, index) -> index
  25
  26   // CHECK: %{{.*}} = arith.addf %[[ARG]], %[[ARG]] : f32
  27   %x = "arith.addf"(%a, %a) : (f32,f32) -> (f32)
  28
  29   // CHECK: return
  30   return
  31 }
  32
  33 // CHECK-LABEL: func @standard_instrs(%arg0: tensor<4x4x?xf32>, %arg1: f32, %arg2: i32, %arg3: index, %arg4: i64, %arg5: f16) {
  34 func @standard_instrs(tensor<4x4x?xf32>, f32, i32, index, i64, f16) {
  35 ^bb42(%t: tensor<4x4x?xf32>, %f: f32, %i: i32, %idx : index, %j: i64, %half: f16):
  36   // CHECK: %[[C2:.*]] = arith.constant 2 : index
  37   // CHECK: %[[A2:.*]] = tensor.dim %arg0, %[[C2]] : tensor<4x4x?xf32>
  38   %c2 = arith.constant 2 : index
  39   %a2 = tensor.dim %t, %c2 : tensor<4x4x?xf32>
  40
  41   // CHECK: %f = constant @func_with_ops : (f32) -> ()
  42   %10 = constant @func_with_ops : (f32) -> ()
  43
  44   // CHECK: %f_0 = constant @affine_apply : () -> ()
  45   %11 = constant @affine_apply : () -> ()
  46
  47   // CHECK: %[[I2:.*]] = arith.addi
  48   %i2 = arith.addi %i, %i: i32
  49   // CHECK: %[[I3:.*]] = arith.addi
  50   %i3 = arith.addi %i2, %i : i32
  51   // CHECK: %[[I4:.*]] = arith.addi
  52   %i4 = arith.addi %i2, %i3 : i32
  53   // CHECK: %[[F3:.*]] = arith.addf
  54   %f3 = arith.addf %f, %f : f32
  55   // CHECK: %[[F4:.*]] = arith.addf
  56   %f4 = arith.addf %f, %f3 : f32
  57
  58   %true = arith.constant true
  59   %tci32 = arith.constant dense<0> : tensor<42xi32>
  60   %vci32 = arith.constant dense<0> : vector<42xi32>
  61   %tci1 = arith.constant dense<1> : tensor<42xi1>
  62   %vci1 = arith.constant dense<1> : vector<42xi1>
  63
  64   // CHECK: %{{.*}} = select %{{.*}}, %arg3, %arg3 : index
  65   %21 = select %true, %idx, %idx : index
  66
  67   // CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : tensor<42xi1>, tensor<42xi32>
  68   %22 = select %tci1, %tci32, %tci32 : tensor<42 x i1>, tensor<42 x i32>
  69
  70   // CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : vector<42xi1>, vector<42xi32>
  71   %23 = select %vci1, %vci32, %vci32 : vector<42 x i1>, vector<42 x i32>
  72
  73   // CHECK: %{{.*}} = select %{{.*}}, %arg3, %arg3 : index
  74   %24 = "std.select"(%true, %idx, %idx) : (i1, index, index) -> index
  75
  76   // CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : tensor<42xi32>
  77   %25 = std.select %true, %tci32, %tci32 : tensor<42 x i32>
  78
  79   %64 = arith.constant dense<0.> : vector<4 x f32>
  80   %tcf32 = arith.constant dense<0.> : tensor<42 x f32>
  81   %vcf32 = arith.constant dense<0.> : vector<4 x f32>
  82
  83   // CHECK: %{{.*}} = arith.cmpf ogt, %{{.*}}, %{{.*}} : f32
  84   %65 = arith.cmpf ogt, %f3, %f4 : f32
  85
  86   // Predicate 0 means ordered equality comparison.
  87   // CHECK: %{{.*}} = arith.cmpf oeq, %{{.*}}, %{{.*}} : f32
  88   %66 = "arith.cmpf"(%f3, %f4) {predicate = 1} : (f32, f32) -> i1
  89
  90   // CHECK: %{{.*}} = arith.cmpf olt, %{{.*}}, %{{.*}}: vector<4xf32>
  91   %67 = arith.cmpf olt, %vcf32, %vcf32 : vector<4 x f32>
  92
  93   // CHECK: %{{.*}} = arith.cmpf oeq, %{{.*}}, %{{.*}}: vector<4xf32>
  94   %68 = "arith.cmpf"(%vcf32, %vcf32) {predicate = 1} : (vector<4 x f32>, vector<4 x f32>) -> vector<4 x i1>
  95
  96   // CHECK: %{{.*}} = arith.cmpf oeq, %{{.*}}, %{{.*}}: tensor<42xf32>
  97   %69 = arith.cmpf oeq, %tcf32, %tcf32 : tensor<42 x f32>
  98
  99   // CHECK: %{{.*}} = arith.cmpf oeq, %{{.*}}, %{{.*}}: vector<4xf32>
 100   %70 = arith.cmpf oeq, %vcf32, %vcf32 : vector<4 x f32>
 101
 102   // CHECK: = constant unit
 103   %73 = constant unit
 104
 105   // CHECK: arith.constant true
 106   %74 = arith.constant true
 107
 108   // CHECK: arith.constant false
 109   %75 = arith.constant false
 110
 111   // CHECK: %{{.*}} = math.abs %arg1 : f32
 112   %100 = "math.abs"(%f) : (f32) -> f32
 113
 114   // CHECK: %{{.*}} = math.abs %arg1 : f32
 115   %101 = math.abs %f : f32
 116
 117   // CHECK: %{{.*}} = math.abs %{{.*}}: vector<4xf32>
 118   %102 = math.abs %vcf32 : vector<4xf32>
 119
 120   // CHECK: %{{.*}} = math.abs %arg0 : tensor<4x4x?xf32>
 121   %103 = math.abs %t : tensor<4x4x?xf32>
 122
 123   // CHECK: %{{.*}} = math.ceil %arg1 : f32
 124   %104 = "math.ceil"(%f) : (f32) -> f32
 125
 126   // CHECK: %{{.*}} = math.ceil %arg1 : f32
 127   %105 = math.ceil %f : f32
 128
 129   // CHECK: %{{.*}} = math.ceil %{{.*}}: vector<4xf32>
 130   %106 = math.ceil %vcf32 : vector<4xf32>
 131
 132   // CHECK: %{{.*}} = math.ceil %arg0 : tensor<4x4x?xf32>
 133   %107 = math.ceil %t : tensor<4x4x?xf32>
 134
 135   // CHECK: %{{.*}} = math.copysign %arg1, %arg1 : f32
 136   %116 = "math.copysign"(%f, %f) : (f32, f32) -> f32
 137
 138   // CHECK: %{{.*}} = math.copysign %arg1, %arg1 : f32
 139   %117 = math.copysign %f, %f : f32
 140
 141   // CHECK: %{{.*}} = math.copysign %{{.*}}, %{{.*}}: vector<4xf32>
 142   %118 = math.copysign %vcf32, %vcf32 : vector<4xf32>
 143
 144   // CHECK: %{{.*}} = math.copysign %arg0, %arg0 : tensor<4x4x?xf32>
 145   %119 = math.copysign %t, %t : tensor<4x4x?xf32>
 146
 147   // CHECK: %{{.*}} = math.rsqrt %arg1 : f32
 148   %145 = math.rsqrt %f : f32
 149
 150   // CHECK: math.floor %arg1 : f32
 151   %163 = "math.floor"(%f) : (f32) -> f32
 152
 153   // CHECK: %{{.*}} = math.floor %arg1 : f32
 154   %164 = math.floor %f : f32
 155
 156   // CHECK: %{{.*}} = math.floor %{{.*}}: vector<4xf32>
 157   %165 = math.floor %vcf32 : vector<4xf32>
 158
 159   // CHECK: %{{.*}} = math.floor %arg0 : tensor<4x4x?xf32>
 160   %166 = math.floor %t : tensor<4x4x?xf32>
 161
 162   return
 163 }
 164
 165 // CHECK-LABEL: func @affine_apply() {
 166 func @affine_apply() {
 167   %i = "arith.constant"() {value = 0: index} : () -> index
 168   %j = "arith.constant"() {value = 1: index} : () -> index
 169
 170   // CHECK: affine.apply #map0(%c0)
 171   %a = "affine.apply" (%i) { map = affine_map<(d0) -> (d0 + 1)> } :
 172     (index) -> (index)
 173
 174   // CHECK: affine.apply #map1()[%c0]
 175   %b = affine.apply affine_map<()[x] -> (x+1)>()[%i]
 176
 177   return
 178 }
 179
 180 // CHECK-LABEL: func @load_store_prefetch
 181 func @load_store_prefetch(memref<4x4xi32>, index) {
 182 ^bb0(%0: memref<4x4xi32>, %1: index):
 183   // CHECK: %0 = memref.load %arg0[%arg1, %arg1] : memref<4x4xi32>
 184   %2 = "memref.load"(%0, %1, %1) : (memref<4x4xi32>, index, index)->i32
 185
 186   // CHECK: %{{.*}} = memref.load %arg0[%arg1, %arg1] : memref<4x4xi32>
 187   %3 = memref.load %0[%1, %1] : memref<4x4xi32>
 188
 189   // CHECK: memref.prefetch %arg0[%arg1, %arg1], write, locality<1>, data : memref<4x4xi32>
 190   memref.prefetch %0[%1, %1], write, locality<1>, data : memref<4x4xi32>
 191
 192   // CHECK: memref.prefetch %arg0[%arg1, %arg1], read, locality<3>, instr : memref<4x4xi32>
 193   memref.prefetch %0[%1, %1], read, locality<3>, instr : memref<4x4xi32>
 194
 195   return
 196 }
 197
 198 // Test with zero-dimensional operands using no index in load/store.
 199 // CHECK-LABEL: func @zero_dim_no_idx
 200 func @zero_dim_no_idx(%arg0 : memref<i32>, %arg1 : memref<i32>, %arg2 : memref<i32>) {
 201   %0 = memref.load %arg0[] : memref<i32>
 202   memref.store %0, %arg1[] : memref<i32>
 203   return
 204   // CHECK: %0 = memref.load %{{.*}}[] : memref<i32>
 205   // CHECK: memref.store %{{.*}}, %{{.*}}[] : memref<i32>
 206 }
 207
 208 // CHECK-LABEL: func @return_op(%arg0: i32) -> i32 {
 209 func @return_op(%a : i32) -> i32 {
 210   // CHECK: return %arg0 : i32
 211   "std.return" (%a) : (i32)->()
 212 }
 213
 214 // CHECK-LABEL: func @calls(%arg0: i32) {
 215 func @calls(%arg0: i32) {
 216   // CHECK: %0 = call @return_op(%arg0) : (i32) -> i32
 217   %x = call @return_op(%arg0) : (i32) -> i32
 218   // CHECK: %1 = call @return_op(%0) : (i32) -> i32
 219   %y = call @return_op(%x) : (i32) -> i32
 220   // CHECK: %2 = call @return_op(%0) : (i32) -> i32
 221   %z = "std.call"(%x) {callee = @return_op} : (i32) -> i32
 222
 223   // CHECK: %f = constant @affine_apply : () -> ()
 224   %f = constant @affine_apply : () -> ()
 225
 226   // CHECK: call_indirect %f() : () -> ()
 227   call_indirect %f() : () -> ()
 228
 229   // CHECK: %f_0 = constant @return_op : (i32) -> i32
 230   %f_0 = constant @return_op : (i32) -> i32
 231
 232   // CHECK: %3 = call_indirect %f_0(%arg0) : (i32) -> i32
 233   %2 = call_indirect %f_0(%arg0) : (i32) -> i32
 234
 235   // CHECK: %4 = call_indirect %f_0(%arg0) : (i32) -> i32
 236   %3 = "std.call_indirect"(%f_0, %arg0) : ((i32) -> i32, i32) -> i32
 237
 238   return
 239 }
 240
 241 // CHECK-LABEL: func @memref_cast(%arg0
 242 func @memref_cast(%arg0: memref<4xf32>, %arg1 : memref<?xf32>, %arg2 : memref<64x16x4xf32, offset: 0, strides: [64, 4, 1]>) {
 243   // CHECK: %0 = memref.cast %arg0 : memref<4xf32> to memref<?xf32>
 244   %0 = memref.cast %arg0 : memref<4xf32> to memref<?xf32>
 245
 246   // CHECK: %1 = memref.cast %arg1 : memref<?xf32> to memref<4xf32>
 247   %1 = memref.cast %arg1 : memref<?xf32> to memref<4xf32>
 248
 249   // CHECK: {{%.*}} = memref.cast %arg2 : memref<64x16x4xf32, #[[$BASE_MAP0]]> to memref<64x16x4xf32, #[[$BASE_MAP3]]>
 250   %2 = memref.cast %arg2 : memref<64x16x4xf32, offset: 0, strides: [64, 4, 1]> to memref<64x16x4xf32, offset: ?, strides: [?, ?, ?]>
 251
 252   // CHECK: {{%.*}} = memref.cast {{%.*}} : memref<64x16x4xf32, #[[$BASE_MAP3]]> to memref<64x16x4xf32, #[[$BASE_MAP0]]>
 253   %3 = memref.cast %2 : memref<64x16x4xf32, offset: ?, strides: [?, ?, ?]> to memref<64x16x4xf32, offset: 0, strides: [64, 4, 1]>
 254
 255   // CHECK: memref.cast %{{.*}} : memref<4xf32> to memref<*xf32>
 256   %4 = memref.cast %1 : memref<4xf32> to memref<*xf32>
 257
 258   // CHECK: memref.cast %{{.*}} : memref<*xf32> to memref<4xf32>
 259   %5 = memref.cast %4 : memref<*xf32> to memref<4xf32>
 260   return
 261 }
 262
 263 // Check that unranked memrefs with non-default memory space roundtrip
 264 // properly.
 265 // CHECK-LABEL: @unranked_memref_roundtrip(memref<*xf32, 4>)
 266 func private @unranked_memref_roundtrip(memref<*xf32, 4>)
 267
 268 // CHECK-LABEL: func @memref_view(%arg0
 269 func @memref_view(%arg0 : index, %arg1 : index, %arg2 : index) {
 270   %0 = memref.alloc() : memref<2048xi8>
 271   // Test two dynamic sizes and dynamic offset.
 272   // CHECK: %{{.*}} = memref.view %0[%arg2][%arg0, %arg1] : memref<2048xi8> to memref<?x?xf32>
 273   %1 = memref.view %0[%arg2][%arg0, %arg1] : memref<2048xi8> to memref<?x?xf32>
 274
 275   // Test one dynamic size and dynamic offset.
 276   // CHECK: %{{.*}} = memref.view %0[%arg2][%arg1] : memref<2048xi8> to memref<4x?xf32>
 277   %3 = memref.view %0[%arg2][%arg1] : memref<2048xi8> to memref<4x?xf32>
 278
 279   // Test static sizes and static offset.
 280   // CHECK: %{{.*}} = memref.view %0[{{.*}}][] : memref<2048xi8> to memref<64x4xf32>
 281   %c0 = arith.constant 0: index
 282   %5 = memref.view %0[%c0][] : memref<2048xi8> to memref<64x4xf32>
 283   return
 284 }
 285
 286 // CHECK-LABEL: func @test_dimop
 287 // CHECK-SAME: %[[ARG:.*]]: tensor<4x4x?xf32>
 288 func @test_dimop(%arg0: tensor<4x4x?xf32>) {
 289   // CHECK: %[[C2:.*]] = arith.constant 2 : index
 290   // CHECK: %{{.*}} = tensor.dim %[[ARG]], %[[C2]] : tensor<4x4x?xf32>
 291   %c2 = arith.constant 2 : index
 292   %0 = tensor.dim %arg0, %c2 : tensor<4x4x?xf32>
 293   // use dim as an index to ensure type correctness
 294   %1 = affine.apply affine_map<(d0) -> (d0)>(%0)
 295   return
 296 }
 297
 298 // CHECK-LABEL: func @test_splat_op
 299 // CHECK-SAME: [[S:%arg[0-9]+]]: f32
 300 func @test_splat_op(%s : f32) {
 301   %v = splat %s : vector<8xf32>
 302   // CHECK: splat [[S]] : vector<8xf32>
 303   %t = splat %s : tensor<8xf32>
 304   // CHECK: splat [[S]] : tensor<8xf32>
 305   %u = "std.splat"(%s) : (f32) -> vector<4xf32>
 306   // CHECK: splat [[S]] : vector<4xf32>
 307   return
 308 }
 309
 310 // CHECK-LABEL: func @tensor_load_store
 311 func @tensor_load_store(%0 : memref<4x4xi32>, %1 : tensor<4x4xi32>) {
 312   // CHECK-SAME: (%[[MEMREF:.*]]: memref<4x4xi32>,
 313   // CHECK-SAME:  %[[TENSOR:.*]]: tensor<4x4xi32>)
 314   // CHECK: memref.tensor_store %[[TENSOR]], %[[MEMREF]] : memref<4x4xi32>
 315   memref.tensor_store %1, %0 : memref<4x4xi32>
 316   return
 317 }
 318
 319 // CHECK-LABEL: func @unranked_tensor_load_store
 320 func @unranked_tensor_load_store(%0 : memref<*xi32>, %1 : tensor<*xi32>) {
 321   // CHECK-SAME: (%[[MEMREF:.*]]: memref<*xi32>,
 322   // CHECK-SAME:  %[[TENSOR:.*]]: tensor<*xi32>)
 323   // CHECK: memref.tensor_store %[[TENSOR]], %[[MEMREF]] : memref<*xi32>
 324   memref.tensor_store %1, %0 : memref<*xi32>
 325   return
 326 }
 327
 328 // CHECK-LABEL: func @generic_atomic_rmw
 329 // CHECK-SAME: ([[BUF:%.*]]: memref<1x2xf32>, [[I:%.*]]: index, [[J:%.*]]: index)
 330 func @generic_atomic_rmw(%I: memref<1x2xf32>, %i : index, %j : index) {
 331   %x = generic_atomic_rmw %I[%i, %j] : memref<1x2xf32> {
 332   // CHECK-NEXT: generic_atomic_rmw [[BUF]]{{\[}}[[I]], [[J]]] : memref
 333     ^bb0(%old_value : f32):
 334       %c1 = arith.constant 1.0 : f32
 335       %out = arith.addf %c1, %old_value : f32
 336       atomic_yield %out : f32
 337   // CHECK: index_attr = 8 : index
 338   } { index_attr = 8 : index }
 339   return
 340 }
 341
 342 // CHECK-LABEL: func @assume_alignment
 343 // CHECK-SAME: %[[MEMREF:.*]]: memref<4x4xf16>
 344 func @assume_alignment(%0: memref<4x4xf16>) {
 345   // CHECK: memref.assume_alignment %[[MEMREF]], 16 : memref<4x4xf16>
 346   memref.assume_alignment %0, 16 : memref<4x4xf16>
 347   return
 348 }