From 2552f65183ebb43cdd6de951c63a3dcbfa021b5e Mon Sep 17 00:00:00 2001
From: Sanjay Patel <spatel@rotateright.com>
Date: Sun, 7 Jun 2020 09:03:45 -0400
Subject: [PATCH] [InstCombine] fold mask op into casted shift (PR46013)

https://rise4fun.com/Alive/Qply8

  Pre: C2 == (-1 u>> zext(C1))
  %a = ashr %x, C1
  %s = sext %a to i16
  %r = and i16 %s, C2
    =>
  %s2 = sext %x to i16
  %r = lshr i16 %s2, zext(C1)

https://bugs.llvm.org/show_bug.cgi?id=46013
---
 .../Transforms/InstCombine/InstCombineAndOrXor.cpp | 12 +++++-
 llvm/test/Transforms/InstCombine/and.ll            | 50 ++++++++++++++++++----
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index 35dcb73..d3c718a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -1848,7 +1848,17 @@ Instruction *InstCombiner::visitAnd(BinaryOperator &I) {
         return BinaryOperator::Create(BinOp, NewLHS, Y);
       }
     }
-
+    const APInt *ShiftC;
+    if (match(Op0, m_OneUse(m_SExt(m_AShr(m_Value(X), m_APInt(ShiftC)))))) {
+      unsigned Width = I.getType()->getScalarSizeInBits();
+      if (*C == APInt::getLowBitsSet(Width, Width - ShiftC->getZExtValue())) {
+        // We are clearing high bits that were potentially set by sext+ashr:
+        // and (sext (ashr X, ShiftC)), C --> lshr (sext X), ShiftC
+        Value *Sext = Builder.CreateSExt(X, I.getType());
+        Constant *ShAmtC = ConstantInt::get(I.getType(), ShiftC->zext(Width));
+        return BinaryOperator::CreateLShr(Sext, ShAmtC);
+      }
+    }
   }
 
   if (ConstantInt *AndRHS = dyn_cast<ConstantInt>(Op1)) {
diff --git a/llvm/test/Transforms/InstCombine/and.ll b/llvm/test/Transforms/InstCombine/and.ll
index 408c3c9..13ce3b8 100644
--- a/llvm/test/Transforms/InstCombine/and.ll
+++ b/llvm/test/Transforms/InstCombine/and.ll
@@ -847,9 +847,8 @@ define i1 @andn_or_cmp_4(i32 %a, i32 %b, i32 %c) {
 
 define i32 @lowbitmask_casted_shift(i8 %x) {
 ; CHECK-LABEL: @lowbitmask_casted_shift(
-; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i32
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[S]], 2147483647
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[X:%.*]] to i32
+; CHECK-NEXT:    [[R:%.*]] = lshr i32 [[TMP1]], 1
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = ashr i8 %x, 1
@@ -858,12 +857,44 @@ define i32 @lowbitmask_casted_shift(i8 %x) {
   ret i32 %r
 }
 
+; Negative test - mask constant is too big.
+
+define i32 @lowbitmask_casted_shift_wrong_mask1(i8 %x) {
+; CHECK-LABEL: @lowbitmask_casted_shift_wrong_mask1(
+; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i32
+; CHECK-NEXT:    [[R:%.*]] = and i32 [[S]], 2147483647
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %a = ashr i8 %x, 2
+  %s = sext i8 %a to i32
+  %r = and i32 %s, 2147483647 ; 0x7fffffff
+  ret i32 %r
+}
+
+; Negative test - mask constant is too small.
+
+define i32 @lowbitmask_casted_shift_wrong_mask2(i8 %x) {
+; CHECK-LABEL: @lowbitmask_casted_shift_wrong_mask2(
+; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i32
+; CHECK-NEXT:    [[R:%.*]] = and i32 [[S]], 536870911
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %a = ashr i8 %x, 2
+  %s = sext i8 %a to i32
+  %r = and i32 %s, 536870911  ; 0x1fffffff
+  ret i32 %r
+}
+
+; Extra use of shift is ok.
+
 define i32 @lowbitmask_casted_shift_use1(i8 %x) {
 ; CHECK-LABEL: @lowbitmask_casted_shift_use1(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 3
 ; CHECK-NEXT:    call void @use8(i8 [[A]])
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i32
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[S]], 536870911
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[X]] to i32
+; CHECK-NEXT:    [[R:%.*]] = lshr i32 [[TMP1]], 3
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = ashr i8 %x, 3
@@ -873,6 +904,8 @@ define i32 @lowbitmask_casted_shift_use1(i8 %x) {
   ret i32 %r
 }
 
+; Negative test - extra use of sext requires more instructions.
+
 define i32 @lowbitmask_casted_shift_use2(i8 %x) {
 ; CHECK-LABEL: @lowbitmask_casted_shift_use2(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 3
@@ -888,11 +921,12 @@ define i32 @lowbitmask_casted_shift_use2(i8 %x) {
   ret i32 %r
 }
 
+; Vectors/weird types are ok.
+
 define <2 x i59> @lowbitmask_casted_shift_vec_splat(<2 x i47> %x) {
 ; CHECK-LABEL: @lowbitmask_casted_shift_vec_splat(
-; CHECK-NEXT:    [[A:%.*]] = ashr <2 x i47> [[X:%.*]], <i47 5, i47 5>
-; CHECK-NEXT:    [[S:%.*]] = sext <2 x i47> [[A]] to <2 x i59>
-; CHECK-NEXT:    [[R:%.*]] = and <2 x i59> [[S]], <i59 18014398509481983, i59 18014398509481983>
+; CHECK-NEXT:    [[TMP1:%.*]] = sext <2 x i47> [[X:%.*]] to <2 x i59>
+; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i59> [[TMP1]], <i59 5, i59 5>
 ; CHECK-NEXT:    ret <2 x i59> [[R]]
 ;
   %a = ashr <2 x i47> %x, <i47 5, i47 5>
-- 
2.7.4