[RegAllocFast] properly handle STATEPOINT instruction.

[dantrushin]

Summary

STATEPOINT is a fancy and complex pseudo instruction which
has both tied defs and regmask operand.

Basic FastRA algorithm is as follows:

1. Mark registers used by defs as free
2. If instruction has regmask operand displace clobbered registers according to regmask.
3. Assign registers for use operands.

In case of tied defs step 1 is replaced with allocation of registers
for them. But regmask is still processed, which may displace already
allocated registers. As a result, tied use and def will get assigned
to different registers.

This patch makes FastRA to process instruction's RegMask (if any) when
checking for physical registers interference.
That way tied operands won't get registers clobbered by regmask.

[skatkov]

Do I understand correctly that with this change tied pair of registers can be assigned to caller saved register?

This is not a correct generally. For statepoint you can probably rely on FixedCallerSaveRegister pass.

[arsenm]

I don't think this should special case statepoints. Any regmask with tied operands should behave the same way

llvm/test/CodeGen/X86/statepoint-fastregalloc.mir

4–8

Can remove this. Should also have a comment explaining the test

26–67

Don't need most of this

[dantrushin]

Address review comments:

• Do this for all tied operands with regmask instructions;
• Strip all unnecessary data from test, add comment;

[arsenm]

llvm/lib/CodeGen/RegAllocFast.cpp

1234

I don't like relying on findLiveVirtReg getting called a second time returning the same result as during the actual assignment

1265

Can you track the register from here instead?

1280–1281

Should get a comment explaining the tied behavior

1281

Checking for exact register matches makes me nervous because of register aliases / subregisters. What if you had something like

FOO tied-def $reg0_reg1, tied-use $reg0_reg1, mask-clobber(reg1)?

What is the correct behavior here and does it matter? Should it be a verifier error?

[skatkov]

I believe this is not right fix for general problem of handling tied-defs wrt reg mask.

Firs of all, running your test I see that tied def/use now uses the same register rax but it is clabbered by the call. So compilation is not correct at all.
As I said before for statepoint instruction you can rely on post fix-up of this case but it is better to get a full support in fast regalloc itself.

I guess that the right way to fix this should look as follows:

1. defineLiveThroughVirtReg for tied-def should check not only whether the register is used in instruction but also it is not in regmask.
2. allocVirtReg should be aware about reg mask and tied operands to avoid allocation of register which is in reg mask.

[dantrushin]

Rework algorithm a bit:

• Store intruction RegMask (if any) in a global;
• Look at RegMask when allocating register if we care about physical register use. (Re)Use LookAtPhysRegUses for this.

With this change, registers clobbered by regmask should not be assigned
to tied defs and AssignedTiedDefs set is not strictly needed. I use
it for assertion check, but can eliminate it.

[dantrushin]

llvm/lib/CodeGen/RegAllocFast.cpp

1281

My understanding is that RegMask containts all registers and subregisters.
I'm not so sure about aliases.
As for verifier, it was recently enhanced to better track tied registers. There is also in-progress review D97054 which tries to enhance it further...

[skatkov]

lgtm. Please wait for @arsenm feedback.

llvm/lib/CodeGen/RegAllocFast.cpp

1148

you do not need to collect all tied-defs assigned. As soon as it is under an assert, you can just iterate over defs of MI and check that if it is def it is not assigned to clobbered register.
Moreover you can move an assert to end of an allocation and also check that tied-use is allocated to the same register.

If you decide to leave it in this way - please ensure that this declaration and inserting in the set is not used in product mode.

1286

consider usage of all_of instead. So you will not need #ifdef then.

[skatkov]

please update the commit message.

[skatkov]

BTW, it seems that previous code was non-optimal or it supports that there are may be ore than one machine operand containing the reg mask.
If the second is true, you should support this in your patch. Can you check that?

[dantrushin]

In D99284#2665877, @skatkov wrote:

BTW, it seems that previous code was non-optimal or it supports that there are may be ore than one machine operand containing the reg mask.
If the second is true, you should support this in your patch. Can you check that?

I've never seen instructions with multiple RegMasks and I do not see a reason for multiple RegMasks. This is bitfield big enough to hold all target's physical registers.
After grepping LLVM sources I have no definitive answer either. One can't tell code handling multiple masks from code handling just one, but iterating over all operands just because there is no other way to find regmask (it has no fixed position in MI).

@arsenm : what do you think?

[skatkov]

ok, I took a look in different usage of MO.getRegMask and it looks like several regmask may happen: see InsteRefBasedImpl.cpp.
So it would be better to re-write the code to support several regmasks - it should not be too complex.

[dantrushin]

Support multiple RegMasks per instruction.

[skatkov]

llvm/lib/CodeGen/RegAllocFast.cpp

164

any_of?

1148

you do not need to collect all tied-defs assigned. As soon as it is under an assert, you can just iterate over defs of MI and check that if it is def it is not assigned to clobbered register.
Moreover you can move an assert to end of an allocation and also check that tied-use is allocated to the same register.

If you decide to leave it in this way - please ensure that this declaration and inserting in the set is not used in product mode.

Did not get any answer on this comment.

1288

all_of?

[dantrushin]

Address comments:

• use any_of
• get rid of AssignedTiedDefs vector

[skatkov]

@arsenm, any comments from your side?

llvm/lib/CodeGen/RegAllocFast.cpp

1287

nit (up to you):

for (const LiveReg &LR : LiveVirtRegs)
  if (MCPhysReg PhysReg = LR.PhysReg)
    if (isClobberedByRegMasks(PhysReg))
      displacePhysReg(MI, PhysReg);

[dantrushin]

@arsenm, Is this OK for you?

[arsenm]

llvm/test/CodeGen/X86/statepoint-fastregalloc.mir

26–67

Can you craft a testcase for multiple regmasks? Other than that LGTM

[dantrushin]

llvm/test/CodeGen/X86/statepoint-fastregalloc.mir

26–67

I don't know any instruction which can have multiple reg masks.
I can try to add another arbitrary reg mask to STATEPOINT instruction in this test and see if it'll work (and will pass through verifier).
Would it be OK?

[arsenm]

llvm/test/CodeGen/X86/statepoint-fastregalloc.mir

26–67

Regmasks are treated like implicit operands, so you can freely add them to any instruction. You could take a NOP and add whatever implicit uses/defs and regmasks you want

[dantrushin]

Added test for instruction with two RegMask operands.

Still it is not quite clear to me what multiple RegMasks mean?
This patch basically ANDs them.

llvm/lib/CodeGen/RegAllocFast.cpp

private:
     SmallVector<MachineInstr *, 32> Coalesced;
 
     using RegUnitSet = SparseSet<uint16_t, identity<uint16_t>>;
     /// Set of register units that are used in the current instruction, and so
     /// cannot be allocated.
     RegUnitSet UsedInInstr;
     RegUnitSet PhysRegUses;
     SmallVector<uint16_t, 8> DefOperandIndexes;
+    // Register masks attached to the current instruction.
+    SmallVector<const uint32_t *> RegMasks;
 
     void setPhysRegState(MCPhysReg PhysReg, unsigned NewState);
     bool isPhysRegFree(MCPhysReg PhysReg) const;
 
     /// Mark a physreg as used in this instruction.
     void markRegUsedInInstr(MCPhysReg PhysReg) {
       for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units)
         UsedInInstr.insert(*Units);
     }
 
+    // Check if physreg is clobbered by instruction's regmask(s).
+    bool isClobberedByRegMasks(MCPhysReg PhysReg) const {
+      for (const auto *RM : RegMasks)
+        if (MachineOperand::clobbersPhysReg(RM, PhysReg))
+          return true;
+      return false;
+    }
+
     /// Check if a physreg or any of its aliases are used in this instruction.
     bool isRegUsedInInstr(MCPhysReg PhysReg, bool LookAtPhysRegUses) const {
+      if (LookAtPhysRegUses && isClobberedByRegMasks(PhysReg))
+        return true;
       for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
         if (UsedInInstr.count(*Units))
           return true;
         if (LookAtPhysRegUses && PhysRegUses.count(*Units))
           return true;
       }
       return false;
     }
void RegAllocFast::allocateInstruction(MachineInstr &MI) {
   // However we need to handle some corner cases:
   // - pre-assigned defs and uses need to be handled before the other def/use
   //   operands are processed to avoid the allocation heuristics clashing with
   //   the pre-assignment.
   // - The "free def operands" step has to come last instead of first for tied
   //   operands and early-clobbers.
 
   UsedInInstr.clear();
+  RegMasks.clear();
   BundleVirtRegsMap.clear();
 
   // Scan for special cases; Apply pre-assigned register defs to state.
   bool HasPhysRegUse = false;
   bool HasRegMask = false;
   bool HasVRegDef = false;
   bool HasDef = false;
   bool HasEarlyClobber = false;
if (MO.isReg()) {
               MO.setIsDead(true);
           }
           if (MO.readsReg())
             HasPhysRegUse = true;
         }
       }
     } else if (MO.isRegMask()) {
       HasRegMask = true;
+      RegMasks.push_back(MO.getRegMask());
     }
   }
 
