[MLIR][NVVM] Add `inline_ptx` op #139923

grypp · 2025-05-14T15:34:25Z

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) -> ()

This op allows using PTX directly within the NVVM dialect, while greatly simplifying llvm.inline_asm generation. 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) -> () ```

llvmbot · 2025-05-14T15:34:56Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Guray Ozen (grypp)

Changes

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) -&gt; ()

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 -&gt; f32

// Lowers to:
%0 = llvm.inline_asm has_side_effects asm_dialect = att "ex2.approx.ftz.f32 $0, $1;", "=f,f" %arg0 : (f32) -&gt; 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) -&gt; ()

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 --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

Copilot

Pull Request Overview

Adds a new nvvm.inline_ptx operation that embeds PTX assembly directly in the NVVM dialect and lowers it to llvm.inline_asm.

Introduces the inline_ptx op in the Dialect definition (NVVMOps.td).
Provides C++ hook (getPtx()) to retrieve the PTX string.
Adds FileCheck-based tests in nvvm-to-llvm.mlir to verify correct lowering.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td	Define `NVVM_InlinePtxOP` and its builder
mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir	Add tests for `nvvm.inline_ptx` lowering

Comments suppressed due to low confidence (1)

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td:243

The op class name NVVM_InlinePtxOP uses uppercase "OP". Per MLIR naming conventions, it should be NVVM_InlinePtxOp (capital ‘O’, lowercase ‘p’).

def NVVM_InlinePtxOP : NVVM_Op<"inline_ptx",

mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir

joker-eph · 2025-05-14T20:12:08Z

The AI bot review seems spot on ;)

joker-eph · 2025-05-14T20:12:53Z

So how does it knows about "Read-only Parameters" and "Read-only and Write-only Parameters" ?

grypp · 2025-05-15T06:32:11Z

So how does it knows about "Read-only Parameters" and "Read-only and Write-only Parameters" ?

Input arguments are read-only, and results are write-only — it's pure SSA. For read-only arguments, no special marker is used, but for write-only ones, we automatically prepend =.

There is also ReadWrite symbols. I support them in NVVM dialect, for example wgmma.mma_async. But I didn't do this for this PR.

It's little bit complicated to support ReadWrite in LLVM. They are supported by +, but this isn't supported by LLVM, I guess because it doesn't align with SSA semantics.

In LLVM, we handle read-write symbols differently. We mark them as =, indicating write access, and then explicitly map the same symbol as a read operand in the input list.

For example here result is readwrite, and valid asm. But llvm doesn't support that.

asm ("OPCODE %0, %1, %2" : "+f" (result));

Instead, we generate the equivalent as:

asm ("OPCODE %0, %1, %2" : "=f" (result) : "0" (result));

This marks result as both written (=f) and read ("0").

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

Co-authored-by: Mehdi Amini <[email protected]>

durga4github · 2025-05-15T11:41:24Z

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

+
+    // Lowers to:
+    llvm.inline_asm has_side_effects asm_dialect = att 
+      "mbarrier.init.b64 [$0], $1;", "l,r" %arg0, %arg2 : (!llvm.ptr, i32) -> ()


nit:
In all these given examples, the mlir looks clean.
The lowered llvm-ir still seems to have llvm.ptr, i32 etc. I think this should reflect what we generate, right?
(may be we could copy-paste from the unit-test below).

In all these given examples, the mlir looks clean.

What do you mean by clean?

In all these given examples, the mlir looks clean.

What do you mean by clean?

I meant the mlir version exactly matches the unit-test but the lowered version still refers to (!llvm.ptr, i32) instead of "l, r"

durga4github

LGTM

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) -> () ``` --------- Co-authored-by: Mehdi Amini <[email protected]>

grypp requested review from Copilot and durga4github May 14, 2025 15:34

llvmbot added mlir:llvm mlir labels May 14, 2025

grypp changed the title ~~[MLIR]NVVM] Add inline_ptx op~~ [MLIR][NVVM] Add inline_ptx op May 14, 2025

a

65dc6f2

Copilot AI reviewed May 14, 2025

View reviewed changes

mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir Show resolved Hide resolved

grypp requested review from joker-eph and schwarzschild-radius May 14, 2025 15:43

joker-eph reviewed May 15, 2025

View reviewed changes

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td Outdated Show resolved Hide resolved

joker-eph approved these changes May 15, 2025

View reviewed changes

durga4github reviewed May 15, 2025

View reviewed changes

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td Show resolved Hide resolved

durga4github reviewed May 15, 2025

View reviewed changes

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td Outdated Show resolved Hide resolved

grypp and others added 2 commits May 15, 2025 13:24

Update mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

5798854

Co-authored-by: Mehdi Amini <[email protected]>

Update NVVMOps.td

7163436

durga4github reviewed May 15, 2025

View reviewed changes

durga4github approved these changes May 15, 2025

View reviewed changes

grypp merged commit d08b176 into llvm:main May 15, 2025
11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][NVVM] Add `inline_ptx` op #139923

[MLIR][NVVM] Add `inline_ptx` op #139923

Uh oh!

grypp commented May 14, 2025

Uh oh!

llvmbot commented May 14, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

joker-eph commented May 14, 2025

Uh oh!

joker-eph commented May 14, 2025

Uh oh!

grypp commented May 15, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

durga4github May 15, 2025

Uh oh!

grypp May 15, 2025 •

edited

Loading

Uh oh!

durga4github May 15, 2025

Uh oh!

durga4github left a comment

Uh oh!

Uh oh!

Uh oh!

[MLIR][NVVM] Add inline_ptx op #139923

[MLIR][NVVM] Add inline_ptx op #139923

Uh oh!

Conversation

grypp commented May 14, 2025

Uh oh!

llvmbot commented May 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

joker-eph commented May 14, 2025

Uh oh!

joker-eph commented May 14, 2025

Uh oh!

grypp commented May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

durga4github May 15, 2025

Choose a reason for hiding this comment

Uh oh!

grypp May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

durga4github May 15, 2025

Choose a reason for hiding this comment

Uh oh!

durga4github left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[MLIR][NVVM] Add `inline_ptx` op #139923

[MLIR][NVVM] Add `inline_ptx` op #139923

llvmbot commented May 14, 2025 •

edited

Loading

grypp commented May 15, 2025 •

edited

Loading

grypp May 15, 2025 •

edited

Loading