[AIEX] Schedule SWP epilogue with "free" instructions #247

andcarminati · 2024-12-10T10:57:19Z

This PR adds support for Epilogue scheduling. In this way, it also adds:

Support for top-down scheduling with explicit emission cycle.
Top-down logic for AIEMachineScheduler.

I recommend reviewing the pull request in the order of commits, although some of them are closely related, so I plan to combine them in the future.

Ongoing work related to EmitFixedSUnits: we current add all WAR and RAW dependencies related to the top-insert and the rest. However, the bot-insert handling can be changed to use bot register events as well.

llvm/lib/CodeGen/MachineScheduler.cpp

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

andcarminati · 2024-12-10T12:08:22Z

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

+            std::optional<unsigned> OptSrcCycle =
+                InstrItins->getOperandCycle(SrcClass, OpNum);
+            assert(OptSrcCycle);
+            int Latency = *OptSrcCycle;


nit: const.

llvm/lib/Target/AIE/AIEMachineScheduler.h

martien-de-jong · 2024-12-10T12:03:19Z

llvm/lib/CodeGen/MachineScheduler.cpp

+  // We want to insert above it.
+  return std::lower_bound(IsTopNode ? begin() : bottom(),
+                          IsTopNode ? top() : end(), *EmissionCycle,
+                          HasGreaterOrLessOrEqEmissionCycle);


yeah just split in two lower_bound calls.

martien-de-jong · 2024-12-10T12:38:34Z

llvm/unittests/CodeGen/ScheduleDAGMITestUtils.cpp

-      IsPreRA(IsPreRA), SchedZone(SchedBoundary::BotQID, "Zone") {}
+      IsPreRA(IsPreRA),
+      SchedZone(IsTopDown ? SchedBoundary::TopQID : SchedBoundary::BotQID,
+                "Zone") {}


It looks as if "Zone" could be more descriptive

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

martien-de-jong · 2024-12-10T13:04:42Z

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

-    ArrayRef<MachineBundle> TopFixedBundles;
+    ArrayRef<MachineBundle> TopFixedBundles =
+        RegionBegin == BB->begin() ? ArrayRef<MachineBundle>(BS.TopInsert)
+                                   : ArrayRef<MachineBundle>();


Check: TopFixedBundles was empty before, triggering no further action.

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

gbossu · 2024-12-10T13:43:40Z

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

