MIR Serialization: Serialize register mask machine operands.

[arphaman]

Summary

This patch is based on a previous serialization patch that global address machine operands (http://reviews.llvm.org/D10671).

This patch implements serialization of register mask machine operands. This patch serializes the call preserved register masks that are defined by a target, it can't serialize arbitrary register masks.

This patch also extends the TargetRegisterInfo class and TableGen so that the users of TRI can get the list of all the call preserved register masks and a list of names that correspond to each mask.

This is what the extended tablegen would produce for the X86 TRI:

ArrayRef<const uint32_t *> X86GenRegisterInfo::getRegMasks() const {
  static const uint32_t *Masks[] = {
    CSR_32_RegMask, 
    CSR_32EHRet_RegMask, 
    CSR_64_RegMask, 
    CSR_64EHRet_RegMask, 
    CSR_64_AllRegs_RegMask, 
    CSR_64_AllRegs_AVX_RegMask, 
    CSR_64_Intel_OCL_BI_RegMask, 
    CSR_64_Intel_OCL_BI_AVX_RegMask, 
    CSR_64_Intel_OCL_BI_AVX512_RegMask, 
    CSR_64_MostRegs_RegMask, 
    CSR_64_RT_AllRegs_RegMask, 
    CSR_64_RT_AllRegs_AVX_RegMask, 
    CSR_64_RT_MostRegs_RegMask, 
    CSR_NoRegs_RegMask, 
    CSR_Win64_RegMask, 
    CSR_Win64_Intel_OCL_BI_AVX_RegMask, 
    CSR_Win64_Intel_OCL_BI_AVX512_RegMask, 
    nullptr
  };
  return ArrayRef<const uint32_t *>(Masks, (size_t)17);
}

ArrayRef<const char *> X86GenRegisterInfo::getRegMaskNames() const {
  static const char *Names[] = {
    "CSR_32",
    "CSR_32EHRet",
    "CSR_64",
    "CSR_64EHRet",
    "CSR_64_AllRegs",
    "CSR_64_AllRegs_AVX",
    "CSR_64_Intel_OCL_BI",
    "CSR_64_Intel_OCL_BI_AVX",
    "CSR_64_Intel_OCL_BI_AVX512",
    "CSR_64_MostRegs",
    "CSR_64_RT_AllRegs",
    "CSR_64_RT_AllRegs_AVX",
    "CSR_64_RT_MostRegs",
    "CSR_NoRegs",
    "CSR_Win64",
    "CSR_Win64_Intel_OCL_BI_AVX",
    "CSR_Win64_Intel_OCL_BI_AVX512",
    nullptr
  };
  return ArrayRef<const char *>(Names, (size_t)17);
}

llvm/trunk/include/llvm/Target/TargetRegisterInfo.h

public:
   /// registers.
   ///
   virtual const uint32_t *getCallPreservedMask(const MachineFunction &MF,
                                                CallingConv::ID) const {
     // The default mask clobbers everything.  All targets should override.
     return nullptr;
   }
 
+  /// Return all the call-preserved register masks defined for this target.
+  virtual ArrayRef<const uint32_t *> getRegMasks() const = 0;
+  virtual ArrayRef<const char *> getRegMaskNames() const = 0;
+
   /// getReservedRegs - Returns a bitset indexed by physical register number
   /// indicating if a register is a special register that has particular uses
   /// and should be considered unavailable at all times, e.g. SP, RA. This is
   /// used by register scavenger to determine what registers are free.
   virtual BitVector getReservedRegs(const MachineFunction &MF) const = 0;
 
   /// Prior to adding the live-out mask to a stackmap or patchpoint
   /// instruction, provide the target the opportunity to adjust it (mainly to

llvm/trunk/lib/CodeGen/MIRParser/MIParser.cpp

class MIParser {
   /// Maps from basic block numbers to MBBs.
   const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots;
   /// Maps from indices to unnamed global values and metadata nodes.
   const SlotMapping &IRSlots;
   /// Maps from instruction names to op codes.
   StringMap<unsigned> Names2InstrOpCodes;
   /// Maps from register names to registers.
   StringMap<unsigned> Names2Regs;
+  /// Maps from register mask names to register masks.
+  StringMap<const uint32_t *> Names2RegMasks;
 
 public:
   MIParser(SourceMgr &SM, MachineFunction &MF, SMDiagnostic &Error,
            StringRef Source,
            const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
            const SlotMapping &IRSlots);
 
   void lex();
private:
 
   bool parseInstruction(unsigned &OpCode);
 
   void initNames2Regs();
 
   /// Try to convert a register name to a register number. Return true if the
   /// register name is invalid.
   bool getRegisterByName(StringRef RegName, unsigned &Reg);
+
+  void initNames2RegMasks();
+
+  /// Check if the given identifier is a name of a register mask.
+  ///
+  /// Return null if the identifier isn't a register mask.
+  const uint32_t *getRegMask(StringRef Identifier);
 };
 
 } // end anonymous namespace
 
 MIParser::MIParser(SourceMgr &SM, MachineFunction &MF, SMDiagnostic &Error,
                    StringRef Source,
                    const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
                    const SlotMapping &IRSlots)
