[AMDGPU] Track occupancy in MFI

[rampitec]

Summary

Keep track of achieved occupancy in SIMachineFunctionInfo.
At the moment we have a lot of duplicated or even missed code to
query and maintain occupancy info. Record it in the MFI and
query in a single call. Interfaces:

• getOccupancy() - returns current recorded achieved occupancy.
• getMinAllowedOccupancy() - returns lesser of the achieved occupancy

and the lowest occupancy we are ready to tolerate. For example if
a kernel is memory bound we are ready to tolerate 4 waves.

• limitOccupancy() - record occupancy level if we have to lower it.
• increaseOccupancy() - record occupancy if scheduler managed to

increase the occupancy.

MFI takes care of integrating different checks affecting occupancy,
including LDS use and waves-per-eu attribute. Note that scheduler
starts with not yet known register pressure, so has to record either
limit or increase in occupancy after it is done. Later passes can
just query a resulting value.

New interface is used in the active scheduler and NFC wrt its work.
Changes are also made to experimental schedulers to use it and record
an occupancy after they are done. Before the change waves-per-eu was
ignored by experimental schedulers and tolerance window for memory
bound kernels was not used.

[arsenm]

The problem with this is until we can serialize MachineFunctionInfo, this is going fo further degrade MIR tests.

lib/Target/AMDGPU/SIISelLowering.cpp

4259–4263

Why is this here?

[rampitec]

In D47509#1115824, @arsenm wrote:

The problem with this is until we can serialize MachineFunctionInfo, this is going fo further degrade MIR tests.

That did not degrade mir tests so far because all what can be recorded from an original function is lost anyway. Default is to have full 10 waves, unless further limited by an attribute (which is lost for mir), static lds usage (which is lost as well) or scheduler register usage (which was never here until this change).

lib/Target/AMDGPU/SIISelLowering.cpp

4259–4263

allocateLDSGlobal() belongs to AMDGPUMachineFunction, not SIMachineFunctionInfo. The choice is to either make it virtual or handle in the single place it is really used.

[rampitec]

lib/Target/AMDGPU/SIISelLowering.cpp

4259–4263

In fact the other choice is to create SIMachineFunctionInfo::allocateLDSGlobal() which will call AMDGPUMachineFunction::allocateLDSGlobal() and then limitOccupancy(), then call it from here. Same thing basically, but more obscure to my taste. The only sound solution is to virtualize MFI, but I am not really sure it is worth it.

[arsenm]

lib/Target/AMDGPU/SIISelLowering.cpp

4259–4263

I don't see a reason to update this during lowering. We don't do anything with occupancy information in the DAG? Why can't you just adjust this once after the function is selected and the LDS size is known?

[rampitec]

Moved LDS processing into finalizeLowering().

lib/Target/AMDGPU/SIISelLowering.cpp

4259–4263

Thanks! Moved to finalizeLowering().

[vpykhtin]

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

185

may be one call to limitOccupancy with min(getMaxWavesPerEU and getOcc..LMS)?

[rampitec]

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

185

It is really the same, it will even produce the same code after compilation. The whole intent of "limit" semantics is that you can take into account just a single factor and omit all others. I would prefer to keep this snippet to emphasize that ;)

lib/Target/AMDGPU/GCNIterativeScheduler.cpp

for (auto R : Regions) {
     NewOcc = std::min(NewOcc, MaxRP.getOccupancy(ST));
     if (NewOcc <= Occ)
       break;
 
     setBestSchedule(*R, MinSchedule, MaxRP);
   }
   LLVM_DEBUG(dbgs() << "New occupancy = " << NewOcc
                     << ", prev occupancy = " << Occ << '\n');
