[LoopInterchange] Fix handling of LCSSA nodes defined in headers and latches.

[fhahn]

Summary

The code to preserve LCSSA PHIs currently only properly supports
reduction PHIs and PHIs for values defined outside the latches.

This patch improves the LCSSA PHI handling to cover PHIs for values
defined in the latches.

Fixes PR41725.

[fhahn]

Add test case without loop exit, clang-format-diff

[jdoerfert]

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

1333

This seems brittle. I don't know all the preconditions here but I imagine various situations where you might end up with multiple uses for escaping LCSSA phi nodes.

One might be two values escape which happen to be equal, in the inner exit block some optimization merges the two LCSSA phi nodes, in the outer that doesn't happen. Now we have one PHI in the inner exit and two uses in the outer exists. Does this make sense (without asserting we have a transformation like this right now)?

1337

Isn't this a job for P.replaceAllUsesWith(...) ?

[fhahn]

Remove restriction for single use PHIs. We can replace such PHIs with their incoming
value, iff all uses are in the loop exit block or in the outer loop header, if
the incoming value is defined in the inner loop header (reduction phis).

[fhahn]

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

1333

Yeah, it could theoretically happen (but I think it is unlikely in practice, as the lcssa PHIs would be trivially equivalent). The single use is actually not important. It is important where those uses actually are. The legality checks should ensure that. I updated the assertion.

[fhahn]

ping

[jdoerfert]

LGTM

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
+#include "llvm/IR/Verifier.h"
 #include <cassert>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-interchange"
 
#endif
       Changed |= Interchanged;
     }
     return Changed;
   }
 
   bool processLoop(LoopVector LoopList, unsigned InnerLoopId,
                    unsigned OuterLoopId, BasicBlock *LoopNestExit,
                    std::vector<std::vector<char>> &DependencyMatrix) {
+
     LLVM_DEBUG(dbgs() << "Processing Inner Loop Id = " << InnerLoopId
                       << " and OuterLoopId = " << OuterLoopId << "\n");
     Loop *InnerLoop = LoopList[InnerLoopId];
     Loop *OuterLoop = LoopList[OuterLoopId];
 
     LoopInterchangeLegality LIL(OuterLoop, InnerLoop, SE, ORE);
     if (!LIL.canInterchangeLoops(InnerLoopId, OuterLoopId, DependencyMatrix)) {
       LLVM_DEBUG(dbgs() << "Not interchanging loops. Cannot prove legality.\n");
if (BI->getSuccessor(i) == OldBB) {
       DTUpdates.push_back(
           {DominatorTree::UpdateKind::Delete, BI->getParent(), OldBB});
       break;
     }
   }
 }
 
 // Move Lcssa PHIs to the right place.
