[Mlir-commits] [mlir] [MLIR][NVVM] Fix `inline_ptx` when there is only clobber register (PR #194380)

[Guray Ozen]

https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/194380

>From 5fda897f6b3b03e1b723fb71b4e680cc58f4bb2e Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Mon, 27 Apr 2026 15:56:47 +0200
Subject: [PATCH 1/2] [MLIR][NVVM] Fix inline_ptx when there is only clobber
 register

The interface crashes when there is only clobber register (read-write) today.There isn't any PTX instruction with only clobber register but the compiler should not crash.
---
 .../LLVMIR/IR/BasicPtxBuilderInterface.cpp    | 34 ++++++++++++-------
 .../Conversion/NVVMToLLVM/nvvm-to-llvm.mlir   | 17 ++++++++++
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
index 9ce413ceeaf6b..0138630ad9dd7 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
@@ -130,7 +130,7 @@ static SmallVector<Value> extractStructElements(PatternRewriter &rewriter,
 
   SmallVector<Value> elems;
   for (unsigned i : llvm::seq<unsigned>(0, structTy.getBody().size()))
-    elems.push_back(LLVM::ExtractValueOp::create(rewriter, loc, structVal, i));
+    elems.push_back(rewriter.create<LLVM::ExtractValueOp>(loc, structVal, i));
 
   return elems;
 }
@@ -505,21 +505,31 @@ void PtxBuilder::buildAndReplaceOp() {
     return;
   }
 
-  // Case 1: Simple path, return single scalar
+  // Case 1: Simple path, single scalar inline asm result.
   if (!needsPackUnpack(interfaceOp, needsManualRegisterMapping,
                        registerModifiers)) {
-    if (inlineAsmOp->getNumResults() > 0) {
+    // Sub-case 1a: the wrapper op has a declared result -- replace it
+    // directly with the inline asm result.
+    if (interfaceOp->getNumResults() > 0) {
       rewriter.replaceOp(interfaceOp, inlineAsmOp->getResults());
-    } else {
-      // RW-only case with no declared results: forward the RW value.
-      SmallVector<Value> results;
-      for (auto [m, v] : llvm::zip(registerModifiers, ptxOperands))
-        if (m == PTXRegisterMod::ReadWrite) {
-          results.push_back(v);
-          break;
-        }
-      rewriter.replaceOp(interfaceOp, results);
+      return;
     }
+    // Sub-case 1b: RW-only, no declared result. The inline asm produces a
+    // single value that represents the post-asm value of the read-write
+    // operand; forward it to that operand's uses and erase the wrapper.
+    if (inlineAsmOp->getNumResults() > 0) {
+      Value postAsm = inlineAsmOp->getResult(0);
+      for (auto [m, v] : llvm::zip(registerModifiers, ptxOperands)) {
+        if (m != PTXRegisterMod::ReadWrite)
+          continue;
+        v.replaceUsesWithIf(postAsm, [&](OpOperand &use) {
+          Operation *owner = use.getOwner();
+          return owner != interfaceOp && owner != inlineAsmOp;
+        });
+        break;
+      }
+    }
+    rewriter.eraseOp(interfaceOp);
     return;
   }
 
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index a188aec18134c..a648295dda52c 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -771,3 +771,20 @@ llvm.func @cvt_i8_bf16(%a : i8)  {
                           -> i16
   llvm.return  
 }
+
+
+// CHECK-LABEL: @inline_ptx_single_rw_no_result(
+// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: f32, %[[arg1:[a-zA-Z0-9_]+]]: f32)
+llvm.func @inline_ptx_single_rw_no_result(%a : f32, %b : f32) -> f32 {
+  // Single read-write operand and no declared results: the inline asm result
+  // value represents the post-asm value of the RW operand and must replace
+  // its uses (without trying to replace the wrapper op's non-existent result).
+  // CHECK: %[[C:.+]] = llvm.fadd %[[arg0]], %[[arg1]] : f32
+  // CHECK: %[[S0:.+]] = llvm.inline_asm has_side_effects asm_dialect = att "asm1 ", "=f,0" %[[C]] : (f32) -> f32
+  // CHECK: %[[R:.+]] = llvm.fadd %[[S0]], %[[arg0]] : f32
+  // CHECK: llvm.return %[[R]] : f32
+  %c = llvm.fadd %a, %b : f32
+  nvvm.inline_ptx "asm1 " rw(%c : f32)
+  %a2 = llvm.fadd %c, %a : f32
+  llvm.return %a2 : f32
+}
\ No newline at end of file

>From 9e5657094879ad29f983a0e5d8669bd5bf563f65 Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Mon, 27 Apr 2026 16:17:15 +0200
Subject: [PATCH 2/2] fz

---
 mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
index 0138630ad9dd7..1747d5e383ca6 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
@@ -130,7 +130,7 @@ static SmallVector<Value> extractStructElements(PatternRewriter &rewriter,
 
   SmallVector<Value> elems;
   for (unsigned i : llvm::seq<unsigned>(0, structTy.getBody().size()))
-    elems.push_back(rewriter.create<LLVM::ExtractValueOp>(loc, structVal, i));
+    elems.push_back(LLVM::ExtractValueOp::create(rewriter, loc, structVal, i));
 
   return elems;
 }