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

[Bastian Hagedorn]

https://github.com/bastianhagedorn created https://github.com/llvm/llvm-project/pull/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.

>From 4e58077eec0fb27fdee84dd4751bed3170cb9f1c Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Thu, 11 Jun 2026 11:16:16 +0000
Subject: [PATCH] [MLIR][NVVM] Preserve PTX special registers in inline_ptx
 lowering

`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>
---
 .../LLVMIR/IR/BasicPtxBuilderInterface.cpp    | 18 ++++++++++---
 .../Conversion/NVVMToLLVM/nvvm-to-llvm.mlir   | 26 +++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

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