[llvm] 664ae7b - [InstCombine] C0 <<{nsw, nuw} (X - C1) --> (C0 >> C1) << X (2nd try)

[Sanjay Patel via llvm-commits]

Author: Sanjay Patel
Date: 2022-04-21T16:18:46-04:00
New Revision: 664ae7bbcc74972531a73d33d7933fc5aa0c0db4

URL: https://github.com/llvm/llvm-project/commit/664ae7bbcc74972531a73d33d7933fc5aa0c0db4
DIFF: https://github.com/llvm/llvm-project/commit/664ae7bbcc74972531a73d33d7933fc5aa0c0db4.diff

LOG: [InstCombine] C0 <<{nsw, nuw} (X - C1) --> (C0 >> C1) << X (2nd try)

The first attempt at this missed a check to make sure the offset
constant was in range and caused many bot failures.

That was missed in the Alive2 proof because on overshift creates
poison rather than the assert from APInt. Here's an alternate
attempt at a proof using count-trailing-zeros:
https://alive2.llvm.org/ce/z/pnXQYR

Original commit message:

This is similar to an existing pre-shift-of-constant fold:
8a9c70fc01e6
...but in this case, we need no-wrap on the shl and a negative
offset:
https://alive2.llvm.org/ce/z/_RVz99

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
    llvm/test/Transforms/InstCombine/shift-add.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index 89558ee06745c..0bd3ac2bf28a1 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -988,6 +988,23 @@ Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
     return BinaryOperator::CreateLShr(
         ConstantInt::get(Ty, APInt::getSignMask(BitWidth)), X);
 
+  // Try to pre-shift a constant shifted by a variable amount:
+  // C << (X + AddC) --> (C >> -AddC) << X
+  // This requires a no-wrap flag and negative offset constant.
+  const APInt *AddC;
+  if ((I.hasNoSignedWrap() || I.hasNoUnsignedWrap()) &&
+      match(Op0, m_APInt(C)) && match(Op1, m_Add(m_Value(X), m_APInt(AddC))) &&
+      AddC->isNegative() && (-*AddC).ult(BitWidth)) {
+    assert(!C->isZero() && "Expected simplify of shifted zero");
+    unsigned PosOffset = (-*AddC).getZExtValue();
+    if (C->eq(C->lshr(PosOffset).shl(PosOffset))) {
+      Constant *NewC = ConstantInt::get(Ty, C->lshr(PosOffset));
+      Instruction *NewShl = BinaryOperator::CreateShl(NewC, X);
+      NewShl->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
+      return NewShl;
+    }
+  }
+
   return nullptr;
 }
 

diff  --git a/llvm/test/Transforms/InstCombine/shift-add.ll b/llvm/test/Transforms/InstCombine/shift-add.ll
index d1618d92e1e07..98786d924309c 100644
--- a/llvm/test/Transforms/InstCombine/shift-add.ll
+++ b/llvm/test/Transforms/InstCombine/shift-add.ll
@@ -174,8 +174,7 @@ define i32 @shl_add_nsw(i32 %x) {
 
 define i32 @shl_nsw_add_negative(i32 %x) {
 ; CHECK-LABEL: @shl_nsw_add_negative(
-; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -1
-; CHECK-NEXT:    [[R:%.*]] = shl nsw i32 2, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = shl i32 1, [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = add i32 %x, -1
@@ -183,11 +182,14 @@ define i32 @shl_nsw_add_negative(i32 %x) {
   ret i32 %r
 }
 
+; vectors and extra uses are allowed
+; nuw propagates to the new shift
+
 define <2 x i8> @shl_nuw_add_negative_splat_uses(<2 x i8> %x, <2 x i8>* %p) {
 ; CHECK-LABEL: @shl_nuw_add_negative_splat_uses(
 ; CHECK-NEXT:    [[A:%.*]] = add <2 x i8> [[X:%.*]], <i8 -2, i8 -2>
 ; CHECK-NEXT:    store <2 x i8> [[A]], <2 x i8>* [[P:%.*]], align 2
-; CHECK-NEXT:    [[R:%.*]] = shl nuw <2 x i8> <i8 12, i8 12>, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = shl nuw <2 x i8> <i8 3, i8 3>, [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %a = add <2 x i8> %x, <i8 -2, i8 -2>
@@ -196,6 +198,8 @@ define <2 x i8> @shl_nuw_add_negative_splat_uses(<2 x i8> %x, <2 x i8>* %p) {
   ret <2 x i8> %r
 }
 
+; negative test - shift constant must have enough trailing zeros to allow the pre-shift
+
 define i32 @shl_nsw_add_negative_invalid_constant(i32 %x) {
 ; CHECK-LABEL: @shl_nsw_add_negative_invalid_constant(
 ; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -2
@@ -207,6 +211,8 @@ define i32 @shl_nsw_add_negative_invalid_constant(i32 %x) {
   ret i32 %r
 }
 
+; negative test - the offset constant must be negative
+
 define i32 @shl_nsw_add_positive_invalid_constant(i32 %x) {
 ; CHECK-LABEL: @shl_nsw_add_positive_invalid_constant(
 ; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], 2
@@ -218,6 +224,8 @@ define i32 @shl_nsw_add_positive_invalid_constant(i32 %x) {
   ret i32 %r
 }
 
+; negative test - a large shift must be detected without crashing
+
 define i32 @shl_nsw_add_negative_invalid_constant2(i32 %x) {
 ; CHECK-LABEL: @shl_nsw_add_negative_invalid_constant2(
 ; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -33
@@ -229,6 +237,9 @@ define i32 @shl_nsw_add_negative_invalid_constant2(i32 %x) {
   ret i32 %r
 }
 
+; negative test - currently transformed to 'xor' before we see it,
+; but INT_MIN should be handled too
+
 define i4 @shl_nsw_add_negative_invalid_constant3(i4 %x) {
 ; CHECK-LABEL: @shl_nsw_add_negative_invalid_constant3(
 ; CHECK-NEXT:    [[A:%.*]] = xor i4 [[X:%.*]], -8