[llvm] [LoopFlatten] Invalidate SCEV after flattening loops (PR #212744)

[via llvm-commits]

https://github.com/Michael-Chen-NJU created https://github.com/llvm/llvm-project/pull/212744

LoopFlatten updates loop structure and replaces the original linear IV expression with the flattened induction variable. SCEV may otherwise keep stale no-wrap information for the scalarized narrowed index, causing later LoopVectorize to omit the required runtime SCEV check.

In the reduced case, the stale SCEV result treats `%flatten.trunciv` as `{0,+,1}<nuw>`, which lets the zext index be modeled as a widened no-wrap AddRec. After invalidation, SCEV recomputes it as a wrapping i8 recurrence and LoopVectorize emits the required `vector.scevcheck`.

Also make `FunctionToLoopPassAdaptor` preserve `ScalarEvolutionAnalysis` only when the loop pipeline's preserved analyses still preserve it, so a loop pass can actually invalidate SCEV at the function level.

Fixes #212176.

>From 41ec6ab12cde0c7b7234e200a73a8e86a071b433 Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <2802328816 at qq.com>
Date: Wed, 29 Jul 2026 19:19:32 +0800
Subject: [PATCH] [LoopFlatten] Invalidate SCEV after flattening loops

---
 llvm/lib/Transforms/Scalar/LoopFlatten.cpp    |  1 +
 .../lib/Transforms/Scalar/LoopPassManager.cpp |  6 ++-
 .../LoopFlatten/scev-invalidation.ll          | 43 +++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopFlatten/scev-invalidation.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopFlatten.cpp b/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
index e48c47f1b4b89..933edd8a5bb43 100644
--- a/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
@@ -1026,6 +1026,7 @@ PreservedAnalyses LoopFlattenPass::run(LoopNest &LN, LoopAnalysisManager &LAM,
     AR.MSSA->verifyMemorySSA();
 
   auto PA = getLoopPassPreservedAnalyses();
+  PA.abandon<ScalarEvolutionAnalysis>();
   if (AR.MSSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;
diff --git a/llvm/lib/Transforms/Scalar/LoopPassManager.cpp b/llvm/lib/Transforms/Scalar/LoopPassManager.cpp
index 978adeedc88d7..5ba47b20ba655 100644
--- a/llvm/lib/Transforms/Scalar/LoopPassManager.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopPassManager.cpp
@@ -332,10 +332,12 @@ PreservedAnalyses FunctionToLoopPassAdaptor::run(Function &F,
   // loop analysis manager incrementally above.
   PA.preserveSet<AllAnalysesOn<Loop>>();
   PA.preserve<LoopAnalysisManagerFunctionProxy>();
-  // We also preserve the set of standard analyses.
+  // We also preserve the set of standard analyses, unless the loop pipeline
+  // explicitly invalidated SCEV.
   PA.preserve<DominatorTreeAnalysis>();
   PA.preserve<LoopAnalysis>();
-  PA.preserve<ScalarEvolutionAnalysis>();
+  if (PA.getChecker<ScalarEvolutionAnalysis>().preserved())
+    PA.preserve<ScalarEvolutionAnalysis>();
   if (UseMemorySSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;
diff --git a/llvm/test/Transforms/LoopFlatten/scev-invalidation.ll b/llvm/test/Transforms/LoopFlatten/scev-invalidation.ll
new file mode 100644
index 0000000000000..875007e5d4688
--- /dev/null
+++ b/llvm/test/Transforms/LoopFlatten/scev-invalidation.ll
@@ -0,0 +1,43 @@
+; RUN: opt < %s -S -passes='module(cgscc(function(loop(loop-flatten)))),function(loop-vectorize)' --force-vector-width=4 --force-vector-interleave=3 | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Make sure loop-flatten does not preserve stale ScalarEvolution results that
+; make loop-vectorize omit the runtime SCEV check for the scalarized zext index.
+define void @zext_i8(i8 %N, ptr %A) {
+; CHECK-LABEL: @zext_i8(
+; CHECK:       for.cond1.preheader.us.preheader:
+; CHECK:         %flatten.tripcount = mul i64
+; CHECK:         br i1 %min.iters.check, label %scalar.ph, label %vector.scevcheck
+; CHECK:       vector.scevcheck:
+; CHECK-NEXT:    {{%.*}} = add nsw i64 %flatten.tripcount, -1
+; CHECK-NEXT:    {{%.*}} = icmp ugt i64 {{%.*}}, 255
+; CHECK-NEXT:    br i1 {{%.*}}, label %scalar.ph, label %vector.ph
+entry:
+  %cmp20.not = icmp eq i8 %N, 0
+  br i1 %cmp20.not, label %common.ret, label %for.cond1.preheader.us
+
+for.cond1.preheader.us:
+  %i.021.us = phi i8 [ %inc8.us, %for.cond1.for.inc7_crit_edge.us ], [ 0, %entry ]
+  br label %for.body3.us
+
+for.body3.us:
+  %j.019.us = phi i8 [ 0, %for.cond1.preheader.us ], [ %inc.us, %for.body3.us ]
+  %mul.us = mul i8 %i.021.us, %N
+  %add.us = add i8 %j.019.us, %mul.us
+  %idxprom.us = zext i8 %add.us to i64
+  %arrayidx.us = getelementptr [2 x i8], ptr %A, i64 %idxprom.us
+  store i16 1, ptr %arrayidx.us, align 2
+  %inc.us = add i8 %j.019.us, 1
+  %cmp2.us = icmp ult i8 %inc.us, %N
+  br i1 %cmp2.us, label %for.body3.us, label %for.cond1.for.inc7_crit_edge.us
+
+for.cond1.for.inc7_crit_edge.us:
+  %inc8.us = add i8 %i.021.us, 1
+  %cmp.us = icmp ult i8 %inc8.us, %N
+  br i1 %cmp.us, label %for.cond1.preheader.us, label %common.ret
+
+common.ret:
+  ret void
+}