[Mlir-commits] [mlir] [MLIR][NVVM] Preserve PTX special registers in inline_ptx lowering (PR #203251)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Bastian Hagedorn (bastianhagedorn)

<details>
<summary>Changes</summary>

`PtxBuilder::build()` converted operand placeholders (written as %0, %1, and the predicate as @%N, since TableGen string attributes cannot contain '$') to the inline-asm operand form with a blanket `replace(ptx, '%', '$')`. That also rewrote literal PTX special-register names such as %tid.x, %laneid and %dynamic_smem_size into $tid.x etc., producing invalid PTX for any `nvvm.inline_ptx` whose body reads a special register.

Convert only a '%' that is immediately followed by a digit (operand placeholders and the @%N predicate); leave %<name> special registers intact. PTX special registers always begin with a letter after '%', so the digit test unambiguously distinguishes them from operand placeholders.

Add an NVVMToLLVM regression test that reads %laneid through nvvm.inline_ptx.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp (+15-3) 
- (modified) mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir (+26) 


``````````diff
diff --git a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
index 1747d5e383ca6..87206ba726765 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
@@ -474,9 +474,21 @@ LLVM::InlineAsmOp PtxBuilder::build() {
     ptxInstruction = predicateStr + " " + ptxInstruction;
   }
 
-  // Tablegen doesn't accept $, so we use %, but inline assembly uses $.
-  // Replace all % with $
-  llvm::replace(ptxInstruction, '%', '$');
+  // 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
+  // immediately followed by a digit; this leaves literal PTX special-register
+  // names such as %tid.x, %laneid or %dynamic_smem_size intact.
+  std::string mapped;
+  mapped.reserve(ptxInstruction.size());
+  for (size_t i = 0, e = ptxInstruction.size(); i < e; ++i) {
+    if (ptxInstruction[i] == '%' && i + 1 < e &&
+        llvm::isDigit(ptxInstruction[i + 1]))
+      mapped.push_back('$');
+    else
+      mapped.push_back(ptxInstruction[i]);
+  }
+  ptxInstruction = std::move(mapped);
 
   return LLVM::InlineAsmOp::create(
       rewriter, interfaceOp->getLoc(),
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index 5a381ce1e679e..c7b0087365398 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -788,3 +788,29 @@ llvm.func @inline_ptx_single_rw_no_result(%a : f32, %b : f32) -> f32 {
   %a2 = llvm.fadd %c, %a : f32
   llvm.return %a2 : f32
 }
+
+// -----
+
+// A literal PTX special register (a %-prefixed name) must survive the operand
+// placeholder rewrite: only the {$w0} placeholder becomes the inline-asm
+// operand $0, while %laneid is preserved verbatim. A blanket '%' -> '$'
+// replacement previously mangled it into "$laneid", emitting invalid PTX.
+// CHECK-LABEL: @inline_ptx_preserves_special_register
+llvm.func @inline_ptx_preserves_special_register() -> i32 {
+  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mov.u32 $0, %laneid;", "=r"
+  %0 = nvvm.inline_ptx "mov.u32 {$w0}, %laneid;" -> i32
+  llvm.return %0 : i32
+}
+
+// -----
+
+// Special-register names that *end* in a digit (e.g. the %pm0..%pm7
+// performance-monitor counters) must survive too: only a '%' immediately
+// followed by a digit is an operand placeholder, so %pm0 stays %pm0 while the
+// {$w0} placeholder becomes $0.
+// CHECK-LABEL: @inline_ptx_special_register_trailing_digit
+llvm.func @inline_ptx_special_register_trailing_digit() -> i32 {
+  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mov.u32 $0, %pm0;", "=r"
+  %0 = nvvm.inline_ptx "mov.u32 {$w0}, %pm0;" -> i32
+  llvm.return %0 : i32
+}

``````````

</details>


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