[AIE] ReservedRegsLICM: hoist/sink save/restore brackets

andcarminati · andcarminati · commit 447ee7987aea · 2026-04-30T03:08:06.000-06:00
Some loops protect a reserved register with a save/restore bracket:

  vreg = COPY $reserved_reg    ; save pre-loop value
  $reserved_reg = MOVX imm     ; set loop-invariant value
  ... uses of $reserved_reg ...
  $reserved_reg = COPY vreg    ; restore pre-loop value

The existing hoisting and sinking passes cannot handle this pattern:
the register is a loop livein (the save reads it before the set
defines it), and it has two defs (set + restore), so both
getUniqueDef and the livein guard block any transformation.

Add findSaveRestoreBracket / processSaveRestoreBracket to detect and
transform the pattern as a unit. The detection checks that:
- the save is a vreg = COPY $reserved_reg with exactly one use (the
  restore),
- no use of the register occurs between the save and the set,
- the set is loop-invariant,
- the restore is the last def of the register in the latch and copies
  exactly the vreg from the save,
- the register is not used in any non-latch block of the loop (which
  would observe the wrong value after the set is hoisted).

When the pattern matches, the save and set are moved to the preheader
(in order, so the save captures the pre-loop value before the set
overwrites it) and the restore is sinked to the exit block. The loop
body then always sees the loop-invariant value with no per-iteration
save/restore overhead.

Add a test (crsat_save_restore_bracket) that verifies the
transformation for $crsat.
diff --git a/llvm/lib/Target/AIE/ReservedRegsLICM.cpp b/llvm/lib/Target/AIE/ReservedRegsLICM.cpp
@@ -88,6 +88,19 @@ struct CandidateInfo {
   MachineInstr *HoistCandidate = nullptr;
 };
 
