https://github.com/schwarzschild-radius updated https://github.com/llvm/llvm-project/pull/217252
>From 3b068062175d7b3dc32b3d1055f776a4b950cb10 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 with layout argument / 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 | 9 ++++ clang/test/CodeGen/builtins-nvptx.c | 4 +- .../Optimizer/Builder/CUDAIntrinsicCall.cpp | 3 +- llvm/docs/NVPTXUsage.md | 39 ++++++++++++-- llvm/include/llvm/IR/IntrinsicsNVVM.td | 52 +++++++++++++------ llvm/lib/IR/AutoUpgrade.cpp | 30 +++++++++++ llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 14 +++++ llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 29 ++++++++--- .../Assembler/auto_upgrade_nvvm_intrinsics.ll | 14 +++++ llvm/test/CodeGen/NVPTX/mbarrier.ll | 8 +-- .../NVPTX/mbarrier_layout_sm90_ptx93.ll | 41 +++++++++++++++ mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 44 ++++++++++++++-- .../Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp | 2 +- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 45 +++++++++++----- .../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 ++++++++++++ 18 files changed, 354 insertions(+), 54 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..dc5b1e4c1be3c 100644 --- a/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp +++ b/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp @@ -1282,6 +1282,15 @@ 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: { + // The intrinsic is overloaded on the pointer, so the two builtins differ + // only in the address space of their first argument. + Value *Ptr = EmitScalarExpr(E->getArg(0)); + return Builder.CreateIntrinsic( + Intrinsic::nvvm_mbarrier_init, {Ptr->getType()}, + {Ptr, EmitScalarExpr(E->getArg(1)), Builder.getInt32(0)}); + } default: return nullptr; } diff --git a/clang/test/CodeGen/builtins-nvptx.c b/clang/test/CodeGen/builtins-nvptx.c index 0eb9eadb251d6..c277e45a776d1 100644 --- a/clang/test/CodeGen/builtins-nvptx.c +++ b/clang/test/CodeGen/builtins-nvptx.c @@ -940,9 +940,9 @@ __device__ void nvvm_nanosleep(int d) { __device__ void nvvm_mbarrier(long long* addr, __attribute__((address_space(3))) long long* sharedAddr, int count, long long state) { #if __CUDA_ARCH__ >= 800 __nvvm_mbarrier_init(addr, count); - // CHECK_PTX70_SM80: call void @llvm.nvvm.mbarrier.init + // CHECK_PTX70_SM80: call void @llvm.nvvm.mbarrier.init.p0 __nvvm_mbarrier_init_shared(sharedAddr, count); - // CHECK_PTX70_SM80: call void @llvm.nvvm.mbarrier.init.shared + // CHECK_PTX70_SM80: call void @llvm.nvvm.mbarrier.init.p3 __nvvm_mbarrier_inval(addr); // CHECK_PTX70_SM80: call void @llvm.nvvm.mbarrier.inval diff --git a/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp b/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp index ca200ac2cd02a..a7267d834cf6a 100644 --- a/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp +++ b/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp @@ -1003,7 +1003,8 @@ void CUDAIntrinsicLibrary::genBarrierInit( mlir::Value barrier = convertPtrToNVVMSpace( builder, loc, fir::getBase(args[0]), mlir::NVVM::NVVMMemorySpace::Shared); mlir::NVVM::MBarrierInitOp::create(builder, loc, barrier, - fir::getBase(args[1]), {}); + fir::getBase(args[1]), /*layout=*/0, + /*predicate=*/{}); auto kind = mlir::NVVM::ProxyKindAttr::get( builder.getContext(), mlir::NVVM::ProxyKind::async_shared); auto space = mlir::NVVM::SharedSpaceAttr::get( diff --git a/llvm/docs/NVPTXUsage.md b/llvm/docs/NVPTXUsage.md index e05af78012691..704f90c8de617 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.p0(ptr %addr, i32 %count, i32 immarg range(i32 0, 2) %layout) +declare void @llvm.nvvm.mbarrier.init.p3(ptr addrspace(3) %addr, i32 %count, i32 immarg range(i32 0, 2) %layout) ``` ##### Overview: @@ -453,16 +453,45 @@ 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: -The `.shared` variant explicitly uses shared memory address space for -the `addr` operand. If the `addr` does not fall within the -shared::cta space, then the behavior of this intrinsic is undefined. +The `addr` operand is overloaded and must point to either generic or +shared::cta memory. When it is generic, the underlying address must fall +within the shared::cta space, otherwise the behavior of this intrinsic +is undefined. 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.p3(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 5fe8051d2b039..2ee8a287b84a6 100644 --- a/llvm/include/llvm/IR/IntrinsicsNVVM.td +++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td @@ -162,6 +162,19 @@ defvar MAX_BLOCK_SIZE_X = 1024; defvar MAX_BLOCK_SIZE_Y = 1024; defvar MAX_BLOCK_SIZE_Z = 64; +class DefaultAttrsIntrinsicFlags<list<LLVMType> ret_types, + list<LLVMType> param_types, + list<LLVMType> flags, + list<IntrinsicProperty> intr_properties, + string name = ""> + : DefaultAttrsIntrinsic< + ret_types, + !listconcat(param_types, flags), + !listconcat(intr_properties, + !foreach(i, !range(flags), + ImmArg<ArgIndex<!add(i, !size(param_types))>>)), + name>; + // Helper class that concatenates list elements with // a given separator 'sep' and returns the result. // Handles empty strings. @@ -2293,14 +2306,24 @@ def int_nvvm_cp_async_bulk_wait_group_read : Intrinsic<[], [llvm_i32_ty], [ImmArg<ArgIndex<0>>]>; // mbarrier + +// The mbar pointer is overloaded and must point to generic or shared::cta +// memory. The layout operand selects the in-memory layout of the mbarrier +// object: 0 is the default layout, emitted as a plain mbarrier.init, and 1 +// selects mbarrier.init.layout::v1. Not a NVVMBuiltin: the __nvvm_mbarrier_init +// builtins keep their two-argument, non-overloaded signature and are handled in +// clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp. +def int_nvvm_mbarrier_init : + DefaultAttrsIntrinsicFlags<[], + [llvm_anyptr_ty, // mbar + llvm_i32_ty], // count + [llvm_i32_ty], // layout + [IntrConvergent, Range<ArgIndex<2>, 0, 2>]>; + 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_inval # shared : NVVMBuiltin, Intrinsic<[], [mbarrier_ptr_ty], [IntrConvergent, IntrWriteMem, IntrArgMemOnly, IntrNoCallback, @@ -2322,6 +2345,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 { @@ -3074,19 +3105,6 @@ foreach op = ["dec", "inc"] in def int_nvvm_exit : NVVMBuiltin, Intrinsic<[], [], [IntrConvergent, IntrInaccessibleMemOnly, IntrNoReturn]>; -class DefaultAttrsIntrinsicFlags<list<LLVMType> ret_types, - list<LLVMType> param_types, - list<LLVMType> flags, - list<IntrinsicProperty> intr_properties, - string name = ""> - : DefaultAttrsIntrinsic< - ret_types, - !listconcat(param_types, flags), - !listconcat(intr_properties, - !foreach(i, !range(flags), - ImmArg<ArgIndex<!add(i, !size(param_types))>>)), - name>; - // TMA Tensor Copy Intrinsics: S2G -> From Shared to Global memory variants foreach dim = 1...5 in { defvar tensor_dim_args = !listsplat(llvm_i32_ty, dim); diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp index cd505273b5bfb..b92771c4fce3b 100644 --- a/llvm/lib/IR/AutoUpgrade.cpp +++ b/llvm/lib/IR/AutoUpgrade.cpp @@ -1400,6 +1400,18 @@ 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) { + // The overloaded replacement is mangled with a pointer suffix, so an exact + // match on the old names never fires for an already-upgraded declaration. + if (Name != "mbarrier.init" && Name != "mbarrier.init.shared") + return Intrinsic::not_intrinsic; + + return Intrinsic::nvvm_mbarrier_init; +} + static bool consumeNVVMPtrAddrSpace(StringRef &Name) { return Name.consume_front("local") || Name.consume_front("shared") || Name.consume_front("global") || Name.consume_front("constant") || @@ -1971,6 +1983,15 @@ 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, + F->getArg(0)->getType()); + 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,15 @@ void llvm::UpgradeIntrinsicCall(CallBase *CI, Function *NewFn) { Builder.CreateCall(NewFn, {CI->getArgOperand(0), CI->getArgOperand(1), Builder.getFalse()}); break; + case Intrinsic::nvvm_mbarrier_init: { + SmallVector<Value *, 3> Args(CI->args()); + // The .shared variant folded into the overloaded form without gaining an + // operand, so only the pre-layout two-argument form needs one appended. + if (Args.size() == 2) + 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/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp index 9df1419d7349d..cf61956642bdc 100644 --- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp +++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp @@ -4749,6 +4749,20 @@ void NVPTXTargetLowering::getTgtMemIntrinsic( return; } + case Intrinsic::nvvm_mbarrier_init: { + // Registered here so that the address space of the mbarrier object is + // preserved in the SelectionDAG, letting ISel pick between the generic + // and the shared::cta form of mbarrier.init. + Info.opc = ISD::INTRINSIC_VOID; + Info.memVT = MVT::i64; + Info.ptrVal = I.getArgOperand(0); + Info.offset = 0; + Info.flags = MachineMemOperand::MOStore; + Info.align = Align(8); + Infos.push_back(Info); + return; + } + case Intrinsic::nvvm_tensormap_replace_global_address: case Intrinsic::nvvm_tensormap_replace_global_stride: { Info.opc = ISD::INTRINSIC_VOID; diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td index 4cf47e32622ab..eb99db065772f 100644 --- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td +++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td @@ -1399,14 +1399,14 @@ def DISCARD_GLOBAL_L2 : DISCARD_L2_INTRS<"global">; //----------------------------------- let Predicates = [SM80] in { - class MBARRIER_INIT<string AddrSpace, Intrinsic Intrin> : + class MBARRIER_INIT<NVPTXAddressSpace as> : BasicNVPTXInst<(outs), (ins ADDR:$addr, B32:$count), - "mbarrier.init" # AddrSpace # ".b64", - [(Intrin addr:$addr, i32:$count)]>; + "mbarrier.init" # as.Suffix # ".b64", + [(IntrinsicInAS<int_nvvm_mbarrier_init, as> + addr:$addr, i32:$count, (i32 0))]>; - def MBARRIER_INIT : MBARRIER_INIT<"", int_nvvm_mbarrier_init>; - def MBARRIER_INIT_SHARED : MBARRIER_INIT<".shared", - int_nvvm_mbarrier_init_shared>; + def MBARRIER_INIT : MBARRIER_INIT<AddrSpaceGeneric>; + def MBARRIER_INIT_SHARED : MBARRIER_INIT<AddrSpaceShared>; class MBARRIER_INVAL<string AddrSpace, Intrinsic Intrin> : BasicNVPTXInst<(outs), (ins ADDR:$addr), @@ -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<NVPTXAddressSpace as> : + BasicNVPTXInst<(outs), (ins ADDR:$addr, B32:$count), + "mbarrier.init.layout::v1" # as.Suffix # ".b64", + [(IntrinsicInAS<int_nvvm_mbarrier_init, as> + addr:$addr, i32:$count, (i32 1))]>; + + def MBARRIER_INIT_LAYOUT : MBARRIER_INIT_LAYOUT<AddrSpaceGeneric>; + def MBARRIER_INIT_LAYOUT_SHARED : MBARRIER_INIT_LAYOUT<AddrSpaceShared>; + + 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..dfa09121d0ac6 100644 --- a/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll +++ b/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll @@ -711,3 +711,17 @@ 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) + +; The layout operand is appended, and the .shared variant folds into the +; pointer-overloaded form. +; 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.p0(ptr %gen, i32 %count, i32 0) +; CHECK: call void @llvm.nvvm.mbarrier.init.p3(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..f723a8d761eff 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.p0(ptr %a, i32 %b, i32 %c) +declare void @llvm.nvvm.mbarrier.init.p3(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.p0(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.p3(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..a105ca65c9268 --- /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.p3(ptr addrspace(3) %a, i32 %b, i32 0) + tail call void @llvm.nvvm.mbarrier.init.p3(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.p3(ptr addrspace(3) %a, i32 0) + %is_layout_v1 = tail call i1 @llvm.nvvm.mbarrier.check_layout.p3(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..295214f6342fb 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); @@ -677,7 +687,7 @@ def NVVM_MBarrierInitOp : NVVM_PTXBuilder_Op<"mbarrier.init">, string llvmBuilder = [{ auto [id, args] = NVVM::MBarrierInitOp::getIntrinsicIDAndArgs( *op, moduleTranslation, builder); - createIntrinsicCall(builder, id, args); + createIntrinsicCall(builder, id, builder.getVoidTy(), args); }]; } @@ -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..27c8f9798638a 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,28 @@ 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)}; + // The intrinsic is overloaded on the mbarrier pointer, so the address space + // selects the generic or shared::cta form on its own. + return {llvm::Intrinsic::nvvm_mbarrier_init, + {mt.lookupValue(thisOp.getAddr()), mt.lookupValue(thisOp.getCount()), + builder.getInt32(thisOp.getLayout())}}; } mlir::NVVM::IDArgPair MBarrierInvalOp::getIntrinsicIDAndArgs( @@ -4106,6 +4118,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..5797dc1672335 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.p0(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.p3(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.p3(ptr addrspace(3) %0, i32 %1, i32 0) + // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.p3(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.p0(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
