[InstCombine] allow sext in fold of mask using signbit, part 2

https://alive2.llvm.org/ce/z/rcbZmx

Sibling tranform to 275aa24

This pattern is seen in the examples in issue #57381.

Author: rotateright

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -2244,12 +2244,11 @@ Instruction *InstCombinerImpl::visitAnd(BinaryOperator &I) {
       A->getType()->isIntOrIntVectorTy(1))
     return SelectInst::Create(A, Op0, Constant::getNullValue(Ty));
 
-  // (iN X s>> (N-1)) & Y --> (X s< 0) ? Y : 0
-  // TODO: Peek through optional sext.
-  unsigned FullShift = Ty->getScalarSizeInBits() - 1;
-  if (match(&I, m_c_And(m_OneUse(m_AShr(m_Value(X),
-                                        m_SpecificIntAllowUndef(FullShift))),
-                        m_Value(Y)))) {
+  // (iN X s>> (N-1)) & Y --> (X s< 0) ? Y : 0 -- with optional sext
+  if (match(&I, m_c_And(m_OneUse(m_SExtOrSelf(
+                            m_AShr(m_Value(X), m_APIntAllowUndef(C)))),
+                        m_Value(Y))) &&
+      *C == X->getType()->getScalarSizeInBits() - 1) {
     Value *IsNeg = Builder.CreateIsNeg(X, "isneg");
     return SelectInst::Create(IsNeg, Y, ConstantInt::getNullValue(Ty));
   }

llvm/test/Transforms/InstCombine/and.ll

@@ -1682,9 +1682,8 @@ define i8 @lshr_bitwidth_mask(i8 %x, i8 %y) {
 
 define i16 @signbit_splat_mask(i8 %x, i16 %y) {
 ; CHECK-LABEL: @signbit_splat_mask(
-; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 7
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i16
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[S]], [[Y:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i8 [[X:%.*]], 0
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISNEG]], i16 [[Y:%.*]], i16 0
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %a = ashr i8 %x, 7
@@ -1696,9 +1695,8 @@ define i16 @signbit_splat_mask(i8 %x, i16 %y) {
 define <2 x i16> @signbit_splat_mask_commute(<2 x i5> %x, <2 x i16> %p) {
 ; CHECK-LABEL: @signbit_splat_mask_commute(
 ; CHECK-NEXT:    [[Y:%.*]] = mul <2 x i16> [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[A:%.*]] = ashr <2 x i5> [[X:%.*]], <i5 4, i5 poison>
-; CHECK-NEXT:    [[S:%.*]] = sext <2 x i5> [[A]] to <2 x i16>
-; CHECK-NEXT:    [[R:%.*]] = and <2 x i16> [[Y]], [[S]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i5> [[X:%.*]], zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[ISNEG]], <2 x i16> [[Y]], <2 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i16> [[R]]
 ;
   %y = mul <2 x i16> %p, %p ; thwart complexity-based canonicalization
@@ -1712,8 +1710,8 @@ define i16 @signbit_splat_mask_use1(i8 %x, i16 %y) {
 ; CHECK-LABEL: @signbit_splat_mask_use1(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 7
 ; CHECK-NEXT:    call void @use8(i8 [[A]])
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i16
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[S]], [[Y:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i8 [[X]], 0
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISNEG]], i16 [[Y:%.*]], i16 0
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %a = ashr i8 %x, 7
@@ -1723,6 +1721,8 @@ define i16 @signbit_splat_mask_use1(i8 %x, i16 %y) {
   ret i16 %r
 }
 
+; negative test - extra use
+
 define i16 @signbit_splat_mask_use2(i8 %x, i16 %y) {
 ; CHECK-LABEL: @signbit_splat_mask_use2(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 7
@@ -1738,6 +1738,8 @@ define i16 @signbit_splat_mask_use2(i8 %x, i16 %y) {
   ret i16 %r
 }
 
+; negative test - wrong extend
+
 define i16 @not_signbit_splat_mask1(i8 %x, i16 %y) {
 ; CHECK-LABEL: @not_signbit_splat_mask1(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 7
@@ -1751,6 +1753,8 @@ define i16 @not_signbit_splat_mask1(i8 %x, i16 %y) {
   ret i16 %r
 }
 
+; negative test - wrong shift amount
+
 define i16 @not_signbit_splat_mask2(i8 %x, i16 %y) {
 ; CHECK-LABEL: @not_signbit_splat_mask2(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 6
@@ -1915,6 +1919,8 @@ define i16 @invert_signbit_splat_mask_use2(i8 %x, i16 %y) {
   ret i16 %r
 }
 
+; negative test - extra use
+
 define i16 @invert_signbit_splat_mask_use3(i8 %x, i16 %y) {
 ; CHECK-LABEL: @invert_signbit_splat_mask_use3(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 7
@@ -1932,6 +1938,8 @@ define i16 @invert_signbit_splat_mask_use3(i8 %x, i16 %y) {
   ret i16 %r
 }
 
+; negative test - wrong extend
+
 define i16 @not_invert_signbit_splat_mask1(i8 %x, i16 %y) {
 ; CHECK-LABEL: @not_invert_signbit_splat_mask1(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 7
@@ -1947,6 +1955,8 @@ define i16 @not_invert_signbit_splat_mask1(i8 %x, i16 %y) {
   ret i16 %r
 }
 
+; negative test - wrong shift amount
+
 define i16 @not_invert_signbit_splat_mask2(i8 %x, i16 %y) {
 ; CHECK-LABEL: @not_invert_signbit_splat_mask2(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 6

llvm/test/Transforms/InstCombine/negated-bitmask.ll

@@ -278,9 +278,8 @@ define i32 @neg_mask(i32 %x, i16 %y) {
 ; CHECK-LABEL: @neg_mask(
 ; CHECK-NEXT:    [[S:%.*]] = sext i16 [[Y:%.*]] to i32
 ; CHECK-NEXT:    [[SUB1:%.*]] = sub nsw i32 [[X:%.*]], [[S]]
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i16 [[Y]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = sext i16 [[TMP1]] to i32
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[SUB1]], [[TMP2]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i16 [[Y]], 0
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISNEG]], i32 [[SUB1]], i32 0
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = sext i16 %y to i32
@@ -296,9 +295,8 @@ define i32 @neg_mask_const(i16 %x) {
 ; CHECK-LABEL: @neg_mask_const(
 ; CHECK-NEXT:    [[S:%.*]] = sext i16 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[SUB1:%.*]] = sub nsw i32 1000, [[S]]
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i16 [[X]], 15
-; CHECK-NEXT:    [[TMP2:%.*]] = sext i16 [[TMP1]] to i32
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[SUB1]], [[TMP2]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i16 [[X]], 0
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISNEG]], i32 [[SUB1]], i32 0
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = sext i16 %x to i32