[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

@@ -451,7 +451,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

@@ -300,6 +300,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;
@@ -330,14 +331,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

@@ -88,6 +88,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;
@@ -269,7 +273,8 @@ void AIEPostRASchedStrategy::initializeBotScoreBoard(ScoreboardTrust Trust) {
   /// by starting in the earliest possible cycle, -Depth
   auto InsertInCycle = [=](MachineInstr &MI, int Cycle) {
     BotHazardRec->emitInScoreboard(
-        MI.getDesc(), BotHazardRec->getMemoryBanks(&MI), MI.operands(),
+        *BotHazardRec->getSelectedAltDescs().getDesc(&MI),
+        BotHazardRec->getMemoryBanks(&MI), MI.operands(),
         MI.getMF()->getRegInfo(), Cycle - Depth);
   };
   auto BlockCycle = [=](int Cycle) {
@@ -536,7 +541,20 @@ void AIEPostRASchedStrategy::enterMBB(MachineBasicBlock *MBB) {
   IsBottomRegion = true;
 }
 
+void AIEPostRASchedStrategy::materializeMultiSlotInstrs() {
+  for (auto &[MI, Desc] : make_range(InterBlock.getSelectedAltDescs().begin(),
+                                     InterBlock.getSelectedAltDescs().end())) {
+    MI->setDesc(*Desc);
+  }
+
+  InterBlock.getSelectedAltDescs().clear();
+}
+
 void AIEPostRASchedStrategy::commitBlockSchedule(MachineBasicBlock *BB) {
+
+  if (ReAssignMultiSlotInstr)
+    materializeMultiSlotInstrs();
+
   auto &BS = InterBlock.getBlockState(BB);
 
   // Safety margin, swp epilogue
@@ -599,8 +617,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,
@@ -616,34 +632,15 @@ void AIEPostRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
   assert(BS.getCurrentRegion().Bundles.empty());
   BS.addBundles(TopBundles);
   BS.addBundles(BotBundles);
+  if (!ReAssignMultiSlotInstr)
+    materializeMultiSlotInstrs();
   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/AIEMachineScheduler.h

@@ -131,7 +131,7 @@ class AIEPostRASchedStrategy : public PostGenericScheduler {
 
   /// Materialize "multi-opcode" instructions into the option that was selected
   /// at schedule time. See AIEHazardRecognizer::getSelectedAltOpcode().
-  void materializeMultiOpcodeInstrs();
+  void materializeMultiSlotInstrs();
 
   /// Returns true if, when "concatenated", the Top and Bot zone have resource
   /// conflicts or timing issues.

llvm/lib/Target/AIE/AIEPostPipeliner.cpp

@@ -112,7 +112,8 @@ int PostPipeliner::getResMII(MachineBasicBlock &LoopBlock) {
   std::vector<uint64_t> Scoreboard(NInstr, 0);
   int MII = 1;
   for (auto &MI : LoopBlock) {
-    auto *SlotInfo = TII->getSlotInfo(TII->getSlotKind(MI.getOpcode()));
+    const unsigned Opcode = HR.getSelectedAltDescs().getOpcode(&MI);
+    auto *SlotInfo = TII->getSlotInfo(TII->getSlotKind(Opcode));
     SlotBits Slots = SlotInfo ? SlotInfo->getSlotSet() : 0;
 
     int C = 0;
@@ -290,6 +291,7 @@ 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;
@@ -299,8 +301,8 @@ bool PostPipeliner::scheduleFirstIteration() {
         return false;
       }
 
-      HR.emitInScoreboard(Scoreboard, MI->getDesc(), MemoryBanks,
-                          MI->operands(), MI->getMF()->getRegInfo(), Cycle);
+      HR.emitInScoreboard(Scoreboard, Desc, MemoryBanks, MI->operands(),
+                          MI->getMF()->getRegInfo(), Cycle);
       Cycle += II;
     }