[LLVM][SCEV] udiv (mul nuw a, vscale), (mul nuw b, vscale) -> udiv a, b (#157836)

Author: paulwalker-arm

llvm/include/llvm/Analysis/ScalarEvolutionPatternMatch.h

@@ -184,6 +184,7 @@ m_scev_PtrToInt(const Op0_t &Op0) {
 
 /// Match a binary SCEV.
 template <typename SCEVTy, typename Op0_t, typename Op1_t,
+          SCEV::NoWrapFlags WrapFlags = SCEV::FlagAnyWrap,
           bool Commutable = false>
 struct SCEVBinaryExpr_match {
   Op0_t Op0;
@@ -192,6 +193,10 @@ struct SCEVBinaryExpr_match {
   SCEVBinaryExpr_match(Op0_t Op0, Op1_t Op1) : Op0(Op0), Op1(Op1) {}
 
   bool match(const SCEV *S) const {
+    if (auto WrappingS = dyn_cast<SCEVNAryExpr>(S))
+      if (WrappingS->getNoWrapFlags(WrapFlags) != WrapFlags)
+        return false;
+
     auto *E = dyn_cast<SCEVTy>(S);
     return E && E->getNumOperands() == 2 &&
            ((Op0.match(E->getOperand(0)) && Op1.match(E->getOperand(1))) ||
@@ -201,10 +206,12 @@ struct SCEVBinaryExpr_match {
 };
 
 template <typename SCEVTy, typename Op0_t, typename Op1_t,
+          SCEV::NoWrapFlags WrapFlags = SCEV::FlagAnyWrap,
           bool Commutable = false>
-inline SCEVBinaryExpr_match<SCEVTy, Op0_t, Op1_t, Commutable>
+inline SCEVBinaryExpr_match<SCEVTy, Op0_t, Op1_t, WrapFlags, Commutable>
 m_scev_Binary(const Op0_t &Op0, const Op1_t &Op1) {
-  return SCEVBinaryExpr_match<SCEVTy, Op0_t, Op1_t, Commutable>(Op0, Op1);
+  return SCEVBinaryExpr_match<SCEVTy, Op0_t, Op1_t, WrapFlags, Commutable>(Op0,
+                                                                           Op1);
 }
 
 template <typename Op0_t, typename Op1_t>
@@ -220,9 +227,17 @@ m_scev_Mul(const Op0_t &Op0, const Op1_t &Op1) {
 }
 
 template <typename Op0_t, typename Op1_t>
-inline SCEVBinaryExpr_match<SCEVMulExpr, Op0_t, Op1_t, true>
+inline SCEVBinaryExpr_match<SCEVMulExpr, Op0_t, Op1_t, SCEV::FlagAnyWrap, true>
 m_scev_c_Mul(const Op0_t &Op0, const Op1_t &Op1) {
-  return m_scev_Binary<SCEVMulExpr, Op0_t, Op1_t, true>(Op0, Op1);
+  return m_scev_Binary<SCEVMulExpr, Op0_t, Op1_t, SCEV::FlagAnyWrap, true>(Op0,
+                                                                           Op1);
+}
+
+template <typename Op0_t, typename Op1_t>
+inline SCEVBinaryExpr_match<SCEVMulExpr, Op0_t, Op1_t, SCEV::FlagNUW, true>
+m_scev_c_NUWMul(const Op0_t &Op0, const Op1_t &Op1) {
+  return m_scev_Binary<SCEVMulExpr, Op0_t, Op1_t, SCEV::FlagNUW, true>(Op0,
+                                                                       Op1);
 }
 
 template <typename Op0_t, typename Op1_t>

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -3598,6 +3598,13 @@ const SCEV *ScalarEvolution::getUDivExpr(const SCEV *LHS,
     }
   }
 
+  // TODO: Generalize to handle any common factors.
+  // udiv (mul nuw a, vscale), (mul nuw b, vscale) --> udiv a, b
+  const SCEV *NewLHS, *NewRHS;
+  if (match(LHS, m_scev_c_NUWMul(m_SCEV(NewLHS), m_SCEVVScale())) &&
+      match(RHS, m_scev_c_NUWMul(m_SCEV(NewRHS), m_SCEVVScale())))
+    return getUDivExpr(NewLHS, NewRHS);
+
   // The Insertion Point (IP) might be invalid by now (due to UniqueSCEVs
   // changes). Make sure we get a new one.
   IP = nullptr;

llvm/test/Analysis/ScalarEvolution/mul-udiv-folds.ll

@@ -188,3 +188,43 @@ loop:
 exit:
   ret void
 }
+
+define noundef i64 @udiv_mul_common_vscale_factor(i64 %a, i64 %b) {
+; CHECK-LABEL: 'udiv_mul_common_vscale_factor'
+; CHECK-NEXT:  Classifying expressions for: @udiv_mul_common_vscale_factor
+; CHECK-NEXT:    %vs = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    --> vscale U: [1,0) S: [1,0)
+; CHECK-NEXT:    %a.vs = mul i64 %a, %vs
+; CHECK-NEXT:    --> (vscale * %a) U: full-set S: full-set
+; CHECK-NEXT:    %b.vs = mul i64 %b, %vs
+; CHECK-NEXT:    --> (vscale * %b) U: full-set S: full-set
+; CHECK-NEXT:    %div = udiv i64 %a.vs, %b.vs
+; CHECK-NEXT:    --> ((vscale * %a) /u (vscale * %b)) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @udiv_mul_common_vscale_factor
+;
+  %vs = call i64 @llvm.vscale()
+  %a.vs = mul i64 %a, %vs
+  %b.vs = mul i64 %b, %vs
+  %div = udiv i64 %a.vs, %b.vs
+  ret i64 %div
+}
+
+define noundef i64 @udiv_mul_nuw_common_vscale_factor(i64 %a, i64 %b) {
+; CHECK-LABEL: 'udiv_mul_nuw_common_vscale_factor'
+; CHECK-NEXT:  Classifying expressions for: @udiv_mul_nuw_common_vscale_factor
+; CHECK-NEXT:    %vs = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    --> vscale U: [1,0) S: [1,0)
+; CHECK-NEXT:    %a.vs = mul nuw i64 %a, %vs
+; CHECK-NEXT:    --> (vscale * %a)<nuw> U: full-set S: full-set
+; CHECK-NEXT:    %b.vs = mul nuw i64 %b, %vs
+; CHECK-NEXT:    --> (vscale * %b)<nuw> U: full-set S: full-set
+; CHECK-NEXT:    %div = udiv i64 %a.vs, %b.vs
+; CHECK-NEXT:    --> (%a /u %b) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @udiv_mul_nuw_common_vscale_factor
+;
+  %vs = call i64 @llvm.vscale()
+  %a.vs = mul nuw i64 %a, %vs
+  %b.vs = mul nuw i64 %b, %vs
+  %div = udiv i64 %a.vs, %b.vs
+  ret i64 %div
+}