diff options
Diffstat (limited to 'llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp')
| -rw-r--r-- | llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp | 835 |
1 files changed, 527 insertions, 308 deletions
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp index 6e693066de10..fce8f36d4596 100644 --- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp +++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp @@ -25,8 +25,13 @@ #include "GCNSchedStrategy.h" #include "AMDGPUIGroupLP.h" +#include "GCNRegPressure.h" #include "SIMachineFunctionInfo.h" +#include "Utils/AMDGPUBaseInfo.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/CodeGen/RegisterClassInfo.h" +#include "llvm/MC/LaneBitmask.h" +#include "llvm/Support/ErrorHandling.h" #define DEBUG_TYPE "machine-scheduler" @@ -94,17 +99,20 @@ void GCNSchedStrategy::initialize(ScheduleDAGMI *DAG) { std::min(ST.getMaxNumSGPRs(TargetOccupancy, true), SGPRExcessLimit); if (!KnownExcessRP) { - VGPRCriticalLimit = - std::min(ST.getMaxNumVGPRs(TargetOccupancy), VGPRExcessLimit); + VGPRCriticalLimit = std::min( + ST.getMaxNumVGPRs(TargetOccupancy, MFI.getDynamicVGPRBlockSize()), + VGPRExcessLimit); } else { // This is similar to ST.getMaxNumVGPRs(TargetOccupancy) result except // returns a reasonably small number for targets with lots of VGPRs, such // as GFX10 and GFX11. LLVM_DEBUG(dbgs() << "Region is known to spill, use alternative " "VGPRCriticalLimit calculation method.\n"); - - unsigned Granule = AMDGPU::IsaInfo::getVGPRAllocGranule(&ST); - unsigned Addressable = AMDGPU::IsaInfo::getAddressableNumVGPRs(&ST); + unsigned DynamicVGPRBlockSize = MFI.getDynamicVGPRBlockSize(); + unsigned Granule = + AMDGPU::IsaInfo::getVGPRAllocGranule(&ST, DynamicVGPRBlockSize); + unsigned Addressable = + AMDGPU::IsaInfo::getAddressableNumVGPRs(&ST, DynamicVGPRBlockSize); unsigned VGPRBudget = alignDown(Addressable / TargetOccupancy, Granule); VGPRBudget = std::max(VGPRBudget, Granule); VGPRCriticalLimit = std::min(VGPRBudget, VGPRExcessLimit); @@ -188,12 +196,6 @@ static void getRegisterPressures( Pressure[AMDGPU::RegisterPressureSets::AGPR_32] = NewPressure.getAGPRNum(); } -// Return true if the instruction is mutually exclusive with all non-IGLP DAG -// mutations, requiring all other mutations to be disabled. -static bool isIGLPMutationOnly(unsigned Opcode) { - return Opcode == AMDGPU::SCHED_GROUP_BARRIER || Opcode == AMDGPU::IGLP_OPT; -} - void GCNSchedStrategy::initCandidate(SchedCandidate &Cand, SUnit *SU, bool AtTop, const RegPressureTracker &RPTracker, @@ -307,11 +309,11 @@ void GCNSchedStrategy::initCandidate(SchedCandidate &Cand, SUnit *SU, HasHighPressure = true; if (SGPRDelta > VGPRDelta) { Cand.RPDelta.CriticalMax = - PressureChange(AMDGPU::RegisterPressureSets::SReg_32); + PressureChange(AMDGPU::RegisterPressureSets::SReg_32); Cand.RPDelta.CriticalMax.setUnitInc(SGPRDelta); } else { Cand.RPDelta.CriticalMax = - PressureChange(AMDGPU::RegisterPressureSets::VGPR_32); + PressureChange(AMDGPU::RegisterPressureSets::VGPR_32); Cand.RPDelta.CriticalMax.setUnitInc(VGPRDelta); } } @@ -324,7 +326,7 @@ void GCNSchedStrategy::pickNodeFromQueue(SchedBoundary &Zone, const RegPressureTracker &RPTracker, SchedCandidate &Cand, bool IsBottomUp) { - const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI); + const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo *>(TRI); ArrayRef<unsigned> Pressure = RPTracker.getRegSetPressureAtPos(); unsigned SGPRPressure = 0; unsigned VGPRPressure = 0; @@ -420,7 +422,7 @@ SUnit *GCNSchedStrategy::pickNodeBidirectional(bool &IsTopNode) { pickNodeFromQueue(Top, TopPolicy, DAG->getTopRPTracker(), TCand, /*IsBottomUp=*/false); assert(TCand.SU == TopCand.SU && - "Last pick result should correspond to re-picking right now"); + "Last pick result should correspond to re-picking right now"); } #endif } @@ -590,12 +592,11 @@ bool GCNMaxILPSchedStrategy::tryCandidate(SchedCandidate &Cand, // This is a best effort to set things up for a post-RA pass. Optimizations // like generating loads of multiple registers should ideally be done within // the scheduler pass by combining the loads during DAG postprocessing. - const SUnit *CandNextClusterSU = - Cand.AtTop ? DAG->getNextClusterSucc() : DAG->getNextClusterPred(); - const SUnit *TryCandNextClusterSU = - TryCand.AtTop ? DAG->getNextClusterSucc() : DAG->getNextClusterPred(); - if (tryGreater(TryCand.SU == TryCandNextClusterSU, - Cand.SU == CandNextClusterSU, TryCand, Cand, Cluster)) + const ClusterInfo *CandCluster = Cand.AtTop ? TopCluster : BotCluster; + const ClusterInfo *TryCandCluster = TryCand.AtTop ? TopCluster : BotCluster; + if (tryGreater(TryCandCluster && TryCandCluster->contains(TryCand.SU), + CandCluster && CandCluster->contains(Cand.SU), TryCand, Cand, + Cluster)) return TryCand.Reason != NoCand; // Avoid increasing the max critical pressure in the scheduled region. @@ -665,12 +666,11 @@ bool GCNMaxMemoryClauseSchedStrategy::tryCandidate(SchedCandidate &Cand, // MaxMemoryClause-specific: We prioritize clustered instructions as we would // get more benefit from clausing these memory instructions. - const SUnit *CandNextClusterSU = - Cand.AtTop ? DAG->getNextClusterSucc() : DAG->getNextClusterPred(); - const SUnit *TryCandNextClusterSU = - TryCand.AtTop ? DAG->getNextClusterSucc() : DAG->getNextClusterPred(); - if (tryGreater(TryCand.SU == TryCandNextClusterSU, - Cand.SU == CandNextClusterSU, TryCand, Cand, Cluster)) + const ClusterInfo *CandCluster = Cand.AtTop ? TopCluster : BotCluster; + const ClusterInfo *TryCandCluster = TryCand.AtTop ? TopCluster : BotCluster; + if (tryGreater(TryCandCluster && TryCandCluster->contains(TryCand.SU), + CandCluster && CandCluster->contains(Cand.SU), TryCand, Cand, + Cluster)) return TryCand.Reason != NoCand; // We only compare a subset of features when comparing nodes between @@ -760,6 +760,10 @@ GCNScheduleDAGMILive::GCNScheduleDAGMILive( StartingOccupancy(MFI.getOccupancy()), MinOccupancy(StartingOccupancy), RegionLiveOuts(this, /*IsLiveOut=*/true) { + // We want regions with a single MI to be scheduled so that we can reason + // about them correctly during scheduling stages that move MIs between regions + // (e.g., rematerialization). + ScheduleSingleMIRegions = true; LLVM_DEBUG(dbgs() << "Starting occupancy is " << StartingOccupancy << ".\n"); if (RelaxedOcc) { MinOccupancy = std::min(MFI.getMinAllowedOccupancy(), StartingOccupancy); @@ -869,6 +873,8 @@ void GCNScheduleDAGMILive::computeBlockPressure(unsigned RegionIdx, Pressure[CurRegion] = RPTracker.moveMaxPressure(); if (CurRegion-- == RegionIdx) break; + auto &Rgn = Regions[CurRegion]; + NonDbgMI = &*skipDebugInstructionsForward(Rgn.first, Rgn.second); } RPTracker.advanceToNext(); RPTracker.advanceBeforeNext(); @@ -890,13 +896,13 @@ GCNScheduleDAGMILive::getRegionLiveInMap() const { std::vector<MachineInstr *> RegionFirstMIs; RegionFirstMIs.reserve(Regions.size()); auto I = Regions.rbegin(), E = Regions.rend(); - auto *BB = I->first->getParent(); do { + const MachineBasicBlock *MBB = I->first->getParent(); auto *MI = &*skipDebugInstructionsForward(I->first, I->second); RegionFirstMIs.push_back(MI); do { ++I; - } while (I != E && I->first->getParent() == BB); + } while (I != E && I->first->getParent() == MBB); } while (I != E); return getLiveRegMap(RegionFirstMIs, /*After=*/false, *LIS); } @@ -932,12 +938,10 @@ void GCNScheduleDAGMILive::finalizeSchedule() { // GCNScheduleDAGMILive::schedule(). LiveIns.resize(Regions.size()); Pressure.resize(Regions.size()); - RescheduleRegions.resize(Regions.size()); RegionsWithHighRP.resize(Regions.size()); RegionsWithExcessRP.resize(Regions.size()); RegionsWithMinOcc.resize(Regions.size()); RegionsWithIGLPInstrs.resize(Regions.size()); - RescheduleRegions.set(); RegionsWithHighRP.reset(); RegionsWithExcessRP.reset(); RegionsWithMinOcc.reset(); @@ -1081,31 +1085,46 @@ bool ClusteredLowOccStage::initGCNSchedStage() { return true; } -bool PreRARematStage::initGCNSchedStage() { - if (!GCNSchedStage::initGCNSchedStage()) - return false; +/// Allows to easily filter for this stage's debug output. +#define REMAT_DEBUG(X) LLVM_DEBUG(dbgs() << "[PreRARemat] "; X;) - if (DAG.RegionsWithMinOcc.none() || DAG.Regions.size() == 1) - return false; +bool PreRARematStage::initGCNSchedStage() { + // FIXME: This pass will invalidate cached BBLiveInMap and MBBLiveIns for + // regions inbetween the defs and region we sinked the def to. Will need to be + // fixed if there is another pass after this pass. + assert(!S.hasNextStage()); - const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo(); - // Rematerialization will not help if occupancy is not limited by reg usage. - if (ST.getOccupancyWithWorkGroupSizes(MF).second == DAG.MinOccupancy) + if (!GCNSchedStage::initGCNSchedStage() || DAG.RegionsWithMinOcc.none() || + DAG.Regions.size() == 1) return false; - // FIXME: This pass will invalidate cached MBBLiveIns for regions - // inbetween the defs and region we sinked the def to. Cached pressure - // for regions where a def is sinked from will also be invalidated. Will - // need to be fixed if there is another pass after this pass. - assert(!S.hasNextStage()); + // Before performing any IR modification record the parent region of each MI + // and the parent MBB of each region. + const unsigned NumRegions = DAG.Regions.size(); + RegionBB.reserve(NumRegions); + for (unsigned I = 0; I < NumRegions; ++I) { + RegionBoundaries Region = DAG.Regions[I]; + for (auto MI = Region.first; MI != Region.second; ++MI) + MIRegion.insert({&*MI, I}); + RegionBB.push_back(Region.first->getParent()); + } - collectRematerializableInstructions(); - if (RematerializableInsts.empty() || !sinkTriviallyRematInsts(ST, TII)) + if (!canIncreaseOccupancyOrReduceSpill()) return false; - LLVM_DEBUG( - dbgs() << "Retrying function scheduling with improved occupancy of " - << DAG.MinOccupancy << " from rematerializing\n"); + // Rematerialize identified instructions and update scheduler's state. + rematerialize(); + if (GCNTrackers) + DAG.RegionLiveOuts.buildLiveRegMap(); + REMAT_DEBUG( + dbgs() << "Retrying function scheduling with new min. occupancy of " + << AchievedOcc << " from rematerializing (original was " + << DAG.MinOccupancy << ", target was " << TargetOcc << ")\n"); + if (AchievedOcc > DAG.MinOccupancy) { + DAG.MinOccupancy = AchievedOcc; + SIMachineFunctionInfo &MFI = *MF.getInfo<SIMachineFunctionInfo>(); + MFI.increaseOccupancy(MF, DAG.MinOccupancy); + } return true; } @@ -1120,7 +1139,8 @@ void UnclusteredHighRPStage::finalizeGCNSchedStage() { if (DAG.MinOccupancy > InitialOccupancy) { for (unsigned IDX = 0; IDX < DAG.Pressure.size(); ++IDX) DAG.RegionsWithMinOcc[IDX] = - DAG.Pressure[IDX].getOccupancy(DAG.ST) == DAG.MinOccupancy; + DAG.Pressure[IDX].getOccupancy( + DAG.ST, DAG.MFI.getDynamicVGPRBlockSize()) == DAG.MinOccupancy; LLVM_DEBUG(dbgs() << StageID << " stage successfully increased occupancy to " @@ -1155,9 +1175,10 @@ bool GCNSchedStage::initGCNRegion() { Unsched.reserve(DAG.NumRegionInstrs); if (StageID == GCNSchedStageID::OccInitialSchedule || StageID == GCNSchedStageID::ILPInitialSchedule) { + const SIInstrInfo *SII = static_cast<const SIInstrInfo *>(DAG.TII); for (auto &I : DAG) { Unsched.push_back(&I); - if (isIGLPMutationOnly(I.getOpcode())) + if (SII->isIGLPMutationOnly(I.getOpcode())) DAG.RegionsWithIGLPInstrs[RegionIdx] = true; } } else { @@ -1215,10 +1236,7 @@ bool ClusteredLowOccStage::initGCNRegion() { } bool PreRARematStage::initGCNRegion() { - if (!DAG.RescheduleRegions[RegionIdx]) - return false; - - return GCNSchedStage::initGCNRegion(); + return RescheduleRegions[RegionIdx] && GCNSchedStage::initGCNRegion(); } void GCNSchedStage::setupNewBlock() { @@ -1237,7 +1255,6 @@ void GCNSchedStage::setupNewBlock() { void GCNSchedStage::finalizeGCNRegion() { DAG.Regions[RegionIdx] = std::pair(DAG.RegionBegin, DAG.RegionEnd); - DAG.RescheduleRegions[RegionIdx] = false; if (S.HasHighPressure) DAG.RegionsWithHighRP[RegionIdx] = true; @@ -1250,7 +1267,7 @@ void GCNSchedStage::finalizeGCNRegion() { SavedMutations.swap(DAG.Mutations); DAG.exitRegion(); - RegionIdx++; + advanceRegion(); } void GCNSchedStage::checkScheduling() { @@ -1260,11 +1277,14 @@ void GCNSchedStage::checkScheduling() { LLVM_DEBUG(dbgs() << "Pressure after scheduling: " << print(PressureAfter)); LLVM_DEBUG(dbgs() << "Region: " << RegionIdx << ".\n"); + unsigned DynamicVGPRBlockSize = DAG.MFI.getDynamicVGPRBlockSize(); + if (PressureAfter.getSGPRNum() <= S.SGPRCriticalLimit && PressureAfter.getVGPRNum(ST.hasGFX90AInsts()) <= S.VGPRCriticalLimit) { DAG.Pressure[RegionIdx] = PressureAfter; DAG.RegionsWithMinOcc[RegionIdx] = - PressureAfter.getOccupancy(ST) == DAG.MinOccupancy; + PressureAfter.getOccupancy(ST, DynamicVGPRBlockSize) == + DAG.MinOccupancy; // Early out if we have achieved the occupancy target. LLVM_DEBUG(dbgs() << "Pressure in desired limits, done.\n"); @@ -1273,10 +1293,10 @@ void GCNSchedStage::checkScheduling() { unsigned TargetOccupancy = std::min( S.getTargetOccupancy(), ST.getOccupancyWithWorkGroupSizes(MF).second); - unsigned WavesAfter = - std::min(TargetOccupancy, PressureAfter.getOccupancy(ST)); - unsigned WavesBefore = - std::min(TargetOccupancy, PressureBefore.getOccupancy(ST)); + unsigned WavesAfter = std::min( + TargetOccupancy, PressureAfter.getOccupancy(ST, DynamicVGPRBlockSize)); + unsigned WavesBefore = std::min( + TargetOccupancy, PressureBefore.getOccupancy(ST, DynamicVGPRBlockSize)); LLVM_DEBUG(dbgs() << "Occupancy before scheduling: " << WavesBefore << ", after " << WavesAfter << ".\n"); @@ -1311,10 +1331,9 @@ void GCNSchedStage::checkScheduling() { unsigned MaxSGPRs = ST.getMaxNumSGPRs(MF); if (PressureAfter.getVGPRNum(ST.hasGFX90AInsts()) > MaxVGPRs || - PressureAfter.getVGPRNum(false) > MaxArchVGPRs || + PressureAfter.getArchVGPRNum() > MaxArchVGPRs || PressureAfter.getAGPRNum() > MaxArchVGPRs || PressureAfter.getSGPRNum() > MaxSGPRs) { - DAG.RescheduleRegions[RegionIdx] = true; DAG.RegionsWithHighRP[RegionIdx] = true; DAG.RegionsWithExcessRP[RegionIdx] = true; } @@ -1326,7 +1345,8 @@ void GCNSchedStage::checkScheduling() { } else { DAG.Pressure[RegionIdx] = PressureAfter; DAG.RegionsWithMinOcc[RegionIdx] = - PressureAfter.getOccupancy(ST) == DAG.MinOccupancy; + PressureAfter.getOccupancy(ST, DynamicVGPRBlockSize) == + DAG.MinOccupancy; } } @@ -1448,6 +1468,18 @@ bool GCNSchedStage::shouldRevertScheduling(unsigned WavesAfter) { if (WavesAfter < DAG.MinOccupancy) return true; + // For dynamic VGPR mode, we don't want to waste any VGPR blocks. + if (DAG.MFI.isDynamicVGPREnabled()) { + unsigned BlocksBefore = AMDGPU::IsaInfo::getAllocatedNumVGPRBlocks( + &ST, DAG.MFI.getDynamicVGPRBlockSize(), + PressureBefore.getVGPRNum(false)); + unsigned BlocksAfter = AMDGPU::IsaInfo::getAllocatedNumVGPRBlocks( + &ST, DAG.MFI.getDynamicVGPRBlockSize(), + PressureAfter.getVGPRNum(false)); + if (BlocksAfter > BlocksBefore) + return true; + } + return false; } @@ -1467,7 +1499,8 @@ bool OccInitialScheduleStage::shouldRevertScheduling(unsigned WavesAfter) { bool UnclusteredHighRPStage::shouldRevertScheduling(unsigned WavesAfter) { // If RP is not reduced in the unclustered reschedule stage, revert to the // old schedule. - if ((WavesAfter <= PressureBefore.getOccupancy(ST) && + if ((WavesAfter <= + PressureBefore.getOccupancy(ST, DAG.MFI.getDynamicVGPRBlockSize()) && mayCauseSpilling(WavesAfter)) || GCNSchedStage::shouldRevertScheduling(WavesAfter)) { LLVM_DEBUG(dbgs() << "Unclustered reschedule did not help.\n"); @@ -1482,16 +1515,16 @@ bool UnclusteredHighRPStage::shouldRevertScheduling(unsigned WavesAfter) { dbgs() << "\n\t *** In shouldRevertScheduling ***\n" << " *********** BEFORE UnclusteredHighRPStage ***********\n"); - ScheduleMetrics MBefore = - getScheduleMetrics(DAG.SUnits); + ScheduleMetrics MBefore = getScheduleMetrics(DAG.SUnits); LLVM_DEBUG( dbgs() << "\n *********** AFTER UnclusteredHighRPStage ***********\n"); ScheduleMetrics MAfter = getScheduleMetrics(DAG); unsigned OldMetric = MBefore.getMetric(); unsigned NewMetric = MAfter.getMetric(); - unsigned WavesBefore = - std::min(S.getTargetOccupancy(), PressureBefore.getOccupancy(ST)); + unsigned WavesBefore = std::min( + S.getTargetOccupancy(), + PressureBefore.getOccupancy(ST, DAG.MFI.getDynamicVGPRBlockSize())); unsigned Profit = ((WavesAfter * ScheduleMetrics::ScaleFactor) / WavesBefore * ((OldMetric + ScheduleMetricBias) * ScheduleMetrics::ScaleFactor) / @@ -1516,13 +1549,9 @@ bool ClusteredLowOccStage::shouldRevertScheduling(unsigned WavesAfter) { } bool PreRARematStage::shouldRevertScheduling(unsigned WavesAfter) { - if (GCNSchedStage::shouldRevertScheduling(WavesAfter)) - return true; - - if (mayCauseSpilling(WavesAfter)) - return true; - - return false; + return GCNSchedStage::shouldRevertScheduling(WavesAfter) || + mayCauseSpilling(WavesAfter) || + (IncreaseOccupancy && WavesAfter < TargetOcc); } bool ILPInitialScheduleStage::shouldRevertScheduling(unsigned WavesAfter) { @@ -1549,11 +1578,9 @@ bool GCNSchedStage::mayCauseSpilling(unsigned WavesAfter) { void GCNSchedStage::revertScheduling() { DAG.RegionsWithMinOcc[RegionIdx] = - PressureBefore.getOccupancy(ST) == DAG.MinOccupancy; + PressureBefore.getOccupancy(ST, DAG.MFI.getDynamicVGPRBlockSize()) == + DAG.MinOccupancy; LLVM_DEBUG(dbgs() << "Attempting to revert scheduling.\n"); - DAG.RescheduleRegions[RegionIdx] = - S.hasNextStage() && - S.getNextStage() != GCNSchedStageID::UnclusteredHighRPReschedule; DAG.RegionEnd = DAG.RegionBegin; int SkippedDebugInstr = 0; for (MachineInstr *MI : Unsched) { @@ -1563,8 +1590,7 @@ void GCNSchedStage::revertScheduling() { } if (MI->getIterator() != DAG.RegionEnd) { - DAG.BB->remove(MI); - DAG.BB->insert(DAG.RegionEnd, MI); + DAG.BB->splice(DAG.RegionEnd, DAG.BB, MI); if (!MI->isDebugInstr()) DAG.LIS->handleMove(*MI, true); } @@ -1615,241 +1641,396 @@ void GCNSchedStage::revertScheduling() { DAG.Regions[RegionIdx] = std::pair(DAG.RegionBegin, DAG.RegionEnd); } -void PreRARematStage::collectRematerializableInstructions() { - const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo *>(DAG.TRI); - for (unsigned I = 0, E = DAG.MRI.getNumVirtRegs(); I != E; ++I) { - Register Reg = Register::index2VirtReg(I); - if (!DAG.LIS->hasInterval(Reg)) - continue; +bool PreRARematStage::allUsesAvailableAt(const MachineInstr *InstToRemat, + SlotIndex OriginalIdx, + SlotIndex RematIdx) const { - // TODO: Handle AGPR and SGPR rematerialization - if (!SRI->isVGPRClass(DAG.MRI.getRegClass(Reg)) || - !DAG.MRI.hasOneDef(Reg) || !DAG.MRI.hasOneNonDBGUse(Reg)) + LiveIntervals *LIS = DAG.LIS; + MachineRegisterInfo &MRI = DAG.MRI; + OriginalIdx = OriginalIdx.getRegSlot(true); + RematIdx = std::max(RematIdx, RematIdx.getRegSlot(true)); + for (const MachineOperand &MO : InstToRemat->operands()) { + if (!MO.isReg() || !MO.getReg() || !MO.readsReg()) continue; - MachineOperand *Op = DAG.MRI.getOneDef(Reg); - MachineInstr *Def = Op->getParent(); - if (Op->getSubReg() != 0 || !isTriviallyReMaterializable(*Def)) - continue; + if (!MO.getReg().isVirtual()) { + // Do not attempt to reason about PhysRegs + // TODO: better analysis of PhysReg livness + if (!DAG.MRI.isConstantPhysReg(MO.getReg()) && + !DAG.TII->isIgnorableUse(MO)) + return false; - MachineInstr *UseI = &*DAG.MRI.use_instr_nodbg_begin(Reg); - if (Def->getParent() == UseI->getParent()) + // Constant PhysRegs and IgnorableUses are okay continue; + } - // We are only collecting defs that are defined in another block and are - // live-through or used inside regions at MinOccupancy. This means that the - // register must be in the live-in set for the region. - bool AddedToRematList = false; - for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) { - auto It = DAG.LiveIns[I].find(Reg); - if (It != DAG.LiveIns[I].end() && !It->second.none()) { - if (DAG.RegionsWithMinOcc[I]) { - RematerializableInsts[I][Def] = UseI; - AddedToRematList = true; - } + LiveInterval &LI = LIS->getInterval(MO.getReg()); + const VNInfo *OVNI = LI.getVNInfoAt(OriginalIdx); + assert(OVNI); + + // Don't allow rematerialization immediately after the original def. + // It would be incorrect if InstToRemat redefines the register. + // See PR14098. + if (SlotIndex::isSameInstr(OriginalIdx, RematIdx)) + return false; + + if (OVNI != LI.getVNInfoAt(RematIdx)) + return false; + + // Check that subrange is live at RematIdx. + if (LI.hasSubRanges()) { + const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo(); + unsigned SubReg = MO.getSubReg(); + LaneBitmask LM = SubReg ? TRI->getSubRegIndexLaneMask(SubReg) + : MRI.getMaxLaneMaskForVReg(MO.getReg()); + for (LiveInterval::SubRange &SR : LI.subranges()) { + if ((SR.LaneMask & LM).none()) + continue; + if (!SR.liveAt(RematIdx)) + return false; - // Collect regions with rematerializable reg as live-in to avoid - // searching later when updating RP. - RematDefToLiveInRegions[Def].push_back(I); + // Early exit if all used lanes are checked. No need to continue. + LM &= ~SR.LaneMask; + if (LM.none()) + break; } } - if (!AddedToRematList) - RematDefToLiveInRegions.erase(Def); } + return true; } -bool PreRARematStage::sinkTriviallyRematInsts(const GCNSubtarget &ST, - const TargetInstrInfo *TII) { - // Temporary copies of cached variables we will be modifying and replacing if - // sinking succeeds. - SmallVector< - std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>, 32> - NewRegions; - DenseMap<unsigned, GCNRPTracker::LiveRegSet> NewLiveIns; - DenseMap<unsigned, GCNRegPressure> NewPressure; - BitVector NewRescheduleRegions; - LiveIntervals *LIS = DAG.LIS; +bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() { + REMAT_DEBUG({ + dbgs() << "Collecting rematerializable instructions in "; + MF.getFunction().printAsOperand(dbgs(), false); + dbgs() << '\n'; + }); - NewRegions.resize(DAG.Regions.size()); - NewRescheduleRegions.resize(DAG.Regions.size()); + // Maps optimizable regions (i.e., regions at minimum and register-limited + // occupancy, or regions with spilling) to the target RP we would like to + // reach. + DenseMap<unsigned, GCNRPTarget> OptRegions; + const Function &F = MF.getFunction(); + unsigned DynamicVGPRBlockSize = + MF.getInfo<SIMachineFunctionInfo>()->getDynamicVGPRBlockSize(); - // Collect only regions that has a rematerializable def as a live-in. - SmallSet<unsigned, 16> ImpactedRegions; - for (const auto &It : RematDefToLiveInRegions) - ImpactedRegions.insert(It.second.begin(), It.second.end()); + std::pair<unsigned, unsigned> WavesPerEU = ST.getWavesPerEU(F); + const unsigned MaxSGPRsNoSpill = ST.getMaxNumSGPRs(F); + const unsigned MaxVGPRsNoSpill = ST.getMaxNumVGPRs(F); + const unsigned MaxSGPRsIncOcc = + ST.getMaxNumSGPRs(DAG.MinOccupancy + 1, false); + const unsigned MaxVGPRsIncOcc = + ST.getMaxNumVGPRs(DAG.MinOccupancy + 1, DynamicVGPRBlockSize); + IncreaseOccupancy = WavesPerEU.second > DAG.MinOccupancy; - // Make copies of register pressure and live-ins cache that will be updated - // as we rematerialize. - for (auto Idx : ImpactedRegions) { - NewPressure[Idx] = DAG.Pressure[Idx]; - NewLiveIns[Idx] = DAG.LiveIns[Idx]; + // Collect optimizable regions. If there is spilling in any region we will + // just try to reduce spilling. Otherwise we will try to increase occupancy by + // one in the whole function. + for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) { + GCNRegPressure &RP = DAG.Pressure[I]; + // We allow ArchVGPR or AGPR savings to count as savings of the other kind + // of VGPR only when trying to eliminate spilling. We cannot do this when + // trying to increase occupancy since VGPR class swaps only occur later in + // the register allocator i.e., the scheduler will not be able to reason + // about these savings and will not report an increase in the achievable + // occupancy, triggering rollbacks. + GCNRPTarget Target(MaxSGPRsNoSpill, MaxVGPRsNoSpill, MF, RP, + /*CombineVGPRSavings=*/true); + if (!Target.satisfied() && IncreaseOccupancy) { + // There is spilling in the region and we were so far trying to increase + // occupancy. Strop trying that and focus on reducing spilling. + IncreaseOccupancy = false; + OptRegions.clear(); + } else if (IncreaseOccupancy) { + // There is no spilling in the region, try to increase occupancy. + Target = GCNRPTarget(MaxSGPRsIncOcc, MaxVGPRsIncOcc, MF, RP, + /*CombineVGPRSavings=*/false); + } + if (!Target.satisfied()) + OptRegions.insert({I, Target}); } - NewRegions = DAG.Regions; - NewRescheduleRegions.reset(); + if (OptRegions.empty()) + return false; - DenseMap<MachineInstr *, MachineInstr *> InsertedMIToOldDef; - bool Improved = false; - for (auto I : ImpactedRegions) { - if (!DAG.RegionsWithMinOcc[I]) - continue; +#ifndef NDEBUG + if (IncreaseOccupancy) { + REMAT_DEBUG(dbgs() << "Occupancy minimal (" << DAG.MinOccupancy + << ") in regions:\n"); + } else { + REMAT_DEBUG(dbgs() << "Spilling w.r.t. minimum target occupancy (" + << WavesPerEU.first << ") in regions:\n"); + } + for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) { + if (auto OptIt = OptRegions.find(I); OptIt != OptRegions.end()) + REMAT_DEBUG(dbgs() << " [" << I << "] " << OptIt->getSecond() << '\n'); + } +#endif - Improved = false; - int VGPRUsage = NewPressure[I].getVGPRNum(ST.hasGFX90AInsts()); - int SGPRUsage = NewPressure[I].getSGPRNum(); + // When we are reducing spilling, the target is the minimum target number of + // waves/EU determined by the subtarget. In cases where either one of + // "amdgpu-num-sgpr" or "amdgpu-num-vgpr" are set on the function, the current + // minimum region occupancy may be higher than the latter. + TargetOcc = IncreaseOccupancy ? DAG.MinOccupancy + 1 + : std::max(DAG.MinOccupancy, WavesPerEU.first); - // TODO: Handle occupancy drop due to AGPR and SGPR. - // Check if cause of occupancy drop is due to VGPR usage and not SGPR. - if (ST.getOccupancyWithNumSGPRs(SGPRUsage) == DAG.MinOccupancy) - break; + // Accounts for a reduction in RP in an optimizable region. Returns whether we + // estimate that we have identified enough rematerialization opportunities to + // achieve our goal, and sets Progress to true when this particular reduction + // in pressure was helpful toward that goal. + auto ReduceRPInRegion = [&](auto OptIt, Register Reg, LaneBitmask Mask, + bool &Progress) -> bool { + GCNRPTarget &Target = OptIt->getSecond(); + if (!Target.isSaveBeneficial(Reg, DAG.MRI)) + return false; + Progress = true; + Target.saveReg(Reg, Mask, DAG.MRI); + if (Target.satisfied()) + OptRegions.erase(OptIt->getFirst()); + return OptRegions.empty(); + }; - // The occupancy of this region could have been improved by a previous - // iteration's sinking of defs. - if (NewPressure[I].getOccupancy(ST) > DAG.MinOccupancy) { - NewRescheduleRegions[I] = true; - Improved = true; - continue; - } + // We need up-to-date live-out info. to query live-out register masks in + // regions containing rematerializable instructions. + DAG.RegionLiveOuts.buildLiveRegMap(); - // First check if we have enough trivially rematerializable instructions to - // improve occupancy. Optimistically assume all instructions we are able to - // sink decreased RP. - int TotalSinkableRegs = 0; - for (const auto &It : RematerializableInsts[I]) { - MachineInstr *Def = It.first; - Register DefReg = Def->getOperand(0).getReg(); - TotalSinkableRegs += - SIRegisterInfo::getNumCoveredRegs(NewLiveIns[I][DefReg]); - } - int VGPRsAfterSink = VGPRUsage - TotalSinkableRegs; - unsigned OptimisticOccupancy = ST.getOccupancyWithNumVGPRs(VGPRsAfterSink); - // If in the most optimistic scenario, we cannot improve occupancy, then do - // not attempt to sink any instructions. - if (OptimisticOccupancy <= DAG.MinOccupancy) - break; + // Cache set of registers that are going to be rematerialized. + DenseSet<unsigned> RematRegs; + + // Identify rematerializable instructions in the function. + for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) { + auto Region = DAG.Regions[I]; + for (auto MI = Region.first; MI != Region.second; ++MI) { + // The instruction must be trivially rematerializable. + MachineInstr &DefMI = *MI; + if (!isTriviallyReMaterializable(DefMI)) + continue; + + // We only support rematerializing virtual registers with one definition. + Register Reg = DefMI.getOperand(0).getReg(); + if (!Reg.isVirtual() || !DAG.MRI.hasOneDef(Reg)) + continue; + + // We only care to rematerialize the instruction if it has a single + // non-debug user in a different region. The using MI may not belong to a + // region if it is a lone region terminator. + MachineInstr *UseMI = DAG.MRI.getOneNonDBGUser(Reg); + if (!UseMI) + continue; + auto UseRegion = MIRegion.find(UseMI); + if (UseRegion != MIRegion.end() && UseRegion->second == I) + 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; - unsigned ImproveOccupancy = 0; - SmallVector<MachineInstr *, 4> SinkedDefs; - for (auto &It : RematerializableInsts[I]) { - MachineInstr *Def = It.first; - MachineBasicBlock::iterator InsertPos = - MachineBasicBlock::iterator(It.second); - Register Reg = Def->getOperand(0).getReg(); - // Rematerialize MI to its use block. Since we are only rematerializing - // instructions that do not have any virtual reg uses, we do not need to - // call LiveRangeEdit::allUsesAvailableAt() and - // LiveRangeEdit::canRematerializeAt(). - TII->reMaterialize(*InsertPos->getParent(), InsertPos, Reg, - Def->getOperand(0).getSubReg(), *Def, *DAG.TRI); - MachineInstr *NewMI = &*std::prev(InsertPos); - LIS->InsertMachineInstrInMaps(*NewMI); - LIS->removeInterval(Reg); - LIS->createAndComputeVirtRegInterval(Reg); - InsertedMIToOldDef[NewMI] = Def; + REMAT_DEBUG(dbgs() << "Region " << I << ": remat instruction " << DefMI); + RematInstruction &Remat = + Rematerializations.try_emplace(&DefMI, UseMI).first->second; - // Update region boundaries in scheduling region we sinked from since we - // may sink an instruction that was at the beginning or end of its region - DAG.updateRegionBoundaries(NewRegions, Def, /*NewMI =*/nullptr, - /*Removing =*/true); + 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. + REMAT_DEBUG(dbgs() << " Defining region is optimizable\n"); + LaneBitmask Mask = DAG.RegionLiveOuts.getLiveRegsForRegionIdx(I)[Reg]; + if (ReduceRPInRegion(It, Reg, Mask, RematUseful)) + return true; + } - // Update region boundaries in region we sinked to. - DAG.updateRegionBoundaries(NewRegions, InsertPos, NewMI); + for (unsigned LIRegion = 0; LIRegion != E; ++LIRegion) { + // We are only collecting regions in which the register is a live-in + // (and may be live-through). + auto It = DAG.LiveIns[LIRegion].find(Reg); + if (It == DAG.LiveIns[LIRegion].end() || It->second.none()) + continue; + Remat.LiveInRegions.insert(LIRegion); - LaneBitmask PrevMask = NewLiveIns[I][Reg]; - // FIXME: Also update cached pressure for where the def was sinked from. - // Update RP for all regions that has this reg as a live-in and remove - // the reg from all regions as a live-in. - for (auto Idx : RematDefToLiveInRegions[Def]) { - NewLiveIns[Idx].erase(Reg); - if (InsertPos->getParent() != DAG.Regions[Idx].first->getParent()) { - // Def is live-through and not used in this block. - NewPressure[Idx].inc(Reg, PrevMask, LaneBitmask::getNone(), DAG.MRI); - } else { - // Def is used and rematerialized into this block. - GCNDownwardRPTracker RPT(*LIS); - auto *NonDbgMI = &*skipDebugInstructionsForward( - NewRegions[Idx].first, NewRegions[Idx].second); - RPT.reset(*NonDbgMI, &NewLiveIns[Idx]); - RPT.advance(NewRegions[Idx].second); - NewPressure[Idx] = RPT.moveMaxPressure(); + // Account for the reduction in RP due to the rematerialization in an + // optimizable region in which the defined register is a live-in. This + // is exact for live-through region but optimistic in the using region, + // where RP is actually reduced only if maximum RP is reached somewhere + // between the beginning of the region and the rematerializable + // instruction's use. + if (auto It = OptRegions.find(LIRegion); It != OptRegions.end()) { + REMAT_DEBUG(dbgs() << " Live-in in region " << LIRegion << '\n'); + if (ReduceRPInRegion(It, Reg, DAG.LiveIns[LIRegion][Reg], + RematUseful)) + return true; } } - SinkedDefs.push_back(Def); - ImproveOccupancy = NewPressure[I].getOccupancy(ST); - if (ImproveOccupancy > DAG.MinOccupancy) - break; + // 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); + } } + } + + if (IncreaseOccupancy) { + // We were trying to increase occupancy but failed, abort the stage. + REMAT_DEBUG(dbgs() << "Cannot increase occupancy\n"); + Rematerializations.clear(); + return false; + } + REMAT_DEBUG(dbgs() << "Can reduce but not eliminate spilling\n"); + return !Rematerializations.empty(); +} - // Remove defs we just sinked from all regions' list of sinkable defs - for (auto &Def : SinkedDefs) - for (auto TrackedIdx : RematDefToLiveInRegions[Def]) - RematerializableInsts[TrackedIdx].erase(Def); +void PreRARematStage::rematerialize() { + const SIInstrInfo *TII = MF.getSubtarget<GCNSubtarget>().getInstrInfo(); - if (ImproveOccupancy <= DAG.MinOccupancy) - break; + // Collect regions whose RP changes in unpredictable way; we will have to + // fully recompute their RP after all rematerailizations. + DenseSet<unsigned> RecomputeRP; - NewRescheduleRegions[I] = true; - Improved = true; - } + // Rematerialize all instructions. + for (auto &[DefMI, Remat] : Rematerializations) { + MachineBasicBlock::iterator InsertPos(Remat.UseMI); + Register Reg = DefMI->getOperand(0).getReg(); + unsigned SubReg = DefMI->getOperand(0).getSubReg(); + unsigned DefRegion = MIRegion.at(DefMI); - if (!Improved) { - // Occupancy was not improved for all regions that were at MinOccupancy. - // Undo sinking and remove newly rematerialized instructions. - for (auto &Entry : InsertedMIToOldDef) { - MachineInstr *MI = Entry.first; - MachineInstr *OldMI = Entry.second; - Register Reg = MI->getOperand(0).getReg(); - LIS->RemoveMachineInstrFromMaps(*MI); - MI->eraseFromParent(); - OldMI->clearRegisterDeads(Reg); - LIS->removeInterval(Reg); - LIS->createAndComputeVirtRegInterval(Reg); + // 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. + DAG.updateRegionBoundaries(DAG.Regions[DefRegion], DefMI, nullptr); + auto UseRegion = MIRegion.find(Remat.UseMI); + if (UseRegion != MIRegion.end()) { + DAG.updateRegionBoundaries(DAG.Regions[UseRegion->second], InsertPos, + Remat.RematMI); } - return false; - } + DAG.LIS->RemoveMachineInstrFromMaps(*DefMI); + DefMI->eraseFromParent(); - // Occupancy was improved for all regions. - for (auto &Entry : InsertedMIToOldDef) { - MachineInstr *MI = Entry.first; - MachineInstr *OldMI = Entry.second; + // Collect all regions impacted by the rematerialization and update their + // live-in/RP information. + for (unsigned I : Remat.LiveInRegions) { + ImpactedRegions.insert({I, DAG.Pressure[I]}); + GCNRPTracker::LiveRegSet &RegionLiveIns = DAG.LiveIns[I]; - // Remove OldMI from BBLiveInMap since we are sinking it from its MBB. - DAG.BBLiveInMap.erase(OldMI); +#ifdef EXPENSIVE_CHECKS + // All uses are known to be available / live at the remat point. Thus, the + // uses should already be live in to the region. + for (MachineOperand &MO : DefMI->operands()) { + if (!MO.isReg() || !MO.getReg() || !MO.readsReg()) + continue; - // Remove OldMI and update LIS - Register Reg = MI->getOperand(0).getReg(); - LIS->RemoveMachineInstrFromMaps(*OldMI); - OldMI->eraseFromParent(); - LIS->removeInterval(Reg); - LIS->createAndComputeVirtRegInterval(Reg); - } + Register UseReg = MO.getReg(); + if (!UseReg.isVirtual()) + continue; - // Update live-ins, register pressure, and regions caches. - for (auto Idx : ImpactedRegions) { - DAG.LiveIns[Idx] = NewLiveIns[Idx]; - DAG.Pressure[Idx] = NewPressure[Idx]; - DAG.MBBLiveIns.erase(DAG.Regions[Idx].first->getParent()); - } - DAG.Regions = NewRegions; - DAG.RescheduleRegions = NewRescheduleRegions; + LiveInterval &LI = DAG.LIS->getInterval(UseReg); + LaneBitmask LM = DAG.MRI.getMaxLaneMaskForVReg(MO.getReg()); + if (LI.hasSubRanges() && MO.getSubReg()) + LM = DAG.TRI->getSubRegIndexLaneMask(MO.getSubReg()); - if (GCNTrackers) - DAG.RegionLiveOuts.buildLiveRegMap(); + LaneBitmask LiveInMask = RegionLiveIns.at(UseReg); + LaneBitmask UncoveredLanes = LM & ~(LiveInMask & LM); + // If this register has lanes not covered by the LiveIns, be sure they + // do not map to any subrange. ref: + // machine-scheduler-sink-trivial-remats.mir::omitted_subrange + if (UncoveredLanes.any()) { + assert(LI.hasSubRanges()); + for (LiveInterval::SubRange &SR : LI.subranges()) + assert((SR.LaneMask & UncoveredLanes).none()); + } + } +#endif + + // The register is no longer a live-in in all regions but the one that + // contains the single use. In live-through regions, maximum register + // pressure decreases predictably so we can directly update it. In the + // using region, maximum RP may or may not decrease, so we will mark it + // for re-computation after all materializations have taken place. + LaneBitmask PrevMask = RegionLiveIns[Reg]; + RegionLiveIns.erase(Reg); + RegMasks.insert({{I, Remat.RematMI->getOperand(0).getReg()}, PrevMask}); + if (Remat.UseMI->getParent() != DAG.Regions[I].first->getParent()) + DAG.Pressure[I].inc(Reg, PrevMask, LaneBitmask::getNone(), DAG.MRI); + else + RecomputeRP.insert(I); + } + // RP in the region from which the instruction was rematerialized may or may + // not decrease. + ImpactedRegions.insert({DefRegion, DAG.Pressure[DefRegion]}); + RecomputeRP.insert(DefRegion); - SIMachineFunctionInfo &MFI = *MF.getInfo<SIMachineFunctionInfo>(); - MFI.increaseOccupancy(MF, ++DAG.MinOccupancy); + // Recompute live interval to reflect the register's rematerialization. + Register RematReg = Remat.RematMI->getOperand(0).getReg(); + DAG.LIS->removeInterval(RematReg); + DAG.LIS->createAndComputeVirtRegInterval(RematReg); + } - return true; + // All regions impacted by at least one rematerialization must be rescheduled. + // Maximum pressure must also be recomputed for all regions where it changed + // non-predictably and checked against the target occupancy. + AchievedOcc = TargetOcc; + for (auto &[I, OriginalRP] : ImpactedRegions) { + bool IsEmptyRegion = DAG.Regions[I].first == DAG.Regions[I].second; + RescheduleRegions[I] = !IsEmptyRegion; + if (!RecomputeRP.contains(I)) + continue; + + GCNRegPressure RP; + if (IsEmptyRegion) { + RP = getRegPressure(DAG.MRI, DAG.LiveIns[I]); + } else { + GCNDownwardRPTracker RPT(*DAG.LIS); + auto *NonDbgMI = &*skipDebugInstructionsForward(DAG.Regions[I].first, + DAG.Regions[I].second); + if (NonDbgMI == DAG.Regions[I].second) { + // Region is non-empty but contains only debug instructions. + RP = getRegPressure(DAG.MRI, DAG.LiveIns[I]); + } else { + RPT.reset(*NonDbgMI, &DAG.LiveIns[I]); + RPT.advance(DAG.Regions[I].second); + RP = RPT.moveMaxPressure(); + } + } + DAG.Pressure[I] = RP; + AchievedOcc = std::min( + AchievedOcc, RP.getOccupancy(ST, MF.getInfo<SIMachineFunctionInfo>() + ->getDynamicVGPRBlockSize())); + } + REMAT_DEBUG(dbgs() << "Achieved occupancy " << AchievedOcc << "\n"); } +// Copied from MachineLICM bool PreRARematStage::isTriviallyReMaterializable(const MachineInstr &MI) { if (!DAG.TII->isTriviallyReMaterializable(MI)) return false; for (const MachineOperand &MO : MI.all_uses()) { - if (MO.getReg().isVirtual()) - return false; - // We can't remat physreg uses, unless it is a constant or an ignorable // use (e.g. implicit exec use on VALU instructions) if (MO.getReg().isPhysical()) { @@ -1862,51 +2043,89 @@ bool PreRARematStage::isTriviallyReMaterializable(const MachineInstr &MI) { return true; } -// When removing, we will have to check both beginning and ending of the region. -// When inserting, we will only have to check if we are inserting NewMI in front -// of a scheduling region and do not need to check the ending since we will only -// ever be inserting before an already existing MI. -void GCNScheduleDAGMILive::updateRegionBoundaries( - SmallVectorImpl<std::pair<MachineBasicBlock::iterator, - MachineBasicBlock::iterator>> &RegionBoundaries, - MachineBasicBlock::iterator MI, MachineInstr *NewMI, bool Removing) { - unsigned I = 0, E = RegionBoundaries.size(); - // Search for first region of the block where MI is located - while (I != E && MI->getParent() != RegionBoundaries[I].first->getParent()) - ++I; +void PreRARematStage::finalizeGCNSchedStage() { + // We consider that reducing spilling is always beneficial so we never + // rollback rematerializations in such cases. It's also possible that + // rescheduling lowers occupancy over the one achieved just through remats, in + // which case we do not want to rollback either (the rescheduling was already + // reverted in PreRARematStage::shouldRevertScheduling in such cases). + unsigned MaxOcc = std::max(AchievedOcc, DAG.MinOccupancy); + if (!IncreaseOccupancy || MaxOcc >= TargetOcc) + return; - for (; I != E; ++I) { - if (MI->getParent() != RegionBoundaries[I].first->getParent()) - return; + REMAT_DEBUG(dbgs() << "Rolling back all rematerializations\n"); + const SIInstrInfo *TII = MF.getSubtarget<GCNSubtarget>().getInstrInfo(); - if (Removing && MI == RegionBoundaries[I].first && - MI == RegionBoundaries[I].second) { - // MI is in a region with size 1, after removing, the region will be - // size 0, set RegionBegin and RegionEnd to pass end of block iterator. - RegionBoundaries[I] = - std::pair(MI->getParent()->end(), MI->getParent()->end()); - return; - } - if (MI == RegionBoundaries[I].first) { - if (Removing) - RegionBoundaries[I] = - std::pair(std::next(MI), RegionBoundaries[I].second); - else - // Inserted NewMI in front of region, set new RegionBegin to NewMI - RegionBoundaries[I] = std::pair(MachineBasicBlock::iterator(NewMI), - RegionBoundaries[I].second); - return; - } - if (Removing && MI == RegionBoundaries[I].second) { - RegionBoundaries[I] = std::pair(RegionBoundaries[I].first, std::prev(MI)); - return; + // Rollback the rematerializations. + for (const auto &[DefMI, Remat] : Rematerializations) { + MachineInstr &RematMI = *Remat.RematMI; + unsigned DefRegion = MIRegion.at(DefMI); + MachineBasicBlock::iterator InsertPos(DAG.Regions[DefRegion].second); + MachineBasicBlock *MBB = RegionBB[DefRegion]; + Register Reg = RematMI.getOperand(0).getReg(); + unsigned SubReg = RematMI.getOperand(0).getSubReg(); + + // Re-rematerialize MI at the end of its original region. Note that it may + // not be rematerialized exactly in the same position as originally within + // the region, but it should not matter much. + TII->reMaterialize(*MBB, InsertPos, Reg, SubReg, RematMI, *DAG.TRI); + MachineInstr *NewMI = &*std::prev(InsertPos); + NewMI->getOperand(0).setSubReg(SubReg); + DAG.LIS->InsertMachineInstrInMaps(*NewMI); + + auto UseRegion = MIRegion.find(Remat.UseMI); + if (UseRegion != MIRegion.end()) { + DAG.updateRegionBoundaries(DAG.Regions[UseRegion->second], RematMI, + nullptr); } + DAG.updateRegionBoundaries(DAG.Regions[DefRegion], InsertPos, NewMI); + + // Erase rematerialized MI. + DAG.LIS->RemoveMachineInstrFromMaps(RematMI); + RematMI.eraseFromParent(); + + // Recompute live interval for the re-rematerialized register + DAG.LIS->removeInterval(Reg); + DAG.LIS->createAndComputeVirtRegInterval(Reg); + + // 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. + for (auto &[I, OriginalRP] : ImpactedRegions) + DAG.Pressure[I] = OriginalRP; + + GCNSchedStage::finalizeGCNSchedStage(); +} + +void GCNScheduleDAGMILive::updateRegionBoundaries( + RegionBoundaries &RegionBounds, MachineBasicBlock::iterator MI, + MachineInstr *NewMI) { + assert((!NewMI || NewMI != RegionBounds.second) && + "cannot remove at region end"); + + if (RegionBounds.first == RegionBounds.second) { + assert(NewMI && "cannot remove from an empty region"); + RegionBounds.first = NewMI; + return; + } + + // We only care for modifications at the beginning of a non-empty region since + // the upper region boundary is exclusive. + if (MI != RegionBounds.first) + return; + if (!NewMI) + RegionBounds.first = std::next(MI); // Removal + else + RegionBounds.first = NewMI; // Insertion } static bool hasIGLPInstrs(ScheduleDAGInstrs *DAG) { - return any_of(*DAG, [](MachineBasicBlock::iterator MI) { - return isIGLPMutationOnly(MI->getOpcode()); + const SIInstrInfo *SII = static_cast<const SIInstrInfo *>(DAG->TII); + return any_of(*DAG, [SII](MachineBasicBlock::iterator MI) { + return SII->isIGLPMutationOnly(MI->getOpcode()); }); } |