[llvm] 7e35cbb - [NVPTX] Cleanup PTX virtual register encoding (NFC) (#212395)

via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 30 14:46:12 PDT 2026


Author: Alex MacLean
Date: 2026-07-30T14:46:07-07:00
New Revision: 7e35cbb1b6f5915ab732a22e83b3afc20e29181a

URL: https://github.com/llvm/llvm-project/commit/7e35cbb1b6f5915ab732a22e83b3afc20e29181a
DIFF: https://github.com/llvm/llvm-project/commit/7e35cbb1b6f5915ab732a22e83b3afc20e29181a.diff

LOG: [NVPTX] Cleanup PTX virtual register encoding (NFC) (#212395)

Unify the scheme used to carry PTX virtual registers through
`MCOperand`s. The register class was encoded and decoded 
by two hand-written switchtables, one in
`NVPTXAsmPrinter::encodeVirtualRegister` and one in
`NVPTXInstPrinter::printRegName`, kept in sync only by
comment. Both now share a `VirtualRegisterKind` enum and
`getVirtualRegisterPrefix` in `NVPTXBaseInfo.h`, which also
subsumes `getNVPTXRegClassStr` and drops the stale
`%f`/`%fd` cases left over from the typed float register
classes.

Co-Authored by Opus-5.

Added: 
    

Modified: 
    llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXBaseInfo.h
    llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXInstPrinter.cpp
    llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
    llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h
    llvm/lib/Target/NVPTX/NVPTXRegisterInfo.cpp
    llvm/lib/Target/NVPTX/NVPTXRegisterInfo.h

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXBaseInfo.h b/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXBaseInfo.h
index ef1f5da34cccd..283b93533b679 100644
--- a/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXBaseInfo.h
+++ b/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXBaseInfo.h
@@ -16,11 +16,60 @@
 #ifndef LLVM_LIB_TARGET_NVPTX_MCTARGETDESC_NVPTXBASEINFO_H
 #define LLVM_LIB_TARGET_NVPTX_MCTARGETDESC_NVPTXBASEINFO_H
 
+#include "llvm/ADT/StringRef.h"
+#include "llvm/MC/MCRegister.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/NVPTXAddrSpace.h"
 namespace llvm {
 
 using namespace NVPTXAS;
 
+namespace NVPTX {
+
+// PTX virtual registers are numbered per register class. The asm printer
+// packs the class into the upper bits of the register number so that the inst
+// printer can recover the name the register was declared with.
+enum class VirtualRegisterKind : unsigned {
+  Physical = 0,
+  B1 = 1,
+  B16 = 2,
+  B32 = 3,
+  B64 = 4,
+  B128 = 5,
+};
+
+constexpr unsigned VirtualRegisterKindShift = 27;
+constexpr unsigned VirtualRegisterNumMask =
+    (1u << VirtualRegisterKindShift) - 1;
+
+// The packed registers are carried in MCOperands, so they must stay inside the
+// range MCRegister reserves for physical registers.
+static_assert(((static_cast<unsigned>(VirtualRegisterKind::B128)
+                << VirtualRegisterKindShift) |
+               VirtualRegisterNumMask) <= MCRegister::LastPhysicalReg,
+              "Packed virtual register does not fit in an MCRegister");
+
+/// The name prefix shared by all virtual registers of \p Kind.
+inline StringRef getVirtualRegisterPrefix(VirtualRegisterKind Kind) {
+  switch (Kind) {
+  case VirtualRegisterKind::B1:
+    return "%p";
+  case VirtualRegisterKind::B16:
+    return "%rs";
+  case VirtualRegisterKind::B32:
+    return "%r";
+  case VirtualRegisterKind::B64:
+    return "%rd";
+  case VirtualRegisterKind::B128:
+    return "%rq";
+  case VirtualRegisterKind::Physical:
+    break;
+  }
+  llvm_unreachable("Invalid virtual register kind");
+}
+
+} // namespace NVPTX
+
 namespace NVPTXII {
 enum {
   // These must be kept in sync with TSFlags in NVPTXInstrFormats.td

diff  --git a/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXInstPrinter.cpp b/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXInstPrinter.cpp
index 108f869effe7c..9db838523fdcd 100644
--- a/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXInstPrinter.cpp
+++ b/llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXInstPrinter.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "MCTargetDesc/NVPTXInstPrinter.h"
+#include "MCTargetDesc/NVPTXBaseInfo.h"
 #include "NVPTX.h"
 #include "NVPTXUtilities.h"
 #include "llvm/ADT/StringRef.h"
@@ -42,41 +43,19 @@ NVPTXInstPrinter::NVPTXInstPrinter(const MCAsmInfo &MAI, const MCInstrInfo &MII,
     : MCInstPrinter(MAI, MII, MRI) {}
 
 void NVPTXInstPrinter::printRegName(raw_ostream &OS, MCRegister Reg) {
-  // Decode the virtual register
-  // Must be kept in sync with NVPTXAsmPrinter::encodeVirtualRegister
-  unsigned RCId = (Reg.id() >> 28);
-  switch (RCId) {
-  default: report_fatal_error("Bad virtual register encoding");
-  case 0:
+  // Decode a register packed by NVPTXAsmPrinter::encodeVirtualRegister.
+  const auto Kind = static_cast<NVPTX::VirtualRegisterKind>(
+      Reg.id() >> NVPTX::VirtualRegisterKindShift);
+
+  if (Kind == NVPTX::VirtualRegisterKind::Physical) {
     // This is actually a physical register, so defer to the autogenerated
     // register printer
     OS << getRegisterName(Reg);
     return;
-  case 1:
-    OS << "%p";
-    break;
-  case 2:
-    OS << "%rs";
-    break;
-  case 3:
-    OS << "%r";
-    break;
-  case 4:
-    OS << "%rd";
-    break;
-  case 5:
-    OS << "%f";
-    break;
-  case 6:
-    OS << "%fd";
-    break;
-  case 7:
-    OS << "%rq";
-    break;
   }
 
-  unsigned VReg = Reg.id() & 0x0FFFFFFF;
-  OS << VReg;
+  OS << NVPTX::getVirtualRegisterPrefix(Kind)
+     << (Reg.id() & NVPTX::VirtualRegisterNumMask);
 }
 
 void NVPTXInstPrinter::printInst(const MCInst *MI, uint64_t Address,

diff  --git a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
index 80979f5594d23..0f011c18736b3 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
@@ -369,38 +369,47 @@ MCOperand NVPTXAsmPrinter::lowerOperand(const MachineOperand &MO) {
   }
 }
 
-unsigned NVPTXAsmPrinter::encodeVirtualRegister(unsigned Reg) {
-  if (Register::isVirtualRegister(Reg)) {
-    const TargetRegisterClass *RC = MRI->getRegClass(Reg);
-
-    DenseMap<unsigned, unsigned> &RegMap = VRegMapping[RC];
-    unsigned RegNum = RegMap[Reg];
-
-    // Encode the register class in the upper 4 bits
-    // Must be kept in sync with NVPTXInstPrinter::printRegName
-    unsigned Ret = 0;
-    if (RC == &NVPTX::B1RegClass) {
-      Ret = (1 << 28);
-    } else if (RC == &NVPTX::B16RegClass) {
-      Ret = (2 << 28);
-    } else if (RC == &NVPTX::B32RegClass) {
-      Ret = (3 << 28);
-    } else if (RC == &NVPTX::B64RegClass) {
-      Ret = (4 << 28);
-    } else if (RC == &NVPTX::B128RegClass) {
-      Ret = (7 << 28);
-    } else {
-      report_fatal_error("Bad register class");
-    }
-
-    // Insert the vreg number
-    Ret |= (RegNum & 0x0FFFFFFF);
-    return Ret;
-  } else {
-    // Some special-use registers are actually physical registers.
-    // Encode this as the register class ID of 0 and the real register ID.
-    return Reg & 0x0FFFFFFF;
-  }
+static NVPTX::VirtualRegisterKind
+getVirtualRegisterKind(const TargetRegisterClass *RC) {
+  if (RC == &NVPTX::B1RegClass)
+    return NVPTX::VirtualRegisterKind::B1;
+  if (RC == &NVPTX::B16RegClass)
+    return NVPTX::VirtualRegisterKind::B16;
+  if (RC == &NVPTX::B32RegClass)
+    return NVPTX::VirtualRegisterKind::B32;
+  if (RC == &NVPTX::B64RegClass)
+    return NVPTX::VirtualRegisterKind::B64;
+  if (RC == &NVPTX::B128RegClass)
+    return NVPTX::VirtualRegisterKind::B128;
+  llvm_unreachable("Bad register class");
+}
+
+unsigned NVPTXAsmPrinter::getVirtualRegisterNumber(Register Reg) const {
+  const auto It = VRegMapping.find(MRI->getRegClass(Reg));
+  assert(It != VRegMapping.end() && "Bad register class");
+
+  const unsigned Num = It->second.lookup(Reg);
+  assert(Num && "Bad virtual register");
+  return Num;
+}
+
+MCRegister NVPTXAsmPrinter::encodeVirtualRegister(Register Reg) {
+  if (Reg.isVirtual()) {
+    // Pack the register class into the upper bits so that
+    // NVPTXInstPrinter::printRegName can recover the declared name.
+    const auto Kind = getVirtualRegisterKind(MRI->getRegClass(Reg));
+    const unsigned Num = getVirtualRegisterNumber(Reg);
+    assert(Num <= NVPTX::VirtualRegisterNumMask &&
+           "Too many virtual registers");
+    return (static_cast<unsigned>(Kind) << NVPTX::VirtualRegisterKindShift) |
+           Num;
+  }
+
+  // Some special-use registers are actually physical registers.
+  // Encode this as the register class ID of 0 and the real register ID.
+  assert(Reg.id() <= NVPTX::VirtualRegisterNumMask &&
+         "Physical register would decode as a virtual register");
+  return Reg.asMCReg();
 }
 
 MCOperand NVPTXAsmPrinter::GetSymbolRef(const MCSymbol *Symbol) {
@@ -754,30 +763,15 @@ void NVPTXAsmPrinter::emitKernelFunctionDirectives(const Function &F,
   }
 }
 
-std::string NVPTXAsmPrinter::getVirtualRegisterName(unsigned Reg) const {
-  const TargetRegisterClass *RC = MRI->getRegClass(Reg);
+std::string NVPTXAsmPrinter::getVirtualRegisterName(Register Reg) const {
+  const auto Kind = getVirtualRegisterKind(MRI->getRegClass(Reg));
 
   std::string Name;
-  raw_string_ostream NameStr(Name);
-
-  VRegRCMap::const_iterator I = VRegMapping.find(RC);
-  assert(I != VRegMapping.end() && "Bad register class");
-  const DenseMap<unsigned, unsigned> &RegMap = I->second;
-
-  VRegMap::const_iterator VI = RegMap.find(Reg);
-  assert(VI != RegMap.end() && "Bad virtual register");
-  unsigned MappedVR = VI->second;
-
-  NameStr << getNVPTXRegClassStr(RC) << MappedVR;
-
+  raw_string_ostream(Name) << NVPTX::getVirtualRegisterPrefix(Kind)
+                           << getVirtualRegisterNumber(Reg);
   return Name;
 }
 
-void NVPTXAsmPrinter::emitVirtualRegister(unsigned int vr,
-                                          raw_ostream &O) {
-  O << getVirtualRegisterName(vr);
-}
-
 void NVPTXAsmPrinter::emitAliasDeclaration(const GlobalAlias *GA,
                                            raw_ostream &O) {
   const Function *F = dyn_cast_or_null<Function>(GA->getAliaseeObject());
@@ -1823,9 +1817,14 @@ void NVPTXAsmPrinter::setAndEmitFunctionVirtualRegisters(
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   for (const TargetRegisterClass &RC : TRI->regclasses()) {
     // Only declare those registers that may be used.
-    if (const unsigned N = VRegMapping[&RC].size())
-      TS->emitRegDirective(TRI->getRegSizeInBits(RC).getFixedValue(),
-                           getNVPTXRegClassStr(&RC), N + 1);
+    const auto It = VRegMapping.find(&RC);
+    if (It == VRegMapping.end() || It->second.empty())
+      continue;
+
+    TS->emitRegDirective(
+        TRI->getRegSizeInBits(RC).getFixedValue(),
+        NVPTX::getVirtualRegisterPrefix(getVirtualRegisterKind(&RC)),
+        It->second.size() + 1);
   }
 }
 
@@ -1834,19 +1833,16 @@ void NVPTXAsmPrinter::setAndEmitFunctionVirtualRegisters(
 void NVPTXAsmPrinter::encodeDebugInfoRegisterNumbers(
     const MachineFunction &MF) {
   const NVPTXSubtarget &STI = MF.getSubtarget<NVPTXSubtarget>();
-  const NVPTXRegisterInfo *registerInfo = STI.getRegisterInfo();
+  const NVPTXRegisterInfo *NRI = STI.getRegisterInfo();
 
   // Clear the old mapping, and add the new one.  This mapping is used after the
   // printing of the current function is complete, but before the next function
   // is printed.
-  registerInfo->clearDebugRegisterMap();
+  NRI->clearDebugRegisterMap();
 
-  for (auto &classMap : VRegMapping) {
-    for (auto &registerMapping : classMap.getSecond()) {
-      auto reg = registerMapping.getFirst();
-      registerInfo->addToDebugRegisterMap(reg, getVirtualRegisterName(reg));
-    }
-  }
+  for (const VRegMap &RegMap : make_second_range(VRegMapping))
+    for (const Register Reg : make_first_range(RegMap))
+      NRI->addToDebugRegisterMap(Reg, getVirtualRegisterName(Reg));
 }
 
 void NVPTXAsmPrinter::printFPConstant(const ConstantFP *Fp,
@@ -2329,7 +2325,7 @@ void NVPTXAsmPrinter::printOperand(const MachineInstr *MI, unsigned OpNum,
       else
         O << NVPTXInstPrinter::getRegisterName(MO.getReg());
     } else {
-      emitVirtualRegister(MO.getReg(), O);
+      O << getVirtualRegisterName(MO.getReg());
     }
     break;
 

diff  --git a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h
index ef04ff375553b..bea7fc9efc8ae 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h
+++ b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h
@@ -176,7 +176,11 @@ class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
   void lowerToMCInst(const MachineInstr *MI, MCInst &OutMI);
   MCOperand lowerOperand(const MachineOperand &MO);
   MCOperand GetSymbolRef(const MCSymbol *Symbol);
-  unsigned encodeVirtualRegister(unsigned Reg);
+  MCRegister encodeVirtualRegister(Register Reg);
+
+  /// The number \p Reg was assigned within its register class, as declared by
+  /// this function's .reg directives.
+  unsigned getVirtualRegisterNumber(Register Reg) const;
 
   void printMemOperand(const MachineInstr *MI, unsigned OpNum, raw_ostream &O,
                        const char *Modifier = nullptr);
@@ -186,7 +190,6 @@ class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
   void emitGlobalAlias(const Module &M, const GlobalAlias &GA) override;
   void emitHeader(Module &M, const NVPTXSubtarget &STI);
   void emitKernelFunctionDirectives(const Function &F, raw_ostream &O) const;
-  void emitVirtualRegister(unsigned int vr, raw_ostream &);
   void emitFunctionParamList(const Function *, raw_ostream &O);
   void setAndEmitFunctionVirtualRegisters(const MachineFunction &MF);
   void encodeDebugInfoRegisterNumbers(const MachineFunction &MF);
@@ -242,11 +245,11 @@ class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
 
   // This is specific per MachineFunction.
   const MachineRegisterInfo *MRI;
-  // The contents are specific for each
-  // MachineFunction. But the size of the
-  // array is not.
-  typedef DenseMap<unsigned, unsigned> VRegMap;
-  typedef DenseMap<const TargetRegisterClass *, VRegMap> VRegRCMap;
+
+  // The number assigned to each virtual register within its class, populated
+  // by setAndEmitFunctionVirtualRegisters and cleared between functions.
+  using VRegMap = DenseMap<Register, unsigned>;
+  using VRegRCMap = DenseMap<const TargetRegisterClass *, VRegMap>;
   VRegRCMap VRegMapping;
 
   // List of variables demoted to a function scope.
@@ -301,7 +304,7 @@ class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
     AsmPrinter::getAnalysisUsage(AU);
   }
 
-  std::string getVirtualRegisterName(unsigned) const;
+  std::string getVirtualRegisterName(Register Reg) const;
 
   const MCSymbol *getFunctionFrameSymbol() const override;
 

diff  --git a/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.cpp b/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.cpp
index bb7a2f41bb3f2..28abf5e108d12 100644
--- a/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.cpp
@@ -25,22 +25,6 @@ using namespace llvm;
 
 #define DEBUG_TYPE "nvptx-reg-info"
 
-StringRef llvm::getNVPTXRegClassStr(TargetRegisterClass const *RC) {
-  if (RC == &NVPTX::B128RegClass)
-    return "%rq";
-  if (RC == &NVPTX::B64RegClass)
-    return "%rd";
-  if (RC == &NVPTX::B32RegClass)
-    return "%r";
-  if (RC == &NVPTX::B16RegClass)
-    return "%rs";
-  if (RC == &NVPTX::B1RegClass)
-    return "%p";
-  if (RC == &NVPTX::SpecialRegsRegClass)
-    return "!Special!";
-  return "INTERNAL";
-}
-
 NVPTXRegisterInfo::NVPTXRegisterInfo() : NVPTXGenRegisterInfo(0) {}
 
 #define GET_REGINFO_TARGET_DESC
@@ -109,7 +93,7 @@ NVPTXRegisterInfo::getFrameLocalRegister(const MachineFunction &MF) const {
 }
 
 void NVPTXRegisterInfo::clearDebugRegisterMap() const {
-  debugRegisterMap.clear();
+  DebugRegisterMap.clear();
 }
 
 static uint64_t encodeRegisterForDwarf(StringRef RegisterName) {
@@ -123,18 +107,16 @@ static uint64_t encodeRegisterForDwarf(StringRef RegisterName) {
   // IE the bytes of the string are concatenated in reverse into a single
   // number, which is stored in ULEB128, but in practice must be no more than 8
   // bytes (excluding null terminator, which is not included).
-  uint64_t result = 0;
-  for (unsigned char c : RegisterName)
-    result = (result << 8) | c;
-  return result;
+  uint64_t Result = 0;
+  for (unsigned char C : RegisterName)
+    Result = (Result << 8) | C;
+  return Result;
 }
 
-void NVPTXRegisterInfo::addToDebugRegisterMap(
-    uint64_t preEncodedVirtualRegister, StringRef RegisterName) const {
-  uint64_t mapped = encodeRegisterForDwarf(RegisterName);
-  if (mapped == 0)
-    return;
-  debugRegisterMap.insert({preEncodedVirtualRegister, mapped});
+void NVPTXRegisterInfo::addToDebugRegisterMap(Register VirtReg,
+                                              StringRef RegisterName) const {
+  if (const uint64_t Encoded = encodeRegisterForDwarf(RegisterName))
+    DebugRegisterMap.insert({VirtReg, Encoded});
 }
 
 int64_t NVPTXRegisterInfo::getDwarfRegNum(MCRegister RegNum, bool isEH) const {
@@ -150,8 +132,7 @@ int64_t NVPTXRegisterInfo::getDwarfRegNum(MCRegister RegNum, bool isEH) const {
 int64_t NVPTXRegisterInfo::getDwarfRegNumForVirtReg(Register RegNum,
                                                     bool isEH) const {
   assert(RegNum.isVirtual());
-  uint64_t lookup = debugRegisterMap.lookup(RegNum.id());
-  if (lookup)
-    return lookup;
+  if (const uint64_t Encoded = DebugRegisterMap.lookup(RegNum))
+    return Encoded;
   return -1;
 }

diff  --git a/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.h b/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.h
index 0bd64e0e48b09..41fd07b2adb32 100644
--- a/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.h
+++ b/llvm/lib/Target/NVPTX/NVPTXRegisterInfo.h
@@ -24,7 +24,7 @@ class NVPTXRegisterInfo : public NVPTXGenRegisterInfo {
   // State for debug register mapping that can be mutated even through a const
   // pointer so that we can get the proper dwarf register encoding during ASM
   // emission.
-  mutable DenseMap<uint64_t, uint64_t> debugRegisterMap;
+  mutable DenseMap<Register, uint64_t> DebugRegisterMap;
 
 public:
   NVPTXRegisterInfo();
@@ -45,7 +45,7 @@ class NVPTXRegisterInfo : public NVPTXGenRegisterInfo {
   Register getFrameRegister(const MachineFunction &MF) const override;
   Register getFrameLocalRegister(const MachineFunction &MF) const;
 
-  // Manage the debugRegisterMap.  PTX virtual registers for DebugInfo are
+  // Manage the DebugRegisterMap.  PTX virtual registers for DebugInfo are
   // encoded using the names used in the emitted text of the PTX assembly. This
   // mapping must be managed during assembly emission.
   //
@@ -53,15 +53,12 @@ class NVPTXRegisterInfo : public NVPTXGenRegisterInfo {
   // RegisterInfo object are all const, but we need to communicate some state
   // here, because the proper encoding for debug registers is available only
   // temporarily during ASM emission.
-  void addToDebugRegisterMap(uint64_t preEncodedVirtualRegister,
-                             StringRef RegisterName) const;
+  void addToDebugRegisterMap(Register VirtReg, StringRef RegisterName) const;
   void clearDebugRegisterMap() const;
   int64_t getDwarfRegNum(MCRegister RegNum, bool isEH) const override;
   int64_t getDwarfRegNumForVirtReg(Register RegNum, bool isEH) const override;
 };
 
-StringRef getNVPTXRegClassStr(const TargetRegisterClass *RC);
-
 } // end namespace llvm
 
 #endif


        


More information about the llvm-commits mailing list