[SCEV] Model ashr(shl(x, n), m) as mul(x, 2^(n-m)) when n > m

[zzheng]

Summary

Given below case:

  %y = shl %x, c0
  %z = ashr %y, c1

when n = m, SCEV models it as sext(trunc(x)). This patch tries to handle
the case where c0 > c1 by using sext(mul(trunc(x), 2^(c0-c1)))) as the SCEV
expression.

[efriedma]

Please upload patches with full context (-U1000000).

lib/Analysis/ScalarEvolution.cpp

5276

Maybe reorganize the code to look more like this (rough outline)?

ConstantInt *LCI = dyn_cast<ConstantInt>(LOp1));
if (!LCI) break;
if (ashr greater than shl) break;
if (ashr less than than shl) {
  // Multiply SCEV
}
return sext(trunc(scev));

test/Analysis/ScalarEvolution/sext-mul.ll

5

The loop induction variable (let's call it i) translates to {0,+,1}<%9>. i*2 translates to {0,+,2}<%9>. sext(trunc(i to i32) to i64) (sext i32 {0,+,1}<%9> to i64). sext(trunc(i*2 to i32) to i64) translates to (sext i32 {0,+,2}<%9> to i64). ashr(shl(i, 33), 32) should also translate to (sext i32 {0,+,2}<%9> to i64).

With your patch, ashr(shl(i, 33), 32) translates to {0,+,2}<%9>, which is wrong because the loop can iterate more that INT_MAX times.

[zzheng]

Addressed comment in 1st version by efriedman.

[efriedma]

lib/Analysis/ScalarEvolution.cpp

5333

You have to check whether Amt == BitWidth before you call getTruncateExpr(); otherwise, you'll get an assertion failure in getTruncateExpr in edge cases like "ashr i32 %x, 0".

5366

Typo: "Handle".

5378

This TODO is suggesting an incorrect transformation; udiv doesn't preserve the sign.

test/Analysis/ScalarEvolution/sext-mul.ll

42

Maybe run opt -instnamer over the testcase, so it's easier to edit if we ever need to change it?

[zzheng]

Added check for edge conditions.

[efriedma]

lib/Analysis/ScalarEvolution.cpp

5266

Maybe rename Amt and LOp1 to something which actually describes the values?

5267

Maybe move this check earlier? You could do it before you even check that the LHS of the ashr is another shift: just "if (CI->isNullValue()) return getSCEV(BO->LHS)".

5290

Hoist this variable definition, so you don't call CI->getZExtValue() twice.

5298

"LShAmt - AShrAmt >= Amt" is equivalent to "LShAmt - AShrAmt >= BitWidth - AshrAmt", which is equivalent to "LShAmt >= BitWidth", which you already check earlier.

5300

APInt::getOneBitSet? (The amount of the multiply could overflow "int".)

[efriedma]

lib/Analysis/ScalarEvolution.cpp

5271

I think you meant to return getSCEV(L) here? Please add a testcase for (ashr (shl x, 1) 0).

5304

64 might not be large enough, and getZExtValue() will assert if the value doesn't fit into a 64-bit integer. The right code is something like APInt Mul = APInt::getOneBitSet(TruncToWidth, LShAmt - AShrAmt); const SCEV *MulSCEV = getConstant(Mul);. Please add a testcase like (ashr (shl i128 x, 100) 1) to show this works correctly.

[efriedma]

lib/Analysis/ScalarEvolution.cpp

5294

To clarify my earlier comment about this: we can in fact transform ashr i32 %x, C to sext(trunc(udiv(x, (1 << C)))).

[zzheng]

Revised per Eli's comments. Added new test case.

[efriedma]

Sanjoy, would you mind looking over this? I'd like a second set of eyes on this.

lib/Analysis/ScalarEvolution.cpp

5323

Whitespace.

5327

"getSCEV(ConstantInt::get(TruncTy, Mul))" can be shortened to "getConstant(Mul)".

[sanjoy]

In D28308#662016, @efriedma wrote:

Sanjoy, would you mind looking over this? I'd like a second set of eyes on this.

I'll take a look at this sometime this week.

[sanjoy]

Did you consider doing this same thing in two more general steps (with separate tests for each case):

• Lower ashr X C as sext(trunc(udiv(X, 1 << C))), irrespective of what X is
• Add a separate rule that (A mul (1 << C0)) udiv (1 << C1) is sext(trunc(A mul (1 << (C0 - C1)))) if C0 > C1

lib/Analysis/ScalarEvolution.cpp

5299

LShOp1 is somewhat ambiguous -- how about ShlOp1 (unique since there is only one Shl in context)?

5316

Please call this ShlAmt, and please call LCI ShlAmtCI or something like that.

test/Analysis/ScalarEvolution/sext-mul.ll

14

Please add some target test cases here (i.e. ones that check specific SCEV expressions for the patterns you want SCEV to catch).

[zzheng]

Sanjoy,

I'm not familiar with SCEV, can you clarify my questions below?

• Lower ashr X C as sext(trunc(udiv(X, 1 << C))), irrespective of what X is

This part should be done in const SCEV *ScalarEvolution::createSCEV(Value *V) as the patch is doing

• Add a separate rule that (A mul (1 << C0)) udiv (1 << C1) is sext(trunc(A mul (1 << (C0 - C1)))) if C0 > C1

Should this be done inside getMulExpr()?

[zzheng]

Adding tranformation: AShr X, C --> sext(trunc(udiv(X, (1<<C))))

Will update test cases later.

[zzheng]

ping..

[efriedma]

lib/Analysis/ScalarEvolution.cpp

5367

dyn_cast<Value>() will always succeed; what are you trying to check here?

test/Analysis/ScalarEvolution/scev-expander-reuse-unroll.ll

6 ↗

(On Diff #90903)

This is defeating the point of the test. I would guess SCEVExpander isn't realizing that sext(trunc(udiv))) is equivalent to an ashr?

[zzheng]

Removed ashr (X, C) -> sext(truc(udiv(X, 1<<C))) from previous revision.
This is causing a test failure and needs more work. I'd like to do it in another patch.

[zzheng]

lib/Analysis/ScalarEvolution.cpp

5368

This is failing test/Analysis/ScalarEvolution/scev-expander-reuse-unroll.ll. See inlined comments of the test.
I'd like to do this in a separated patch.

test/Analysis/ScalarEvolution/scev-expander-reuse-unroll.ll

6 ↗

(On Diff #90903)

"%shr = ashr i32 %add, 16" becomes computable and SCEVExpander hoists its InsertPt to outer loop, causing FindValueInExprValueMap() to return {nullptr, nullptr}.
SCEVExpander then expands it literally and creates another 'select' instruction, failing the test.

[zzheng]

In D28308#665317, @sanjoy wrote:

Did you consider doing this same thing in two more general steps (with separate tests for each case):

• Lower ashr X C as sext(trunc(udiv(X, 1 << C))), irrespective of what X is

This is failing test/Analysis/ScalarEvolution/scev-expander-reuse-unroll.ll.
"%shr = ashr i32 %add, 16" becomes computable and SCEVExpander hoists its InsertPt to outer loop, causing FindValueInExprValueMap() to return {nullptr, nullptr}.
SCEVExpander then expands it literally and creates another 'select' instruction, failing the test.

• Add a separate rule that (A mul (1 << C0)) udiv (1 << C1) is sext(trunc(A mul (1 << (C0 - C1)))) if C0 > C1

For my test case, Shl is already lowered to a mul expr before processing AShr.

AShr can be lowered to udiv(1<<C1) but it's not easy to look for the pattern "(A mul (1 << C0)) udiv (1 << C1)".

[zzheng]

ping..

[zzheng]

ping?

[sanjoy]

Sorry for the delay, I'll make sure to review this in the next one / two days.

[sanjoy]

lgtm modulo minor comments

lib/Analysis/ScalarEvolution.cpp

5322

I'd call this OuterTy, since at this point it isn't obvious what this has to do with sext.

5336

IMO things would read a bit clearer if you just s/TruncToWidth/BitWidth - AShrAmt/ everywhere. But this is a minor stylistic thing, and I'd understand if you did not want to make this change.

5346

If you want to use n and m, please denote that in the legend above, i.e:

// X = Shl A, n
// Y = AShr X, m

or use C0 and C1 in this comment.

5357

Typo: LhlAmt.

test/Analysis/ScalarEvolution/sext-mul.ll

9

Why not CHECK-NEXT here?

test/Analysis/ScalarEvolution/sext-zero.ll

7

Please remove the S: [-9223372036854775808,9223372028264841217) Exits: (-8589934592 + (8589934592 * (zext i32 %arg2 to i64))) LoopDispositions: { %bb7: Computable } bits, here and elsewhere, unless you specifically want to check them. They're pretty distracting.

[zzheng]

addressed Sanjoy's comments

[sanjoy]

lgtm again, in case you were waiting for it.

[zzheng]

rebased and commited as r298631

lib/Analysis/ScalarEvolution.cpp

case Instruction::Xor:
         // the case where instcombine has trimmed non-demanded bits out
         // of an xor with -1.
         if (auto *LBO = dyn_cast<BinaryOperator>(BO->LHS))
           if (ConstantInt *LCI = dyn_cast<ConstantInt>(LBO->getOperand(1)))
             if (LBO->getOpcode() == Instruction::And &&
                 LCI->getValue() == CI->getValue())
               if (const SCEVZeroExtendExpr *Z =
                       dyn_cast<SCEVZeroExtendExpr>(getSCEV(BO->LHS))) {
                 Type *UTy = BO->LHS->getType();
                 const SCEV *Z0 = Z->getOperand();
                 Type *Z0Ty = Z0->getType();
                 unsigned Z0TySize = getTypeSizeInBits(Z0Ty);
 
                 // If C is a low-bits mask, the zero extend is serving to
                 // mask off the high bits. Complement the operand and
                 // re-apply the zext.
                 if (APIntOps::isMask(Z0TySize, CI->getValue()))
                   return getZeroExtendExpr(getNotSCEV(Z0), UTy);
 
                 // If C is a single bit, it may be in the sign-bit position
                 // before the zero-extend. In this case, represent the xor
                 // using an add, which is equivalent, and re-apply the zext.
                 APInt Trunc = CI->getValue().trunc(Z0TySize);
                 if (Trunc.zext(getTypeSizeInBits(UTy)) == CI->getValue() &&
                     Trunc.isSignBit())
                   return getZeroExtendExpr(getAddExpr(Z0, getConstant(Trunc)),
                                            UTy);
               }
       }
       break;
 
   case Instruction::Shl:
     // Turn shift left of a constant amount into a multiply.
     if (ConstantInt *SA = dyn_cast<ConstantInt>(BO->RHS)) {
       uint32_t BitWidth = cast<IntegerType>(SA->getType())->getBitWidth();
 
       // If the shift count is not less than the bitwidth, the result of
       // the shift is undefined. Don't try to analyze it, because the
       // resolution chosen here may differ from the resolution chosen in
       // other parts of the compiler.
       if (SA->getValue().uge(BitWidth))
         break;
 
       // It is currently not resolved how to interpret NSW for left
       // shift by BitWidth - 1, so we avoid applying flags in that
       // case. Remove this check (or this comment) once the situation
       // is resolved. See
       // http://lists.llvm.org/pipermail/llvm-dev/2015-April/084195.html
       // and http://reviews.llvm.org/D8890 .
       auto Flags = SCEV::FlagAnyWrap;
       if (BO->Op && SA->getValue().ult(BitWidth - 1))
         Flags = getNoWrapFlagsFromUB(BO->Op);
 
       Constant *X = ConstantInt::get(getContext(),
         APInt::getOneBitSet(BitWidth, SA->getZExtValue()));
       return getMulExpr(getSCEV(BO->LHS), getSCEV(X), Flags);
     }
     break;
 
     case Instruction::AShr:
-      // For a two-shift sext-inreg, use sext(trunc(x)) as the SCEV expression.
-      if (ConstantInt *CI = dyn_cast<ConstantInt>(BO->RHS))
-        if (Operator *L = dyn_cast<Operator>(BO->LHS))
-          if (L->getOpcode() == Instruction::Shl &&
-              L->getOperand(1) == BO->RHS) {
-            uint64_t BitWidth = getTypeSizeInBits(BO->LHS->getType());
+      // AShr X, C, where C is a constant.
+      ConstantInt *CI = dyn_cast<ConstantInt>(BO->RHS);
+      if (!CI)
+        break;
 
-            // If the shift count is not less than the bitwidth, the result of
-            // the shift is undefined. Don't try to analyze it, because the
-            // resolution chosen here may differ from the resolution chosen in
-            // other parts of the compiler.
-            if (CI->getValue().uge(BitWidth))
-              break;
+      Type *OuterTy = BO->LHS->getType();
+      uint64_t BitWidth = getTypeSizeInBits(OuterTy);
 
-            uint64_t Amt = BitWidth - CI->getZExtValue();
-            if (Amt == BitWidth)
-              return getSCEV(L->getOperand(0)); // shift by zero --> noop
-            return getSignExtendExpr(
-                getTruncateExpr(getSCEV(L->getOperand(0)),
-                                IntegerType::get(getContext(), Amt)),
-                BO->LHS->getType());
-          }
+      if (CI->isNullValue())
+        return getSCEV(BO->LHS); // shift by zero --> noop
+
+      uint64_t AShrAmt = CI->getZExtValue();
+      Type *TruncTy = IntegerType::get(getContext(), BitWidth - AShrAmt);
+
+      Operator *L = dyn_cast<Operator>(BO->LHS);
+      if (L && L->getOpcode() == Instruction::Shl) {
+        // X = Shl A, n
+        // Y = AShr X, m
+        // Both n and m are constant.
+
+        const SCEV *ShlOp0SCEV = getSCEV(L->getOperand(0));
+        if (L->getOperand(1) == BO->RHS)
+          // For a two-shift sext-inreg, i.e. n = m,
+          // use sext(trunc(x)) as the SCEV expression.
+              getTruncateExpr(ShlOp0SCEV, TruncTy), OuterTy);
+
+        ConstantInt *ShlAmtCI = dyn_cast<ConstantInt>(L->getOperand(1));
+        if (ShlAmtCI && ShlAmtCI->getValue().ult(BitWidth)) {
+          uint64_t ShlAmt = ShlAmtCI->getZExtValue();
+          if (ShlAmt > AShrAmt) {
+            // When n > m, use sext(mul(trunc(x), 2^(n-m)))) as the SCEV
+            // expression. We already checked that ShlAmt < BitWidth, so
+            // the multiplier, 1 << (ShlAmt - AShrAmt), fits into TruncTy as
+            // ShlAmt - AShrAmt < Amt.
+            APInt Mul = APInt::getOneBitSet(BitWidth - AShrAmt,
+                                            ShlAmt - AShrAmt);
+            return getSignExtendExpr(
+                getMulExpr(getTruncateExpr(ShlOp0SCEV, TruncTy),
+                getConstant(Mul)), OuterTy);
+          }
+        }
       break;
     }
   }
 
   switch (U->getOpcode()) {
   case Instruction::Trunc:
     return getTruncateExpr(getSCEV(U->getOperand(0)), U->getType());
 
   case Instruction::ZExt:
     return getZeroExtendExpr(getSCEV(U->getOperand(0)), U->getType());
 
   case Instruction::SExt:
     return getSignExtendExpr(getSCEV(U->getOperand(0)), U->getType());
 
   case Instruction::BitCast:
     // BitCasts are no-op casts so we just eliminate the cast.
     if (isSCEVable(U->getType()) && isSCEVable(U->getOperand(0)->getType()))
       return getSCEV(U->getOperand(0));
     break;
 

test/Analysis/ScalarEvolution/sext-mul.ll

+; RUN: opt < %s -analyze -scalar-evolution | FileCheck %s
+
+; CHECK: %tmp9 = shl i64 %tmp8, 33
+; CHECK-NEXT: --> {{.*}} Exits: (-8589934592 + (8589934592 * (zext i32 %arg2 to i64)))
+; CHECK: %tmp10 = ashr exact i64 %tmp9, 32
+; CHECK-NEXT: --> {{.*}} Exits: (sext i32 (-2 + (2 * %arg2)) to i64)
+; CHECK: %tmp11 = getelementptr inbounds i32, i32* %arg, i64 %tmp10
+; CHECK-NEXT: --> {{.*}} Exits: ((4 * (sext i32 (-2 + (2 * %arg2)) to i64)) + %arg)
+; CHECK:  %tmp14 = or i64 %tmp10, 1
+; CHECK-NEXT: --> {{.*}} Exits: (1 + (sext i32 (-2 + (2 * %arg2)) to i64))<nsw>
+; CHECK: %tmp15 = getelementptr inbounds i32, i32* %arg, i64 %tmp14
+; CHECK-NEXT: --> {{.*}} Exits: (4 + (4 * (sext i32 (-2 + (2 * %arg2)) to i64)) + %arg)
+; CHECK:Loop %bb7: backedge-taken count is (-1 + (zext i32 %arg2 to i64))<nsw>
+; CHECK-NEXT:Loop %bb7: max backedge-taken count is -1
+; CHECK-NEXT:Loop %bb7: Predicated backedge-taken count is (-1 + (zext i32 %arg2 to i64))<nsw>
+
+define void @foo(i32* nocapture %arg, i32 %arg1, i32 %arg2) {
+bb:
+  %tmp = icmp sgt i32 %arg2, 0
+  br i1 %tmp, label %bb3, label %bb6
+
+bb3:                                              ; preds = %bb
+  %tmp4 = zext i32 %arg2 to i64
+  br label %bb7
+
+bb5:                                              ; preds = %bb7
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  ret void
+
+bb7:                                              ; preds = %bb7, %bb3
+  %tmp8 = phi i64 [ %tmp18, %bb7 ], [ 0, %bb3 ]
+  %tmp9 = shl i64 %tmp8, 33
+  %tmp10 = ashr exact i64 %tmp9, 32
+  %tmp11 = getelementptr inbounds i32, i32* %arg, i64 %tmp10
+  %tmp12 = load i32, i32* %tmp11, align 4
+  %tmp13 = sub nsw i32 %tmp12, %arg1
+  store i32 %tmp13, i32* %tmp11, align 4
+  %tmp14 = or i64 %tmp10, 1
+  %tmp15 = getelementptr inbounds i32, i32* %arg, i64 %tmp14
+  %tmp16 = load i32, i32* %tmp15, align 4
+  %tmp17 = mul nsw i32 %tmp16, %arg1
+  store i32 %tmp17, i32* %tmp15, align 4
+  %tmp18 = add nuw nsw i64 %tmp8, 1
+  %tmp19 = icmp eq i64 %tmp18, %tmp4
+  br i1 %tmp19, label %bb5, label %bb7
+}
+
+; CHECK: %t10 = ashr exact i128 %t9, 1
+; CHECK-NEXT: --> {{.*}} Exits: (sext i127 (-633825300114114700748351602688 + (633825300114114700748351602688 * (zext i32 %arg5 to i127))) to i128)
+; CHECK: %t14 = or i128 %t10, 1
+; CHECK-NEXT: --> {{.*}} Exits: (1 + (sext i127 (-633825300114114700748351602688 + (633825300114114700748351602688 * (zext i32 %arg5 to i127))) to i128))<nsw>
+; CHECK: Loop %bb7: backedge-taken count is (-1 + (zext i32 %arg5 to i128))<nsw>
+; CHECK-NEXT: Loop %bb7: max backedge-taken count is -1
+; CHECK-NEXT: Loop %bb7: Predicated backedge-taken count is (-1 + (zext i32 %arg5 to i128))<nsw>
+
+define void @goo(i32* nocapture %arg3, i32 %arg4, i32 %arg5) {
+bb:
+  %t = icmp sgt i32 %arg5, 0
+  br i1 %t, label %bb3, label %bb6
+
+bb3:                                              ; preds = %bb
+  %t4 = zext i32 %arg5 to i128
+  br label %bb7
+
+bb5:                                              ; preds = %bb7
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  ret void
+
+bb7:                                              ; preds = %bb7, %bb3
+  %t8 = phi i128 [ %t18, %bb7 ], [ 0, %bb3 ]
+  %t9 = shl i128 %t8, 100
+  %t10 = ashr exact i128 %t9, 1
+  %t11 = getelementptr inbounds i32, i32* %arg3, i128 %t10
+  %t12 = load i32, i32* %t11, align 4
+  %t13 = sub nsw i32 %t12, %arg4
+  store i32 %t13, i32* %t11, align 4
+  %t14 = or i128 %t10, 1
+  %t15 = getelementptr inbounds i32, i32* %arg3, i128 %t14
+  %t16 = load i32, i32* %t15, align 4
+  %t17 = mul nsw i32 %t16, %arg4
+  store i32 %t17, i32* %t15, align 4
+  %t18 = add nuw nsw i128 %t8, 1
+  %t19 = icmp eq i128 %t18, %t4
+  br i1 %t19, label %bb5, label %bb7
+}

test/Analysis/ScalarEvolution/sext-zero.ll

+; RUN: opt < %s -analyze -scalar-evolution | FileCheck %s
+
+; CHECK:  %tmp9 = shl i64 %tmp8, 33
+; CHECK-NEXT:  -->  {{.*}} Exits: (-8589934592 + (8589934592 * (zext i32 %arg2 to i64)))
+; CHECK-NEXT:  %tmp10 = ashr exact i64 %tmp9, 0
+; CHECK-NEXT:  -->  {{.*}} Exits: (-8589934592 + (8589934592 * (zext i32 %arg2 to i64)))
+
+define void @foo(i32* nocapture %arg, i32 %arg1, i32 %arg2) {
+bb:
+  %tmp = icmp sgt i32 %arg2, 0
+  br i1 %tmp, label %bb3, label %bb6
+
+bb3:                                              ; preds = %bb
+  %tmp4 = zext i32 %arg2 to i64
+  br label %bb7
+
+bb5:                                              ; preds = %bb7
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  ret void
+
+bb7:                                              ; preds = %bb7, %bb3
+  %tmp8 = phi i64 [ %tmp18, %bb7 ], [ 0, %bb3 ]
+  %tmp9 = shl i64 %tmp8, 33
+  %tmp10 = ashr exact i64 %tmp9, 0
+  %tmp11 = getelementptr inbounds i32, i32* %arg, i64 %tmp10
+  %tmp12 = load i32, i32* %tmp11, align 4
+  %tmp13 = sub nsw i32 %tmp12, %arg1
+  store i32 %tmp13, i32* %tmp11, align 4
+  %tmp14 = or i64 %tmp10, 1
+  %tmp15 = getelementptr inbounds i32, i32* %arg, i64 %tmp14
+  %tmp16 = load i32, i32* %tmp15, align 4
+  %tmp17 = mul nsw i32 %tmp16, %arg1
+  store i32 %tmp17, i32* %tmp15, align 4
+  %tmp18 = add nuw nsw i64 %tmp8, 1
+  %tmp19 = icmp eq i64 %tmp18, %tmp4
+  br i1 %tmp19, label %bb5, label %bb7
+}