[SCEV] Simplify zext/trunc idiom that appears when handling bitmasks.

[jlebar]

Summary

Specifically, we transform

zext(2^K * (trunc X to iN)) to iM ->
2^K * (zext(trunc X to i{N-K}) to iM)<nuw>

This is helpful because pulling the 2^K out of the zext allows further
optimizations.

[sanjoy]

lgtm

llvm/lib/Analysis/ScalarEvolution.cpp

1771

SA is a bad name for a SCEVMulExpr -- should be SM or Mul. Do you mind fixing this in a followup CL while you're here?

1800

Let's s/TruncTy/NewTruncTy for consistency with NewTruncBits

1801

Super minor issue - if NewTruncBits < 1 then IntegerType::get would have asserted already. If you want to keep this assert for the better error message it is probably better to put it before IntegerType::get.

[jlebar]

Thank you for the reviews, Sanjoy!

Will push with your suggested changes.

llvm/lib/Analysis/ScalarEvolution.cpp

1771

Will do.

1801

Sure, let's just remove this assert.

llvm/lib/Analysis/ScalarEvolution.cpp

if (SA->hasNoUnsignedWrap()) {
       // commute the zero extension with the addition operation.
       SmallVector<const SCEV *, 4> Ops;
       for (const auto *Op : SA->operands())
         Ops.push_back(getZeroExtendExpr(Op, Ty, Depth + 1));
       return getAddExpr(Ops, SCEV::FlagNUW, Depth + 1);
     }
   }
 
   if (auto *SA = dyn_cast<SCEVMulExpr>(Op)) {
     // zext((A * B * ...)<nuw>) --> (zext(A) * zext(B) * ...)<nuw>
     if (SA->hasNoUnsignedWrap()) {
       // If the multiply does not unsign overflow then we can, by definition,
       // commute the zero extension with the multiply operation.
       SmallVector<const SCEV *, 4> Ops;
       for (const auto *Op : SA->operands())
         Ops.push_back(getZeroExtendExpr(Op, Ty, Depth + 1));
       return getMulExpr(Ops, SCEV::FlagNUW, Depth + 1);
     }
+
+    // zext(2^K * (trunc X to iN)) to iM ->
+    // 2^K * (zext(trunc X to i{N-K}) to iM)<nuw>
+    //
+    // Proof:
+    //
+    //     zext(2^K * (trunc X to iN)) to iM
+    //   = zext((trunc X to iN) << K) to iM
+    //   = zext((trunc X to i{N-K}) << K)<nuw> to iM
+    //     (because shl removes the top K bits)
+    //   = zext((2^K * (trunc X to i{N-K}))<nuw>) to iM
+    //   = (2^K * (zext(trunc X to i{N-K}) to iM))<nuw>.
+    //
+    if (SA->getNumOperands() == 2)
+      if (auto *MulLHS = dyn_cast<SCEVConstant>(SA->getOperand(0)))
+        if (MulLHS->getAPInt().isPowerOf2())
+          if (auto *TruncRHS = dyn_cast<SCEVTruncateExpr>(SA->getOperand(1))) {
+            int NewTruncBits = getTypeSizeInBits(TruncRHS->getType()) -
+                               MulLHS->getAPInt().logBase2();
+            Type *TruncTy = IntegerType::get(getContext(), NewTruncBits);
+            assert(NewTruncBits > 0 && "Can't represent 2^N in 2^(N-1) bits.");
+            return getMulExpr(
+                getZeroExtendExpr(MulLHS, Ty),
+                getZeroExtendExpr(
+                    getTruncateExpr(TruncRHS->getOperand(), TruncTy), Ty),
+                SCEV::FlagNUW, Depth + 1);
+          }
   }
 
   // The cast wasn't folded; create an explicit cast node.
   // Recompute the insert position, as it may have been invalidated.
   if (const SCEV *S = UniqueSCEVs.FindNodeOrInsertPos(ID, IP)) return S;
   SCEV *S = new (SCEVAllocator) SCEVZeroExtendExpr(ID.Intern(SCEVAllocator),
                                                    Op, Ty);
   UniqueSCEVs.InsertNode(S, IP);

llvm/test/Analysis/LoopAccessAnalysis/wrapping-pointer-versioning.ll

 ;
 ; This loop has a negative stride for A, and the nusw flag is required in
 ; order to properly extend the increment from i32 -4 to i64 -4.
 
 ; LAA-LABEL: f2
 ; LAA: Memory dependences are safe{{$}}
 ; LAA: SCEV assumptions:
 ; LAA-NEXT: {(2 * (trunc i64 %N to i32)),+,-2}<%for.body> Added Flags: <nusw>
-; LAA-NEXT: {((2 * (zext i32 (2 * (trunc i64 %N to i32)) to i64))<nuw> + %a),+,-4}<%for.body> Added Flags: <nusw>
+; LAA-NEXT: {((4 * (zext i31 (trunc i64 %N to i31) to i64)) + %a),+,-4}<%for.body> Added Flags: <nusw>
 
 ; The expression for %mul_ext as analyzed by SCEV is
 ;     (zext i32 {(2 * (trunc i64 %N to i32)),+,-2}<%for.body> to i64)
 ; We have added the nusw flag to turn this expression into the following SCEV:
 ;     i64 {zext i32 (2 * (trunc i64 %N to i32)) to i64,+,-2}<%for.body>
 
 ; LAA: [PSE]  %arrayidxA = getelementptr i16, i16* %a, i64 %mul_ext:
 ; LAA-NEXT: ((2 * (zext i32 {(2 * (trunc i64 %N to i32)),+,-2}<%for.body> to i64))<nuw> + %a)
-; LAA-NEXT: --> {((2 * (zext i32 (2 * (trunc i64 %N to i32)) to i64))<nuw> + %a),+,-4}<%for.body>
+; LAA-NEXT: --> {((4 * (zext i31 (trunc i64 %N to i31) to i64)) + %a),+,-4}<%for.body>
 
 ; LV-LABEL: f2
 ; LV-LABEL: for.body.lver.check
 
 ; LV: [[OFMul:%[^ ]*]] = call { i32, i1 } @llvm.umul.with.overflow.i32(i32 2, i32 [[BETrunc:%[^ ]*]])
 ; LV-NEXT: [[OFMulResult:%[^ ]*]] = extractvalue { i32, i1 } [[OFMul]], 0
 ; LV-NEXT: [[OFMulOverflow:%[^ ]*]] = extractvalue { i32, i1 } [[OFMul]], 1
 ; LV-NEXT: [[AddEnd:%[^ ]*]] = add i32 [[Start:%[^ ]*]], [[OFMulResult]]

llvm/test/Analysis/ScalarEvolution/and-xor.ll

 ; CHECK: -->  (8 * (zext i1 (trunc i64 %x to i1) to i64))
 
 define i64 @test2(i64 %x) {
   %a = shl i64 %x, 3
   %t = and i64 %a, 8
   %z = xor i64 %t, 8
   ret i64 %z
 }
+
+; Check that we transform the naive lowering of the sequence below,
+;   (4 * (zext i5 (2 * (trunc i32 %x to i5)) to i32)),
+; to
+;   (8 * (zext i4 (trunc i32 %x to i4) to i32))
+;
+; CHECK-LABEL: @test3
+define i32 @test3(i32 %x) {
+  %a = mul i32 %x, 8
+; CHECK: %b
+; CHECK-NEXT: --> (8 * (zext i4 (trunc i32 %x to i4) to i32))
+  %b = and i32 %a, 124
+  ret i32 %b
+}