Merge pull request #50 from Xilinx/sanubola.waw.mutator.liveness

Supply interblock liveness information to WAW rewriter

Author: abhinay-anubola

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

@@ -425,6 +425,8 @@ void dumpDependencies(ScheduleDAGInstrs *DAG, SDep::Kind depType,
 /// live set of MBB, backtrack the DAG and update the live set. Whenever an edge
 /// points to a non-live write, it is updated to the subsequent live write.
 class WAWEdges : public ScheduleDAGMutation {
+
+  AIEPostRASchedStrategy *Scheduler = nullptr;
   // Collect all edges in a separate vector. This allows modifying SU.Preds
   // without invalidating iterators.
   SmallVector<SDep, 4> getPreds(SUnit &SU) {
@@ -446,17 +448,41 @@ class WAWEdges : public ScheduleDAGMutation {
       }
     }
   }
+
+public:
+  void setScheduler(AIEPostRASchedStrategy *Scheduler) {
+    this->Scheduler = Scheduler;
+  }
+
   void apply(ScheduleDAGInstrs *DAG) override {
     MachineFunction &MF = DAG->MF;
     MachineRegisterInfo &MRI = MF.getRegInfo();
     const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
     auto *RI = static_cast<const AIEBaseRegisterInfo *>(TRI);
     LivePhysRegs LiveRegs;
     LiveRegs.init(*TRI);
-    // Reserved registers are considered always live
-    for (MCPhysReg PhysReg : MRI.getReservedRegs().set_bits()) {
-      if (RI->isSimplifiableReservedReg(PhysReg))
-        LiveRegs.addReg(PhysReg);
+    bool AddReservedRegs = true;
+    if (Scheduler) {
+      MachineBasicBlock *MBB = DAG->getBB();
+      const BlockState &BS = Scheduler->getInterBlock().getBlockState(MBB);
+      auto Region = BS.getCurrentRegion();
+      auto BottomRegion = BS.getBottom();
+      if (*Region.begin() == *BottomRegion.begin()) {
+        // If the region is bottom region, liveouts of region are same as
+        // liveouts of the MBB
+        for (const MCPhysReg Reg : BS.LiveOuts) {
+          LiveRegs.addReg(Reg);
+        }
+        AddReservedRegs = false;
+      }
+    }
+
+    if (AddReservedRegs) {
+      // Reserved registers are considered always live
+      for (const MCPhysReg PhysReg : MRI.getReservedRegs().set_bits()) {
+        if (RI->isSimplifiableReservedReg(PhysReg))
+          LiveRegs.addReg(PhysReg);
+      }
     }
     // Stores latest live write of physical register.
     std::map<Register, SUnit *> PhysRegWriters;
@@ -482,6 +508,26 @@ class WAWEdges : public ScheduleDAGMutation {
   };
 };
 
+// Adds WAW edges for scheduling in the context of the Scheduler.
+// This class extends WAWEdges to apply WAW edges using a Scheduler if available
+// It overrides the apply method to retrieve the Scheduler from the DAG if a
+// BasicBlock is present, otherwise, it uses nullptr.
+class MachineSchedWAWEdges : public WAWEdges {
+  void apply(ScheduleDAGInstrs *DAG) override {
+    AIEPostRASchedStrategy *Scheduler =
+        DAG->getBB() ? static_cast<AIEScheduleDAGMI *>(DAG)->getSchedImpl()
+                     : nullptr;
+    setScheduler(Scheduler);
+    WAWEdges::apply(DAG);
+  }
+};
+
+// This class extends WAWEdges to apply WAW edges without using a Scheduler.
+// This is useful for scenarios where the SWP (Software Pipelining) is performed
+// independently of the Scheduler.
+class SWPWAWEdges : public WAWEdges {
+  void apply(ScheduleDAGInstrs *DAG) override { WAWEdges::apply(DAG); }
+};
 } // namespace
 
 std::vector<std::unique_ptr<ScheduleDAGMutation>>
@@ -491,7 +537,7 @@ AIEBaseSubtarget::getPostRAMutationsImpl(const Triple &TT) {
   if (!TT.isAIE1()) {
     Mutations.emplace_back(std::make_unique<RegionEndEdges>());
     Mutations.emplace_back(std::make_unique<MemoryEdges>());
-    Mutations.emplace_back(std::make_unique<WAWEdges>());
+    Mutations.emplace_back(std::make_unique<MachineSchedWAWEdges>());
   }
   return Mutations;
 }
@@ -504,7 +550,7 @@ AIEBaseSubtarget::getInterBlockMutationsImpl(const Triple &TT) {
   if (!TT.isAIE1()) {
     Mutations.emplace_back(std::make_unique<RegionEndEdges>());
     Mutations.emplace_back(std::make_unique<MemoryEdges>());
-    Mutations.emplace_back(std::make_unique<WAWEdges>());
+    Mutations.emplace_back(std::make_unique<MachineSchedWAWEdges>());
   }
   return Mutations;
 }
@@ -523,7 +569,7 @@ std::vector<std::unique_ptr<ScheduleDAGMutation>>
 AIEBaseSubtarget::getSMSMutationsImpl(const Triple &TT) {
   std::vector<std::unique_ptr<ScheduleDAGMutation>> Mutations;
   if (!TT.isAIE1()) {
-    Mutations.emplace_back(std::make_unique<WAWEdges>());
+    Mutations.emplace_back(std::make_unique<SWPWAWEdges>());
     if (EnablePipelinerSchedPropagateIncomingLatencies)
       Mutations.emplace_back(std::make_unique<PropagateIncomingLatencies>());
   }

llvm/lib/Target/AIE/AIEInterBlockScheduling.cpp

@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "AIEInterBlockScheduling.h"
+#include "AIELiveRegs.h"
 #include "AIEMaxLatencyFinder.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
@@ -111,9 +112,28 @@ void InterBlockScheduling::enterFunction(MachineFunction *MF) {
   // Get ourselves a hazard recognizer
   HR = std::make_unique<AIEHazardRecognizer>(MF->getSubtarget());
 
+  const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
+  LiveRegs MBBLiveness(MF);
+  const std::map<const MachineBasicBlock *, LivePhysRegs> &LiveIns =
+      MBBLiveness.getLiveIns();
+
   // Define our universe of blocks
   for (MachineBasicBlock &MBB : *MF) {
-    Blocks.emplace(&MBB, &MBB);
+    auto Itr = Blocks.emplace(&MBB, &MBB).first;
+    BlockState &BS = Itr->second;
+    BS.LiveOuts.init(*TRI);
+    // Calculating LiveOuts by iterating over each successor of the MBB and
+    // adding each successor's LiveIns to LiveOuts.
+    for (const MachineBasicBlock *Succ : MBB.successors()) {
+      const LivePhysRegs &MBBLiveins = LiveIns.at(Succ);
+      for (const MCPhysReg Reg : MBBLiveins) {
+        BS.LiveOuts.addReg(Reg);
+      }
+    }
+    LLVM_DEBUG({
+      dbgs() << MBB.getFullName() << " LiveOuts\n";
+      BS.LiveOuts.dump();
+    });
   }
   if (LoopAware) {
     // Mark epilogues of the loops we found. This is only necessary if

llvm/lib/Target/AIE/AIEInterBlockScheduling.h

@@ -165,6 +165,7 @@ class BlockState {
   MachineBasicBlock *TheBlock = nullptr;
   FixedpointState FixPoint;
   BlockType Kind = BlockType::Regular;
+  LivePhysRegs LiveOuts;
   void initInterBlock(const MachineSchedContext &Context);
 
   // Concatenate Bundles to the current region

llvm/lib/Target/AIE/AIELiveRegs.cpp

@@ -0,0 +1,112 @@
+//===- AIELiveRegs.cpp - Liveness Analysis infrastructure -----------------===//
+//
+// 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 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+//
+//===----------------------------------------------------------------------===//
+// Implementations of the classes used to support Liveness Analysis of all
+// physical registers, including reserved registers.
+//===----------------------------------------------------------------------===//
+
+#include "AIELiveRegs.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+
+using namespace llvm;
+
+namespace llvm::AIE {
+
+void LiveRegs::addToWorkList(const MachineBasicBlock *MBB) {
+  auto It = InWorkList.find(MBB);
+  if (It == InWorkList.end() || !It->second) {
+    WorkList.push(MBB);
+    It->second = true;
+  }
+}
+
+const MachineBasicBlock *LiveRegs::popFromWorkList() {
+  const MachineBasicBlock *MBB = WorkList.front();
+  WorkList.pop();
+  InWorkList[MBB] = false;
+  return MBB;
+}
+
+void LiveRegs::computeBlockLiveIns(const MachineBasicBlock *MBB,
+                                   LivePhysRegs &CurrentLive) {
+  CurrentLive.init(*TRI);
+
+  // Calculates CurrentLive by iterating over each successor of the current
+  // MBB and adding each successor's LiveIns to CurrentLive.
+  // This ensures CurrentLive is the union of LiveIns of all successor MBBs.
+  for (const MachineBasicBlock *Succ : MBB->successors()) {
+    for (MCPhysReg Reg : LiveIns[Succ]) {
+      CurrentLive.addReg(Reg);
+    }
+  }
+  // Calculates CurrentLive by bottom-up traversal of the MBB
+  for (const MachineInstr &MI : llvm::reverse(*MBB)) {
+    CurrentLive.stepBackward(MI);
+  }
+}
+
+bool LiveRegs::equal(LivePhysRegs &CurrentLive, LivePhysRegs &OldLive) {
+  int LiveInCount = 0;
+  bool Equal = true;
+
+  for (const MCPhysReg Reg : OldLive) {
+    if (!CurrentLive.contains(Reg)) {
+      Equal = false;
+    }
+    LiveInCount++;
+  }
+
+  for (const MCPhysReg Reg : CurrentLive) {
+    LiveInCount--;
+  }
+  return LiveInCount == 0 && Equal;
+}
+
+void LiveRegs::updateLiveRegs(LivePhysRegs &CurrentLive,
+                              LivePhysRegs &OldLive) {
+  OldLive.init(*TRI);
+  for (const MCPhysReg Reg : CurrentLive) {
+    OldLive.addReg(Reg);
+  }
+}
+
+const std::map<const MachineBasicBlock *, LivePhysRegs> &
+LiveRegs::getLiveIns() const {
+  return LiveIns;
+}
+
+LiveRegs::LiveRegs(const MachineFunction *MF) {
+  const TargetRegisterInfo *RegisterInfo = MF->getSubtarget().getRegisterInfo();
+  TRI = RegisterInfo;
+
+  // Using post_order optimizes by minimizing re-runs, though it's not required
+  // for correctness. Post-order ensures that we process the basic blocks in a
+  // way that naturally accommodates dependencies and minimizes redundant work.
+  for (const MachineBasicBlock *MBB : post_order(MF)) {
+    LiveIns[MBB].init(*TRI);
+    WorkList.push(MBB);
+    InWorkList[MBB] = true;
+  }
+  while (!WorkList.empty()) {
+    const MachineBasicBlock *MBB = popFromWorkList();
+    LivePhysRegs CurrentLive;
+    computeBlockLiveIns(MBB, CurrentLive);
+
+    // If MBB's liveins changed, force a recomputation for MBB's preds
+    if (equal(CurrentLive, LiveIns[MBB]))
+      continue;
+
+    updateLiveRegs(CurrentLive, LiveIns[MBB]);
+    for (const MachineBasicBlock *Pred : MBB->predecessors())
+      addToWorkList(Pred);
+  }
+}
+
+} // namespace llvm::AIE

llvm/lib/Target/AIE/AIELiveRegs.h

@@ -0,0 +1,66 @@
+//===- AIELiveRegs.h - Liveness Analysis logic -*- C++ -*------------------===//
+//
+// 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 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+//
+//===----------------------------------------------------------------------===//
+//
+// Class providing services for Liveness Analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_AIE_AIELIVEREGS_H
+#define LLVM_LIB_TARGET_AIE_AIELIVEREGS_H
+
+#include "llvm/CodeGen/LivePhysRegs.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
+#include <map>
+#include <queue>
+#include <set>
+
+namespace llvm::AIE {
+
+class LiveRegs {
+  // Mapping from Machine Basic Blocks to their livein registers.
+  std::map<const MachineBasicBlock *, LivePhysRegs> LiveIns;
+
+  // Queue to manage the order of MBBs to be processed.
+  std::queue<const MachineBasicBlock *> WorkList;
+
+  // Map to track whether an MBB is in the WorkList.
+  std::map<const MachineBasicBlock *, bool> InWorkList;
+
+  // Pointer to Target Register Information, used for initializing livein
+  // registers.
+  const TargetRegisterInfo *TRI;
+
+  // Adds a Machine Basic Block to the work list if it is not already present.
+  void addToWorkList(const MachineBasicBlock *MBB);
+
+  // Pops a Machine Basic Block from the work list, removes it from the work
+  // list set, and returns the popped MBB.
+  const MachineBasicBlock *popFromWorkList();
+
+  // Computes the live-in registers for a given Machine Basic Block.
+  void computeBlockLiveIns(const MachineBasicBlock *MBB,
+                           LivePhysRegs &CurrentLive);
+
+  // Checks if the CurrentLive is the same as the liveins of the given MBB.
+  static bool equal(LivePhysRegs &CurrentLive, LivePhysRegs &OldLive);
+
+  // Updates OldLive registers with CurrentLive registers.
+  void updateLiveRegs(LivePhysRegs &CurrentLive, LivePhysRegs &OldLive);
+
+public:
+  const std::map<const MachineBasicBlock *, LivePhysRegs> &getLiveIns() const;
+  LiveRegs(const MachineFunction *MF);
+};
+
+} // end namespace llvm::AIE
+
+#endif // LLVM_LIB_TARGET_AIE_AIELIVEREGS_H

llvm/lib/Target/AIE/CMakeLists.txt

@@ -76,6 +76,7 @@ add_llvm_target(AIECodeGen
    AIEInterBlockScheduling.cpp
    AIEISelDAGToDAG.cpp
    AIELegalizerInfo.cpp
+   AIELiveRegs.cpp
    AIEMachineAlignment.cpp
    AIEMachineBlockPlacement.cpp
    AIEMachineFunctionInfo.cpp

llvm/test/CodeGen/AIE/aie2/hardware-loops/nested.ll

@@ -39,32 +39,26 @@ define void @nested(ptr nocapture %out, ptr nocapture readonly %in, i32 noundef
 ; CHECK-NEXT:  .LBB0_2: // %for.body6
 ; CHECK-NEXT:    // Parent Loop BB0_1 Depth=1
 ; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
-; CHECK-NEXT:    nopa ; lshl r7, r6, r4
+; CHECK-NEXT:    nopb ; nopa ; nops ; lshl r7, r6, r4; nopm ; nopv
 ; CHECK-NEXT:    mov dj0, r7
 ; CHECK-NEXT:    lda r7, [p4, dj0]
 ; CHECK-NEXT:    nop
 ; CHECK-NEXT:    nop
-; CHECK-NEXT:    nop
-; CHECK-NEXT:    nop
-; CHECK-NEXT:    nop
-; CHECK-NEXT:    add r6, r6, #1
-; CHECK-NEXT:    add r2, r2, r7
 ; CHECK-NEXT:    jnzd r5, r5, p2
 ; CHECK-NEXT:    nop // Delay Slot 5
 ; CHECK-NEXT:    nop // Delay Slot 4
-; CHECK-NEXT:    nop // Delay Slot 3
-; CHECK-NEXT:    nop // Delay Slot 2
+; CHECK-NEXT:    add r6, r6, #1 // Delay Slot 3
+; CHECK-NEXT:    add r2, r2, r7 // Delay Slot 2
 ; CHECK-NEXT:    st r2, [p0, #0] // Delay Slot 1
 ; CHECK-NEXT:    .p2align 4
 ; CHECK-NEXT:  // %bb.3: // %for.cond3.for.cond.cleanup5_crit_edge
 ; CHECK-NEXT:    // in Loop: Header=BB0_1 Depth=1
-; CHECK-NEXT:    nopb ; nopa ; nops ; add r3, r3, #1; nopm ; nopv
-; CHECK-NEXT:    nopa ; jnzd r0, r0, p3
+; CHECK-NEXT:    jnzd r0, r0, p3
 ; CHECK-NEXT:    nop // Delay Slot 5
 ; CHECK-NEXT:    nop // Delay Slot 4
 ; CHECK-NEXT:    nop // Delay Slot 3
 ; CHECK-NEXT:    nop // Delay Slot 2
-; CHECK-NEXT:    nop // Delay Slot 1
+; CHECK-NEXT:    add r3, r3, #1 // Delay Slot 1
 ; CHECK-NEXT:    .p2align 4
 ; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
 ; CHECK-NEXT:    nopa ; ret lr