[Mlir-commits] [llvm] [mlir] [LLVM][NVPTX][MLIR] Add mbarrier layout support (PR #217252)

Pradeep Kumar llvmlistbot at llvm.org
Wed Aug 19 02:30:24 PDT 2026


https://github.com/schwarzschild-radius updated https://github.com/llvm/llvm-project/pull/217252

>From b21535f34de14a4b05211233f55c7e1946db62e9 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.layout / nvvm.mbarrier.init's optional
    `layout` attribute, lowering to mbarrier.init.layout::v{0,1}.
  - llvm.nvvm.mbarrier.check_layout / nvvm.mbarrier.check_layout,
    lowering to mbarrier.check_layout.layout::v{0,1}.

Both require PTX ISA 9.3 and sm_90.

Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
 llvm/docs/NVPTXUsage.md                       | 31 ++++++++++++++
 llvm/include/llvm/IR/IntrinsicsNVVM.td        | 14 +++++++
 llvm/lib/Target/NVPTX/NVPTXIntrinsics.td      | 13 ++++++
 .../NVPTX/mbarrier_layout_sm90_ptx93.ll       | 41 +++++++++++++++++++
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td   | 36 +++++++++++++++-
 .../Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp    |  4 +-
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp    | 41 ++++++++++++++-----
 .../Conversion/NVVMToLLVM/nvvm-to-llvm.mlir   |  4 +-
 .../Target/LLVMIR/nvvm/mbar_check_layout.mlir | 14 +++++++
 mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir   | 22 ++++++++++
 .../test/Target/LLVMIR/nvvm/mbar_invalid.mlir | 32 +++++++++++++++
 11 files changed, 237 insertions(+), 15 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/llvm/docs/NVPTXUsage.md b/llvm/docs/NVPTXUsage.md
index 8924a44e43a8e..7dfd2672b644e 100644
--- a/llvm/docs/NVPTXUsage.md
+++ b/llvm/docs/NVPTXUsage.md
@@ -442,6 +442,7 @@ For more information, refer [PTX ISA](https://docs.nvidia.com/cuda/parallel-thre
 ```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.layout(ptr addrspace(3) %addr, i32 %count, i32 immarg %layout)
 ```
 
 ##### Overview:
@@ -454,6 +455,10 @@ 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`.
 
+The `.layout` variant additionally selects the in-memory layout of the
+mbarrier object. `%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
@@ -463,6 +468,32 @@ Performing `mbarrier.init` on a valid mbarrier object is undefined;
 use `mbarrier.inval` before reusing the memory for another mbarrier
 or any other purpose.
 
+An mbarrier object initialized with a particular layout must only be
+used with operations that support that layout; the layout of an existing
+mbarrier object can be queried with `llvm.nvvm.mbarrier.check.layout.*`.
+
+#### '`llvm.nvvm.mbarrier.check.layout`'
+
+##### Syntax:
+
+```llvm
+declare i1 @llvm.nvvm.mbarrier.check.layout(ptr addrspace(3) %addr, i32 immarg %layout)
+```
+
+##### Overview:
+
+The '`@llvm.nvvm.mbarrier.check.layout.*`' intrinsics test whether the
+mbarrier object at `addr` was initialized with the layout named by
+`%layout`. They return `true` when the layout matches and `false`
+otherwise. `%layout` is an immediate argument that accepts only
+`0` (`layout::v0`) or `1` (`layout::v1`).
+
+##### Semantics:
+
+If the `addr` does not fall within the shared::cta space, then the behavior of
+this intrinsic is undefined. It is expected that `addr` was previously
+initialized using `mbarrier.init`; otherwise, the behavior is undefined.
+
 #### '`llvm.nvvm.mbarrier.inval`'
 
 ##### Syntax:
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index 4357ad367d269..80f0345c16850 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -2157,9 +2157,23 @@ foreach is_shared = [true, false] in {
   }
 }
 
+def int_nvvm_mbarrier_init_layout :
+  DefaultAttrsIntrinsic<[], [llvm_shared_ptr_ty,  // mbar
+                              llvm_i32_ty,        // count
+                              llvm_i32_ty],       // layout
+                        [IntrConvergent, IntrNoCallback, ImmArg<ArgIndex<2>>,
+                          Range<ArgIndex<2>, 0, 2>]>;
+
 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_shared_ptr_ty, // mbar
+                           llvm_i32_ty],       // layout
+                          [IntrReadMem, IntrNoCallback, ImmArg<ArgIndex<1>>,
+                           Range<ArgIndex<1>, 0, 2>]>;
+
 // mbarrier.{expect_tx/complete_tx}
 foreach op = ["expect_tx", "complete_tx"] in {
   foreach scope = ["scope_cta", "scope_cluster"] in {
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 02bcfcb6af926..cb3a1f85672fc 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -1282,6 +1282,19 @@ let Predicates = [SM80] in {
             [(set i32:$res, (int_nvvm_mbarrier_pending_count i64:$state))]>;
 }
 
+let Predicates = [PTX93, SM90] in {
+  def MBARRIER_INIT_LAYOUT :
+      NVPTXInst<(outs), (ins ADDR:$addr, B32:$count, B32:$layout),
+                "mbarrier.init.layout::v${layout}.shared.b64 [$addr], $count;",
+                [(int_nvvm_mbarrier_init_layout addr:$addr, i32:$count, i32:$layout)]>;
+
+  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/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll b/llvm/test/CodeGen/NVPTX/mbarrier_layout_sm90_ptx93.ll
new file mode 100644
index 0000000000000..3ed2a263f810d
--- /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.layout::v0.shared.b64 [%rd1], %r1;
+; CHECK-NEXT:    mbarrier.init.layout::v1.shared.b64 [%rd1], %r1;
+; CHECK-NEXT:    ret;
+  tail call void @llvm.nvvm.mbarrier.init.layout(ptr addrspace(3) %a, i32 %b, i32 0)
+  tail call void @llvm.nvvm.mbarrier.init.layout(ptr addrspace(3) %a, i32 %b, i32 1)
+  ret void
+}
+
+define i1 @mbarrier_check_layout(ptr addrspace(3) %a) {
+; CHECK-LABEL: mbarrier_check_layout(
+; CHECK:       {
+; CHECK-NEXT:    .reg .pred %p<4>;
+; CHECK-NEXT:    .reg .b32 %r<2>;
+; CHECK-NEXT:    .reg .b64 %rd<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param::func.b64 %rd1, [mbarrier_check_layout_param_0];
+; CHECK-NEXT:    mbarrier.check_layout.layout::v0.shared.b64 %p1, [%rd1];
+; CHECK-NEXT:    mbarrier.check_layout.layout::v1.shared.b64 %p2, [%rd1];
+; CHECK-NEXT:    or.pred %p3, %p1, %p2;
+; CHECK-NEXT:    selp.b32 %r1, -1, 0, %p3;
+; CHECK-NEXT:    st.param::func.b32 [func_retval0], %r1;
+; CHECK-NEXT:    ret;
+  %is_layout_v0 = tail call i1 @llvm.nvvm.mbarrier.check.layout(ptr addrspace(3) %a, i32 0)
+  %is_layout_v1 = tail call i1 @llvm.nvvm.mbarrier.check.layout(ptr addrspace(3) %a, i32 1)
+  %ret = or i1 %is_layout_v0, %is_layout_v1
+  ret i1 %ret
+}
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 6241606122f43..95508013def8f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -635,7 +635,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,
+                 OptionalAttr<ConfinedAttr<I32Attr,
+                   [IntMinValue<0>, IntMaxValue<1>]>>:$layout,
+                 PtxPredicate:$predicate)> {
   let summary = "MBarrier Initialization Op";
   let description = [{
     The `nvvm.mbarrier.init` operation initializes an *mbarrier object* at the specified 
@@ -654,6 +657,10 @@ 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 absent, the default layout is used. 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)
@@ -698,6 +705,33 @@ 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 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 b155505a2dae1..3b44d2a8ccd67 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -846,8 +846,8 @@ 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,
-                                                      adaptor.getPredicate());
+    rewriter.replaceOpWithNewOp<NVVM::MBarrierInitOp>(
+        op, barrier, count, /*layout=*/nullptr, adaptor.getPredicate());
     return success();
   }
 };
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 3cb03e297d03c..160c237fa2add 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3398,9 +3398,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() {
@@ -3727,16 +3730,23 @@ PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
 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;
+  llvm::Value *mbar = mt.lookupValue(thisOp.getAddr());
+  llvm::Value *count = mt.lookupValue(thisOp.getCount());
 
-  // Fill the Intrinsic Args
-  llvm::SmallVector<llvm::Value *> args;
-  args.push_back(mt.lookupValue(thisOp.getAddr()));
-  args.push_back(mt.lookupValue(thisOp.getCount()));
+  // Only mbarrier.init.layout carries a layout, and it exists for shared::cta
+  // only, so a generic pointer has to be cast before it can be passed along.
+  if (std::optional<uint32_t> layout = thisOp.getLayout()) {
+    if (isPtrInGenericSpace(thisOp.getAddr()))
+      mbar = castPtrToAddrSpace(builder, mbar, NVVMMemorySpace::Shared);
 
-  return {id, std::move(args)};
+    return {llvm::Intrinsic::nvvm_mbarrier_init_layout,
+            {mbar, count, builder.getInt32(*layout)}};
+  }
+
+  llvm::Intrinsic::ID id = isPtrInSharedCTASpace(thisOp.getAddr())
+                               ? llvm::Intrinsic::nvvm_mbarrier_init_shared
+                               : llvm::Intrinsic::nvvm_mbarrier_init;
+  return {id, {mbar, count}};
 }
 
 mlir::NVVM::IDArgPair MBarrierInvalOp::getIntrinsicIDAndArgs(
@@ -3750,6 +3760,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 807681bf828d4..788a99c4b8a63 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, predicate = %pred {layout = 1 : i32} : !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..8c35f66dc3dbb
--- /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(ptr addrspace(3) %0, i32 0)
+  // CHECK-NEXT: %[[V1:.+]] = call i1 @llvm.nvvm.mbarrier.check.layout(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 {layout = 0 : i32} : !llvm.ptr<3> -> i1
+  %v1 = nvvm.mbarrier.check_layout %barrier {layout = 1 : i32} : !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 9c1d1cc0cdc31..aec5029ac38fd 100644
--- a/mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/mbar_init.mlir
@@ -37,6 +37,28 @@ 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.layout(ptr addrspace(3) %0, i32 %1, i32 0)
+  // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.layout(ptr addrspace(3) %0, i32 %1, i32 1)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.mbarrier.init %barrier, %count {layout = 0 : i32} : !llvm.ptr<3>, i32
+  nvvm.mbarrier.init %barrier, %count {layout = 1 : i32} : !llvm.ptr<3>, i32
+  llvm.return
+}
+
+// mbarrier.init.layout is shared::cta only, so a generic pointer is cast.
+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: %[[CAST:.+]] = addrspacecast ptr %0 to ptr addrspace(3)
+  // CHECK-NEXT: call void @llvm.nvvm.mbarrier.init.layout(ptr addrspace(3) %[[CAST]], i32 %1, i32 1)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.mbarrier.init %barrier, %count {layout = 1 : i32} : !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 4a7776d86b28e..6a9bc87239e15 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 : i32} : !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 : i32} : !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 : i32} : !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 : i32} : !llvm.ptr<3> -> i1
+  llvm.return
+}
+



More information about the Mlir-commits mailing list