1// RUN: mlir-opt %s -convert-gpu-to-nvvm -split-input-file | FileCheck %s 2// RUN: mlir-opt %s -convert-gpu-to-nvvm='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 builtin.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: = nvvm.read.ptx.sreg.tid.x : i32 13 // CHECK: = llvm.sext %{{.*}} : i32 to i64 14 %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index) 15 // CHECK: = nvvm.read.ptx.sreg.tid.y : i32 16 // CHECK: = llvm.sext %{{.*}} : i32 to i64 17 %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index) 18 // CHECK: = nvvm.read.ptx.sreg.tid.z : i32 19 // CHECK: = llvm.sext %{{.*}} : i32 to i64 20 %tIdZ = "gpu.thread_id"() {dimension = "z"} : () -> (index) 21 22 // CHECK: = nvvm.read.ptx.sreg.ntid.x : i32 23 // CHECK: = llvm.sext %{{.*}} : i32 to i64 24 %bDimX = "gpu.block_dim"() {dimension = "x"} : () -> (index) 25 // CHECK: = nvvm.read.ptx.sreg.ntid.y : i32 26 // CHECK: = llvm.sext %{{.*}} : i32 to i64 27 %bDimY = "gpu.block_dim"() {dimension = "y"} : () -> (index) 28 // CHECK: = nvvm.read.ptx.sreg.ntid.z : i32 29 // CHECK: = llvm.sext %{{.*}} : i32 to i64 30 %bDimZ = "gpu.block_dim"() {dimension = "z"} : () -> (index) 31 32 // CHECK: = nvvm.read.ptx.sreg.ctaid.x : i32 33 // CHECK: = llvm.sext %{{.*}} : i32 to i64 34 %bIdX = "gpu.block_id"() {dimension = "x"} : () -> (index) 35 // CHECK: = nvvm.read.ptx.sreg.ctaid.y : i32 36 // CHECK: = llvm.sext %{{.*}} : i32 to i64 37 %bIdY = "gpu.block_id"() {dimension = "y"} : () -> (index) 38 // CHECK: = nvvm.read.ptx.sreg.ctaid.z : i32 39 // CHECK: = llvm.sext %{{.*}} : i32 to i64 40 %bIdZ = "gpu.block_id"() {dimension = "z"} : () -> (index) 41 42 // CHECK: = nvvm.read.ptx.sreg.nctaid.x : i32 43 // CHECK: = llvm.sext %{{.*}} : i32 to i64 44 %gDimX = "gpu.grid_dim"() {dimension = "x"} : () -> (index) 45 // CHECK: = nvvm.read.ptx.sreg.nctaid.y : i32 46 // CHECK: = llvm.sext %{{.*}} : i32 to i64 47 %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index) 48 // CHECK: = nvvm.read.ptx.sreg.nctaid.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 builtin.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_all_reduce_op() 78 gpu.func @gpu_all_reduce_op() { 79 %arg0 = constant 1.0 : f32 80 // TODO: Check full IR expansion once lowering has settled. 81 // CHECK: nvvm.shfl.sync.bfly 82 // CHECK: nvvm.barrier0 83 // CHECK: llvm.fadd 84 %result = "gpu.all_reduce"(%arg0) ({}) {op = "add"} : (f32) -> (f32) 85 86 gpu.return 87 } 88} 89 90// ----- 91 92gpu.module @test_module { 93 // CHECK-LABEL: func @gpu_all_reduce_region() 94 gpu.func @gpu_all_reduce_region() { 95 %arg0 = constant 1 : i32 96 // TODO: Check full IR expansion once lowering has settled. 97 // CHECK: nvvm.shfl.sync.bfly 98 // CHECK: nvvm.barrier0 99 %result = "gpu.all_reduce"(%arg0) ({ 100 ^bb(%lhs : i32, %rhs : i32): 101 %xor = xor %lhs, %rhs : i32 102 "gpu.yield"(%xor) : (i32) -> () 103 }) : (i32) -> (i32) 104 gpu.return 105 } 106} 107 108// ----- 109 110gpu.module @test_module { 111 // CHECK-LABEL: func @gpu_shuffle() 112 builtin.func @gpu_shuffle() -> (f32) { 113 // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : f32 114 %arg0 = constant 1.0 : f32 115 // CHECK: %[[#OFFSET:]] = llvm.mlir.constant(4 : i32) : i32 116 %arg1 = constant 4 : i32 117 // CHECK: %[[#WIDTH:]] = llvm.mlir.constant(23 : i32) : i32 118 %arg2 = constant 23 : i32 119 // CHECK: %[[#ONE:]] = llvm.mlir.constant(1 : i32) : i32 120 // CHECK: %[[#SHL:]] = llvm.shl %[[#ONE]], %[[#WIDTH]] : i32 121 // CHECK: %[[#MASK:]] = llvm.sub %[[#SHL]], %[[#ONE]] : i32 122 // CHECK: %[[#CLAMP:]] = llvm.sub %[[#WIDTH]], %[[#ONE]] : i32 123 // CHECK: %[[#SHFL:]] = nvvm.shfl.sync.bfly %[[#MASK]], %[[#VALUE]], %[[#OFFSET]], %[[#CLAMP]] : !llvm.struct<(f32, i1)> 124 // CHECK: llvm.extractvalue %[[#SHFL]][0 : index] : !llvm.struct<(f32, i1)> 125 // CHECK: llvm.extractvalue %[[#SHFL]][1 : index] : !llvm.struct<(f32, i1)> 126 %shfl, %pred = "gpu.shuffle"(%arg0, %arg1, %arg2) { mode = "xor" } : (f32, i32, i32) -> (f32, i1) 127 128 std.return %shfl : f32 129 } 130} 131 132// ----- 133 134gpu.module @test_module { 135 // CHECK-LABEL: func @gpu_sync() 136 builtin.func @gpu_sync() { 137 // CHECK: nvvm.barrier0 138 gpu.barrier 139 std.return 140 } 141} 142 143// ----- 144 145gpu.module @test_module { 146 // CHECK: llvm.func @__nv_fabsf(f32) -> f32 147 // CHECK: llvm.func @__nv_fabs(f64) -> f64 148 // CHECK-LABEL: func @gpu_fabs 149 builtin.func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 150 %result32 = std.absf %arg_f32 : f32 151 // CHECK: llvm.call @__nv_fabsf(%{{.*}}) : (f32) -> f32 152 %result64 = std.absf %arg_f64 : f64 153 // CHECK: llvm.call @__nv_fabs(%{{.*}}) : (f64) -> f64 154 std.return %result32, %result64 : f32, f64 155 } 156} 157 158// ----- 159 160gpu.module @test_module { 161 // CHECK: llvm.func @__nv_ceilf(f32) -> f32 162 // CHECK: llvm.func @__nv_ceil(f64) -> f64 163 // CHECK-LABEL: func @gpu_ceil 164 builtin.func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 165 %result32 = std.ceilf %arg_f32 : f32 166 // CHECK: llvm.call @__nv_ceilf(%{{.*}}) : (f32) -> f32 167 %result64 = std.ceilf %arg_f64 : f64 168 // CHECK: llvm.call @__nv_ceil(%{{.*}}) : (f64) -> f64 169 std.return %result32, %result64 : f32, f64 170 } 171} 172 173// ----- 174 175gpu.module @test_module { 176 // CHECK: llvm.func @__nv_floorf(f32) -> f32 177 // CHECK: llvm.func @__nv_floor(f64) -> f64 178 // CHECK-LABEL: func @gpu_floor 179 builtin.func @gpu_floor(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 180 %result32 = std.floorf %arg_f32 : f32 181 // CHECK: llvm.call @__nv_floorf(%{{.*}}) : (f32) -> f32 182 %result64 = std.floorf %arg_f64 : f64 183 // CHECK: llvm.call @__nv_floor(%{{.*}}) : (f64) -> f64 184 std.return %result32, %result64 : f32, f64 185 } 186} 187 188// ----- 189 190gpu.module @test_module { 191 // CHECK: llvm.func @__nv_cosf(f32) -> f32 192 // CHECK: llvm.func @__nv_cos(f64) -> f64 193 // CHECK-LABEL: func @gpu_cos 194 builtin.func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 195 %result32 = math.cos %arg_f32 : f32 196 // CHECK: llvm.call @__nv_cosf(%{{.*}}) : (f32) -> f32 197 %result64 = math.cos %arg_f64 : f64 198 // CHECK: llvm.call @__nv_cos(%{{.*}}) : (f64) -> f64 199 std.return %result32, %result64 : f32, f64 200 } 201} 202 203// ----- 204gpu.module @test_module { 205 // CHECK: llvm.func @__nv_expf(f32) -> f32 206 // CHECK: llvm.func @__nv_exp(f64) -> f64 207 // CHECK-LABEL: func @gpu_exp 208 builtin.func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 209 %result32 = math.exp %arg_f32 : f32 210 // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32 211 %result64 = math.exp %arg_f64 : f64 212 // CHECK: llvm.call @__nv_exp(%{{.*}}) : (f64) -> f64 213 std.return %result32, %result64 : f32, f64 214 } 215} 216 217// ----- 218gpu.module @test_module { 219 // CHECK: llvm.func @__nv_exp2f(f32) -> f32 220 // CHECK: llvm.func @__nv_exp2(f64) -> f64 221 // CHECK-LABEL: func @gpu_exp2 222 builtin.func @gpu_exp2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 223 %result32 = math.exp2 %arg_f32 : f32 224 // CHECK: llvm.call @__nv_exp2f(%{{.*}}) : (f32) -> f32 225 %result64 = math.exp2 %arg_f64 : f64 226 // CHECK: llvm.call @__nv_exp2(%{{.*}}) : (f64) -> f64 227 std.return %result32, %result64 : f32, f64 228 } 229} 230 231// ----- 232 233gpu.module @test_module { 234 // CHECK: llvm.func @__nv_logf(f32) -> f32 235 // CHECK: llvm.func @__nv_log(f64) -> f64 236 // CHECK-LABEL: func @gpu_log 237 builtin.func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 238 %result32 = math.log %arg_f32 : f32 239 // CHECK: llvm.call @__nv_logf(%{{.*}}) : (f32) -> f32 240 %result64 = math.log %arg_f64 : f64 241 // CHECK: llvm.call @__nv_log(%{{.*}}) : (f64) -> f64 242 std.return %result32, %result64 : f32, f64 243 } 244} 245 246// ----- 247 248gpu.module @test_module { 249 // CHECK: llvm.func @__nv_log10f(f32) -> f32 250 // CHECK: llvm.func @__nv_log10(f64) -> f64 251 // CHECK-LABEL: func @gpu_log10 252 builtin.func @gpu_log10(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 253 %result32 = math.log10 %arg_f32 : f32 254 // CHECK: llvm.call @__nv_log10f(%{{.*}}) : (f32) -> f32 255 %result64 = math.log10 %arg_f64 : f64 256 // CHECK: llvm.call @__nv_log10(%{{.*}}) : (f64) -> f64 257 std.return %result32, %result64 : f32, f64 258 } 259} 260 261// ----- 262 263gpu.module @test_module { 264 // CHECK: llvm.func @__nv_log1pf(f32) -> f32 265 // CHECK: llvm.func @__nv_log1p(f64) -> f64 266 // CHECK-LABEL: func @gpu_log1p 267 builtin.func @gpu_log1p(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 268 %result32 = math.log1p %arg_f32 : f32 269 // CHECK: llvm.call @__nv_log1pf(%{{.*}}) : (f32) -> f32 270 %result64 = math.log1p %arg_f64 : f64 271 // CHECK: llvm.call @__nv_log1p(%{{.*}}) : (f64) -> f64 272 std.return %result32, %result64 : f32, f64 273 } 274} 275 276// ----- 277 278gpu.module @test_module { 279 // CHECK: llvm.func @__nv_log2f(f32) -> f32 280 // CHECK: llvm.func @__nv_log2(f64) -> f64 281 // CHECK-LABEL: func @gpu_log2 282 builtin.func @gpu_log2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 283 %result32 = math.log2 %arg_f32 : f32 284 // CHECK: llvm.call @__nv_log2f(%{{.*}}) : (f32) -> f32 285 %result64 = math.log2 %arg_f64 : f64 286 // CHECK: llvm.call @__nv_log2(%{{.*}}) : (f64) -> f64 287 std.return %result32, %result64 : f32, f64 288 } 289} 290 291// ----- 292 293gpu.module @test_module { 294 // CHECK: llvm.func @__nv_sinf(f32) -> f32 295 // CHECK: llvm.func @__nv_sin(f64) -> f64 296 // CHECK-LABEL: func @gpu_sin 297 builtin.func @gpu_sin(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 298 %result32 = math.sin %arg_f32 : f32 299 // CHECK: llvm.call @__nv_sinf(%{{.*}}) : (f32) -> f32 300 %result64 = math.sin %arg_f64 : f64 301 // CHECK: llvm.call @__nv_sin(%{{.*}}) : (f64) -> f64 302 std.return %result32, %result64 : f32, f64 303 } 304} 305 306// ----- 307 308gpu.module @test_module { 309 // CHECK: llvm.func @__nv_tanhf(f32) -> f32 310 // CHECK: llvm.func @__nv_tanh(f64) -> f64 311 // CHECK-LABEL: func @gpu_tanh 312 builtin.func @gpu_tanh(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) { 313 %result16 = math.tanh %arg_f16 : f16 314 // CHECK: llvm.fpext %{{.*}} : f16 to f32 315 // CHECK-NEXT: llvm.call @__nv_tanhf(%{{.*}}) : (f32) -> f32 316 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 317 %result32 = math.tanh %arg_f32 : f32 318 // CHECK: llvm.call @__nv_tanhf(%{{.*}}) : (f32) -> f32 319 %result64 = math.tanh %arg_f64 : f64 320 // CHECK: llvm.call @__nv_tanh(%{{.*}}) : (f64) -> f64 321 std.return %result16, %result32, %result64 : f16, f32, f64 322 } 323} 324 325// ----- 326 327gpu.module @test_module { 328 // CHECK: llvm.func @__nv_rsqrtf(f32) -> f32 329 // CHECK: llvm.func @__nv_rsqrt(f64) -> f64 330 // CHECK-LABEL: func @gpu_rsqrt 331 builtin.func @gpu_rsqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) 332 -> (f16, f32, f64) { 333 %result16 = math.rsqrt %arg_f16 : f16 334 // CHECK: llvm.fpext %{{.*}} : f16 to f32 335 // CHECK-NEXT: llvm.call @__nv_rsqrtf(%{{.*}}) : (f32) -> f32 336 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 337 %result32 = math.rsqrt %arg_f32 : f32 338 // CHECK: llvm.call @__nv_rsqrtf(%{{.*}}) : (f32) -> f32 339 %result64 = math.rsqrt %arg_f64 : f64 340 // CHECK: llvm.call @__nv_rsqrt(%{{.*}}) : (f64) -> f64 341 std.return %result16, %result32, %result64 : f16, f32, f64 342 } 343} 344 345// ----- 346 347gpu.module @test_module { 348 // CHECK: llvm.func @__nv_sqrtf(f32) -> f32 349 // CHECK: llvm.func @__nv_sqrt(f64) -> f64 350 // CHECK-LABEL: func @gpu_sqrt 351 builtin.func @gpu_sqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) 352 -> (f16, f32, f64) { 353 %result16 = math.sqrt %arg_f16 : f16 354 // CHECK: llvm.fpext %{{.*}} : f16 to f32 355 // CHECK-NEXT: llvm.call @__nv_sqrtf(%{{.*}}) : (f32) -> f32 356 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 357 %result32 = math.sqrt %arg_f32 : f32 358 // CHECK: llvm.call @__nv_sqrtf(%{{.*}}) : (f32) -> f32 359 %result64 = math.sqrt %arg_f64 : f64 360 // CHECK: llvm.call @__nv_sqrt(%{{.*}}) : (f64) -> f64 361 std.return %result16, %result32, %result64 : f16, f32, f64 362 } 363} 364 365// ----- 366 367gpu.module @test_module { 368 // CHECK: llvm.func @__nv_atanf(f32) -> f32 369 // CHECK: llvm.func @__nv_atan(f64) -> f64 370 // CHECK-LABEL: func @gpu_atan 371 builtin.func @gpu_atan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) 372 -> (f16, f32, f64) { 373 %result16 = math.atan %arg_f16 : f16 374 // CHECK: llvm.fpext %{{.*}} : f16 to f32 375 // CHECK-NEXT: llvm.call @__nv_atanf(%{{.*}}) : (f32) -> f32 376 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 377 %result32 = math.atan %arg_f32 : f32 378 // CHECK: llvm.call @__nv_atanf(%{{.*}}) : (f32) -> f32 379 %result64 = math.atan %arg_f64 : f64 380 // CHECK: llvm.call @__nv_atan(%{{.*}}) : (f64) -> f64 381 std.return %result16, %result32, %result64 : f16, f32, f64 382 } 383} 384 385// ----- 386 387gpu.module @test_module { 388 // CHECK: llvm.func @__nv_atan2f(f32, f32) -> f32 389 // CHECK: llvm.func @__nv_atan2(f64, f64) -> f64 390 // CHECK-LABEL: func @gpu_atan2 391 builtin.func @gpu_atan2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) 392 -> (f16, f32, f64) { 393 %result16 = math.atan2 %arg_f16, %arg_f16 : f16 394 // CHECK: llvm.fpext %{{.*}} : f16 to f32 395 // CHECK: llvm.fpext %{{.*}} : f16 to f32 396 // CHECK-NEXT: llvm.call @__nv_atan2f(%{{.*}}) : (f32, f32) -> f32 397 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16 398 %result32 = math.atan2 %arg_f32, %arg_f32 : f32 399 // CHECK: llvm.call @__nv_atan2f(%{{.*}}) : (f32, f32) -> f32 400 %result64 = math.atan2 %arg_f64, %arg_f64 : f64 401 // CHECK: llvm.call @__nv_atan2(%{{.*}}) : (f64, f64) -> f64 402 std.return %result16, %result32, %result64 : f16, f32, f64 403 } 404} 405 406// ----- 407 408// Test that we handled properly operation with SymbolTable other than module op 409gpu.module @test_module { 410 "test.symbol_scope"() ({ 411 // CHECK: test.symbol_scope 412 // CHECK: llvm.func @__nv_expf(f32) -> f32 413 // CHECK: llvm.func @__nv_exp(f64) -> f64 414 // CHECK-LABEL: func @gpu_exp 415 builtin.func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 416 %result32 = math.exp %arg_f32 : f32 417 // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32 418 %result64 = math.exp %arg_f64 : f64 419 // CHECK: llvm.call @__nv_exp(%{{.*}}) : (f64) -> f64 420 std.return %result32, %result64 : f32, f64 421 } 422 "test.finish" () : () -> () 423 }) : () -> () 424} 425 426// ----- 427 428gpu.module @test_module { 429 // CHECK: llvm.func @__nv_expm1f(f32) -> f32 430 // CHECK: llvm.func @__nv_expm1(f64) -> f64 431 // CHECK-LABEL: func @gpu_expm1 432 builtin.func @gpu_expm1(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 433 %result32 = math.expm1 %arg_f32 : f32 434 // CHECK: llvm.call @__nv_expm1f(%{{.*}}) : (f32) -> f32 435 %result64 = math.expm1 %arg_f64 : f64 436 // CHECK: llvm.call @__nv_expm1(%{{.*}}) : (f64) -> f64 437 std.return %result32, %result64 : f32, f64 438 } 439} 440 441// ----- 442 443gpu.module @test_module { 444 // CHECK: llvm.func @__nv_powf(f32, f32) -> f32 445 // CHECK: llvm.func @__nv_pow(f64, f64) -> f64 446 // CHECK-LABEL: func @gpu_pow 447 builtin.func @gpu_pow(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) { 448 %result32 = math.powf %arg_f32, %arg_f32 : f32 449 // CHECK: llvm.call @__nv_powf(%{{.*}}, %{{.*}}) : (f32, f32) -> f32 450 %result64 = math.powf %arg_f64, %arg_f64 : f64 451 // CHECK: llvm.call @__nv_pow(%{{.*}}, %{{.*}}) : (f64, f64) -> f64 452 std.return %result32, %result64 : f32, f64 453 } 454} 455 456// ----- 457 458gpu.module @test_module { 459 // CHECK-LABEL: @kernel_func 460 // CHECK: attributes 461 // CHECK: gpu.kernel 462 // CHECK: nvvm.kernel 463 gpu.func @kernel_func() kernel { 464 gpu.return 465 } 466} 467