[llvm] [SCEV] Add dedicated AffineAddRec matcher + loop matchers (NFC). (PR #141141)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/141141

None

>From 1d8731fbc21d0f7373b6ddba64049dedc1abc0fe Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Thu, 22 May 2025 22:14:13 +0100
Subject: [PATCH] [SCEV] Add dedicated AffineAddRec matcher + loop matchers
 (NFC).

---
 .../Analysis/ScalarEvolutionPatternMatch.h    | 41 ++++++++++++++++++-
 llvm/lib/Analysis/ScalarEvolution.cpp         | 10 ++---
 llvm/lib/Transforms/Scalar/IndVarSimplify.cpp |  3 +-
 .../Transforms/Scalar/LoopStrengthReduce.cpp  |  7 ++--
 4 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ScalarEvolutionPatternMatch.h b/llvm/include/llvm/Analysis/ScalarEvolutionPatternMatch.h
index cfb1b4c6ea6b4..8e9d7e0b72142 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolutionPatternMatch.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolutionPatternMatch.h
@@ -196,11 +196,48 @@ m_scev_UDiv(const Op0_t &Op0, const Op1_t &Op1) {
   return m_scev_Binary<SCEVUDivExpr>(Op0, Op1);
 }
 
+inline class_match<const Loop> m_Loop() { return class_match<const Loop>(); }
+
+/// Match an affine SCEVAddRecExpr.
+template <typename Op0_t, typename Op1_t, typename Loop_t>
+struct SCEVAffineAddRec_match {
+  SCEVBinaryExpr_match<SCEVAddRecExpr, Op0_t, Op1_t> Ops;
+  Loop_t Loop;
+
+  SCEVAffineAddRec_match(Op0_t Op0, Op1_t Op1, Loop_t Loop)
+      : Ops(Op0, Op1), Loop(Loop) {}
+
+  bool match(const SCEV *S) const {
+    return Ops.match(S) && Loop.match(cast<SCEVAddRecExpr>(S)->getLoop());
+  }
+};
+
+/// Match a specified const Loop*.
+struct specificloop_ty {
+  const Loop *L;
+
+  specificloop_ty(const Loop *L) : L(L) {}
+
+  bool match(const Loop *L) const { return L == this->L; }
+};
+
+inline specificloop_ty m_SpecificLoop(const Loop *L) { return L; }
+
+inline bind_ty<const Loop> m_Loop(const Loop *&L) { return L; }
+
 template <typename Op0_t, typename Op1_t>
-inline SCEVBinaryExpr_match<SCEVAddRecExpr, Op0_t, Op1_t>
+inline SCEVAffineAddRec_match<Op0_t, Op1_t, class_match<const Loop>>
 m_scev_AffineAddRec(const Op0_t &Op0, const Op1_t &Op1) {
-  return m_scev_Binary<SCEVAddRecExpr>(Op0, Op1);
+  return SCEVAffineAddRec_match<Op0_t, Op1_t, class_match<const Loop>>(
+      Op0, Op1, m_Loop());
+}
+
+template <typename Op0_t, typename Op1_t, typename Loop_t>
+inline SCEVAffineAddRec_match<Op0_t, Op1_t, Loop_t>
+m_scev_AffineAddRec(const Op0_t &Op0, const Op1_t &Op1, const Loop_t &L) {
+  return SCEVAffineAddRec_match<Op0_t, Op1_t, Loop_t>(Op0, Op1, L);
 }
+
 } // namespace SCEVPatternMatch
 } // namespace llvm
 
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 5e01c29615fab..74384494c9df4 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -12531,14 +12531,14 @@ static bool IsKnownPredicateViaAddRecStart(ScalarEvolution &SE,
     return false;
 
   const SCEV *LStart, *RStart, *Step;
-  if (!match(LHS, m_scev_AffineAddRec(m_SCEV(LStart), m_SCEV(Step))) ||
-      !match(RHS, m_scev_AffineAddRec(m_SCEV(RStart), m_scev_Specific(Step))))
+  const Loop *L;
+  if (!match(LHS,
+             m_scev_AffineAddRec(m_SCEV(LStart), m_SCEV(Step), m_Loop(L))) ||
+      !match(RHS, m_scev_AffineAddRec(m_SCEV(RStart), m_scev_Specific(Step),
+                                      m_SpecificLoop(L))))
     return false;
   const SCEVAddRecExpr *LAR = cast<SCEVAddRecExpr>(LHS);
   const SCEVAddRecExpr *RAR = cast<SCEVAddRecExpr>(RHS);
-  if (LAR->getLoop() != RAR->getLoop())
-    return false;
-
   SCEV::NoWrapFlags NW = ICmpInst::isSigned(Pred) ?
                          SCEV::FlagNSW : SCEV::FlagNUW;
   if (!LAR->getNoWrapFlags(NW) || !RAR->getNoWrapFlags(NW))
diff --git a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
index e774e5fd99cbb..3116b3ccd6170 100644
--- a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
+++ b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -808,7 +808,8 @@ static bool isLoopCounter(PHINode* Phi, Loop *L,
     return false;
 
   const SCEV *S = SE->getSCEV(Phi);
-  if (!match(S, m_scev_AffineAddRec(m_SCEV(), m_scev_One())) ||
+  if (!match(S,
+             m_scev_AffineAddRec(m_SCEV(), m_scev_One(), m_SpecificLoop(L))) ||
       cast<SCEVAddRecExpr>(S)->getLoop() != L)
     return false;
 
diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 1373d940f61b6..7df8d66d81af8 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -557,13 +557,14 @@ static void DoInitialMatch(const SCEV *S, Loop *L,
 
   // Look at addrec operands.
   const SCEV *Start, *Step;
-  if (match(S, m_scev_AffineAddRec(m_SCEV(Start), m_SCEV(Step))) &&
+  const Loop *ARLoop;
+  if (match(S,
+            m_scev_AffineAddRec(m_SCEV(Start), m_SCEV(Step), m_Loop(ARLoop))) &&
       !Start->isZero()) {
     DoInitialMatch(Start, L, Good, Bad, SE);
     DoInitialMatch(SE.getAddRecExpr(SE.getConstant(S->getType(), 0), Step,
                                     // FIXME: AR->getNoWrapFlags()
-                                    cast<SCEVAddRecExpr>(S)->getLoop(),
-                                    SCEV::FlagAnyWrap),
+                                    ARLoop, SCEV::FlagAnyWrap),
                    L, Good, Bad, SE);
     return;
   }