[LoopInterchange] Support reductions across inner and outer loop.

[fhahn]

Summary

This patch adds logic to detect reductions across the inner and outer
loop by following the incoming values of PHI nodes in the outer loop. If
the incoming values take part in a reduction in the inner loop or come
from outside the outer loop, we found a reduction spanning across inner
and outer loop.

With this change, ~10% more loops are interchanged in the LLVM
test-suite + SPEC2006.

Fixes https://bugs.llvm.org/show_bug.cgi?id=30472

[fhahn]

Rebased, simplified and moved parts to D46198. This patch now only contains the code to detect reduction PHIs across inner and outer loops.

[fhahn]

ping

[fhahn]

Move PHI handling for reductions across inner and outer loop to this patch from D46198 and add an additional test case

[fhahn]

ping

[fhahn]

ping

[fhahn]

ping. I would very much appreciate any feedback on the overall approach of following LCSSA nodes and classification of PHI nodes across inner & outer loop. Maybe there is some existing infrastructure for that, which I missed?

[fhahn]

ping

[xbolva00]

In D43245#1187495, @fhahn wrote:

ping

@efriedma @hfinkel ?

Maybe you can try to write a RFC to llvm-dev

[fhahn]

In D43245#1195421, @xbolva00 wrote:

Maybe you can try to write a RFC to llvm-dev

I would be happy to write something, but I am not sure if that would make it easier to review this.

I think it is a fairly small change (thanks to Eli reviewing some patches leading up to this :)) and it would be great to get some feedback on whether this approach makes sense overall.

[fhahn]

ping :)

[efriedma]

lib/Transforms/Scalar/LoopInterchange.cpp

814

Given you're deleting this, is there some check that disallows interchanging if a PHI is a reduction in the outer loop, but not the inner loop? Something like:

for( int i=1;i<N;i++) {
  for( int j=1;j<N;j++) {
    A[i][j]++;
  }
  X += j;
}

1382–1383

one kind?

1388

Checking isa<PHINode>() here seems sort of confusing... is this supposed to be checking for a PHI node from the outer loop? If so, is there some guarantee it isn't a PHI node from somewhere else?

[fhahn]

Thanks Eli!

lib/Transforms/Scalar/LoopInterchange.cpp

814

Right! We should keep the check and could just not add PHIs across inner and outer loop there I think

1388

Yeah, it seems like inner loop reductions are not handled properly to start with. I will look into that first. It might be helpful to keep track of reduction/induction phis after the initial checks and use that info here (and other places, rather than recomputing)

[xbolva00]

How this patch now interacts with LoopInterchange as LoopPass?

[fhahn]

In D43245#1248907, @xbolva00 wrote:

How this patch now interacts with LoopInterchange as LoopPass?

The change to being a loop pass should not really impact this patch. It needs a small change to work with the preserved LCSSA phis now (which I have locally), but before updating the patch, I need to find some time to improve the PHI handling in general and address Eli's comment properly.

[fhahn]

Updated based on D53027, to keep sets PHIs that are part of inner loop only and across the inner and outer loop. This makes the checks during transform trivial. It also ensures that we do not deal with unclassified PHIs.

[fhahn]

Sorry for the delay Eli! I've restructured the code a bit to fix inner loop only reductions (D53027) and to keep 2 sets for reduction PHIs we can use to find out which kind of PHI we are dealing with.

lib/Transforms/Scalar/LoopInterchange.cpp

1388

The code now just checks OuterInnerReductions and InnerReductions.

[fhahn]

Updated after removing support for inner loop only reductions in D53027. The code is now slightly simpler. What do you think?

[fhahn]

lib/Transforms/Scalar/LoopInterchange.cpp

814

With the new structure, those kinds of loops are ruled out, because reductions in the outer loop need to have a corresponding reduction in the inner loop. I've added test/Transforms/LoopInterchange/outer-only-reductions.ll, which should show this case.

[fhahn]

ping. Eli it would be great if you could have another look now that the patch removing support for inner loop only reductions is committed and this patch is simpler.

[efriedma]

lib/Transforms/Scalar/LoopInterchange.cpp

764

I think "PHI.getNumIncomingValues() == 2" must be true? If it weren't true, it would imply the loop header had more than two predecessors (and I think you disallow that elsewhere).

[fhahn]

Turned check into assertion, added a comment explaining the if (!InnerLoop) {} block and another comment to findInductionAndReductions.

