224 lines · python
1# RUN: %PYTHON %s | FileCheck %s2# This is just a smoke test that the dialect is functional.3 4from mlir.ir import *5from mlir.dialects import nvvm6from mlir.dialects import llvm7from mlir.dialects import func8import mlir.extras.types as T9from mlir.dialects import arith10 11 12def constructAndPrintInModule(f):13 print("\nTEST:", f.__name__)14 with Context(), Location.unknown():15 module = Module.create()16 with InsertionPoint(module.body):17 f()18 19 print(module)20 module.operation.verify()21 return f22 23 24# CHECK-LABEL: testSmoke25@constructAndPrintInModule26def testSmoke():27 i64 = IntegerType.get_signless(64)28 mat64f32_t = Type.parse(29 "!llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>"30 )31 shape_attr = Attribute.parse("#nvvm.shape<m = 64, n = 32, k = 16>")32 33 # CHECK-LABEL: func @wgmma_f32_f16_f16(%arg0: i64, %arg1: i64)34 @func.FuncOp.from_py_func(i64, i64)35 def wgmma_f32_f16_f16(desc_a, desc_b):36 # CHECK: nvvm.cp.async.wait.group 537 nvvm.CpAsyncWaitGroupOp(5)38 # CHECK: %0 = llvm.mlir.undef : [[MAT_T:.*]]39 result = llvm.UndefOp(mat64f32_t)40 # CHECK: %1 = nvvm.wgmma.mma_async %arg0, %arg1, %0, <m = 64, n = 32, k = 16>, D[<f32>, <zero>], A[<f16>, <neg>, <col>], B[<f16>, <neg>, <col>] : [[MAT_T]] -> [[MAT_T]]41 result1 = nvvm.WgmmaMmaAsyncOp(42 results_=mat64f32_t,43 inouts=result,44 descriptorA=desc_a,45 descriptorB=desc_b,46 shape=shape_attr,47 typeA=nvvm.WGMMATypes.f16,48 typeB=nvvm.WGMMATypes.f16,49 typeD=nvvm.WGMMATypes.f32,50 scaleD=nvvm.WGMMAScaleOut.zero,51 scaleA=nvvm.WGMMAScaleIn.neg,52 scaleB=nvvm.WGMMAScaleIn.neg,53 layoutA=nvvm.MMALayout.col,54 layoutB=nvvm.MMALayout.col,55 )56 57 58# CHECK-LABEL: TEST: test_inline_ptx59# CHECK-LABEL: func.func @my_inline_ptx(60# CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: f32, %[[arg1:[a-zA-Z0-9_]+]]: f32, %[[arg2:[a-zA-Z0-9_]+]]: i32, %[[arg3:[a-zA-Z0-9_]+]]: i32)61# CHECK: %[[S0:.+]]:2 = nvvm.inline_ptx62# CHECK-SAME: ro(%[[arg0]], %[[arg1]] : f32, f32) rw(%[[arg2]], %[[arg3]] : i32, i32) -> f32, f3263# CHECK: %[[S1:.+]] = arith.addf %[[arg0]], %[[arg1]] : f3264# CHECK: %[[S2:.+]] = arith.addi %[[arg2]], %[[arg3]] : i3265# CHECK: %[[S3:.+]] = arith.addf %[[S0]]#0, %[[S0]]#1 : f3266 67 68@constructAndPrintInModule69def test_inline_ptx():70 i32 = T.i32()71 f32 = T.f32()72 73 @func.FuncOp.from_py_func(f32, f32, i32, i32)74 def my_inline_ptx(a, b, c, d):75 ptx = r"""76 {77 .reg .pred p;78 setp.ge.s32 p, {$r0}, {$r1};79 selp.s32 {$r0}, {$r0}, {$r1}, p;80 selp.s32 {$r1}, {$r0}, {$r1}, p;81 selp.s32 {$rw0}, {$r0}, {$r1}, p;82 selp.s32 {$rw1}, {$r0}, {$r1}, p;83 }84 """85 wo0, wo1 = nvvm.inline_ptx(86 read_only_args=[a, b],87 read_write_args=[c, d],88 write_only_args=[f32, f32],89 ptx_code=ptx,90 )91 arith.addf(a, b)92 arith.addi(c, d)93 arith.addf(wo0, wo1)94 95 96@constructAndPrintInModule97def test_barriers():98 i32 = T.i32()99 f32 = T.f32()100 101 @func.FuncOp.from_py_func(i32, i32, f32)102 def barriers(mask, vi32, vf32):103 c0 = arith.constant(T.i32(), 0)104 cffff = arith.constant(T.i32(), 0xFFFF)105 res = nvvm.barrier(106 res=i32,107 barrier_id=c0,108 number_of_threads=cffff,109 )110 111 for reduction in (112 nvvm.BarrierReduction.AND,113 nvvm.BarrierReduction.OR,114 nvvm.BarrierReduction.POPC,115 ):116 res = nvvm.barrier(117 res=i32,118 reduction_op=reduction,119 reduction_predicate=res,120 )121 122 nvvm.barrier0()123 nvvm.bar_warp_sync(mask)124 nvvm.cluster_arrive()125 nvvm.cluster_arrive(aligned=True)126 nvvm.cluster_arrive_relaxed()127 nvvm.cluster_arrive_relaxed(aligned=True)128 nvvm.cluster_wait()129 nvvm.cluster_wait(aligned=True)130 nvvm.fence_mbarrier_init()131 nvvm.bar_warp_sync(mask)132 return res133 134 135# CHECK-LABEL: func.func @barriers(136# CHECK: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: f32) -> i32 {137# CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : i32138# CHECK: %[[CONSTANT_1:.*]] = arith.constant 65535 : i32139# CHECK: %[[BARRIER_0:.*]] = nvvm.barrier id = %[[CONSTANT_0]] number_of_threads = %[[CONSTANT_1]] -> i32140# CHECK: %[[BARRIER_1:.*]] = nvvm.barrier #nvvm.reduction<and> %[[BARRIER_0]] -> i32141# CHECK: %[[BARRIER_2:.*]] = nvvm.barrier #nvvm.reduction<or> %[[BARRIER_1]] -> i32142# CHECK: %[[BARRIER_3:.*]] = nvvm.barrier #nvvm.reduction<popc> %[[BARRIER_2]] -> i32143# CHECK: nvvm.barrier0144# CHECK: nvvm.bar.warp.sync %[[ARG0]] : i32145# CHECK: nvvm.cluster.arrive146# CHECK: nvvm.cluster.arrive {aligned}147# CHECK: nvvm.cluster.arrive.relaxed148# CHECK: nvvm.cluster.arrive.relaxed {aligned}149# CHECK: nvvm.cluster.wait150# CHECK: nvvm.cluster.wait {aligned}151# CHECK: nvvm.fence.mbarrier.init152# CHECK: nvvm.bar.warp.sync %[[ARG0]] : i32153# CHECK: return %[[BARRIER_3]] : i32154# CHECK: }155 156 157@constructAndPrintInModule158def test_reductions():159 i32 = T.i32()160 f32 = T.f32()161 162 @func.FuncOp.from_py_func(i32, i32, f32)163 def reductions(mask, vi32, vf32):164 for abs in (True, False):165 for nan in (True, False):166 for kind in (167 nvvm.ReduxKind.AND,168 nvvm.ReduxKind.MAX,169 nvvm.ReduxKind.MIN,170 nvvm.ReduxKind.OR,171 nvvm.ReduxKind.UMAX,172 nvvm.ReduxKind.UMIN,173 nvvm.ReduxKind.XOR,174 ):175 nvvm.redux_sync(i32, vi32, kind, vi32)176 177 for kind in (178 nvvm.ReduxKind.FMIN,179 nvvm.ReduxKind.FMAX,180 ):181 nvvm.redux_sync(f32, vf32, kind, vi32, abs=abs, nan=nan)182 183 184# CHECK-LABEL: func.func @reductions(185# CHECK: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: f32) {186# CHECK: %[[REDUX_0:.*]] = nvvm.redux.sync and %[[ARG1]], %[[ARG1]] : i32 -> i32187# CHECK: %[[REDUX_1:.*]] = nvvm.redux.sync max %[[ARG1]], %[[ARG1]] : i32 -> i32188# CHECK: %[[REDUX_2:.*]] = nvvm.redux.sync min %[[ARG1]], %[[ARG1]] : i32 -> i32189# CHECK: %[[REDUX_3:.*]] = nvvm.redux.sync or %[[ARG1]], %[[ARG1]] : i32 -> i32190# CHECK: %[[REDUX_4:.*]] = nvvm.redux.sync umax %[[ARG1]], %[[ARG1]] : i32 -> i32191# CHECK: %[[REDUX_5:.*]] = nvvm.redux.sync umin %[[ARG1]], %[[ARG1]] : i32 -> i32192# CHECK: %[[REDUX_6:.*]] = nvvm.redux.sync xor %[[ARG1]], %[[ARG1]] : i32 -> i32193# CHECK: %[[REDUX_7:.*]] = nvvm.redux.sync fmin %[[ARG2]], %[[ARG1]] {abs = true, nan = true} : f32 -> f32194# CHECK: %[[REDUX_8:.*]] = nvvm.redux.sync fmax %[[ARG2]], %[[ARG1]] {abs = true, nan = true} : f32 -> f32195# CHECK: %[[REDUX_9:.*]] = nvvm.redux.sync and %[[ARG1]], %[[ARG1]] : i32 -> i32196# CHECK: %[[REDUX_10:.*]] = nvvm.redux.sync max %[[ARG1]], %[[ARG1]] : i32 -> i32197# CHECK: %[[REDUX_11:.*]] = nvvm.redux.sync min %[[ARG1]], %[[ARG1]] : i32 -> i32198# CHECK: %[[REDUX_12:.*]] = nvvm.redux.sync or %[[ARG1]], %[[ARG1]] : i32 -> i32199# CHECK: %[[REDUX_13:.*]] = nvvm.redux.sync umax %[[ARG1]], %[[ARG1]] : i32 -> i32200# CHECK: %[[REDUX_14:.*]] = nvvm.redux.sync umin %[[ARG1]], %[[ARG1]] : i32 -> i32201# CHECK: %[[REDUX_15:.*]] = nvvm.redux.sync xor %[[ARG1]], %[[ARG1]] : i32 -> i32202# CHECK: %[[REDUX_16:.*]] = nvvm.redux.sync fmin %[[ARG2]], %[[ARG1]] {abs = true} : f32 -> f32203# CHECK: %[[REDUX_17:.*]] = nvvm.redux.sync fmax %[[ARG2]], %[[ARG1]] {abs = true} : f32 -> f32204# CHECK: %[[REDUX_18:.*]] = nvvm.redux.sync and %[[ARG1]], %[[ARG1]] : i32 -> i32205# CHECK: %[[REDUX_19:.*]] = nvvm.redux.sync max %[[ARG1]], %[[ARG1]] : i32 -> i32206# CHECK: %[[REDUX_20:.*]] = nvvm.redux.sync min %[[ARG1]], %[[ARG1]] : i32 -> i32207# CHECK: %[[REDUX_21:.*]] = nvvm.redux.sync or %[[ARG1]], %[[ARG1]] : i32 -> i32208# CHECK: %[[REDUX_22:.*]] = nvvm.redux.sync umax %[[ARG1]], %[[ARG1]] : i32 -> i32209# CHECK: %[[REDUX_23:.*]] = nvvm.redux.sync umin %[[ARG1]], %[[ARG1]] : i32 -> i32210# CHECK: %[[REDUX_24:.*]] = nvvm.redux.sync xor %[[ARG1]], %[[ARG1]] : i32 -> i32211# CHECK: %[[REDUX_25:.*]] = nvvm.redux.sync fmin %[[ARG2]], %[[ARG1]] {nan = true} : f32 -> f32212# CHECK: %[[REDUX_26:.*]] = nvvm.redux.sync fmax %[[ARG2]], %[[ARG1]] {nan = true} : f32 -> f32213# CHECK: %[[REDUX_27:.*]] = nvvm.redux.sync and %[[ARG1]], %[[ARG1]] : i32 -> i32214# CHECK: %[[REDUX_28:.*]] = nvvm.redux.sync max %[[ARG1]], %[[ARG1]] : i32 -> i32215# CHECK: %[[REDUX_29:.*]] = nvvm.redux.sync min %[[ARG1]], %[[ARG1]] : i32 -> i32216# CHECK: %[[REDUX_30:.*]] = nvvm.redux.sync or %[[ARG1]], %[[ARG1]] : i32 -> i32217# CHECK: %[[REDUX_31:.*]] = nvvm.redux.sync umax %[[ARG1]], %[[ARG1]] : i32 -> i32218# CHECK: %[[REDUX_32:.*]] = nvvm.redux.sync umin %[[ARG1]], %[[ARG1]] : i32 -> i32219# CHECK: %[[REDUX_33:.*]] = nvvm.redux.sync xor %[[ARG1]], %[[ARG1]] : i32 -> i32220# CHECK: %[[REDUX_34:.*]] = nvvm.redux.sync fmin %[[ARG2]], %[[ARG1]] : f32 -> f32221# CHECK: %[[REDUX_35:.*]] = nvvm.redux.sync fmax %[[ARG2]], %[[ARG1]] : f32 -> f32222# CHECK: return223# CHECK: }224