[Mlir-commits] [mlir] [mlir][NVVM] Add memory clobber support to inline_ptx and BasicPtxBui… (PR #212937)
Siyi Xu
llvmlistbot at llvm.org
Thu Jul 30 01:07:41 PDT 2026
https://github.com/siyix-ace updated https://github.com/llvm/llvm-project/pull/212937
>From a1425ff16741427686eb54adf7b1636e2a12422c Mon Sep 17 00:00:00 2001
From: Siyi Xu <siyix at nvidia.com>
Date: Wed, 29 Jul 2026 22:26:17 -0700
Subject: [PATCH 1/2] [mlir][NVVM] Add memory clobber support to inline_ptx and
BasicPtxBuilder
PTX with memory side effects (stores, atomics, mbarrier operations with
acquire/release semantics) emitted through nvvm.inline_ptx lowers to
llvm.inline_asm with register constraints only. Without a ~{memory}
clobber, LLVM may reorder memory accesses across the inline assembly.
There was no way to express the clobber through this op or the
BasicPtxBuilderInterface machinery.
Add a hasMemoryClobber interface method (default false) to
BasicPtxBuilderOpInterface; when it returns true, PtxBuilder appends
~{memory} to the constraints of the generated inline assembly. Expose
it on nvvm.inline_ptx as an opt-in memory_clobber unit attribute.
---
.../LLVMIR/BasicPtxBuilderInterface.td | 15 ++++++++-
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 31 ++++++++++-------
.../LLVMIR/IR/BasicPtxBuilderInterface.cpp | 7 ++++
.../Conversion/NVVMToLLVM/nvvm-to-llvm.mlir | 33 +++++++++++++++++++
4 files changed, 74 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td b/mlir/include/mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td
index 6ac4cca01a588..2834ae160595d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td
@@ -105,7 +105,20 @@ def BasicPtxBuilderOpInterface : OpInterface<"BasicPtxBuilderInterface"> {
/*methodBody=*/"",
/*defaultImplementation=*/"return true;"
>,
-
+ InterfaceMethod<
+ /*desc=*/[{
+ Return whether a memory clobber ("~{memory}") is appended to the
+ constraints of the generated inline assembly. It informs LLVM that
+ the PTX may read or write memory beyond its listed operands, so
+ optimizations must not reorder memory accesses across it.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"hasMemoryClobber",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/"return false;"
+ >,
+
InterfaceMethod<
/*desc=*/[{Helper function to generate i32 constant value.}],
/*retType=*/"::mlir::Value",
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 1f76218936e16..e9680703b808c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -383,16 +383,22 @@ def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx",
`BasicPtxBuilderInterface` to abstract away low-level details of
PTX assembly formatting.
- The `predicate` attribute is used to specify a predicate for the
+ The `predicate` attribute is used to specify a predicate for the
PTX instruction.
+ The `memory_clobber` attribute appends a "~{memory}" clobber to the
+ constraints of the generated inline assembly. Set it when the PTX reads
+ or writes memory beyond its listed operands (e.g. stores, atomics, or
+ instructions with acquire/release semantics such as mbarrier), so that
+ LLVM does not reorder memory accesses across the inline assembly.
+
Example 1: Read-only Parameters
```mlir
- nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count) : !llvm.ptr, i32
+ nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" ro(%barrier_gen, %count : !llvm.ptr, i32) memory_clobber
// Lowers to:
- llvm.inline_asm has_side_effects asm_dialect = att
- "mbarrier.init.b64 [$0], $1;", "l,r" %arg0, %arg2 : (!llvm.ptr, i32) -> ()
+ llvm.inline_asm has_side_effects asm_dialect = att
+ "mbarrier.init.b64 [$0], $1;", "l,r,~{memory}" %arg0, %arg2 : (!llvm.ptr, i32) -> ()
```
Example 2: Read-only and Write-only Parameters
@@ -416,22 +422,24 @@ def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx",
```
}];
- let arguments = (ins Variadic<AnyType>:$readOnlyArgs,
- Variadic<AnyType>:$readWriteArgs,
+ let arguments = (ins Variadic<AnyType>:$readOnlyArgs,
+ Variadic<AnyType>:$readWriteArgs,
StrAttr:$ptxCode,
- PtxPredicate:$predicate);
-
+ PtxPredicate:$predicate,
+ UnitAttr:$memoryClobber);
+
let results = (outs Variadic<AnyType>:$writeOnlyArgs);
-
+
let assemblyFormat = [{
$ptxCode
( `ro` `(` $readOnlyArgs^ `:` type($readOnlyArgs) `)` )?
( `rw` `(` $readWriteArgs^ `:` type($readWriteArgs) `)` )?
- (`,` `predicate` `=` $predicate^)?
+ ( `memory_clobber` $memoryClobber^ )?
+ (`,` `predicate` `=` $predicate^)?
attr-dict
( `->` type($writeOnlyArgs)^ )?
}];
-
+
let extraClassDefinition = [{
std::string $cppClass::getPtx() {
StringRef ptxInstStr = getPtxCode();
@@ -441,6 +449,7 @@ def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx",
let extraClassDeclaration = [{
bool getAsmValues(RewriterBase &, llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &);
+ bool hasMemoryClobber() { return getMemoryClobber(); }
}];
}
diff --git a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
index 87206ba726765..a8767611adb64 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
@@ -474,6 +474,13 @@ LLVM::InlineAsmOp PtxBuilder::build() {
ptxInstruction = predicateStr + " " + ptxInstruction;
}
+ // Clobbers go after all register constraints.
+ if (interfaceOp.hasMemoryClobber()) {
+ if (!registerConstraints.empty())
+ registerConstraints += ",";
+ registerConstraints += "~{memory}";
+ }
+
// Operand placeholders are written as %0, %1, ... (and the predicate as
// @%N), because TableGen string attributes cannot contain '$', which inline
// assembly uses for operand substitution. Convert only a '%' that is
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index 1886e72d72f41..0b925f9cf679f 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -597,6 +597,25 @@ llvm.func @init_mbarrier(
}
// -----
+llvm.func @init_mbarrier_memory_clobber(
+ %barrier_gen : !llvm.ptr,
+ %count : i32,
+ %pred : i1) {
+ // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.init.b64 [$0], $1;", "l,r,~{memory}"
+ nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32) memory_clobber
+ // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b,~{memory}"
+ nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32) memory_clobber, predicate = %pred
+ llvm.return
+}
+// -----
+
+llvm.func @memory_clobber_no_operands() {
+ // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.sc.cta;", "~{memory}"
+ nvvm.inline_ptx "fence.sc.cta;" memory_clobber
+ llvm.return
+}
+// -----
+
llvm.func @ex2(%input : f32, %pred : i1) {
// CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "ex2.approx.ftz.f32 $0, $1;", "=f,f" %{{.*}} : (f32) -> f32
%0 = nvvm.inline_ptx "ex2.approx.ftz.f32 {$w0}, {$r0};" ro (%input : f32) -> f32
@@ -670,6 +689,20 @@ llvm.func @inline_ptx_multi_rw_pred(%a : i32, %b : i32, %rw_c : f32, %rw_d : f32
llvm.return %r4 : f32
}
+// CHECK-LABEL: @inline_ptx_multi_rw_memory_clobber(
+// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: i32, %[[arg1:[a-zA-Z0-9_]+]]: i32, %[[arg2:[a-zA-Z0-9_]+]]: f32, %[[arg3:[a-zA-Z0-9_]+]]: f32)
+llvm.func @inline_ptx_multi_rw_memory_clobber(%a : i32, %b : i32, %rw_c : f32, %rw_d : f32) -> f32 {
+// CHECK: %[[S0:.+]] = llvm.inline_asm has_side_effects asm_dialect = att "{.reg .pred p; setp.ge.s32 p, $2, $3; selp.s32 $0, $2,$3, p; selp.s32 $1, $2,$3, p;}",
+// CHECK-SAME: "=f,=f,r,r,0,1,~{memory}"
+// CHECK-SAME: %[[arg2]], %[[arg3]], %[[arg0]], %[[arg1]]
+// CHECK-SAME: : (f32, f32, i32, i32) -> !llvm.struct<(f32, f32)>
+ nvvm.inline_ptx "{.reg .pred p; setp.ge.s32 p, {$r0}, {$r1}; selp.s32 {$rw0}, {$r0},{$r1}, p; selp.s32 {$rw1}, {$r0},{$r1}, p;}"
+ ro (%a, %b : i32,i32)
+ rw (%rw_c, %rw_d: f32,f32) memory_clobber
+ %r4 = llvm.fadd %rw_c, %rw_d : f32
+ llvm.return %r4 : f32
+}
+
// CHECK-LABEL: @inline_ptx_multi_rw_r(
// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: i32, %[[arg1:[a-zA-Z0-9_]+]]: i32, %[[arg2:[a-zA-Z0-9_]+]]: f32, %[[arg3:[a-zA-Z0-9_]+]]: f32)
llvm.func @inline_ptx_multi_rw_r(%a : i32, %b : i32, %rw_c : f32, %rw_d : f32) -> f32 {
>From 12d8f7335a778cc6758cc64fd2271333ad63bebb Mon Sep 17 00:00:00 2001
From: Siyi Xu <siyix at nvidia.com>
Date: Thu, 30 Jul 2026 01:06:39 -0700
Subject: [PATCH 2/2] [mlir][NVVM] Use DefaultValuedAttr<BoolAttr> for
memory_clobber
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 6 +++---
mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir | 8 ++++----
2 files changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index e9680703b808c..75bc7527aba52 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -394,7 +394,7 @@ def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx",
Example 1: Read-only Parameters
```mlir
- nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" ro(%barrier_gen, %count : !llvm.ptr, i32) memory_clobber
+ nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" ro(%barrier_gen, %count : !llvm.ptr, i32) memory_clobber = true
// Lowers to:
llvm.inline_asm has_side_effects asm_dialect = att
@@ -426,7 +426,7 @@ def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx",
Variadic<AnyType>:$readWriteArgs,
StrAttr:$ptxCode,
PtxPredicate:$predicate,
- UnitAttr:$memoryClobber);
+ DefaultValuedAttr<BoolAttr, "false">:$memoryClobber);
let results = (outs Variadic<AnyType>:$writeOnlyArgs);
@@ -434,7 +434,7 @@ def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx",
$ptxCode
( `ro` `(` $readOnlyArgs^ `:` type($readOnlyArgs) `)` )?
( `rw` `(` $readWriteArgs^ `:` type($readWriteArgs) `)` )?
- ( `memory_clobber` $memoryClobber^ )?
+ ( `memory_clobber` `=` $memoryClobber^ )?
(`,` `predicate` `=` $predicate^)?
attr-dict
( `->` type($writeOnlyArgs)^ )?
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index 0b925f9cf679f..807681bf828d4 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -602,16 +602,16 @@ llvm.func @init_mbarrier_memory_clobber(
%count : i32,
%pred : i1) {
// CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.init.b64 [$0], $1;", "l,r,~{memory}"
- nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32) memory_clobber
+ nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32) memory_clobber = true
// CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b,~{memory}"
- nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32) memory_clobber, predicate = %pred
+ nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32) memory_clobber = true, predicate = %pred
llvm.return
}
// -----
llvm.func @memory_clobber_no_operands() {
// CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.sc.cta;", "~{memory}"
- nvvm.inline_ptx "fence.sc.cta;" memory_clobber
+ nvvm.inline_ptx "fence.sc.cta;" memory_clobber = true
llvm.return
}
// -----
@@ -698,7 +698,7 @@ llvm.func @inline_ptx_multi_rw_memory_clobber(%a : i32, %b : i32, %rw_c : f32,
// CHECK-SAME: : (f32, f32, i32, i32) -> !llvm.struct<(f32, f32)>
nvvm.inline_ptx "{.reg .pred p; setp.ge.s32 p, {$r0}, {$r1}; selp.s32 {$rw0}, {$r0},{$r1}, p; selp.s32 {$rw1}, {$r0},{$r1}, p;}"
ro (%a, %b : i32,i32)
- rw (%rw_c, %rw_d: f32,f32) memory_clobber
+ rw (%rw_c, %rw_d: f32,f32) memory_clobber = true
%r4 = llvm.fadd %rw_c, %rw_d : f32
llvm.return %r4 : f32
}
More information about the Mlir-commits
mailing list