[Mlir-commits] [mlir] [mlir][NVVM] Add nvvm.membar operation (PR #166698)
Valentin Clement バレンタイン クレメン
llvmlistbot at llvm.org
Fri Nov 7 10:21:17 PST 2025
https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/166698
>From 02ae9ea5840313910a8ccd9deffedad5ea533f70 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/6] [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 0308d0d5e22a5c6b780cea940b184602be89d06a 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/6] 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 4f3be47365018b2571553d74cfec49dcc8a0bf92 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/6] 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
-}
>From 00a45d2cb024c1a957d48e1a5b1f44d7f94b8239 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 6 Nov 2025 11:33:53 -0800
Subject: [PATCH 4/6] Rename getMembarLevelID to getMembarIntrinsicID
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +-
mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index f21f80bd6dce5..407d8f8dbae87 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1249,7 +1249,7 @@ def NVVM_MembarOp : NVVM_Op<"memory_barrier">,
let assemblyFormat = "$scope attr-dict";
let llvmBuilder = [{
- createIntrinsicCall(builder, getMemoryBarrierLevelID($scope), {});
+ createIntrinsicCall(builder, getMembarIntrinsicID($scope), {});
}];
}
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index df0bbe0e2b866..7fdf4957a3c93 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -291,7 +291,7 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
llvm_unreachable("Unsupported proxy kinds");
}
-static unsigned getMemoryBarrierLevelID(NVVM::MemScopeKind scope) {
+static unsigned getMembarIntrinsicID(NVVM::MemScopeKind scope) {
switch (scope) {
case NVVM::MemScopeKind::CTA: {
return llvm::Intrinsic::nvvm_membar_cta;
>From f3d62f8fdabea1c52332235d3787ffc41fb6857a Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 7 Nov 2025 10:20:23 -0800
Subject: [PATCH 5/6] Address nits comments
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +-
.../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 14 +++++---------
mlir/test/Target/LLVMIR/nvvm/membar.mlir | 8 ++++----
3 files changed, 10 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 407d8f8dbae87..fda0e6bb3c01f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1236,7 +1236,7 @@ def NVVM_FenceProxyAcquireOp : NVVM_Op<"fence.proxy.acquire">,
let hasVerifier = 1;
}
-def NVVM_MembarOp : NVVM_Op<"memory_barrier">,
+def NVVM_MembarOp : NVVM_Op<"memory.barrier">,
Arguments<(ins MemScopeKindAttr:$scope)> {
let summary = "Memory barrier operation";
let description = [{
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index 7fdf4957a3c93..cecff51e637a5 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -293,20 +293,16 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
static unsigned getMembarIntrinsicID(NVVM::MemScopeKind scope) {
switch (scope) {
- case NVVM::MemScopeKind::CTA: {
+ case NVVM::MemScopeKind::CTA:
return llvm::Intrinsic::nvvm_membar_cta;
- }
- case NVVM::MemScopeKind::CLUSTER: {
+ case NVVM::MemScopeKind::CLUSTER:
return llvm::Intrinsic::nvvm_fence_sc_cluster;
- }
- case NVVM::MemScopeKind::GPU: {
+ case NVVM::MemScopeKind::GPU:
return llvm::Intrinsic::nvvm_membar_gl;
- }
- case NVVM::MemScopeKind::SYS: {
+ case NVVM::MemScopeKind::SYS:
return llvm::Intrinsic::nvvm_membar_sys;
}
- }
- llvm_unreachable("Unknown level for memory barrier");
+ llvm_unreachable("Unknown scope for memory barrier");
}
#define TCGEN05LD(SHAPE, NUM) llvm::Intrinsic::nvvm_tcgen05_ld_##SHAPE##_##NUM
diff --git a/mlir/test/Target/LLVMIR/nvvm/membar.mlir b/mlir/test/Target/LLVMIR/nvvm/membar.mlir
index edaeaccdc57c2..1b794f663b573 100644
--- a/mlir/test/Target/LLVMIR/nvvm/membar.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/membar.mlir
@@ -3,12 +3,12 @@
// CHECK-lABEL: @memorybarrier()
llvm.func @memorybarrier() {
// CHECK: call void @llvm.nvvm.membar.cta()
- nvvm.memory_barrier #nvvm.mem_scope<cta>
+ nvvm.memory.barrier #nvvm.mem_scope<cta>
// CHECK: call void @llvm.nvvm.fence.sc.cluster()
- nvvm.memory_barrier #nvvm.mem_scope<cluster>
+ nvvm.memory.barrier #nvvm.mem_scope<cluster>
// CHECK: call void @llvm.nvvm.membar.gl()
- nvvm.memory_barrier #nvvm.mem_scope<gpu>
+ nvvm.memory.barrier #nvvm.mem_scope<gpu>
// CHECK: call void @llvm.nvvm.membar.sys()
- nvvm.memory_barrier #nvvm.mem_scope<sys>
+ nvvm.memory.barrier #nvvm.mem_scope<sys>
llvm.return
}
>From 8aaf2d2a501e00d380630c32600d0be9a3e32b21 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 7 Nov 2025 10:20:59 -0800
Subject: [PATCH 6/6] 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 fda0e6bb3c01f..1cc5b74a3cb67 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1241,7 +1241,7 @@ def NVVM_MembarOp : NVVM_Op<"memory.barrier">,
let summary = "Memory barrier operation";
let description = [{
`membar` operation guarantees that prior memory accesses requested by this
- thread are performed at the specified `level`, before later memory
+ thread are performed at the specified `scope`, 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)
More information about the Mlir-commits
mailing list