[Mlir-commits] [mlir] [mlir][NVVM] Add nvvm.membar operation (PR #166698)
Valentin Clement バレンタイン クレメン
llvmlistbot at llvm.org
Thu Nov 6 11:32:48 PST 2025
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/3] [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/3] 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.
>From 3a0da14fe9194d86f26d10976b653c3937dfe66e Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 6 Nov 2025 11:32:29 -0800
Subject: [PATCH 3/3] Address review comments
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 24 ++++---------------
.../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 13 ++++++----
mlir/test/Target/LLVMIR/nvvm/membar.mlir | 14 +++++++++++
mlir/test/Target/LLVMIR/nvvmir.mlir | 13 ----------
4 files changed, 26 insertions(+), 38 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/nvvm/membar.mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index dea15629a3900..f21f80bd6dce5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1236,24 +1236,8 @@ 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)> {
+def NVVM_MembarOp : NVVM_Op<"memory_barrier">,
+ Arguments<(ins MemScopeKindAttr:$scope)> {
let summary = "Memory barrier operation";
let description = [{
`membar` operation guarantees that prior memory accesses requested by this
@@ -1263,9 +1247,9 @@ def NVVM_MembarOp : NVVM_Op<"membar">,
[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 assemblyFormat = "$scope attr-dict";
let llvmBuilder = [{
- createIntrinsicCall(builder, getMembarLevelID($level), {});
+ createIntrinsicCall(builder, getMemoryBarrierLevelID($scope), {});
}];
}
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index 9d6ccd90b2060..df0bbe0e2b866 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -291,15 +291,18 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
llvm_unreachable("Unsupported proxy kinds");
}
-static unsigned getMembarLevelID(NVVM::MemLevelKind level) {
- switch (level) {
- case NVVM::MemLevelKind::CTA: {
+static unsigned getMemoryBarrierLevelID(NVVM::MemScopeKind scope) {
+ switch (scope) {
+ case NVVM::MemScopeKind::CTA: {
return llvm::Intrinsic::nvvm_membar_cta;
}
- case NVVM::MemLevelKind::GL: {
+ case NVVM::MemScopeKind::CLUSTER: {
+ return llvm::Intrinsic::nvvm_fence_sc_cluster;
+ }
+ case NVVM::MemScopeKind::GPU: {
return llvm::Intrinsic::nvvm_membar_gl;
}
- case NVVM::MemLevelKind::SYS: {
+ case NVVM::MemScopeKind::SYS: {
return llvm::Intrinsic::nvvm_membar_sys;
}
}
diff --git a/mlir/test/Target/LLVMIR/nvvm/membar.mlir b/mlir/test/Target/LLVMIR/nvvm/membar.mlir
new file mode 100644
index 0000000000000..edaeaccdc57c2
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/membar.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-translate -mlir-to-llvmir %s -split-input-file --verify-diagnostics | FileCheck %s
+
+// CHECK-lABEL: @memorybarrier()
+llvm.func @memorybarrier() {
+ // CHECK: call void @llvm.nvvm.membar.cta()
+ nvvm.memory_barrier #nvvm.mem_scope<cta>
+ // CHECK: call void @llvm.nvvm.fence.sc.cluster()
+ nvvm.memory_barrier #nvvm.mem_scope<cluster>
+ // CHECK: call void @llvm.nvvm.membar.gl()
+ nvvm.memory_barrier #nvvm.mem_scope<gpu>
+ // CHECK: call void @llvm.nvvm.membar.sys()
+ nvvm.memory_barrier #nvvm.mem_scope<sys>
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 04b2d791188c1..1ec55408e97a5 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -975,16 +975,3 @@ 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
-}
More information about the Mlir-commits
mailing list