MachineInstr *MIParser::parse() {
   // Verify machine operands.
   if (!MCID.isVariadic()) {
     for (size_t I = 0, E = Operands.size(); I < E; ++I) {
       if (I < MCID.getNumOperands())
         continue;
       // Mark this register as implicit to prevent an assertion when it's added
       // to an instruction. This is a temporary workaround until the implicit
       // register flag can be parsed.
+      if (Operands[I].isReg())
         Operands[I].setImplicit();
     }
   }
 
   // TODO: Determine the implicit behaviour when implicit register flags are
   // parsed.
   auto *MI = MF.CreateMachineInstr(MCID, DebugLoc(), /*NoImplicit=*/true);
   for (const auto &Operand : Operands)
     MI->addOperand(MF, Operand);
case MIToken::IntegerLiteral:
     return parseImmediateOperand(Dest);
   case MIToken::MachineBasicBlock:
     return parseMBBOperand(Dest);
   case MIToken::GlobalValue:
   case MIToken::NamedGlobalValue:
     return parseGlobalAddressOperand(Dest);
   case MIToken::Error:
     return true;
+  case MIToken::Identifier:
+    if (const auto *RegMask = getRegMask(Token.stringValue())) {
+      Dest = MachineOperand::CreateRegMask(RegMask);
+      lex();
+      break;
+    }
+  // fallthrough
   default:
     // TODO: parse the other machine operands.
     return error("expected a machine operand");
   }
   return false;
 }
 
 void MIParser::initNames2InstrOpCodes() {
bool MIParser::getRegisterByName(StringRef RegName, unsigned &Reg) {
   initNames2Regs();
   auto RegInfo = Names2Regs.find(RegName);
   if (RegInfo == Names2Regs.end())
     return true;
   Reg = RegInfo->getValue();
   return false;
 }
 
+void MIParser::initNames2RegMasks() {
+  if (!Names2RegMasks.empty())
+    return;
+  const auto *TRI = MF.getSubtarget().getRegisterInfo();
+  assert(TRI && "Expected target register info");
+  ArrayRef<const uint32_t *> RegMasks = TRI->getRegMasks();
+  ArrayRef<const char *> RegMaskNames = TRI->getRegMaskNames();
+  assert(RegMasks.size() == RegMaskNames.size());
+  for (size_t I = 0, E = RegMasks.size(); I < E; ++I)
+    Names2RegMasks.insert(
+        std::make_pair(StringRef(RegMaskNames[I]).lower(), RegMasks[I]));
+}
+
+const uint32_t *MIParser::getRegMask(StringRef Identifier) {
+  initNames2RegMasks();
+  auto RegMaskInfo = Names2RegMasks.find(Identifier);
+  if (RegMaskInfo == Names2RegMasks.end())
+    return nullptr;
+  return RegMaskInfo->getValue();
+}
+
 MachineInstr *
 llvm::parseMachineInstr(SourceMgr &SM, MachineFunction &MF, StringRef Src,
                         const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
                         const SlotMapping &IRSlots, SMDiagnostic &Error) {
   return MIParser(SM, MF, Error, Src, MBBSlots, IRSlots).parse();
 }

llvm/trunk/lib/CodeGen/MIRPrinter.cpp

 using namespace llvm;
 
 namespace {
 
 /// This class prints out the machine functions using the MIR serialization
 /// format.
 class MIRPrinter {
   raw_ostream &OS;
+  DenseMap<const uint32_t *, unsigned> RegisterMaskIds;
 
 public:
   MIRPrinter(raw_ostream &OS) : OS(OS) {}
 
   void print(const MachineFunction &MF);
 
   void convert(yaml::MachineFunction &MF, const MachineRegisterInfo &RegInfo);
   void convert(const Module &M, yaml::MachineBasicBlock &YamlMBB,
                const MachineBasicBlock &MBB);
+
+private:
+  void initRegisterMaskIds(const MachineFunction &MF);
 };
 
 /// This class prints out the machine instructions using the MIR serialization
 /// format.
 class MIPrinter {
   const Module &M;
   raw_ostream &OS;
+  const DenseMap<const uint32_t *, unsigned> &RegisterMaskIds;
 
 public:
-  MIPrinter(const Module &M, raw_ostream &OS) : M(M), OS(OS) {}
+  MIPrinter(const Module &M, raw_ostream &OS,
+            const DenseMap<const uint32_t *, unsigned> &RegisterMaskIds)
+      : M(M), OS(OS), RegisterMaskIds(RegisterMaskIds) {}
 
   void print(const MachineInstr &MI);
   void print(const MachineOperand &Op, const TargetRegisterInfo *TRI);
 };
 
 } // end anonymous namespace
 
 namespace llvm {
static StringRef input(StringRef Str, void *Ctxt, Module &Mod) {
     return "";
   }
 };
 
 } // end namespace yaml
 } // end namespace llvm
 
 void MIRPrinter::print(const MachineFunction &MF) {
+  initRegisterMaskIds(MF);
+
   yaml::MachineFunction YamlMF;
   YamlMF.Name = MF.getName();
   YamlMF.Alignment = MF.getAlignment();
   YamlMF.ExposesReturnsTwice = MF.exposesReturnsTwice();
   YamlMF.HasInlineAsm = MF.hasInlineAsm();
   convert(YamlMF, MF.getRegInfo());
 
   int I = 0;
void MIRPrinter::convert(const Module &M, yaml::MachineBasicBlock &YamlMBB,
   YamlMBB.AddressTaken = MBB.hasAddressTaken();
   YamlMBB.IsLandingPad = MBB.isLandingPad();
 
   // Print the machine instructions.
   YamlMBB.Instructions.reserve(MBB.size());
   std::string Str;
   for (const auto &MI : MBB) {
     raw_string_ostream StrOS(Str);
-    MIPrinter(M, StrOS).print(MI);
+    MIPrinter(M, StrOS, RegisterMaskIds).print(MI);
     YamlMBB.Instructions.push_back(StrOS.str());
     Str.clear();
   }
 }
 
+void MIRPrinter::initRegisterMaskIds(const MachineFunction &MF) {
+  const auto *TRI = MF.getSubtarget().getRegisterInfo();
+  unsigned I = 0;
+  for (const uint32_t *Mask : TRI->getRegMasks())
+    RegisterMaskIds.insert(std::make_pair(Mask, I++));
+}
+
 void MIPrinter::print(const MachineInstr &MI) {
   const auto &SubTarget = MI.getParent()->getParent()->getSubtarget();
   const auto *TRI = SubTarget.getRegisterInfo();
   assert(TRI && "Expected target register info");
   const auto *TII = SubTarget.getInstrInfo();
   assert(TII && "Expected target instruction info");
 
   unsigned I = 0, E = MI.getNumOperands();
case MachineOperand::MO_MachineBasicBlock:
     break;
   case MachineOperand::MO_GlobalAddress:
     // FIXME: Make this faster - print as operand will create a slot tracker to
     // print unnamed values for the whole module every time it's called, which
     // is inefficient.
     Op.getGlobal()->printAsOperand(OS, /*PrintType=*/false, &M);
     // TODO: Print offset and target flags.
     break;
+  case MachineOperand::MO_RegisterMask: {
+    auto RegMaskInfo = RegisterMaskIds.find(Op.getRegMask());
+    if (RegMaskInfo != RegisterMaskIds.end())
+      OS << StringRef(TRI->getRegMaskNames()[RegMaskInfo->second]).lower();
+    else
+      llvm_unreachable("Can't print this machine register mask yet.");
+    break;
+  }
   default:
     // TODO: Print the other machine operands.
     llvm_unreachable("Can't print this machine operand at the moment");
   }
 }
 
 void llvm::printMIR(raw_ostream &OS, const Module &M) {
   yaml::Output Out(OS);
   Out << const_cast<Module &>(M);
 }
 
 void llvm::printMIR(raw_ostream &OS, const MachineFunction &MF) {
   MIRPrinter Printer(OS);
   Printer.print(MF);
 }

