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

llvm/lib/Target/AIE/AIEHazardRecognizer.h

@@ -161,6 +161,11 @@ class AIEHazardRecognizer : public ScheduleHazardRecognizer {
   /// the opcode selected during scheduling.
   std::optional<unsigned> getSelectedAltOpcode(MachineInstr *MI) const;
 
+  /// Function to return a pointer to the SelectedAltDescs map
+  std::map<MachineInstr *, const MCInstrDesc *> *getSelectedAltDescs() {
+    return &SelectedAltDescs;
+  }
+
   /// The instructions with memory bank attribute return the address space
   /// number
   MemoryBankBits getMemoryBanks(const MachineInstr *MI) const;

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

@@ -51,26 +51,32 @@ void dumpInterBlock(const InterBlockEdges &Edges) {
   }
 }
 
-void emitBundlesTopDown(const std::vector<MachineBundle> &Bundles,
-                        ResourceScoreboard<FuncUnitWrapper> &Scoreboard,
-                        AIEHazardRecognizer *HR) {
+void emitBundlesTopDown(
+    const std::vector<MachineBundle> &Bundles,
+    std::map<MachineInstr *, const MCInstrDesc *> SelectedAltDesc,
+    ResourceScoreboard<FuncUnitWrapper> &Scoreboard, AIEHazardRecognizer *HR) {
 
   const int TotalBundles = Bundles.size();
   const int AmountToEmit = std::min(TotalBundles, HR->getConflictHorizon());
   // Do not emit more than the specified by the conflict horizon. More
   // 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),
+    for (MachineInstr *MI : Bundles[i].getInstrs()) {
+      auto Desc = SelectedAltDesc.find(MI) != SelectedAltDesc.end()
+                      ? *SelectedAltDesc[MI]
+                      : MI->getDesc();
+      HR->emitInScoreboard(Scoreboard, Desc, HR->getMemoryBanks(MI),
                            MI->operands(), MI->getMF()->getRegInfo(), 0);
+    }
 
     Scoreboard.advance();
   }
 }
 
-ResourceScoreboard<FuncUnitWrapper>
-createBottomUpScoreboard(ArrayRef<MachineBundle> Bundles,
-                         const AIEHazardRecognizer &HR) {
+ResourceScoreboard<FuncUnitWrapper> createBottomUpScoreboard(
+    ArrayRef<MachineBundle> Bundles,
+    std::map<MachineInstr *, const MCInstrDesc *> SelectedAltDesc,
+    const AIEHazardRecognizer &HR) {
   const unsigned NumBundles = Bundles.size();
   const unsigned RequiredCycles = HR.getConflictHorizon();
 
@@ -91,9 +97,13 @@ createBottomUpScoreboard(ArrayRef<MachineBundle> Bundles,
   ArrayRef<MachineBundle> MinBundles(
       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),
+    for (MachineInstr *MI : B.getInstrs()) {
+      auto Desc = SelectedAltDesc.find(MI) != SelectedAltDesc.end()
+                      ? *SelectedAltDesc[MI]
+                      : MI->getDesc();
+      HR.emitInScoreboard(Scoreboard, Desc, HR.getMemoryBanks(MI),
                           MI->operands(), MI->getMF()->getRegInfo(), 0);
+    }
     Scoreboard.recede();
   }
   return Scoreboard;
@@ -104,10 +114,11 @@ createBottomUpScoreboard(ArrayRef<MachineBundle> Bundles,
 /// from \p PredBundles that is responsible for it.
 ///
 /// \pre The bundles contain no multi-slot pseudo.
-MachineInstr *
-checkResourceConflicts(const ResourceScoreboard<FuncUnitWrapper> &Scoreboard,
-                       const std::vector<MachineBundle> &PredBundles,
-                       const AIEHazardRecognizer &HR) {
+MachineInstr *checkResourceConflicts(
+    const ResourceScoreboard<FuncUnitWrapper> &Scoreboard,
+    const std::vector<MachineBundle> &PredBundles,
+    std::map<MachineInstr *, const MCInstrDesc *> SelectedAltDesc,
+    const AIEHazardRecognizer &HR) {
   DEBUG_LOOPAWARE(dbgs() << "Interblock Successor scoreboard:\n";
                   Scoreboard.dump());
 
@@ -116,7 +127,10 @@ checkResourceConflicts(const ResourceScoreboard<FuncUnitWrapper> &Scoreboard,
     for (MachineInstr *MI : B.getInstrs()) {
       if (BottomUpCycle >= HR.getConflictHorizon())
         break;
-      if (HR.getHazardType(Scoreboard, MI->getDesc(), HR.getMemoryBanks(MI),
+      auto Desc = SelectedAltDesc.find(MI) != SelectedAltDesc.end()
+                      ? *SelectedAltDesc[MI]
+                      : MI->getDesc();
+      if (HR.getHazardType(Scoreboard, Desc, HR.getMemoryBanks(MI),
                            MI->operands(), MI->getMF()->getRegInfo(),
                            -BottomUpCycle))
         return MI;
@@ -246,16 +260,18 @@ MachineInstr *InterBlockScheduling::resourcesConverged(BlockState &BS) const {
   // We are a single-block loop body. Check that there is no resource conflict
   // on the backedge, by overlaying top and bottom region
   if (MachineInstr *MICausingConflict = checkResourceConflicts(
-          createBottomUpScoreboard(BS.getTop().Bundles, *HR),
-          BS.getBottom().Bundles, *HR))
+          createBottomUpScoreboard(BS.getTop().Bundles, BS.SelectedAltDesc,
+                                   *HR),
+          BS.getBottom().Bundles, BS.SelectedAltDesc, *HR))
     return MICausingConflict;
 
   // Bottom represents the resources that are sticking out of the block.
   // The last non-empty cycle is a safe upperbound for the resource
   // safety margin.
   ResourceScoreboard<FuncUnitWrapper> Bottom;
   Bottom.reset(HR->getMaxLookAhead());
-  emitBundlesTopDown(BS.getBottom().Bundles, Bottom, HR.get());
+  emitBundlesTopDown(BS.getBottom().Bundles, BS.SelectedAltDesc, Bottom,
+                     HR.get());
   BS.FixPoint.MaxResourceExtent = Bottom.lastOccupied();
   return nullptr;
 }
@@ -642,8 +658,8 @@ int InterBlockScheduling::getCyclesToAvoidResourceConflicts(
                          << " Original Loop " << *LoopMBB << " Original Epilog "
                          << EpilogueMBB << "\n");
 
-  ResourceScoreboard<FuncUnitWrapper> Scoreboard =
-      createBottomUpScoreboard(EpilogueBS.getTop().Bundles, *HR);
+  ResourceScoreboard<FuncUnitWrapper> Scoreboard = createBottomUpScoreboard(
+      EpilogueBS.getTop().Bundles, EpilogueBS.SelectedAltDesc, *HR);
 
   // We know how many latency cycles we need to respect, and we can advance
   // the scoreboard to the first possible cycle that can accommodate another
@@ -656,7 +672,8 @@ int InterBlockScheduling::getCyclesToAvoidResourceConflicts(
 
   // Increment the number of intermediate nops until there are no resource
   // conflicts between the last iteration of the loop and the epilogue.
-  while (checkResourceConflicts(Scoreboard, LoopBS.getBottom().Bundles, *HR)) {
+  while (checkResourceConflicts(Scoreboard, LoopBS.getBottom().Bundles,
+                                LoopBS.SelectedAltDesc, *HR)) {
     Scoreboard.recede();
     ++NopCounter;
   }
@@ -714,6 +731,7 @@ void BlockState::clearSchedule() {
   for (auto &R : Regions) {
     R.Bundles.clear();
   }
+  SelectedAltDesc.clear();
   CurrentRegion = 0;
 }
 

llvm/lib/Target/AIE/AIEInterBlockScheduling.h

@@ -39,7 +39,7 @@ class DataDependenceHelper : public ScheduleDAGInstrs {
   /// Ordered list of DAG postprocessing steps.
   std::vector<std::unique_ptr<ScheduleDAGMutation>> Mutations;
   const MachineSchedContext &Context;
-  void schedule() override{};
+  void schedule() override {};
 
 public:
   DataDependenceHelper(const MachineSchedContext &Context)
@@ -171,12 +171,19 @@ class BlockState {
   BlockType Kind = BlockType::Regular;
   LivePhysRegs LiveOuts;
   void initInterBlock(const MachineSchedContext &Context);
+  // Multi-Opcode Instr and their selected Desc
+  std::map<MachineInstr *, const MCInstrDesc *> SelectedAltDesc;
 
   // Concatenate Bundles to the current region
   void addBundles(const std::vector<MachineBundle> &Bundles) {
     auto &TheBundles = Regions.at(CurrentRegion).Bundles;
     TheBundles.insert(TheBundles.end(), Bundles.begin(), Bundles.end());
   }
+  void addSelectedAltDesc(
+      std::map<MachineInstr *, const MCInstrDesc *> &SelecteOpcodeMap) {
+    for (const auto &[MI, SelectedDesc] : SelecteOpcodeMap)
+      SelectedAltDesc[MI] = SelectedDesc;
+  }
   void addRegion(MachineBasicBlock *BB, MachineBasicBlock::iterator RegionBegin,
                  MachineBasicBlock::iterator RegionEnd) {
     CurrentRegion = Regions.size();

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

@@ -543,6 +543,16 @@ void AIEPostRASchedStrategy::enterMBB(MachineBasicBlock *MBB) {
   IsBottomRegion = true;
 }
 
+void AIEPostRASchedStrategy::materializeMultiOpcodeInstrs(
+    MachineBasicBlock *BB) {
+  auto &BS = InterBlock.getBlockState(BB);
+  for (auto &Region : BS.getRegions())
+    for (const AIE::MachineBundle &Bundle : Region.Bundles)
+      for (MachineInstr *MI : Bundle.getInstrs())
+        if (BS.SelectedAltDesc.find(MI) != BS.SelectedAltDesc.end())
+          MI->setDesc(*BS.SelectedAltDesc.at(MI));
+}
+
 void AIEPostRASchedStrategy::commitBlockSchedule(MachineBasicBlock *BB) {
   auto &BS = InterBlock.getBlockState(BB);
   MachineBasicBlock::iterator It = BB->begin();
@@ -575,6 +585,7 @@ void AIEPostRASchedStrategy::commitBlockSchedule(MachineBasicBlock *BB) {
 void AIEPostRASchedStrategy::leaveMBB() {
   if (InterBlock.leaveBlock()) {
     // Finish it off and move to the next block.
+    materializeMultiOpcodeInstrs(CurMBB);
     commitBlockSchedule(CurMBB);
   }
   CurMBB = nullptr;
@@ -595,7 +606,6 @@ void AIEPostRASchedStrategy::enterRegion(MachineBasicBlock *BB,
 
 void AIEPostRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
   LLVM_DEBUG(dbgs() << "    << leaveRegion\n");
-  materializeMultiOpcodeInstrs();
 
   auto &BS = InterBlock.getBlockState(CurMBB);
   if (IsBottomRegion) {
@@ -613,34 +623,15 @@ void AIEPostRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
   assert(BS.getCurrentRegion().Bundles.empty());
   BS.addBundles(TopBundles);
   BS.addBundles(BotBundles);
+  BS.addSelectedAltDesc(*getAIEHazardRecognizer(Top)->getSelectedAltDescs());
+  BS.addSelectedAltDesc(*getAIEHazardRecognizer(Bot)->getSelectedAltDescs());
   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<const MCInstrDesc *> AltOpcode =
-            HazardRec.getSelectedAltDesc(&MI)) {
-      MI.setDesc(*AltOpcode.value());
-    }
-  };
-
-  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

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