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

[llvmlistbot at llvm.org]

Author: Bastian Hagedorn
Date: 2026-06-12T08:55:19+02:00
New Revision: 62a2af17a33394ac6f286dc110633013cb01ed7d

URL: https://github.com/llvm/llvm-project/commit/62a2af17a33394ac6f286dc110633013cb01ed7d
DIFF: https://github.com/llvm/llvm-project/commit/62a2af17a33394ac6f286dc110633013cb01ed7d.diff

LOG: [MLIR][NVVM] Preserve PTX special registers in inline_ptx lowering (#203251)

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

Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
    mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir

Removed: 
    


################################################################################
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 b012e695072a8..1886e72d72f41 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -775,3 +775,21 @@ llvm.func @inline_ptx_single_rw_no_result(%a : f32, %b : f32) -> f32 {
   %a2 = llvm.fadd %c, %a : f32
   llvm.return %a2 : f32
 }
+
+// -----
+
+// 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
+}
+
+// -----
+
+// 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
+}