[fhahn]

lib/Transforms/Scalar/LoopInterchange.cpp

764

Yep, I turned it into an assertion.

[efriedma]

LGTM

[fhahn]

Thank you very much for all the reviews, Eli!

lib/Transforms/Scalar/LoopInterchange.cpp

 
 private:
   bool tightlyNested(Loop *Outer, Loop *Inner);
   bool containsUnsafeInstructionsInHeader(BasicBlock *BB);
   bool areAllUsesReductions(Instruction *Ins, Loop *L);
   bool containsUnsafeInstructionsInLatch(BasicBlock *BB);
   bool findInductionAndReductions(Loop *L,
                                   SmallVector<PHINode *, 8> &Inductions,
-                                  SmallVector<PHINode *, 8> &Reductions);
+                                  SmallVector<PHINode *, 8> &Reductions,
+                                  Loop *InnerLoop);
 
   Loop *OuterLoop;
   Loop *InnerLoop;
 
   ScalarEvolution *SE;
   LoopInfo *LI;
   DominatorTree *DT;
   bool PreserveLCSSA;
if (InnerInduction->getIncomingBlock(IncomBlockIndx) ==
             InnerLoopPreheader &&
         !OuterLoop->isLoopInvariant(I)) {
       return false;
     }
   }
   return true;
 }
 
