[llvm] [WIP][LSR][term-fold] Configure expander to avoid introducing new IVs (PR #80197)

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/80197

It was reported to me that LSR term-fold was generating in IVs in an outer loop when running on the innermost loop.  This is undesirable and suspected to be a side effect of having the expander in the wrong mode.

Patch is currently WIP due to lack of test case.

>From 2edad7448924439dcc099676415f1d5e98bb0e3c Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Wed, 31 Jan 2024 13:15:53 -0800
Subject: [PATCH] [LSR][term-fold] Configure expander to avoid introducing new
 IVs

It was reported to me that LSR term-fold was generating in IVs in an
outer loop when running on the innermost loop.  This is undesirable
and suspected to be a side effect of having the expander in the wrong
mode.

Patch is currently WIP due to lack of test case.
---
 .../Transforms/Scalar/LoopStrengthReduce.cpp  | 23 +++++++++----------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 7ebc5da8b25a5..3c8c12f443b64 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -6761,8 +6761,8 @@ static llvm::PHINode *GetInductionVariable(const Loop &L, ScalarEvolution &SE,
 }
 
 static std::optional<std::tuple<PHINode *, PHINode *, const SCEV *, bool>>
-canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
-                      const LoopInfo &LI) {
+canFoldTermCondOfLoop(Loop *L, SCEVExpander &Expander, ScalarEvolution &SE,
+                      DominatorTree &DT) {
   if (!L->isInnermost()) {
     LLVM_DEBUG(dbgs() << "Cannot fold on non-innermost loop\n");
     return std::nullopt;
@@ -6814,8 +6814,6 @@ canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
     return std::nullopt;
 
   const SCEV *BECount = SE.getBackedgeTakenCount(L);
-  const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
-  SCEVExpander Expander(SE, DL, "lsr_fold_term_cond");
 
   PHINode *ToHelpFold = nullptr;
   const SCEV *TermValueS = nullptr;
@@ -6986,7 +6984,15 @@ static bool ReduceLoopStrength(Loop *L, IVUsers &IU, ScalarEvolution &SE,
   }();
 
   if (EnableFormTerm) {
-    if (auto Opt = canFoldTermCondOfLoop(L, SE, DT, LI)) {
+    // SCEVExpander for both use in preheader and latch
+    const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
+    SCEVExpander Expander(SE, DL, "lsr_fold_term_cond");
+#ifndef NDEBUG
+    Expander.setDebugType(DEBUG_TYPE);
+#endif
+    Expander.disableCanonicalMode();
+    Expander.enableLSRMode();
+    if (auto Opt = canFoldTermCondOfLoop(L, Expander, SE, DT)) {
       auto [ToFold, ToHelpFold, TermValueS, MustDrop] = *Opt;
 
       Changed = true;
@@ -7009,11 +7015,6 @@ static bool ReduceLoopStrength(Loop *L, IVUsers &IU, ScalarEvolution &SE,
       if (MustDrop)
         cast<Instruction>(LoopValue)->dropPoisonGeneratingFlags();
 
-      // SCEVExpander for both use in preheader and latch
-      const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
-      SCEVExpander Expander(SE, DL, "lsr_fold_term_cond");
-      SCEVExpanderCleaner ExpCleaner(Expander);
-
       assert(Expander.isSafeToExpand(TermValueS) &&
              "Terminating value was checked safe in canFoldTerminatingCondition");
 
@@ -7045,8 +7046,6 @@ static bool ReduceLoopStrength(Loop *L, IVUsers &IU, ScalarEvolution &SE,
 
       OldTermCond->eraseFromParent();
       DeleteDeadPHIs(L->getHeader(), &TLI, MSSAU.get());
-
-      ExpCleaner.markResultUsed();
     }
   }