[SCEV] Constant fold MultExpr before applying depth limit.

[dantrushin]

Summary

Users of SCEV reasonably assume that multiplication of two constant
SCEVs will in turn be constant.
However, that is not always the case:
First, we can get here with reached depth limit, and will create
MultExpr SCEV C1 * C2 and cache it.
Then, we can get here with the same operands, but with small depth
level. But this time we will find existing MultExpr SCEV and return
it, instead of expected constant SCEV.

This patch changes getMultExpr to first try constant folding of SCEV
and then apply depth limit, just like getAddExpr does.
It might result in further get{Add|Mult}Expr calls, but they would
perform constant folding only and should not be too expensive.

[efriedma]

llvm/lib/Analysis/ScalarEvolution.cpp

3006

I think you need to rearrange the code a bit more: with the current version of the patch, we increase the depth before checking the current depth.

The general idea makes sense: we should perform optimizations that aren't recursive before the depth check.

[dantrushin]

llvm/lib/Analysis/ScalarEvolution.cpp

3006

Eli, could you explain?
Do you mean Depth + 1 used in calls above on constant folding path?
My reasoning was following:
If we're already at depth limit (Depth == MaxArithDepth) then we don't want to recurse into possible expensive folding.
On the other hand, constant folding will be attempted anyway, which is what we want here.

Does it makes sense?

[dantrushin]

Fix clang-format error

[mkazantsev]

I like the idea, however I would strongly avoid any situation where we can, even theoretically, recurse with more depth than limit.

Do you plan to make similar changes for other SCEV types (e.g. Add)?

llvm/lib/Analysis/ScalarEvolution.cpp

2960

Is this change realy necessary? Looks unrelated, please commit it separately unless there is a reason to have it here. I'm also sure that similar change can be done in Add and some other SCEVs.

2981

We can potentially recurse here with unlimited depth, which is undersiable. Can you please avoid this?

[dantrushin]

llvm/lib/Analysis/ScalarEvolution.cpp

2960

This exactly is copied from getAddExpr (the only other place where depth limits are applied)

2981

Again, the idea was that just constant folding should not take too much time even if going deep. This also allows to keep code clean and simple.
Ok, I'll change it to check for all constant ops

[dantrushin]

Only ignore depth limit for all-constant expressions

[mkazantsev]

Please add a test which clearly shows the impact (e.g. running -analyze scalar-evolution on single mul instruction with zero depth).

llvm/lib/Analysis/ScalarEvolution.cpp

2935

This is too expensive. Provided that ops are sorted by type, it's enough to check the last one.

[dantrushin]

Address comments

[mkazantsev]

LGTM

[mkazantsev]

Please consider making same thing for add as a follow-up.

llvm/lib/Analysis/ScalarEvolution.cpp

 //===- ScalarEvolution.cpp - Scalar Evolution Analysis --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the implementation of the scalar evolution analysis
assert(getEffectiveSCEVType(Ops[i]->getType()) == ETy &&
            "SCEVMulExpr operand types don't match!");
 #endif
 
   // Sort by complexity, this groups all similar expression types together.
   GroupByComplexity(Ops, &LI, DT);
 
   Flags = StrengthenNoWrapFlags(this, scMulExpr, Ops, Flags);
 
