[AIEX] Improve pre-sched decisions for Post-SWP candidates

Author: andcarminati

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

@@ -59,6 +59,11 @@ static cl::opt<bool> ForcePostPipeliner(
     cl::desc(
         "Force using AIE's post-pipeliner instead of the MachinePipeliner"),
     cl::init(false), cl::Hidden);
+
+static cl::opt<bool> PreSchedPostSWPCandidates(
+    "aie-presched-postpipeliner-candidates",
+    cl::desc("Run pre-scheduler over potential postpipeliner candidates"),
+    cl::init(true), cl::Hidden);
 // These are debugging/testing options.
 
 // aie-latency-margin defines the latency that will be given to ExitSU edges.
@@ -864,3 +869,8 @@ AIEBaseSubtarget::getSMSMutationsImpl(const Triple &TT) {
 bool AIEBaseSubtarget::enableMachinePipeliner() const {
   return !ForcePostPipeliner;
 }
+
+/// Whether to enable the pre-RA MachineScheduler for Post SWP candidates.
+bool AIEBaseSubtarget::shouldPreSchedPostSWPCandidates() const {
+  return PreSchedPostSWPCandidates;
+}

llvm/lib/Target/AIE/AIEBaseSubtarget.h

@@ -89,6 +89,9 @@ class AIEBaseSubtarget {
   /// Whether to enable the pre-RA MachinePipeliner. This can be disabled to let
   /// the post-RA pipeliner handle the scheduling.
   bool enableMachinePipeliner() const;
+
+  /// Whether to enable the pre-RA MachineScheduler for Post SWP candidates.
+  virtual bool shouldPreSchedPostSWPCandidates() const;
 };
 } // namespace llvm
 

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

@@ -91,10 +91,6 @@ 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> PreSchedFollowsSkipPipeliner(
-    "aie-presched-follows-skip-pipeliner", cl::init(true),
-    cl::desc("Don't run the prescheduler if the pipeliner is skipped"));
-
 namespace {
 // A sentinel value to represent an unknown SUnit.
 const constexpr unsigned UnknownSUNum = ~0;
@@ -1082,11 +1078,19 @@ MachineBasicBlock *AIEPreRASchedStrategy::nextBlock() {
   auto Skip = [](MachineBasicBlock *Block) {
     if (!Block)
       return false;
-    bool PrePipelinerDisabled =
-        AIELoopUtils::getPipelinerDisabled(*Block) ||
-        !Block->getParent()->getSubtarget().enableMachinePipeliner();
-    return PreSchedFollowsSkipPipeliner &&
-           AIELoopUtils::isSingleMBBLoop(Block) && PrePipelinerDisabled;
+
+    auto *TII = static_cast<const AIEBaseInstrInfo *>(
+        Block->getParent()->getSubtarget().getInstrInfo());
+    const AIEBaseSubtarget &STI = AIEBaseSubtarget::get(*Block->getParent());
+
+    if (AIELoopUtils::isPostSWPCandidate(*TII, Block)) {
+      bool ShouldSkip = !STI.shouldPreSchedPostSWPCandidates();
+      LLVM_DEBUG(dbgs() << "Skip pre-sched of post-SWP candidate "
+                        << Block->getName() << ":" << ShouldSkip << "\n");
+      return ShouldSkip;
+    }
+
+    return false;
   };
 
   do {

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

@@ -190,4 +190,51 @@ bool hasUnrollPragma(const Loop *L) {
   return false;
 }
 
+/// Check if this block is a post-SWP candidate.
+bool isPostSWPCandidate(const AIEBaseInstrInfo &TII,
+                        const MachineBasicBlock *MBB) {
+
+  if (!isSingleMBBLoop(MBB))
+    return false;
+
+  const MachineInstr &Terminator = *MBB->getFirstTerminator();
+  if (!TII.isHardwareLoopEnd(Terminator.getOpcode()))
+    return false;
+
+  if (Terminator.getOperand(1).getMBB() != MBB)
+    return false;
+
+  auto GetLoopStartBlock =
+      [&](const MachineBasicBlock *LoopBlock) -> const MachineBasicBlock * {
+    const MachineBasicBlock *LoopStartBlock = nullptr;
+    for (auto *Pred : LoopBlock->predecessors()) {
+      if (Pred == LoopBlock)
+        continue;
+      if (LoopStartBlock)
+        return nullptr;
+      LoopStartBlock = Pred;
+    }
+    return LoopStartBlock;
+  };
+
+  auto LoopStartBlock = GetLoopStartBlock(MBB);
+  if (!LoopStartBlock)
+    return false;
+
+  auto FindLoopStart =
+      [&](const MachineBasicBlock &Block) -> const MachineInstr * {
+    for (auto &MI : reverse(Block)) {
+      if (TII.isHardwareLoopStart(MI.getOpcode()))
+        return &MI;
+    }
+    return nullptr;
+  };
+
+  auto Init = FindLoopStart(*LoopStartBlock);
+  if (!Init)
+    return false;
+
+  return Init->getOperand(1).getImm() == 0;
+}
+
 } // namespace llvm::AIELoopUtils

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

@@ -14,6 +14,7 @@
 #ifndef LLVM_LIB_TARGET_AIE_UTILS_AIELOOPUTILS_H
 #define LLVM_LIB_TARGET_AIE_UTILS_AIELOOPUTILS_H
 
+#include "AIEBaseInstrInfo.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 
@@ -70,6 +71,10 @@ bool hasUnrollCountPragma(const MDNode *LoopID);
 
 bool hasUnrollPragma(const Loop *L);
 
+/// Check if this block is a post-SWP candidate.
+bool isPostSWPCandidate(const AIEBaseInstrInfo &TII,
+                        const MachineBasicBlock *MBB);
+
 } // namespace llvm::AIELoopUtils
 
 #endif

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:   [[XOR:%[0-9]+]]:er = XOR [[MOV_RLC_imm10_pseudo]], [[MOV_RLC_imm10_pseudo1]]
   ; 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:   [[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:   [[XOR1:%[0-9]+]]:er = XOR [[LDA_dms_lda_pstm_nrm_imm]], [[XOR]]
   ; CHECK-NEXT:   PseudoJNZ [[MOV_RLC_imm10_pseudo]], %bb.1
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2.for.cond.cleanup: