[MIR] Support Customed Register Mask and CSRs

The MIR printer dumps a string that describe the register mask of a function.
A static predefined list of register masks matches a static list of strings.
However when the register mask is not from the static predefined list, there is no descriptor string and the printer fails.
This patch adds support to custom register mask printing and dumping.
Also the list of callee saved registers (describing the registers that must be preserved for the caller) might be dynamic.
As such this data needs to be dumped and parsed back to the Machine Register Info.

Differential Revision: https://reviews.llvm.org/D30971

llvm-svn: 298207

include/llvm/CodeGen/MachineRegisterInfo.h

@@ -216,6 +216,9 @@ public:
   // Register Info
   //===--------------------------------------------------------------------===//
 
+  /// Returns true if the updated CSR list was initialized and false otherwise.
+  bool isUpdatedCSRsInitialized() const { return IsUpdatedCSRsInitialized; }
+
   /// Disables the register from the list of CSRs.
   /// I.e. the register will not appear as part of the CSR mask.
   /// \see UpdatedCalleeSavedRegs.
@@ -226,6 +229,10 @@ public:
   /// registers that are disabled from the CSR list).
   const MCPhysReg *getCalleeSavedRegs() const;
 
+  /// Sets the updated Callee Saved Registers list.
+  /// Notice that it will override ant previously disabled/saved CSRs.
+  void setCalleeSavedRegs(ArrayRef<MCPhysReg> CSRs);
+
   // Strictly for use by MachineInstr.cpp.
   void addRegOperandToUseList(MachineOperand *MO);
 
@@ -755,8 +762,6 @@ public:
 
   const BitVector &getUsedPhysRegsMask() const { return UsedPhysRegMask; }
 
-  void setUsedPhysRegMask(BitVector &Mask) { UsedPhysRegMask = Mask; }
-
   //===--------------------------------------------------------------------===//
   // Reserved Register Info
   //===--------------------------------------------------------------------===//

lib/CodeGen/MIRParser/MIParser.cpp

@@ -176,6 +176,7 @@ public:
   bool parseIntrinsicOperand(MachineOperand &Dest);
   bool parsePredicateOperand(MachineOperand &Dest);
   bool parseTargetIndexOperand(MachineOperand &Dest);
+  bool parseCustomRegisterMaskOperand(MachineOperand &Dest);
   bool parseLiveoutRegisterMaskOperand(MachineOperand &Dest);
   bool parseMachineOperand(MachineOperand &Dest,
                            Optional<unsigned> &TiedDefIdx);
@@ -1670,6 +1671,35 @@ bool MIParser::parseTargetIndexOperand(MachineOperand &Dest) {
   return false;
 }
 
+bool MIParser::parseCustomRegisterMaskOperand(MachineOperand &Dest) {
+  assert(Token.stringValue() == "CustomRegMask" && "Expected a custom RegMask");
+  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
+  assert(TRI && "Expected target register info");
+  lex();
+  if (expectAndConsume(MIToken::lparen))
+    return true;
+
+  uint32_t *Mask = MF.allocateRegisterMask(TRI->getNumRegs());
+  while (true) {
+    if (Token.isNot(MIToken::NamedRegister))
+      return error("expected a named register");
+    unsigned Reg;
+    if (parseNamedRegister(Reg))
+      return true;
+    lex();
+    Mask[Reg / 32] |= 1U << (Reg % 32);
+    // TODO: Report an error if the same register is used more than once.
+    if (Token.isNot(MIToken::comma))
+      break;
+    lex();
+  }
+
+  if (expectAndConsume(MIToken::rparen))
+    return true;
+  Dest = MachineOperand::CreateRegMask(Mask);
+  return false;
+}
+
 bool MIParser::parseLiveoutRegisterMaskOperand(MachineOperand &Dest) {
   assert(Token.is(MIToken::kw_liveout));
   const auto *TRI = MF.getSubtarget().getRegisterInfo();
@@ -1767,8 +1797,8 @@ bool MIParser::parseMachineOperand(MachineOperand &Dest,
       Dest = MachineOperand::CreateRegMask(RegMask);
       lex();
       break;
-    }
-    LLVM_FALLTHROUGH;
+    } else
+      return parseCustomRegisterMaskOperand(Dest);
   default:
     // FIXME: Parse the MCSymbol machine operand.
     return error("expected a machine operand");

lib/CodeGen/MIRParser/MIRParser.cpp

@@ -467,17 +467,19 @@ bool MIRParserImpl::parseRegisterInfo(PerFunctionMIParsingState &PFS,
     RegInfo.addLiveIn(Reg, VReg);
   }
 
