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

[Ryotaro Kasuga via llvm-commits]

================
@@ -1633,10 +1882,68 @@ 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() {
+
+  ArrayRef<LoopInterchangeLegality::SimpleReduction> InnerSimpleReductions =
+      LIL.getInnerSimpleReductions();
+
+  assert(InnerSimpleReductions.size() == 1 &&
+         "So far we only support at most one reduction.");
+
+  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.
+  PHINode *IV = SR.CounterIV;
+  Value *IVInit = IV->getIncomingValueForBlock(InnerLoop->getLoopPreheader());
+  Value *FirstIter = Builder.CreateICmpNE(IV, IVInit, "first.iter");
+
+  // Init new_var to MEM_REF or CONST depending on if it is the first iteration.
+  Value *NewVar = Builder.CreateSelect(FirstIter, LoadMem, SR.Init, "new.var");
+
+  // Replace all uses of reduction var with new variable.
+  SR.Re->replaceAllUsesWith(NewVar);
+
+  // Move store instruction into inner loop, just after reduction next's def.
+  SR.LcssaStorer->setOperand(0, SR.Next);
+  SR.LcssaStorer->moveAfter(dyn_cast<Instruction>(SR.Next));
+}
+
 bool LoopInterchangeTransform::transform(
     ArrayRef<Instruction *> DropNoWrapInsts) {
   bool Transformed = false;
 
+  auto &InnerSimpleReductions = LIL.getInnerSimpleReductions();
----------------
kasuga-fj wrote:

```suggestion
  ArrayRef<SimpleReduction> InnerSimpleReductions = LIL.getInnerSimpleReductions();
```

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