llvm/trunk/test/CodeGen/MIR/X86/register-mask-operands.mir

+# RUN: llc -march=x86-64 -start-after branch-folder -stop-after branch-folder -o /dev/null %s | FileCheck %s
+# This test ensures that the MIR parser parses register mask operands correctly.
+
+--- |
+
+  define i32 @compute(i32 %a) #0 {
+  body:
+    %c = mul i32 %a, 11
+    ret i32 %c
+  }
+
+  define i32 @foo(i32 %a) #0 {
+  entry:
+    %b = call i32 @compute(i32 %a)
+    ret i32 %b
+  }
+
+  attributes #0 = { "no-frame-pointer-elim"="false" }
+
+...
+---
+name:            compute
+body:
+  - id:          0
+    name:        body
+    instructions:
+      - '%eax = IMUL32rri8 %edi, 11'
+      - 'RETQ %eax'
+...
+---
+# CHECK: name: foo
+name:            foo
+body:
+  - id:          0
+    name:        entry
+    instructions:
+      # CHECK:      - 'PUSH64r %rax
+      # CHECK-NEXT: - 'CALL64pcrel32 @compute, csr_64, %rsp, %edi, %rsp, %eax'
+      - 'PUSH64r %rax'
+      - 'CALL64pcrel32 @compute, csr_64, %rsp, %edi, %rsp, %eax'
+      - '%rdx = POP64r'
+      - 'RETQ %eax'
+...

