From 37ecc39dbfabf8516f9f4f15d37de4460634504a Mon Sep 17 00:00:00 2001
From: Niwin Anto <niwin.anto@amd.com>
Date: Tue, 11 Mar 2025 12:58:59 +0000
Subject: [PATCH] [AIEX] Create dependency between lock instr and fixed instr
 for SWP prologue merging

---
 llvm/lib/Target/AIE/AIEBaseSubtarget.cpp        | 17 +++++++++++++++--
 llvm/lib/Target/AIE/AIEMaxLatencyFinder.cpp     | 11 +++++++++--
 .../postpipeliner/interleave-prologue.mir       |  7 +++----
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/AIE/AIEBaseSubtarget.cpp b/llvm/lib/Target/AIE/AIEBaseSubtarget.cpp
index c78300a0a2f6..5c1d6b691198 100644
--- a/llvm/lib/Target/AIE/AIEBaseSubtarget.cpp
+++ b/llvm/lib/Target/AIE/AIEBaseSubtarget.cpp
@@ -396,8 +396,21 @@ class EmitFixedSUnits : public ScheduleDAGMutation {
           DAG->getSUnit(&*getBundleStart(FixedDepMI->getIterator()));
       assert(FixedDepSU && "Fixed Bundle has no corresponding SU.");
       SDep Dep(&FreeSU, SDep::Artificial);
-      Dep.setLatency(
-          AIE::maxLatency(&MI, *TII, *ItinData, /*IncludeStages=*/true));
+      auto Latency =
+          AIE::maxLatency(&MI, *TII, *ItinData, /*IncludeStages=*/true);
+      if (TII->isLock(MI.getOpcode())) {
+        Dep.setLatency(std::max(
+            TII->getCoreResumeCycleAfterLock() -
+                *TII->getFirstMemoryCycle(FixedDepMI->getDesc().SchedClass) + 1,
+            Latency));
+      } else if (TII->isLock(FixedDepMI->getOpcode())) {
+        Dep.setLatency(
+            std::max(*TII->getLastMemoryCycle(MI.getDesc().SchedClass) -
+                         TII->getCoreStallCycleAfterLock() + 1,
+                     Latency));
+      } else {
+        Dep.setLatency(Latency);
+      }
       FixedDepSU->addPred(Dep, /*Required=*/true);
     }
 
diff --git a/llvm/lib/Target/AIE/AIEMaxLatencyFinder.cpp b/llvm/lib/Target/AIE/AIEMaxLatencyFinder.cpp
index d481baf95634..61f1eb2283f5 100644
--- a/llvm/lib/Target/AIE/AIEMaxLatencyFinder.cpp
+++ b/llvm/lib/Target/AIE/AIEMaxLatencyFinder.cpp
@@ -90,6 +90,14 @@ static bool overlap(const MachineOperand &SrcOp, const MachineOperand &DstOp,
 /// Check whether Dst depends on Src
 static bool depends(const MachineInstr &Src, const MachineInstr &Dst,
                     const TargetRegisterInfo *TRI) {
+
+  const AIEBaseInstrInfo *const TII = static_cast<const AIEBaseInstrInfo *>(
+      Src.getMF()->getSubtarget().getInstrInfo());
+  // Detect dependency between lock and ld/st intructions.
+  if ((TII->isLock(Src.getOpcode()) && (Dst.mayLoadOrStore())) ||
+      (TII->isLock(Dst.getOpcode()) && (Src.mayLoadOrStore()))) {
+    return true;
+  }
   // We don't try anything clever in terms of alias analysis
   // The memory latency is accounted for by maxLatency() and any
   // possible dependence will be corrected for by its scheduled cycle.
@@ -126,7 +134,6 @@ InstrAndCycle findEarliestRef(const MachineInstr &SrcMI,
                               ArrayRef<MachineBundle> Bundles, int Prune) {
   const TargetRegisterInfo *TRI =
       SrcMI.getMF()->getSubtarget().getRegisterInfo();
-
   int Cycle = 0;
   for (const auto &Bundle : Bundles) {
     if (Cycle >= Prune) {
@@ -134,7 +141,7 @@ InstrAndCycle findEarliestRef(const MachineInstr &SrcMI,
       return {/*MI=*/nullptr, Cycle};
     }
     for (MachineInstr *DstMI : Bundle.getInstrs()) {
-      LLVM_DEBUG(dbgs() << "  " << *DstMI);
+      LLVM_DEBUG(dbgs() << " " << *DstMI);
       if (depends(SrcMI, *DstMI, TRI)) {
         LLVM_DEBUG(dbgs() << "    depends in cycle=" << Cycle << "\n");
         return {DstMI, Cycle};
diff --git a/llvm/test/CodeGen/AIE/aie2/schedule/postpipeliner/interleave-prologue.mir b/llvm/test/CodeGen/AIE/aie2/schedule/postpipeliner/interleave-prologue.mir
index 0e35797781ac..9b1be7c00997 100644
--- a/llvm/test/CodeGen/AIE/aie2/schedule/postpipeliner/interleave-prologue.mir
+++ b/llvm/test/CodeGen/AIE/aie2/schedule/postpipeliner/interleave-prologue.mir
@@ -448,16 +448,15 @@ body:             |
   ; CHECK-NEXT:   successors: %bb.1(0x80000000)
   ; CHECK-NEXT:   liveins: $p0, $p1, $r0, $r18
   ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   ACQ_mLockId_imm 49, killed renamable $r18
+  ; CHECK-NEXT:   NOP
   ; CHECK-NEXT:   $lc = ADD_NC $r0, -7
   ; CHECK-NEXT:   $ls = MOVXM_lng_cg %bb.2
   ; CHECK-NEXT:   $le = MOVXM_lng_cg <mcsymbol .L_LEnd0>
   ; CHECK-NEXT:   $x6, $p0 = VLDB_UNPACK_S16_S8_ag_pstm_nrm killed $p0, $m0
   ; CHECK-NEXT:   $x6, $p0 = VLDB_UNPACK_S16_S8_ag_pstm_nrm killed $p0, $m0
   ; CHECK-NEXT:   $x6, $p0 = VLDB_UNPACK_S16_S8_ag_pstm_nrm killed $p0, $m0
-  ; CHECK-NEXT:   BUNDLE implicit-def $x6, implicit-def $wl6, implicit-def $wh6, implicit-def $p0, implicit killed $p0, implicit $m0, implicit killed $r18 {
-  ; CHECK-NEXT:     $x6, $p0 = VLDB_UNPACK_S16_S8_ag_pstm_nrm killed $p0, $m0
-  ; CHECK-NEXT:     ACQ_mLockId_imm 49, killed renamable $r18
-  ; CHECK-NEXT:   }
+  ; CHECK-NEXT:   $x6, $p0 = VLDB_UNPACK_S16_S8_ag_pstm_nrm killed $p0, $m0
   ; CHECK-NEXT:   BUNDLE implicit-def $r1, implicit-def $x6, implicit-def $wl6, implicit-def $wh6, implicit-def $p0, implicit killed $p0, implicit $m0 {
   ; CHECK-NEXT:     $r1 = MOVA_lda_cg 0
   ; CHECK-NEXT:     $x6, $p0 = VLDB_UNPACK_S16_S8_ag_pstm_nrm killed $p0, $m0