+/// Describes a matched save/restore bracket for a reserved register:
+///   save:    vreg = COPY $reserved_reg   (first use in latch, before set)
+///   set:     $reserved_reg = <imm>       (loop-invariant, immediately after)
+///   restore: $reserved_reg = COPY vreg   (last def in latch)
+/// The transformation hoists save+set to the preheader and sinks restore to
+/// the exit block, so the loop body always sees the loop-invariant value.
+struct SaveRestoreBracket {
+  MachineInstr *Save = nullptr;
+  MachineInstr *Set = nullptr;
+  MachineInstr *Restore = nullptr;
+  MCPhysReg Reg = MCRegister::NoRegister;
+};
+
 /// Used to collect candidates for hoisting/sinking
 class Candidates {
   DenseMap<MCPhysReg, CandidateInfo> Candidates;
@@ -163,6 +176,19 @@ class ReservedRegsLICM : public MachineFunctionPass {
   /// Hoist \p Cand to \p L's preheader if it is safe to do so.
   bool tryHoistToPreHeader(const CandidateInfo &Cand, MachineLoop &L);
 
+  /// Detect and transform a save/restore bracket in \p L's latch block.
+  /// A bracket has the form:
+  ///   vreg = COPY $reserved_reg   (save)
+  ///   $reserved_reg = <imm>       (set, loop-invariant)
+  ///   ...uses of $reserved_reg...
+  ///   $reserved_reg = COPY vreg   (restore)
+  /// Returns the matched bracket, or std::nullopt if not found.
+  std::optional<SaveRestoreBracket> findSaveRestoreBracket(MachineLoop &L);
+
+  /// Apply the save/restore bracket transformation: hoist save+set to the
+  /// preheader and sink restore to the exit block.
+  bool processSaveRestoreBracket(MachineLoop &L);
+
   /// Verify if \p Cand is loop invariant and can be safely hoisted.
   /// \pre Cand->DefinedReg has a unique live value within the loop. This is
   /// verified by processForExitSink() or processForPreheaderHoist().
@@ -257,8 +283,14 @@ void ReservedRegsLICM::runOnLoop(MachineLoop &L) {
   BitVector ReservedLiveins = collectLoopReservedLiveins(L);
   processForExitSink(L, ReservedLiveins);
   processForPreheaderHoist(L, ReservedLiveins);
+  Changed |= processSaveRestoreBracket(L);
 }
 
+// Forward declaration — defined after processForPreheaderHoist.
+static void moveInstruction(const CandidateInfo &Cand,
+                            MachineBasicBlock::iterator InsertBefore,
+                            MachineBasicBlock &InsertMBB);
+
 void ReservedRegsLICM::processForExitSink(MachineLoop &L,
                                           const BitVector &ReservedLiveins) {
   // Pre-compute what's live at the entry of the latch block by walking it
@@ -305,6 +337,145 @@ void ReservedRegsLICM::processForExitSink(MachineLoop &L,
   }
 }
 
+std::optional<SaveRestoreBracket>
+ReservedRegsLICM::findSaveRestoreBracket(MachineLoop &L) {
+  MachineBasicBlock *Latch = L.getLoopLatch();
+  assert(Latch);
+
+  for (MachineInstr &SaveMI : *Latch) {
+    // Look for: vreg = COPY $reserved_reg
+    if (!SaveMI.isCopy())
+      continue;
+    Register SaveDst = SaveMI.getOperand(0).getReg();
+    Register SaveSrc = SaveMI.getOperand(1).getReg();
+    if (!SaveDst.isVirtual() || !SaveSrc.isPhysical())
+      continue;
+    MCPhysReg PhysReg = SaveSrc.asMCReg();
+    if (!TRI->isSimplifiableReservedReg(PhysReg))
+      continue;
+
+    // The saved vreg must have exactly one non-debug use (the restore).
+    if (!MRI->hasOneNonDBGUse(SaveDst))
+      continue;
+
+    // Helper: does MI use PhysReg?
+    auto UsesPhysReg = [&](const MachineInstr &MI) {
+      return any_of(MI.operands(), [&](const MachineOperand &MO) {
+        return MO.isReg() && MO.isUse() && MO.getReg() == PhysReg;
+      });
+    };
+
+    // No uses of PhysReg before the save. Instructions before the save
+    // rely on the original value of PhysReg; hoisting the set would make
+    // them see the loop-invariant value instead.
+    auto BeforeSave =
+        make_range(Latch->begin(), MachineBasicBlock::iterator(SaveMI));
+    if (any_of(BeforeSave, UsesPhysReg))
+      continue;
+
+    // Find the set: the next def of PhysReg after the save.
+    // No use of PhysReg is allowed between the save and the set.
+    MachineInstr *SetMI = nullptr;
+    bool PhysRegUsedBeforeSet = false;
+    for (MachineInstr *Next = SaveMI.getNextNode(); Next;
+         Next = Next->getNextNode()) {
+      if (getSinglePhysRegDef(*Next) == PhysReg) {
+        SetMI = Next;
+        break;
+      }
+      if (UsesPhysReg(*Next)) {
+        PhysRegUsedBeforeSet = true;
+        break;
+      }
+    }
+    if (!SetMI || PhysRegUsedBeforeSet)
+      continue;
+
+    // The set must be loop-invariant (e.g. MOVX imm).
+    CandidateInfo SetCand(PhysReg);
+    SetCand.HoistCandidate = SetMI;
+    if (!isLoopInvariantInst(SetCand, L))
+      continue;
+
+    // Find the restore: the last def of PhysReg in the latch.
+    // It must be: $reserved_reg = COPY SaveDst.
+    MachineInstr *RestoreMI = nullptr;
+    for (MachineInstr &MI2 : reverse(*Latch)) {
+      if (getSinglePhysRegDef(MI2) == PhysReg) {
+        if (MI2.isCopy() && MI2.getOperand(1).getReg() == SaveDst)
+          RestoreMI = &MI2;
+        break;
+      }
+    }
+    if (!RestoreMI)
+      continue;
+
+    // No uses of PhysReg after the restore. Once the restore is sinked to
+    // the exit block, instructions after the restore position in the loop
+    // body would see the loop-invariant value (from the set) instead of
+    // the restored value.
+    auto AfterRestore = make_range(
+        std::next(MachineBasicBlock::iterator(RestoreMI)), Latch->end());
+    if (any_of(AfterRestore, UsesPhysReg))
+      continue;
+
+    // PhysReg must not be used in any non-latch block of the loop.
+    // If it were, those uses might see the wrong value after we hoist
+    // the set to the preheader.
+    if (any_of(L.getBlocks(), [&](MachineBasicBlock *MBB) {
+          return MBB != Latch && any_of(*MBB, UsesPhysReg);
+        }))
+      continue;
+
+    return SaveRestoreBracket{&SaveMI, SetMI, RestoreMI, PhysReg};
+  }
+  return std::nullopt;
+}
+
+bool ReservedRegsLICM::processSaveRestoreBracket(MachineLoop &L) {
+  std::optional<SaveRestoreBracket> BracketOpt = findSaveRestoreBracket(L);
+  if (!BracketOpt)
+    return false;
+
+  const SaveRestoreBracket &Bracket = *BracketOpt;
+  MachineBasicBlock *Preheader = L.getLoopPreheader();
+  MachineBasicBlock *ExitMBB = L.getExitBlock();
+  assert(Preheader && ExitMBB);
+
+  // Ensure the exit block is a dedicated exit (single predecessor).
+  // runOnLoop already verified that the critical edge can be split if needed.
+  if (!ExitMBB->getSinglePredecessor()) {
+    MachineBasicBlock *ExitingBlock = L.getExitingBlock();
+    ExitMBB = ExitingBlock->SplitCriticalEdge(ExitMBB, *this);
+    assert(ExitMBB);
+    LLVM_DEBUG(dbgs() << "Created dedicated exit: "
+                      << printMBBReference(*ExitMBB) << "\n");
+  }
+
+  LLVM_DEBUG(dbgs() << "Save/restore bracket for " << TRI->getName(Bracket.Reg)
+                    << ":\n"
+                    << "  Save:    " << *Bracket.Save << "  Set:     "
+                    << *Bracket.Set << "  Restore: " << *Bracket.Restore);
+
+  // Move save + set to the preheader (save first so it captures the
+  // pre-loop value before the set overwrites it).
+  auto InsertPt = Preheader->getFirstTerminator();
+  CandidateInfo SaveCand(Bracket.Reg);
+  SaveCand.HoistCandidate = Bracket.Save;
+  moveInstruction(SaveCand, InsertPt, *Preheader);
+
+  CandidateInfo SetCand(Bracket.Reg);
+  SetCand.HoistCandidate = Bracket.Set;
+  moveInstruction(SetCand, InsertPt, *Preheader);
+
+  // Sink restore to the exit block.
+  CandidateInfo RestoreCand(Bracket.Reg);
+  RestoreCand.HoistCandidate = Bracket.Restore;
+  moveInstruction(RestoreCand, ExitMBB->getFirstNonPHI(), *ExitMBB);
+
+  return true;
+}
+
 void ReservedRegsLICM::processForPreheaderHoist(
     MachineLoop &L, const BitVector &ReservedLiveins) {
   Candidates HoistCandidates;
@@ -333,9 +504,9 @@ void ReservedRegsLICM::processForPreheaderHoist(
 /// When an instruction is found to only use loop invariant operands that is
 /// safe to hoist/sink, this function is called to actually move the MI out of
 /// the loop.
-void moveInstruction(const CandidateInfo &Cand,
-                     MachineBasicBlock::iterator InsertBefore,
-                     MachineBasicBlock &InsertMBB) {
+static void moveInstruction(const CandidateInfo &Cand,
+                            MachineBasicBlock::iterator InsertBefore,
+                            MachineBasicBlock &InsertMBB) {
   MachineInstr &MI = *Cand.HoistCandidate;
   LLVM_DEBUG(dbgs() << "Moving to " << printMBBReference(InsertMBB) << " from "
                     << printMBBReference(*MI.getParent()) << ": " << MI);
diff --git a/llvm/test/CodeGen/AIE/aie2ps/reserved-reg-licm.mir b/llvm/test/CodeGen/AIE/aie2ps/reserved-reg-licm.mir
@@ -330,6 +330,112 @@ body:             |
     PseudoJ_jump_imm %bb.1
 ...
 
+# $crsat save/restore bracket: the loop saves $crsat, sets it to 1, uses it,
+# then restores it. The save+set should be hoisted to the preheader and the
+# restore should be sinked to the exit block.
+---
+name:            crsat_save_restore_bracket
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: crsat_save_restore_bracket
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT:   liveins: $r0, $r1, $s0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:er = COPY $r0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:acc1024 = COPY $r1
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:es = COPY $s0
+  ; CHECK-NEXT:   [[COPY3:%[0-9]+]]:er = COPY $crsat
+  ; CHECK-NEXT:   $crsat = MOVX_mvx_cr_imm 1
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   $crsat = COPY [[COPY3]]
+  ; CHECK-NEXT:   PseudoRET implicit $lr
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.2(0x7c000000), %bb.1(0x04000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[VSRS_4x_mv_w_srs_cm_srsSign1_:%[0-9]+]]:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 [[COPY1]], [[COPY2]], implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+  ; CHECK-NEXT:   PseudoJNZ [[COPY]], %bb.2
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.1
+  bb.0:
+    successors: %bb.2(0x80000000)
+    liveins: $r0, $r1, $s0
+
+    %0:er = COPY $r0
+    %1:acc1024 = COPY $r1
+    %2:es = COPY $s0
+    PseudoJ_jump_imm %bb.2
+
+  bb.1:
+    PseudoRET implicit $lr
+
+  bb.2:
+    successors: %bb.2(0x7c000000), %bb.1(0x04000000)
+    %3:er = COPY $crsat
+    $crsat = MOVX_mvx_cr_imm 1
+    %4:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 %1, %2, implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+    $crsat = COPY %3
+    PseudoJNZ %0, %bb.2
+    PseudoJ_jump_imm %bb.1
+...
+
+# $crrnd save/restore bracket with uses before the save and after the restore.
+# The bracket must NOT be transformed because $crrnd is used outside the
+# [save, restore] window: once before the save and once after the restore.
+---
+name:            crrnd_bracket_uses_outside_window
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: crrnd_bracket_uses_outside_window
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT:   liveins: $r0, $r1, $s0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:er = COPY $r0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:acc1024 = COPY $r1
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:es = COPY $s0
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   PseudoRET implicit $lr
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.2(0x7c000000), %bb.1(0x04000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[VSRS_4x_mv_w_srs_cm_srsSign1_:%[0-9]+]]:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 [[COPY1]], [[COPY2]], implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+  ; CHECK-NEXT:   [[COPY3:%[0-9]+]]:er = COPY $crrnd
+  ; CHECK-NEXT:   $crrnd = MOVX_mvx_cr_imm 0
+  ; CHECK-NEXT:   [[VSRS_4x_mv_w_srs_cm_srsSign1_1:%[0-9]+]]:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 [[COPY1]], [[COPY2]], implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+  ; CHECK-NEXT:   $crrnd = COPY [[COPY3]]
+  ; CHECK-NEXT:   [[VSRS_4x_mv_w_srs_cm_srsSign1_2:%[0-9]+]]:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 [[COPY1]], [[COPY2]], implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+  ; CHECK-NEXT:   PseudoJNZ [[COPY]], %bb.2
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.1
+  bb.0:
+    successors: %bb.2(0x80000000)
+    liveins: $r0, $r1, $s0
+
+    %0:er = COPY $r0
+    %1:acc1024 = COPY $r1
+    %2:es = COPY $s0
+    PseudoJ_jump_imm %bb.2
+
+  bb.1:
+    PseudoRET implicit $lr
+
+  bb.2:
+    successors: %bb.2(0x7c000000), %bb.1(0x04000000)
+    %5:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 %1, %2, implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+    %3:er = COPY $crrnd
+    $crrnd = MOVX_mvx_cr_imm 0
+    %4:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 %1, %2, implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+    $crrnd = COPY %3
+    %6:vec256 = VSRS_4x_mv_w_srs_cm_srsSign1 %1, %2, implicit-def dead $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign1
+    PseudoJNZ %0, %bb.2
+    PseudoJ_jump_imm %bb.1
+...
+
 # $srssign0 is defined in one conditional branch (bb.2) but not the other
 # (bb.3). The latch (bb.4) uses $srssign0 before defining it, so $srssign0 is
 # live at the latch entry on the path through bb.3. Sinking MOVX_mvx_cr_imm 0
