[Mlir-commits] [mlir] 9b734e2 - [mlir][NVVM] Add memory clobber support to inline_ptx and BasicPtxBui… (#212937)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 31 01:38:17 PDT 2026
Author: Siyi Xu
Date: 2026-07-31T10:38:13+02:00
New Revision: 9b734e25473bc9e299a4b58575b9b0e618925b78
URL: https://github.com/llvm/llvm-project/commit/9b734e25473bc9e299a4b58575b9b0e618925b78
DIFF: https://github.com/llvm/llvm-project/commit/9b734e25473bc9e299a4b58575b9b0e618925b78.diff
LOG: [mlir][NVVM] Add memory clobber support to inline_ptx and BasicPtxBui… (#212937)
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.
This patch adds, as discussed with @grypp:
- A `hasMemoryClobber` interface method (default `false`) on
`BasicPtxBuilderOpInterface`; when it returns `true`, `PtxBuilder`
appends `~{memory}` to the constraints of the generated inline assembly,
after all register constraints and tied indices.
- An opt-in `memory_clobber` boolean attribute (default `false`) on
`nvvm.inline_ptx` exposing this. Existing behavior is unchanged unless
the attribute is set to `true`.
For example:
```mlir
nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" ro(%p, %c : !llvm.ptr, i32) memory_clobber = true
```
now lowers to inline assembly with constraints `"l,r,~{memory}"`.
Tests: constraint emission with/without the attribute, combination with
predicate (`"l,r,b,~{memory}"`), read-write operands with tied indices
(`"=f,=f,r,r,0,1,~{memory}"`), and the no-operand case (`"~{memory}"`).
The op documentation's first example (mbarrier.init) is updated to show
correct usage.
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
Removed:
################################################################################
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..75bc7527aba52 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 = true
// 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,
+ DefaultValuedAttr<BoolAttr, "false">:$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..807681bf828d4 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 = 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 = 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 = true
+ 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 = true
+ %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 {
More information about the Mlir-commits
mailing list