-  // Limit recursion calls depth.
-  if (Depth > MaxArithDepth || hasHugeExpression(Ops))
-    return getOrCreateMulExpr(Ops, Flags);
-
-  if (SCEV *S = std::get<0>(findExistingSCEVInCache(scMulExpr, Ops))) {
-    static_cast<SCEVMulExpr *>(S)->setNoWrapFlags(Flags);
-    return S;
-  }
-
   // If there are any constants, fold them together.
   unsigned Idx = 0;
   if (const SCEVConstant *LHSC = dyn_cast<SCEVConstant>(Ops[0])) {
 
     if (Ops.size() == 2)
       // C1*(C2+V) -> C1*C2 + C1*V
       if (const SCEVAddExpr *Add = dyn_cast<SCEVAddExpr>(Ops[1]))
         // If any of Add's ops are Adds or Muls with a constant, apply this
         // transformation as well.
         //
         // TODO: There are some cases where this transformation is not
         // profitable; for example, Add = (C0 + X) * Y + Z.  Maybe the scope of
         // this transformation should be narrowed down.
         if (Add->getNumOperands() == 2 && containsConstantInAddMulChain(Add))
           return getAddExpr(getMulExpr(LHSC, Add->getOperand(0),
                                        SCEV::FlagAnyWrap, Depth + 1),
                             getMulExpr(LHSC, Add->getOperand(1),
                                        SCEV::FlagAnyWrap, Depth + 1),
                             SCEV::FlagAnyWrap, Depth + 1);
 
     ++Idx;
     while (const SCEVConstant *RHSC = dyn_cast<SCEVConstant>(Ops[Idx])) {
       // We found two constants, fold them together!
       ConstantInt *Fold =
           ConstantInt::get(getContext(), LHSC->getAPInt() * RHSC->getAPInt());
       Ops[0] = getConstant(Fold);
+      if (Ops.size() == 2) return Ops[0];
       Ops.erase(Ops.begin()+1);  // Erase the folded element
-      if (Ops.size() == 1) return Ops[0];
       LHSC = cast<SCEVConstant>(Ops[0]);
     }
 
     // If we are left with a constant one being multiplied, strip it off.
     if (cast<SCEVConstant>(Ops[0])->getValue()->isOne()) {
       Ops.erase(Ops.begin());
       --Idx;
     } else if (cast<SCEVConstant>(Ops[0])->getValue()->isZero()) {
       // If we have a multiply of zero, it will always be zero.
       return Ops[0];
     } else if (Ops[0]->isAllOnesValue()) {
       // If we have a mul by -1 of an add, try distributing the -1 among the
       // add operands.
       if (Ops.size() == 2) {
         if (const SCEVAddExpr *Add = dyn_cast<SCEVAddExpr>(Ops[1])) {
           SmallVector<const SCEV *, 4> NewOps;
           bool AnyFolded = false;
           for (const SCEV *AddOp : Add->operands()) {
             const SCEV *Mul = getMulExpr(Ops[0], AddOp, SCEV::FlagAnyWrap,
                                          Depth + 1);
             if (!isa<SCEVMulExpr>(Mul)) AnyFolded = true;
             NewOps.push_back(Mul);
           }
           if (AnyFolded)
             return getAddExpr(NewOps, SCEV::FlagAnyWrap, Depth + 1);
         } else if (const auto *AddRec = dyn_cast<SCEVAddRecExpr>(Ops[1])) {
           // Negation preserves a recurrence's no self-wrap property.
           SmallVector<const SCEV *, 4> Operands;
           for (const SCEV *AddRecOp : AddRec->operands())
             Operands.push_back(getMulExpr(Ops[0], AddRecOp, SCEV::FlagAnyWrap,
                                           Depth + 1));
 
           return getAddRecExpr(Operands, AddRec->getLoop(),
                                AddRec->getNoWrapFlags(SCEV::FlagNW));
         }
       }
     }
 
     if (Ops.size() == 1)
       return Ops[0];
   }
 
