[MIR] Support Customed Register Mask and CSRs

[oren_ben_simhon]

Summary

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.

[MatzeB]

Thanks for working on this!

• Printing the regmask as hexnumbers seems like a bad idea: Adding or removing registers in between would render old .mir files invalid. Even though it's unwieldy I think we should print lists of register names instead. Printing the list of preserved registers instead of the clobbered ones may also make the lists shorter.

• I only very recently learned about the thing that was called "CalleeSavedRegisters" (or UnusedPhysRegs after your patch). It is silly that we serialize this information in the first place as you can trivially recompute it after loading a mir file: Loop over all blocks, instructions, operands and call MachineRegisterInfo::addPhysRegsUsedFromRegMask on every regmask operand you find. That way we have correct information and don't have to deal with printing and parsing this information in the first place. Maybe you could change this as part of this patch?

• Browsing around I found "@todo Need to add support in MIPrinter and MIParser to represent..." in the middle of X86 code which you can remove with this.

include/llvm/CodeGen/MachineRegisterInfo.h

78–82 ↗

(On Diff #91837)

Unrelated and so trivial that you can just commmit this independently without review.

lib/CodeGen/MIRParser/MIParser.cpp

1676 ↗

(On Diff #91837)

Writing the type instead of auto is friendlier to readers (unless the type is immediately obvious, which it is not here).

1682 ↗

(On Diff #91837)

typo

1683 ↗

(On Diff #91837)

Move the declaration to where it is used to avoid the impression that it needs to hold values accross loop iterations.

lib/CodeGen/MIRPrinter.cpp

889–892 ↗

(On Diff #91837)

RegMask operands with a nullptr in the masks are illegal and it is perfectly reasonable not to test for that here and just crash in that case.

lib/CodeGen/MachineRegisterInfo.cpp

593 ↗

(On Diff #91837)

Use ArrayRef<MCPhysReg> here for increased flexibility (smallvectors are among the things that automatically convert to ArrayRefs).

597 ↗

(On Diff #91837)

Please use MCPhysReg instead of auto.

[oren_ben_simhon]

Implemented comments posted until 03/15 (Thank you Matthias).

[MatzeB]

LGTM, thanks.

lib/CodeGen/MIRParser/MIRParser.cpp

513–523 ↗

(On Diff #91995)

Oh wow, we even already had the code to compute the UsedPhysRegMask :)

llvm/trunk/include/llvm/CodeGen/MachineRegisterInfo.h

public:
   bool subRegLivenessEnabled() const {
     return TracksSubRegLiveness;
   }
 
   //===--------------------------------------------------------------------===//
   // 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.
   void disableCalleeSavedRegister(unsigned Reg);
 
   /// Returns list of callee saved registers.
   /// The function returns the updated CSR list (after taking into account
   /// 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);
 
   // Strictly for use by MachineInstr.cpp.
   void removeRegOperandFromUseList(MachineOperand *MO);
 
   // Strictly for use by MachineInstr.cpp.
   void moveOperands(MachineOperand *Dst, MachineOperand *Src, unsigned NumOps);
#endif
   /// addPhysRegsUsedFromRegMask - Mark any registers not in RegMask as used.
   /// This corresponds to the bit mask attached to register mask operands.
   void addPhysRegsUsedFromRegMask(const uint32_t *RegMask) {
     UsedPhysRegMask.setBitsNotInMask(RegMask);
   }
 
   const BitVector &getUsedPhysRegsMask() const { return UsedPhysRegMask; }
 
-  void setUsedPhysRegMask(BitVector &Mask) { UsedPhysRegMask = Mask; }
-
   //===--------------------------------------------------------------------===//
   // Reserved Register Info
   //===--------------------------------------------------------------------===//
   //
   // The set of reserved registers must be invariant during register
   // allocation.  For example, the target cannot suddenly decide it needs a
   // frame pointer when the register allocator has already used the frame
   // pointer register for something else.

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

public:
   bool parseCFIOffset(int &Offset);
   bool parseCFIRegister(unsigned &Reg);
   bool parseCFIOperand(MachineOperand &Dest);
   bool parseIRBlock(BasicBlock *&BB, const Function &F);
   bool parseBlockAddressOperand(MachineOperand &Dest);
   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);
   bool parseMachineOperandAndTargetFlags(MachineOperand &Dest,
                                          Optional<unsigned> &TiedDefIdx);
   bool parseOffset(int64_t &Offset);
   bool parseAlignment(unsigned &Alignment);
   bool parseOperandsOffset(MachineOperand &Op);
bool MIParser::parseTargetIndexOperand(MachineOperand &Dest) {
   if (expectAndConsume(MIToken::rparen))
     return true;
   Dest = MachineOperand::CreateTargetIndex(unsigned(Index), /*Offset=*/0);
   if (parseOperandsOffset(Dest))
     return true;
   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();
   assert(TRI && "Expected target register info");
   uint32_t *Mask = MF.allocateRegisterMask(TRI->getNumRegs());
   lex();
   if (expectAndConsume(MIToken::lparen))
     return true;
case MIToken::kw_intpred:
     return parsePredicateOperand(Dest);
   case MIToken::Error:
     return true;
   case MIToken::Identifier:
     if (const auto *RegMask = getRegMask(Token.stringValue())) {
       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");
   }
   return false;
 }
 
 bool MIParser::parseMachineOperandAndTargetFlags(

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

if (!LiveIn.VirtualRegister.Value.empty()) {
       if (parseVirtualRegisterReference(PFS, Info, LiveIn.VirtualRegister.Value,
                                         Error))
         return error(Error, LiveIn.VirtualRegister.SourceRange);
       VReg = Info->VReg;
     }
     RegInfo.addLiveIn(Reg, VReg);
   }
 
-  // Parse the callee saved register mask.
-  BitVector CalleeSavedRegisterMask(RegInfo.getUsedPhysRegsMask().size());
-  if (!YamlMF.CalleeSavedRegisters)
-    return false;
+  // 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);
-    CalleeSavedRegisterMask[Reg] = true;
+      CalleeSavedRegisters.push_back(Reg);
+    }
+    RegInfo.setCalleeSavedRegs(CalleeSavedRegisters);
   }
-  RegInfo.setUsedPhysRegMask(CalleeSavedRegisterMask.flip());
+
   return false;
 }
 
 bool MIRParserImpl::setupRegisterInfo(const PerFunctionMIParsingState &PFS,
                                       const yaml::MachineFunction &YamlMF) {
   MachineFunction &MF = PFS.MF;
   MachineRegisterInfo &MRI = MF.getRegInfo();
   bool Error = false;
case VRegInfo::GENERIC:
       break;
     case VRegInfo::REGBANK:
       MRI.setRegBank(Reg, *Info.D.RegBank);
       break;
     }
   }
 
   // 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());
-        }
-      }
-    }
-  }
 
   // FIXME: This is a temporary workaround until the reserved registers can be
   // serialized.
   MRI.freezeReservedRegs(MF);
   return Error;
 }
 
 bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,

