[AMDGPU] Dynamically set load latency in the scheduler

Author: kerbowa

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

@@ -1130,6 +1130,32 @@ void UnclusteredHighRPStage::finalizeGCNSchedStage() {
   GCNSchedStage::finalizeGCNSchedStage();
 }
 
+bool ILPInitialScheduleStage::initGCNSchedStage() {
+  if (!GCNSchedStage::initGCNSchedStage())
+    return false;
+
+  const SIInstrInfo *TII = ST.getInstrInfo();
+  OriginalLoadLatencyScaleFactor = TII->getLoadLatencyScaleFactor();
+  const unsigned ILPLoadLatencyScaleFactorDefault = 300;
+  if (ILPLoadLatencyScaleFactorDefault > TII->getLoadLatencyScaleFactor())
+    TII->setLoadLatencyScaleFactor(ILPLoadLatencyScaleFactorDefault);
+
+  LLVM_DEBUG(dbgs() << "ILP Initial Schedule: Set load latency scale factor to "
+                    << TII->getLoadLatencyScaleFactor() << '\n');
+  return true;
+}
+
+void ILPInitialScheduleStage::finalizeGCNSchedStage() {
+  const SIInstrInfo *TII = ST.getInstrInfo();
+  TII->setLoadLatencyScaleFactor(OriginalLoadLatencyScaleFactor);
+
+  LLVM_DEBUG(
+      dbgs() << "ILP Initial Schedule: Restored load latency scale factor to "
+             << OriginalLoadLatencyScaleFactor << "\n");
+
+  GCNSchedStage::finalizeGCNSchedStage();
+}
+
 bool GCNSchedStage::initGCNRegion() {
   // Check whether this new region is also a new block.
   if (DAG.RegionBegin->getParent() != CurrentMBB)

llvm/lib/Target/AMDGPU/GCNSchedStrategy.h

@@ -474,8 +474,13 @@ class PreRARematStage : public GCNSchedStage {
 };
 
 class ILPInitialScheduleStage : public GCNSchedStage {
+private:
+  unsigned OriginalLoadLatencyScaleFactor = 0;
+
 public:
   bool shouldRevertScheduling(unsigned WavesAfter) override;
+  bool initGCNSchedStage() override;
+  void finalizeGCNSchedStage() override;
 
   ILPInitialScheduleStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG)
       : GCNSchedStage(StageID, DAG) {}

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -60,9 +60,15 @@ static cl::opt<bool> Fix16BitCopies(
   cl::init(true),
   cl::ReallyHidden);
 
+static cl::opt<unsigned> AMDGPULoadLatencyScaleFactor(
+    "amdgpu-load-latency-scale-factor",
+    cl::desc("Scale factor for load instruction latency. Final latency is "
+             "scalled by `Factor / 100 * Latency`."),
+    cl::init(100), cl::ReallyHidden);
+
 SIInstrInfo::SIInstrInfo(const GCNSubtarget &ST)
-  : AMDGPUGenInstrInfo(AMDGPU::ADJCALLSTACKUP, AMDGPU::ADJCALLSTACKDOWN),
-    RI(ST), ST(ST) {
+    : AMDGPUGenInstrInfo(AMDGPU::ADJCALLSTACKUP, AMDGPU::ADJCALLSTACKDOWN),
+      RI(ST), ST(ST), LoadLatencyScaleFactor(AMDGPULoadLatencyScaleFactor) {
   SchedModel.init(&ST);
 }
 
@@ -9792,6 +9798,15 @@ unsigned SIInstrInfo::getInstrLatency(const InstrItineraryData *ItinData,
   return SchedModel.computeInstrLatency(&MI);
 }
 
+unsigned SIInstrInfo::getInstrLatency(const TargetSchedModel &TargetSchedModel,
+                                      const MachineInstr &MI) const {
+  unsigned Latency = TargetInstrInfo::getInstrLatency(TargetSchedModel, MI);
+  if (MI.mayLoad())
+    Latency *= LoadLatencyScaleFactor / 100;
+
+  return Latency;
+}
+
 InstructionUniformity
 SIInstrInfo::getGenericInstructionUniformity(const MachineInstr &MI) const {
   unsigned opcode = MI.getOpcode();

llvm/lib/Target/AMDGPU/SIInstrInfo.h

@@ -88,6 +88,9 @@ class SIInstrInfo final : public AMDGPUGenInstrInfo {
   const GCNSubtarget &ST;
   TargetSchedModel SchedModel;
   mutable std::unique_ptr<AMDGPUMIRFormatter> Formatter;
+  // Final load latency in the machine model is scalled by
+  // `Factor / 100 * Latency`
+  mutable unsigned LoadLatencyScaleFactor = 100;
 
   // The inverse predicate should have the negative value.
   enum BranchPredicate {
@@ -106,6 +109,12 @@ class SIInstrInfo final : public AMDGPUGenInstrInfo {
   static BranchPredicate getBranchPredicate(unsigned Opcode);
 
 public:
+  void setLoadLatencyScaleFactor(unsigned Factor) const {
+    LoadLatencyScaleFactor = Factor;
+  }
+
+  unsigned getLoadLatencyScaleFactor() const { return LoadLatencyScaleFactor; }
+
   unsigned buildExtractSubReg(MachineBasicBlock::iterator MI,
                               MachineRegisterInfo &MRI,
                               const MachineOperand &SuperReg,
@@ -1462,6 +1471,9 @@ class SIInstrInfo final : public AMDGPUGenInstrInfo {
                            const MachineInstr &MI,
                            unsigned *PredCost = nullptr) const override;
 
+  unsigned getInstrLatency(const TargetSchedModel &TargetSchedModel,
+                           const MachineInstr &MI) const override;
+
   InstructionUniformity
   getInstructionUniformity(const MachineInstr &MI) const override final;