-static void moveLCSSAPhis(BasicBlock *InnerExit, BasicBlock *InnerLatch,
-                          BasicBlock *OuterLatch) {
+static void moveLCSSAPhis(BasicBlock *InnerExit, BasicBlock *InnerHeader, BasicBlock *InnerLatch, BasicBlock *OuterHeader,
+                          BasicBlock *OuterLatch, BasicBlock *OuterExit) {
+
+  // Deal with LCSSA PHI nodes in the exit block of the inner loop, that are
+  // defined either in the header or latch. Those blocks will become header and
+  // latch of the new outer loop, and the only possible user can be a LCSSA node
+  // in the exit block of the loop nest. We can just substitute the user with
+  // the incoming value and remove the PHI.
+  for (PHINode &P : make_early_inc_range(InnerExit->phis())) {
+    if (P.getNumIncomingValues() != 1)
+      continue;
+    auto I = dyn_cast<Instruction>(P.getIncomingValue(0));
+    if (!I || (I->getParent() != InnerLatch && I->getParent() != InnerHeader))
+      continue;
+    assert(P.hasOneUse() && "Can only deal with PHI nodes with single users here.");
+    PHINode *UP = dyn_cast<PHINode>(*P.user_begin());
+    assert(UP && UP->getParent() == OuterExit && "Cannot preserve PHI node.");
+    (void) UP;
+    P.use_begin()->set(P.getIncomingValue(0));
+    P.eraseFromParent();
+  }
+
   SmallVector<PHINode *, 8> LcssaInnerExit;
-  for (PHINode &P : InnerExit->phis())
+  for (PHINode &P : InnerExit->phis()) {
+    assert(all_of(P.users(), [OuterHeader, OuterExit](User *U) {
+                  return cast<PHINode>(U)->getParent() == OuterExit || cast<PHINode>(U)->getParent() == OuterHeader;}));
     LcssaInnerExit.push_back(&P);
+  }
 
   SmallVector<PHINode *, 8> LcssaInnerLatch;
   for (PHINode &P : InnerLatch->phis())
     LcssaInnerLatch.push_back(&P);
 
   // Lcssa PHIs for values used outside the inner loop are in InnerExit.
   // If a PHI node has users outside of InnerExit, it has a use outside the
   // interchanged loop and we have to preserve it. We move these to
   // InnerLatch, which will become the new exit block for the innermost
-  // loop after interchanging. For PHIs only used in InnerExit, we can just
-  // replace them with the incoming value.
-  for (PHINode *P : LcssaInnerExit) {
-    bool hasUsersOutside = false;
-    for (auto UI = P->use_begin(), E = P->use_end(); UI != E;) {
-      Use &U = *UI;
-      ++UI;
-      auto *Usr = cast<Instruction>(U.getUser());
-      if (Usr->getParent() != InnerExit) {
-        hasUsersOutside = true;
-        continue;
-      }
-      U.set(P->getIncomingValueForBlock(InnerLatch));
-    }
-    if (hasUsersOutside)
-      P->moveBefore(InnerLatch->getFirstNonPHI());
-    else
-      P->eraseFromParent();
-  }
+  // loop after interchanging.
+  for (PHINode *P : LcssaInnerExit)
 
   // If the inner loop latch contains LCSSA PHIs, those come from a child loop
   // and we have to move them to the new inner latch.
   for (PHINode *P : LcssaInnerLatch)
     P->moveBefore(InnerExit->getFirstNonPHI());
 
+  // Deal with LCSSA PHI nodes in the loop nest exit block. For PHIs that have
+  // incoming values from the outer latch or header, we have to add a new PHI
+  // in the inner loop latch, which became the exit block of the outer loop,
+  // after interchanging.
+  if (OuterExit) {
+    for (PHINode &P : OuterExit->phis()) {
+      if (P.getNumIncomingValues() != 1)
+        continue;
+      // Skip Phis with incoming values not defined in the outer loop's header and latch.
+      // Also skip incoming phis defined in the latch. Those should already have been updated.
+     auto I = dyn_cast<Instruction>(P.getIncomingValue(0));
+      if (!I || ((I->getParent() != OuterLatch || isa<PHINode>(I)) && I->getParent() != OuterHeader)) 
+        continue;
+
+      PHINode *NewPhi = dyn_cast<PHINode>(P.clone());
+      NewPhi->setIncomingValue(0, P.getIncomingValue(0));
+      NewPhi->setIncomingBlock(0, OuterLatch);
+      NewPhi->insertBefore(InnerLatch->getFirstNonPHI());
+      P.setIncomingValue(0, NewPhi);
+    }
+  }
+
   // Now adjust the incoming blocks for the LCSSA PHIs.
   // For PHIs moved from Inner's exit block, we need to replace Inner's latch
   // with the new latch.
   updateIncomingBlock(InnerLatch, InnerLatch, OuterLatch);
 }
 
 bool LoopInterchangeTransform::adjustLoopBranches() {
   LLVM_DEBUG(dbgs() << "adjustLoopBranches called\n");
updateSuccessor(InnerLoopLatchBI, InnerLoopLatchSuccessor,
                   OuterLoopLatchSuccessor, DTUpdates);
   updateSuccessor(OuterLoopLatchBI, OuterLoopLatchSuccessor, InnerLoopLatch,
                   DTUpdates);
 
   DT->applyUpdates(DTUpdates);
   restructureLoops(OuterLoop, InnerLoop, InnerLoopPreHeader,
                    OuterLoopPreHeader);
 
-  moveLCSSAPhis(InnerLoopLatchSuccessor, InnerLoopLatch, OuterLoopLatch);
+  moveLCSSAPhis(InnerLoopLatchSuccessor, InnerLoopHeader, InnerLoopLatch, OuterLoopHeader, OuterLoopLatch, InnerLoop->getExitBlock());
   // For PHIs in the exit block of the outer loop, outer's latch has been
   // replaced by Inners'.
   updateIncomingBlock(OuterLoopLatchSuccessor, OuterLoopLatch, InnerLoopLatch);
 
   // Now update the reduction PHIs in the inner and outer loop headers.
   SmallVector<PHINode *, 4> InnerLoopPHIs, OuterLoopPHIs;
   for (PHINode &PHI : drop_begin(InnerLoopHeader->phis(), 1))
     InnerLoopPHIs.push_back(cast<PHINode>(&PHI));

llvm/test/Transforms/LoopInterchange/perserve-lcssa.ll

+; RUN: opt < %s -loop-interchange -verify-loop-lcssa -S | FileCheck %s
+
+; Test case for PR41725. The induction variables in the latches escape the
+; loops and we must move some PHIs around.
+
+@a = common dso_local global i64 0, align 4
+@b = common dso_local global i64 0, align 4
+@c = common dso_local global [10 x [1 x i32 ]] zeroinitializer, align 16
+
+
+; CHECK-LABEL: @test_lcssa_indvars1()
+; CHECK-LABEL: inner.body.split:
+; CHECK-NEXT:    %0 = phi i64 [ %iv.outer.next, %outer.latch ]
+; CHECK-NEXT:    %iv.inner.next = add nsw i64 %iv.inner, -1
+
+; CHECK-LABEL: exit:
+; CHECK-NEXT:    %v4.lcssa = phi i64 [ %0, %inner.body.split ]
+; CHECK-NEXT:    %v8.lcssa.lcssa = phi i64 [ %iv.inner.next, %inner.body.split ]
+; CHECK-NEXT:    store i64 %v8.lcssa.lcssa, i64* @b, align 4
+; CHECK-NEXT:    store i64 %v4.lcssa, i64* @a, align 4
+
+define void @test_lcssa_indvars1()  {
+entry:
+  br label %outer.header
+
+outer.header:                                     ; preds = %outer.latch, %entry
+  %iv.outer = phi i64 [ 0, %entry ], [ %iv.outer.next, %outer.latch ]
+  br label %inner.body
+
+inner.body:                                       ; preds = %inner.body, %outer.header
+  %iv.inner = phi i64 [ 5, %outer.header ], [ %iv.inner.next, %inner.body ]
+  %v7 = getelementptr inbounds [10 x [1 x i32]], [10 x [1 x i32]]* @c, i64 0, i64 %iv.inner, i64 %iv.outer
+  store i32 0, i32* %v7, align 4
+  %iv.inner.next = add nsw i64 %iv.inner, -1
+  %v9 = icmp eq i64 %iv.inner, 0
+  br i1 %v9, label %outer.latch, label %inner.body
+
+outer.latch:                                      ; preds = %inner.body
+  %v8.lcssa = phi i64 [ %iv.inner.next, %inner.body ]
+  %iv.outer.next = add nuw nsw i64 %iv.outer, 1
+  %v5 = icmp ult i64 %iv.outer, 2
+  br i1 %v5, label %outer.header, label %exit
+
+exit:                                             ; preds = %outer.latch
+  %v4.lcssa = phi i64 [ %iv.outer.next, %outer.latch ]
+  %v8.lcssa.lcssa = phi i64 [ %v8.lcssa, %outer.latch ]
+  store i64 %v8.lcssa.lcssa, i64* @b, align 4
+  store i64 %v4.lcssa, i64* @a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @test_lcssa_indvars2()
+; CHECK-LABEL: inner.body.split:
+; CHECK-NEXT:    %0 = phi i64 [ %iv.outer, %outer.latch ]
+; CHECK-NEXT:    %iv.inner.next = add nsw i64 %iv.inner, -1
+
+; CHECK-LABEL: exit:
+; CHECK-NEXT:    %v4.lcssa = phi i64 [ %0, %inner.body.split ]
+; CHECK-NEXT:    %v8.lcssa.lcssa = phi i64 [ %iv.inner, %inner.body.split ]
+; CHECK-NEXT:    store i64 %v8.lcssa.lcssa, i64* @b, align 4
+; CHECK-NEXT:    store i64 %v4.lcssa, i64* @a, align 4
+
+
+define void @test_lcssa_indvars2()  {
+entry:
+  br label %outer.header
+
+outer.header:                                     ; preds = %outer.latch, %entry
+  %iv.outer = phi i64 [ 0, %entry ], [ %iv.outer.next, %outer.latch ]
+  br label %inner.body
+
+inner.body:                                       ; preds = %inner.body, %outer.header
+  %iv.inner = phi i64 [ 5, %outer.header ], [ %iv.inner.next, %inner.body ]
+  %v7 = getelementptr inbounds [10 x [1 x i32]], [10 x [1 x i32]]* @c, i64 0, i64 %iv.inner, i64 %iv.outer
+  store i32 0, i32* %v7, align 4
+  %iv.inner.next = add nsw i64 %iv.inner, -1
+  %v9 = icmp eq i64 %iv.inner.next, 0
+  br i1 %v9, label %outer.latch, label %inner.body
+
+outer.latch:                                      ; preds = %inner.body
+  %v8.lcssa = phi i64 [ %iv.inner, %inner.body ]
+  %iv.outer.next = add nuw nsw i64 %iv.outer, 1
+  %v5 = icmp ult i64 %iv.outer.next, 2
+  br i1 %v5, label %outer.header, label %exit
+
+exit:                                             ; preds = %outer.latch
+  %v4.lcssa = phi i64 [ %iv.outer, %outer.latch ]
+  %v8.lcssa.lcssa = phi i64 [ %v8.lcssa, %outer.latch ]
+  store i64 %v8.lcssa.lcssa, i64* @b, align 4
+  store i64 %v4.lcssa, i64* @a, align 4
+  ret void
+}