mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

   1 // RUN: mlir-opt %s -convert-gpu-to-rocdl -split-input-file | FileCheck %s
   2 // RUN: mlir-opt %s -convert-gpu-to-rocdl='index-bitwidth=32' -split-input-file | FileCheck --check-prefix=CHECK32 %s
   3
   4 // CHECK-LABEL: @test_module
   5 // CHECK-SAME: llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"
   6
   7 gpu.module @test_module {
   8   // CHECK-LABEL: func @gpu_index_ops()
   9   // CHECK32-LABEL: func @gpu_index_ops()
  10   func.func @gpu_index_ops()
  11       -> (index, index, index, index, index, index,
  12           index, index, index, index, index, index,
  13           index) {
  14     // CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i64
  15
  16     // CHECK: rocdl.workitem.id.x : i32
  17     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  18     %tIdX = gpu.thread_id x
  19     // CHECK: rocdl.workitem.id.y : i32
  20     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  21     %tIdY = gpu.thread_id y
  22     // CHECK: rocdl.workitem.id.z : i32
  23     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  24     %tIdZ = gpu.thread_id z
  25
  26     // CHECK: rocdl.workgroup.dim.x : i32
  27     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  28     %bDimX = gpu.block_dim x
  29     // CHECK: rocdl.workgroup.dim.y : i32
  30     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  31     %bDimY = gpu.block_dim y
  32     // CHECK: rocdl.workgroup.dim.z : i32
  33     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  34     %bDimZ = gpu.block_dim z
  35
  36     // CHECK: rocdl.workgroup.id.x : i32
  37     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  38     %bIdX = gpu.block_id x
  39     // CHECK: rocdl.workgroup.id.y : i32
  40     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  41     %bIdY = gpu.block_id y
  42     // CHECK: rocdl.workgroup.id.z : i32
  43     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  44     %bIdZ = gpu.block_id z
  45
  46     // CHECK: rocdl.grid.dim.x : i32
  47     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  48     %gDimX = gpu.grid_dim x
  49     // CHECK: rocdl.grid.dim.y : i32
  50     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  51     %gDimY = gpu.grid_dim y
  52     // CHECK: rocdl.grid.dim.z : i32
  53     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  54     %gDimZ = gpu.grid_dim z
  55
  56     // CHECK: = rocdl.mbcnt.lo %{{.*}}, %{{.*}} : (i32, i32) -> i32
  57     // CHECK: = rocdl.mbcnt.hi %{{.*}}, %{{.*}} : (i32, i32) -> i32
  58     // CHECK: = llvm.sext %{{.*}} : i32 to i64
  59     %laneId = gpu.lane_id
  60
  61     func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
  62                %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
  63                %laneId
  64         : index, index, index, index, index, index,
  65           index, index, index, index, index, index,
  66           index
  67   }
  68 }
  69
  70 // -----
  71
  72 gpu.module @test_module {
  73   // CHECK-LABEL: func @gpu_index_ops_range
  74   // CHECK-SAME: rocdl.flat_work_group_size = "1536,1536"
  75   // CHECK-SAME: rocdl.reqd_work_group_size = array<i32: 8, 12, 16>
  76   gpu.func @gpu_index_ops_range(%place: memref<i32>)  kernel attributes
  77       {known_block_size = array<i32: 8, 12, 16>,
  78        known_grid_size = array<i32: 20, 24, 28>} {
  79
  80     // CHECK: rocdl.workitem.id.x range <i32, 0, 8> : i32
  81     %tIdX = gpu.thread_id x
  82     // CHECK: rocdl.workitem.id.y range <i32, 0, 12> : i32
  83     %tIdY = gpu.thread_id y
  84     // CHECK: rocdl.workitem.id.z range <i32, 0, 16> : i32
  85     %tIdZ = gpu.thread_id z
  86
  87     // CHECK: rocdl.workgroup.id.x range <i32, 0, 20> : i32
  88     %bIdX = gpu.block_id x
  89     // CHECK: rocdl.workgroup.id.y range <i32, 0, 24> : i32
  90     %bIdY = gpu.block_id y
  91     // CHECK: rocdl.workgroup.id.z range <i32, 0, 28> : i32
  92     %bIdZ = gpu.block_id z
  93
  94     // "Usage" to make the ID calls not die
  95     %0 = arith.addi %tIdX, %tIdY : index
  96     %1 = arith.addi %0, %tIdZ : index
  97     %2 = arith.addi %1, %bIdX : index
  98     %3 = arith.addi %2, %bIdY : index
  99     %4 = arith.addi %3, %bIdZ : index
 100     %5 = arith.index_cast %4 : index to i32
 101     memref.store %5, %place[] : memref<i32>
 102     gpu.return
 103   }
 104 }
 105
 106 // -----
 107
 108 gpu.module @test_module {
 109   // CHECK-LABEL: func @gpu_index_comp
 110   // CHECK32-LABEL: func @gpu_index_comp
 111   func.func @gpu_index_comp(%idx : index) -> index {
 112     // CHECK: = llvm.add %{{.*}}, %{{.*}} : i64
 113     // CHECK32: = llvm.add %{{.*}}, %{{.*}} : i32
 114     %0 = arith.addi %idx, %idx : index
 115     // CHECK: llvm.return %{{.*}} : i64
 116     // CHECK32: llvm.return %{{.*}} : i32
 117     func.return %0 : index
 118   }
 119 }
 120
 121 // -----
 122
 123 gpu.module @test_module {
 124   // CHECK-LABEL: func @gpu_sync()
 125   func.func @gpu_sync() {
 126     // CHECK: rocdl.barrier
 127     gpu.barrier
 128     func.return
 129   }
 130 }
 131
 132 // -----
 133
 134 gpu.module @test_module {
 135   // CHECK-LABEL: func @gpu_sqrt
 136   func.func @gpu_sqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 137     %result16 = math.sqrt %arg_f16 : f16
 138     // CHECK: llvm.intr.sqrt(%{{.*}})  : (f16) -> f16
 139     %result32 = math.sqrt %arg_f32 : f32
 140     // CHECK: llvm.intr.sqrt(%{{.*}})  : (f32) -> f32
 141     %result64 = math.sqrt %arg_f64 : f64
 142     // CHECK: llvm.intr.sqrt(%{{.*}})  : (f64) -> f64
 143     func.return  %result16, %result32, %result64 : f16, f32, f64
 144   }
 145 }
 146
 147 // -----
 148
 149 gpu.module @test_module {
 150   // CHECK-LABEL: func @gpu_fabs
 151   func.func @gpu_fabs(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 152     %result16 = math.absf %arg_f16 : f16
 153     // CHECK: llvm.intr.fabs(%{{.*}})  : (f16) -> f16
 154     %result32 = math.absf %arg_f32 : f32
 155     // CHECK: llvm.intr.fabs(%{{.*}})  : (f32) -> f32
 156     %result64 = math.absf %arg_f64 : f64
 157     // CHECK: llvm.intr.fabs(%{{.*}})  : (f64) -> f64
 158     func.return  %result16, %result32, %result64 : f16, f32, f64
 159   }
 160 }
 161
 162 // -----
 163
 164 gpu.module @test_module {
 165   // CHECK: llvm.func @__ocml_exp_f16(f16) -> f16
 166   // CHECK: llvm.func @__ocml_exp_f64(f64) -> f64
 167   // CHECK-LABEL: func @gpu_exp
 168   func.func @gpu_exp(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 169     %result16 = math.exp %arg_f16 : f16
 170     // CHECK: llvm.call @__ocml_exp_f16(%{{.*}}) : (f16) -> f16
 171     %result32 = math.exp %arg_f32 : f32
 172     // CHECK: llvm.intr.exp(%{{.*}})  : (f32) -> f32
 173     %result64 = math.exp %arg_f64 : f64
 174     // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (f64) -> f64
 175     func.return  %result16, %result32, %result64 : f16, f32, f64
 176   }
 177 }
 178
 179 // -----
 180
 181 gpu.module @test_module {
 182   // CHECK: llvm.func @__ocml_log_f16(f16) -> f16
 183   // CHECK: llvm.func @__ocml_log_f64(f64) -> f64
 184   // CHECK-LABEL: func @gpu_log
 185   func.func @gpu_log(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 186     %result16 = math.log %arg_f16 : f16
 187     // CHECK: llvm.call @__ocml_log_f16(%{{.*}}) : (f16) -> f16
 188     %result32 = math.log %arg_f32 : f32
 189     // CHECK: llvm.intr.log(%{{.*}})  : (f32) -> f32
 190     %result64 = math.log %arg_f64 : f64
 191     // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (f64) -> f64
 192     func.return  %result16, %result32, %result64 : f16, f32, f64
 193   }
 194 }
 195
 196 // -----
 197
 198 gpu.module @test_module {
 199   // CHECK: llvm.func @__ocml_cbrt_f16(f16) -> f16
 200   // CHECK: llvm.func @__ocml_cbrt_f32(f32) -> f32
 201   // CHECK: llvm.func @__ocml_cbrt_f64(f64) -> f64
 202   // CHECK-LABEL: func @gpu_cbrt
 203   func.func @gpu_cbrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 204     %result16 = math.cbrt %arg_f16 : f16
 205     // CHECK: llvm.call @__ocml_cbrt_f16(%{{.*}}) : (f16) -> f16
 206     %result32 = math.cbrt %arg_f32 : f32
 207     // CHECK: llvm.call @__ocml_cbrt_f32(%{{.*}}) : (f32) -> f32
 208     %result64 = math.cbrt %arg_f64 : f64
 209     // CHECK: llvm.call @__ocml_cbrt_f64(%{{.*}}) : (f64) -> f64
 210     func.return %result16, %result32, %result64 : f16, f32, f64
 211   }
 212 }
 213
 214 // -----
 215
 216 gpu.module @test_module {
 217   // CHECK: llvm.func @__ocml_ceil_f16(f16) -> f16
 218   // CHECK: llvm.func @__ocml_ceil_f32(f32) -> f32
 219   // CHECK: llvm.func @__ocml_ceil_f64(f64) -> f64
 220   // CHECK-LABEL: func @gpu_ceil
 221   func.func @gpu_ceil(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 222     %result16 = math.ceil %arg_f16 : f16
 223     // CHECK: llvm.call @__ocml_ceil_f16(%{{.*}}) : (f16) -> f16
 224     %result32 = math.ceil %arg_f32 : f32
 225     // CHECK: llvm.call @__ocml_ceil_f32(%{{.*}}) : (f32) -> f32
 226     %result64 = math.ceil %arg_f64 : f64
 227     // CHECK: llvm.call @__ocml_ceil_f64(%{{.*}}) : (f64) -> f64
 228     func.return %result16, %result32, %result64 : f16, f32, f64
 229   }
 230 }
 231
 232 // -----
 233
 234 gpu.module @test_module {
 235   // CHECK: llvm.func @__ocml_floor_f16(f16) -> f16
 236   // CHECK: llvm.func @__ocml_floor_f32(f32) -> f32
 237   // CHECK: llvm.func @__ocml_floor_f64(f64) -> f64
 238   // CHECK-LABEL: func @gpu_floor
 239   func.func @gpu_floor(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 240     %result16 = math.floor %arg_f16 : f16
 241     // CHECK: llvm.call @__ocml_floor_f16(%{{.*}}) : (f16) -> f16
 242     %result32 = math.floor %arg_f32 : f32
 243     // CHECK: llvm.call @__ocml_floor_f32(%{{.*}}) : (f32) -> f32
 244     %result64 = math.floor %arg_f64 : f64
 245     // CHECK: llvm.call @__ocml_floor_f64(%{{.*}}) : (f64) -> f64
 246     func.return %result16, %result32, %result64 : f16, f32, f64
 247   }
 248 }
 249
 250 // -----
 251
 252 gpu.module @test_module {
 253   // CHECK: llvm.func @__ocml_cos_f16(f16) -> f16
 254   // CHECK: llvm.func @__ocml_cos_f32(f32) -> f32
 255   // CHECK: llvm.func @__ocml_cos_f64(f64) -> f64
 256   // CHECK-LABEL: func @gpu_cos
 257   func.func @gpu_cos(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 258     %result16 = math.cos %arg_f16 : f16
 259     // CHECK: llvm.call @__ocml_cos_f16(%{{.*}}) : (f16) -> f16
 260     %result32 = math.cos %arg_f32 : f32
 261     // CHECK: llvm.call @__ocml_cos_f32(%{{.*}}) : (f32) -> f32
 262     %result64 = math.cos %arg_f64 : f64
 263     // CHECK: llvm.call @__ocml_cos_f64(%{{.*}}) : (f64) -> f64
 264     func.return %result16, %result32, %result64 : f16, f32, f64
 265   }
 266 }
 267
 268 // -----
 269
 270 gpu.module @test_module {
 271   // CHECK: llvm.func @__ocml_exp2_f16(f16) -> f16
 272   // CHECK: llvm.func @__ocml_exp2_f32(f32) -> f32
 273   // CHECK: llvm.func @__ocml_exp2_f64(f64) -> f64
 274   // CHECK-LABEL: func @gpu_exp2
 275   func.func @gpu_exp2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 276     %result16 = math.exp2 %arg_f16 : f16
 277     // CHECK: llvm.call @__ocml_exp2_f16(%{{.*}}) : (f16) -> f16
 278     %exp2_f32 = math.exp2 %arg_f32 : f32
 279     // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32
 280     %result32 = math.exp2 %exp2_f32 : f32
 281     // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32
 282     %result64 = math.exp2 %arg_f64 : f64
 283     // CHECK: llvm.call @__ocml_exp2_f64(%{{.*}}) : (f64) -> f64
 284     func.return %result16, %result32, %result64 : f16, f32, f64
 285   }
 286 }
 287
 288 // -----
 289
 290 // Test that we handled properly operation with SymbolTable other than module op
 291 gpu.module @test_module {
 292   "test.symbol_scope"() ({
 293     // CHECK: test.symbol_scope
 294     // CHECK: llvm.func @__ocml_sin_f16(f16) -> f16
 295     // CHECK: llvm.func @__ocml_sin_f32(f32) -> f32
 296     // CHECK: llvm.func @__ocml_sin_f64(f64) -> f64
 297     // CHECK-LABEL: func @gpu_sin
 298     func.func @gpu_sin(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 299       // CHECK: llvm.call @__ocml_sin_f16(%{{.*}}) : (f16) -> f16
 300       %result16 = math.sin %arg_f16 : f16
 301       // CHECK: llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32
 302       %result32 = math.sin %arg_f32 : f32
 303       // CHECK: llvm.call @__ocml_sin_f64(%{{.*}}) : (f64) -> f64
 304       %result64 = math.sin %arg_f64 : f64
 305       func.return %result16, %result32, %result64 : f16, f32, f64
 306     }
 307     "test.finish" () : () -> ()
 308   }) : () -> ()
 309 }
 310
 311 // -----
 312
 313 gpu.module @test_module {
 314   // CHECK: llvm.func @__ocml_expm1_f16(f16) -> f16
 315   // CHECK: llvm.func @__ocml_expm1_f32(f32) -> f32
 316   // CHECK: llvm.func @__ocml_expm1_f64(f64) -> f64
 317   // CHECK-LABEL: func @gpu_expm1
 318   func.func @gpu_expm1(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 319     %result16 = math.expm1 %arg_f16 : f16
 320     // CHECK: llvm.call @__ocml_expm1_f16(%{{.*}}) : (f16) -> f16
 321     %expm1_f32 = math.expm1 %arg_f32 : f32
 322     // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32
 323     %result32 = math.expm1 %expm1_f32 : f32
 324     // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32
 325     %result64 = math.expm1 %arg_f64 : f64
 326     // CHECK: llvm.call @__ocml_expm1_f64(%{{.*}}) : (f64) -> f64
 327     func.return %result16, %result32, %result64 : f16, f32, f64
 328   }
 329 }
 330
 331 // -----
 332
 333 gpu.module @test_module {
 334   // CHECK: llvm.func @__ocml_log_f16(f16) -> f16
 335   // CHECK: llvm.func @__ocml_log_f64(f64) -> f64
 336   // CHECK-LABEL: func @gpu_log
 337   func.func @gpu_log(%arg_f16 : f16, %arg_f64 : f64) -> (f16, f64) {
 338     %result16 = math.log %arg_f16 : f16
 339     // CHECK: llvm.call @__ocml_log_f16(%{{.*}}) : (f16) -> f16
 340     %result64 = math.log %arg_f64 : f64
 341     // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (f64) -> f64
 342     func.return %result16, %result64 : f16, f64
 343   }
 344 }
 345
 346 // -----
 347
 348 gpu.module @test_module {
 349   // CHECK: llvm.func @__ocml_log1p_f16(f16) -> f16
 350   // CHECK: llvm.func @__ocml_log1p_f32(f32) -> f32
 351   // CHECK: llvm.func @__ocml_log1p_f64(f64) -> f64
 352   // CHECK-LABEL: func @gpu_log1p
 353   func.func @gpu_log1p(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 354     %result16 = math.log1p %arg_f16 : f16
 355     // CHECK: llvm.call @__ocml_log1p_f16(%{{.*}}) : (f16) -> f16
 356     %result32 = math.log1p %arg_f32 : f32
 357     // CHECK: llvm.call @__ocml_log1p_f32(%{{.*}}) : (f32) -> f32
 358     %result64 = math.log1p %arg_f64 : f64
 359     // CHECK: llvm.call @__ocml_log1p_f64(%{{.*}}) : (f64) -> f64
 360     func.return %result16, %result32, %result64 : f16, f32, f64
 361   }
 362 }
 363
 364 // -----
 365
 366 gpu.module @test_module {
 367   // CHECK: llvm.func @__ocml_log10_f16(f16) -> f16
 368   // CHECK: llvm.func @__ocml_log10_f32(f32) -> f32
 369   // CHECK: llvm.func @__ocml_log10_f64(f64) -> f64
 370   // CHECK-LABEL: func @gpu_log10
 371   func.func @gpu_log10(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 372     %result16 = math.log10 %arg_f16 : f16
 373     // CHECK: llvm.call @__ocml_log10_f16(%{{.*}}) : (f16) -> f16
 374     %result32 = math.log10 %arg_f32 : f32
 375     // CHECK: llvm.call @__ocml_log10_f32(%{{.*}}) : (f32) -> f32
 376     %result64 = math.log10 %arg_f64 : f64
 377     // CHECK: llvm.call @__ocml_log10_f64(%{{.*}}) : (f64) -> f64
 378     func.return %result16, %result32, %result64 : f16, f32, f64
 379   }
 380 }
 381
 382 // -----
 383
 384 gpu.module @test_module {
 385   // CHECK: llvm.func @__ocml_log2_f16(f16) -> f16
 386   // CHECK: llvm.func @__ocml_log2_f32(f32) -> f32
 387   // CHECK: llvm.func @__ocml_log2_f64(f64) -> f64
 388   // CHECK-LABEL: func @gpu_log2
 389   func.func @gpu_log2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 390     %result16 = math.log2 %arg_f16 : f16
 391     // CHECK: llvm.call @__ocml_log2_f16(%{{.*}}) : (f16) -> f16
 392     %result32 = math.log2 %arg_f32 : f32
 393     // CHECK: llvm.call @__ocml_log2_f32(%{{.*}}) : (f32) -> f32
 394     %result64 = math.log2 %arg_f64 : f64
 395     // CHECK: llvm.call @__ocml_log2_f64(%{{.*}}) : (f64) -> f64
 396     func.return %result16, %result32, %result64 : f16, f32, f64
 397   }
 398 }
 399
 400 // -----
 401
 402 gpu.module @test_module {
 403   // CHECK: llvm.func @__ocml_rsqrt_f16(f16) -> f16
 404   // CHECK: llvm.func @__ocml_rsqrt_f32(f32) -> f32
 405   // CHECK: llvm.func @__ocml_rsqrt_f64(f64) -> f64
 406   // CHECK-LABEL: func @gpu_rsqrt
 407   func.func @gpu_rsqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 408     %result16 = math.rsqrt %arg_f16 : f16
 409     // CHECK: llvm.call @__ocml_rsqrt_f16(%{{.*}}) : (f16) -> f16
 410     %result32 = math.rsqrt %arg_f32 : f32
 411     // CHECK: llvm.call @__ocml_rsqrt_f32(%{{.*}}) : (f32) -> f32
 412     %result64 = math.rsqrt %arg_f64 : f64
 413     // CHECK: llvm.call @__ocml_rsqrt_f64(%{{.*}}) : (f64) -> f64
 414     func.return %result16, %result32, %result64 : f16, f32, f64
 415   }
 416 }
 417
 418 // -----
 419
 420 gpu.module @test_module {
 421   // CHECK: llvm.func @__ocml_tan_f16(f16) -> f16
 422   // CHECK: llvm.func @__ocml_tan_f32(f32) -> f32
 423   // CHECK: llvm.func @__ocml_tan_f64(f64) -> f64
 424   // CHECK-LABEL: func @gpu_tan
 425   func.func @gpu_tan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 426     %result16 = math.tan %arg_f16 : f16
 427     // CHECK: llvm.call @__ocml_tan_f16(%{{.*}}) : (f16) -> f16
 428     %result32 = math.tan %arg_f32 : f32
 429     // CHECK: llvm.call @__ocml_tan_f32(%{{.*}}) : (f32) -> f32
 430     %result64 = math.tan %arg_f64 : f64
 431     // CHECK: llvm.call @__ocml_tan_f64(%{{.*}}) : (f64) -> f64
 432     func.return %result16, %result32, %result64 : f16, f32, f64
 433   }
 434 }
 435
 436 // -----
 437
 438 gpu.module @test_module {
 439   // CHECK: llvm.func @__ocml_tanh_f16(f16) -> f16
 440   // CHECK: llvm.func @__ocml_tanh_f32(f32) -> f32
 441   // CHECK: llvm.func @__ocml_tanh_f64(f64) -> f64
 442   // CHECK-LABEL: func @gpu_tanh
 443   func.func @gpu_tanh(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 444     %result16 = math.tanh %arg_f16 : f16
 445     // CHECK: llvm.call @__ocml_tanh_f16(%{{.*}}) : (f16) -> f16
 446     %result32 = math.tanh %arg_f32 : f32
 447     // CHECK: llvm.call @__ocml_tanh_f32(%{{.*}}) : (f32) -> f32
 448     %result64 = math.tanh %arg_f64 : f64
 449     // CHECK: llvm.call @__ocml_tanh_f64(%{{.*}}) : (f64) -> f64
 450     func.return %result16, %result32, %result64 : f16, f32, f64
 451   }
 452 }
 453
 454 // -----
 455
 456 gpu.module @test_module {
 457   // CHECK: llvm.func @__ocml_atan_f16(f16) -> f16
 458   // CHECK: llvm.func @__ocml_atan_f32(f32) -> f32
 459   // CHECK: llvm.func @__ocml_atan_f64(f64) -> f64
 460   // CHECK-LABEL: func @gpu_atan
 461   func.func @gpu_atan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 462     %result16 = math.atan %arg_f16 : f16
 463     // CHECK: llvm.call @__ocml_atan_f16(%{{.*}}) : (f16) -> f16
 464     %result32 = math.atan %arg_f32 : f32
 465     // CHECK: llvm.call @__ocml_atan_f32(%{{.*}}) : (f32) -> f32
 466     %result64 = math.atan %arg_f64 : f64
 467     // CHECK: llvm.call @__ocml_atan_f64(%{{.*}}) : (f64) -> f64
 468     func.return %result16, %result32, %result64 : f16, f32, f64
 469   }
 470 }
 471
 472 // -----
 473
 474 gpu.module @test_module {
 475   // CHECK: llvm.func @__ocml_atan2_f16(f16, f16) -> f16
 476   // CHECK: llvm.func @__ocml_atan2_f32(f32, f32) -> f32
 477   // CHECK: llvm.func @__ocml_atan2_f64(f64, f64) -> f64
 478   // CHECK-LABEL: func @gpu_atan2
 479   func.func @gpu_atan2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 480     %result16 = math.atan2 %arg_f16, %arg_f16 : f16
 481     // CHECK: llvm.call @__ocml_atan2_f16(%{{.*}}) : (f16, f16) -> f16
 482     %result32 = math.atan2 %arg_f32, %arg_f32 : f32
 483     // CHECK: llvm.call @__ocml_atan2_f32(%{{.*}}) : (f32, f32) -> f32
 484     %result64 = math.atan2 %arg_f64, %arg_f64 : f64
 485     // CHECK: llvm.call @__ocml_atan2_f64(%{{.*}}) : (f64, f64) -> f64
 486     func.return %result16, %result32, %result64 : f16, f32, f64
 487   }
 488 }
 489
 490 // -----
 491
 492 gpu.module @test_module {
 493   // CHECK: llvm.func @__ocml_pow_f16(f16, f16) -> f16
 494   // CHECK: llvm.func @__ocml_pow_f32(f32, f32) -> f32
 495   // CHECK: llvm.func @__ocml_pow_f64(f64, f64) -> f64
 496   // CHECK-LABEL: func @gpu_pow
 497   func.func @gpu_pow(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 498     %result16 = math.powf %arg_f16, %arg_f16 : f16
 499     // CHECK: llvm.call @__ocml_pow_f16(%{{.*}}, %{{.*}}) : (f16, f16) -> f16
 500     %result32 = math.powf %arg_f32, %arg_f32 : f32
 501     // CHECK: llvm.call @__ocml_pow_f32(%{{.*}}, %{{.*}}) : (f32, f32) -> f32
 502     %result64 = math.powf %arg_f64, %arg_f64 : f64
 503     // CHECK: llvm.call @__ocml_pow_f64(%{{.*}}, %{{.*}}) : (f64, f64) -> f64
 504     func.return %result16, %result32, %result64 : f16, f32, f64
 505   }
 506 }
 507
 508 // -----
 509
 510 gpu.module @test_module {
 511   // CHECK: llvm.func @__ocml_erf_f16(f16) -> f16
 512   // CHECK: llvm.func @__ocml_erf_f32(f32) -> f32
 513   // CHECK: llvm.func @__ocml_erf_f64(f64) -> f64
 514   // CHECK-LABEL: func @gpu_erf
 515   func.func @gpu_erf(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 516     %result16 = math.erf %arg_f16 : f16
 517     // CHECK: llvm.call @__ocml_erf_f16(%{{.*}}) : (f16) -> f16
 518     %result32 = math.erf %arg_f32 : f32
 519     // CHECK: llvm.call @__ocml_erf_f32(%{{.*}}) : (f32) -> f32
 520     %result64 = math.erf %arg_f64 : f64
 521     // CHECK: llvm.call @__ocml_erf_f64(%{{.*}}) : (f64) -> f64
 522     func.return %result16, %result32, %result64 : f16, f32, f64
 523   }
 524 }
 525
 526 // -----
 527
 528 gpu.module @test_module {
 529   // CHECK-LABEL: func @gpu_unroll
 530   func.func @gpu_unroll(%arg0 : vector<4xf32>) -> vector<4xf32> {
 531     %result = math.sin %arg0 : vector<4xf32>
 532     // CHECK: %[[V0:.+]] = llvm.mlir.undef : vector<4xf32>
 533     // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32
 534     // CHECK: %[[V1:.+]] = llvm.insertelement %[[CL]], %[[V0]]
 535     // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32
 536     // CHECK: %[[V2:.+]] = llvm.insertelement %[[CL]], %[[V1]]
 537     // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32
 538     // CHECK: %[[V3:.+]] = llvm.insertelement %[[CL]], %[[V2]]
 539     // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32
 540     // CHECK: %[[V4:.+]] = llvm.insertelement %[[CL]], %[[V3]]
 541     // CHECK: return %[[V4]]
 542     func.return %result : vector<4xf32>
 543   }
 544 }
 545
 546 // -----
 547
 548 // Test that the bf16 type is passed through to LLVM.
 549
 550 gpu.module @test_module {
 551   // CHECK-LABEL: func @bf16_id
 552   func.func @bf16_id(%arg0 : bf16) -> bf16 {
 553     // CHECK-SAME: (%[[ARG0:.+]]: bf16)
 554     // CHECK-SAME: -> bf16
 555     // CHECK: return %[[ARG0]] : bf16
 556     func.return %arg0 : bf16
 557   }
 558
 559   // CHECK-LABEL: func @bf16x4_id
 560   func.func @bf16x4_id(%arg0 : vector<4xbf16>) -> vector<4xbf16> {
 561     // CHECK-SAME: (%[[ARG0:.+]]: vector<4xbf16>)
 562     // CHECK-SAME: -> vector<4xbf16>
 563     // CHECK: return %[[ARG0]] : vector<4xbf16>
 564     func.return %arg0 : vector<4xbf16>
 565   }
 566
 567 }
 568
 569 // -----
 570
 571 gpu.module @test_module {
 572   // CHECK-LABEL: @kernel_func
 573   // CHECK: attributes
 574   // CHECK: gpu.kernel
 575   // CHECK: rocdl.kernel
 576   gpu.func @kernel_func() kernel {
 577     gpu.return
 578   }
 579 }
 580
 581 // -----
 582
 583 gpu.module @module {
 584 // CHECK-LABEL: @spirv_sin
 585 // CHECK: llvm.call @__ocml_sin_f32
 586   spirv.func @spirv_sin(%arg0: vector<4xf32>) -> vector<4xf32> "None" {
 587     %0 = math.sin %arg0 : vector<4xf32>
 588     spirv.ReturnValue %0 : vector<4xf32>
 589   }
 590 }
 591
 592 // -----
 593
 594 gpu.module @test_module {
 595   // CHECK-LABEL: func @gpu_all_reduce_op()
 596   gpu.func @gpu_all_reduce_op() {
 597     %arg0 = arith.constant 1.0 : f32
 598     // TODO: Check full IR expansion once lowering has settled.
 599     // CHECK: llvm.add
 600     // CHECK: llvm.and
 601     // CHECK: llvm.xor
 602     // CHECK: llvm.icmp "slt"
 603     // CHECK: llvm.select
 604     // CHECK: llvm.shl
 605     // CHECK: rocdl.ds_bpermute {{.*}}
 606     // CHECK: rocdl.barrier
 607     // CHECK: llvm.bitcast
 608     // CHECK: llvm.fadd
 609     %result = gpu.all_reduce add %arg0 uniform {} : (f32) -> (f32)
 610
 611     gpu.return
 612   }
 613 }
 614
 615
 616 // -----
 617
 618 gpu.module @test_module {
 619   // CHECK-LABEL: func @gpu_all_reduce_region()
 620   gpu.func @gpu_all_reduce_region() {
 621     %arg0 = arith.constant 1 : i32
 622     // TODO: Check full IR expansion once lowering has settled.
 623     // CHECK: llvm.add
 624     // CHECK: llvm.and
 625     // CHECK: llvm.xor
 626     // CHECK: llvm.icmp "slt"
 627     // CHECK: llvm.select
 628     // CHECK: llvm.shl
 629     // CHECK: rocdl.ds_bpermute {{.*}}
 630     // CHECK: rocdl.barrier
 631     %result = gpu.all_reduce %arg0 uniform {
 632     ^bb(%lhs : i32, %rhs : i32):
 633       %xor = arith.xori %lhs, %rhs : i32
 634       "gpu.yield"(%xor) : (i32) -> ()
 635     } : (i32) -> (i32)
 636     gpu.return
 637   }
 638 }
 639
 640 // -----
 641
 642 gpu.module @test_module {
 643   // CHECK: llvm.func @__ocml_fmod_f16(f16, f16) -> f16
 644   // CHECK: llvm.func @__ocml_fmod_f32(f32, f32) -> f32
 645   // CHECK: llvm.func @__ocml_fmod_f64(f64, f64) -> f64
 646   // CHECK-LABEL: func @gpu_fmod
 647   func.func @gpu_fmod(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {
 648     %result16 = arith.remf %arg_f16, %arg_f16 : f16
 649     // CHECK: llvm.call @__ocml_fmod_f16(%{{.*}}, %{{.*}}) : (f16, f16) -> f16
 650     %result32 = arith.remf %arg_f32, %arg_f32 : f32
 651     // CHECK: llvm.call @__ocml_fmod_f32(%{{.*}}, %{{.*}}) : (f32, f32) -> f32
 652     %result64 = arith.remf %arg_f64, %arg_f64 : f64
 653     // CHECK: llvm.call @__ocml_fmod_f64(%{{.*}}, %{{.*}}) : (f64, f64) -> f64
 654     func.return %result16, %result32, %result64 : f16, f32, f64
 655   }
 656 }
 657
 658 // -----
 659
 660 gpu.module @test_module {
 661   // CHECK-LABEL: func @gpu_shuffle()
 662   func.func @gpu_shuffle() -> (f32, f32, f32) {
 663     // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : f32
 664     %arg0 = arith.constant 1.0 : f32
 665     // CHECK: %[[#OFFSET:]] = llvm.mlir.constant(4 : i32) : i32
 666     %arg1 = arith.constant 4 : i32
 667     // CHECK: %[[#WIDTH:]] = llvm.mlir.constant(23 : i32) : i32
 668     %arg2 = arith.constant 23 : i32
 669     // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi
 670     // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32
 671     // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32
 672     // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32
 673     // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32
 674     // CHECK: %[[#XOR:]] = llvm.xor %[[#LANE_ID]], %{{.*}} : i32
 675     // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#XOR]], %[[#WARP_OR_ZERO]] : i32
 676     // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#XOR]], %{{.*}} : i1, i32
 677     // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32
 678     // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32
 679     // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32
 680     // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32
 681     // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32
 682     %shfl, %pred = gpu.shuffle xor %arg0, %arg1, %arg2 : f32
 683     // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi
 684     // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32
 685     // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32
 686     // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32
 687     // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32
 688     // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#OFFSET]], %[[#WARP_OR_ZERO]] : i32
 689     // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#OFFSET]], %{{.*}} : i1, i32
 690     // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32
 691     // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32
 692     // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32
 693     // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32
 694     // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32
 695     %shfli, %predi = gpu.shuffle idx %arg0, %arg1, %arg2 : f32
 696     // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi
 697     // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32
 698     // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32
 699     // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32
 700     // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32
 701     // CHECK: %[[#DOWN:]] = llvm.add %[[#LANE_ID]], %{{.*}} : i32
 702     // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#DOWN]], %[[#WARP_OR_ZERO]] : i32
 703     // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#DOWN]], %{{.*}} : i1, i32
 704     // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32
 705     // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32
 706     // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32
 707     // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32
 708     // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32
 709     %shfld, %predd = gpu.shuffle down %arg0, %arg1, %arg2 : f32
 710     func.return %shfl, %shfli, %shfld : f32, f32, f32
 711   }
 712 }
 713
 714 // -----
 715
 716 // CHECK-LABEL: @test_custom_data_layout
 717 // CHECK-SAME: llvm.data_layout = "e"
 718 gpu.module @test_custom_data_layout attributes {llvm.data_layout = "e"} {
 719
 720 }