1152 lines · plain
1// RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1' -split-input-file | FileCheck %s2// RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1 allow-pattern-rollback=0' -split-input-file | FileCheck %s3// RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1 allowed-dialects=func,arith,cf' -split-input-file | FileCheck %s4// RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1 use-bare-ptr-memref-call-conv=1' -split-input-file | FileCheck %s --check-prefix=CHECK-BARE5// RUN: mlir-opt %s -transform-interpreter | FileCheck %s6 7gpu.module @test_module_0 {8 // CHECK-LABEL: func @gpu_index_ops()9 func.func @gpu_index_ops()10 -> (index, index, index, index, index, index,11 index, index, index, index, index, index,12 index) {13 14 // CHECK: = nvvm.read.ptx.sreg.tid.x : i3215 // CHECK: = llvm.sext %{{.*}} : i32 to i6416 %tIdX = gpu.thread_id x17 // CHECK: = nvvm.read.ptx.sreg.tid.y : i3218 // CHECK: = llvm.sext %{{.*}} : i32 to i6419 %tIdY = gpu.thread_id y20 // CHECK: = nvvm.read.ptx.sreg.tid.z : i3221 // CHECK: = llvm.sext %{{.*}} : i32 to i6422 %tIdZ = gpu.thread_id z23 24 // CHECK: = nvvm.read.ptx.sreg.ntid.x : i3225 // CHECK: = llvm.sext %{{.*}} : i32 to i6426 %bDimX = gpu.block_dim x27 // CHECK: = nvvm.read.ptx.sreg.ntid.y : i3228 // CHECK: = llvm.sext %{{.*}} : i32 to i6429 %bDimY = gpu.block_dim y30 // CHECK: = nvvm.read.ptx.sreg.ntid.z : i3231 // CHECK: = llvm.sext %{{.*}} : i32 to i6432 %bDimZ = gpu.block_dim z33 34 // CHECK: = nvvm.read.ptx.sreg.ctaid.x : i3235 // CHECK: = llvm.sext %{{.*}} : i32 to i6436 %bIdX = gpu.block_id x37 // CHECK: = nvvm.read.ptx.sreg.ctaid.y : i3238 // CHECK: = llvm.sext %{{.*}} : i32 to i6439 %bIdY = gpu.block_id y40 // CHECK: = nvvm.read.ptx.sreg.ctaid.z : i3241 // CHECK: = llvm.sext %{{.*}} : i32 to i6442 %bIdZ = gpu.block_id z43 44 // CHECK: = nvvm.read.ptx.sreg.nctaid.x : i3245 // CHECK: = llvm.sext %{{.*}} : i32 to i6446 %gDimX = gpu.grid_dim x47 // CHECK: = nvvm.read.ptx.sreg.nctaid.y : i3248 // CHECK: = llvm.sext %{{.*}} : i32 to i6449 %gDimY = gpu.grid_dim y50 // CHECK: = nvvm.read.ptx.sreg.nctaid.z : i3251 // CHECK: = llvm.sext %{{.*}} : i32 to i6452 %gDimZ = gpu.grid_dim z53 54 55 // CHECK: = nvvm.read.ptx.sreg.laneid range <i32, 0, 32> : i3256 // CHECK: = llvm.sext %{{.*}} : i32 to i6457 %laneId = gpu.lane_id58 59 func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,60 %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,61 %laneId62 : index, index, index, index, index, index,63 index, index, index, index, index, index,64 index65 }66}67 68 69 70gpu.module @test_module_1 {71 // CHECK-LABEL: func @gpu_index_comp72 func.func @gpu_index_comp(%idx : index) -> index {73 // CHECK: = llvm.add %{{.*}}, %{{.*}} : i6474 %0 = arith.addi %idx, %idx : index75 // CHECK: llvm.return %{{.*}} : i6476 func.return %0 : index77 }78}79 80 81 82gpu.module @test_module_2 {83 // CHECK-LABEL: func @gpu_all_reduce_op()84 gpu.func @gpu_all_reduce_op() {85 %arg0 = arith.constant 1.0 : f3286 // TODO: Check full IR expansion once lowering has settled.87 // CHECK: nvvm.shfl.sync bfly {{.*}}88 // CHECK: nvvm.barrier089 // CHECK: llvm.fadd90 %result = gpu.all_reduce add %arg0 uniform {} : (f32) -> (f32)91 92 gpu.return93 }94}95 96 97 98gpu.module @test_module_3 {99 // CHECK-LABEL: func @gpu_all_reduce_region()100 gpu.func @gpu_all_reduce_region() {101 %arg0 = arith.constant 1 : i32102 // TODO: Check full IR expansion once lowering has settled.103 // CHECK: nvvm.shfl.sync bfly {{.*}}104 // CHECK: nvvm.barrier0105 %result = gpu.all_reduce %arg0 uniform {106 ^bb(%lhs : i32, %rhs : i32):107 %xor = arith.xori %lhs, %rhs : i32108 "gpu.yield"(%xor) : (i32) -> ()109 } : (i32) -> (i32)110 gpu.return111 }112}113 114 115 116gpu.module @test_module_4 {117 // CHECK-LABEL: func @gpu_shuffle()118 func.func @gpu_shuffle() -> (f32, f32, f32, f32, i1, i1, i1, i1) {119 // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : f32120 %arg0 = arith.constant 1.0 : f32121 // CHECK: %[[#OFFSET:]] = llvm.mlir.constant(4 : i32) : i32122 %arg1 = arith.constant 4 : i32123 // CHECK: %[[#WIDTH:]] = llvm.mlir.constant(23 : i32) : i32124 %arg2 = arith.constant 23 : i32125 // CHECK: %[[#ONE:]] = llvm.mlir.constant(1 : i32) : i32126 // CHECK: %[[#MINUS_ONE:]] = llvm.mlir.constant(-1 : i32) : i32127 // CHECK: %[[#THIRTY_TWO:]] = llvm.mlir.constant(32 : i32) : i32128 // CHECK: %[[#NUM_LANES:]] = llvm.sub %[[#THIRTY_TWO]], %[[#WIDTH]] : i32129 // CHECK: %[[#MASK:]] = llvm.lshr %[[#MINUS_ONE]], %[[#NUM_LANES]] : i32130 // CHECK: %[[#CLAMP:]] = llvm.sub %[[#WIDTH]], %[[#ONE]] : i32131 // CHECK: %[[#SHFL:]] = nvvm.shfl.sync bfly %[[#MASK]], %[[#VALUE]], %[[#OFFSET]], %[[#CLAMP]] {return_value_and_is_valid} : f32 -> !llvm.struct<(f32, i1)>132 // CHECK: llvm.extractvalue %[[#SHFL]][0] : !llvm.struct<(f32, i1)>133 // CHECK: llvm.extractvalue %[[#SHFL]][1] : !llvm.struct<(f32, i1)>134 %shfl, %pred = gpu.shuffle xor %arg0, %arg1, %arg2 : f32135 // CHECK: %[[#ONE:]] = llvm.mlir.constant(1 : i32) : i32136 // CHECK: %[[#MINUS_ONE:]] = llvm.mlir.constant(-1 : i32) : i32137 // CHECK: %[[#THIRTY_TWO:]] = llvm.mlir.constant(32 : i32) : i32138 // CHECK: %[[#NUM_LANES:]] = llvm.sub %[[#THIRTY_TWO]], %[[#WIDTH]] : i32139 // CHECK: %[[#MASK:]] = llvm.lshr %[[#MINUS_ONE]], %[[#NUM_LANES]] : i32140 // CHECK: %[[#SHFL:]] = nvvm.shfl.sync up %[[#MASK]], %[[#VALUE]], %[[#OFFSET]], %[[#NUM_LANES]] {return_value_and_is_valid} : f32 -> !llvm.struct<(f32, i1)>141 // CHECK: llvm.extractvalue %[[#SHFL]][0] : !llvm.struct<(f32, i1)>142 // CHECK: llvm.extractvalue %[[#SHFL]][1] : !llvm.struct<(f32, i1)>143 %shflu, %predu = gpu.shuffle up %arg0, %arg1, %arg2 : f32144 // CHECK: nvvm.shfl.sync down {{.*}} {return_value_and_is_valid} : f32 -> !llvm.struct<(f32, i1)>145 %shfld, %predd = gpu.shuffle down %arg0, %arg1, %arg2 : f32146 // CHECK: nvvm.shfl.sync idx {{.*}} {return_value_and_is_valid} : f32 -> !llvm.struct<(f32, i1)>147 %shfli, %predi = gpu.shuffle idx %arg0, %arg1, %arg2 : f32148 149 func.return %shfl, %shflu, %shfld, %shfli, %pred, %predu, %predd, %predi150 : f32, f32,f32, f32, i1, i1, i1, i1151 }152 153 // CHECK-LABEL: func @gpu_shuffle_unused_pred()154 func.func @gpu_shuffle_unused_pred() -> (f32, f32, f32, f32) {155 // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : f32156 %arg0 = arith.constant 1.0 : f32157 // CHECK: %[[#OFFSET:]] = llvm.mlir.constant(4 : i32) : i32158 %arg1 = arith.constant 4 : i32159 // CHECK: %[[#WIDTH:]] = llvm.mlir.constant(23 : i32) : i32160 %arg2 = arith.constant 23 : i32161 // CHECK: %[[#ONE:]] = llvm.mlir.constant(1 : i32) : i32162 // CHECK: %[[#MINUS_ONE:]] = llvm.mlir.constant(-1 : i32) : i32163 // CHECK: %[[#THIRTY_TWO:]] = llvm.mlir.constant(32 : i32) : i32164 // CHECK: %[[#NUM_LANES:]] = llvm.sub %[[#THIRTY_TWO]], %[[#WIDTH]] : i32165 // CHECK: %[[#MASK:]] = llvm.lshr %[[#MINUS_ONE]], %[[#NUM_LANES]] : i32166 // CHECK: %[[#CLAMP:]] = llvm.sub %[[#WIDTH]], %[[#ONE]] : i32167 // CHECK: %[[#SHFL:]] = nvvm.shfl.sync bfly %[[#MASK]], %[[#VALUE]], %[[#OFFSET]], %[[#CLAMP]] : f32 -> f32168 %shfl, %pred = gpu.shuffle xor %arg0, %arg1, %arg2 : f32169 // CHECK: %[[#ONE:]] = llvm.mlir.constant(1 : i32) : i32170 // CHECK: %[[#MINUS_ONE:]] = llvm.mlir.constant(-1 : i32) : i32171 // CHECK: %[[#THIRTY_TWO:]] = llvm.mlir.constant(32 : i32) : i32172 // CHECK: %[[#NUM_LANES:]] = llvm.sub %[[#THIRTY_TWO]], %[[#WIDTH]] : i32173 // CHECK: %[[#MASK:]] = llvm.lshr %[[#MINUS_ONE]], %[[#NUM_LANES]] : i32174 // CHECK: %[[#SHFL:]] = nvvm.shfl.sync up %[[#MASK]], %[[#VALUE]], %[[#OFFSET]], %[[#NUM_LANES]] : f32 -> f32175 %shflu, %predu = gpu.shuffle up %arg0, %arg1, %arg2 : f32176 // CHECK: nvvm.shfl.sync down {{.*}} : f32 -> f32177 %shfld, %predd = gpu.shuffle down %arg0, %arg1, %arg2 : f32178 // CHECK: nvvm.shfl.sync idx {{.*}} : f32 -> f32179 %shfli, %predi = gpu.shuffle idx %arg0, %arg1, %arg2 : f32180 181 func.return %shfl, %shflu, %shfld, %shfli : f32, f32,f32, f32182 }183}184 185gpu.module @test_module_5 {186 // CHECK-LABEL: func @gpu_sync()187 func.func @gpu_sync() {188 // CHECK: nvvm.barrier0189 gpu.barrier190 func.return191 }192}193 194 195 196gpu.module @test_module_6 {197 // CHECK: llvm.func @__nv_fabsf(f32) -> f32198 // CHECK: llvm.func @__nv_fabs(f64) -> f64199 // CHECK-LABEL: func @gpu_fabs200 func.func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {201 %result32 = math.absf %arg_f32 : f32202 // CHECK: llvm.call @__nv_fabsf(%{{.*}}) : (f32) -> f32203 %result64 = math.absf %arg_f64 : f64204 // CHECK: llvm.call @__nv_fabs(%{{.*}}) : (f64) -> f64205 func.return %result32, %result64 : f32, f64206 }207}208 209 210 211gpu.module @test_module_7 {212 // CHECK: llvm.func @__nv_cbrtf(f32) -> f32213 // CHECK: llvm.func @__nv_cbrt(f64) -> f64214 // CHECK-LABEL: func @gpu_cbrt215 func.func @gpu_cbrt(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {216 %result32 = math.cbrt %arg_f32 : f32217 // CHECK: llvm.call @__nv_cbrtf(%{{.*}}) : (f32) -> f32218 %result64 = math.cbrt %arg_f64 : f64219 // CHECK: llvm.call @__nv_cbrt(%{{.*}}) : (f64) -> f64220 func.return %result32, %result64 : f32, f64221 }222}223 224 225 226gpu.module @test_module_8 {227 // CHECK: llvm.func @__nv_ceilf(f32) -> f32228 // CHECK: llvm.func @__nv_ceil(f64) -> f64229 // CHECK-LABEL: func @gpu_ceil230 func.func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {231 %result32 = math.ceil %arg_f32 : f32232 // CHECK: llvm.call @__nv_ceilf(%{{.*}}) : (f32) -> f32233 %result64 = math.ceil %arg_f64 : f64234 // CHECK: llvm.call @__nv_ceil(%{{.*}}) : (f64) -> f64235 func.return %result32, %result64 : f32, f64236 }237}238 239 240 241gpu.module @test_module_9 {242 // CHECK: llvm.func @__nv_floorf(f32) -> f32243 // CHECK: llvm.func @__nv_floor(f64) -> f64244 // CHECK-LABEL: func @gpu_floor245 func.func @gpu_floor(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {246 %result32 = math.floor %arg_f32 : f32247 // CHECK: llvm.call @__nv_floorf(%{{.*}}) : (f32) -> f32248 %result64 = math.floor %arg_f64 : f64249 // CHECK: llvm.call @__nv_floor(%{{.*}}) : (f64) -> f64250 func.return %result32, %result64 : f32, f64251 }252}253 254 255 256gpu.module @test_module_10 {257 // CHECK: llvm.func @__nv_cosf(f32) -> f32258 // CHECK: llvm.func @__nv_cos(f64) -> f64259 // CHECK: llvm.func @__nv_fast_cosf(f32) -> f32260 // CHECK-LABEL: func @gpu_cos261 func.func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {262 %result32 = math.cos %arg_f32 : f32263 // CHECK: llvm.call @__nv_cosf(%{{.*}}) : (f32) -> f32264 %result64 = math.cos %arg_f64 : f64265 // CHECK: llvm.call @__nv_cos(%{{.*}}) : (f64) -> f64266 %result32Fast = math.cos %arg_f32 fastmath<afn> : f32267 // CHECK: llvm.call @__nv_fast_cosf(%{{.*}}) : (f32) -> f32268 func.return %result32, %result64, %result32Fast : f32, f64, f32269 }270}271 272 273gpu.module @test_module_11 {274 // CHECK: llvm.func @__nv_expf(f32) -> f32275 // CHECK: llvm.func @__nv_exp(f64) -> f64276 // CHECK: llvm.func @__nv_fast_expf(f32) -> f32277 // CHECK-LABEL: func @gpu_exp278 func.func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {279 %result32 = math.exp %arg_f32 : f32280 // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32281 %result64 = math.exp %arg_f64 : f64282 // CHECK: llvm.call @__nv_exp(%{{.*}}) : (f64) -> f64283 %result32Fast = math.exp %arg_f32 fastmath<fast> : f32284 // CHECK: llvm.call @__nv_fast_expf(%{{.*}}) : (f32) -> f32285 func.return %result32, %result64, %result32Fast : f32, f64, f32286 }287}288 289 290gpu.module @test_module_12 {291 // CHECK: llvm.func @__nv_exp2f(f32) -> f32292 // CHECK: llvm.func @__nv_exp2(f64) -> f64293 // CHECK-LABEL: func @gpu_exp2294 func.func @gpu_exp2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {295 %result32 = math.exp2 %arg_f32 : f32296 // CHECK: llvm.call @__nv_exp2f(%{{.*}}) : (f32) -> f32297 %result64 = math.exp2 %arg_f64 : f64298 // CHECK: llvm.call @__nv_exp2(%{{.*}}) : (f64) -> f64299 func.return %result32, %result64 : f32, f64300 }301}302 303 304 305gpu.module @test_module_13 {306 // CHECK: llvm.func @__nv_logf(f32) -> f32307 // CHECK: llvm.func @__nv_log(f64) -> f64308 // CHECK: llvm.func @__nv_fast_logf(f32) -> f32309 // CHECK-LABEL: func @gpu_log310 func.func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {311 %result32 = math.log %arg_f32 : f32312 // CHECK: llvm.call @__nv_logf(%{{.*}}) : (f32) -> f32313 %result64 = math.log %arg_f64 : f64314 // CHECK: llvm.call @__nv_log(%{{.*}}) : (f64) -> f64315 %result32Fast = math.log %arg_f32 fastmath<afn> : f32316 // CHECK: llvm.call @__nv_fast_logf(%{{.*}}) : (f32) -> f32317 func.return %result32, %result64, %result32Fast : f32, f64, f32318 }319}320 321 322 323gpu.module @test_module_14 {324 // CHECK: llvm.func @__nv_log10f(f32) -> f32325 // CHECK: llvm.func @__nv_log10(f64) -> f64326 // CHECK: llvm.func @__nv_fast_log10f(f32) -> f32327 // CHECK-LABEL: func @gpu_log10328 func.func @gpu_log10(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {329 %result32 = math.log10 %arg_f32 : f32330 // CHECK: llvm.call @__nv_log10f(%{{.*}}) : (f32) -> f32331 %result64 = math.log10 %arg_f64 : f64332 // CHECK: llvm.call @__nv_log10(%{{.*}}) : (f64) -> f64333 %result32Fast = math.log10 %arg_f32 fastmath<afn> : f32334 // CHECK: llvm.call @__nv_fast_log10f(%{{.*}}) : (f32) -> f32335 func.return %result32, %result64, %result32Fast : f32, f64, f32336 }337}338 339 340 341gpu.module @test_module_15 {342 // CHECK: llvm.func @__nv_log1pf(f32) -> f32343 // CHECK: llvm.func @__nv_log1p(f64) -> f64344 // CHECK-LABEL: func @gpu_log1p345 func.func @gpu_log1p(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {346 %result32 = math.log1p %arg_f32 : f32347 // CHECK: llvm.call @__nv_log1pf(%{{.*}}) : (f32) -> f32348 %result64 = math.log1p %arg_f64 : f64349 // CHECK: llvm.call @__nv_log1p(%{{.*}}) : (f64) -> f64350 func.return %result32, %result64 : f32, f64351 }352}353 354 355 356gpu.module @test_module_16 {357 // CHECK: llvm.func @__nv_log2f(f32) -> f32358 // CHECK: llvm.func @__nv_log2(f64) -> f64359 // CHECK: llvm.func @__nv_fast_log2f(f32) -> f32360 // CHECK-LABEL: func @gpu_log2361 func.func @gpu_log2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {362 %result32 = math.log2 %arg_f32 : f32363 // CHECK: llvm.call @__nv_log2f(%{{.*}}) : (f32) -> f32364 %result64 = math.log2 %arg_f64 : f64365 // CHECK: llvm.call @__nv_log2(%{{.*}}) : (f64) -> f64366 %result32Fast = math.log2 %arg_f32 fastmath<fast> : f32367 // CHECK: llvm.call @__nv_fast_log2f(%{{.*}}) : (f32) -> f32368 func.return %result32, %result64, %result32Fast : f32, f64, f32369 }370}371 372 373 374gpu.module @test_module_17 {375 // CHECK: llvm.func @__nv_sinf(f32) -> f32376 // CHECK: llvm.func @__nv_sin(f64) -> f64377 // CHECK: llvm.func @__nv_fast_sinf(f32) -> f32378 // CHECK-LABEL: func @gpu_sin379 func.func @gpu_sin(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {380 %result32 = math.sin %arg_f32 : f32381 // CHECK: llvm.call @__nv_sinf(%{{.*}}) : (f32) -> f32382 %result64 = math.sin %arg_f64 : f64383 // CHECK: llvm.call @__nv_sin(%{{.*}}) : (f64) -> f64384 %result32Fast = math.sin %arg_f32 fastmath<fast> : f32385 // CHECK: llvm.call @__nv_fast_sinf(%{{.*}}) : (f32) -> f32386 func.return %result32, %result64, %result32Fast : f32, f64, f32387 }388}389 390 391 392gpu.module @test_module_18 {393 // CHECK: llvm.func @__nv_tanf(f32) -> f32394 // CHECK: llvm.func @__nv_tan(f64) -> f64395 // CHECK: llvm.func @__nv_fast_tanf(f32) -> f32396 // CHECK-LABEL: func @gpu_tan397 func.func @gpu_tan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64, f32) {398 %result16 = math.tan %arg_f16 : f16399 // CHECK: llvm.fpext %{{.*}} : f16 to f32400 // CHECK-NEXT: llvm.call @__nv_tanf(%{{.*}}) : (f32) -> f32401 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16402 %result32 = math.tan %arg_f32 : f32403 // CHECK: llvm.call @__nv_tanf(%{{.*}}) : (f32) -> f32404 %result64 = math.tan %arg_f64 : f64405 // CHECK: llvm.call @__nv_tan(%{{.*}}) : (f64) -> f64406 %result32Fast = math.tan %arg_f32 fastmath<fast> : f32407 // CHECK: llvm.call @__nv_fast_tanf(%{{.*}}) : (f32) -> f32408 func.return %result16, %result32, %result64, %result32Fast : f16, f32, f64, f32409 }410}411 412 413 414gpu.module @test_module_19 {415 // CHECK: llvm.func @__nv_tanhf(f32) -> f32416 // CHECK: llvm.func @__nv_tanh(f64) -> f64417 // CHECK-LABEL: func @gpu_tanh418 func.func @gpu_tanh(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {419 %result16 = math.tanh %arg_f16 : f16420 // CHECK: llvm.fpext %{{.*}} : f16 to f32421 // CHECK-NEXT: llvm.call @__nv_tanhf(%{{.*}}) : (f32) -> f32422 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16423 %result32 = math.tanh %arg_f32 : f32424 // CHECK: llvm.call @__nv_tanhf(%{{.*}}) : (f32) -> f32425 %result64 = math.tanh %arg_f64 : f64426 // CHECK: llvm.call @__nv_tanh(%{{.*}}) : (f64) -> f64427 func.return %result16, %result32, %result64 : f16, f32, f64428 }429}430 431 432 433gpu.module @test_module_20 {434 // CHECK: llvm.func @__nv_rsqrtf(f32) -> f32435 // CHECK: llvm.func @__nv_rsqrt(f64) -> f64436 // CHECK-LABEL: func @gpu_rsqrt437 func.func @gpu_rsqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64)438 -> (f16, f32, f64) {439 %result16 = math.rsqrt %arg_f16 : f16440 // CHECK: llvm.fpext %{{.*}} : f16 to f32441 // CHECK-NEXT: llvm.call @__nv_rsqrtf(%{{.*}}) : (f32) -> f32442 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16443 %result32 = math.rsqrt %arg_f32 : f32444 // CHECK: llvm.call @__nv_rsqrtf(%{{.*}}) : (f32) -> f32445 %result64 = math.rsqrt %arg_f64 : f64446 // CHECK: llvm.call @__nv_rsqrt(%{{.*}}) : (f64) -> f64447 func.return %result16, %result32, %result64 : f16, f32, f64448 }449}450 451 452 453gpu.module @test_module_21 {454 // CHECK: llvm.func @__nv_sqrtf(f32) -> f32455 // CHECK: llvm.func @__nv_sqrt(f64) -> f64456 // CHECK-LABEL: func @gpu_sqrt457 func.func @gpu_sqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64)458 -> (f16, f32, f64) {459 %result16 = math.sqrt %arg_f16 : f16460 // CHECK: llvm.fpext %{{.*}} : f16 to f32461 // CHECK-NEXT: llvm.call @__nv_sqrtf(%{{.*}}) : (f32) -> f32462 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16463 %result32 = math.sqrt %arg_f32 : f32464 // CHECK: llvm.call @__nv_sqrtf(%{{.*}}) : (f32) -> f32465 %result64 = math.sqrt %arg_f64 : f64466 // CHECK: llvm.call @__nv_sqrt(%{{.*}}) : (f64) -> f64467 func.return %result16, %result32, %result64 : f16, f32, f64468 }469}470 471 472 473gpu.module @test_module_22 {474 // CHECK: llvm.func @__nv_atanf(f32) -> f32475 // CHECK: llvm.func @__nv_atan(f64) -> f64476 // CHECK-LABEL: func @gpu_atan477 func.func @gpu_atan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64)478 -> (f16, f32, f64) {479 %result16 = math.atan %arg_f16 : f16480 // CHECK: llvm.fpext %{{.*}} : f16 to f32481 // CHECK-NEXT: llvm.call @__nv_atanf(%{{.*}}) : (f32) -> f32482 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16483 %result32 = math.atan %arg_f32 : f32484 // CHECK: llvm.call @__nv_atanf(%{{.*}}) : (f32) -> f32485 %result64 = math.atan %arg_f64 : f64486 // CHECK: llvm.call @__nv_atan(%{{.*}}) : (f64) -> f64487 func.return %result16, %result32, %result64 : f16, f32, f64488 }489}490 491 492 493gpu.module @test_module_23 {494 // CHECK: llvm.func @__nv_atan2f(f32, f32) -> f32495 // CHECK: llvm.func @__nv_atan2(f64, f64) -> f64496 // CHECK-LABEL: func @gpu_atan2497 func.func @gpu_atan2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64)498 -> (f16, f32, f64) {499 %result16 = math.atan2 %arg_f16, %arg_f16 : f16500 // CHECK: llvm.fpext %{{.*}} : f16 to f32501 // CHECK: llvm.fpext %{{.*}} : f16 to f32502 // CHECK-NEXT: llvm.call @__nv_atan2f(%{{.*}}) : (f32, f32) -> f32503 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16504 %result32 = math.atan2 %arg_f32, %arg_f32 : f32505 // CHECK: llvm.call @__nv_atan2f(%{{.*}}) : (f32, f32) -> f32506 %result64 = math.atan2 %arg_f64, %arg_f64 : f64507 // CHECK: llvm.call @__nv_atan2(%{{.*}}) : (f64, f64) -> f64508 func.return %result16, %result32, %result64 : f16, f32, f64509 }510}511 512 513 514// Test that we handled properly operation with SymbolTable other than module op515gpu.module @test_module_24 {516 "test.symbol_scope"() ({517 // CHECK: test.symbol_scope518 // CHECK: llvm.func @__nv_expf(f32) -> f32519 // CHECK: llvm.func @__nv_exp(f64) -> f64520 // CHECK: llvm.func @__nv_fast_expf(f32) -> f32521 // CHECK-LABEL: func @gpu_exp522 func.func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {523 %result32 = math.exp %arg_f32 : f32524 // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32525 %result64 = math.exp %arg_f64 : f64526 // CHECK: llvm.call @__nv_exp(%{{.*}}) : (f64) -> f64527 %result32Fast = math.exp %arg_f32 fastmath<afn> : f32528 // CHECK: llvm.call @__nv_fast_expf(%{{.*}}) : (f32) -> f32529 func.return %result32, %result64, %result32Fast : f32, f64, f32530 }531 "test.finish" () : () -> ()532 }) : () -> ()533}534 535 536 537gpu.module @test_module_25 {538 // CHECK: llvm.func @__nv_expm1f(f32) -> f32539 // CHECK: llvm.func @__nv_expm1(f64) -> f64540 // CHECK-LABEL: func @gpu_expm1541 func.func @gpu_expm1(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {542 %result32 = math.expm1 %arg_f32 : f32543 // CHECK: llvm.call @__nv_expm1f(%{{.*}}) : (f32) -> f32544 %result64 = math.expm1 %arg_f64 : f64545 // CHECK: llvm.call @__nv_expm1(%{{.*}}) : (f64) -> f64546 func.return %result32, %result64 : f32, f64547 }548}549 550 551 552gpu.module @test_module_26 {553 // CHECK: llvm.func @__nv_powf(f32, f32) -> f32554 // CHECK: llvm.func @__nv_pow(f64, f64) -> f64555 // CHECK: llvm.func @__nv_fast_powf(f32, f32) -> f32556 // CHECK-LABEL: func @gpu_pow557 func.func @gpu_pow(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {558 %result32 = math.powf %arg_f32, %arg_f32 : f32559 // CHECK: llvm.call @__nv_powf(%{{.*}}, %{{.*}}) : (f32, f32) -> f32560 %result64 = math.powf %arg_f64, %arg_f64 : f64561 // CHECK: llvm.call @__nv_pow(%{{.*}}, %{{.*}}) : (f64, f64) -> f64562 %result32Fast = math.powf %arg_f32, %arg_f32 fastmath<fast> : f32563 // CHECK: llvm.call @__nv_fast_powf(%{{.*}}, %{{.*}}) : (f32, f32) -> f32564 func.return %result32, %result64, %result32Fast : f32, f64, f32565 }566}567 568 569 570gpu.module @test_module_27 {571 // CHECK-LABEL: func @gpu_unroll572 func.func @gpu_unroll(%arg0 : vector<4xf32>) -> vector<4xf32> {573 %result = math.exp %arg0 : vector<4xf32>574 // CHECK: %[[V0:.+]] = llvm.mlir.poison : vector<4xf32>575 // CHECK: %[[CL:.+]] = llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32576 // CHECK: %[[V1:.+]] = llvm.insertelement %[[CL]], %[[V0]]577 // CHECK: %[[CL:.+]] = llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32578 // CHECK: %[[V2:.+]] = llvm.insertelement %[[CL]], %[[V1]]579 // CHECK: %[[CL:.+]] = llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32580 // CHECK: %[[V3:.+]] = llvm.insertelement %[[CL]], %[[V2]]581 // CHECK: %[[CL:.+]] = llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32582 // CHECK: %[[V4:.+]] = llvm.insertelement %[[CL]], %[[V3]]583 // CHECK: return %[[V4]]584 func.return %result : vector<4xf32>585 }586}587 588 589 590gpu.module @test_module_28 {591 // CHECK-LABEL: @kernel_func592 // CHECK: attributes593 // CHECK: gpu.kernel594 // CHECK: nvvm.kernel595 gpu.func @kernel_func() kernel {596 gpu.return597 }598}599 600 601 602gpu.module @test_module_29 {603 // CHECK-DAG: llvm.mlir.global internal constant @[[$PRINT_GLOBAL0:[A-Za-z0-9_]+]]("Hello, world\0A\00")604 // CHECK-DAG: llvm.mlir.global internal constant @[[$PRINT_GLOBAL1:[A-Za-z0-9_]+]]("Hello: %d\0A\00")605 // CHECK-DAG: llvm.func @vprintf(!llvm.ptr, !llvm.ptr) -> i32606 607 // CHECK-LABEL: func @test_const_printf608 gpu.func @test_const_printf() {609 // CHECK-NEXT: %[[FORMATSTR:.*]] = llvm.mlir.addressof @[[$PRINT_GLOBAL0]] : !llvm.ptr610 // CHECK-NEXT: %[[FORMATSTART:.*]] = llvm.getelementptr %[[FORMATSTR]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<14 x i8>611 // CHECK-NEXT: %[[O:.*]] = llvm.mlir.constant(1 : index) : i64612 // CHECK-NEXT: %[[ALLOC:.*]] = llvm.alloca %[[O]] x !llvm.struct<()> : (i64) -> !llvm.ptr613 // CHECK-NEXT: llvm.call @vprintf(%[[FORMATSTART]], %[[ALLOC]]) : (!llvm.ptr, !llvm.ptr) -> i32614 gpu.printf "Hello, world\n"615 616 // Make sure that the same global is reused.617 // CHECK: %[[FORMATSTR2:.*]] = llvm.mlir.addressof @[[$PRINT_GLOBAL0]] : !llvm.ptr618 // CHECK: %[[FORMATSTART2:.*]] = llvm.getelementptr %[[FORMATSTR2]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<14 x i8>619 // CHECK: llvm.call @vprintf(%[[FORMATSTART2]], %{{.*}}) : (!llvm.ptr, !llvm.ptr) -> i32620 gpu.printf "Hello, world\n"621 622 gpu.return623 }624 625 // CHECK-LABEL: func @test_printf626 // CHECK: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: f32)627 gpu.func @test_printf(%arg0: i32, %arg1: f32) {628 // CHECK-NEXT: %[[FORMATSTR:.*]] = llvm.mlir.addressof @[[$PRINT_GLOBAL1]] : !llvm.ptr629 // CHECK-NEXT: %[[FORMATSTART:.*]] = llvm.getelementptr %[[FORMATSTR]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<11 x i8>630 // CHECK-NEXT: %[[EXT:.+]] = llvm.fpext %[[ARG1]] : f32 to f64631 // CHECK-NEXT: %[[O:.*]] = llvm.mlir.constant(1 : index) : i64632 // CHECK-NEXT: %[[ALLOC:.*]] = llvm.alloca %[[O]] x !llvm.struct<(i32, f64)> : (i64) -> !llvm.ptr633 // CHECK-NEXT: %[[EL0:.*]] = llvm.getelementptr %[[ALLOC]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, f64)>634 // CHECK-NEXT: llvm.store %[[ARG0]], %[[EL0]] : i32, !llvm.ptr635 // CHECK-NEXT: %[[EL1:.*]] = llvm.getelementptr %[[ALLOC]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, f64)>636 // CHECK-NEXT: llvm.store %[[EXT]], %[[EL1]] : f64, !llvm.ptr637 // CHECK-NEXT: llvm.call @vprintf(%[[FORMATSTART]], %[[ALLOC]]) : (!llvm.ptr, !llvm.ptr) -> i32638 gpu.printf "Hello: %d\n", %arg0, %arg1 : i32, f32639 gpu.return640 }641}642 643 644 645gpu.module @test_module_30 {646 // CHECK-LABEL: func @subgroup_reduce_add647 gpu.func @subgroup_reduce_add(%arg0 : i32) {648 // CHECK: nvvm.redux.sync add {{.*}}649 %result = gpu.subgroup_reduce add %arg0 uniform {} : (i32) -> (i32)650 gpu.return651 }652 // CHECK-LABEL: @subgroup_reduce_minsi653 gpu.func @subgroup_reduce_minsi(%arg0 : i32) {654 // CHECK: nvvm.redux.sync min {{.*}}655 %result = gpu.subgroup_reduce minsi %arg0 uniform {} : (i32) -> (i32)656 gpu.return657 }658 // CHECK-LABEL: @subgroup_reduce_maxsi659 gpu.func @subgroup_reduce_maxsi(%arg0 : i32) {660 // CHECK: nvvm.redux.sync max {{.*}}661 %result = gpu.subgroup_reduce maxsi %arg0 uniform {} : (i32) -> (i32)662 gpu.return663 }664 // CHECK-LABEL: func @subgroup_reduce_and665 gpu.func @subgroup_reduce_and(%arg0 : i32) {666 // CHECK: nvvm.redux.sync and {{.*}}667 %result = gpu.subgroup_reduce and %arg0 uniform {} : (i32) -> (i32)668 gpu.return669 }670 // CHECK-LABEL: @subgroup_reduce_or671 gpu.func @subgroup_reduce_or(%arg0 : i32) {672 // CHECK: nvvm.redux.sync or {{.*}}673 %result = gpu.subgroup_reduce or %arg0 uniform {} : (i32) -> (i32)674 gpu.return675 }676 // CHECK-LABEL: @subgroup_reduce_xor677 gpu.func @subgroup_reduce_xor(%arg0 : i32) {678 // CHECK: nvvm.redux.sync xor {{.*}}679 %result = gpu.subgroup_reduce xor %arg0 uniform {} : (i32) -> (i32)680 gpu.return681 }682}683 684gpu.module @test_module_31 {685 // CHECK: llvm.func @__nv_fmodf(f32, f32) -> f32686 // CHECK: llvm.func @__nv_fmod(f64, f64) -> f64687 // CHECK-LABEL: func @gpu_fmod688 func.func @gpu_fmod(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {689 %result32 = arith.remf %arg_f32, %arg_f32 : f32690 // CHECK: llvm.call @__nv_fmodf(%{{.*}}, %{{.*}}) : (f32, f32) -> f32691 %result64 = arith.remf %arg_f64, %arg_f64 : f64692 // CHECK: llvm.call @__nv_fmod(%{{.*}}, %{{.*}}) : (f64, f64) -> f64693 func.return %result32, %result64 : f32, f64694 }695}696 697gpu.module @test_module_32 {698 // CHECK: llvm.func @__nv_erff(f32) -> f32699 // CHECK: llvm.func @__nv_erf(f64) -> f64700 // CHECK-LABEL: func @gpu_erf701 func.func @gpu_erf(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {702 %result32 = math.erf %arg_f32 : f32703 // CHECK: llvm.call @__nv_erff(%{{.*}}) : (f32) -> f32704 %result64 = math.erf %arg_f64 : f64705 // CHECK: llvm.call @__nv_erf(%{{.*}}) : (f64) -> f64706 func.return %result32, %result64 : f32, f64707 }708}709 710gpu.module @test_module_33 {711// CHECK-LABEL: func @kernel_with_block_size(712// CHECK: attributes {gpu.kernel, gpu.known_block_size = array<i32: 32, 4, 2>, nvvm.kernel, nvvm.maxntid = array<i32: 32, 4, 2>}713 gpu.func @kernel_with_block_size(%arg0: !llvm.ptr) kernel attributes {known_block_size = array<i32: 32, 4, 2>} {714 // CHECK: = nvvm.read.ptx.sreg.tid.x range <i32, 0, 32> : i32715 %0 = gpu.thread_id x716 // CHECK: = nvvm.read.ptx.sreg.tid.y range <i32, 0, 4> : i32717 %1 = gpu.thread_id y718 // CHECK: = nvvm.read.ptx.sreg.tid.z range <i32, 0, 2> : i32719 %2 = gpu.thread_id z720 721 // Fake usage to prevent dead code elimination722 %3 = arith.addi %0, %1 : index723 %4 = arith.addi %3, %2 : index724 %5 = arith.index_cast %4 : index to i64725 llvm.store %5, %arg0 : i64, !llvm.ptr726 gpu.return727 }728}729 730 731gpu.module @test_module_34 {732 // CHECK-LABEL: llvm.func @memref_signature(733 // CHECK-SAME: %{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: f32) -> !llvm.struct<(struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, f32)>734 // CHECK: llvm.mlir.poison735 // CHECK: llvm.insertvalue736 // CHECK: llvm.insertvalue737 // CHECK: llvm.insertvalue738 // CHECK: llvm.insertvalue739 // CHECK: llvm.insertvalue740 // CHECK: llvm.mlir.poison741 // CHECK: llvm.insertvalue742 // CHECK: llvm.insertvalue743 // CHECK: llvm.return744 745 // CHECK-BARE-LABEL: llvm.func @memref_signature(746 // CHECK-BARE-SAME: %{{.*}}: !llvm.ptr, %{{.*}}: f32) -> !llvm.struct<(ptr, f32)>747 gpu.func @memref_signature(%m: memref<2xf32>, %f: f32) -> (memref<2xf32>, f32) {748 gpu.return %m, %f : memref<2xf32>, f32749 }750}751 752gpu.module @test_module_35 {753 // CHECK: llvm.func @__nv_acosf(f32) -> f32754 // CHECK: llvm.func @__nv_acos(f64) -> f64755 // CHECK-LABEL: func @gpu_acos756 func.func @gpu_acos(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {757 %result32 = math.acos %arg_f32 : f32758 // CHECK: llvm.call @__nv_acosf(%{{.*}}) : (f32) -> f32759 %result64 = math.acos %arg_f64 : f64760 // CHECK: llvm.call @__nv_acos(%{{.*}}) : (f64) -> f64761 func.return %result32, %result64 : f32, f64762 }763}764 765gpu.module @test_module_36 {766 // CHECK: llvm.func @__nv_acoshf(f32) -> f32767 // CHECK: llvm.func @__nv_acosh(f64) -> f64768 // CHECK-LABEL: func @gpu_acosh769 func.func @gpu_acosh(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {770 %result32 = math.acosh %arg_f32 : f32771 // CHECK: llvm.call @__nv_acoshf(%{{.*}}) : (f32) -> f32772 %result64 = math.acosh %arg_f64 : f64773 // CHECK: llvm.call @__nv_acosh(%{{.*}}) : (f64) -> f64774 func.return %result32, %result64 : f32, f64775 }776}777 778gpu.module @test_module_37 {779 // CHECK: llvm.func @__nv_asinf(f32) -> f32780 // CHECK: llvm.func @__nv_asin(f64) -> f64781 // CHECK-LABEL: func @gpu_asin782 func.func @gpu_asin(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {783 %result32 = math.asin %arg_f32 : f32784 // CHECK: llvm.call @__nv_asinf(%{{.*}}) : (f32) -> f32785 %result64 = math.asin %arg_f64 : f64786 // CHECK: llvm.call @__nv_asin(%{{.*}}) : (f64) -> f64787 func.return %result32, %result64 : f32, f64788 }789}790 791gpu.module @test_module_38 {792 // CHECK: llvm.func @__nv_asinhf(f32) -> f32793 // CHECK: llvm.func @__nv_asinh(f64) -> f64794 // CHECK-LABEL: func @gpu_asinh795 func.func @gpu_asinh(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {796 %result32 = math.asinh %arg_f32 : f32797 // CHECK: llvm.call @__nv_asinhf(%{{.*}}) : (f32) -> f32798 %result64 = math.asinh %arg_f64 : f64799 // CHECK: llvm.call @__nv_asinh(%{{.*}}) : (f64) -> f64800 func.return %result32, %result64 : f32, f64801 }802}803 804gpu.module @test_module_39 {805 // CHECK: llvm.func @__nv_atanhf(f32) -> f32806 // CHECK: llvm.func @__nv_atanh(f64) -> f64807 // CHECK-LABEL: func @gpu_atanh808 func.func @gpu_atanh(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64)809 -> (f16, f32, f64) {810 %result16 = math.atanh %arg_f16 : f16811 // CHECK: llvm.fpext %{{.*}} : f16 to f32812 // CHECK-NEXT: llvm.call @__nv_atanhf(%{{.*}}) : (f32) -> f32813 // CHECK-NEXT: llvm.fptrunc %{{.*}} : f32 to f16814 %result32 = math.atanh %arg_f32 : f32815 // CHECK: llvm.call @__nv_atanhf(%{{.*}}) : (f32) -> f32816 %result64 = math.atanh %arg_f64 : f64817 // CHECK: llvm.call @__nv_atanh(%{{.*}}) : (f64) -> f64818 func.return %result16, %result32, %result64 : f16, f32, f64819 }820}821 822gpu.module @test_module_40 {823 // CHECK: llvm.func @__nv_copysignf(f32, f32) -> f32824 // CHECK: llvm.func @__nv_copysign(f64, f64) -> f64825 // CHECK-LABEL: func @gpu_copysign826 func.func @gpu_copysign(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {827 %result32 = math.copysign %arg_f32, %arg_f32 : f32828 // CHECK: llvm.call @__nv_copysignf(%{{.*}}, %{{.*}}) : (f32, f32) -> f32829 %result64 = math.copysign %arg_f64, %arg_f64 : f64830 // CHECK: llvm.call @__nv_copysign(%{{.*}}, %{{.*}}) : (f64, f64) -> f64831 func.return %result32, %result64 : f32, f64832 }833}834 835gpu.module @test_module_41 {836 // CHECK: llvm.func @__nv_coshf(f32) -> f32837 // CHECK: llvm.func @__nv_cosh(f64) -> f64838 // CHECK-LABEL: func @gpu_cosh839 func.func @gpu_cosh(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {840 %result32 = math.cosh %arg_f32 : f32841 // CHECK: llvm.call @__nv_coshf(%{{.*}}) : (f32) -> f32842 %result64 = math.cosh %arg_f64 : f64843 // CHECK: llvm.call @__nv_cosh(%{{.*}}) : (f64) -> f64844 func.return %result32, %result64 : f32, f64845 }846}847 848gpu.module @test_module_42 {849 // CHECK: llvm.func @__nv_fmaf(f32, f32, f32) -> f32850 // CHECK: llvm.func @__nv_fma(f64, f64, f64) -> f64851 // CHECK-LABEL: func @gpu_fma852 func.func @gpu_fma(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {853 %result32 = math.fma %arg_f32, %arg_f32, %arg_f32 : f32854 // CHECK: llvm.call @__nv_fmaf(%{{.*}}, %{{.*}}, %{{.*}}) : (f32, f32, f32) -> f32855 %result64 = math.fma %arg_f64, %arg_f64, %arg_f64 : f64856 // CHECK: llvm.call @__nv_fma(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, f64, f64) -> f64857 func.return %result32, %result64 : f32, f64858 }859}860 861gpu.module @test_module_43 {862 // CHECK: llvm.func @__nv_roundf(f32) -> f32863 // CHECK: llvm.func @__nv_round(f64) -> f64864 // CHECK-LABEL: func @gpu_round865 func.func @gpu_round(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {866 %result32 = math.round %arg_f32 : f32867 // CHECK: llvm.call @__nv_roundf(%{{.*}}) : (f32) -> f32868 %result64 = math.round %arg_f64 : f64869 // CHECK: llvm.call @__nv_round(%{{.*}}) : (f64) -> f64870 func.return %result32, %result64 : f32, f64871 }872}873 874gpu.module @test_module_44 {875 // CHECK: llvm.func @__nv_rintf(f32) -> f32876 // CHECK: llvm.func @__nv_rint(f64) -> f64877 // CHECK-LABEL: func @gpu_roundeven878 func.func @gpu_roundeven(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {879 %result32 = math.roundeven %arg_f32 : f32880 // CHECK: llvm.call @__nv_rintf(%{{.*}}) : (f32) -> f32881 %result64 = math.roundeven %arg_f64 : f64882 // CHECK: llvm.call @__nv_rint(%{{.*}}) : (f64) -> f64883 func.return %result32, %result64 : f32, f64884 }885}886 887gpu.module @test_module_45 {888 // CHECK: llvm.func @__nv_sinhf(f32) -> f32889 // CHECK: llvm.func @__nv_sinh(f64) -> f64890 // CHECK-LABEL: func @gpu_sinh891 func.func @gpu_sinh(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {892 %result32 = math.sinh %arg_f32 : f32893 // CHECK: llvm.call @__nv_sinhf(%{{.*}}) : (f32) -> f32894 %result64 = math.sinh %arg_f64 : f64895 // CHECK: llvm.call @__nv_sinh(%{{.*}}) : (f64) -> f64896 func.return %result32, %result64 : f32, f64897 }898}899 900gpu.module @test_module_46 {901 // CHECK: llvm.func @__nv_coshf(f32) -> f32902 // CHECK: llvm.func @__nv_cosh(f64) -> f64903 // CHECK-LABEL: func @gpu_cosh_with_fastmath904 func.func @gpu_cosh_with_fastmath(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {905 %result32 = math.cosh %arg_f32 fastmath<fast> : f32906 // CHECK: llvm.call @__nv_coshf(%{{.*}}) : (f32) -> f32907 %result64 = math.cosh %arg_f64 fastmath<afn> : f64908 // CHECK: llvm.call @__nv_cosh(%{{.*}}) : (f64) -> f64909 func.return %result32, %result64 : f32, f64910 }911}912 913gpu.module @test_module_47 {914 // CHECK: llvm.func @__nv_sinhf(f32) -> f32915 // CHECK: llvm.func @__nv_sinh(f64) -> f64916 // CHECK-LABEL: func @gpu_sinh_with_fastmath917 func.func @gpu_sinh_with_fastmath(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {918 %result32 = math.sinh %arg_f32 fastmath<contract> : f32919 // CHECK: llvm.call @__nv_sinhf(%{{.*}}) : (f32) -> f32920 %result64 = math.sinh %arg_f64 fastmath<none> : f64921 // CHECK: llvm.call @__nv_sinh(%{{.*}}) : (f64) -> f64922 func.return %result32, %result64 : f32, f64923 }924}925 926gpu.module @test_module_48 {927 // CHECK: llvm.func @__nv_expf(f32) -> f32928 // CHECK: llvm.func @__nv_exp(f64) -> f64929 // CHECK-LABEL: func @gpu_exp_with_fastmath930 func.func @gpu_exp_with_fastmath(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64, f32) {931 %result32 = math.exp %arg_f32 fastmath<reassoc>: f32932 // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32933 %result64 = math.exp %arg_f64 fastmath<contract>: f64934 // CHECK: llvm.call @__nv_exp(%{{.*}}) : (f64) -> f64935 %result32Fast = math.exp %arg_f32 fastmath<ninf> : f32936 // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32937 func.return %result32, %result64, %result32Fast : f32, f64, f32938 }939}940 941gpu.module @test_module_49 {942// CHECK-LABEL: func @explicit_id_bounds()943 func.func @explicit_id_bounds() -> (index, index, index) {944 // CHECK: = nvvm.read.ptx.sreg.tid.x range <i32, 0, 32> : i32945 %0 = gpu.thread_id x upper_bound 32946 // CHECK: = nvvm.read.ptx.sreg.ntid.x range <i32, 1, 33> : i32947 %1 = gpu.block_dim x upper_bound 32948 // CHECK: = nvvm.read.ptx.sreg.laneid range <i32, 0, 16> : i32949 %2 = gpu.lane_id upper_bound 16950 951 return %0, %1, %2 : index, index, index952 }953}954 955gpu.module @test_module_50 {956// CHECK-LABEL: func @kernel_with_grid_size(957 gpu.func @kernel_with_grid_size(%arg0: !llvm.ptr) kernel attributes {known_grid_size = array<i32: 32, 4, 2>} {958 // CHECK: = nvvm.read.ptx.sreg.ctaid.x range <i32, 0, 32> : i32959 %0 = gpu.block_id x960 // CHECK: = nvvm.read.ptx.sreg.ctaid.y range <i32, 0, 4> : i32961 %1 = gpu.block_id y962 // CHECK: = nvvm.read.ptx.sreg.ctaid.z range <i32, 0, 2> : i32963 %2 = gpu.block_id z964 965 // Fake usage to prevent dead code elimination966 %3 = arith.addi %0, %1 : index967 %4 = arith.addi %3, %2 : index968 %5 = arith.index_cast %4 : index to i64969 llvm.store %5, %arg0 : i64, !llvm.ptr970 gpu.return971 }972}973 974// CHECK-LABEL: gpu.module @test_module_51975// CHECK: llvm.mlir.global internal constant @[[func_name:.*]]("(unknown)\00") {addr_space = 0 : i32}976// CHECK: llvm.mlir.global internal constant @[[file_name:.*]]("{{.*}}gpu-to-nvvm.mlir{{.*}}") {addr_space = 0 : i32}977// CHECK: llvm.mlir.global internal constant @[[message:.*]]("assert message\00") {addr_space = 0 : i32}978// CHECK: llvm.func @__assertfail(!llvm.ptr, !llvm.ptr, i32, !llvm.ptr, i64) attributes {passthrough = ["noreturn"]}979// CHECK: llvm.func @test_assert(%[[cond:.*]]: i1) attributes {gpu.kernel, nvvm.kernel} {980// CHECK: llvm.cond_br %[[cond]], ^[[after_block:.*]], ^[[assert_block:.*]]981// CHECK: ^[[assert_block]]:982// CHECK: %[[message_ptr:.*]] = llvm.mlir.addressof @[[message]] : !llvm.ptr983// CHECK: %[[message_start:.*]] = llvm.getelementptr %[[message_ptr]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<15 x i8>984// CHECK: %[[file_ptr:.*]] = llvm.mlir.addressof @[[file_name]] : !llvm.ptr985// CHECK: %[[file_start:.*]] = llvm.getelementptr %[[file_ptr]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<{{.*}} x i8>986// CHECK: %[[func_ptr:.*]] = llvm.mlir.addressof @[[func_name]] : !llvm.ptr987// CHECK: %[[func_start:.*]] = llvm.getelementptr %[[func_ptr]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<{{.*}} x i8>988// CHECK: %[[line_num:.*]] = llvm.mlir.constant({{.*}} : i32) : i32989// CHECK: %[[ptr:.*]] = llvm.mlir.constant(1 : i64) : i64990// CHECK: llvm.call @__assertfail(%[[message_start]], %[[file_start]], %[[line_num]], %[[func_start]], %[[ptr]]) : (!llvm.ptr, !llvm.ptr, i32, !llvm.ptr, i64) -> ()991// CHECK: llvm.br ^[[after_block]]992// CHECK: ^[[after_block]]:993// CHECK: llvm.return994// CHECK: }995 996gpu.module @test_module_51 {997 gpu.func @test_assert(%arg0: i1) kernel {998 cf.assert %arg0, "assert message"999 gpu.return1000 }1001}1002 1003module attributes {transform.with_named_sequence} {1004 transform.named_sequence @__transform_main(%toplevel_module: !transform.any_op {transform.readonly}) {1005 %gpu_module = transform.structured.match ops{["gpu.module"]} in %toplevel_module1006 : (!transform.any_op) -> !transform.any_op1007 1008 transform.apply_patterns to %gpu_module {1009 transform.apply_patterns.gpu.gpu_rewrite_patterns1010 } : !transform.any_op1011 1012 transform.apply_conversion_patterns to %gpu_module {1013 transform.apply_conversion_patterns.dialect_to_llvm "arith"1014 transform.apply_conversion_patterns.dialect_to_llvm "cf"1015 transform.apply_conversion_patterns.vector.vector_to_llvm1016 transform.apply_conversion_patterns.func.func_to_llvm1017 transform.apply_conversion_patterns.dialect_to_llvm "memref"1018 transform.apply_conversion_patterns.gpu.gpu_to_nvvm {benefit = 10 : i16}1019 transform.apply_conversion_patterns.gpu.gpu_wmma_to_nvvm1020 transform.apply_conversion_patterns.gpu.gpu_subgroup_reduce_to_nvvm1021 transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm1022 } with type_converter {1023 transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter1024 {index_bitwidth = 64,1025 use_bare_ptr_call_conv = false}1026 } {1027 legal_dialects = ["llvm", "memref", "nvvm", "test"],1028 legal_ops = ["func.func", "gpu.module", "gpu.yield"],1029 illegal_dialects = ["gpu"],1030 illegal_ops = ["llvm.copysign", "llvm.cos", "llvm.exp", "llvm.exp2", "llvm.fabs", "llvm.fceil",1031 "llvm.ffloor", "llvm.frem", "llvm.log", "llvm.log10", "llvm.log2", "llvm.pow",1032 "llvm.roundeven", "llvm.round", "llvm.sin", "llvm.sqrt"],1033 partial_conversion1034 } : !transform.any_op1035 transform.yield1036 }1037}1038 1039 1040gpu.module @test_module_52 {1041 // CHECK: llvm.func @__nv_abs(i32) -> i321042 // CHECK-LABEL: func @gpu_abs1043 func.func @gpu_abs(%arg_i32 : i32) -> (i32) {1044 %result32 = math.absi %arg_i32 : i321045 // CHECK: llvm.call @__nv_abs(%{{.*}}) : (i32) -> i321046 func.return %result32 : i321047 }1048}1049 1050gpu.module @test_module_53 {1051 // CHECK: llvm.func @__nv_powif(f32, i32) -> f321052 // CHECK: llvm.func @__nv_powi(f64, i32) -> f641053 // CHECK-LABEL: func @gpu_powi1054 func.func @gpu_powi(%arg_f32 : f32, %arg_f64 : f64, %arg_i32 : i32) -> (f32, f64) {1055 %result32 = math.fpowi %arg_f32, %arg_i32 : f32, i321056 // CHECK: llvm.call @__nv_powif(%{{.*}}, %{{.*}}) : (f32, i32) -> f321057 %result64 = math.fpowi %arg_f64, %arg_i32 : f64, i321058 // CHECK: llvm.call @__nv_powi(%{{.*}}, %{{.*}}) : (f64, i32) -> f641059 func.return %result32, %result64 : f32, f641060 }1061}1062 1063gpu.module @test_module_54 {1064 // CHECK: llvm.func @__nv_isinff(f32) -> i321065 // CHECK: llvm.func @__nv_isinfd(f64) -> i321066 // CHECK: llvm.func @__nv_isnanf(f32) -> i321067 // CHECK: llvm.func @__nv_isnand(f64) -> i321068 // CHECK: llvm.func @__nv_finitef(f32) -> i321069 // CHECK: llvm.func @__nv_isfinited(f64) -> i321070 // CHECK-LABEL: @fpclassify1071 func.func @fpclassify(%f32: f32, %f64: f64) -> (i1, i1, i1, i1, i1, i1) {1072 // CHECK: %[[INFF:.+]] = llvm.call @__nv_isinff(%{{.*}}) : (f32) -> i321073 // CHECK: %[[ZERO:.+]] = llvm.mlir.constant(0 : i32) : i321074 // CHECK: %[[R0:.+]] = llvm.icmp "ne" %[[INFF]], %[[ZERO]]1075 %0 = math.isinf %f32 : f321076 // CHECK: llvm.call @__nv_isinfd(%{{.*}}) : (f64) -> i321077 // CHECK: llvm.mlir.constant(01078 // CHECK: llvm.icmp "ne"1079 %1 = math.isinf %f64 : f641080 // CHECK: llvm.call @__nv_isnanf(%{{.*}}) : (f32) -> i321081 // CHECK: llvm.mlir.constant(01082 // CHECK: llvm.icmp "ne"1083 %2 = math.isnan %f32 : f321084 // CHECK: llvm.call @__nv_isnand(%{{.*}}) : (f64) -> i321085 // CHECK: llvm.mlir.constant(01086 // CHECK: llvm.icmp "ne"1087 %3 = math.isnan %f64 : f641088 // CHECK: llvm.call @__nv_finitef(%{{.*}}) : (f32) -> i321089 // CHECK: llvm.mlir.constant(01090 // CHECK: llvm.icmp "ne"1091 %4 = math.isfinite %f32 : f321092 // CHECK: llvm.call @__nv_isfinited(%{{.*}}) : (f64) -> i321093 // CHECK: llvm.mlir.constant(01094 // CHECK: llvm.icmp "ne"1095 %5 = math.isfinite %f64 : f641096 // CHECK: llvm.return %[[R0]]1097 return %0, %1, %2, %3, %4, %5 : i1, i1, i1, i1, i1, i11098 }1099}1100 1101gpu.module @test_module_55 {1102 // CHECK: llvm.func @__nv_erfcf(f32) -> f321103 // CHECK: llvm.func @__nv_erfc(f64) -> f641104 // CHECK-LABEL: func @gpu_erf1105 func.func @gpu_erfc(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {1106 %result32 = math.erfc %arg_f32 : f321107 // CHECK: llvm.call @__nv_erfcf(%{{.*}}) : (f32) -> f321108 %result64 = math.erfc %arg_f64 : f641109 // CHECK: llvm.call @__nv_erfc(%{{.*}}) : (f64) -> f641110 func.return %result32, %result64 : f32, f641111 }1112}1113 1114gpu.module @test_module_56 {1115 // CHECK: gpu.module @test_module_561116 1117 // CHECK-DAG: llvm.func @__nv_sincosf(f32, !llvm.ptr, !llvm.ptr)1118 // CHECK-DAG: llvm.func @__nv_sincos(f64, !llvm.ptr, !llvm.ptr)1119 1120 // CHECK-LABEL: func @gpu_sincos1121 // CHECK-SAME: %[[ARG_f16:.*]]: f16, %[[ARG_f32:.*]]: f32, %[[ARG_f64:.*]]: f641122 func.func @gpu_sincos(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f16, f32, f32, f64, f64) {1123 // CHECK-COUNT-6: llvm.alloca1124 // CHECK: %[[ARG_f16_ext:.*]] = llvm.fpext %[[ARG_f16]] : f16 to f321125 // CHECK: llvm.call @__nv_sincosf(%[[ARG_f16_ext]], %{{.+}}, %{{.+}}) : (f32, !llvm.ptr, !llvm.ptr) -> ()1126 // CHECK-COUNT-2: llvm.fptrunc1127 // CHECK: llvm.call @__nv_sincosf(%[[ARG_f32]], %{{.+}}, %{{.+}}) : (f32, !llvm.ptr, !llvm.ptr) -> ()1128 // CHECK: llvm.call @__nv_sincos(%[[ARG_f64]], %{{.+}}, %{{.+}}) : (f64, !llvm.ptr, !llvm.ptr) -> ()1129 %sin16, %cos16 = math.sincos %arg_f16 : f161130 %sin32, %cos32 = math.sincos %arg_f32 : f321131 %sin64, %cos64 = math.sincos %arg_f64 : f641132 func.return %sin16, %cos16, %sin32, %cos32, %sin64, %cos64 : f16, f16, f32, f32, f64, f641133 }1134 1135 // CHECK: llvm.func @__nv_fast_sincosf(f32, !llvm.ptr, !llvm.ptr)1136 1137 // CHECK-LABEL: func @gpu_sincos_fastmath1138 // CHECK-SAME: %[[ARG_f16:.*]]: f16, %[[ARG_f32:.*]]: f32, %[[ARG_f64:.*]]: f641139 func.func @gpu_sincos_fastmath(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f16, f32, f32, f64, f64) {1140 // CHECK-COUNT-6: llvm.alloca1141 // CHECK: %[[ARG_f16_ext:.*]] = llvm.fpext %[[ARG_f16]] : f16 to f321142 // CHECK: llvm.call @__nv_fast_sincosf(%[[ARG_f16_ext]], %{{.+}}, %{{.+}}) : (f32, !llvm.ptr, !llvm.ptr) -> ()1143 // CHECK-COUNT-2: llvm.fptrunc1144 // CHECK: llvm.call @__nv_fast_sincosf(%[[ARG_f32]], %{{.+}}, %{{.+}}) : (f32, !llvm.ptr, !llvm.ptr) -> ()1145 // CHECK: llvm.call @__nv_sincos(%[[ARG_f64]], %{{.+}}, %{{.+}}) : (f64, !llvm.ptr, !llvm.ptr) -> ()1146 %sin16, %cos16 = math.sincos %arg_f16 fastmath<afn> : f161147 %sin32, %cos32 = math.sincos %arg_f32 fastmath<afn> : f321148 %sin64, %cos64 = math.sincos %arg_f64 fastmath<afn> : f641149 func.return %sin16, %cos16, %sin32, %cos32, %sin64, %cos64 : f16, f16, f32, f32, f64, f641150 }1151}1152