[Mlir-commits] [mlir] [mlir][NVVM] Add nvvm.membar operation (PR #166698)

[Valentin Clement バレンタイン クレメン]

https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/166698

>From c80d759f2f924b0a7a48e2e14be83385f91b5d90 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 5 Nov 2025 21:07:30 -0800
Subject: [PATCH 1/2] [mlir][NVVM] Add nvvm.membar operation

---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td   | 33 +++++++++++++++++++
 .../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp  | 15 +++++++++
 mlir/test/Target/LLVMIR/nvvmir.mlir           | 13 ++++++++
 3 files changed, 61 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 80bc0e5986e51..f00aba15bfcae 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1236,6 +1236,39 @@ def NVVM_FenceProxyAcquireOp : NVVM_Op<"fence.proxy.acquire">,
   let hasVerifier = 1;
 }
 
+// Attrs describing the level of the Memory Operation
+def MemLevelCTA : I32EnumAttrCase<"CTA", 0, "cta">;
+def MemLevelGL : I32EnumAttrCase<"GL", 1, "gl">;
+def MemLevelSys : I32EnumAttrCase<"SYS", 2, "sys">;
+
+def MemLevelKind
+    : I32EnumAttr<
+          "MemLevelKind",
+          "NVVM Memory Level kind", [MemLevelCTA, MemLevelGL, MemLevelSys]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::NVVM";
+}
+def MemLevelKindAttr : EnumAttr<NVVM_Dialect, MemLevelKind, "mem_level"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
+def NVVM_MembarOp : NVVM_Op<"membar">,
+                    Arguments<(ins MemLevelKindAttr:$level)> {
+  let summary = "Memory barrier operation";
+  let description = [{
+    `member` operation guarantees that prior memory accesses requested by this
+    thread are performed at the specified `level`, before later memory
+    operations requested by this thread following the membar instruction.
+
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar)
+  }];
+
+  let assemblyFormat = "$level attr-dict";
+  let llvmBuilder = [{
+    createIntrinsicCall(builder, getMembarLevelID($level), {});
+  }];
+}
+
 def NVVM_FenceProxyReleaseOp : NVVM_Op<"fence.proxy.release">,
       Arguments<(ins MemScopeKindAttr:$scope,
                      DefaultValuedAttr<ProxyKindAttr,
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index 0964e1b8c5ef3..9d6ccd90b2060 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -291,6 +291,21 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
   llvm_unreachable("Unsupported proxy kinds");
 }
 
+static unsigned getMembarLevelID(NVVM::MemLevelKind level) {
+  switch (level) {
+  case NVVM::MemLevelKind::CTA: {
+    return llvm::Intrinsic::nvvm_membar_cta;
+  }
+  case NVVM::MemLevelKind::GL: {
+    return llvm::Intrinsic::nvvm_membar_gl;
+  }
+  case NVVM::MemLevelKind::SYS: {
+    return llvm::Intrinsic::nvvm_membar_sys;
+  }
+  }
+  llvm_unreachable("Unknown level for memory barrier");
+}
+
 #define TCGEN05LD(SHAPE, NUM) llvm::Intrinsic::nvvm_tcgen05_ld_##SHAPE##_##NUM
 
 static llvm::Intrinsic::ID
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 1ec55408e97a5..04b2d791188c1 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -975,3 +975,16 @@ llvm.func @nanosleep() {
   nvvm.nanosleep 4000
   llvm.return
 }
+
+// -----
+
+// CHECK-lABEL: @memorybarrier()
+llvm.func @memorybarrier() {
+  // CHECK: call void @llvm.nvvm.membar.cta()
+  nvvm.membar #nvvm.mem_level<cta>
+  // CHECK: call void @llvm.nvvm.membar.gl()
+  nvvm.membar #nvvm.mem_level<gl>
+  // CHECK: call void @llvm.nvvm.membar.sys()
+  nvvm.membar #nvvm.mem_level<sys>
+  llvm.return
+}

>From f21b256460d97fda1cbebc48a7a1beb0e954ce8f Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 5 Nov 2025 21:15:39 -0800
Subject: [PATCH 2/2] Fix typo

---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index f00aba15bfcae..dea15629a3900 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1256,7 +1256,7 @@ def NVVM_MembarOp : NVVM_Op<"membar">,
                     Arguments<(ins MemLevelKindAttr:$level)> {
   let summary = "Memory barrier operation";
   let description = [{
-    `member` operation guarantees that prior memory accesses requested by this
+    `membar` operation guarantees that prior memory accesses requested by this
     thread are performed at the specified `level`, before later memory
     operations requested by this thread following the membar instruction.