[AIE] skip pre-scheduler for loops that skip pre-pipeliner

SKipping is based on loop metadata
Xilinx · Nov 7, 2024 · ddf19aa · ddf19aa
1 parent 21fe5d8
commit ddf19aa
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 2 deletions.
diff --git a/llvm/lib/Target/AIE/AIEMachineScheduler.cpp b/llvm/lib/Target/AIE/AIEMachineScheduler.cpp
@@ -15,6 +15,7 @@
 #include "AIEInterBlockScheduling.h"
 #include "AIEMaxLatencyFinder.h"
 #include "AIEPostPipeliner.h"
+#include "Utils/AIELoopUtils.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
@@ -852,6 +853,23 @@ void AIEPreRASchedStrategy::initialize(ScheduleDAGMI *DAG) {
   }
 }
 
+MachineBasicBlock *AIEPreRASchedStrategy::nextBlock() {
+  MachineBasicBlock *Next = nullptr;
+
+  // The pipeliner is usually disabled to give the postpipeliner a chance.
+  // The prescheduler also clutters the view of the postpipeliner, so we skip
+  // such blocks here.
+  auto Skip = [](MachineBasicBlock *Block) {
+    return Block && AIELoopUtils::isSingleMBBLoop(Block) &&
+           AIELoopUtils::getPipelinerDisabled(*Block);
+  };
+
+  do {
+    Next = MachineSchedStrategy::nextBlock();
+  } while (Skip(Next));
+  return Next;
+}
+
 void AIEPreRASchedStrategy::enterRegion(MachineBasicBlock *BB,
                                         MachineBasicBlock::iterator Begin,
                                         MachineBasicBlock::iterator End,

diff --git a/llvm/lib/Target/AIE/AIEMachineScheduler.h b/llvm/lib/Target/AIE/AIEMachineScheduler.h
@@ -163,6 +163,9 @@ class AIEPreRASchedStrategy : public GenericScheduler {
 
   void initialize(ScheduleDAGMI *DAG) override;
 
+  // We override to be able to skip the prescheduler for specific blocks
+  MachineBasicBlock *nextBlock() override;
+
   void enterRegion(MachineBasicBlock *BB, MachineBasicBlock::iterator Begin,
                    MachineBasicBlock::iterator End, unsigned RegionInstrs);
   void leaveRegion(const SUnit &ExitSU);

diff --git a/llvm/lib/Target/AIE/Utils/AIELoopUtils.cpp b/llvm/lib/Target/AIE/Utils/AIELoopUtils.cpp
@@ -32,6 +32,29 @@ std::optional<int64_t> getMinTripCount(const MachineBasicBlock &LoopBlock) {
   return getMinTripCount(getLoopID(LoopBlock));
 }
 
+std::optional<bool> getPipelinerDisabled(const MachineBasicBlock &LoopBlock) {
+  auto *LoopID = getLoopID(LoopBlock);
+  if (!LoopID) {
+    return {};
+  }
+  for (const MDOperand &MDO : llvm::drop_begin(LoopID->operands())) {
+    MDNode *MD = dyn_cast<MDNode>(MDO);
+    if (MD == nullptr) {
+      continue;
+    }
+
+    MDString *S = dyn_cast<MDString>(MD->getOperand(0));
+    if (S == nullptr) {
+      continue;
+    }
+
+    if (S->getString() == "llvm.loop.pipeline.disable") {
+      return true;
+    }
+  }
+  return {};
+}
+
 MachineBasicBlock *
 getDedicatedFallThroughPreheader(const MachineBasicBlock &LoopBlock) {
   MachineBasicBlock *Candidate = nullptr;

diff --git a/llvm/lib/Target/AIE/Utils/AIELoopUtils.h b/llvm/lib/Target/AIE/Utils/AIELoopUtils.h
@@ -30,6 +30,10 @@ const MDNode *getLoopID(const MachineBasicBlock &LoopBlock);
 /// single block loops, that's automatically true.
 std::optional<int64_t> getMinTripCount(const MachineBasicBlock &LoopBlock);
 
+/// Returns true if this is a loop latch that has a pipeliner disable pragma,
+/// none otherwise.
+std::optional<bool> getPipelinerDisabled(const MachineBasicBlock &LoopBlock);
+
 /// Check that the single block loop represented by LoopBlock has a fallthrough
 /// preheader. Return the preheader if true, nullptr otherwise
 MachineBasicBlock *

diff --git a/llvm/test/CodeGen/AIE/aie2/schedule/swp/disable.mir b/llvm/test/CodeGen/AIE/aie2/schedule/swp/disable.mir
@@ -158,10 +158,10 @@ body:             |
   ; CHECK-NEXT: bb.1.for.body:
   ; CHECK-NEXT:   successors: %bb.2(0x04000000), %bb.1(0x7c000000)
   ; CHECK-NEXT: {{  $}}
-  ; CHECK-NEXT:   [[LDA_dms_lda_pstm_nrm_imm:%[0-9]+]]:er, [[COPY:%[0-9]+]]:ep_as_32bit = LDA_dms_lda_pstm_nrm_imm [[COPY]], 4 :: (load (s32) from %ir.p.addr.04)
   ; CHECK-NEXT:   [[XOR:%[0-9]+]]:er = XOR [[MOV_RLC_imm10_pseudo]], [[MOV_RLC_imm10_pseudo1]]
-  ; CHECK-NEXT:   [[MOV_RLC_imm10_pseudo:%[0-9]+]]:er = nuw nsw ADD_add_r_ri [[MOV_RLC_imm10_pseudo]], -1, implicit-def dead $srcarry
+  ; CHECK-NEXT:   [[LDA_dms_lda_pstm_nrm_imm:%[0-9]+]]:er, [[COPY:%[0-9]+]]:ep_as_32bit = LDA_dms_lda_pstm_nrm_imm [[COPY]], 4 :: (load (s32) from %ir.p.addr.04)
   ; CHECK-NEXT:   [[XOR1:%[0-9]+]]:er = XOR [[LDA_dms_lda_pstm_nrm_imm]], [[XOR]]
+  ; CHECK-NEXT:   [[MOV_RLC_imm10_pseudo:%[0-9]+]]:er = nuw nsw ADD_add_r_ri [[MOV_RLC_imm10_pseudo]], -1, implicit-def dead $srcarry
   ; CHECK-NEXT:   PseudoJNZ [[MOV_RLC_imm10_pseudo]], %bb.1
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2.for.cond.cleanup: