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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Guray Ozen (grypp)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/194380.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp (+22-12) 
- (modified) mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir (+17) 


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

``````````

</details>


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