[AIEX] premisched: do not delay SUs if this can form a cycle

We now check whether a SU can be safely delayed without forming a cycle
of SUs delaying each other indefinitely.

Author: gbossu

llvm/lib/Target/AIE/AIEMachineScheduler.cpp

@@ -80,6 +80,11 @@ static cl::opt<bool>
     InterBlockAlignment("aie-interblock-alignment", cl::init(true),
                         cl::desc("Allow for alignment of successor blocks"));
 
+namespace {
+// A sentinel value to represent an unknown SUnit.
+const constexpr unsigned UnknownSUNum = ~0;
+} // namespace
+
 static AIEHazardRecognizer *getAIEHazardRecognizer(const SchedBoundary &Zone) {
   return static_cast<AIEHazardRecognizer *>(Zone.HazardRec);
 }
@@ -771,6 +776,7 @@ void AIEPreRASchedStrategy::enterRegion(MachineBasicBlock *BB,
   CurMBB = BB;
   RegionBegin = Begin;
   RegionEnd = End;
+  SUDelayerMap.resize(std::distance(Begin, End), UnknownSUNum);
 }
 
 void AIEPreRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
@@ -795,6 +801,7 @@ void AIEPreRASchedStrategy::leaveRegion(const SUnit &ExitSU) {
   CurMBB = nullptr;
   RegionBegin = nullptr;
   RegionEnd = nullptr;
+  SUDelayerMap.clear();
 }
 
 PressureDiff estimatePressureDiff(const SUnit &SU,
@@ -883,8 +890,34 @@ bool AIEPreRASchedStrategy::isAvailableNode(SUnit &SU, SchedBoundary &Zone,
 
   // The node will likely cause a spill, only consider it schedule-able if
   // there is no pending node that can reduce the register pressure.
-  return findPressureReducer(WorstPC.getPSet(), Zone.Pending.elements(),
-                             BotRPT) == nullptr;
+  if (const SUnit *PendingPressureReducer = findPressureReducer(
+          WorstPC.getPSet(), Zone.Pending.elements(), BotRPT);
+      PendingPressureReducer && canBeDelayed(SU, *PendingPressureReducer)) {
+    LLVM_DEBUG(dbgs() << "** Delaying SU(" << SU.NodeNum << "): Waiting for SU("
+                      << PendingPressureReducer->NodeNum << ")\n");
+
+    // Keep track of PendingPressureReducer to avoid cycles of SUs
+    // delaying each other.
+    SUDelayerMap[SU.NodeNum] = PendingPressureReducer->NodeNum;
+    return false;
+  }
+
+  // Can't prove a pending SU will help reduce reg pressure, keep as available.
+  return true;
+}
+
+bool AIEPreRASchedStrategy::canBeDelayed(const SUnit &DelayedSU,
+                                         const SUnit &Delayer) const {
+  std::function<bool(unsigned)> Impl = [&](unsigned SUNum) {
+    if (SUNum == UnknownSUNum)
+      return true;
+    if (SUNum == DelayedSU.NodeNum)
+      return false;
+    return Impl(SUDelayerMap[SUNum]);
+  };
+  // If SU is delayed by another instruction that is eventually waiting on SU
+  // itself, do not keep delaying SU otherwise this creates an infinite loop.
+  return Impl(Delayer.NodeNum);
 }
 
 bool AIEPreRASchedStrategy::tryCandidate(SchedCandidate &Cand,

llvm/lib/Target/AIE/AIEMachineScheduler.h

@@ -166,6 +166,10 @@ class AIEPreRASchedStrategy : public GenericScheduler {
                        bool VerifyReadyCycle) override;
 
 protected:
+  /// Whether \p DelayedSU can be safely delayed without forming a cycle
+  /// of SUs delaying each other indefinitely.
+  bool canBeDelayed(const SUnit &DelayedSU, const SUnit &Delayer) const;
+
   /// Apply a set of heuristics to a new candidate for scheduling.
   ///
   /// \param Cand provides the policy and current best candidate.
@@ -180,6 +184,11 @@ class AIEPreRASchedStrategy : public GenericScheduler {
   MachineBasicBlock *CurMBB = nullptr;
   MachineBasicBlock::iterator RegionBegin = nullptr;
   MachineBasicBlock::iterator RegionEnd = nullptr;
+
+  /// Keeps track of SUs that have been delayed, waiting on another
+  /// pressure-reducing SU to be scheduled first.
+  /// SUDelayerMap[0] = 2 means that SU(0) is waiting on SU(2).
+  std::vector<unsigned> SUDelayerMap;
 };
 
 /// An extension to ScheduleDAGMI that provides callbacks on region entry/exit

llvm/test/CodeGen/AIE/aie2/schedule/pre_ra/reduce_pressure_cycle.mir

@@ -0,0 +1,50 @@
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
+
+# RUN: llc -march=aie2 -run-pass=machine-scheduler %s -o -
+
+
+# This is testing a complex case where the reg pressure is critical for
+# both accumulators and vectors. In bb.1, the scheduler will try to delay
+# scheduling "%31:vec256 = VCONV_BF16_FP32 %21:acc512" because it increases the
+# pressure on accumulators and "%30:acc512 = VCONV_FP32_BF16 %20:vec256" can
+# help reducing that pressure. But the latter increases the pressure on vectors
+# so the scheduler will try to delay it because the former can help reduce that
+# pressure on vectors.
+# This can be an infinite loop if not careful.
+---
+name: tied_pressure_reducers
+tracksRegLiveness: true
+body: |
+  bb.0.entry:
+    liveins: $y2, $bml0, $wl0
+
+    %1:vec1024 = COPY $y2
+    %2:vec1024 = COPY $y2
+    %3:vec1024 = COPY $y2
+    %4:vec1024 = COPY $y2
+    %5:vec1024 = COPY $y2
+    %6:vec1024 = COPY $y2
+    %11:acc1024 = COPY $y2
+    %12:acc1024 = COPY $y2
+    %13:acc1024 = COPY $y2
+    %14:acc1024 = COPY $y2
+    %15:acc1024 = COPY $y2
+    %16:acc1024 = COPY $y2
+    %17:acc1024 = COPY $y2
+    %18:acc1024 = COPY $y2
+    %19:acc1024 = COPY $y2
+
+    %20:vec256 = COPY $wl0
+    %21:acc512 = COPY $bml0
+    PseudoJ_jump_imm %bb.1
+
+  bb.1:
+    %30:acc512 = VCONV_FP32_BF16 %20
+    %31:vec256 = VCONV_BF16_FP32 %21, implicit-def $srf2fflags, implicit $crf2fmask, implicit $crrnd
+
+    PseudoRET implicit $lr, implicit %30, implicit %31, implicit %1, implicit %2, implicit %3, implicit %4, implicit %5, implicit %6, implicit %11, implicit %12, implicit %13, implicit %14, implicit %15, implicit %16, implicit %17, implicit %18, implicit %19
+...