[AMDGPU] Reschedule loads in clauses to improve throughput

After clauses are formed their internal loads can be reordered
to facilitate some additional opportunities for overlapping
computation.
This late stage rescheduling causes no change register pressure.

Author: perlfu

llvm/lib/Target/AMDGPU/SIPostRABundler.cpp

@@ -44,6 +44,7 @@ class SIPostRABundler : public MachineFunctionPass {
   }
 
 private:
+  const SIInstrInfo *TII = nullptr;
   const SIRegisterInfo *TRI;
 
   SmallSet<Register, 16> Defs;
@@ -54,6 +55,9 @@ class SIPostRABundler : public MachineFunctionPass {
   bool isBundleCandidate(const MachineInstr &MI) const;
   bool isDependentLoad(const MachineInstr &MI) const;
   bool canBundle(const MachineInstr &MI, const MachineInstr &NextMI) const;
+  void reorderLoads(MachineBasicBlock &MBB,
+                    MachineBasicBlock::instr_iterator &BundleStart,
+                    MachineBasicBlock::instr_iterator Next);
 };
 
 constexpr uint64_t MemFlags = SIInstrFlags::MTBUF | SIInstrFlags::MUBUF |
@@ -121,10 +125,132 @@ bool SIPostRABundler::canBundle(const MachineInstr &MI,
           !isDependentLoad(NextMI));
 }
 
+void SIPostRABundler::reorderLoads(
+    MachineBasicBlock &MBB, MachineBasicBlock::instr_iterator &BundleStart,
+    MachineBasicBlock::instr_iterator Next) {
+  auto II = BundleStart;
+  if (!TII->isMIMG(II->getOpcode()) || II->mayStore())
+    return;
+
+  LLVM_DEBUG(dbgs() << "Begin bundle reorder\n");
+
+  // Collect clause
+  SmallVector<MachineInstr *> Clause;
+  for (auto II = BundleStart; II != Next; ++II)
+    Clause.push_back(&*II);
+
+  // Search to find the usage distance of each defined register in the clause.
+  const int MaxSearch = 100;
+  SmallSet<Register, 16> DefRegs(Defs);
+  SmallSet<unsigned, 16> Distances;
+  DenseMap<Register, unsigned> UseDistance;
+  unsigned Dist = 0;
+  for (MachineBasicBlock::iterator SearchI = Next;
+       SearchI != MBB.end() && Dist < MaxSearch && !DefRegs.empty();
+       ++SearchI, ++Dist) {
+    SmallVector<Register, 4> Found;
+    // FIXME: fix search efficiency
+    for (Register DefReg : DefRegs) {
+      if (SearchI->readsRegister(DefReg, TRI))
+        Found.push_back(DefReg);
+    }
+    for (Register Reg : Found) {
+      UseDistance[Reg] = Dist;
+      DefRegs.erase(Reg);
+      Distances.insert(Dist);
+    }
+  }
+
+  if (Distances.size() <= 1)
+    return;
+
+  std::vector<std::pair<MachineInstr *, unsigned>> Schedule;
+  unsigned TotalOrder = Dist + 1;
+  bool Reorder = false;
+  for (MachineInstr *MI : Clause) {
+    unsigned Order = TotalOrder++;
+    if (MI->getNumExplicitDefs() >= 0) {
+      Register Reg = MI->defs().begin()->getReg();
+      if (!UseDistance.contains(Reg))
+        continue;
+      Order = std::min(Order, UseDistance[Reg]);
+      Reorder = true;
+    }
+    LLVM_DEBUG(dbgs() << "Order: " << Order << ", MI: " << *MI);
+    Schedule.push_back(std::pair(MI, Order));
+  }
+
+  if (!Reorder)
+    return;
+
+  std::sort(Schedule.begin(), Schedule.end(),
+            [](std::pair<MachineInstr *, unsigned> A,
+               std::pair<MachineInstr *, unsigned> B) {
+              return A.second < B.second;
+            });
+
+  // Rebuild clause order.
+  // Schedule holds ideal order for the load operations; however, each def
+  // can only be scheduled when it will no longer clobber any uses.
+  Clause.clear();
+  while (!Schedule.empty()) {
+    auto It = Schedule.begin();
+    while (It != Schedule.end()) {
+      MachineInstr *MI = It->first;
+
+      LLVM_DEBUG(dbgs() << "Try schedule: " << *MI);
+
+      if (MI->getNumExplicitDefs() == 0) {
+        // No defs, always schedule.
+        Clause.push_back(MI);
+        break;
+      }
+
+      // FIXME: make this scan more efficient
+      Register Reg = MI->defs().begin()->getReg();
+      bool ClobbersUse = false;
+      for (auto SearchIt = Schedule.begin(); SearchIt != Schedule.end();
+           ++SearchIt) {
+        // We are allowed to clobber our own uses.
+        if (SearchIt == It)
+          continue;
+        if (SearchIt->first->readsRegister(Reg, TRI)) {
+          ClobbersUse = true;
+          break;
+        }
+      }
+      if (ClobbersUse) {
+        // Use is clobbered; try next def in the schedule.
+        It++;
+        LLVM_DEBUG(dbgs() << "  Clobbers uses\n");
+        continue;
+      }
+
+      // Safe to schedule.
+      LLVM_DEBUG(dbgs() << "  OK!\n");
+      Clause.push_back(MI);
+      break;
+    }
+    assert(It != Schedule.end());
+    Schedule.erase(It);
+  }
+
+  // Apply order to instructions.
+  for (MachineInstr *MI : Clause)
+    MI->moveBefore(&*Next);
+
+  // FIXME: update kill flags
+
+  // Update start of bundle.
+  BundleStart = Clause[0]->getIterator();
+}
+
 bool SIPostRABundler::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  TII = ST.getInstrInfo();
   TRI = MF.getSubtarget<GCNSubtarget>().getRegisterInfo();
   BitVector BundleUsedRegUnits(TRI->getNumRegUnits());
   BitVector KillUsedRegUnits(TRI->getNumRegUnits());