@@ -639,7 +640,7 @@ void AIEPostRASchedStrategy::commitBlockSchedule(MachineBasicBlock *BB) {

  // Safety margin, swp epilogue
  // Note that the prologue is handled in a different way. See enterMBB.


Comment is out of date, we now only handle the safety margin here.

Curious: Can we have both a safety margin and a top-fixed region? If not, can we assert it doesn't happen?

We cannot have! If we need to supply safety margin for swp loop, it means an incorrect schedule. We cannot calculate the safety margin for swp-loop without triggering this assert:

llc: ../llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp:903: auto llvm::AIE::InterBlockScheduling::getCyclesToRespectTiming(const llvm::AIE::BlockState &, const llvm::AIE::BlockState &)::(anonymous class)::operator()(const llvm::AIE::Region &) const: Assertion `R.top_fixed_instrs().empty() && "SWP epilogue already emitted?"' failed. + /scratch/llvm-aie/build-public-mem/bin/FileCheck /scratch/llvm-aie/llvm/test/CodeGen/AIE/aie2/schedule/postpipeliner/add-store.mir

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

gbossu · 2025-01-08T10:41:46Z

llvm/lib/Target/AIE/AIERegisterEventTracker.cpp

+        }
+      }
+      return CurrMaxLatency;
+    };


Note for the future: It's not clear to me why we look at operand latencies in computeUseDefForward() but not here. We could get more accurate numbers by using Itineratries here as well for each operand we visit.

We conservatively assume that all operands are read on the first cycle, what is actually the case for most of instructions so the implementation is targeting mostly RAW dependencies that is the common case for merging epilogues. For WAR we are too conservative, but we can improve in a follow-up PR.

For WAR, we also assume a write on the first cycle by not looking to itineraries, but this is more then safe because it pushes the instruction even more far away.

gbossu

I think this looks good, I left lots of nits (as usual), but the only real problem is with memory dependencies. Once we solve that, I think it's good to go :)

F-Stuckmann · 2025-01-09T12:36:00Z

llvm/unittests/CodeGen/ScheduleDAGMITest.cpp

+  EXPECT_POS_IN_BB(Scheduler->bottom(), MBB, 4);
+
+  // Insert MI1 in top-down cycle 2. This is less than MI3's cycle so MI1
+  // should actually be inserted aboce MI3.


F-Stuckmann · 2025-01-09T12:37:58Z

llvm/unittests/CodeGen/ScheduleDAGMITest.cpp

+
+/// Schedule an instruction (MI3)  in the Top zone.
+/// Then move MI1 above because it has a lower EmissionCycle.
+TEST_F(ScheduleDAGMITest, MoveInTopAboveDbg) {


what does this unit-test test, that isn't covered by the previous one?

Indeed they are quite similar, but MI is above a dbg instruction. I would prefer to keep this test because the handling of debug instructions is a bit tricky ;-)

F-Stuckmann · 2025-01-10T14:42:31Z

llvm/lib/Target/AIE/Utils/AIELoopUtils.cpp

+MachineBasicBlock *getLoopPredecessor(const MachineBasicBlock &MBB) {
+  if (MBB.pred_size() == 1) {


maybe rename the Argument to EpilogueMBB and the function to getLoopBody?

F-Stuckmann · 2025-01-10T15:39:51Z

llvm/lib/Target/AIE/Utils/AIELoopUtils.cpp

+MachineBasicBlock *getLoopPredecessor(const MachineBasicBlock &MBB) {
+  if (MBB.pred_size() == 1) {
+    // if we have only one, it must be the loop
+    return *MBB.predecessors().begin();


Should you also check for isSingleMBBLoop(*MBB.predecessors().begin())?

I changed the code to assert in both cases.

F-Stuckmann · 2025-01-10T15:44:06Z

llvm/lib/Target/AIE/Utils/AIELoopUtils.cpp

+  // Otherwise, the loop is the fallthrough predecessor by construction
+  for (auto *Pred : MBB.predecessors()) {
+    if (Pred->isLayoutSuccessor(&MBB)) {
+      assert(isSingleMBBLoop(Pred) && "Layout predecessor is not a loop");


How would we handle loops consisting of multiple MBBs?

We don't handle this case. This function was refactored out from the InterBlockScheduler, were we handle this specific case. This function is not intended for general use. I also updated the comment in LoopUtils.h to make this a bit more clear.

F-Stuckmann · 2025-01-10T15:45:58Z

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

@@ -784,43 +788,80 @@ void InterBlockScheduling::emitBundles(
  AIEHazardRecognizer::applyBundles(Bundles, BB);
 }

-void InterBlockScheduling::emitInterBlockTop(const BlockState &BS) const {
+void InterBlockScheduling::emitTopSafetyMargin(const BlockState &BS) const {
  if (BS.Kind != BlockType::Epilogue) {
    return;
  }

  MachineBasicBlock *BB = BS.TheBlock;


NIT: rename BB to EpiloqueMBB to clarify the nature of the MachineBasicBlock

F-Stuckmann · 2025-01-10T15:58:01Z

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

+  }
+
+  MachineBasicBlock *BB = BS.TheBlock;
+  MachineBasicBlock *Loop = AIELoopUtils::getLoopPredecessor(*BB);


The epiloque check and the Loop BlockState extraction is copied in emitTopSafetyMargin, you may want to move it into a dedicated function

I refactored the code to remove this duplication.

martien-de-jong · 2025-01-10T16:11:04Z

llvm/lib/CodeGen/MachineScheduler.cpp

+                            HasLessOrEqEmissionCycle);
+  else
+    return std::lower_bound(bottom(), end(), *EmissionCycle,
+                            HasGreaterEmissionCycle);


Check: we seem to search in the same direction, but we reverse the polarity of the comparison. Aren't we finding the last such cycle in one of the cases?
-- I guess that it is compensated by the fact that we have LessEq vs Greater. LessEq vs GreaterEq would have looked more symmetric.

martien-de-jong · 2025-01-10T16:13:13Z

llvm/lib/CodeGen/MachineScheduler.cpp

-                          HasGreaterEmissionCycle);
+  // For top-down: less or equal emission cycle.
+  auto HasLessOrEqEmissionCycle = [&](const MachineInstr &MI,
+                                      unsigned EmissionCycle) {


nit: reusing/hiding EmissionCycle for something with a different type

Oh, indeed!

F-Stuckmann · 2025-01-13T09:42:42Z

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

+    return;
+  }
+
+  MachineBasicBlock *BB = BS.TheBlock;


NIT: rename BB to EpiloqueBB to clarify the nature of BB

F-Stuckmann · 2025-01-13T09:59:46Z

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

+    // Trim excedent empty bundles. Empty TopInsert means 1-stage pipeline.
+    if (!BS.TopInsert.empty()) {
+      while (BS.TopInsert.back().empty()) {
+        assert(BS.TopInsert.back().getMetaInstrs().empty());


How can BS.TopInsert.back().emptyand !BS.TopInsert.empty() be true at the same time?

In this situation we have a vector of bundles (actually, MachineBundle). In this case, if our vector is not empty, we drop reversely all empty bundles until we find a non-empty one. We do this because SWP most of times extract a sequence of empty bundles in the top insert. It can be a bit confusing but BS.TopInsert.back().empty is testing the last bundle of the vector.

In the past, without the epilog merging, those empty bundles helped to create a safety margin between the top-timed region and the epilogue instructions.

F-Stuckmann · 2025-01-13T10:02:20Z

llvm/lib/Target/AIE/AIEInterBlockScheduling.h

@@ -427,7 +427,8 @@ class InterBlockScheduling {

  /// Emit extra code induced by interblock scheduling:


NIT: introduced

martien-de-jong · 2025-01-13T13:46:08Z

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

+  // If the zone still has unscheduled fixed instructions, the next one to
+  // pick is (DAG->bottom() - 1) for bottom-up, or DAG->top() for top-down.
+  if (Zone.isTop()) {
+    if (NumEmitted < Reg.getTopFixedBundles().size()) {


nit: could we make this CurRegion, or something else that doesn't resemble a Register?

Similar to what is done for bottom-up.

Only for free instructions.

This commit prepares to schedule top fixed bundles. We also create dedicated loop exits early, handling new blocks along with their corresponding block states.

If we have TopFixed instructions, we start top-down and we change to bottom-up when we fill as much as possible the slots related to those instructions. Special care is needed for instructions with delay slot and bottom-fixed instructions.

Also add all related dependencies to have safety margins.

This is necessay because epilogue blocks of SWP loops are constrained between two already scheduled regions. When we can't se a full conflict horizon for the epilogue, we block suffucient cycles.

andcarminati requested review from abhinay-anubola, abnikant, gbossu, khallouh, konstantinschwarz, martien-de-jong, SagarMaheshwari99 and stephenneuendorffer as code owners December 10, 2024 10:57

andcarminati marked this pull request as draft December 10, 2024 10:57

andcarminati force-pushed the andreu.swp.epilogue.scheduling branch 2 times, most recently from a905dec to 4776f1d Compare December 10, 2024 12:05

andcarminati commented Dec 10, 2024

View reviewed changes

llvm/lib/CodeGen/MachineScheduler.cpp Outdated Show resolved Hide resolved

andcarminati commented Dec 10, 2024

View reviewed changes

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp Outdated Show resolved Hide resolved

andcarminati commented Dec 10, 2024

View reviewed changes

llvm/lib/Target/AIE/AIEMachineScheduler.h Outdated Show resolved Hide resolved