+  if (NewOcc > Occ) {
+    SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+    MFI->increaseOccupancy(MF, NewOcc);
+  }
+
   return std::max(NewOcc, Occ);
 }
 
 void GCNIterativeScheduler::scheduleLegacyMaxOccupancy(
   bool TryMaximizeOccupancy) {
   const auto &ST = MF.getSubtarget<SISubtarget>();
-  auto TgtOcc = ST.getOccupancyWithLocalMemSize(MF);
+  SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  auto TgtOcc = MFI->getMinAllowedOccupancy();
 
   sortRegionsByPressure(TgtOcc);
   auto Occ = Regions.front()->MaxPressure.getOccupancy(ST);
 
   if (TryMaximizeOccupancy && Occ < TgtOcc)
     Occ = tryMaximizeOccupancy(TgtOcc);
 
   // This is really weird but for some magic scheduling regions twice
   // gives performance improvement
   const int NumPasses = Occ < TgtOcc ? 2 : 1;
 
   TgtOcc = std::min(Occ, TgtOcc);
   LLVM_DEBUG(dbgs() << "Scheduling using default scheduler, "
                        "target occupancy = "
                     << TgtOcc << '\n');
   GCNMaxOccupancySchedStrategy LStrgy(Context);
+  unsigned FinalOccupancy = std::min(Occ, MFI->getOccupancy());
 
   for (int I = 0; I < NumPasses; ++I) {
     // running first pass with TargetOccupancy = 0 mimics previous scheduling
     // approach and is a performance magic
     LStrgy.setTargetOccupancy(I == 0 ? 0 : TgtOcc);
     for (auto R : Regions) {
       OverrideLegacyStrategy Ovr(*R, LStrgy, *this);
 
       Ovr.schedule();
       const auto RP = getRegionPressure(*R);
       LLVM_DEBUG(printSchedRP(dbgs(), R->MaxPressure, RP));
 
       if (RP.getOccupancy(ST) < TgtOcc) {
         LLVM_DEBUG(dbgs() << "Didn't fit into target occupancy O" << TgtOcc);
         if (R->BestSchedule.get() &&
             R->BestSchedule->MaxPressure.getOccupancy(ST) >= TgtOcc) {
           LLVM_DEBUG(dbgs() << ", scheduling minimal register\n");
           scheduleBest(*R);
         } else {
           LLVM_DEBUG(dbgs() << ", restoring\n");
           Ovr.restoreOrder();
           assert(R->MaxPressure.getOccupancy(ST) >= TgtOcc);
         }
       }
+      FinalOccupancy = std::min(FinalOccupancy, RP.getOccupancy(ST));
     }
   }
+  MFI->limitOccupancy(FinalOccupancy);
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // Minimal Register Strategy
 
 void GCNIterativeScheduler::scheduleMinReg(bool force) {
   const auto &ST = MF.getSubtarget<SISubtarget>();
-  const auto TgtOcc = ST.getOccupancyWithLocalMemSize(MF);
+  const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  const auto TgtOcc = MFI->getOccupancy();
   sortRegionsByPressure(TgtOcc);
 
   auto MaxPressure = Regions.front()->MaxPressure;
   for (auto R : Regions) {
     if (!force && R->MaxPressure.less(ST, MaxPressure, TgtOcc))
       break;
 
     BuildDAG DAG(*R, *this);

 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // ILP scheduler port
 
 void GCNIterativeScheduler::scheduleILP(
   bool TryMaximizeOccupancy) {
   const auto &ST = MF.getSubtarget<SISubtarget>();
-  const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
-  auto TgtOcc = std::min(ST.getOccupancyWithLocalMemSize(MF),
-                         MFI->getMaxWavesPerEU());
+  SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  auto TgtOcc = MFI->getMinAllowedOccupancy();
 
   sortRegionsByPressure(TgtOcc);
   auto Occ = Regions.front()->MaxPressure.getOccupancy(ST);
 
   if (TryMaximizeOccupancy && Occ < TgtOcc)
     Occ = tryMaximizeOccupancy(TgtOcc);
 
   TgtOcc = std::min(Occ, TgtOcc);
   LLVM_DEBUG(dbgs() << "Scheduling using default scheduler, "
                        "target occupancy = "
                     << TgtOcc << '\n');
 
+  unsigned FinalOccupancy = std::min(Occ, MFI->getOccupancy());
   for (auto R : Regions) {
     BuildDAG DAG(*R, *this);
     const auto ILPSchedule = makeGCNILPScheduler(DAG.getBottomRoots(), *this);
 
     const auto RP = getSchedulePressure(*R, ILPSchedule);
     LLVM_DEBUG(printSchedRP(dbgs(), R->MaxPressure, RP));
 
     if (RP.getOccupancy(ST) < TgtOcc) {
       LLVM_DEBUG(dbgs() << "Didn't fit into target occupancy O" << TgtOcc);
       if (R->BestSchedule.get() &&
         R->BestSchedule->MaxPressure.getOccupancy(ST) >= TgtOcc) {
         LLVM_DEBUG(dbgs() << ", scheduling minimal register\n");
         scheduleBest(*R);
       }
     } else {
       scheduleRegion(*R, ILPSchedule, RP);
       LLVM_DEBUG(printSchedResult(dbgs(), R, RP));
+      FinalOccupancy = std::min(FinalOccupancy, RP.getOccupancy(ST));
     }
   }
+  MFI->limitOccupancy(FinalOccupancy);
 }

lib/Target/AMDGPU/GCNSchedStrategy.h

public:
 
   void setTargetOccupancy(unsigned Occ) { TargetOccupancy = Occ; }
 };
 
 class GCNScheduleDAGMILive : public ScheduleDAGMILive {
 
   const SISubtarget &ST;
 
-  const SIMachineFunctionInfo &MFI;
+  SIMachineFunctionInfo &MFI;
 
   // Occupancy target at the beginning of function scheduling cycle.
   unsigned StartingOccupancy;
 
   // Minimal real occupancy recorder for the function.
   unsigned MinOccupancy;
 
   // Scheduling stage number.

lib/Target/AMDGPU/GCNSchedStrategy.cpp

SUnit *GCNMaxOccupancySchedStrategy::pickNode(bool &IsTopNode) {
   return SU;
 }
 
 GCNScheduleDAGMILive::GCNScheduleDAGMILive(MachineSchedContext *C,
                         std::unique_ptr<MachineSchedStrategy> S) :
   ScheduleDAGMILive(C, std::move(S)),
   ST(MF.getSubtarget<SISubtarget>()),
   MFI(*MF.getInfo<SIMachineFunctionInfo>()),
-  StartingOccupancy(std::min(ST.getOccupancyWithLocalMemSize(MFI.getLDSSize(),
-                                                             MF.getFunction()),
-                             MFI.getMaxWavesPerEU())),
+  StartingOccupancy(MFI.getOccupancy()),
   MinOccupancy(StartingOccupancy), Stage(0), RegionIdx(0) {
 
   LLVM_DEBUG(dbgs() << "Starting occupancy is " << StartingOccupancy << ".\n");
 }
 
 void GCNScheduleDAGMILive::schedule() {
   if (Stage == 0) {
     // Just record regions at the first pass.
void GCNScheduleDAGMILive::schedule() {
   LLVM_DEBUG(dbgs() << "Occupancy before scheduling: " << WavesBefore
                     << ", after " << WavesAfter << ".\n");
 
   // We could not keep current target occupancy because of the just scheduled
   // region. Record new occupancy for next scheduling cycle.
   unsigned NewOccupancy = std::max(WavesAfter, WavesBefore);
   // Allow memory bound functions to drop to 4 waves if not limited by an
   // attribute.
-  unsigned MinMemBoundWaves = std::max(MFI.getMinWavesPerEU(), 4u);
   if (WavesAfter < WavesBefore && WavesAfter < MinOccupancy &&
-      WavesAfter >= MinMemBoundWaves &&
-      (MFI.isMemoryBound() || MFI.needsWaveLimiter())) {
+      WavesAfter >= MFI.getMinAllowedOccupancy()) {
     LLVM_DEBUG(dbgs() << "Function is memory bound, allow occupancy drop up to "
-                      << MinMemBoundWaves << " waves\n");
+                      << MFI.getMinAllowedOccupancy() << " waves\n");
     NewOccupancy = WavesAfter;
   }
   if (NewOccupancy < MinOccupancy) {
     MinOccupancy = NewOccupancy;
+    MFI.limitOccupancy(MinOccupancy);
     LLVM_DEBUG(dbgs() << "Occupancy lowered for the function to "
                       << MinOccupancy << ".\n");
   }
 
   if (WavesAfter >= MinOccupancy) {
     Pressure[RegionIdx] = PressureAfter;
     return;
   }

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::LowerGlobalAddress(AMDGPUMachineFunction *MFI,
 
   if (GSD->getAddressSpace() != AMDGPUASI.CONSTANT_ADDRESS &&
       GSD->getAddressSpace() != AMDGPUASI.CONSTANT_ADDRESS_32BIT &&
       GSD->getAddressSpace() != AMDGPUASI.GLOBAL_ADDRESS &&
       // FIXME: It isn't correct to rely on the type of the pointer. This should
       // be removed when address space 0 is 64-bit.
       !GV->getType()->getElementType()->isFunctionTy())
     return AMDGPUTargetLowering::LowerGlobalAddress(MFI, Op, DAG);
 
   SDLoc DL(GSD);
   EVT PtrVT = Op.getValueType();
 
   if (shouldEmitFixup(GV))
     return buildPCRelGlobalAddress(DAG, GV, DL, GSD->getOffset(), PtrVT);
   else if (shouldEmitPCReloc(GV))
     return buildPCRelGlobalAddress(DAG, GV, DL, GSD->getOffset(), PtrVT,
                                    SIInstrInfo::MO_REL32);
 
   SDValue GOTAddr = buildPCRelGlobalAddress(DAG, GV, DL, 0, PtrVT,
                                             SIInstrInfo::MO_GOTPCREL32);
 
if (NeedSP) {
     MRI.replaceRegWith(AMDGPU::SP_REG, Info->getStackPtrOffsetReg());
   }
 
   MRI.replaceRegWith(AMDGPU::PRIVATE_RSRC_REG, Info->getScratchRSrcReg());
   MRI.replaceRegWith(AMDGPU::FP_REG, Info->getFrameOffsetReg());
   MRI.replaceRegWith(AMDGPU::SCRATCH_WAVE_OFFSET_REG,
                      Info->getScratchWaveOffsetReg());
 
+  Info->limitOccupancy(MF);
+
   TargetLoweringBase::finalizeLowering(MF);
 }
 
 void SITargetLowering::computeKnownBitsForFrameIndex(const SDValue Op,
                                                      KnownBits &Known,
                                                      const APInt &DemandedElts,
                                                      const SelectionDAG &DAG,
                                                      unsigned Depth) const {

lib/Target/AMDGPU/SIMachineFunctionInfo.h

private:
 
   // The hard-wired high half of the address of the global information table
   // for AMDPAL OS type. 0xffffffff represents no hard-wired high half, since
   // current hardware only allows a 16 bit value.
   unsigned GITPtrHigh;
 
   unsigned HighBitsOf32BitAddress;
 
+  // Current recorded maximum possible occupancy.
+  unsigned Occupancy;
+
   MCPhysReg getNextUserSGPR() const;
 
   MCPhysReg getNextSystemSGPR() const;
 
 public:
   struct SpilledReg {
     unsigned VGPR = 0;
     int Lane = -1;
public:
   const AMDGPUImagePseudoSourceValue *getImagePSV(const SIInstrInfo &TII,
                                                   const Value *ImgRsrc) {
     assert(ImgRsrc);
     auto PSV = ImagePSVs.try_emplace(
       ImgRsrc,
       llvm::make_unique<AMDGPUImagePseudoSourceValue>(TII));
     return PSV.first->second.get();
   }
+
+  unsigned getOccupancy() const {
+    return Occupancy;
+  }
+
+  unsigned getMinAllowedOccupancy() const {
+    if (!isMemoryBound() && !needsWaveLimiter())
+      return Occupancy;
+    return (Occupancy < 4) ? Occupancy : 4;
+  }
+
+  void limitOccupancy(const MachineFunction &MF);
+
+  void limitOccupancy(unsigned Limit) {
+    if (Occupancy > Limit)
+      Occupancy = Limit;
+  }
+
+  void increaseOccupancy(const MachineFunction &MF, unsigned Limit) {
+    if (Occupancy < Limit)
+      Occupancy = Limit;
+    limitOccupancy(MF);
+  }
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_SIMACHINEFUNCTIONINFO_H

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

: AMDGPUMachineFunction(MF),
     ImplicitArgPtr(false),
     GITPtrHigh(0xffffffff),
     HighBitsOf32BitAddress(0) {
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
   const Function &F = MF.getFunction();
   FlatWorkGroupSizes = ST.getFlatWorkGroupSizes(F);
   WavesPerEU = ST.getWavesPerEU(F);
 
+  Occupancy = getMaxWavesPerEU();
+  limitOccupancy(MF);
+
   if (!isEntryFunction()) {
     // Non-entry functions have no special inputs for now, other registers
     // required for scratch access.
     ScratchRSrcReg = AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3;
     ScratchWaveOffsetReg = AMDGPU::SGPR4;
     FrameOffsetReg = AMDGPU::SGPR5;
     StackPtrOffsetReg = AMDGPU::SGPR32;
 
if (!S.empty())
     S.consumeInteger(0, GITPtrHigh);
 
   A = F.getFnAttribute("amdgpu-32bit-address-high-bits");
   S = A.getValueAsString();
   if (!S.empty())
     S.consumeInteger(0, HighBitsOf32BitAddress);
 }
 
+void SIMachineFunctionInfo::limitOccupancy(const MachineFunction &MF) {
+  limitOccupancy(getMaxWavesPerEU());
+  const SISubtarget& ST = MF.getSubtarget<SISubtarget>();
+  limitOccupancy(ST.getOccupancyWithLocalMemSize(getLDSSize(),
+                 MF.getFunction()));
+}
+
 unsigned SIMachineFunctionInfo::addPrivateSegmentBuffer(
   const SIRegisterInfo &TRI) {
   ArgInfo.PrivateSegmentBuffer =
     ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_128RegClass));
   NumUserSGPRs += 4;
   return ArgInfo.PrivateSegmentBuffer.getRegister();
 }