[Mlir-commits] [mlir] [MLIR]NVVM] Add `inline_ptx` op (PR #139923)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 14 08:35:02 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm
Author: Guray Ozen (grypp)
<details>
<summary>Changes</summary>
This op allows using PTX directly within the NVVM dialect, while greatly simplifying llvm.inline_asm generation.
**Example 1: Read-only Parameters**
Sets `"l,r"` automatically.
```
nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count) : !llvm.ptr, i32
// Lowers to:
llvm.inline_asm has_side_effects asm_dialect = att
"mbarrier.init.b64 [$0], $1;", "l,r" %arg0, %arg2 : (!llvm.ptr, i32) -> ()
```
**Example 2: Read-only and Write-only Parameters**
Sets `=f,f"` automatically. `=` is set because there is store.
```
%0 = nvvm.inline_ptx "ex2.approx.ftz.f32 $0, $1;" (%input) : f32 -> f32
// Lowers to:
%0 = llvm.inline_asm has_side_effects asm_dialect = att "ex2.approx.ftz.f32 $0, $1;", "=f,f" %arg0 : (f32) -> f32
```
**Example 3: Predicate Usage**
Now `@$2` is set automatically for predication.
```
nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count), predicate = %pred : !llvm.ptr, i32, i1
// Lowers to:
llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b" %arg0, %arg2, %arg3 : (!llvm.ptr, i32, i1) -> ()
```
---
Full diff: https://github.com/llvm/llvm-project/pull/139923.diff
2 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+70)
- (modified) mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir (+25)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 654aff71f25be..4ba54fa3c1ca7 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -236,6 +236,76 @@ foreach index = !range(0, 32) in {
def NVVM_EnvReg # index # Op : NVVM_SpecialRegisterOp<"read.ptx.sreg.envreg" # index>;
}
+//===----------------------------------------------------------------------===//
+// Inline PTX op definition
+//===----------------------------------------------------------------------===//
+
+def NVVM_InlinePtxOP : NVVM_Op<"inline_ptx",
+ [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>,
+ AttrSizedOperandSegments]>
+{
+ let summary = "Inline PTX Op";
+ let description = [{This op allows using PTX directly within the NVVM
+ dialect, while greatly simplifying llvm.inline_asm generation. It
+ automatically handles register size selection and sets the correct
+ read/write access for each operand. The operation leverages the
+ `BasicPtxBuilderInterface` to abstract away low-level details of
+ PTX assembly formatting.
+
+ The `predicate` attribute is used to specify a predicate for the
+ PTX instruction.
+
+ Example 1: Read-only Parameters
+ ```mlir
+ nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count) : !llvm.ptr, i32
+
+ // Lowers to:
+ llvm.inline_asm has_side_effects asm_dialect = att
+ "mbarrier.init.b64 [$0], $1;", "l,r" %arg0, %arg2 : (!llvm.ptr, i32) -> ()
+ ```
+
+ Example 2: Read-only and Write-only Parameters
+ ```mlir
+ %0 = nvvm.inline_ptx "ex2.approx.ftz.f32 $0, $1;" (%input) : f32 -> f32
+
+ // Lowers to:
+ %0 = llvm.inline_asm has_side_effects asm_dialect = att
+ "ex2.approx.ftz.f32 $0, $1;", "=f,f" %arg0 : (f32) -> f32
+ ```
+
+ Example 3: Predicate Usage
+ ```mlir
+ nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count),
+ predicate = %pred : !llvm.ptr, i32, i1
+
+ // Lowers to:
+ llvm.inline_asm has_side_effects asm_dialect = att
+ "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b" %arg0, %arg2, %arg3
+ : (!llvm.ptr, i32, i1) -> ()
+ ```
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$readOnlyArgs,
+ StrAttr:$ptxCode,
+ PtxPredicate:$predicate);
+
+ let results = (outs Variadic<AnyType>:$writeOnlyArgs);
+
+ let assemblyFormat = [{
+ $ptxCode `(` $readOnlyArgs `)`
+ (`,` `predicate` `=` $predicate^)? attr-dict
+ `:` type(operands)
+ (`->` type($writeOnlyArgs)^)?
+ }];
+
+ let extraClassDefinition = [{
+ std::string $cppClass::getPtx() {
+ StringRef eventName = getPtxCode();
+ return std::string(eventName.data());
+ }
+ }];
+}
+
//===----------------------------------------------------------------------===//
// NVVM approximate op definitions
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index c7a6eca158276..1d9164ac94d76 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -680,3 +680,28 @@ llvm.func @llvm_nvvm_barrier_arrive(%barID : i32, %numberOfThreads : i32) {
nvvm.barrier.arrive id = %barID number_of_threads = %numberOfThreads
llvm.return
}
+
+
+// -----
+
+llvm.func @init_mbarrier(
+ %barrier_gen : !llvm.ptr,
+ %barrier : !llvm.ptr<3>,
+ %count : i32,
+ %pred : i1) {
+ // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.init.b64 [$0], $1;", "l,r"
+ nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count) : !llvm.ptr, i32
+ // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b"
+ nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count), predicate = %pred : !llvm.ptr, i32, i1
+ 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 $0, $1;" (%input) : f32 -> f32
+
+ // CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "@$1 ex2.approx.ftz.f32 $0, $1;", "=f,f,b" %{{.*}}, %{{.*}} : (f32, i1) -> f32
+ %1 = nvvm.inline_ptx "ex2.approx.ftz.f32 $0, $1;" (%input), predicate = %pred : f32, i1 -> f32
+ llvm.return
+}
\ No newline at end of file
``````````
</details>
https://github.com/llvm/llvm-project/pull/139923
More information about the Mlir-commits
mailing list