+  // Limit recursion calls depth.
+  if (Depth > MaxArithDepth || hasHugeExpression(Ops))
+    return getOrCreateMulExpr(Ops, Flags);
+
+  if (SCEV *S = std::get<0>(findExistingSCEVInCache(scMulExpr, Ops))) {
+    static_cast<SCEVMulExpr *>(S)->setNoWrapFlags(Flags);
+    return S;
+  }
+
   // Skip over the add expression until we get to a multiply.
   while (Idx < Ops.size() && Ops[Idx]->getSCEVType() < scMulExpr)
     ++Idx;
 
   // If there are mul operands inline them all into this expression.
   if (Idx < Ops.size()) {
     bool DeletedMul = false;
     while (const SCEVMulExpr *Mul = dyn_cast<SCEVMulExpr>(Ops[Idx])) {

llvm/test/Analysis/ScalarEvolution/depth-limit-overrun.ll

+; RUN: opt -passes 'strength-reduce' -scalar-evolution-max-arith-depth=2 -S < %s | FileCheck %s
+; RUN: opt -loop-reduce -scalar-evolution-max-arith-depth=2 -S < %s | FileCheck %s
+
+; This test should just compile cleanly without assertions.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:1-p2:32:8:8:32-ni:2"
+
+define void @test(i32 %A, i32 %B, i32 %C) {
+; CHECK-LABEL: @test(
+; CHECK:       inner_loop:
+; CHECK-NEXT:    [[LSR_IV3:%.*]] = phi i32
+; CHECK-NEXT:    [[LSR_IV1:%.*]] = phi i32
+; CHECK-NEXT:    [[LSR_IV:%.*]] = phi i32
+; CHECK:         [[LSR_IV_NEXT:%.*]] = add i32 [[LSR_IV]], 3
+; CHECK-NEXT:    [[LSR_IV_NEXT2:%.*]] = add i32 [[LSR_IV1]], 3
+; CHECK-NEXT:    [[LSR_IV_NEXT4:%.*]] = add i32 [[LSR_IV3]], -3
+;
+entry:
+  br label %outer_loop
+
+outer_loop:
+  %phi2 = phi i32 [ %A, %entry ], [ 204, %outer_tail ]
+  %phi3 = phi i32 [ %A, %entry ], [ 243, %outer_tail ]
+  %phi4 = phi i32 [ %B, %entry ], [ %i35, %outer_tail ]
+  br label %guard
+
+guard:
+  %lcmp.mod = icmp eq i32 %C, 0
+  br i1 %lcmp.mod, label %outer_tail, label %preheader
+
+preheader:
+  %i15 = shl i32 %B, 1
+  br label %inner_loop
+
+inner_loop:
+  %phi5 = phi i32 [ %phi3, %preheader ], [ %i30, %inner_loop ]
+  %phi6 = phi i32 [ %phi2, %preheader ], [ %i33, %inner_loop ]
+  %iter = phi i32 [ %C, %preheader ], [ %iter.sub, %inner_loop ]
+  %i17 = sub i32 %phi4, %phi6
+  %i18 = sub i32 14, %phi5
+  %i19 = mul i32 %i18, %C
+  %factor.prol = shl i32 %phi5, 1
+  %i20 = add i32 %i17, %factor.prol
+  %i21 = add i32 %i20, %B
+  %i22 = add i32 %i21, %i19
+  %i23 = sub i32 14, %i22
+  %i24 = mul i32 %i23, %C
+  %factor.1.prol = shl i32 %i22, 1
+  %i25 = add i32 %i17, %factor.1.prol
+  %i27 = add i32 %i25, %i24
+  %i29 = mul i32 %i25, %C
+  %factor.2.prol = shl i32 %i27, 1
+  %i30 = add i32 %i17, %factor.2.prol
+  %i33 = add nsw i32 %phi6, -3
+  %iter.sub = add i32 %iter, -1
+  %iter.cmp = icmp eq i32 %iter.sub, 0
+  br i1 %iter.cmp, label %outer_tail, label %inner_loop
+
+outer_tail:
+  %phi7 = phi i32 [ %phi2, %guard ], [ %i33, %inner_loop ]
+  %i35 = sub i32 %A, %phi7
+  %cmp = icmp sgt i32 %i35, 9876
+  br i1 %cmp, label %exit, label %outer_loop
+
+exit:
+  ret void
+
+}

llvm/test/Analysis/ScalarEvolution/limit-depth.ll

; CHECK-NEXT:   -->  (%a + %a + %b + %b + %c + %c + %d + %d + %e + %e + %f + %f)
   %s1 = add i32 %p1, %p2
   %s2 = add i32 %s1, %p3
   ret void
 }
 
 define void @test_mul(i32 %a, i32 %b, i32 %c, i32 %d, i32 %e, i32 %f) {
 ; CHECK-LABEL: @test_mul
 ; CHECK:       %s2 = mul i32 %s1, %p3
-; CHECK-NEXT:  -->  (2 * 3 * 4 * 5 * 6 * 7 * %a * %b * %c * %d * %e * %f)
+; CHECK-NEXT:  -->  (5040 * %a * %b * %c * %d * %e * %f)
   %tmp0 = mul i32 %a, 2
   %tmp1 = mul i32 %b, 3
   %tmp2 = mul i32 %c, 4
   %tmp3 = mul i32 %d, 5
   %tmp4 = mul i32 %e, 6
   %tmp5 = mul i32 %f, 7
 
   %p1 = mul i32 %tmp0, %tmp3