[AIE] Improve Zero-overhead loop 112-bytes padding

Author: Abnikant Singh

llvm/lib/Target/AIE/AIEBaseInstrInfo.cpp

@@ -376,6 +376,7 @@ void AIEBaseInstrInfo::adjustTripCount(MachineInstr &MI, int Adjustment) const {
   auto &Imm = MI.getOperand(2);
   Imm.setImm(Imm.getImm() + Adjustment);
 }
+
 bool AIEBaseInstrInfo::isHardwareLoopStart(unsigned Opcode) const {
   const auto ZOLSupport = getZOLSupport();
   return ZOLSupport && Opcode == ZOLSupport->LoopStartOpcode;
@@ -397,13 +398,57 @@ bool AIEBaseInstrInfo::isLastZOLSetupBundleInMBB(
   return true;
 }
 
-unsigned AIEBaseInstrInfo::getRegionSize(
+// Compute the total size (in bytes) of all instruction bundles in the
+// pre-header that follow the last ZOL setup instruction.
+unsigned AIEBaseInstrInfo::getPostZOLRegionSize(MachineBasicBlock &MBB) const {
+  unsigned Size = 0;
+  for (auto &MI : llvm::reverse(MBB)) {
+    if (MI.isDebugInstr())
+      continue;
+
+    if (isZOLSetupBundle(&MI) && isLastZOLSetupBundleInMBB(&MI))
+      break;
+    if (MI.isBundle()) {
+      AIE::MachineBundle Bundle = getAIEMachineBundle(MI);
+      const VLIWFormat *Format = Bundle.getFormatOrNull();
+      assert(Format);
+      Size += Format->getSize();
+    }
+  }
+  return Size;
+}
+
+// Return true if this is ZeroOverhead loop body.
+bool AIEBaseInstrInfo::isZOLBody(const MachineBasicBlock &MBB) const {
+  auto Last = MBB.getLastNonDebugInstr();
+
+  // If MBB is empty or has no non-debug instructions, return false.
+  if (Last == MBB.end())
+    return false;
+
+  return isHardwareLoopEnd(Last->getOpcode());
+}
+
+// Count the number of Machine Bundles in a MachineBasicBlock.
+unsigned
+AIEBaseInstrInfo::getZOLBundlesCount(const MachineBasicBlock &MBB) const {
+  if (!isZOLBody(MBB))
+    return 0;
+
+  auto First = MBB.getFirstNonDebugInstr();
+  auto Last = MBB.getLastNonDebugInstr();
+
+  return std::count_if(
+      First, Last, [](const MachineInstr &MI) { return !MI.isDebugInstr(); });
+}
+
+unsigned AIEBaseInstrInfo::getRegionSizeInBytes(
     llvm::iterator_range<MachineBasicBlock::iterator> Region) const {
   unsigned Size = 0;
   LLVM_DEBUG(dbgs() << "---Region Begin---\n");
-  for (auto it = Region.begin(), end = Region.end(); it != end; ++it) {
-    if (it->isBundle()) {
-      AIE::MachineBundle Bundle = getAIEMachineBundle(it);
+  for (auto It = Region.begin(), End = Region.end(); It != End; ++It) {
+    if (It->isBundle()) {
+      AIE::MachineBundle Bundle = getAIEMachineBundle(It);
       const VLIWFormat *Format = Bundle.getFormatOrNull();
       assert(Format);
       Size += Format->getSize();
@@ -1082,17 +1127,75 @@ const PacketFormats &AIEBaseInstrInfo::getPacketFormats() const {
 std::vector<MachineBasicBlock::iterator>
 AIEBaseInstrInfo::getAlignmentBoundaries(MachineBasicBlock &MBB) const {
   std::vector<MachineBasicBlock::iterator> AlgnCandidates;
-  unsigned DelaySlot = 0;
 
+  unsigned DelaySlot = 0;
   // LoopSetupDistance will be set to number of instructions (7). In
   // PostRAScheduler, this is enforced by setting the exit latency in the
-  // schduler dag mutator
+  // schduler dag mutator.
   unsigned LoopSetupDistance = 0;
+  unsigned ZOLBundlesCount = 0;
   bool IsCall = false;
   auto ZOLSupport = getZOLSupport();
+
+  const bool IsZOLBody = isZOLBody(MBB);
+  if (IsZOLBody) {
+    assert(ZOLSupport);
+    auto LoopSizeExcludingLastBundle = [&](MachineBasicBlock &MBB) -> unsigned {
+      if (MBB.empty())
+        return 0;
+
+      auto It = MBB.getLastNonDebugInstr();
+      if (It == MBB.begin())
+        return 0;
+      // Step before the PseudoLoopEnd.
+      --It;
+      while (It != MBB.begin()) {
+        if (It->isBundle())
+          return getRegionSizeInBytes(llvm::make_range(MBB.begin(), It));
+        --It;
+      }
+      return 0;
+    };
+
+    auto getPostZOLSetupRegionSize =
+        [&](MachineBasicBlock &LoopMBB) -> unsigned {
+      for (auto *Pred : LoopMBB.predecessors()) {
+        if (Pred == &LoopMBB)
+          continue;
+
+        const unsigned Size = getPostZOLRegionSize(*Pred);
+        if (Size > 0)
+          return Size;
+      }
+      return 0;
+    };
+    const unsigned ZOLSetupToLoopEndDist = ZOLSupport->LoopSetupDistance;
+    // Exclude the LoopEnd bundle as it must be placed in its own standalone
+    // region to guarantee 128-bit instruction alignment. Additionally, there
+    // must be a 112-byte gap (in PM address space) between writing to the ls,
+    // le, and lc registers and the LoopEnd instruction.
+    ZOLBundlesCount = getZOLBundlesCount(MBB) - 1;
+    if (ZOLBundlesCount < ZOLSetupToLoopEndDist)
+      LoopSetupDistance = ZOLBundlesCount;
+    else {
+      // Elongate the ZOL loop body only if the distance from the end of the
+      // ZOL setup instruction to the last bundle in the loop (excluding the
+      // final bundle) is less than 112 bytes.
+      const unsigned LoopSetupSizeInBytes = 16 * ZOLSetupToLoopEndDist;
+      const unsigned LoopSize = LoopSizeExcludingLastBundle(MBB);
+      if (LoopSize >= LoopSetupSizeInBytes)
+        LoopSetupDistance = 0;
+      else {
+        const unsigned PostZOLRegionSize = getPostZOLSetupRegionSize(MBB);
+        const bool DistanceConstraintMet =
+            (LoopSize + PostZOLRegionSize) >= LoopSetupSizeInBytes;
+        LoopSetupDistance = DistanceConstraintMet ? 0 : ZOLSetupToLoopEndDist;
+      }
+    }
+  }
   for (auto MI = MBB.begin(), End = MBB.end(); MI != End; ++MI) {
     if (MI->isBundle()) {
-      // Return Address Candidate
+      // Return Address Candidate.
       IsCall = isCallBundle(MI);
       if (IsCall && DelaySlot > 0)
         llvm_unreachable("Cannot have branch in branch delay slot!\n");
@@ -1118,8 +1221,15 @@ AIEBaseInstrInfo::getAlignmentBoundaries(MachineBasicBlock &MBB) const {
       // Distance in terms of fully-expanded 128-bit bundles that
       // loop setup should maintain. We force each of these bundles to an
       // alignment boundary, so that they will occupy 16 bytes.
-      if (ZOLSupport && isZOLSetupBundle(MI) && isLastZOLSetupBundleInMBB(MI))
-        LoopSetupDistance = ZOLSupport->LoopSetupDistance;
+      if (ZOLSupport && isZOLSetupBundle(MI) && isLastZOLSetupBundleInMBB(MI)) {
+        // if we have only one MBB, it must be the loop.
+        if (MBB.succ_size() == 1) {
+          const MachineBasicBlock *LoopSucc = *MBB.successors().begin();
+          ZOLBundlesCount = getZOLBundlesCount(*LoopSucc) - 1;
+        }
+        if (ZOLBundlesCount < ZOLSupport->LoopSetupDistance)
+          LoopSetupDistance = ZOLSupport->LoopSetupDistance - ZOLBundlesCount;
+      }
     } else if (isHardwareLoopEnd(MI->getOpcode())) {
       if (DelaySlot > 0)
         llvm_unreachable("Cannot have HWLoopEnd in branch delay slot!\n");
@@ -1128,7 +1238,7 @@ AIEBaseInstrInfo::getAlignmentBoundaries(MachineBasicBlock &MBB) const {
       AlgnCandidates.emplace_back(std::prev(MI));
     } else if (!MI->isMetaInstruction()) {
       // single instruction, there should not be any
-      // after Bundle Finalization Pass
+      // after Bundle Finalization Pass.
       llvm_unreachable("Found an un-expected standalone instruction !");
     }
   }

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -16,6 +16,7 @@
 #define LLVM_LIB_TARGET_AIE_AIEBASEINSTRRINFO_H
 
 #include "AIE.h"
+#include "AIEBundle.h"
 #include "AIEMIRFormatter.h"
 #include "AIETiedRegOperands.h"
 #include "MCTargetDesc/AIEFormat.h"
@@ -35,7 +36,6 @@ template <class I> class Bundle;
 
 struct AIEBaseInstrInfo : public TargetInstrInfo {
   using TargetInstrInfo::TargetInstrInfo;
-
   // This codifies the model of ZeroOverheadLoops
   class ZOLSupport {
   public:
@@ -317,6 +317,12 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
   // registers(lc, le, ls, etc.) and the end of the loop,
   virtual unsigned getLoopSetupDistance() const;
 
+  virtual unsigned getZOLBundlesCount(const MachineBasicBlock &MBB) const;
+
+  virtual unsigned getPostZOLRegionSize(MachineBasicBlock &MBB) const;
+
+  virtual bool isZOLBody(const MachineBasicBlock &MBB) const;
+
   // Return the vector of Alignment Region Boundaries.
   virtual std::vector<MachineBasicBlock::iterator>
   getAlignmentBoundaries(MachineBasicBlock &MBB) const;
@@ -441,8 +447,8 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
   bool isLastZOLSetupBundleInMBB(MachineBasicBlock::iterator MII) const;
   virtual const AIE::Bundle<MachineInstr>
   getAIEMachineBundle(MachineBasicBlock::iterator MII) const;
-  virtual unsigned
-  getRegionSize(llvm::iterator_range<MachineBasicBlock::iterator> Region) const;
+  virtual unsigned getRegionSizeInBytes(
+      llvm::iterator_range<MachineBasicBlock::iterator> Region) const;
 
   /// Central place to compute RAW/WAR/WAW operand latencies.
   /// This uses itineraries when they exist. It returns std::nullopt for

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

@@ -260,6 +260,8 @@ class RegionEndEdges : public ScheduleDAGMutation {
   }
   void apply(ScheduleDAGInstrs *DAG) override {
     AIE::MaxLatencyFinder MaxLatency(DAG);
+    MachineBasicBlock *PrologueMBB = DAG->getBB();
+    unsigned int ZOLBundlesCount = 0;
 
     // Default edges to ExitSU are conservative, and can't be shrunk.
     // We really should know what we're doing here, so just remove and
@@ -296,9 +298,18 @@ class RegionEndEdges : public ScheduleDAGMutation {
       if (TII->isZeroOverheadLoopSetupInstr(MI)) {
         auto ZOLSupport = TII->getZOLSupport();
         assert(ZOLSupport);
-        EdgeLatency = std::max(EdgeLatency, ZOLSupport->LoopSetupDistance + 1);
+        if (PrologueMBB && PrologueMBB->succ_size() == 1) {
+          // if we have only one MBB, it must be the loop.
+          MachineBasicBlock *LoopSucc = *PrologueMBB->successors().begin();
+          // Exclude the LoopEnd bundle since it must reside in its own
+          // standalone region to ensure it points to a 128-bit aligned
+          // instruction.
+          ZOLBundlesCount = TII->getZOLBundlesCount(*LoopSucc) - 1;
+        }
+        if (ZOLBundlesCount < ZOLSupport->LoopSetupDistance)
+          EdgeLatency = std::max(EdgeLatency, ZOLSupport->LoopSetupDistance +
+                                                  1 - ZOLBundlesCount);
       }
-
       ExitDep.setLatency(EdgeLatency);
       DAG->ExitSU.addPred(ExitDep, /*Required=*/true);
     }

llvm/lib/Target/AIE/AIEBaseSubtarget.h

@@ -18,6 +18,7 @@
 #include "AIEBaseAddrSpaceInfo.h"
 #include "AIEBaseInstrInfo.h"
 #include "Utils/AIEBaseInfo.h"
+#include "llvm/CodeGen/MachineScheduler.h"
 #include "llvm/CodeGen/ScheduleDAGMutation.h"
 #include "llvm/CodeGenTypes/MachineValueType.h"
 #include "llvm/MC/MCInstrItineraries.h"

llvm/lib/Target/AIE/AIEMachineAlignment.cpp

@@ -150,7 +150,7 @@ void AIEMachineAlignment::applyBundlesAlignment(
   for (auto Region : Regions) {
     unsigned Size = 0;
     unsigned PadBytes = 0;
-    Size = TII->getRegionSize(Region);
+    Size = TII->getRegionSizeInBytes(Region);
     if ((Size % 16) == 0)
       continue;
     PadBytes = 16 - (Size % 16);
@@ -167,8 +167,8 @@ void AIEMachineAlignment::applyBundlesAlignment(
 
 // Find Regions for Alignment Candidate e.g. Region ending with Return Address,
 // End of BB, etc.
-static std::vector<llvm::iterator_range<MachineBasicBlock::iterator>>
-findRegions(MachineBasicBlock &MBB) {
+std::vector<llvm::iterator_range<MachineBasicBlock::iterator>>
+AIEMachineAlignment::findRegions(MachineBasicBlock &MBB) {
   auto *TII = static_cast<const AIEBaseInstrInfo *>(
       MBB.getParent()->getSubtarget().getInstrInfo());
   MachineBasicBlock::iterator RegionBegin = MBB.begin();
@@ -192,11 +192,12 @@ bool AIEMachineAlignment::runOnMachineFunction(MachineFunction &MF) {
   // ordering of the blocks within the function.
   MF.RenumberBlocks();
 
+  auto *TII =
+      static_cast<const AIEBaseInstrInfo *>(MF.getSubtarget().getInstrInfo());
+
   for (auto &MBB : MF) {
     std::vector<llvm::iterator_range<MachineBasicBlock::iterator>> Regions =
         findRegions(MBB);
-    auto *TII = static_cast<const AIEBaseInstrInfo *>(
-        MBB.getParent()->getSubtarget().getInstrInfo());
     applyBundlesAlignment(Regions, TII);
     // Clean up BB local Regions
     Regions.clear();

llvm/lib/Target/AIE/AIEMachineAlignment.h

@@ -18,14 +18,17 @@
 namespace llvm {
 
 class AIEMachineAlignment : public llvm::MachineFunctionPass {
-
 public:
   static char ID;
   AIEMachineAlignment() : MachineFunctionPass(ID) {}
   llvm::StringRef getPassName() const override {
     return "AIE Machine Alignment";
   }
   bool runOnMachineFunction(llvm::MachineFunction &MF) override;
+
+  std::vector<llvm::iterator_range<MachineBasicBlock::iterator>>
+  findRegions(MachineBasicBlock &MBB);
+
   void applyBundlesAlignment(
       const std::vector<llvm::iterator_range<MachineBasicBlock::iterator>>
           &Regions,