[RegAllocFast] properly handle STATEPOINT instruction.

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.

Reviewed By: arsenm, skatkov
Differential Revision: https://reviews.llvm.org/D99284

llvm/lib/CodeGen/RegAllocFast.cpp

@@ -147,6 +147,8 @@ namespace {
     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;
@@ -157,8 +159,17 @@ namespace {
         UsedInInstr.insert(*Units);
     }
 
+    // Check if physreg is clobbered by instruction's regmask(s).
+    bool isClobberedByRegMasks(MCPhysReg PhysReg) const {
+      return llvm::any_of(RegMasks, [PhysReg](const uint32_t *Mask) {
+        return MachineOperand::clobbersPhysReg(Mask, PhysReg);
+      });
+    }
+
     /// 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;
@@ -1088,6 +1099,7 @@ void RegAllocFast::allocateInstruction(MachineInstr &MI) {
   //   operands and early-clobbers.
 
   UsedInInstr.clear();
+  RegMasks.clear();
   BundleVirtRegsMap.clear();
 
   // Scan for special cases; Apply pre-assigned register defs to state.
@@ -1127,6 +1139,7 @@ void RegAllocFast::allocateInstruction(MachineInstr &MI) {
       }
     } else if (MO.isRegMask()) {
       HasRegMask = true;
+      RegMasks.push_back(MO.getRegMask());
     }
   }
 
@@ -1242,6 +1255,9 @@ void RegAllocFast::allocateInstruction(MachineInstr &MI) {
         continue;
       }
 
+      assert((!MO.isTied() || !isClobberedByRegMasks(MO.getReg())) &&
+             "tied def assigned to clobbered register");
+
       // Do not free tied operands and early clobbers.
       if (MO.isTied() || MO.isEarlyClobber())
         continue;
@@ -1258,19 +1274,16 @@ void RegAllocFast::allocateInstruction(MachineInstr &MI) {
 
   // Displace clobbered registers.
   if (HasRegMask) {
-    for (const MachineOperand &MO : MI.operands()) {
-      if (MO.isRegMask()) {
-        // MRI bookkeeping.
-        MRI->addPhysRegsUsedFromRegMask(MO.getRegMask());
-
-        // 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);
-        }
-      }
+    assert(!RegMasks.empty() && "expected RegMask");
+    // MRI bookkeeping.
+    for (const auto *RM : RegMasks)
+      MRI->addPhysRegsUsedFromRegMask(RM);
+
+    // Displace clobbered registers.
+    for (const LiveReg &LR : LiveVirtRegs) {
+      MCPhysReg PhysReg = LR.PhysReg;
+      if (PhysReg != 0 && isClobberedByRegMasks(PhysReg))
+        displacePhysReg(MI, PhysReg);
     }
   }
 

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

@@ -0,0 +1,40 @@
+# 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
+...
+
+# Same as above but with multiple RegMask operands per instruction.
+# These regmasks have no real meaning and chosen to allow only single register to be assignable ($r12)
+---
+name:            test_relocate_multi_regmasks
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $rdi
+
+    ; CHECK: renamable $r12 = STATEPOINT 0, 0, 0, target-flags(x86-plt) 0, 2, 0, 2, 0, 2, 0, 2, 1, renamable $r12(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_rt_allregs, csr_64_hhvm, 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
+...