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 4gpu.module @test_module { 5 // CHECK-LABEL: func @gpu_index_ops() 6 // CHECK32-LABEL: func @gpu_index_ops() 7 func @gpu_index_ops() 8 -> (index, index, index, index, index, index, 9 index, index, index, index, index, index) { 10 // CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i64 11 12 // CHECK: rocdl.workitem.id.x : i32 13 // CHECK: = llvm.sext %{{.*}} : i32 to i64 14 %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index) 15 // CHECK: rocdl.workitem.id.y : i32 16 // CHECK: = llvm.sext %{{.*}} : i32 to i64 17 %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index) 18 // CHECK: rocdl.workitem.id.z : i32 19 // CHECK: = llvm.sext %{{.*}} : i32 to i64 20 %tIdZ = "gpu.thread_id"() {dimension = "z"} : () -> (index) 21 22 // CHECK: rocdl.workgroup.dim.x : i32 23 // CHECK: = llvm.sext %{{.*}} : i32 to i64 24 %bDimX = "gpu.block_dim"() {dimension = "x"} : () -> (index) 25 // CHECK: rocdl.workgroup.dim.y : i32 26 // CHECK: = llvm.sext %{{.*}} : i32 to i64 27 %bDimY = "gpu.block_dim"() {dimension = "y"} : () -> (index) 28 // CHECK: rocdl.workgroup.dim.z : i32 29 // CHECK: = llvm.sext %{{.*}} : i32 to i64 30 %bDimZ = "gpu.block_dim"() {dimension = "z"} : () -> (index) 31 32 // CHECK: rocdl.workgroup.id.x : i32 33 // CHECK: = llvm.sext %{{.*}} : i32 to i64 34 %bIdX = "gpu.block_id"() {dimension = "x"} : () -> (index) 35 // CHECK: rocdl.workgroup.id.y : i32 36 // CHECK: = llvm.sext %{{.*}} : i32 to i64 37 %bIdY = "gpu.block_id"() {dimension = "y"} : () -> (index) 38 // CHECK: rocdl.workgroup.id.z : i32 39 // CHECK: = llvm.sext %{{.*}} : i32 to i64 40 %bIdZ = "gpu.block_id"() {dimension = "z"} : () -> (index) 41 42 // CHECK: rocdl.grid.dim.x : i32 43 // CHECK: = llvm.sext %{{.*}} : i32 to i64 44 %gDimX = "gpu.grid_dim"() {dimension = "x"} : () -> (index) 45 // CHECK: rocdl.grid.dim.y : i32 46 // CHECK: = llvm.sext %{{.*}} : i32 to i64 47 %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index) 48 // CHECK: rocdl.grid.dim.z : i32 49 // CHECK: = llvm.sext %{{.*}} : i32 to i64 50 %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index) 51 52 std.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ, 53 %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ 54 : index, index, index, index, index, index, 55 index, index, index, index, index, index 56 } 57} 58 59// ----- 60 61gpu.module @test_module { 62 // CHECK-LABEL: func @gpu_index_comp 63 // CHECK32-LABEL: func @gpu_index_comp 64 func @gpu_index_comp(%idx : index) -> index { 65 // CHECK: = llvm.add %{{.*}}, %{{.*}} : i64 66 // CHECK32: = llvm.add %{{.*}}, %{{.*}} : i32 67 %0 = addi %idx, %idx : index 68 // CHECK: llvm.return %{{.*}} : i64 69 // CHECK32: llvm.return %{{.*}} : i32 70 std.return %0 : index 71 } 72} 73 74// ----- 75 76gpu.module @test_module { 77 // CHECK-LABEL: func @gpu_sync() 78 func @gpu_sync() { 79 // CHECK: rocdl.barrier 80 gpu.barrier 81 std.return 82 } 83} 84 85// ----- 86 87gpu.module @test_module { 88 // CHECK: llvm.func @__ocml_fabs_f32(f32) -> f32 89 // CHECK: llvm.func @__ocml_fabs_f64(f64) -> f64 90 // CHECK-LABEL: func @gpu_fabs 91 func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 92 %result32 = std.absf %arg_f32 : f32 93 // CHECK: llvm.call @__ocml_fabs_f32(%{{.*}}) : (f32) -> f32 94 %result64 = std.absf %arg_f64 : f64 95 // CHECK: llvm.call @__ocml_fabs_f64(%{{.*}}) : (f64) -> f64 96 std.return %result32, %result64 : f32, f64 97 } 98} 99 100// ----- 101 102gpu.module @test_module { 103 // CHECK: llvm.func @__ocml_ceil_f32(f32) -> f32 104 // CHECK: llvm.func @__ocml_ceil_f64(f64) -> f64 105 // CHECK-LABEL: func @gpu_ceil 106 func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 107 %result32 = std.ceilf %arg_f32 : f32 108 // CHECK: llvm.call @__ocml_ceil_f32(%{{.*}}) : (f32) -> f32 109 %result64 = std.ceilf %arg_f64 : f64 110 // CHECK: llvm.call @__ocml_ceil_f64(%{{.*}}) : (f64) -> f64 111 std.return %result32, %result64 : f32, f64 112 } 113} 114 115// ----- 116 117gpu.module @test_module { 118 // CHECK: llvm.func @__ocml_floor_f32(f32) -> f32 119 // CHECK: llvm.func @__ocml_floor_f64(f64) -> f64 120 // CHECK-LABEL: func @gpu_floor 121 func @gpu_floor(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 122 %result32 = std.floorf %arg_f32 : f32 123 // CHECK: llvm.call @__ocml_floor_f32(%{{.*}}) : (f32) -> f32 124 %result64 = std.floorf %arg_f64 : f64 125 // CHECK: llvm.call @__ocml_floor_f64(%{{.*}}) : (f64) -> f64 126 std.return %result32, %result64 : f32, f64 127 } 128} 129 130// ----- 131 132gpu.module @test_module { 133 // CHECK: llvm.func @__ocml_cos_f32(f32) -> f32 134 // CHECK: llvm.func @__ocml_cos_f64(f64) -> f64 135 // CHECK-LABEL: func @gpu_cos 136 func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 137 %result32 = math.cos %arg_f32 : f32 138 // CHECK: llvm.call @__ocml_cos_f32(%{{.*}}) : (f32) -> f32 139 %result64 = math.cos %arg_f64 : f64 140 // CHECK: llvm.call @__ocml_cos_f64(%{{.*}}) : (f64) -> f64 141 std.return %result32, %result64 : f32, f64 142 } 143} 144 145// ----- 146 147gpu.module @test_module { 148 // CHECK: llvm.func @__ocml_exp_f32(f32) -> f32 149 // CHECK: llvm.func @__ocml_exp_f64(f64) -> f64 150 // CHECK-LABEL: func @gpu_exp 151 func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 152 %exp_f32 = math.exp %arg_f32 : f32 153 // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32 154 %result32 = math.exp %exp_f32 : f32 155 // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32 156 %result64 = math.exp %arg_f64 : f64 157 // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (f64) -> f64 158 std.return %result32, %result64 : f32, f64 159 } 160} 161 162// ----- 163 164gpu.module @test_module { 165 // CHECK: llvm.func @__ocml_exp2_f32(f32) -> f32 166 // CHECK: llvm.func @__ocml_exp2_f64(f64) -> f64 167 // CHECK-LABEL: func @gpu_exp2 168 func @gpu_exp2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 169 %exp2_f32 = math.exp2 %arg_f32 : f32 170 // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32 171 %result32 = math.exp2 %exp2_f32 : f32 172 // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32 173 %result64 = math.exp2 %arg_f64 : f64 174 // CHECK: llvm.call @__ocml_exp2_f64(%{{.*}}) : (f64) -> f64 175 std.return %result32, %result64 : f32, f64 176 } 177} 178 179// ----- 180 181// Test that we handled properly operation with SymbolTable other than module op 182gpu.module @test_module { 183 "test.symbol_scope"() ({ 184 // CHECK: test.symbol_scope 185 // CHECK: llvm.func @__ocml_exp_f32(f32) -> f32 186 // CHECK: llvm.func @__ocml_exp_f64(f64) -> f64 187 // CHECK-LABEL: func @gpu_exp 188 func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 189 %exp_f32 = math.exp %arg_f32 : f32 190 // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32 191 %result32 = math.exp %exp_f32 : f32 192 // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32 193 %result64 = math.exp %arg_f64 : f64 194 // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (f64) -> f64 195 std.return %result32, %result64 : f32, f64 196 } 197 "test.finish" () : () -> () 198 }) : () -> () 199} 200 201// ----- 202 203gpu.module @test_module { 204 // CHECK: llvm.func @__ocml_expm1_f32(f32) -> f32 205 // CHECK: llvm.func @__ocml_expm1_f64(f64) -> f64 206 // CHECK-LABEL: func @gpu_expm1 207 func @gpu_expm1(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 208 %expm1_f32 = math.expm1 %arg_f32 : f32 209 // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32 210 %result32 = math.expm1 %expm1_f32 : f32 211 // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32 212 %result64 = math.expm1 %arg_f64 : f64 213 // CHECK: llvm.call @__ocml_expm1_f64(%{{.*}}) : (f64) -> f64 214 std.return %result32, %result64 : f32, f64 215 } 216} 217 218// ----- 219 220gpu.module @test_module { 221 // CHECK: llvm.func @__ocml_log_f32(f32) -> f32 222 // CHECK: llvm.func @__ocml_log_f64(f64) -> f64 223 // CHECK-LABEL: func @gpu_log 224 func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 225 %result32 = math.log %arg_f32 : f32 226 // CHECK: llvm.call @__ocml_log_f32(%{{.*}}) : (f32) -> f32 227 %result64 = math.log %arg_f64 : f64 228 // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (f64) -> f64 229 std.return %result32, %result64 : f32, f64 230 } 231} 232 233// ----- 234 235gpu.module @test_module { 236 // CHECK: llvm.func @__ocml_log1p_f32(f32) -> f32 237 // CHECK: llvm.func @__ocml_log1p_f64(f64) -> f64 238 // CHECK-LABEL: func @gpu_log1p 239 func @gpu_log1p(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 240 %result32 = math.log1p %arg_f32 : f32 241 // CHECK: llvm.call @__ocml_log1p_f32(%{{.*}}) : (f32) -> f32 242 %result64 = math.log1p %arg_f64 : f64 243 // CHECK: llvm.call @__ocml_log1p_f64(%{{.*}}) : (f64) -> f64 244 std.return %result32, %result64 : f32, f64 245 } 246} 247 248// ----- 249 250gpu.module @test_module { 251 // CHECK: llvm.func @__ocml_log10_f32(f32) -> f32 252 // CHECK: llvm.func @__ocml_log10_f64(f64) -> f64 253 // CHECK-LABEL: func @gpu_log10 254 func @gpu_log10(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 255 %result32 = math.log10 %arg_f32 : f32 256 // CHECK: llvm.call @__ocml_log10_f32(%{{.*}}) : (f32) -> f32 257 %result64 = math.log10 %arg_f64 : f64 258 // CHECK: llvm.call @__ocml_log10_f64(%{{.*}}) : (f64) -> f64 259 std.return %result32, %result64 : f32, f64 260 } 261} 262 263// ----- 264 265gpu.module @test_module { 266 // CHECK: llvm.func @__ocml_log2_f32(f32) -> f32 267 // CHECK: llvm.func @__ocml_log2_f64(f64) -> f64 268 // CHECK-LABEL: func @gpu_log2 269 func @gpu_log2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 270 %result32 = math.log2 %arg_f32 : f32 271 // CHECK: llvm.call @__ocml_log2_f32(%{{.*}}) : (f32) -> f32 272 %result64 = math.log2 %arg_f64 : f64 273 // CHECK: llvm.call @__ocml_log2_f64(%{{.*}}) : (f64) -> f64 274 std.return %result32, %result64 : f32, f64 275 } 276} 277 278// ----- 279 280gpu.module @test_module { 281 // CHECK: llvm.func @__ocml_rsqrt_f32(f32) -> f32 282 // CHECK: llvm.func @__ocml_rsqrt_f64(f64) -> f64 283 // CHECK-LABEL: func @gpu_rsqrt 284 func @gpu_rsqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) 285 -> (f16, f32, f64) { 286 %result16 = math.rsqrt %arg_f16 : f16 287 // CHECK: llvm.fpext %{{.*}} : f16 to f32 288 // CHECK-NEXT: llvm.call @__ocml_rsqrt_f32(%{{.*}}) : (f32) -> f32 289 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 290 %result32 = math.rsqrt %arg_f32 : f32 291 // CHECK: llvm.call @__ocml_rsqrt_f32(%{{.*}}) : (f32) -> f32 292 %result64 = math.rsqrt %arg_f64 : f64 293 // CHECK: llvm.call @__ocml_rsqrt_f64(%{{.*}}) : (f64) -> f64 294 std.return %result16, %result32, %result64 : f16, f32, f64 295 } 296} 297 298// ----- 299 300gpu.module @test_module { 301 // CHECK: llvm.func @__ocml_sqrt_f32(f32) -> f32 302 // CHECK: llvm.func @__ocml_sqrt_f64(f64) -> f64 303 // CHECK-LABEL: func @gpu_sqrt 304 func @gpu_sqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) 305 -> (f16, f32, f64) { 306 %result16 = math.sqrt %arg_f16 : f16 307 // CHECK: llvm.fpext %{{.*}} : f16 to f32 308 // CHECK-NEXT: llvm.call @__ocml_sqrt_f32(%{{.*}}) : (f32) -> f32 309 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 310 %result32 = math.sqrt %arg_f32 : f32 311 // CHECK: llvm.call @__ocml_sqrt_f32(%{{.*}}) : (f32) -> f32 312 %result64 = math.sqrt %arg_f64 : f64 313 // CHECK: llvm.call @__ocml_sqrt_f64(%{{.*}}) : (f64) -> f64 314 std.return %result16, %result32, %result64 : f16, f32, f64 315 } 316} 317 318// ----- 319 320gpu.module @test_module { 321 // CHECK: llvm.func @__ocml_tanh_f32(f32) -> f32 322 // CHECK: llvm.func @__ocml_tanh_f64(f64) -> f64 323 // CHECK-LABEL: func @gpu_tanh 324 func @gpu_tanh(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 325 %result32 = math.tanh %arg_f32 : f32 326 // CHECK: llvm.call @__ocml_tanh_f32(%{{.*}}) : (f32) -> f32 327 %result64 = math.tanh %arg_f64 : f64 328 // CHECK: llvm.call @__ocml_tanh_f64(%{{.*}}) : (f64) -> f64 329 std.return %result32, %result64 : f32, f64 330 } 331} 332 333// ----- 334 335gpu.module @test_module { 336 // CHECK: llvm.func @__ocml_atan_f32(f32) -> f32 337 // CHECK: llvm.func @__ocml_atan_f64(f64) -> f64 338 // CHECK-LABEL: func @gpu_atan 339 func @gpu_atan(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 340 %result32 = math.atan %arg_f32 : f32 341 // CHECK: llvm.call @__ocml_atan_f32(%{{.*}}) : (f32) -> f32 342 %result64 = math.atan %arg_f64 : f64 343 // CHECK: llvm.call @__ocml_atan_f64(%{{.*}}) : (f64) -> f64 344 std.return %result32, %result64 : f32, f64 345 } 346} 347 348// ----- 349 350gpu.module @test_module { 351 // CHECK: llvm.func @__ocml_atan2_f32(f32, f32) -> f32 352 // CHECK: llvm.func @__ocml_atan2_f64(f64, f64) -> f64 353 // CHECK-LABEL: func @gpu_atan2 354 func @gpu_atan2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 355 %result32 = math.atan2 %arg_f32, %arg_f32 : f32 356 // CHECK: llvm.call @__ocml_atan2_f32(%{{.*}}) : (f32, f32) -> f32 357 %result64 = math.atan2 %arg_f64, %arg_f64 : f64 358 // CHECK: llvm.call @__ocml_atan2_f64(%{{.*}}) : (f64, f64) -> f64 359 std.return %result32, %result64 : f32, f64 360 } 361} 362 363// ----- 364 365gpu.module @test_module { 366 // CHECK: llvm.func @__ocml_pow_f32(f32, f32) -> f32 367 // CHECK: llvm.func @__ocml_pow_f64(f64, f64) -> f64 368 // CHECK-LABEL: func @gpu_pow 369 func @gpu_pow(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 370 %result32 = math.powf %arg_f32, %arg_f32 : f32 371 // CHECK: llvm.call @__ocml_pow_f32(%{{.*}}, %{{.*}}) : (f32, f32) -> f32 372 %result64 = math.powf %arg_f64, %arg_f64 : f64 373 // CHECK: llvm.call @__ocml_pow_f64(%{{.*}}, %{{.*}}) : (f64, f64) -> f64 374 std.return %result32, %result64 : f32, f64 375 } 376} 377 378// ----- 379 380gpu.module @test_module { 381 // CHECK-LABEL: @kernel_func 382 // CHECK: attributes 383 // CHECK: gpu.kernel 384 // CHECK: rocdl.kernel 385 gpu.func @kernel_func() kernel { 386 gpu.return 387 } 388} 389