-  // Parse the callee saved register mask.
-  BitVector CalleeSavedRegisterMask(RegInfo.getUsedPhysRegsMask().size());
-  if (!YamlMF.CalleeSavedRegisters)
-    return false;
-  for (const auto &RegSource : YamlMF.CalleeSavedRegisters.getValue()) {
-    unsigned Reg = 0;
-    if (parseNamedRegisterReference(PFS, Reg, RegSource.Value, Error))
-      return error(Error, RegSource.SourceRange);
-    CalleeSavedRegisterMask[Reg] = true;
+  // Parse the callee saved registers (Registers that will
+  // be saved for the caller).
+  if (YamlMF.CalleeSavedRegisters) {
+    SmallVector<MCPhysReg, 16> CalleeSavedRegisters;
+    for (const auto &RegSource : YamlMF.CalleeSavedRegisters.getValue()) {
+      unsigned Reg = 0;
+      if (parseNamedRegisterReference(PFS, Reg, RegSource.Value, Error))
+        return error(Error, RegSource.SourceRange);
+      CalleeSavedRegisters.push_back(Reg);
+    }
+    RegInfo.setCalleeSavedRegs(CalleeSavedRegisters);
   }
-  RegInfo.setUsedPhysRegMask(CalleeSavedRegisterMask.flip());
+
   return false;
 }
 
@@ -510,14 +512,12 @@ bool MIRParserImpl::setupRegisterInfo(const PerFunctionMIParsingState &PFS,
   }
 
   // Compute MachineRegisterInfo::UsedPhysRegMask
