[llvm] [SCEV] Try to push op into ZExt: C * zext (A + B) -> zext (A*C + B*C) (PR #155300)

[Nikita Popov via llvm-commits]

================
@@ -3199,6 +3199,22 @@ const SCEV *ScalarEvolution::getMulExpr(SmallVectorImpl<const SCEV *> &Ops,
                                AddRec->getNoWrapFlags(FlagsMask));
         }
       }
+
+      // Try to push the constant operand into a ZExt: C + zext (A + B) ->
+      // zext (C*A + C*B) if trunc (C) * (A + B)  does not unsigned-wrap.
+      const SCEVAddExpr *InnerAdd;
+      if (match(Ops[1], m_scev_ZExt(m_scev_Add(InnerAdd)))) {
+        const SCEV *NarrowC = getTruncateExpr(LHSC, InnerAdd->getType());
+        if (isa<SCEVConstant>(InnerAdd->getOperand(0)) &&
+            getZeroExtendExpr(NarrowC, Ops[1]->getType()) == LHSC &&
+            hasFlags(StrengthenNoWrapFlags(this, scMulExpr, {NarrowC, InnerAdd},
+                                           SCEV::FlagAnyWrap),
+                     SCEV::FlagNUW)) {
+          auto *Res =
+              getMulExpr(NarrowC, InnerAdd, SCEV::FlagAnyWrap, Depth + 1);
----------------
nikic wrote:

May as well pass NUW here, as you already proved it?

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