[clang] [flang] [llvm] [mlir] [LLVM][NVPTX][MLIR] Support layouts in mbarrier.init and add mbarrier.check_layout (PR #217252)
Pradeep Kumar via cfe-commits
cfe-commits at lists.llvm.org
Wed Sep 9 05:15:39 PDT 2026
https://github.com/schwarzschild-radius updated https://github.com/llvm/llvm-project/pull/217252
>From e89fb27bed2cd51a345fe8e060f4a847d0a2bc3e 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] Support layouts in mbarrier.init and add
mbarrier.check_layout
This commit adds LLVM NVPTX and MLIR NVVM support for the mbarrier layout
extensions:
- llvm.nvvm.mbarrier.init with trailing immarg `layout` operand, and
nvvm.mbarrier.init a matching `layout` attribute (default to 0)
Layout 0 is the default in-memory layout and is emitted as a plain
mbarrier.init; layout 1 is emitted as mbarrier.init.layout::v1.
- llvm.nvvm.mbarrier.check_layout / nvvm.mbarrier.check_layout, lowering
to mbarrier.check_layout.layout::v{0,1}.
Only the layout::v1 form of mbarrier.init and mbarrier.check_layout require
PTX ISA 9.3 and sm_90
llvm.nvvm.mbarrier.init and llvm.nvvm.mbarrier.init.shared are also merged
into a single pointer-overloaded llvm.nvvm.mbarrier.init, so the address
space now selects between the generic and shared::cta forms. Calls to either
old name, with or without the layout operand, are auto-upgraded.
Assisted by: Claude Code (Opus 5)
---
clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp | 9 +++
clang/test/CodeGen/builtins-nvptx.c | 4 +-
.../Optimizer/Builder/CUDAIntrinsicCall.cpp | 3 +-
llvm/docs/NVPTXUsage.md | 42 ++++++++--
llvm/include/llvm/IR/IntrinsicsNVVM.td | 51 ++++++++----
llvm/include/llvm/IR/NVVMIntrinsicUtils.h | 9 +++
llvm/lib/IR/AutoUpgrade.cpp | 25 ++++++
llvm/lib/IR/NVVMIntrinsicUtils.cpp | 13 +++
llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 22 ++++++
llvm/lib/Target/NVPTX/NVPTXInstrInfo.td | 1 +
llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 37 +++++++--
.../Assembler/auto_upgrade_nvvm_intrinsics.ll | 12 +++
llvm/test/CodeGen/NVPTX/mbarrier.ll | 8 +-
.../NVPTX/mbarrier_layout_sm90_ptx93.ll | 79 +++++++++++++++++++
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 49 +++++++++++-
.../Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp | 2 +-
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 46 ++++++++---
.../Conversion/NVVMToLLVM/nvvm-to-llvm.mlir | 6 +-
mlir/test/Dialect/LLVMIR/nvvm.mlir | 40 ++++++++++
.../LLVMIR/nvvm_check_target_sm_trait.mlir | 19 +++++
.../Target/LLVMIR/nvvm/mbar_check_layout.mlir | 27 +++++++
mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir | 24 +++++-
.../test/Target/LLVMIR/nvvm/mbar_invalid.mlir | 32 ++++++++
23 files changed, 506 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 06c5069d6f984..c78814a2fbc97 100644
--- a/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp
+++ b/clang/lib/CodeGen/TargetBuiltins/NVPTX.cpp
@@ -1393,6 +1393,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 bed1498236b06..de9ec1f878d66 100644
--- a/clang/test/CodeGen/builtins-nvptx.c
+++ b/clang/test/CodeGen/builtins-nvptx.c
@@ -946,9 +946,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 6e0ca158266ca..ee12c940a5e72 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,48 @@ 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.p0(ptr %addr, i32 immarg range(i32 0, 2) %layout)
+declare i1 @llvm.nvvm.mbarrier.check_layout.p3(ptr addrspace(3) %addr, i32 immarg range(i32 0, 2) %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:
+
+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. 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 87ae664ce4c17..4fb9aef1596e0 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.
@@ -2346,14 +2359,20 @@ def int_nvvm_cp_async_bulk_wait_group_read :
Intrinsic<[], [llvm_i32_ty], [ImmArg<ArgIndex<0>>]>;
// mbarrier
+
+def int_nvvm_mbarrier_init :
+ DefaultAttrsIntrinsicFlags<[],
+ [llvm_anyptr_ty, // mbar (generic or shared AS)
+ llvm_i32_ty], // count
+ [llvm_i32_ty], // layout
+ [IntrConvergent, Range<ArgIndex<2>, 0, 2>,
+ ArgInfo<ArgIndex<2>, [ArgName<"layout">,
+ ImmArgPrinter<"printMBarrierLayout">]>]>;
+
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,
@@ -2375,6 +2394,17 @@ 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 :
+ DefaultAttrsIntrinsicFlags<[llvm_i1_ty],
+ [llvm_anyptr_ty], // mbar
+ [llvm_i32_ty], // layout
+ [IntrReadMem,
+ Range<ArgIndex<1>, 0, 2>,
+ ArgInfo<ArgIndex<1>,
+ [ArgName<"layout">,
+ ImmArgPrinter<"printMBarrierLayout">]>],
+ "llvm.nvvm.mbarrier.check_layout">;
+
// mbarrier.{expect_tx/complete_tx}
foreach op = ["expect_tx", "complete_tx"] in {
foreach scope = ["scope_cta", "scope_cluster"] in {
@@ -3127,19 +3157,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/include/llvm/IR/NVVMIntrinsicUtils.h b/llvm/include/llvm/IR/NVVMIntrinsicUtils.h
index 6fdaf631a12f9..c6f3567791944 100644
--- a/llvm/include/llvm/IR/NVVMIntrinsicUtils.h
+++ b/llvm/include/llvm/IR/NVVMIntrinsicUtils.h
@@ -186,8 +186,17 @@ enum class TensormapFillMode : uint8_t {
OOB_NAN_FILL = 1,
};
+// In-memory layout of an mbarrier object, as selected by the layout operand of
+// the llvm.nvvm.mbarrier.init and llvm.nvvm.mbarrier.check_layout intrinsics.
+enum class MBarrierLayout : uint8_t {
+ V0 = 0,
+ V1 = 1,
+};
+
LLVM_ABI void printTcgen05MMAKind(raw_ostream &OS, const Constant *ImmArgVal);
+LLVM_ABI void printMBarrierLayout(raw_ostream &OS, const Constant *ImmArgVal);
+
LLVM_ABI void printEvictPolicyType(raw_ostream &OS, const Constant *ImmArgVal);
LLVM_ABI void printTMAReductionOp(raw_ostream &OS, const Constant *ImmArgVal);
diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp
index 530ee5b7b6047..fb756f9d1078c 100644
--- a/llvm/lib/IR/AutoUpgrade.cpp
+++ b/llvm/lib/IR/AutoUpgrade.cpp
@@ -1498,6 +1498,13 @@ getNVVMFAddUpgrade(StringRef Name) {
return std::make_pair(IID, *RoundingMode);
}
+static Intrinsic::ID shouldUpgradeNVPTXMBarrierInitIntrinsic(StringRef Name) {
+ 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") ||
@@ -2088,6 +2095,15 @@ static bool upgradeIntrinsicFunction1(Function *F, Function *&NewFn,
return NewFn != F;
}
+ // Upgrade mbarrier.init intrinsics missing the layout operand.
+ IID = shouldUpgradeNVPTXMBarrierInitIntrinsic(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.
//
@@ -6201,6 +6217,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/IR/NVVMIntrinsicUtils.cpp b/llvm/lib/IR/NVVMIntrinsicUtils.cpp
index f6816fa0e706b..d0433040fc7ab 100644
--- a/llvm/lib/IR/NVVMIntrinsicUtils.cpp
+++ b/llvm/lib/IR/NVVMIntrinsicUtils.cpp
@@ -49,6 +49,19 @@ void nvvm::printTMAValidateDataPattern(raw_ostream &OS,
static_cast<TMAValidateDataPattern>(CI->getZExtValue()));
}
+void nvvm::printMBarrierLayout(raw_ostream &OS, const Constant *ImmArgVal) {
+ if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
+ switch (static_cast<MBarrierLayout>(CI->getZExtValue())) {
+ case MBarrierLayout::V0:
+ OS << "v0";
+ return;
+ case MBarrierLayout::V1:
+ OS << "v1";
+ return;
+ }
+ }
+}
+
void nvvm::printTcgen05MMAKind(raw_ostream &OS, const Constant *ImmArgVal) {
if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
uint64_t Val = CI->getZExtValue();
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index a946ceaa35d81..a21af472576cd 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -4755,6 +4755,28 @@ void NVPTXTargetLowering::getTgtMemIntrinsic(
return;
}
+ case Intrinsic::nvvm_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_mbarrier_check_layout: {
+ Info.opc = ISD::INTRINSIC_W_CHAIN;
+ Info.memVT = MVT::i64;
+ Info.ptrVal = I.getArgOperand(0);
+ Info.offset = 0;
+ Info.flags = MachineMemOperand::MOLoad;
+ 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/NVPTXInstrInfo.td b/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td
index 846bcc0557464..03090c5aaf964 100644
--- a/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td
+++ b/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td
@@ -251,6 +251,7 @@ def hasConvertWithStochasticRounding
def hasConvertWithPZOSupport : PredOr<[hasRubinFamilySupport]>;
+def hasMBarrierLayoutSupport : PredAnd<[PTX93, SM90]>;
//===----------------------------------------------------------------------===//
// Some Common Instruction Class Templates
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 220ef64732830..e73d1bee2dea7 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -1577,14 +1577,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),
@@ -1653,6 +1653,31 @@ let Predicates = [SM80] in {
[(set i32:$res, (int_nvvm_mbarrier_pending_count i64:$state))]>;
}
+let Predicates = [hasMBarrierLayoutSupport] 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>;
+
+ class MBARRIER_CHECK_LAYOUT<NVPTXAddressSpace as, int layout> :
+ BasicNVPTXInst<(outs B1:$res), (ins ADDR:$addr),
+ "mbarrier.check_layout.layout::v" # layout # as.Suffix # ".b64",
+ [(set i1:$res,
+ (IntrinsicInAS<int_nvvm_mbarrier_check_layout, as>
+ addr:$addr, (i32 layout)))]>;
+
+ foreach layout = [0, 1] in {
+ def MBARRIER_CHECK_LAYOUT_V # layout :
+ MBARRIER_CHECK_LAYOUT<AddrSpaceGeneric, layout>;
+ def MBARRIER_CHECK_LAYOUT_V # layout # _SHARED :
+ MBARRIER_CHECK_LAYOUT<AddrSpaceShared, 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 f95215834320e..390d82b4668e4 100644
--- a/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll
+++ b/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll
@@ -776,3 +776,15 @@ define void @nvvm_add(float %a, double %b, half %c, <2 x half> %d) {
%r10 = call <2 x half> @llvm.nvvm.add.rn.ftz.sat.v2f16(<2 x half> %d, <2 x half> %d)
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.p0(ptr %gen, i32 %count, /* layout=v0 */ i32 0)
+; CHECK: call void @llvm.nvvm.mbarrier.init.p3(ptr addrspace(3) %shared, i32 %count, /* layout=v0 */ 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..39e1a13ea2ae9
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll
@@ -0,0 +1,79 @@
+; 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 void @mbarrier_init_generic(ptr %a, i32 %b) {
+; CHECK-LABEL: mbarrier_init_generic(
+; 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_generic_param_0];
+; CHECK-NEXT: ld.param::func.b32 %r1, [mbarrier_init_generic_param_1];
+; CHECK-NEXT: mbarrier.init.b64 [%rd1], %r1;
+; CHECK-NEXT: mbarrier.init.layout::v1.b64 [%rd1], %r1;
+; CHECK-NEXT: ret;
+ tail call void @llvm.nvvm.mbarrier.init.p0(ptr %a, i32 %b, i32 0)
+ tail call void @llvm.nvvm.mbarrier.init.p0(ptr %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
+}
+
+define i1 @mbarrier_check_layout_generic(ptr %a) {
+; CHECK-LABEL: mbarrier_check_layout_generic(
+; 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_generic_param_0];
+; CHECK-NEXT: mbarrier.check_layout.layout::v0.b64 %p1, [%rd1];
+; CHECK-NEXT: mbarrier.check_layout.layout::v1.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.p0(ptr %a, i32 0)
+ %is_layout_v1 = tail call i1 @llvm.nvvm.mbarrier.check_layout.p0(ptr %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 80a0094190383..c5e99d999920d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -632,7 +632,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
@@ -651,15 +654,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);
@@ -668,7 +678,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);
}];
}
@@ -695,6 +705,39 @@ def NVVM_MBarrierInvalOp : NVVM_VoidIntrinsicOp<"mbarrier.inval">,
let assemblyFormat = "$addr attr-dict `:` type(operands)";
}
+def NVVM_MBarrierCheckLayoutOp :
+ NVVM_SingleResultIntrinsicOp<"mbarrier.check_layout",
+ [NVVMRequiresSM<90>]> {
+ 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 generic or shared::cta memory. When it is
+ generic, the underlying address must be within the shared::cta memory
+ space; otherwise the behavior is undefined.
+ - `layout`: The mbarrier layout version to test for. Only `0`
+ (`layout::v0`) and `1` (`layout::v1`) are valid values. When it is
+ omitted, it defaults to `0`.
+
+ [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
+ AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
+ DefaultValuedAttr<ConfinedAttr<I32Attr,
+ [IntMinValue<0>, IntMaxValue<1>]>, "0">:$layout);
+
+ let assemblyFormat =
+ "$addr (`layout` `=` $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 7c393e8cf2b66..fe73f72812f24 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -847,7 +847,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 68f2a79fa7e36..6705f423e3332 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3755,9 +3755,13 @@ 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" : "";
+ // Layout v0 is the default, so it is emitted as a plain mbarrier.init.
+ std::string layout =
+ getLayout() == 1 ? std::string(".layout::v1") : std::string();
+
+ return llvm::formatv("mbarrier.init{0}{1}.b64 [%0], %1;", layout, space)
+ .str();
}
std::string NVVM::MBarrierArriveExpectTxOp::getPtx() {
@@ -4081,19 +4085,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(
@@ -4107,6 +4120,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..2ff915aa461dd 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -9,8 +9,12 @@
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
+ //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.layout::v1.b64 [$0], $1;", "l,r,b"
+ nvvm.mbarrier.init %barrier_gen, %count layout = 1, predicate = %pred : !llvm.ptr, i32, i1
llvm.return
}
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 7003cfe319f8f..e619f6e3092a3 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -452,6 +452,46 @@ llvm.func private @mbarrier_init_shared(%barrier: !llvm.ptr<3>) {
}
+// The `layout` attribute and the optional `predicate` operand are independent,
+// so all four combinations must round-trip.
+llvm.func private @mbarrier_init_layout_predicate(%barrier: !llvm.ptr<3>,
+ %count: i32, %pred: i1) {
+ // CHECK: nvvm.mbarrier.init %{{.*}}, %{{.*}} : !llvm.ptr<3>, i32
+ nvvm.mbarrier.init %barrier, %count : !llvm.ptr<3>, i32
+ // CHECK: nvvm.mbarrier.init %{{.*}}, %{{.*}} layout = 1 : !llvm.ptr<3>, i32
+ nvvm.mbarrier.init %barrier, %count layout = 1 : !llvm.ptr<3>, i32
+ // CHECK: nvvm.mbarrier.init %{{.*}}, %{{.*}}, predicate = %{{.*}} : !llvm.ptr<3>, i32, i1
+ nvvm.mbarrier.init %barrier, %count, predicate = %pred : !llvm.ptr<3>, i32, i1
+ // CHECK: nvvm.mbarrier.init %{{.*}}, %{{.*}} layout = 1, predicate = %{{.*}} : !llvm.ptr<3>, i32, i1
+ nvvm.mbarrier.init %barrier, %count layout = 1, predicate = %pred : !llvm.ptr<3>, i32, i1
+ llvm.return
+}
+
+
+llvm.func private @mbarrier_check_layout_generic(%barrier: !llvm.ptr) {
+ // CHECK: nvvm.mbarrier.check_layout %{{.*}} layout = 1 : !llvm.ptr -> i1
+ %0 = nvvm.mbarrier.check_layout %barrier layout = 1 : !llvm.ptr -> i1
+ llvm.return
+}
+
+
+llvm.func private @mbarrier_check_layout_shared(%barrier: !llvm.ptr<3>) {
+ // CHECK: nvvm.mbarrier.check_layout %{{.*}} layout = 1 : !llvm.ptr<3> -> i1
+ %0 = nvvm.mbarrier.check_layout %barrier layout = 1 : !llvm.ptr<3> -> i1
+ llvm.return
+}
+
+
+// `layout` defaults to 0, so it is elided when absent and when written out.
+llvm.func private @mbarrier_check_layout_default(%barrier: !llvm.ptr<3>) {
+ // CHECK: nvvm.mbarrier.check_layout %{{.*}} : !llvm.ptr<3> -> i1
+ %0 = nvvm.mbarrier.check_layout %barrier : !llvm.ptr<3> -> i1
+ // CHECK: nvvm.mbarrier.check_layout %{{.*}} : !llvm.ptr<3> -> i1
+ %1 = nvvm.mbarrier.check_layout %barrier layout = 0 : !llvm.ptr<3> -> i1
+ llvm.return
+}
+
+
llvm.func private @mbarrier_inval_generic(%barrier: !llvm.ptr) {
// CHECK: nvvm.mbarrier.inval %{{.*}} : !llvm.ptr
nvvm.mbarrier.inval %barrier : !llvm.ptr
diff --git a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm_trait.mlir b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm_trait.mlir
index 686b9671fb4ee..d95a42f8c94ea 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm_trait.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm_trait.mlir
@@ -158,3 +158,22 @@ gpu.module @check_invalid_SM_arch_or_family [#nvvm.target<chip = "sm_100">] {
// expected-error @below {{is not supported on sm_100}}
test.nvvm_requires_sm_90a_or_sm_100f
}
+
+// -----
+
+gpu.module @check_valid_SM_mbarrier_check_layout [#nvvm.target<chip = "sm_90">] {
+ llvm.func @mbarrier_check_layout(%barrier: !llvm.ptr<3>) {
+ %0 = nvvm.mbarrier.check_layout %barrier layout = 1 : !llvm.ptr<3> -> i1
+ llvm.return
+ }
+}
+
+// -----
+
+gpu.module @check_invalid_SM_mbarrier_check_layout [#nvvm.target<chip = "sm_80">] {
+ llvm.func @mbarrier_check_layout(%barrier: !llvm.ptr<3>) {
+ // expected-error @below {{is not supported on sm_80}}
+ %0 = nvvm.mbarrier.check_layout %barrier layout = 1 : !llvm.ptr<3> -> 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..b073aa745022b
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/mbar_check_layout.mlir
@@ -0,0 +1,27 @@
+// 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, /* layout=v0 */ i32 0)
+ // CHECK-NEXT: %[[V1:.+]] = call i1 @llvm.nvvm.mbarrier.check_layout.p3(ptr addrspace(3) %0, /* layout=v1 */ i32 1)
+ // CHECK-NEXT: %[[RES:.+]] = or i1 %[[V0]], %[[V1]]
+ // CHECK-NEXT: ret i1 %[[RES]]
+ // CHECK-NEXT: }
+ %v0 = nvvm.mbarrier.check_layout %barrier layout = 0 : !llvm.ptr<3> -> i1
+ %v1 = nvvm.mbarrier.check_layout %barrier layout = 1 : !llvm.ptr<3> -> i1
+ %res = llvm.or %v0, %v1 : i1
+ llvm.return %res : i1
+}
+
+llvm.func @mbarrier_check_layout_generic(%barrier: !llvm.ptr) -> i1 {
+ // CHECK-LABEL: define i1 @mbarrier_check_layout_generic(ptr %0) {
+ // CHECK-NEXT: %[[V0:.+]] = call i1 @llvm.nvvm.mbarrier.check_layout.p0(ptr %0, /* layout=v0 */ i32 0)
+ // CHECK-NEXT: %[[V1:.+]] = call i1 @llvm.nvvm.mbarrier.check_layout.p0(ptr %0, /* layout=v1 */ i32 1)
+ // CHECK-NEXT: %[[RES:.+]] = or i1 %[[V0]], %[[V1]]
+ // CHECK-NEXT: ret i1 %[[RES]]
+ // CHECK-NEXT: }
+ %v0 = nvvm.mbarrier.check_layout %barrier layout = 0 : !llvm.ptr -> i1
+ %v1 = nvvm.mbarrier.check_layout %barrier layout = 1 : !llvm.ptr -> 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..3934b8187073a 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, /* layout=v0 */ 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, /* layout=v0 */ 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, /* layout=v0 */ i32 0)
+ // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.p3(ptr addrspace(3) %0, i32 %1, /* layout=v1 */ 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, /* layout=v1 */ 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..af441a492c954 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 layout = 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 layout = -1 : !llvm.ptr<3> -> i1
+ llvm.return
+}
+
More information about the cfe-commits
mailing list