-  if (!YamlMF.CalleeSavedRegisters) {
-    for (const MachineBasicBlock &MBB : MF) {
-      for (const MachineInstr &MI : MBB) {
-        for (const MachineOperand &MO : MI.operands()) {
-          if (!MO.isRegMask())
-            continue;
-          MRI.addPhysRegsUsedFromRegMask(MO.getRegMask());
-        }
+  for (const MachineBasicBlock &MBB : MF) {
+    for (const MachineInstr &MI : MBB) {
+      for (const MachineOperand &MO : MI.operands()) {
+        if (!MO.isRegMask())
+          continue;
+        MRI.addPhysRegsUsedFromRegMask(MO.getRegMask());
       }
     }
   }

lib/CodeGen/MIRPrinter.cpp

@@ -207,6 +207,25 @@ void MIRPrinter::print(const MachineFunction &MF) {
   Out << YamlMF;
 }
 
+static void printCustomRegMask(const uint32_t *RegMask, raw_ostream &OS,
+                               const TargetRegisterInfo *TRI) {
+  assert(RegMask && "Can't print an empty register mask");
+  OS << StringRef("CustomRegMask(");
+
+  bool IsRegInRegMaskFound = false;
+  for (int I = 0, E = TRI->getNumRegs(); I < E; I++) {
+    // Check whether the register is asserted in regmask.
+    if (RegMask[I / 32] & (1u << (I % 32))) {
+      if (IsRegInRegMaskFound)
+        OS << ',';
+      printReg(I, OS, TRI);
+      IsRegInRegMaskFound = true;
+    }
+  }
+
+  OS << ')';
+}
+
 void MIRPrinter::convert(yaml::MachineFunction &MF,
                          const MachineRegisterInfo &RegInfo,
                          const TargetRegisterInfo *TRI) {
@@ -241,20 +260,18 @@ void MIRPrinter::convert(yaml::MachineFunction &MF,
       printReg(I->second, LiveIn.VirtualRegister, TRI);
     MF.LiveIns.push_back(LiveIn);
   }
-  // The used physical register mask is printed as an inverted callee saved
-  // register mask.
-  const BitVector &UsedPhysRegMask = RegInfo.getUsedPhysRegsMask();
-  if (UsedPhysRegMask.none())
-    return;
-  std::vector<yaml::FlowStringValue> CalleeSavedRegisters;
-  for (unsigned I = 0, E = UsedPhysRegMask.size(); I != E; ++I) {
-    if (!UsedPhysRegMask[I]) {
+
+  // Prints the callee saved registers.
+  if (RegInfo.isUpdatedCSRsInitialized()) {
+    const MCPhysReg *CalleeSavedRegs = RegInfo.getCalleeSavedRegs();
+    std::vector<yaml::FlowStringValue> CalleeSavedRegisters;
+    for (const MCPhysReg *I = CalleeSavedRegs; *I; ++I) {
       yaml::FlowStringValue Reg;
-      printReg(I, Reg, TRI);
+      printReg(*I, Reg, TRI);
       CalleeSavedRegisters.push_back(Reg);
     }
+    MF.CalleeSavedRegisters = CalleeSavedRegisters;
   }
-  MF.CalleeSavedRegisters = CalleeSavedRegisters;
 }
 
 void MIRPrinter::convert(ModuleSlotTracker &MST,
@@ -862,7 +879,7 @@ void MIPrinter::print(const MachineOperand &Op, const TargetRegisterInfo *TRI,
     if (RegMaskInfo != RegisterMaskIds.end())
       OS << StringRef(TRI->getRegMaskNames()[RegMaskInfo->second]).lower();
     else
-      llvm_unreachable("Can't print this machine register mask yet.");
+      printCustomRegMask(Op.getRegMask(), OS, TRI);
     break;
   }
   case MachineOperand::MO_RegisterLiveOut: {

lib/CodeGen/MachineRegisterInfo.cpp

@@ -589,3 +589,15 @@ const MCPhysReg *MachineRegisterInfo::getCalleeSavedRegs() const {
   return getTargetRegisterInfo()->getCalleeSavedRegs(MF);
 }
 
+void MachineRegisterInfo::setCalleeSavedRegs(ArrayRef<MCPhysReg> CSRs) {
+  if (IsUpdatedCSRsInitialized)
+    UpdatedCSRs.clear();
+
+  for (MCPhysReg Reg : CSRs)
+    UpdatedCSRs.push_back(Reg);
+
+  // Zero value represents the end of the register list
+  // (no more registers should be pushed).
+  UpdatedCSRs.push_back(0);
+  IsUpdatedCSRsInitialized = true;
+}

lib/Target/X86/X86ISelLowering.cpp

@@ -3816,8 +3816,6 @@ X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   // In some calling conventions we need to remove the used physical registers
   // from the reg mask.
   if (CallConv == CallingConv::X86_RegCall) {
-    /// @todo Need to add support in MIPrinter and MIParser to represent
-    ///       the customed RegMask.
     const TargetRegisterInfo *TRI = Subtarget.getRegisterInfo();
 
     // Allocate a new Reg Mask and copy Mask.

test/CodeGen/MIR/Generic/dynamic-regmask.ll

@@ -0,0 +1,30 @@
+; RUN: llc -stop-after machine-sink %s -o %t.mir
+; RUN: FileCheck %s < %t.mir
+; RUN: llc %t.mir -run-pass machine-sink
+; Check that callee saved registers are printed in a format that can then be parsed.
+
+declare x86_regcallcc i32 @callee(i32 %a0, i32 %b0, i32 %c0, i32 %d0, i32 %e0)
+
+define i32 @caller(i32 %a0) nounwind {
+  %b1 = call x86_regcallcc i32 @callee(i32 %a0, i32 %a0, i32 %a0, i32 %a0, i32 %a0)
+  %b2 = add i32 %b1, %a0
+  ret i32 %b2
+}
+; CHECK:    name: caller
+; CHECK:    CALL64pcrel32 @callee, CustomRegMask(%bh,%bl,%bp,%bpl,%bx,%ebp,%ebx,%esp,%rbp,%rbx,%rsp,%sp,%spl,%r10,%r11,%r12,%r13,%r14,%r15,%xmm8,%xmm9,%xmm10,%xmm11,%xmm12,%xmm13,%xmm14,%xmm15,%r10b,%r11b,%r12b,%r13b,%r14b,%r15b,%r10d,%r11d,%r12d,%r13d,%r14d,%r15d,%r10w,%r11w,%r12w,%r13w,%r14w,%r15w)
+; CHECK:    RET 0, %eax
+
+define x86_regcallcc {i32, i32, i32} @test_callee(i32 %a0, i32 %b0, i32 %c0, i32 %d0, i32 %e0) nounwind {
+  %b1 = mul i32 7, %e0
+  %b2 = udiv i32 5, %e0
+  %b3 = mul i32 7, %d0
+  %b4 = insertvalue {i32, i32, i32} undef, i32 %b1, 0
+  %b5 = insertvalue {i32, i32, i32} %b4, i32 %b2, 1
+  %b6 = insertvalue {i32, i32, i32} %b5, i32 %b3, 2
+  ret {i32, i32, i32} %b6
+}
+; CHECK: name:            test_callee
+; CHECK: calleeSavedRegisters: [ '%rbx', '%rbp', '%rsp', '%r10', '%r11', '%r12',
+; CHECK:                         '%r13', '%r14', '%r15', '%xmm8', '%xmm9', '%xmm10',
+; CHECK:                         '%xmm11', '%xmm12', '%xmm13', '%xmm14', '%xmm15' ]
+; CHECK: RET 0, %eax, %ecx, %edx

test/CodeGen/MIR/X86/used-physical-register-info.mir

@@ -1,109 +0,0 @@
-# RUN: llc -march=x86-64 -run-pass none -o - %s | FileCheck %s
-# This test ensures that the MIR parser parses the callee saved register mask
-# correctly and that the MIR parser can infer it as well.
-
---- |
-
-  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
-  }
-
-  define i32 @bar(i32 %a) #0 {
-  entry:
-    %b = call i32 @compute(i32 %a)
-    ret i32 %b
-  }
-
-  define i32 @empty(i32 %a) #0 {
-  entry:
-    %b = call i32 @compute(i32 %a)
-    ret i32 %b
-  }
-
-  attributes #0 = { "no-frame-pointer-elim"="false" }
-
-...
----
-# CHECK: name: compute
-# CHECK: liveins:
-# CHECK-NEXT: - { reg: '%edi' }
-# CHECK-NEXT: frameInfo:
-name:            compute
-liveins:
-  - { reg: '%edi' }
-frameInfo:
-  stackSize:     8
-body: |
-  bb.0.body:
-    liveins: %edi
-
-    %eax = IMUL32rri8 %edi, 11, implicit-def %eflags
-    RETQ %eax
-...
----
-name:            foo
-liveins:
-  - { reg: '%edi' }
-# CHECK: name: foo
-# CHECK: calleeSavedRegisters: [ '%bh', '%bl', '%bp', '%bpl', '%bx', '%ebp', '%ebx',
-# CHECK-NEXT:                    '%rbp', '%rbx', '%r12', '%r13', '%r14', '%r15',
-# CHECK-NEXT:                    '%r12b', '%r13b', '%r14b', '%r15b', '%r12d', '%r13d',
-# CHECK-NEXT:                    '%r14d', '%r15d', '%r12w', '%r13w', '%r14w', '%r15w' ]
-calleeSavedRegisters: [ '%bh', '%bl', '%bp', '%bpl', '%bx', '%ebp', '%ebx',
-                        '%rbp', '%rbx', '%r12', '%r13', '%r14', '%r15',
-                        '%r12b', '%r13b', '%r14b', '%r15b', '%r12d', '%r13d',
-                        '%r14d', '%r15d', '%r12w', '%r13w', '%r14w', '%r15w' ]
-body: |
-  bb.0.entry:
-    liveins: %edi
-
-    PUSH64r %rax, implicit-def %rsp, implicit %rsp
-    CALL64pcrel32 @compute, csr_64, implicit %rsp, implicit %edi, implicit-def %rsp, implicit-def %eax
-    %rdx = POP64r implicit-def %rsp, implicit %rsp
-    RETQ %eax
-...
----
-name:            bar
-liveins:
-  - { reg: '%edi' }
-# Verify that the callee saved register can be inferred from register mask
-# machine operands:
-# CHECK: name: bar
-# CHECK: calleeSavedRegisters: [ '%bh', '%bl', '%bp', '%bpl', '%bx', '%ebp', '%ebx',
-# CHECK-NEXT:                    '%rbp', '%rbx', '%r12', '%r13', '%r14', '%r15',
-# CHECK-NEXT:                    '%r12b', '%r13b', '%r14b', '%r15b', '%r12d', '%r13d',
-# CHECK-NEXT:                    '%r14d', '%r15d', '%r12w', '%r13w', '%r14w', '%r15w' ]
-body: |
-  bb.0.entry:
-    liveins: %edi
-
-    PUSH64r %rax, implicit-def %rsp, implicit %rsp
-    CALL64pcrel32 @compute, csr_64, implicit %rsp, implicit %edi, implicit-def %rsp, implicit-def %eax
-    %rdx = POP64r implicit-def %rsp, implicit %rsp
-    RETQ %eax
-...
----
-name:            empty
-liveins:
-  - { reg: '%edi' }
-# Verify that the callee saved register can be empty.
-# CHECK: name: empty
-# CHECK: calleeSavedRegisters: [ ]
-calleeSavedRegisters: [ ]
-body: |
-  bb.0.entry:
-    liveins: %edi
-
-    PUSH64r %rax, implicit-def %rsp, implicit %rsp
-    CALL64pcrel32 @compute, csr_64, implicit %rsp, implicit %edi, implicit-def %rsp, implicit-def %eax
-    %rdx = POP64r implicit-def %rsp, implicit %rsp
-    RETQ %eax
-...