[flang-commits] [clang] [flang] [llvm] [mlir] [LLVM][NVPTX][MLIR] Add mbarrier layout support (PR #217252)
Pradeep Kumar via flang-commits
flang-commits at lists.llvm.org
Mon Aug 31 04:16:44 PDT 2026
https://github.com/schwarzschild-radius updated https://github.com/llvm/llvm-project/pull/217252
>From 1500876b421db3f5ea515fc08f567a04a055f60a Mon Sep 17 00:00:00 2001
From: Pradeep Kumar <pradeepku at nvidia.com>
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 <noreply at anthropic.com>
---
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
+}
+
More information about the flang-commits
mailing list