llvm/trunk/utils/TableGen/RegisterInfoEmitter.cpp

OS << "  const RegClassWeight &getRegClassWeight("
      << "  unsigned getNumRegPressureSets() const override;\n"
      << "  const char *getRegPressureSetName(unsigned Idx) const override;\n"
      << "  unsigned getRegPressureSetLimit(const MachineFunction &MF, unsigned "
         "Idx) const override;\n"
      << "  const int *getRegClassPressureSets("
      << "const TargetRegisterClass *RC) const override;\n"
      << "  const int *getRegUnitPressureSets("
      << "unsigned RegUnit) const override;\n"
+     << "  ArrayRef<const char *> getRegMaskNames() const override;\n"
+     << "  ArrayRef<const uint32_t *> getRegMasks() const override;\n"
      << "};\n\n";
 
   const auto &RegisterClasses = RegBank.getRegClasses();
 
   if (!RegisterClasses.empty()) {
     OS << "namespace " << RegisterClasses.front().Namespace
        << " { // Register classes\n";
 
for (unsigned i = 0, e = CSRSets.size(); i != e; ++i) {
 
     OS << "static const uint32_t " << CSRSet->getName()
        << "_RegMask[] = { ";
     printBitVectorAsHex(OS, Covered, 32);
     OS << "};\n";
   }
   OS << "\n\n";
 
+  OS << "ArrayRef<const uint32_t *> " << ClassName
+     << "::getRegMasks() const {\n";
+  OS << "  static const uint32_t *Masks[] = {\n";
+  for (Record *CSRSet : CSRSets)
+    OS << "    " << CSRSet->getName() << "_RegMask, \n";
+  OS << "    nullptr\n  };\n";
+  OS << "  return ArrayRef<const uint32_t *>(Masks, (size_t)" << CSRSets.size()
+     << ");\n";
+  OS << "}\n\n";
+
+  OS << "ArrayRef<const char *> " << ClassName
+     << "::getRegMaskNames() const {\n";
+  OS << "  static const char *Names[] = {\n";
+  for (Record *CSRSet : CSRSets)
+    OS << "    " << '"' << CSRSet->getName() << '"' << ",\n";
+  OS << "    nullptr\n  };\n";
+  OS << "  return ArrayRef<const char *>(Names, (size_t)" << CSRSets.size()
+     << ");\n";
+  OS << "}\n\n";
+
   OS << "} // End llvm namespace\n";
   OS << "#endif // GET_REGINFO_TARGET_DESC\n\n";
 }
 
 void RegisterInfoEmitter::run(raw_ostream &OS) {
   CodeGenTarget Target(Records);
   CodeGenRegBank &RegBank = Target.getRegBank();
   RegBank.computeDerivedInfo();