PTX: Encode registers as unsigned values in the MC asm printer instead of using external symbols

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@145946 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 88 insertions, 14 deletions

diff --git a/lib/Target/PTX/InstPrinter/PTXInstPrinter.cpp b/lib/Target/PTX/InstPrinter/PTXInstPrinter.cpp
index 2f6c92d11c..a274348788 100644
--- a/lib/Target/PTX/InstPrinter/PTXInstPrinter.cpp
+++ b/lib/Target/PTX/InstPrinter/PTXInstPrinter.cpp
@@ -38,7 +38,37 @@ StringRef PTXInstPrinter::getOpcodeName(unsigned Opcode) const {
 }
 
 void PTXInstPrinter::printRegName(raw_ostream &OS, unsigned RegNo) const {
-  OS << getRegisterName(RegNo);
+  // Decode the register number into type and offset
+  unsigned RegType   = RegNo & 0xF;
+  unsigned RegOffset = RegNo >> 4;
+
+  // Print the register
+  OS << "%";
+
+  switch (RegType) {
+  default:
+    llvm_unreachable("Unknown register type!");
+  case PTXRegisterType::Pred:
+    OS << "p";
+    break;
+  case PTXRegisterType::B16:
+    OS << "rh";
+    break;
+  case PTXRegisterType::B32:
+    OS << "r";
+    break;
+  case PTXRegisterType::B64:
+    OS << "rd";
+    break;
+  case PTXRegisterType::F32:
+    OS << "f";
+    break;
+  case PTXRegisterType::F64:
+    OS << "fd";
+    break;
+  }
+
+  OS << RegOffset;
 }
 
 void PTXInstPrinter::printInst(const MCInst *MI, raw_ostream &O,
@@ -139,6 +169,8 @@ void PTXInstPrinter::printOperand(const MCInst *MI, unsigned OpNo,
     } else {
       O << "0000000000000000";
     }
+  } else if (Op.isReg()) {
+    printRegName(O, Op.getReg());
   } else {
     assert(Op.isExpr() && "unknown operand kind in printOperand");
     const MCExpr *Expr = Op.getExpr();
diff --git a/lib/Target/PTX/MCTargetDesc/PTXBaseInfo.h b/lib/Target/PTX/MCTargetDesc/PTXBaseInfo.h
index c6094be4d1..d376d367ab 100644
--- a/lib/Target/PTX/MCTargetDesc/PTXBaseInfo.h
+++ b/lib/Target/PTX/MCTargetDesc/PTXBaseInfo.h
@@ -57,6 +57,18 @@ namespace llvm {
       RndPosInfInt        = 10  // .rpi
     };
   } // namespace PTXII
+
+  namespace PTXRegisterType {
+    // Register type encoded in MCOperands
+    enum {
+      Pred  = 0,
+      B16,
+      B32,
+      B64,
+      F32,
+      F64
+    };
+  } // namespace PTXRegisterType
 } // namespace llvm
 
 #endif
diff --git a/lib/Target/PTX/PTXAsmPrinter.cpp b/lib/Target/PTX/PTXAsmPrinter.cpp
index bdf238b1b0..db614ad36a 100644
--- a/lib/Target/PTX/PTXAsmPrinter.cpp
+++ b/lib/Target/PTX/PTXAsmPrinter.cpp
@@ -521,20 +521,38 @@ MCOperand PTXAsmPrinter::GetSymbolRef(const MachineOperand &MO,
 MCOperand PTXAsmPrinter::lowerOperand(const MachineOperand &MO) {
   MCOperand MCOp;
   const PTXMachineFunctionInfo *MFI = MF->getInfo<PTXMachineFunctionInfo>();
-  const MCExpr *Expr;
-  const char *RegSymbolName;
+  const MachineRegisterInfo& MRI = MF->getRegInfo();
+  const TargetRegisterClass* TRC;
+  unsigned RegType;
+  unsigned RegOffset;
+  unsigned EncodedReg;
   switch (MO.getType()) {
   default:
     llvm_unreachable("Unknown operand type");
   case MachineOperand::MO_Register:
-    // We create register operands as symbols, since the PTXInstPrinter class
-    // has no way to map virtual registers back to a name without some ugly
-    // hacks.
-    // FIXME: Figure out a better way to handle virtual register naming.
-    RegSymbolName = MFI->getRegisterName(MO.getReg());
-    Expr = MCSymbolRefExpr::Create(RegSymbolName, MCSymbolRefExpr::VK_None,
-                                   OutContext);
-    MCOp = MCOperand::CreateExpr(Expr);
+    if (MO.getReg() > 0) {
+      TRC = MRI.getRegClass(MO.getReg());
+      // Determine which PTX register type to use
+      if (TRC == PTX::RegPredRegisterClass)
+        RegType = PTXRegisterType::Pred;
+      else if (TRC == PTX::RegI16RegisterClass)
+        RegType = PTXRegisterType::B16;
+      else if (TRC == PTX::RegI32RegisterClass)
+        RegType = PTXRegisterType::B32;
+      else if (TRC == PTX::RegI64RegisterClass)
+        RegType = PTXRegisterType::B64;
+      else if (TRC == PTX::RegF32RegisterClass)
+        RegType = PTXRegisterType::F32;
+      else if (TRC == PTX::RegF64RegisterClass)
+        RegType = PTXRegisterType::F64;
+      // Determine our virtual register offset
+      RegOffset = MFI->getOffsetForRegister(TRC, MO.getReg());
+      // Encode the register
+      EncodedReg = (RegOffset << 4) | RegType;
+    } else {
+      EncodedReg = 0;
+    }
+    MCOp = MCOperand::CreateReg(EncodedReg);
     break;
   case MachineOperand::MO_Immediate:
     MCOp = MCOperand::CreateImm(MO.getImm());
diff --git a/lib/Target/PTX/PTXMachineFunctionInfo.h b/lib/Target/PTX/PTXMachineFunctionInfo.h
index 3b985f7dd6..6f3b826db9 100644
--- a/lib/Target/PTX/PTXMachineFunctionInfo.h
+++ b/lib/Target/PTX/PTXMachineFunctionInfo.h
@@ -143,18 +143,30 @@ public:
     return UsedRegs.lookup(TRC).size();
   }
 
+  /// getOffsetForRegister - Returns the offset of the virtual register
+  unsigned getOffsetForRegister(const TargetRegisterClass *TRC,
+                                unsigned Reg) const {
+    const RegisterList &RegList = UsedRegs.lookup(TRC);
+    for (unsigned i = 0, e = RegList.size(); i != e; ++i) {
+      if (RegList[i] == Reg)
+        return i;
+    }
+    //llvm_unreachable("Unknown virtual register");
+    return 0;
+  }
+
   /// getFrameSymbol - Returns the symbol name for the given FrameIndex.
   const char* getFrameSymbol(int FrameIndex) {
     if (FrameSymbols.count(FrameIndex)) {
       return FrameSymbols.lookup(FrameIndex).c_str();
     } else {
-      std::string Name = "__local";
-      Name += utostr(FrameIndex);
+      std::string Name          = "__local";
+      Name                     += utostr(FrameIndex);
       // The whole point of caching this name is to ensure the pointer we pass
       // to any getExternalSymbol() calls will remain valid for the lifetime of
       // the back-end instance. This is to work around an issue in SelectionDAG
       // where symbol names are expected to be life-long strings.
-      FrameSymbols[FrameIndex] = Name;
+      FrameSymbols[FrameIndex]  = Name;
       return FrameSymbols[FrameIndex].c_str();
     }
   }