[LoopInterchange] Support multi-level reduction loops

[congzhe]

Summary

This is a follow-up to D120386 where we enabled interchange to get the optimal access pattern for multi-level loopnests. But currently loop interchange fails to interchange multi-level loopnests that have reduction operations inside the loopnest. This patch enables interchange of reduction loops with multi-level loopnests, for example the following loopnest should be interchanged three times to get the optimal access pattern:

; Triply nested loop with reduction, should be able to do interchange three times
; to get the ideal access pattern.
 unsigned f(int e[100][100][100], int f[100][100][100]) {
   unsigned x = 0;
   for (int a = 0; a < 100; a++) {
     for (int b = 0; b < 100; b++) {
       for (int c = 0; c < 100; c++) {
          x += e[c][b][a];
       }
     }
   }
   return x;
 }

Currently interchange only supports reduction loops within 2-level loopnests. For loopnests with more than 2-levels like the case above, (after interchanging the b-loop and c-loop) the pass cannot recognize reduction phis and would bail out. This patch adds capability to recognize reduction phis for multi-level loopnests.

[congzhe]

My apologies for the late update. Last time during a loop meeting our conclusion on this patch was that we need to 1) add more comments to explain the code, and 2) avoid construction of dominator tree and loop info locally.

I've updated the patch accordingly, i.e., added more comments and directly leveraged the loop info that is already available in loop interchange. I'd appreciate it if you could let me know your comments, thanks a lot. @bmahjour @Meinersbur

[Meinersbur]

Consider being less specific to only recognize nested loops, but PHINodes in the loop body in general.

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

789

[Nit] innermost could be considered a single word

790–792

[style] https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

I think the name findInnerReductionPhi implies that the search for the innermost loop happens in that function. Either that, or rename it to findReductionPhi.

findInnerReductionPhi is already a static helper function, you could wrap the new functionality into such a helper function as well.

[serious] This uses InnerMostLoop->getSubLoops().front() even if there are multiple nested loops inside (that hence do not even particpate in the loop interchange). Both loops might participate in the reduction.

797–798

The while condition and its if-guard is essentially the same condition and can be integrated.

I was also thinking about how InnerRedPhi->getParent() represents the CurLoop and can be directly compared InnerLoop instead of its header blocks, but it's the header of the parent loop.
Its dependence on just loops may cause it to break when a condition is involved:

unsigned f(int e[100][100][100], int f[100][100][100]) {
   unsigned x = 0;
   for (int a = 0; a < 100; a++) {
     for (int b = 0; b < 100; b++) {
       if (c)
         for (int c = 0; c < 100; c++) // Not participating in the loop interchange
            x += e[c][b][a];
         } else {
           for (int c = 0; c < 100; c++) 
              x += 1;
        }
     }
   }
   return x;
 }

[serious] Here may want to interchange a and b. b has two subloops where your first while will just randomly pick one of them. Then, there is another PHINode joining the result of the two x where getIncomingValueForBlock will be called with the wrong incoming BB and fail.

[serious] Another problem is that it does not check for what the value for the other incoming block is, it could be something unrelated to a reduction

for (int a = 0; a < 100; a++) 
  for (int b = 0; b < 100; b++) {
    for (int c = 0; c < n; c++)
      new_x = x = e[c][b][a];
    x = PHINode(42, new_x); // should be: PHINode(x, new_x); where `x` is the value when the c-loop has zero iterations and hence skipped by the loop guard.
  }

[hiraditya]