@@ -214,6 +340,7 @@ bool SIPostRABundler::runOnMachineFunction(MachineFunction &MF) {
           BundleUsedRegUnits.reset();
         }
 
+        reorderLoads(MBB, BundleStart, Next);
         finalizeBundle(MBB, BundleStart, Next);
       }
 

llvm/test/CodeGen/AMDGPU/issue92561-restore-undef-scc-verifier-error.ll

@@ -61,15 +61,16 @@ define void @issue92561(ptr addrspace(1) %arg) {
 ; SDAG-NEXT:    s_mov_b32 s7, s12
 ; SDAG-NEXT:    s_clause 0x2
 ; SDAG-NEXT:    image_sample_c_lz v0, [v1, v1, v0, v1], s[0:7], s[12:15] dmask:0x1 dim:SQ_RSRC_IMG_2D_ARRAY
-; SDAG-NEXT:    image_sample_c_lz v3, [v1, v1, v1, v1], s[0:7], s[12:15] dmask:0x1 dim:SQ_RSRC_IMG_2D_ARRAY
 ; SDAG-NEXT:    image_sample_c_lz v2, [v1, v2, v1, v1], s[0:7], s[12:15] dmask:0x1 dim:SQ_RSRC_IMG_2D_ARRAY
+; SDAG-NEXT:    image_sample_c_lz v3, [v1, v1, v1, v1], s[0:7], s[12:15] dmask:0x1 dim:SQ_RSRC_IMG_2D_ARRAY
 ; SDAG-NEXT:    v_mov_b32_e32 v4, v1
 ; SDAG-NEXT:    s_waitcnt vmcnt(2)
 ; SDAG-NEXT:    v_add_f32_e32 v0, v9, v0
-; SDAG-NEXT:    s_waitcnt vmcnt(0)
-; SDAG-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_2)
+; SDAG-NEXT:    s_waitcnt vmcnt(1)
+; SDAG-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_2) | instid1(VALU_DEP_2)
 ; SDAG-NEXT:    v_add_f32_e32 v0, v2, v0
 ; SDAG-NEXT:    v_mov_b32_e32 v2, v1
+; SDAG-NEXT:    s_waitcnt vmcnt(0)
 ; SDAG-NEXT:    v_dual_add_f32 v0, v3, v0 :: v_dual_mov_b32 v3, v1
 ; SDAG-NEXT:    s_delay_alu instid0(VALU_DEP_1)
 ; SDAG-NEXT:    v_mul_f32_e32 v0, 0x3e800000, v0