+static Value *followLCSSA(Value *SV) {
+  PHINode *PHI = dyn_cast<PHINode>(SV);
+  if (!PHI)
+    return SV;
+
+  if (PHI->getNumIncomingValues() != 1)
+    return SV;
+  return followLCSSA(PHI->getIncomingValue(0));
+}
+
+/// \brief Checks if V is part of a reduction in L by following its operands
+/// and trying to find an reduction PHI in L.
+static bool isPartOfReduction(Loop *L, Value *V,
+                              SmallPtrSetImpl<Value *> &Visited) {
+  if (!Visited.insert(V).second)
+    return false;
+
+  Instruction *I = dyn_cast<Instruction>(V);
+  if (!I || !L->contains(I->getParent()))
+    return false;
+
+  if (PHINode *PHI = dyn_cast<PHINode>(I)) {
+    RecurrenceDescriptor RD;
+    return RecurrenceDescriptor::isReductionPHI(PHI, L, RD);
+  }
+  return llvm::any_of(I->operands(), [L, &Visited](Value *Op) {
+    return isPartOfReduction(L, Op, Visited);
+  });
+}
+
 bool LoopInterchangeLegality::findInductionAndReductions(
     Loop *L, SmallVector<PHINode *, 8> &Inductions,
-    SmallVector<PHINode *, 8> &Reductions) {
+    SmallVector<PHINode *, 8> &Reductions, Loop *InnerLoop) {
   if (!L->getLoopLatch() || !L->getLoopPredecessor())
     return false;
   for (BasicBlock::iterator I = L->getHeader()->begin(); isa<PHINode>(I); ++I) {
     RecurrenceDescriptor RD;
     InductionDescriptor ID;
     PHINode *PHI = cast<PHINode>(I);
     if (InductionDescriptor::isInductionPHI(PHI, L, SE, ID))
       Inductions.push_back(PHI);
     else if (RecurrenceDescriptor::isReductionPHI(PHI, L, RD))
       Reductions.push_back(PHI);
-    else {
+    else if (InnerLoop && PHI->getNumIncomingValues() == 2) {
+      // Check if we have a PHI node in the outer loop that has a reduction
+      // result from the inner loop as an incoming value.
+      Value *V = followLCSSA(PHI->getIncomingValueForBlock(L->getLoopLatch()));
+      SmallPtrSet<Value *, 10> Visited;
+      if (!isPartOfReduction(InnerLoop, V, Visited))
+        return false;
+      Reductions.push_back(PHI);
+    } else {
       DEBUG(
           dbgs() << "Failed to recognize PHI as an induction or reduction.\n");
       return false;
     }
   }
   return true;
 }
 
 static bool containsSafePHI(BasicBlock *Block, bool isOuterLoopExitBlock) {
   for (auto I = Block->begin(); isa<PHINode>(I); ++I) {
     PHINode *PHI = cast<PHINode>(I);
     // Reduction lcssa phi will have only 1 incoming block that from loop latch.
     if (PHI->getNumIncomingValues() > 1)
       return false;

 bool LoopInterchangeLegality::currentLimitations() {
   BasicBlock *InnerLoopPreHeader = InnerLoop->getLoopPreheader();
   BasicBlock *InnerLoopHeader = InnerLoop->getHeader();
   BasicBlock *InnerLoopLatch = InnerLoop->getLoopLatch();
 
   PHINode *InnerInductionVar;
   SmallVector<PHINode *, 8> Inductions;
   SmallVector<PHINode *, 8> Reductions;
-  if (!findInductionAndReductions(InnerLoop, Inductions, Reductions)) {
+  if (!findInductionAndReductions(InnerLoop, Inductions, Reductions, nullptr)) {
     DEBUG(dbgs() << "Only inner loops with induction or reduction PHI nodes "
                  << "are supported currently.\n");
     ORE->emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "UnsupportedPHIInner",
                                       InnerLoop->getStartLoc(),
                                       InnerLoop->getHeader())
              << "Only inner loops with induction or reduction PHI nodes can be"
                 " interchange currently.";
if (Inductions.size() != 1) {
     });
     return true;
   }
   if (Reductions.size() > 0)
     InnerLoopHasReduction = true;
 
   InnerInductionVar = Inductions.pop_back_val();
   Reductions.clear();
-  if (!findInductionAndReductions(OuterLoop, Inductions, Reductions)) {
+  if (!findInductionAndReductions(OuterLoop, Inductions, Reductions,
+                                  InnerLoop)) {
     DEBUG(dbgs() << "Only outer loops with induction or reduction PHI nodes "
                  << "are supported currently.\n");
     ORE->emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "UnsupportedPHIOuter",
                                       OuterLoop->getStartLoc(),
                                       OuterLoop->getHeader())
              << "Only outer loops with induction or reduction PHI nodes can be"
                 " interchanged currently.";
     });
     return true;
   }
 
-  // Outer loop cannot have reduction because then loops will not be tightly
-  // nested.
-  if (!Reductions.empty()) {
-    DEBUG(dbgs() << "Outer loops with reductions are not supported "
-                 << "currently.\n");
-    ORE->emit([&]() {
-      return OptimizationRemarkMissed(DEBUG_TYPE, "ReductionsOuter",
-                                      OuterLoop->getStartLoc(),
-                                      OuterLoop->getHeader())
-             << "Outer loops with reductions cannot be interchangeed "
-                "currently.";
-    });
-    return true;
-  }
   // TODO: Currently we handle only loops with 1 induction variable.
   if (Inductions.size() != 1) {
     DEBUG(dbgs() << "Loops with more than 1 induction variables are not "
                  << "supported currently.\n");
     ORE->emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "MultiIndutionOuter",
                                       OuterLoop->getStartLoc(),
                                       OuterLoop->getHeader())
if (!InductionPHI) {
       return false;
     }
 
     if (InductionPHI->getIncomingBlock(0) == InnerLoopPreHeader)
       InnerIndexVar = dyn_cast<Instruction>(InductionPHI->getIncomingValue(1));
     else
       InnerIndexVar = dyn_cast<Instruction>(InductionPHI->getIncomingValue(0));
 
-    // Ensure that InductionPHI is the first Phi node as required by
-    // splitInnerLoopHeader
+    // Ensure that InductionPHI is the first Phi node.
     if (&InductionPHI->getParent()->front() != InductionPHI)
       InductionPHI->moveBefore(&InductionPHI->getParent()->front());
 
     // Split at the place were the induction variable is
     // incremented/decremented.
     // TODO: This splitting logic may not work always. Fix this.
     splitInnerLoopLatch(InnerIndexVar);
     DEBUG(dbgs() << "splitInnerLoopLatch done\n");
 
     // Splits the inner loops phi nodes out into a separate basic block.
-    splitInnerLoopHeader();
-    DEBUG(dbgs() << "splitInnerLoopHeader done\n");
+    BasicBlock *InnerLoopHeader = InnerLoop->getHeader();
+    SplitBlock(InnerLoopHeader, InnerLoopHeader->getFirstNonPHI(), DT, LI);
+    DEBUG(dbgs() << "splitting InnerLoopHeader done\n");
   }
 
   Transformed |= adjustLoopLinks();
   if (!Transformed) {
     DEBUG(dbgs() << "adjustLoopLinks failed\n");
     return false;
   }
 
   restructureLoops(InnerLoop, OuterLoop);
   return true;
 }
 
 void LoopInterchangeTransform::splitInnerLoopLatch(Instruction *Inc) {
   BasicBlock *InnerLoopLatch = InnerLoop->getLoopLatch();
   BasicBlock *InnerLoopLatchPred = InnerLoopLatch;
   InnerLoopLatch = SplitBlock(InnerLoopLatchPred, Inc, DT, LI);
 }
 