Looks pretty interesting. I'm curious if there is a popular workload/benchmark that would improve with this patch?

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

 }
 
 namespace {
 
 /// LoopInterchangeLegality checks if it is legal to interchange the loop.
 class LoopInterchangeLegality {
 public:
   LoopInterchangeLegality(Loop *Outer, Loop *Inner, ScalarEvolution *SE,
-                          OptimizationRemarkEmitter *ORE)
-      : OuterLoop(Outer), InnerLoop(Inner), SE(SE), ORE(ORE) {}
+                          LoopInfo *LI, OptimizationRemarkEmitter *ORE)
+      : OuterLoop(Outer), InnerLoop(Inner), SE(SE), LI(LI), ORE(ORE) {}
 
   /// Check if the loops can be interchanged.
   bool canInterchangeLoops(unsigned InnerLoopId, unsigned OuterLoopId,
                            CharMatrix &DepMatrix);
 
   /// Discover induction PHIs in the header of \p L. Induction
   /// PHIs are added to \p Inductions.
   bool findInductions(Loop *L, SmallVectorImpl<PHINode *> &Inductions);
private:
   /// OuterInnerReductions. When the outer loop is passed, the inner loop needs
   /// to be passed as \p InnerLoop.
   bool findInductionAndReductions(Loop *L,
                                   SmallVector<PHINode *, 8> &Inductions,
                                   Loop *InnerLoop);
 
   Loop *OuterLoop;
   Loop *InnerLoop;
-
+  LoopInfo *LI;
   ScalarEvolution *SE;
 
   /// Interface to emit optimization remarks.
   OptimizationRemarkEmitter *ORE;
 
   /// Set of reduction PHIs taking part of a reduction across the inner and
   /// outer loop.
   SmallPtrSet<PHINode *, 4> OuterInnerReductions;
#endif
     return Changed;
   }
 
   bool processLoop(Loop *InnerLoop, Loop *OuterLoop, unsigned InnerLoopId,
                    unsigned OuterLoopId,
                    std::vector<std::vector<char>> &DependencyMatrix) {
     LLVM_DEBUG(dbgs() << "Processing InnerLoopId = " << InnerLoopId
                       << " and OuterLoopId = " << OuterLoopId << "\n");
-    LoopInterchangeLegality LIL(OuterLoop, InnerLoop, SE, ORE);
+    LoopInterchangeLegality LIL(OuterLoop, InnerLoop, SE, LI, ORE);
     if (!LIL.canInterchangeLoops(InnerLoopId, OuterLoopId, DependencyMatrix)) {
       LLVM_DEBUG(dbgs() << "Not interchanging loops. Cannot prove legality.\n");
       return false;
     }
     LLVM_DEBUG(dbgs() << "Loops are legal to interchange\n");
     LoopInterchangeProfitability LIP(OuterLoop, InnerLoop, SE, ORE);
     if (!LIP.isProfitable(InnerLoopId, OuterLoopId, DependencyMatrix)) {
       LLVM_DEBUG(dbgs() << "Interchanging loops not profitable.\n");
else {
                                "across the outer loop.\n");
           return false;
         }
       } else {
         assert(PHI.getNumIncomingValues() == 2 &&
                "Phis in loop header should have exactly 2 incoming values");
         // Check if we have a PHI node in the outer loop that has a reduction
         // result from the inner loop as an incoming value.
+        //
+        // For multi-level nested loops, in order to find the reduction phi in
+        // InnerLoop, we need to start from the innermost loop and find the
+        // reduction phi there. Note that InnerLoop might not be the innermost
+        // loop for multi-level nested loops.
         Value *V = followLCSSA(PHI.getIncomingValueForBlock(L->getLoopLatch()));
-        PHINode *InnerRedPhi = findInnerReductionPhi(InnerLoop, V);
+        Loop *InnerMostLoop = InnerLoop;
+        while (!InnerMostLoop->getSubLoops().empty()) {
+          InnerMostLoop = InnerMostLoop->getSubLoops().front();
+        }
+        PHINode *InnerRedPhi = findInnerReductionPhi(InnerMostLoop, V);
+        // Start with the reduction phi in the innermost loop and iteratively
+        // find the target reduction phi in InnerLoop.
+        Loop *CurLoop = InnerMostLoop;
+        if (InnerRedPhi && InnerRedPhi->getParent() != InnerLoop->getHeader()) {
+          while (InnerRedPhi->getParent() != InnerLoop->getHeader()) {
+            InnerRedPhi = cast<PHINode>(InnerRedPhi->getIncomingValueForBlock(
+                CurLoop->getLoopPreheader()));
+            CurLoop = LI->getLoopFor(InnerRedPhi->getParent());
+          }
+        }
         if (!InnerRedPhi ||
             !llvm::is_contained(InnerRedPhi->incoming_values(), &PHI)) {
           LLVM_DEBUG(
               dbgs()
               << "Failed to recognize PHI as an induction or reduction.\n");
           return false;
         }
         OuterInnerReductions.insert(&PHI);

llvm/test/Transforms/LoopInterchange/pr43326-ideal-access-pattern.ll

for.innermost:                                        ; preds = %for.middle.header, %for.innermost
   %indvars.innermost = phi i64 [ 0, %for.middle.header ], [ %indvars.innermost.next, %for.innermost ]
   %arrayidx12 = getelementptr inbounds [10 x [10 x i32]], [10 x [10 x i32]]* %e, i64 %indvars.innermost, i64 %indvars.middle, i64 %indvars.outermost
   %0 = load i32, i32* %arrayidx12
   %arrayidx18 = getelementptr inbounds [10 x [10 x i32]], [10 x [10 x i32]]* %f, i64 %indvars.innermost, i64 %indvars.middle, i64 %indvars.outermost
   store i32 %0, i32* %arrayidx18
   %indvars.innermost.next = add nuw nsw i64 %indvars.innermost, 1
   %exitcond.innermost = icmp ne i64 %indvars.innermost.next, 10
   br i1 %exitcond.innermost, label %for.innermost, label %for.middle.latch
+}
+
+
+; Triply nested loop with reduction operations inside the loopnest,
+; should be able to do interchange three times to get the ideal access pattern.
+;
+; unsigned f(int e[10][10][10], int f[10][10][10]) {
+;   unsigned x = 0;
+;   for (int a = 0; a < 10; a++) {
+;     for (int b = 0; b < 10; b++) {
+;       for (int c = 0; c < 10; c++) {
+;          x += e[c][b][a];
+;       }
+;     }
+;   }
+;   return x;
+; }
+
+; REMARKS: --- !Passed
+; REMARKS-NEXT: Pass:            loop-interchange
+; REMARKS-NEXT: Name:            Interchanged
+; REMARKS-NEXT: Function:        pr43326-triply-nested-reduction
+; REMARKS: --- !Passed
+; REMARKS-NEXT: Pass:            loop-interchange
+; REMARKS-NEXT: Name:            Interchanged
+; REMARKS-NEXT: Function:        pr43326-triply-nested-reduction
+; REMARKS: --- !Passed
+; REMARKS-NEXT: Pass:            loop-interchange
+; REMARKS-NEXT: Name:            Interchanged
+; REMARKS-NEXT: Function:        pr43326-triply-nested-reduction
+define i32 @pr43326-triply-nested-reduction([10 x [10 x i32]]* %e, [10 x [10 x i32]]* %f) {
+entry:
+  br label %for.outermost.header
+
+for.outermost.header:                              ; preds = %entry, %for.outermost.latch
+  %indvars.outermost = phi i64 [ 0, %entry ], [ %indvars.outermost.next, %for.outermost.latch ]
+  %xor.outermost = phi i32 [ 0, %entry ], [ %xor.outermost.lcssa, %for.outermost.latch ]
+  br label %for.middle.header
+
+for.cond.cleanup:                                 ; preds = %for.outermost.latch
+  %xor.outermost.lcssa.lcssa = phi i32 [ %xor.outermost.lcssa, %for.outermost.latch ]
+  ret i32 %xor.outermost.lcssa.lcssa
+
+for.middle.header:                              ; preds = %for.outermost.header, %for.middle.latch
+  %indvars.middle = phi i64 [ 0, %for.outermost.header ], [ %indvars.middle.next, %for.middle.latch ]
+  %xor.middle = phi i32 [ %xor.outermost, %for.outermost.header ], [ %xor.middle.lcssa, %for.middle.latch ]
+  br label %for.innermost
+
+for.outermost.latch:                                ; preds = %for.middle.latch
+  %xor.outermost.lcssa = phi i32 [ %xor.middle.lcssa, %for.middle.latch ]
+  %indvars.outermost.next = add nuw nsw i64 %indvars.outermost, 1
+  %exitcond.outermost = icmp ne i64 %indvars.outermost.next, 10
+  br i1 %exitcond.outermost, label %for.outermost.header, label %for.cond.cleanup
+
+for.middle.latch:                                ; preds = %for.innermost
+  %xor.middle.lcssa = phi i32 [ %xor.reduction, %for.innermost ]
+  %indvars.middle.next = add nuw nsw i64 %indvars.middle, 1
+  %exitcond.middle = icmp ne i64 %indvars.middle.next, 10
+  br i1 %exitcond.middle, label %for.middle.header, label %for.outermost.latch
+
+for.innermost:                                        ; preds = %for.middle.header, %for.innermost
+  %xor.innermost = phi i32 [ %xor.middle, %for.middle.header ], [ %xor.reduction, %for.innermost ]
+  %indvars.innermost = phi i64 [ 0, %for.middle.header ], [ %indvars.innermost.next, %for.innermost ]
+  %arrayidx12 = getelementptr inbounds [10 x [10 x i32]], [10 x [10 x i32]]* %e, i64 %indvars.innermost, i64 %indvars.middle, i64 %indvars.outermost
+  %0 = load i32, i32* %arrayidx12  
+  %xor.reduction = add i32 %xor.innermost, %0
+  %indvars.innermost.next = add nuw nsw i64 %indvars.innermost, 1
+  %exitcond.innermost = icmp ne i64 %indvars.innermost.next, 10
+  br i1 %exitcond.innermost, label %for.innermost, label %for.middle.latch
 }
 No newline at end of file