[llvm] [LoopInterchange] Support inner-loop simple reductions via UndoSimpleReduction (PR #172970)

[Yingying Wang via llvm-commits]

================
@@ -1633,10 +1821,66 @@ void LoopInterchangeTransform::restructureLoops(
   SE->forgetLoop(NewOuter);
 }
 
+/*
+  User can write, optimizers can generate simple reduction for inner loop. In
+  order to make interchange valid, we have to undo reduction by moving th
+  initialization and store instructions into the inner loop. So far we only
+  handle cases where the reduction variable is initialized to a constant.
+  For example, below code:
+
+  loop:
+    re = phi<0.0, next>
+    next = re op ...
+  reduc_sum = phi<next>       // lcssa phi
+  MEM_REF[idx] = reduc_sum		// LcssaStorer
+
+  is transformed into:
+
+  loop:
+    tmp = MEM_REF[idx];
+    new_var = !first_iteration ? tmp : 0.0;
+    next = new_var op ...
+    MEM_REF[idx] = next;		// after moving
+
+  In this way the initial const is used in the first iteration of loop.
+*/
+void LoopInterchangeTransform::undoSimpleReduction() {
+
+  auto &InnerSimpleReductions = LIL.getInnerSimpleReductions();
+  LoopInterchangeLegality::SimpleReduction *SR = InnerSimpleReductions[0];
+  BasicBlock *InnerLoopHeader = InnerLoop->getHeader();
+  IRBuilder<> Builder(&*(InnerLoopHeader->getFirstNonPHIIt()));
+
+  // When the reduction is intialized from constant value, we need to add
+  // a stmt loading from the memory object to target basic block in inner
+  // loop during undoing the reduction.
+  Instruction *LoadMem = Builder.CreateLoad(SR->ElemTy, SR->MemRef);
+
+  // Check if it's the first iteration.
+  auto &InductionPHIs = LIL.getInnerLoopInductions();
+  PHINode *IV = InductionPHIs[0];
+  Value *IVInit = IV->getIncomingValueForBlock(InnerLoop->getLoopPreheader());
+  Value *FirstIter = Builder.CreateICmpNE(IV, IVInit, "first.iter");
----------------
buggfg wrote:

That makes sense! This should perhaps be the IV used for the loop exit condition.

https://github.com/llvm/llvm-project/pull/172970