-void LoopInterchangeTransform::splitInnerLoopHeader() {
-  // Split the inner loop header out. Here make sure that the reduction PHI's
-  // stay in the innerloop body.
-  BasicBlock *InnerLoopHeader = InnerLoop->getHeader();
-  BasicBlock *InnerLoopPreHeader = InnerLoop->getLoopPreheader();
-  SplitBlock(InnerLoopHeader, InnerLoopHeader->getFirstNonPHI(), DT, LI);
-  if (InnerLoopHasReduction) {
-    // Adjust Reduction PHI's in the block. The induction PHI must be the first
-    // PHI in InnerLoopHeader for this to work.
-    SmallVector<PHINode *, 8> PHIVec;
-    for (auto I = std::next(InnerLoopHeader->begin()); isa<PHINode>(I); ++I) {
-      PHINode *PHI = dyn_cast<PHINode>(I);
-      Value *V = PHI->getIncomingValueForBlock(InnerLoopPreHeader);
-      PHI->replaceAllUsesWith(V);
-      PHIVec.push_back((PHI));
-    }
-    for (PHINode *P : PHIVec) {
-      P->eraseFromParent();
-    }
-  }
-
-  DEBUG(dbgs() << "Output of splitInnerLoopHeader InnerLoopHeaderSucc & "
-                  "InnerLoopHeader\n");
-}
-
 /// \brief Move all instructions except the terminator from FromBB right before
 /// InsertBefore
 static void moveBBContents(BasicBlock *FromBB, Instruction *InsertBefore) {
   auto &ToList = InsertBefore->getParent()->getInstList();
   auto &FromList = FromBB->getInstList();
 
   ToList.splice(InsertBefore->getIterator(), FromList, FromList.begin(),
                 FromBB->getTerminator()->getIterator());
else
     OuterLoopLatchSuccessor = OuterLoopLatchBI->getSuccessor(0);
 
   if (InnerLoopLatchBI->getSuccessor(1) == InnerLoopLatchSuccessor)
     InnerLoopLatchBI->setSuccessor(1, OuterLoopLatchSuccessor);
   else
     InnerLoopLatchBI->setSuccessor(0, OuterLoopLatchSuccessor);
 
   updateIncomingBlock(OuterLoopLatchSuccessor, OuterLoopLatch, InnerLoopLatch);
 
   if (OuterLoopLatchBI->getSuccessor(0) == OuterLoopLatchSuccessor) {
     OuterLoopLatchBI->setSuccessor(0, InnerLoopLatch);
   } else {
     OuterLoopLatchBI->setSuccessor(1, 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));
+  for (PHINode &PHI : drop_begin(OuterLoopHeader->phis(), 1))
+    OuterLoopPHIs.push_back(cast<PHINode>(&PHI));
+
+  for (PHINode *PHI : OuterLoopPHIs)
+    PHI->moveBefore(InnerLoopHeader->getFirstNonPHI());
+
+  // Move the PHI nodes from the inner loop header to the outer loop header.
+  // We have to deal with 2 different kind of PHI nodes:
+  //  1) PHI nodes that are part of a reduction in the outer loop
+  //  2) PHI nodes that are part of inner loop-only reductions.
+  //  For 1) we only have to move the PHI node and update the incoming blocks.
+  //  For 2) we can replace the PHI node with the value coming from outside the
+  //  inner loop.
+  for (PHINode *PHI : InnerLoopPHIs) {
+    PHI->moveBefore(OuterLoopHeader->getFirstNonPHI());
+    for (BasicBlock *InBB : PHI->blocks()) {
+      if (InnerLoop->contains(InBB) ||
+          isa<PHINode>(PHI->getIncomingValueForBlock(InBB)))
+        continue;
+
+      PHI->replaceAllUsesWith(PHI->getIncomingValueForBlock(InBB));
+      PHI->eraseFromParent();
+      break;
+    }
+  }
+
+  // Update the incoming blocks for moved PHI nodes.
+  updateIncomingBlock(OuterLoopHeader, InnerLoopPreHeader, OuterLoopPreHeader);
+  updateIncomingBlock(OuterLoopHeader, InnerLoopLatch, OuterLoopLatch);
+  updateIncomingBlock(InnerLoopHeader, OuterLoopPreHeader, InnerLoopPreHeader);
+  updateIncomingBlock(InnerLoopHeader, OuterLoopLatch, InnerLoopLatch);
+
   return true;
 }
 
 void LoopInterchangeTransform::adjustLoopPreheaders() {
   // We have interchanged the preheaders so we need to interchange the data in
   // the preheader as well.
   // This is because the content of inner preheader was previously executed
   // inside the outer loop.

test/Transforms/LoopInterchange/reductions-across-inner-and-outer-loop.ll

+; RUN: opt < %s -basicaa -loop-interchange -simplifycfg -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @test(i32** nocapture readonly %Arr, i32 %k) local_unnamed_addr #0 {
+entry:
+  br label %for.body
+
+
+for.body:                                         ; preds = %for.cond.cleanup3, %entry
+  %indvars.iv23 = phi i64 [ 0, %entry ], [ %indvars.iv.next24, %for.cond.cleanup3 ]
+  %sum.021 = phi i32 [ 0, %entry ], [ %add7.lcssa, %for.cond.cleanup3 ]
+  br label %for.body4
+
+
+for.body4:                                        ; preds = %for.body4, %for.body
+  %indvars.iv = phi i64 [ 0, %for.body ], [ %indvars.iv.next.3, %for.body4 ]
+  %sum.119 = phi i32 [ %sum.021, %for.body ], [ %add7, %for.body4 ]
+  %arrayidx = getelementptr inbounds i32*, i32** %Arr, i64 %indvars.iv
+  %0 = load i32*, i32** %arrayidx, align 8
+  %arrayidx6 = getelementptr inbounds i32, i32* %0, i64 %indvars.iv23
+  %1 = load i32, i32* %arrayidx6, align 4
+  %add = add i32 %sum.119, %k
+  %add7 = add i32 %add, %1
+  %indvars.iv.next.3 = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.3 = icmp eq i64 %indvars.iv.next.3, 1024
+  br i1 %exitcond.3, label %for.cond.cleanup3, label %for.body4
+
+for.cond.cleanup3:                                ; preds = %for.body4
+  %add7.lcssa = phi i32 [ %add7, %for.body4 ]
+  %indvars.iv.next24 = add nuw nsw i64 %indvars.iv23, 1
+  %exitcond25 = icmp eq i64 %indvars.iv.next24, 1024
+  br i1 %exitcond25, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup3
+  %add7.lcssa.lcssa = phi i32 [ %add7.lcssa, %for.cond.cleanup3 ]
+  ret i32 %add7.lcssa.lcssa
+
+
+}
+
+; CHECK-LABEL: @test
+; CHECK-LABEL: entry:
+; CHECK-NEXT: br label %for.body4
+
+; CHECK-LABEL: for.body:                                         ; preds = %for.body4, %for.body
+; CHECK: %indvars.iv23 = phi i64 [ %indvars.iv.next24, %for.body ], [ 0, %for.body4 ]
+; CHECK: %sum.119 = phi i32 [ %add7, %for.body ], [ %sum.021, %for.body4 ]
+; CHECK: br i1 %exitcond25, label %for.body4.split, label %for.body
+
+; CHECK-LABEL: for.body4:
+; CHECK: %indvars.iv = phi i64 [ %indvars.iv.next.3, %for.body4.split ], [ 0, %entry ]
+; CHECK: %sum.021 = phi i32 [ %add7, %for.body4.split ], [ 0, %entry ]
+; CHECK: br label %for.body
+
+; CHECK-LABEL: for.body4.split:                                  ; preds = %for.body
+; CHECK: %indvars.iv.next.3 = add nuw nsw i64 %indvars.iv, 1
+; CHECK: %exitcond.3 = icmp eq i64 %indvars.iv.next.3, 1024
+; CHECK: br i1 %exitcond.3, label %for.cond.cleanup, label %for.body4
+
+; CHECK-LABEL: for.cond.cleanup:                                 ; preds = %for.body4.split
+; CHECK: %add7.lcssa.lcssa = phi i32 [ %add7, %for.body4.split ]
+; CHECK: ret i32 %add7.lcssa.lcssa