llvm/trunk/lib/CodeGen/MIRPrinter.cpp

MIPrinter(StrOS, MST, RegisterMaskIds, StackObjectOperandMapping)
         .print(MBB);
     IsNewlineNeeded = true;
   }
   StrOS.flush();
   yaml::Output Out(OS);
   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) {
   MF.TracksRegLiveness = RegInfo.tracksLiveness();
 
   // Print the virtual register definitions.
   for (unsigned I = 0, E = RegInfo.getNumVirtRegs(); I < E; ++I) {
     unsigned Reg = TargetRegisterInfo::index2VirtReg(I);
void MIRPrinter::convert(yaml::MachineFunction &MF,
   // Print the live ins.
   for (auto I = RegInfo.livein_begin(), E = RegInfo.livein_end(); I != E; ++I) {
     yaml::MachineFunctionLiveIn LiveIn;
     printReg(I->first, LiveIn.Register, TRI);
     if (I->second)
       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;
+
+  // Prints the callee saved registers.
+  if (RegInfo.isUpdatedCSRsInitialized()) {
+    const MCPhysReg *CalleeSavedRegs = RegInfo.getCalleeSavedRegs();
     std::vector<yaml::FlowStringValue> CalleeSavedRegisters;
-  for (unsigned I = 0, E = UsedPhysRegMask.size(); I != E; ++I) {
-    if (!UsedPhysRegMask[I]) {
+    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;
   }
+}
 
 void MIRPrinter::convert(ModuleSlotTracker &MST,
                          yaml::MachineFrameInfo &YamlMFI,
                          const MachineFrameInfo &MFI) {
   YamlMFI.IsFrameAddressTaken = MFI.isFrameAddressTaken();
   YamlMFI.IsReturnAddressTaken = MFI.isReturnAddressTaken();
   YamlMFI.HasStackMap = MFI.hasStackMap();
   YamlMFI.HasPatchPoint = MFI.hasPatchPoint();
case MachineOperand::MO_BlockAddress:
     OS << ')';
     printOffset(Op.getOffset());
     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.");
+      printCustomRegMask(Op.getRegMask(), OS, TRI);
     break;
   }
   case MachineOperand::MO_RegisterLiveOut: {
     const uint32_t *RegMask = Op.getRegLiveOut();
     OS << "liveout(";
     bool IsCommaNeeded = false;
     for (unsigned Reg = 0, E = TRI->getNumRegs(); Reg < E; ++Reg) {
       if (RegMask[Reg / 32] & (1U << (Reg % 32))) {

llvm/trunk/lib/CodeGen/MachineRegisterInfo.cpp

 
 const MCPhysReg *MachineRegisterInfo::getCalleeSavedRegs() const {
   if (IsUpdatedCSRsInitialized)
     return UpdatedCSRs.data();
 
   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;
+}

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   }
 
   // Define a new register mask from the existing mask.
   uint32_t *RegMask = nullptr;
 
   // 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.
     RegMask = MF.allocateRegisterMask(TRI->getNumRegs());
     unsigned RegMaskSize = (TRI->getNumRegs() + 31) / 32;
     memcpy(RegMask, Mask, sizeof(uint32_t) * RegMaskSize);
 
     // Make sure all sub registers of the argument registers are reset

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

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

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

-# 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
-...