[RegAlloc] Relax the split constrain on MBB prolog (#168259)

https://reviews.llvm.org/D52052 is to prevent register split on the MBB
which have prolog instructions defining the exec register (or mask register
that activate the threads of a warp in GPU). The constrain seems too
strict, because 1) If the split is allowed, it may fit the free live range
of a physical register, and no spill will happen; 2) The register class of
register that is under splitting may not be the same to the register that
is defined in prolog, so there is no interference with the register being
defined in prolog. 
The current code has another small issue. The MBB->getFirstNonDebugInstr()
just skip debug instructions, but SA->getFirstSplitPoint(Number) would skip
label and phi instructions. This cause some MBB with label instruction
being taken as prolog.
This patch is to relax the split constrain on MMB with prolog by checking
if the register defined in prolog has the common register class with the
register being split. It allow the split if the register defined in prolog
is physical register or there is no common register class.

---------

Co-authored-by: Yuanke Luo <[email protected]>

Author: LuoYuanke

llvm/lib/CodeGen/RegAllocGreedy.cpp

@@ -774,8 +774,7 @@ bool RAGreedy::addSplitConstraints(InterferenceCache::Cursor Intf,
       // Abort if the spill cannot be inserted at the MBB' start
       if (((BC.Entry == SpillPlacement::MustSpill) ||
            (BC.Entry == SpillPlacement::PrefSpill)) &&
-          SlotIndex::isEarlierInstr(BI.FirstInstr,
-                                    SA->getFirstSplitPoint(BC.Number)))
+          !SA->canSplitBeforeProlog(BC.Number))
         return false;
     }
 
@@ -830,11 +829,7 @@ bool RAGreedy::addThroughConstraints(InterferenceCache::Cursor Intf,
     BCS[B].Number = Number;
 
     // Abort if the spill cannot be inserted at the MBB' start
-    MachineBasicBlock *MBB = MF->getBlockNumbered(Number);
-    auto FirstNonDebugInstr = MBB->getFirstNonDebugInstr();
-    if (FirstNonDebugInstr != MBB->end() &&
-        SlotIndex::isEarlierInstr(LIS->getInstructionIndex(*FirstNonDebugInstr),
-                                  SA->getFirstSplitPoint(Number)))
+    if (!SA->canSplitBeforeProlog(Number))
       return false;
     // Interference for the live-in value.
     if (Intf.first() <= Indexes->getMBBStartIdx(Number))

llvm/lib/CodeGen/SplitKit.cpp

@@ -147,6 +147,54 @@ InsertPointAnalysis::getLastInsertPointIter(const LiveInterval &CurLI,
   return LIS.getInstructionFromIndex(LIP);
 }
 
+bool InsertPointAnalysis::canSplitBeforeProlog(const LiveInterval &CurLI,
+                                               const MachineBasicBlock &MBB) {
+  const TargetInstrInfo *TII = MBB.getParent()->getSubtarget().getInstrInfo();
+
+  for (auto &MI : MBB) {
+    if (MI.isPHI() || MI.isPosition() || MI.isDebugInstr() ||
+        MI.isPseudoProbe())
+      continue;
+
+    if (!TII->isBasicBlockPrologue(MI))
+      return true;
+
+    for (auto &MO : MI.operands()) {
+      if (!MO.isReg() || !MO.isDef() || !MO.getReg().isVirtual())
+        continue;
+
+      // For the AMDGPU target if a MBB contains exec mask restore preamble,
+      // SplitEditor may get state when it cannot insert a spill instruction
+      // at the begin of the MBB.
+      // E.g. for a MIR
+      // bb.100:
+      //     %1 = S_OR_SAVEEXEC_B64 %2, implicit-def $exec, implicit-def $scc,
+      //          implicit $exec
+      //     ...
+      //     use %1
+      // If the regalloc try to allocate a virtreg to the physreg already
+      // assigned to virtreg %1 and the pyhsreg is computed as the best
+      // candidate for split, it may insert COPY instruction.
+      //  bb.100:
+      //     %1 = S_OR_SAVEEXEC_B64 %2, implicit-def $exec, implicit-def $scc,
+      //          implicit $exec
+      //     %2 = COPY %orig
+      //     ...
+      //     use %1
+      // Thus %1 and %orig still have interference. We may add cost for the
+      // physreg candidate or abandon the candidate.
+      const MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
+      const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
+      const TargetRegisterClass *RC = MRI.getRegClass(MO.getReg());
+      const TargetRegisterClass *CurRC = MRI.getRegClass(CurLI.reg());
+      if (TRI->getCommonSubClass(RC, CurRC))
+        return false;
+    }
+  }
+
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 //                                 Split Analysis
 //===----------------------------------------------------------------------===//

llvm/lib/CodeGen/SplitKit.h

@@ -89,6 +89,9 @@ class LLVM_LIBRARY_VISIBILITY InsertPointAnalysis {
     return Res;
   }
 
+  /// Return true if we can split \pCurLI before \pMBB's prolog.
+  bool canSplitBeforeProlog(const LiveInterval &CurLI,
+                            const MachineBasicBlock &MBB);
 };
 
 /// SplitAnalysis - Analyze a LiveInterval, looking for live range splitting
@@ -247,6 +250,11 @@ class LLVM_LIBRARY_VISIBILITY SplitAnalysis {
   SlotIndex getFirstSplitPoint(unsigned Num) {
     return IPA.getFirstInsertPoint(*MF.getBlockNumbered(Num));
   }
+
+  bool canSplitBeforeProlog(unsigned Num) {
+    MachineBasicBlock *BB = MF.getBlockNumbered(Num);
+    return IPA.canSplitBeforeProlog(*CurLI, *BB);
+  }
 };
 
 /// SplitEditor - Edit machine code and LiveIntervals for live range