Added configuration flag for StackSlotColoring & refactored comments.

Author: vg0204

llvm/include/llvm/CodeGen/Passes.h

@@ -371,6 +371,8 @@ namespace llvm {
 
   /// StackSlotColoring - This pass performs stack slot coloring.
   extern char &StackSlotColoringID;
+  FunctionPass *
+  createStackSlotColoring(bool preserveRegAllocNeededAnalysis = false);
 
   /// This pass lays out funclets contiguously.
   extern char &FuncletLayoutID;

llvm/lib/CodeGen/StackSlotColoring.cpp

@@ -67,6 +67,9 @@ namespace {
     const MachineBlockFrequencyInfo *MBFI = nullptr;
     SlotIndexes *Indexes = nullptr;
 
+    // - preserves Analysis passes in case RA may be called afterwards.
+    bool preserveRegAllocNeededAnalysis = false;
+
     // SSIntervals - Spill slot intervals.
     std::vector<LiveInterval*> SSIntervals;
 
@@ -142,7 +145,9 @@ namespace {
   public:
     static char ID; // Pass identification
 
-    StackSlotColoring() : MachineFunctionPass(ID) {
+    StackSlotColoring(bool preserveRegAllocNeededAnalysis_ = false)
+        : MachineFunctionPass(ID),
+          preserveRegAllocNeededAnalysis(preserveRegAllocNeededAnalysis_) {
       initializeStackSlotColoringPass(*PassRegistry::getPassRegistry());
     }
 
@@ -155,12 +160,14 @@ namespace {
       AU.addPreserved<MachineBlockFrequencyInfo>();
       AU.addPreservedID(MachineDominatorsID);
 
-      /// NOTE: As in AMDGPU pass pipeline, reg alloc is spillted into 2 phases
-      /// and StackSlotColoring is invoked after each phase, it becomes
-      /// important to preserve additional analyses result to be used by VGPR
-      /// regAlloc, after being done with SGPR regAlloc and its related passes.
-      AU.addPreserved<LiveIntervals>();
-      AU.addPreserved<LiveDebugVariables>();
+      // As in some Target's pipeline, register allocation (RA) might be
+      // splitted into multiple phases based on register class. So, this pass
+      // may be invoked multiple times requiring it to save these analyses to be
+      // used by RA later.
+      if (preserveRegAllocNeededAnalysis) {
+        AU.addPreserved<LiveIntervals>();
+        AU.addPreserved<LiveDebugVariables>();
+      }
 
       MachineFunctionPass::getAnalysisUsage(AU);
     }
@@ -506,9 +513,6 @@ bool StackSlotColoring::RemoveDeadStores(MachineBasicBlock* MBB) {
     ++I;
   }
 
-  /// FIXED: As this pass preserves SlotIndexesAnalysis result, any
-  /// addition/removal of MI needs corresponding update in SlotIndexAnalysis,
-  /// to avoid corruption of SlotIndexesAnalysis result.
   for (MachineInstr *MI : toErase) {
     MI->eraseFromParent();
     Indexes->removeMachineInstrFromMaps(*MI);
@@ -565,3 +569,8 @@ bool StackSlotColoring::runOnMachineFunction(MachineFunction &MF) {
 
   return Changed;
 }
+
+FunctionPass *
+llvm::createStackSlotColoring(bool preserveRegAllocNeededAnalysis) {
+  return new StackSlotColoring(preserveRegAllocNeededAnalysis);
+}

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -1406,8 +1406,11 @@ bool GCNPassConfig::addRegAssignAndRewriteOptimized() {
   // since FastRegAlloc does the replacements itself.
   addPass(createVirtRegRewriter(false));
 
-  // Optimizes SGPR spills into VGPR lanes for non-interferring spill-ranges.
-  addPass(&StackSlotColoringID);
+  // As by this point SGPR's RA is done introducing SGPR spills to stack frame
+  // through SGPRAllocPass. So, invoking StackSlotColoring here, may allow these
+  // SGPR spills to re-use stack slots, before these spills is further lowered
+  // down via SILowerSGPRSpills(i.e. equivalent of PEI for SGPRs).
+  addPass(createStackSlotColoring(/*preserveRegAllocNeededAnalysis*/ true));
 
   // Equivalent of PEI for SGPRs.
   addPass(&SILowerSGPRSpillsID);

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

@@ -52,8 +52,7 @@ class SILowerSGPRSpills : public MachineFunctionPass {
   void calculateSaveRestoreBlocks(MachineFunction &MF);
   bool spillCalleeSavedRegs(MachineFunction &MF,
                             SmallVectorImpl<int> &CalleeSavedFIs);
-  void extendWWMVirtRegLiveness(MachineFunction &MF, SlotIndexes *Indexes,
-                                LiveIntervals *LIS);
+  void extendWWMVirtRegLiveness(MachineFunction &MF, LiveIntervals *LIS);
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
@@ -261,7 +260,6 @@ bool SILowerSGPRSpills::spillCalleeSavedRegs(
 }
 
 void SILowerSGPRSpills::extendWWMVirtRegLiveness(MachineFunction &MF,
-                                                 SlotIndexes *Indexes,
                                                  LiveIntervals *LIS) {
   // TODO: This is a workaround to avoid the unmodelled liveness computed with
   // whole-wave virtual registers when allocated together with the regular VGPR
@@ -280,7 +278,6 @@ void SILowerSGPRSpills::extendWWMVirtRegLiveness(MachineFunction &MF,
   for (auto Reg : MFI->getSGPRSpillVGPRs()) {
     for (MachineBasicBlock *SaveBlock : SaveBlocks) {
       MachineBasicBlock::iterator InsertBefore = SaveBlock->begin();
-      MachineInstrSpan MIS(InsertBefore, SaveBlock);
 
       DebugLoc DL = SaveBlock->findDebugLoc(InsertBefore);
       auto MIB = BuildMI(*SaveBlock, InsertBefore, DL,
@@ -289,13 +286,8 @@ void SILowerSGPRSpills::extendWWMVirtRegLiveness(MachineFunction &MF,
       // Set SGPR_SPILL asm printer flag
       MIB->setAsmPrinterFlag(AMDGPU::SGPR_SPILL);
 
-      if (LIS) {
+      if (LIS)
         LIS->InsertMachineInstrInMaps(*MIB);
-      } else if (Indexes) {
-        assert(std::distance(MIS.begin(), InsertBefore) == 1);
-        MachineInstr &Inst = *std::prev(InsertBefore);
-        Indexes->insertMachineInstrInMaps(Inst);
-      }
     }
   }
 
@@ -310,12 +302,8 @@ void SILowerSGPRSpills::extendWWMVirtRegLiveness(MachineFunction &MF,
                          TII->get(TargetOpcode::KILL));
       MIB.addReg(Reg);
 
-      if (LIS) {
+      if (LIS)
         LIS->InsertMachineInstrInMaps(*MIB);
-      } else if (Indexes) {
-        MachineInstr &Inst = *std::prev(InsertBefore);
-        Indexes->insertMachineInstrInMaps(Inst);
-      }
     }
   }
 }
@@ -406,7 +394,7 @@ bool SILowerSGPRSpills::runOnMachineFunction(MachineFunction &MF) {
     }
 
     if (SpilledToVirtVGPRLanes) {
-      extendWWMVirtRegLiveness(MF, Indexes, LIS);
+      extendWWMVirtRegLiveness(MF, LIS);
       if (LIS) {
         // Compute the LiveInterval for the newly created virtual registers.
         for (auto Reg : FuncInfo->getSGPRSpillVGPRs())