+  // Set of registers allocated to tied defs.
+  SmallSet<Register, 16> AssignedTiedDefs;
+
   // Allocate virtreg defs.
   if (HasDef) {
     if (HasVRegDef) {
       // Special handling for early clobbers, tied operands or subregister defs:
       // Compared to "normal" defs these:
       // - Must not use a register that is pre-assigned for a use operand.
       // - In order to solve tricky inline assembly constraints we change the
       //   heuristic to figure out a good operand order before doing
if (HasVRegDef) {
           unsigned Reg = MO.getReg();
           if (MO.isEarlyClobber() || MO.isTied() ||
               (MO.getSubReg() && !MO.isUndef())) {
             defineLiveThroughVirtReg(MI, OpIdx, Reg);
           } else {
             defineVirtReg(MI, OpIdx, Reg);
           }
         }
       } else {
         // Assign virtual register defs.
         for (unsigned I = 0, E = MI.getNumOperands(); I < E; ++I) {
           MachineOperand &MO = MI.getOperand(I);
           if (!MO.isReg() || !MO.isDef())
             continue;
           Register Reg = MO.getReg();
           if (Reg.isVirtual())
             defineVirtReg(MI, I, Reg);
for (unsigned I = MI.getNumOperands(); I-- > 0;) {
 
       // subreg defs don't free the full register. We left the subreg number
       // around as a marker in setPhysReg() to recognize this case here.
       if (MO.getSubReg() != 0) {
         MO.setSubReg(0);
         continue;
       }
 
+      if (MO.isTied()) {
+        AssignedTiedDefs.insert(MO.getReg());
+      }
       // Do not free tied operands and early clobbers.
       if (MO.isTied() || MO.isEarlyClobber())
         continue;
       Register Reg = MO.getReg();
       if (!Reg)
         continue;
       assert(Reg.isPhysical());
       if (MRI->isReserved(Reg))
         continue;
       freePhysReg(Reg);
       unmarkRegUsedInInstr(Reg);
     }
   }
 
   // Displace clobbered registers.
   if (HasRegMask) {
-    for (const MachineOperand &MO : MI.operands()) {
-      if (MO.isRegMask()) {
-        // MRI bookkeeping.
-        MRI->addPhysRegsUsedFromRegMask(MO.getRegMask());
+    assert(!RegMasks.empty() && "expected RegMask");
+    for (const auto *RM : RegMasks)
+      MRI->addPhysRegsUsedFromRegMask(RM);
+
+#ifndef NDEBUG
+    for (auto R : AssignedTiedDefs) {
+      assert(!isClobberedByRegMasks(R) &&
+             "tied def allocated to clobbered register");
+    }
+#endif
 
-        // Displace clobbered registers.
-        const uint32_t *Mask = MO.getRegMask();
-        for (const LiveReg &LR : LiveVirtRegs) {
-          MCPhysReg PhysReg = LR.PhysReg;
-          if (PhysReg != 0 && MachineOperand::clobbersPhysReg(Mask, PhysReg))
-            displacePhysReg(MI, PhysReg);
-        }
-      }
-    }
-  }
+      if (PhysReg != 0 && isClobberedByRegMasks(PhysReg))
 
   // Apply pre-assigned register uses to state.
   if (HasPhysRegUse) {
     for (MachineOperand &MO : MI.operands()) {
       if (!MO.isReg() || !MO.readsReg())
         continue;
       Register Reg = MO.getReg();
       if (!Reg.isPhysical())

llvm/test/CodeGen/X86/statepoint-fastregalloc.mir

+# RUN: llc -mtriple=x86_64-- -run-pass=regallocfast -o - %s | FileCheck %s
+
+# Check that fastregalloc does not displace register assigned to tied def when
+# RegMask operand is present. STATEPOINT is an example of such instruction.
+# Tied def/use must be assigned to the same register.
+---
+name:            test_relocate
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $rdi
+
+    ; CHECK: renamable [[REG:\$[a-z0-9]+]] = STATEPOINT 0, 0, 0, target-flags(x86-plt) 0, 2, 0, 2, 0, 2, 0, 2, 1, renamable [[REG]](tied-def 0)
+  
+    %1:gr64 = COPY $rdi
+    ADJCALLSTACKDOWN64 0, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    %1:gr64 = STATEPOINT 0, 0, 0, target-flags(x86-plt) 0, 2, 0, 2, 0, 2, 0, 2, 1, %1(tied-def 0), 2, 0, 2, 1, 0, 0, csr_64, implicit-def $rsp, implicit-def $ssp
+    ADJCALLSTACKUP64 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    $rax = COPY %1
+    RET 0, killed $rax
+...