[AIEX] Re-assign multi-slot instructions during iterative scheduling

llvm/lib/Target/AIE/AIEAlternateDescriptors.h

@@ -31,6 +31,9 @@ class AIEAlternateDescriptors {
   AIEAlternateDescriptors() = default;
   ~AIEAlternateDescriptors() = default;
 
+  MIAltDescsMap::const_iterator begin() const { return AlternateDescs.begin(); }
+  MIAltDescsMap::const_iterator end() const { return AlternateDescs.end(); }
+
   // Construct an alternate descriptor with the given alternate descriptors.
   AIEAlternateDescriptors(const MIAltDescsMap &AltDescs)
       : AlternateDescs(AltDescs) {}
@@ -43,6 +46,10 @@ class AIEAlternateDescriptors {
     AlternateDescs[MI] = &TII->get(AltInstOpcode);
   }
 
+  void setAlternateDescriptor(MachineInstr *MI, const MCInstrDesc *AltDesc) {
+    AlternateDescs[MI] = AltDesc;
+  }
+
   // Return the alternate descriptor for the given multi-opcode instruction.
   std::optional<const MCInstrDesc *>
   getSelectedDescriptor(MachineInstr *MI) const {

llvm/lib/Target/AIE/AIEHazardRecognizer.cpp

@@ -447,7 +447,7 @@ ScheduleHazardRecognizer::HazardType AIEHazardRecognizer::getHazardType(
 bool AIEHazardRecognizer::checkConflict(
     const ResourceScoreboard<FuncUnitWrapper> &Scoreboard, MachineInstr &MI,
     int DeltaCycles) const {
-  const MCInstrDesc &Desc = MI.getDesc();
+  const MCInstrDesc &Desc = *SelectedAltDescs.getDesc(&MI);
   const unsigned SchedClass =
       TII->getSchedClass(Desc, MI.operands(), MI.getMF()->getRegInfo());
   const MemoryBankBits MemoryBanks = getMemoryBanks(&MI);

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

@@ -70,8 +70,9 @@ void emitBundlesTopDown(const std::vector<MachineBundle> &Bundles,
   // then this will not cause conflicts.
   for (int I = TotalBundles - AmountToEmit; I < TotalBundles; I++) {
     for (MachineInstr *MI : Bundles[I].getInstrs())
-      HR->emitInScoreboard(Scoreboard, MI->getDesc(), HR->getMemoryBanks(MI),
-                           MI->operands(), MI->getMF()->getRegInfo(), 0);
+      HR->emitInScoreboard(Scoreboard, *HR->getSelectedAltDescs().getDesc(MI),
+                           HR->getMemoryBanks(MI), MI->operands(),
+                           MI->getMF()->getRegInfo(), 0);
     Scoreboard.advance();
   }
 }
@@ -100,8 +101,9 @@ createBottomUpScoreboard(ArrayRef<MachineBundle> Bundles,
       Bundles.begin(), Bundles.begin() + std::min(NumBundles, RequiredCycles));
   for (const MachineBundle &B : reverse(MinBundles)) {
     for (MachineInstr *MI : B.getInstrs())
-      HR.emitInScoreboard(Scoreboard, MI->getDesc(), HR.getMemoryBanks(MI),
-                          MI->operands(), MI->getMF()->getRegInfo(), 0);
+      HR.emitInScoreboard(Scoreboard, *HR.getSelectedAltDescs().getDesc(MI),
+                          HR.getMemoryBanks(MI), MI->operands(),
+                          MI->getMF()->getRegInfo(), 0);
     Scoreboard.recede();
   }
   return Scoreboard;
@@ -124,9 +126,9 @@ checkResourceConflicts(const ResourceScoreboard<FuncUnitWrapper> &Scoreboard,
     for (MachineInstr *MI : B.getInstrs()) {
       if (BottomUpCycle >= HR.getConflictHorizon())
         break;
-      if (HR.getHazardType(Scoreboard, MI->getDesc(), HR.getMemoryBanks(MI),
-                           MI->operands(), MI->getMF()->getRegInfo(),
-                           -BottomUpCycle))
+      if (HR.getHazardType(Scoreboard, *HR.getSelectedAltDescs().getDesc(MI),
+                           HR.getMemoryBanks(MI), MI->operands(),
+                           MI->getMF()->getRegInfo(), -BottomUpCycle))
         return MI;
     }
     ++BottomUpCycle;
@@ -233,6 +235,7 @@ namespace {
 /// into the appropriate blockstate region.
 /// TimedRegion is built one bundle at the time
 class PipelineExtractor : public PipelineScheduleVisitor {
+  AIEAlternateDescriptors &AlternateDesc;
   BlockState &Loop;
   BlockState *Prologue = nullptr;
   BlockState *Epilogue = nullptr;
@@ -263,14 +266,19 @@ class PipelineExtractor : public PipelineScheduleVisitor {
     // Prologue and epilogue obtain copies.
     MachineInstr *ToBeEmitted =
         InLoop ? MI : Loop.TheBlock->getParent()->CloneMachineInstr(MI);
-    CurrentBundle.add(ToBeEmitted);
+    if (auto AltDesc = AlternateDesc.getSelectedDescriptor(MI);
+        AltDesc.has_value())
+      AlternateDesc.setAlternateDescriptor(ToBeEmitted, AltDesc.value());
+
+    CurrentBundle.add(ToBeEmitted, AlternateDesc.getOpcode(MI));
   }
   void endBundle() override { TimedRegion.emplace_back(CurrentBundle); }
 
 public:
   PipelineExtractor(InterBlockScheduling &InterBlock, BlockState &BS,
                     const AIEBaseInstrInfo &TII)
-      : Loop(BS), CurrentBundle(TII.getFormatInterface()) {
+      : AlternateDesc(InterBlock.getSelectedAltDescs()), Loop(BS),
+        CurrentBundle(TII.getFormatInterface()) {
     MachineBasicBlock *LoopBlock = Loop.TheBlock;
     for (auto *P : LoopBlock->predecessors()) {
       if (P == LoopBlock) {

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

@@ -87,6 +87,10 @@ static cl::opt<bool> UseLoopHeuristics(
     "aie-loop-sched-heuristics", cl::init(true),
     cl::desc("Use special picking heuristics when scheduling a loop region"));
 
+static cl::opt<bool> ReAssignMultiSlotInstr(
+    "aie-reassign-multislot-instr", cl::init(true),
+    cl::desc("Re-assign multi-slot instructions during iterative scheduling"));
+
 namespace {
 // A sentinel value to represent an unknown SUnit.
 const constexpr unsigned UnknownSUNum = ~0;
@@ -535,7 +539,21 @@ void AIEPostRASchedStrategy::enterMBB(MachineBasicBlock *MBB) {
   IsBottomRegion = true;
 }
 
+void AIEPostRASchedStrategy::materializeMultiOpcodeInstrs() {
+  for (auto [MI, Desc] : make_range(
+           InterBlock.getSelectedAltDescs().getAlternateDescriptors().begin(),
+           InterBlock.getSelectedAltDescs().getAlternateDescriptors().end())) {
+    MI->setDesc(*Desc);
+  }
+
+  InterBlock.getSelectedAltDescs().clear();
+}
+
 void AIEPostRASchedStrategy::commitBlockSchedule(MachineBasicBlock *BB) {
+
+  if (ReAssignMultiSlotInstr)
+    materializeMultiOpcodeInstrs();
+
   auto &BS = InterBlock.getBlockState(BB);
 
   // Safety margin, swp epilogue
@@ -598,8 +616,6 @@ void AIEPostRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
   if (BS.FixPoint.Stage != SchedulingStage::Scheduling) {
     return;
   }
-  materializeMultiOpcodeInstrs();
-  InterBlock.getSelectedAltDescs().clear();
   if (IsBottomRegion) {
     // This is the earliest point where we can destroy the recorded
     // schedule in iterative scheduling. enterMBB and enterRegion are too early,
@@ -615,34 +631,15 @@ void AIEPostRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
   assert(BS.getCurrentRegion().Bundles.empty());
   BS.addBundles(TopBundles);
   BS.addBundles(BotBundles);
+  if (!ReAssignMultiSlotInstr)
+    materializeMultiOpcodeInstrs();
   RegionBegin = nullptr;
   RegionEnd = nullptr;
   IsBottomRegion = false;
   BS.advanceRegion();
   DEBUG_BLOCKS(dbgs() << "    << leaveRegion\n");
 }
 
-void AIEPostRASchedStrategy::materializeMultiOpcodeInstrs() {
-  const TargetInstrInfo *TII = getTII(CurMBB);
-  const AIEHazardRecognizer &TopHazardRec = *getAIEHazardRecognizer(Top);
-  const AIEHazardRecognizer &BotHazardRec = *getAIEHazardRecognizer(Bot);
-
-  auto MaterializePseudo = [&TII](MachineInstr &MI,
-                                  const AIEHazardRecognizer &HazardRec) {
-    // Materialize instructions with multiple opcode options
-    if (std::optional<unsigned> AltOpcode =
-            HazardRec.getSelectedAltDescs().getSelectedOpcode(&MI)) {
-      MI.setDesc(TII->get(*AltOpcode));
-    }
-  };
-
-  assert(DAG->top() == DAG->bottom());
-  for (MachineInstr &MI : make_range(DAG->begin(), DAG->top()))
-    MaterializePseudo(MI, TopHazardRec);
-  for (MachineInstr &MI : make_range(DAG->bottom(), DAG->end()))
-    MaterializePseudo(MI, BotHazardRec);
-}
-
 bool AIEPostRASchedStrategy::checkInterZoneConflicts(
     const std::vector<AIE::MachineBundle> &BotBundles) const {
   const AIEHazardRecognizer *TopHazardRec = getAIEHazardRecognizer(Top);

llvm/lib/Target/AIE/AIEPostPipeliner.cpp

@@ -266,11 +266,12 @@ bool PostPipeliner::scheduleFirstIteration() {
       return false;
     }
     const int LocalCycle = Actual % II;
+    const MCInstrDesc &Desc = *HR.getSelectedAltDescs().getDesc(MI);
     const MemoryBankBits MemoryBanks = HR.getMemoryBanks(MI);
     LLVM_DEBUG(dbgs() << "  Emit in " << -Depth + LocalCycle << "\n");
     int Cycle = -Depth + LocalCycle;
     LLVM_DEBUG(dbgs() << "  Emit in " << Cycle << "\n");
-    HR.emitInScoreboard(Scoreboard, MI->getDesc(), MemoryBanks, MI->operands(),
+    HR.emitInScoreboard(Scoreboard, Desc, MemoryBanks, MI->operands(),
                         MI->getMF()->getRegInfo(), Cycle);
 
     scheduleNode(SU, Actual);
@@ -317,12 +318,12 @@ bool PostPipeliner::scheduleOtherIterations() {
         LLVM_DEBUG(dbgs() << "  Resource conflict\n");
         return false;
       }
+      const MCInstrDesc &Desc = *HR.getSelectedAltDescs().getDesc(MI);
       const MemoryBankBits MemoryBanks = HR.getMemoryBanks(MI);
       const int LocalCycle = (Insert - CurrentCycle) % II;
       LLVM_DEBUG(dbgs() << "  Emit in " << -Depth + LocalCycle << "\n");
-      HR.emitInScoreboard(Scoreboard, MI->getDesc(), MemoryBanks,
-                          MI->operands(), MI->getMF()->getRegInfo(),
-                          -Depth + LocalCycle);
+      HR.emitInScoreboard(Scoreboard, Desc, MemoryBanks, MI->operands(),
+                          MI->getMF()->getRegInfo(), -Depth + LocalCycle);
       scheduleNode(SU, Insert);
     }
   }