https://github.com/schwarzschild-radius updated https://github.com/llvm/llvm-project/pull/217252
>From 282483b53d25345c7cf1a57f9068485f905357db Mon Sep 17 00:00:00 2001 From: Pradeep Kumar <[email protected]> Date: Wed, 19 Aug 2026 06:55:06 +0000 Subject: [PATCH] [LLVM][NVPTX][MLIR] Add mbarrier layout support This commit adds LLVM NVPTX and MLIR NVVM support for the mbarrier layout extensions: - llvm.nvvm.mbarrier.init.layout / nvvm.mbarrier.init's optional `layout` attribute, lowering to mbarrier.init.layout::v{0,1}. - llvm.nvvm.mbarrier.check_layout / nvvm.mbarrier.check_layout, lowering to mbarrier.check_layout.layout::v{0,1}. Both require PTX ISA 9.3 and sm_90. Co-Authored-By: Claude Opus 5 <[email protected]> --- clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp | 10 ++++ llvm/docs/NVPTXUsage.md | 32 ++++++++++++- llvm/include/llvm/IR/IntrinsicsNVVM.td | 17 +++++-- llvm/lib/IR/AutoUpgrade.cpp | 28 +++++++++++ llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 19 +++++++- .../Assembler/auto_upgrade_nvvm_intrinsics.ll | 12 +++++ llvm/test/CodeGen/NVPTX/mbarrier.ll | 8 ++-- .../NVPTX/mbarrier_layout_sm90_ptx93.ll | 41 +++++++++++++++++ mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 42 ++++++++++++++++- .../Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp | 2 +- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 46 ++++++++++++++----- .../Conversion/NVVMToLLVM/nvvm-to-llvm.mlir | 4 +- .../Target/LLVMIR/nvvm/mbar_check_layout.mlir | 14 ++++++ mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir | 24 +++++++++- .../test/Target/LLVMIR/nvvm/mbar_invalid.mlir | 32 +++++++++++++ 15 files changed, 303 insertions(+), 28 deletions(-) create mode 100644 llvm/test/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll create mode 100644 mlir/test/Target/LLVMIR/nvvm/mbar_check_layout.mlir diff --git a/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp b/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp index 64fdae9d8934d..614b0e72df13c 100644 --- a/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp +++ b/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp @@ -1282,6 +1282,16 @@ Value *CodeGenFunction::EmitNVPTXBuiltinExpr(unsigned BuiltinID, Intrinsic::nvvm_barrier_cta_red_popc_aligned_all, {}, {Builder.getInt32(0), Builder.CreateICmpNE(EmitScalarExpr(E->getArg(0)), Builder.getInt32(0))}); + case NVPTX::BI__nvvm_mbarrier_init: + case NVPTX::BI__nvvm_mbarrier_init_shared: { + Intrinsic::ID IID = BuiltinID == NVPTX::BI__nvvm_mbarrier_init + ? Intrinsic::nvvm_mbarrier_init + : Intrinsic::nvvm_mbarrier_init_shared; + return Builder.CreateIntrinsic(IID, {}, + {EmitScalarExpr(E->getArg(0)), + EmitScalarExpr(E->getArg(1)), + Builder.getInt32(0)}); + } default: return nullptr; } diff --git a/llvm/docs/NVPTXUsage.md b/llvm/docs/NVPTXUsage.md index 5f6a00945a666..42316cf4c7746 100644 --- a/llvm/docs/NVPTXUsage.md +++ b/llvm/docs/NVPTXUsage.md @@ -440,8 +440,8 @@ For more information, refer [PTX ISA](https://docs.nvidia.com/cuda/parallel-thre ##### Syntax: ```llvm -declare void @llvm.nvvm.mbarrier.init(ptr %addr, i32 %count) -declare void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %addr, i32 %count) +declare void @llvm.nvvm.mbarrier.init(ptr %addr, i32 %count, i32 immarg range(i32 0, 2) %layout) +declare void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %addr, i32 %count, i32 immarg range(i32 0, 2) %layout) ``` ##### Overview: @@ -453,6 +453,8 @@ the range [1...2^20-1]. During initialization: - The tx-count and the current phase of the mbarrier object are set to 0. - The expected and pending arrival counts are set to `count`. +- `%layout` is an immediate argument that accepts only +`0` (`layout::v0`) or `1` (`layout::v1`). ##### Semantics: @@ -463,6 +465,32 @@ Performing `mbarrier.init` on a valid mbarrier object is undefined; use `mbarrier.inval` before reusing the memory for another mbarrier or any other purpose. +An mbarrier object initialized with a particular layout must only be +used with operations that support that layout; the layout of an existing +mbarrier object can be queried with `llvm.nvvm.mbarrier.check_layout.*`. + +#### '`llvm.nvvm.mbarrier.check_layout`' + +##### Syntax: + +```llvm +declare i1 @llvm.nvvm.mbarrier.check_layout(ptr addrspace(3) %addr, i32 immarg %layout) +``` + +##### Overview: + +The '`@llvm.nvvm.mbarrier.check_layout.*`' intrinsics test whether the +mbarrier object at `addr` was initialized with the layout named by +`%layout`. They return `true` when the layout matches and `false` +otherwise. `%layout` is an immediate argument that accepts only +`0` (`layout::v0`) or `1` (`layout::v1`). + +##### Semantics: + +If the `addr` does not fall within the shared::cta space, then the behavior of +this intrinsic is undefined. It is expected that `addr` was previously +initialized using `mbarrier.init`; otherwise, the behavior is undefined. + #### '`llvm.nvvm.mbarrier.inval`' ##### Syntax: diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td index 58204732a6b10..3b41b98bcfc54 100644 --- a/llvm/include/llvm/IR/IntrinsicsNVVM.td +++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td @@ -2269,9 +2269,12 @@ foreach is_shared = [true, false] in { defvar mbarrier_ptr_ty = !if(is_shared, llvm_shared_ptr_ty, llvm_ptr_ty); defvar shared = !if(is_shared, "_shared", ""); - def int_nvvm_mbarrier_init # shared : NVVMBuiltin, - Intrinsic<[], [mbarrier_ptr_ty, llvm_i32_ty], - [IntrConvergent, IntrNoCallback]>; + def int_nvvm_mbarrier_init # shared : + DefaultAttrsIntrinsic<[], + [mbarrier_ptr_ty, // mbar + llvm_i32_ty, // count + llvm_i32_ty], // layout + [IntrConvergent, ImmArg<ArgIndex<2>>, Range<ArgIndex<2>, 0, 2>]>; def int_nvvm_mbarrier_inval # shared : NVVMBuiltin, Intrinsic<[], [mbarrier_ptr_ty], @@ -2294,6 +2297,14 @@ foreach is_shared = [true, false] in { def int_nvvm_mbarrier_pending_count : NVVMBuiltin, NVVMPureIntrinsic<[llvm_i32_ty], [llvm_i64_ty]>; +def int_nvvm_mbarrier_check_layout : + DefaultAttrsIntrinsic<[llvm_i1_ty], + [llvm_anyptr_ty, // mbar + llvm_i32_ty], // layout + [IntrReadMem, IntrNoCallback, ImmArg<ArgIndex<1>>, + Range<ArgIndex<1>, 0, 2>], + "llvm.nvvm.mbarrier.check_layout">; + // mbarrier.{expect_tx/complete_tx} foreach op = ["expect_tx", "complete_tx"] in { foreach scope = ["scope_cta", "scope_cluster"] in { diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp index cd505273b5bfb..589e28b1e1f7a 100644 --- a/llvm/lib/IR/AutoUpgrade.cpp +++ b/llvm/lib/IR/AutoUpgrade.cpp @@ -1400,6 +1400,19 @@ static Intrinsic::ID shouldUpgradeNVPTXTcgen05MMAIntrinsic(Function *F, return F->getIntrinsicID(); } +// The mbarrier.init intrinsics gained a trailing layout operand. Calls to the +// older two-argument form request the default layout. +static Intrinsic::ID shouldUpgradeNVPTXMBarrierInitIntrinsic(Function *F, + StringRef Name) { + if (F->arg_size() != 2) + return Intrinsic::not_intrinsic; + + return StringSwitch<Intrinsic::ID>(Name) + .Case("mbarrier.init", Intrinsic::nvvm_mbarrier_init) + .Case("mbarrier.init.shared", Intrinsic::nvvm_mbarrier_init_shared) + .Default(Intrinsic::not_intrinsic); +} + static bool consumeNVVMPtrAddrSpace(StringRef &Name) { return Name.consume_front("local") || Name.consume_front("shared") || Name.consume_front("global") || Name.consume_front("constant") || @@ -1971,6 +1984,14 @@ static bool upgradeIntrinsicFunction1(Function *F, Function *&NewFn, return NewFn != F; } + // Upgrade mbarrier.init intrinsics missing the layout operand. + IID = shouldUpgradeNVPTXMBarrierInitIntrinsic(F, Name); + if (IID != Intrinsic::not_intrinsic) { + rename(F); + NewFn = Intrinsic::getOrInsertDeclaration(F->getParent(), IID); + return true; + } + // The following nvvm intrinsics correspond exactly to an LLVM idiom, but // not to an intrinsic alone. We expand them in UpgradeIntrinsicCall. // @@ -6044,6 +6065,13 @@ void llvm::UpgradeIntrinsicCall(CallBase *CI, Function *NewFn) { Builder.CreateCall(NewFn, {CI->getArgOperand(0), CI->getArgOperand(1), Builder.getFalse()}); break; + case Intrinsic::nvvm_mbarrier_init: + case Intrinsic::nvvm_mbarrier_init_shared: { + SmallVector<Value *, 3> Args(CI->args()); + Args.push_back(Builder.getInt32(0)); // layout = default(0) + NewCall = Builder.CreateCall(NewFn, Args); + break; + } case Intrinsic::riscv_sha256sig0: case Intrinsic::riscv_sha256sig1: case Intrinsic::riscv_sha256sum0: diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td index 907c6cbb178fd..b3c5a66e4fe3c 100644 --- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td +++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td @@ -1402,7 +1402,7 @@ let Predicates = [SM80] in { class MBARRIER_INIT<string AddrSpace, Intrinsic Intrin> : BasicNVPTXInst<(outs), (ins ADDR:$addr, B32:$count), "mbarrier.init" # AddrSpace # ".b64", - [(Intrin addr:$addr, i32:$count)]>; + [(Intrin addr:$addr, i32:$count, (i32 0))]>; def MBARRIER_INIT : MBARRIER_INIT<"", int_nvvm_mbarrier_init>; def MBARRIER_INIT_SHARED : MBARRIER_INIT<".shared", @@ -1475,6 +1475,23 @@ let Predicates = [SM80] in { [(set i32:$res, (int_nvvm_mbarrier_pending_count i64:$state))]>; } +let Predicates = [PTX93, SM90] in { + class MBARRIER_INIT_LAYOUT<string AddrSpace, Intrinsic Intrin> : + BasicNVPTXInst<(outs), (ins ADDR:$addr, B32:$count), + "mbarrier.init.layout::v1" # AddrSpace # ".b64", + [(Intrin addr:$addr, i32:$count, (i32 1))]>; + + def MBARRIER_INIT_LAYOUT : MBARRIER_INIT_LAYOUT<"", int_nvvm_mbarrier_init>; + def MBARRIER_INIT_LAYOUT_SHARED : + MBARRIER_INIT_LAYOUT<".shared", int_nvvm_mbarrier_init_shared>; + + def MBARRIER_CHECK_LAYOUT_SHARED : + NVPTXInst<(outs B1:$res), (ins ADDR:$addr, B32:$layout), + "mbarrier.check_layout.layout::v${layout}.shared.b64 $res, [$addr];", + [(set i1:$res, + (int_nvvm_mbarrier_check_layout addr:$addr, i32:$layout))]>; +} + class MBAR_UTIL<string op, string scope, string space = "", string sem = "", bit tl = 0, bit parity = 0> { diff --git a/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll b/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll index bcbeef260f2bb..d3e30c5ea838b 100644 --- a/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll +++ b/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll @@ -711,3 +711,15 @@ define void @nvvm_ex2_approx(float %a, double %b, half %c, <2 x half> %d) { %r4 = call float @llvm.nvvm.ex2.approx.ftz.f(float %a) ret void } + +declare void @llvm.nvvm.mbarrier.init(ptr, i32) +declare void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3), i32) + +; CHECK-LABEL: @nvvm_mbarrier_init_default_layout +define void @nvvm_mbarrier_init_default_layout(ptr %gen, ptr addrspace(3) %shared, i32 %count) { +; CHECK: call void @llvm.nvvm.mbarrier.init(ptr %gen, i32 %count, i32 0) +; CHECK: call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %shared, i32 %count, i32 0) + call void @llvm.nvvm.mbarrier.init(ptr %gen, i32 %count) + call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %shared, i32 %count) + ret void +} diff --git a/llvm/test/CodeGen/NVPTX/mbarrier.ll b/llvm/test/CodeGen/NVPTX/mbarrier.ll index 78edc0aa2db56..e8eaa2809d45d 100644 --- a/llvm/test/CodeGen/NVPTX/mbarrier.ll +++ b/llvm/test/CodeGen/NVPTX/mbarrier.ll @@ -3,14 +3,14 @@ ; RUN: %if ptxas-sm_80 && ptxas-ptr32 %{ llc < %s -mtriple=nvptx -mcpu=sm_80 | %ptxas-verify -arch=sm_80 %} ; RUN: %if ptxas-sm_80 %{ llc < %s -mtriple=nvptx64 -mcpu=sm_80 | %ptxas-verify -arch=sm_80 %} -declare void @llvm.nvvm.mbarrier.init(ptr %a, i32 %b) -declare void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %a, i32 %b) +declare void @llvm.nvvm.mbarrier.init(ptr %a, i32 %b, i32 %c) +declare void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %a, i32 %b, i32 %c) ; CHECK-LABEL: barrierinit define void @barrierinit(ptr %a, i32 %b) { ; CHECK_PTX32: mbarrier.init.b64 [%r{{[0-9]+}}], %r{{[0-9]+}}; ; CHECK_PTX64: mbarrier.init.b64 [%rd{{[0-9]+}}], %r{{[0-9]+}}; - tail call void @llvm.nvvm.mbarrier.init(ptr %a, i32 %b) + tail call void @llvm.nvvm.mbarrier.init(ptr %a, i32 %b, i32 0) ret void } @@ -18,7 +18,7 @@ define void @barrierinit(ptr %a, i32 %b) { define void @barrierinitshared(ptr addrspace(3) %a, i32 %b) { ; CHECK_PTX32: mbarrier.init.shared.b64 [%r{{[0-9]+}}], %r{{[0-9]+}}; ; CHECK_PTX64: mbarrier.init.shared.b64 [%rd{{[0-9]+}}], %r{{[0-9]+}}; - tail call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %a, i32 %b) + tail call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %a, i32 %b, i32 0) ret void } diff --git a/llvm/test/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll b/llvm/test/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll new file mode 100644 index 0000000000000..521c1fa4a002a --- /dev/null +++ b/llvm/test/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll @@ -0,0 +1,41 @@ +; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6 +; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_90 -mattr=+ptx93 | FileCheck %s +; RUN: %if ptxas-sm_90 && ptxas-isa-9.3 %{ llc < %s -mtriple=nvptx64 -mcpu=sm_90 -mattr=+ptx93| %ptxas-verify -arch=sm_90 %} + +define void @mbarrier_init(ptr addrspace(3) %a, i32 %b) { +; CHECK-LABEL: mbarrier_init( +; CHECK: { +; CHECK-NEXT: .reg .b32 %r<2>; +; CHECK-NEXT: .reg .b64 %rd<2>; +; CHECK-EMPTY: +; CHECK-NEXT: // %bb.0: +; CHECK-NEXT: ld.param::func.b64 %rd1, [mbarrier_init_param_0]; +; CHECK-NEXT: ld.param::func.b32 %r1, [mbarrier_init_param_1]; +; CHECK-NEXT: mbarrier.init.shared.b64 [%rd1], %r1; +; CHECK-NEXT: mbarrier.init.layout::v1.shared.b64 [%rd1], %r1; +; CHECK-NEXT: ret; + tail call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %a, i32 %b, i32 0) + tail call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %a, i32 %b, i32 1) + ret void +} + +define i1 @mbarrier_check_layout(ptr addrspace(3) %a) { +; CHECK-LABEL: mbarrier_check_layout( +; CHECK: { +; CHECK-NEXT: .reg .pred %p<4>; +; CHECK-NEXT: .reg .b32 %r<2>; +; CHECK-NEXT: .reg .b64 %rd<2>; +; CHECK-EMPTY: +; CHECK-NEXT: // %bb.0: +; CHECK-NEXT: ld.param::func.b64 %rd1, [mbarrier_check_layout_param_0]; +; CHECK-NEXT: mbarrier.check_layout.layout::v0.shared.b64 %p1, [%rd1]; +; CHECK-NEXT: mbarrier.check_layout.layout::v1.shared.b64 %p2, [%rd1]; +; CHECK-NEXT: or.pred %p3, %p1, %p2; +; CHECK-NEXT: selp.b32 %r1, -1, 0, %p3; +; CHECK-NEXT: st.param::func.b32 [func_retval0], %r1; +; CHECK-NEXT: ret; + %is_layout_v0 = tail call i1 @llvm.nvvm.mbarrier.check_layout(ptr addrspace(3) %a, i32 0) + %is_layout_v1 = tail call i1 @llvm.nvvm.mbarrier.check_layout(ptr addrspace(3) %a, i32 1) + %ret = or i1 %is_layout_v0, %is_layout_v1 + ret i1 %ret +} diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index 634918b1b21f2..f961ea2e4668f 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -641,7 +641,10 @@ def NVVM_PMEventOp : NVVM_VoidIntrinsicOp<"pmevent">, /// mbarrier.init instruction with generic pointer type def NVVM_MBarrierInitOp : NVVM_PTXBuilder_Op<"mbarrier.init">, Arguments<(ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr, - I32:$count, PtxPredicate:$predicate)> { + I32:$count, + DefaultValuedAttr<ConfinedAttr<I32Attr, + [IntMinValue<0>, IntMaxValue<1>]>, "0">:$layout, + PtxPredicate:$predicate)> { let summary = "MBarrier Initialization Op"; let description = [{ The `nvvm.mbarrier.init` operation initializes an *mbarrier object* at the specified @@ -660,15 +663,22 @@ def NVVM_MBarrierInitOp : NVVM_PTXBuilder_Op<"mbarrier.init">, the behavior is undefined. - `count`: Integer specifying the number of threads that will participate in barrier synchronization. Must be in the range [1, 2²⁰ - 1]. + - `layout`: Optional in-memory layout to initialize the *mbarrier object* + with. Only `0` (`layout::v0`) and `1` (`layout::v1`) are valid values. + When it is omitted, it defaults to `0`. The layout of an existing + *mbarrier object* can be queried with `nvvm.mbarrier.check_layout`. - `predicate`: Optional predicate for conditional execution. [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init) }]; - let assemblyFormat = "$addr `,` $count (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)"; + let assemblyFormat = "$addr `,` $count (`layout` `=` $layout^)? (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)"; let extraClassDeclaration = [{ bool hasIntrinsic() { if(getPredicate()) return false; return true; } + bool getAsmValues(RewriterBase &rewriter, + llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &asmValues); + static mlir::NVVM::IDArgPair getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase& builder); @@ -704,6 +714,34 @@ def NVVM_MBarrierInvalOp : NVVM_VoidIntrinsicOp<"mbarrier.inval">, let assemblyFormat = "$addr attr-dict `:` type(operands)"; } +def NVVM_MBarrierCheckLayoutOp : + NVVM_SingleResultIntrinsicOp<"mbarrier.check_layout"> { + let summary = "MBarrier Check-Layout Operation"; + let description = [{ + The `nvvm.mbarrier.check_layout` operation tests whether the *mbarrier + object* at `addr` was initialized with the layout named by `layout`. + + - `res`: An `i1` that is `true` when the *mbarrier object* has the queried + layout and `false` otherwise. + + The operation takes the following operand and attribute: + - `addr`: A pointer to the memory location of the *mbarrier object*. The + `addr` must be a pointer to shared::cta memory. + - `layout`: The mbarrier layout version to test for. Only `0` and `1` are + valid values. + + [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-check-layout) + }]; + + let results = (outs I1:$res); + let arguments = (ins + LLVM_PointerShared:$addr, + ConfinedAttr<I32Attr, [IntMinValue<0>, IntMaxValue<1>]>:$layout); + + let assemblyFormat = + "$addr `,` $layout attr-dict `:` type($addr) `->` type($res)"; +} + def NVVM_MBarrierExpectTxOp : NVVM_VoidIntrinsicOp<"mbarrier.expect_tx"> { let summary = "MBarrier expect-tx Operation"; let description = [{ diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp index 7207c960985f3..d73bbebbf395b 100644 --- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp +++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp @@ -846,7 +846,7 @@ struct NVGPUMBarrierInitLowering Value barrier = getMbarrierPtr(b, mbarrierType, adaptor.getBarriers(), adaptor.getMbarId(), rewriter); Value count = truncToI32(b, adaptor.getCount()); - rewriter.replaceOpWithNewOp<NVVM::MBarrierInitOp>(op, barrier, count, + rewriter.replaceOpWithNewOp<NVVM::MBarrierInitOp>(op, barrier, count, 0, adaptor.getPredicate()); return success(); } diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index aa45bba341240..e2d2c34b09aa7 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -3754,9 +3754,12 @@ void Tcgen05MmaSmemDescOp::createSmemDescriptor(Operation &op, //===----------------------------------------------------------------------===// std::string NVVM::MBarrierInitOp::getPtx() { - bool isShared = isPtrInSharedCTASpace(getAddr()); - return isShared ? std::string("mbarrier.init.shared.b64 [%0], %1;") - : std::string("mbarrier.init.b64 [%0], %1;"); + std::string space = isPtrInSharedCTASpace(getAddr()) ? ".shared" : ""; + std::string layout = + getLayout() ? llvm::formatv(".layout::v{0}", getLayout()).str() : ""; + + return llvm::formatv("mbarrier.init{0}{1}.b64 [%0], %1;", layout, space) + .str(); } std::string NVVM::MBarrierArriveExpectTxOp::getPtx() { @@ -4080,19 +4083,29 @@ PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt, return {llvm::Intrinsic::nvvm_pm_event_mask, {maskVal}}; } +bool MBarrierInitOp::getAsmValues( + RewriterBase &rewriter, + llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> + &asmValues) { + // Add all the operands but not the attrs to the asmValues list. + // The layout attr is already baked into the PTX string by getPtx(), so + // passing it along here too would shift the operand numbering. + for (auto val : getOperands()) + asmValues.push_back({val, mlir::NVVM::PTXRegisterMod::Read}); + + return false; +} + mlir::NVVM::IDArgPair MBarrierInitOp::getIntrinsicIDAndArgs( Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { auto thisOp = cast<NVVM::MBarrierInitOp>(op); - bool isShared = isPtrInSharedCTASpace(thisOp.getAddr()); - llvm::Intrinsic::ID id = isShared ? llvm::Intrinsic::nvvm_mbarrier_init_shared - : llvm::Intrinsic::nvvm_mbarrier_init; - - // Fill the Intrinsic Args - llvm::SmallVector<llvm::Value *> args; - args.push_back(mt.lookupValue(thisOp.getAddr())); - args.push_back(mt.lookupValue(thisOp.getCount())); - return {id, std::move(args)}; + llvm::Intrinsic::ID id = isPtrInSharedCTASpace(thisOp.getAddr()) + ? llvm::Intrinsic::nvvm_mbarrier_init_shared + : llvm::Intrinsic::nvvm_mbarrier_init; + return {id, + {mt.lookupValue(thisOp.getAddr()), mt.lookupValue(thisOp.getCount()), + builder.getInt32(thisOp.getLayout())}}; } mlir::NVVM::IDArgPair MBarrierInvalOp::getIntrinsicIDAndArgs( @@ -4106,6 +4119,15 @@ mlir::NVVM::IDArgPair MBarrierInvalOp::getIntrinsicIDAndArgs( return {id, {mt.lookupValue(thisOp.getAddr())}}; } +mlir::NVVM::IDArgPair MBarrierCheckLayoutOp::getIntrinsicIDAndArgs( + Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { + auto thisOp = cast<NVVM::MBarrierCheckLayoutOp>(op); + + return { + llvm::Intrinsic::nvvm_mbarrier_check_layout, + {mt.lookupValue(thisOp.getAddr()), builder.getInt32(thisOp.getLayout())}}; +} + mlir::NVVM::IDArgPair MBarrierExpectTxOp::getIntrinsicIDAndArgs( Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { auto thisOp = cast<NVVM::MBarrierExpectTxOp>(op); diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir index 83b1eb232fa85..ded3ecc6279b5 100644 --- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir +++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir @@ -9,8 +9,10 @@ llvm.func @init_mbarrier(%barrier_gen : !llvm.ptr, %barrier : !llvm.ptr<3>, %count : i32, %pred : i1) { //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.shared.b64 [$0], $1;", "r,r,b" nvvm.mbarrier.init %barrier, %count, predicate = %pred : !llvm.ptr<3>, i32, i1 - //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b" + //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b" nvvm.mbarrier.init %barrier_gen, %count, predicate = %pred : !llvm.ptr, i32, i1 + //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.layout::v1.shared.b64 [$0], $1;", "r,r,b" + nvvm.mbarrier.init %barrier, %count layout = 1, predicate = %pred : !llvm.ptr<3>, i32, i1 llvm.return } diff --git a/mlir/test/Target/LLVMIR/nvvm/mbar_check_layout.mlir b/mlir/test/Target/LLVMIR/nvvm/mbar_check_layout.mlir new file mode 100644 index 0000000000000..5eac32fbbfabf --- /dev/null +++ b/mlir/test/Target/LLVMIR/nvvm/mbar_check_layout.mlir @@ -0,0 +1,14 @@ +// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s + +llvm.func @mbarrier_check_layout(%barrier: !llvm.ptr<3>) -> i1 { + // CHECK-LABEL: define i1 @mbarrier_check_layout(ptr addrspace(3) %0) { + // CHECK-NEXT: %[[V0:.+]] = call i1 @llvm.nvvm.mbarrier.check_layout.p3(ptr addrspace(3) %0, i32 0) + // CHECK-NEXT: %[[V1:.+]] = call i1 @llvm.nvvm.mbarrier.check_layout.p3(ptr addrspace(3) %0, i32 1) + // CHECK-NEXT: %[[RES:.+]] = or i1 %[[V0]], %[[V1]] + // CHECK-NEXT: ret i1 %[[RES]] + // CHECK-NEXT: } + %v0 = nvvm.mbarrier.check_layout %barrier, 0 : !llvm.ptr<3> -> i1 + %v1 = nvvm.mbarrier.check_layout %barrier, 1 : !llvm.ptr<3> -> i1 + %res = llvm.or %v0, %v1 : i1 + llvm.return %res : i1 +} diff --git a/mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir b/mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir index 5ea1f4915142d..0495c2a9b63be 100644 --- a/mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir +++ b/mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir @@ -18,7 +18,7 @@ llvm.func @cp_async_mbarrier_arrive(%bar_shared: !llvm.ptr<3>, %bar_gen: !llvm.p llvm.func @mbarrier_init_generic(%barrier: !llvm.ptr) { // CHECK-LABEL: define void @mbarrier_init_generic(ptr %0) { // CHECK-NEXT: %2 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() - // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init(ptr %0, i32 %2) + // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init(ptr %0, i32 %2, i32 0) // CHECK-NEXT: ret void // CHECK-NEXT: } %count = nvvm.read.ptx.sreg.ntid.x : i32 @@ -29,7 +29,7 @@ llvm.func @mbarrier_init_generic(%barrier: !llvm.ptr) { llvm.func @mbarrier_init_shared(%barrier: !llvm.ptr<3>) { // CHECK-LABEL: define void @mbarrier_init_shared(ptr addrspace(3) %0) { // CHECK-NEXT: %2 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x() - // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %0, i32 %2) + // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %0, i32 %2, i32 0) // CHECK-NEXT: ret void // CHECK-NEXT: } %count = nvvm.read.ptx.sreg.ntid.x : i32 @@ -37,6 +37,26 @@ llvm.func @mbarrier_init_shared(%barrier: !llvm.ptr<3>) { llvm.return } +llvm.func @mbarrier_init_layout_shared(%barrier: !llvm.ptr<3>, %count: i32) { + // CHECK-LABEL: define void @mbarrier_init_layout_shared(ptr addrspace(3) %0, i32 %1) { + // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %0, i32 %1, i32 0) + // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3) %0, i32 %1, i32 1) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + nvvm.mbarrier.init %barrier, %count layout = 0 : !llvm.ptr<3>, i32 + nvvm.mbarrier.init %barrier, %count layout = 1 : !llvm.ptr<3>, i32 + llvm.return +} + +llvm.func @mbarrier_init_layout_generic(%barrier: !llvm.ptr, %count: i32) { + // CHECK-LABEL: define void @mbarrier_init_layout_generic(ptr %0, i32 %1) { + // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init(ptr %0, i32 %1, i32 1) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + nvvm.mbarrier.init %barrier, %count layout = 1 : !llvm.ptr, i32 + llvm.return +} + llvm.func @mbarrier_inval_generic(%barrier: !llvm.ptr) { // CHECK-LABEL: define void @mbarrier_inval_generic(ptr %0) { // CHECK-NEXT: call void @llvm.nvvm.mbarrier.inval(ptr %0) diff --git a/mlir/test/Target/LLVMIR/nvvm/mbar_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/mbar_invalid.mlir index 32954d1d860ec..17945acaf02fb 100644 --- a/mlir/test/Target/LLVMIR/nvvm/mbar_invalid.mlir +++ b/mlir/test/Target/LLVMIR/nvvm/mbar_invalid.mlir @@ -136,3 +136,35 @@ llvm.func @mbarrier_try_wait_with_timelimit(%barrier: !llvm.ptr<3>, %phase: i32, llvm.return } +// ----- + +llvm.func @mbarrier_init_layout_too_large(%barrier: !llvm.ptr<3>, %count: i32) { + // expected-error @below {{attribute 'layout' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 0 whose maximum value is 1}} + nvvm.mbarrier.init %barrier, %count layout = 2 : !llvm.ptr<3>, i32 + llvm.return +} + +// ----- + +llvm.func @mbarrier_init_layout_negative(%barrier: !llvm.ptr<3>, %count: i32) { + // expected-error @below {{attribute 'layout' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 0 whose maximum value is 1}} + nvvm.mbarrier.init %barrier, %count layout = -1 : !llvm.ptr<3>, i32 + llvm.return +} + +// ----- + +llvm.func @mbarrier_check_layout_too_large(%barrier: !llvm.ptr<3>) { + // expected-error @below {{attribute 'layout' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 0 whose maximum value is 1}} + %0 = nvvm.mbarrier.check_layout %barrier, 2 : !llvm.ptr<3> -> i1 + llvm.return +} + +// ----- + +llvm.func @mbarrier_check_layout_negative(%barrier: !llvm.ptr<3>) { + // expected-error @below {{attribute 'layout' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 0 whose maximum value is 1}} + %0 = nvvm.mbarrier.check_layout %barrier, -1 : !llvm.ptr<3> -> i1 + llvm.return +} + _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
