Piece 1

Author: jrbyrnes

llvm/lib/Target/AMDGPU/GCNRegPressure.cpp

@@ -533,6 +533,7 @@ bool GCNDownwardRPTracker::advanceBeforeNext(MachineInstr *MI,
 
   // Remove dead registers or mask bits.
   SmallSet<Register, 8> SeenRegs;
+  //errs() << "Advance before next: "; CurrMI->dump();
   for (auto &MO : CurrMI->operands()) {
     if (!MO.isReg() || !MO.getReg().isVirtual())
       continue;
@@ -549,8 +550,10 @@ bool GCNDownwardRPTracker::advanceBeforeNext(MachineInstr *MI,
         if (!S.liveAt(SI)) {
           if (It == LiveRegs.end()) {
             It = LiveRegs.find(MO.getReg());
-            if (It == LiveRegs.end())
+            if (It == LiveRegs.end()) {
+              errs() << "BadReg: " << printReg(MO.getReg()) << "\n";
               llvm_unreachable("register isn't live");
+            }
           }
           auto PrevMask = It->second;
           It->second &= ~S.LaneMask;
@@ -561,8 +564,10 @@ bool GCNDownwardRPTracker::advanceBeforeNext(MachineInstr *MI,
         LiveRegs.erase(It);
     } else if (!LI.liveAt(SI)) {
       auto It = LiveRegs.find(MO.getReg());
-      if (It == LiveRegs.end())
+      if (It == LiveRegs.end()) {
+        errs() << "BadReg: " << printReg(MO.getReg()) << "\n";
         llvm_unreachable("register isn't live");
+      }
       CurPressure.inc(MO.getReg(), It->second, LaneBitmask::getNone(), *MRI);
       LiveRegs.erase(It);
     }

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

@@ -988,7 +988,6 @@ void GCNScheduleDAGMILive::runSchedStages() {
             ->reset(MRI, RegionLiveOuts.getLiveRegsForRegionIdx(
                              Stage->getRegionIdx()));
       }
-
       ScheduleDAGMILive::schedule();
       Stage->finalizeGCNRegion();
     }
@@ -1818,7 +1817,7 @@ bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
     auto NumRegs = SIRegisterInfo::getNumCoveredRegs(Mask);
     unsigned I = OptIt->getFirst();
     unsigned &Excess = OptIt->getSecond();
-    if (NumRegs >= Excess)
+    if (NumRegs >= Excess) 
       OptRegions.erase(I);
     else
       Excess -= NumRegs;
@@ -1854,41 +1853,32 @@ bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
       if (!UseMI || DefMI.getParent() == UseMI->getParent())
         continue;
 
-      // Do not rematerialize an instruction if it uses or is used by an
-      // instruction that we have designated for rematerialization.
-      // FIXME: Allow for rematerialization chains: this requires 1. updating
-      // remat points to account for uses that are rematerialized, and 2. either
-      // rematerializing the candidates in careful ordering, or deferring the
-      // MBB RP walk until the entire chain has been rematerialized.
-      if (Rematerializations.contains(UseMI) ||
-          llvm::any_of(DefMI.operands(), [&RematRegs](MachineOperand &MO) {
-            return MO.isReg() && RematRegs.contains(MO.getReg());
-          }))
-        continue;
-
       // Do not rematerialize an instruction it it uses registers that aren't
       // available at its use. This ensures that we are not extending any live
       // range while rematerializing.
       SlotIndex DefIdx = DAG.LIS->getInstructionIndex(DefMI);
       SlotIndex UseIdx = DAG.LIS->getInstructionIndex(*UseMI).getRegSlot(true);
+
       if (!allUsesAvailableAt(&DefMI, DefIdx, UseIdx))
         continue;
 
       REMAT_DEBUG(dbgs() << "Region " << I << ": remat instruction " << DefMI);
-      RematInstruction &Remat =
-          Rematerializations.try_emplace(&DefMI, I, UseMI).first->second;
+      RematInstruction &Temp = *Remats.insert({&DefMI, I, UseMI});
 
       bool RematUseful = false;
       if (auto It = OptRegions.find(I); It != OptRegions.end()) {
         // Optimistically consider that moving the instruction out of its
         // defining region will reduce RP in the latter; this assumes that
         // maximum RP in the region is reached somewhere between the defining
         // instruction and the end of the region.
+        // Since we only remat instructions with one use, we can assume that we
+        // adding a new remat instead of merely updating the remat position.
         REMAT_DEBUG(dbgs() << "  Defining region is optimizable\n");
         RematUseful = true;
         LaneBitmask Mask = DAG.RegionLiveOuts.getLiveRegsForRegionIdx(I)[Reg];
-        if (ReduceRPInRegion(It, Mask))
-          return true;
+        if (ReduceRPInRegion(It, Mask)) {
+          return  Remats.resolveInsertPos(&DAG.MRI, DAG.LIS);
+        }
       }
 
       for (unsigned LIRegion = 0; LIRegion != E; ++LIRegion) {
@@ -1897,7 +1887,7 @@ bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
         auto It = DAG.LiveIns[LIRegion].find(Reg);
         if (It == DAG.LiveIns[LIRegion].end() || It->second.none())
           continue;
-        Remat.LiveInRegions.insert(LIRegion);
+        Temp.LiveInRegions.insert(LIRegion);
 
         // Account for the reduction in RP due to the rematerialization in an
         // optimizable region in which the defined register is a live-in. This
@@ -1909,14 +1899,13 @@ bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
           REMAT_DEBUG(dbgs() << "  Live-in in region " << LIRegion << '\n');
           RematUseful = true;
           if (ReduceRPInRegion(It, DAG.LiveIns[LIRegion][Reg]))
-            return true;
+            return Remats.resolveInsertPos(&DAG.MRI, DAG.LIS);
         }
       }
 
       // If the instruction is not a live-in or live-out in any optimizable
       // region then there is no point in rematerializing it.
       if (!RematUseful) {
-        Rematerializations.pop_back();
         REMAT_DEBUG(dbgs() << "  No impact, not rematerializing instruction\n");
       } else {
         RematRegs.insert(Reg);
@@ -1927,11 +1916,11 @@ bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
   if (IncreaseOccupancy) {
     // We were trying to increase occupancy but failed, abort the stage.
     REMAT_DEBUG(dbgs() << "Cannot increase occupancy\n");
-    Rematerializations.clear();
+    Remats.clear();
     return false;
   }
   REMAT_DEBUG(dbgs() << "Can reduce but not eliminate spilling\n");
-  return !Rematerializations.empty();
+  return !Remats.empty() && Remats.resolveInsertPos(&DAG.MRI, DAG.LIS);
 }
 
 void PreRARematStage::rematerialize() {
@@ -1943,19 +1932,23 @@ void PreRARematStage::rematerialize() {
   DenseSet<unsigned> RecomputeRP;
   SlotIndexes *Slots = DAG.LIS->getSlotIndexes();
 
+  // Remat the dependencies last
+  Remats.sort(&DAG.MRI);
+
   // Rematerialize all instructions.
-  for (auto &[DefMI, Remat] : Rematerializations) {
-    MachineBasicBlock::iterator InsertPos(Remat.UseMI);
+  for (auto &Remat : Remats) {
+    MachineInstr *DefMI = Remat.DefMI;
+    MachineBasicBlock::iterator InsertPos = Remat.InsertPos;
     Register Reg = DefMI->getOperand(0).getReg();
     unsigned SubReg = DefMI->getOperand(0).getSubReg();
-
     // Rematerialize DefMI to its use block.
     TII->reMaterialize(*InsertPos->getParent(), InsertPos, Reg, SubReg, *DefMI,
                        *DAG.TRI);
     Remat.RematMI = &*std::prev(InsertPos);
     Remat.RematMI->getOperand(0).setSubReg(SubReg);
     DAG.LIS->InsertMachineInstrInMaps(*Remat.RematMI);
 
+
     // Update region boundaries in regions we sinked from (remove defining MI)
     // and to (insert MI rematerialized in use block). Only then we can erase
     // the original MI.
@@ -2054,6 +2047,20 @@ void PreRARematStage::rematerialize() {
       UpdateLiveRange(SubRange);
     }
   }
+  for (auto &Remat : reverse(Remats)) {
+      if (Remat.HasDependency) {
+      for (auto &ROp : Remat.RematMI->operands()) {
+        if (!ROp.isReg() || !ROp.getReg() || !ROp.readsReg())
+          continue;
+        auto UseReg = ROp.getReg();
+        if (!UseReg.isVirtual())
+          continue;
+        
+        DAG.LIS->removeInterval(UseReg);
+        DAG.LIS->createAndComputeVirtRegInterval(UseReg);
+      }
+    }
+  }
 
   // All regions impacted by at least one rematerialization must be rescheduled.
   // Maximum pressure must also be recomputed for all regions where it changed
@@ -2066,6 +2073,7 @@ void PreRARematStage::rematerialize() {
       continue;
 
     GCNRegPressure RP;
+
     if (IsEmptyRegion) {
       RP = getRegPressure(DAG.MRI, DAG.LiveIns[I]);
     } else {
@@ -2140,10 +2148,11 @@ void PreRARematStage::finalizeGCNSchedStage() {
       static_cast<const SIInstrInfo *>(MF.getSubtarget().getInstrInfo());
 
   // Rollback the rematerializations.
-  for (const auto &[_, Remat] : Rematerializations) {
+  for (auto &Remat : Remats) {
     MachineInstr &RematMI = *Remat.RematMI;
-    MachineBasicBlock::iterator InsertPos(DAG.Regions[Remat.DefRegion].second);
     MachineBasicBlock *MBB = getRegionMBB(MF, DAG.Regions[Remat.DefRegion]);
+    MachineBasicBlock::iterator InsertPos(MBB->end());
+
     Register Reg = RematMI.getOperand(0).getReg();
     unsigned SubReg = RematMI.getOperand(0).getSubReg();
 
@@ -2156,6 +2165,7 @@ void PreRARematStage::finalizeGCNSchedStage() {
     DAG.LIS->InsertMachineInstrInMaps(*NewMI);
 
     // Erase rematerialized MI.
+    DAG.updateRegionBoundaries(DAG.Regions, RematMI, nullptr);
     RematMI.eraseFromParent();
     DAG.LIS->RemoveMachineInstrFromMaps(RematMI);
 
@@ -2166,6 +2176,7 @@ void PreRARematStage::finalizeGCNSchedStage() {
     // Re-add the register as a live-in in all regions it used to be one in.
     for (unsigned LIRegion : Remat.LiveInRegions)
       DAG.LiveIns[LIRegion].insert({Reg, RegMasks.at({LIRegion, Reg})});
+    
   }
 
   // Reset RP in all impacted regions.

llvm/lib/Target/AMDGPU/GCNSchedStrategy.h

@@ -451,6 +451,9 @@ class PreRARematStage : public GCNSchedStage {
 private:
   /// Useful information about a rematerializable instruction.
   struct RematInstruction {
+    /// Single use of the rematerializable instruction's defined register,
+    /// located in a different block.
+    MachineInstr *DefMI;
     /// Single use of the rematerializable instruction's defined register,
     /// located in a different block.
     MachineInstr *UseMI;
@@ -463,12 +466,151 @@ class PreRARematStage : public GCNSchedStage {
     /// Region containing the rematerializable instruction.
     unsigned DefRegion;
 
-    RematInstruction(unsigned DefRegion, MachineInstr *UseMI)
-        : UseMI(UseMI), DefRegion(DefRegion) {}
+    /// The position at which to insert the remat.
+    MachineBasicBlock::iterator InsertPos;
+
+    /// The position at which we should rollback
+    MachineBasicBlock::iterator RollbackPos;
+
+    bool HasDependency = false;
+
+    RematInstruction(MachineInstr *DefMI, unsigned DefRegion, MachineInstr *UseMI)
+        : DefMI(DefMI), UseMI(UseMI), DefRegion(DefRegion), InsertPos(UseMI) {}
+    
+        RematInstruction(MachineInstr *DefMI, unsigned DefRegion, MachineInstr *UseMI, MachineBasicBlock::iterator InsertPos)
+        : DefMI(DefMI), UseMI(UseMI), DefRegion(DefRegion), InsertPos(InsertPos) {}
+  };
+
+  class RematInstructions {
+  private:
+    SmallVector<RematInstruction, 8> Rematerializations;
+
+  public:
+    using iterator = typename SmallVectorImpl<RematInstruction>::iterator;
+    using const_iterator = typename SmallVectorImpl<RematInstruction>::const_iterator;
+    using reverse_iterator = typename SmallVectorImpl<RematInstruction>::reverse_iterator;
+    using const_reverse_iterator = typename SmallVectorImpl<RematInstruction>::const_reverse_iterator;
+
+    iterator begin() { return Rematerializations.begin(); }
+    const_iterator begin() const { return Rematerializations.begin(); }
+    iterator end() { return Rematerializations.end(); }
+    const_iterator end() const { return Rematerializations.end(); }
+
+    reverse_iterator rbegin() { return Rematerializations.rbegin(); }
+    const_reverse_iterator rbegin() const { return Rematerializations.rbegin(); }
+    reverse_iterator rend() { return Rematerializations.rend(); }
+    const_reverse_iterator rend() const { return Rematerializations.rend(); }
+
+    unsigned size() {return Rematerializations.size();}
+
+    RematInstruction *insert(const RematInstruction &R) {
+      Rematerializations.push_back(R);
+      return &Rematerializations[size() - 1];
+    }
+
+    RematInstruction *insert(MachineInstr *DefMI, unsigned DefRegion, MachineInstr *UseMI) {
+      Rematerializations.push_back({DefMI, DefRegion, UseMI});
+      return &Rematerializations[size() - 1];
+    }
+
+    bool erase(const RematInstruction &R) {
+      auto Match = find_if(Rematerializations, [&R](const RematInstruction &Other){
+        return R.DefMI == Other.DefMI && R.UseMI == Other.UseMI;
+      });
+      if (Match == Rematerializations.end())
+        return false;
+      return !Rematerializations.erase(Match);
+    }
+
+    void clear() {
+      Rematerializations.clear();
+    }
+
+    bool empty() {
+      return Rematerializations.empty();
+    }
+
+    // We may be rematiarlizing an instruction used by another instruciton we are rematerializing. Be
+    // sure that we insert the user remats after -- the user remats and def remats will have the same InsertPt,
+    // by inserting the users last, they will occur after the defs. Thus, we must sort the remats so
+    // the users occur after the defs. 
+    void sort(const MachineRegisterInfo *MRI) {
+      std::sort(Rematerializations.begin(), Rematerializations.end(), [MRI](RematInstruction &A, RematInstruction &B) {
+        if (A.HasDependency && B.HasDependency) {
+          for (auto BOp : B.DefMI->operands()) {
+            if (!BOp.isReg() || !BOp.getReg() || !BOp.readsReg())
+              continue;
+            auto UseReg = BOp.getReg();
+            if (!UseReg.isVirtual())
+              continue;
+            MachineInstr *DefInst = &*MRI->def_instr_begin(UseReg);
+            if (DefInst == A.DefMI)
+              return true;
+          }
+          return false;
+        }
+
+        return !A.HasDependency;
+      });
+    }
+
+    bool resolveInsertPos(const MachineRegisterInfo *MRI,
+                              const LiveIntervals *LIS) {
+      // We may have added remat candidates which are used by other remat
+      // candidates -- be sure that we have correct insert points for this
+      bool FixedPoint = false;
+      unsigned IterCount = 0;
+      while (!FixedPoint && IterCount < 5) {
+        ++IterCount;
+        FixedPoint = true;
+        for (auto &Remat : Rematerializations) {
+          MachineInstr *RematInst = Remat.DefMI;
+
+          for (auto MO : RematInst->operands()) {
+            if (!MO.isReg() || !MO.getReg() || !MO.readsReg())
+              continue;
+            auto UseReg = MO.getReg();
+            if (!UseReg.isVirtual())
+              continue;
+            for (MachineInstr &DefInst : MRI->def_instructions(UseReg)) {
+
+
+              auto Match = find_if(Rematerializations, [&DefInst](const RematInstruction &R) {
+                return R.DefMI == &DefInst;
+              });
+
+              if (Match == Rematerializations.end())
+                continue;
+              
+              RematInstruction &TheMatch = *const_cast<RematInstruction *>(&*Match);
+              Remat.HasDependency = true;
+
+              // Since the remats 1. only have one use, and 2. the operands of the remat are live at the remat point
+              // we do not need further analysis to check whether changing the TheMatch.InsertPos will effect the condition
+              // of allUsesAvailableAt.
+              // The questionable case is when we are moving a remat pt to a later MBB. For such a condition to occur, we must have
+              // instructions A, B and C where we have previously set the remat point of A to B, but B is being remat to C so we 
+              // would like to update the remat point of A to C. We will remat B to C iff the use operands (i.e. A) are live at C. 
+              // since there is only 1 use of A, this can occur iff A is defined outside a loop and is live-in / live-out. Thus, B
+              // must be in the body of a loop. Moreover, since we know A can be remat at B, then the use operands of A must also be defined
+              // outside a loop and are live-in / live-out. Since these operands are live throughout the body of the loop, we are safe to 
+              // remat A to C without further checking.
+              // TODO: handle mulit-use case.
+              if (TheMatch.InsertPos == Remat.InsertPos)
+                continue;
+              FixedPoint = false;
+              TheMatch.InsertPos = Remat.InsertPos;
+            }
+          }
+
+        }
+      }
+      return FixedPoint;
+    }
   };
 
-  /// Collects instructions to rematerialize.
-  MapVector<MachineInstr *, RematInstruction> Rematerializations;
+  RematInstructions Remats;
+
   /// Collect regions whose live-ins or register pressure will change due to
   /// rematerializations.
   DenseMap<unsigned, GCNRegPressure> ImpactedRegions;