[AIE] ReservedRegsLICM: don't sink if reg is live at latch entry

In a multi-block loop with conditional branches, a candidate register
may be defined on some paths to the latch but not others. On the paths
that skip the definition the register is live at the latch entry,
carrying a value from a previous iteration. Sinking the last def from
the latch to the exit block would expose that stale value to the
instructions that use the register before the def in the latch.

Fix this by pre-computing LatchEntryLive — the set of registers live
at the entry of the latch block — and rejecting any sink candidate
whose register appears in that set.

Add a test (srs_dyn_sign_conditional_no_sink) where $srssign0 is
defined in one conditional branch but not the other, and the latch
uses it before redefining it. The pass must leave both the COPY and
the MOVX_mvx_cr_imm in place.

Author: andcarminati

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

@@ -261,22 +261,36 @@ void ReservedRegsLICM::runOnLoop(MachineLoop &L) {
 
 void ReservedRegsLICM::processForExitSink(MachineLoop &L,
                                           const BitVector &ReservedLiveins) {
+  // Pre-compute what's live at the entry of the latch block by walking it
+  // fully backward. For multi-block loops with conditional paths, a candidate
+  // register may be defined in some branches but not others, making it live at
+  // the latch entry on those paths. If the register is live at the latch entry
+  // (i.e., used before its first def in the latch), sinking its last def to
+  // the exit block would expose the wrong value on those paths.
+  assert(L.getLoopLatch());
+  LivePhysRegs LatchEntryLive(*TRI);
+  for (const MachineInstr &MI : reverse(*L.getLoopLatch()))
+    LatchEntryLive.stepBackward(MI);
+
   RegDefMap PhysRegChanged(*TRI);
   LivePhysRegs LiveRegs(*TRI);
   Candidates SinkCandidates;
 
   // Walk the latch block, track defs for each register, and
   // collect potential LICM candidates.
-  assert(L.getLoopLatch());
   for (MachineInstr &MI : reverse(*L.getLoopLatch())) {
     CandidateInfo *CandInfo = SinkCandidates.getInfo(MI);
 
     // First time we meet a reserved reg definition while iterating upwards.
-    // If that def is not a loop livein and it isn't used in this block either,
-    // then one can move the instruction to the exit BB of the loop.
+    // If that def is not a loop livein, it isn't used in this block after the
+    // def, and it isn't live at the latch entry (which would indicate it is
+    // used before its def in the latch, possibly carrying a value from a
+    // conditional path that skips the def), then one can move the instruction
+    // to the exit BB of the loop.
     if (CandInfo && !PhysRegChanged.hasChanged(CandInfo->DefinedReg) &&
         !ReservedLiveins.test(CandInfo->DefinedReg) &&
-        !LiveRegs.contains(CandInfo->DefinedReg)) {
+        !LiveRegs.contains(CandInfo->DefinedReg) &&
+        !LatchEntryLive.contains(CandInfo->DefinedReg)) {
       assert(!CandInfo->HoistCandidate);
       CandInfo->HoistCandidate = &MI;
     }

llvm/test/CodeGen/AIE/aie2ps/reserved-reg-licm.mir

@@ -329,3 +329,84 @@ body:             |
     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
+# from the latch to the exit would expose the wrong value on that path.
+# Neither the COPY nor the MOVX should be moved.
+---
+name:            srs_dyn_sign_conditional_no_sink
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: srs_dyn_sign_conditional_no_sink
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $p0, $r0, $r1, $dm0, $s0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:er = COPY $r0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:ep = COPY $p0
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:er = COPY $r1
+  ; CHECK-NEXT:   [[COPY3:%[0-9]+]]:edm = COPY $dm0
+  ; CHECK-NEXT:   [[COPY4:%[0-9]+]]:es = COPY $s0
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.2(0x40000000), %bb.3(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   PseudoJNZ [[COPY]], %bb.3
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.4(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $srssign0 = COPY [[COPY2]]
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.4
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3:
+  ; CHECK-NEXT:   successors: %bb.4(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.4
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.4:
+  ; CHECK-NEXT:   successors: %bb.1(0x7c000000), %bb.5(0x04000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[VSRS_4x_mv_x_srs_dm_srsSign0_:%[0-9]+]]:mxs = VSRS_4x_mv_x_srs_dm_srsSign0 [[COPY3]], [[COPY4]], implicit-def $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign0
+  ; CHECK-NEXT:   $srssign0 = MOVX_mvx_cr_imm 0
+  ; CHECK-NEXT:   PseudoJNZ [[COPY]], %bb.1
+  ; CHECK-NEXT:   PseudoJ_jump_imm %bb.5
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.5:
+  ; CHECK-NEXT:   PseudoRET implicit $lr
+  bb.0:
+    successors: %bb.1(0x80000000)
+    liveins: $p0, $r0, $r1, $dm0, $s0
+    %0:er = COPY $r0
+    %1:ep = COPY $p0
+    %2:er = COPY $r1
+    %3:edm = COPY $dm0
+    %4:es = COPY $s0
+    PseudoJ_jump_imm %bb.1
+
+  bb.1:
+    successors: %bb.2(0x40000000), %bb.3(0x40000000)
+    PseudoJNZ %0, %bb.3
+
+  bb.2:
+    successors: %bb.4(0x80000000)
+    $srssign0 = COPY %2
+    PseudoJ_jump_imm %bb.4
+
+  bb.3:
+    successors: %bb.4(0x80000000)
+    PseudoJ_jump_imm %bb.4
+
+  bb.4:
+    successors: %bb.1(0x7c000000), %bb.5(0x04000000)
+    %5:mxs = VSRS_4x_mv_x_srs_dm_srsSign0 %3, %4, implicit-def $srsrs_of, implicit $crrnd, implicit $crsrsmode, implicit $crsat, implicit $srssign0
+    $srssign0 = MOVX_mvx_cr_imm 0
+    PseudoJNZ %0, %bb.1
+    PseudoJ_jump_imm %bb.5
+
+  bb.5:
+    PseudoRET implicit $lr
+...