[InstCombine] generalize fold for mask-with-signbit-splat

[spatel]

Summary

(iN X s>> (N-1)) & Y --> (X < 0) ? Y : 0

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

I was looking at a missing abs() transform and found my way to this generalization of an existing fold that was added with D67799. As discussed in that review, we want to make sure codegen handles this difference well, and for all of the targets/types that I spot-checked, it looks good.

There's a difference on an existing test that shows a potentially unnecessary use limit on an icmp fold.

That fold is in InstCombinerImpl::foldICmpSubConstant(), and IIRC there was some back-and-forth on it and similar folds because it could cause analysis/passes (SCEV, LSR?) to miss optimizations.

So I think this patch makes things more consistent. But if the more specific transform was here as a work-around, we could leave it in as a special-case before we do the more general fold (...but that would be pretty hacky).

[spatel]

llvm/test/Transforms/InstCombine/sub-ashr-and-to-icmp-select.ll

134 ↗

(On Diff #378208)

This is the missing icmp fold - it checks hasOneUse(), but that is not strictly necessary.

[lebedev.ri]

Hm, looks fine to me.
I guess there's also the 'commutative' variant: https://alive2.llvm.org/ce/z/t8iHTy

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2072

[spatel]

In D111410#3051136, @lebedev.ri wrote:

Hm, looks fine to me.
I guess there's also the 'commutative' variant: https://alive2.llvm.org/ce/z/t8iHTy

Thanks! I'll create some tests and make that a follow-up.
I'll leave this up for a bit in case there are other comments about the existing fold.

[RKSimon]

LGTM

So I think this patch makes things more consistent. But if the more specific transform was here as a work-around, we could leave it in as a special-case before we do the more general fold (...but that would be pretty hacky).

Maybe remove the current fold as a followup commit to make bisection a little easier if somebody was depending upon it?

[spatel]

In D111410#3051834, @RKSimon wrote:

LGTM

So I think this patch makes things more consistent. But if the more specific transform was here as a work-around, we could leave it in as a special-case before we do the more general fold (...but that would be pretty hacky).

Maybe remove the current fold as a followup commit to make bisection a little easier if somebody was depending upon it?

Yes - that will make it clearer. Also, I found the existing 'or' sibling to this was added with D67800, so I'll do the same sequence with that one.

[spatel]

The sibling 'or' pattern to this 'and' patch uncovered many potential regressions, so I have been fixing those and referencing back to this review. I think I'm down to the last 1 or 2 regressions on that one. I'll go ahead and push this in 2 parts and see if that uncovers any regressions that are not immediately visible.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Instruction *InstCombinerImpl::visitAnd(BinaryOperator &I) {
   if (match(Op0, m_OneUse(m_SExt(m_Value(A)))) &&
       A->getType()->isIntOrIntVectorTy(1))
     return SelectInst::Create(A, Op1, Constant::getNullValue(Ty));
   if (match(Op1, m_OneUse(m_SExt(m_Value(A)))) &&
       A->getType()->isIntOrIntVectorTy(1))
     return SelectInst::Create(A, Op0, Constant::getNullValue(Ty));
 
   // and(ashr(subNSW(Y, X), ScalarSizeInBits(Y)-1), X) --> X s> Y ? X : 0.
-  if (match(&I, m_c_And(m_OneUse(m_AShr(
-                            m_NSWSub(m_Value(Y), m_Value(X)),
-                            m_SpecificInt(Ty->getScalarSizeInBits() - 1))),
+  // TODO: This is a specific case of the more general pattern below, so it
+  //       should be removed.
+  unsigned FullShift = Ty->getScalarSizeInBits() - 1;
+  if (match(&I, m_c_And(m_OneUse(m_AShr(m_NSWSub(m_Value(Y), m_Value(X)),
+                                        m_SpecificInt(FullShift))),
                         m_Deferred(X)))) {
     Value *NewICmpInst = Builder.CreateICmpSGT(X, Y);
     return SelectInst::Create(NewICmpInst, X, ConstantInt::getNullValue(Ty));
   }
 
+  // (iN X s>> (N-1)) & Y --> (X < 0) ? Y : 0
+  if (match(&I, m_c_And(m_OneUse(m_AShr(m_Value(X), m_SpecificInt(FullShift))),
+                        m_Value(Y)))) {
+    Constant *Zero = ConstantInt::getNullValue(Ty);
+    Value *Cmp = Builder.CreateICmpSLT(X, Zero, "isneg");
+    return SelectInst::Create(Cmp, Y, Zero);
+  }
+
   // (~x) & y  -->  ~(x | (~y))  iff that gets rid of inversions
   if (sinkNotIntoOtherHandOfAndOrOr(I))
     return &I;
 
   // An and recurrence w/loop invariant step is equivelent to (and start, step)
   PHINode *PN = nullptr;
   Value *Start = nullptr, *Step = nullptr;
   if (matchSimpleRecurrence(&I, PN, Start, Step) && DT.dominates(Step, PN))

llvm/test/Transforms/InstCombine/and.ll

 ;
   %B = add <2 x i8> %A, <i8 16, i8 4>
   %C = and <2 x i8> %B, <i8 16, i8 4>
   ret <2 x i8> %C
 }
 
 define i8 @ashr_bitwidth_mask(i8 %x, i8 %y) {
 ; CHECK-LABEL: @ashr_bitwidth_mask(
-; CHECK-NEXT:    [[SIGN:%.*]] = ashr i8 [[X:%.*]], 7
-; CHECK-NEXT:    [[NEG_OR_ZERO:%.*]] = and i8 [[SIGN]], [[Y:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i8 [[X:%.*]], 0
+; CHECK-NEXT:    [[NEG_OR_ZERO:%.*]] = select i1 [[ISNEG]], i8 [[Y:%.*]], i8 0
 ; CHECK-NEXT:    ret i8 [[NEG_OR_ZERO]]
 ;
   %sign = ashr i8 %x, 7
   %neg_or_zero = and i8 %sign, %y
   ret i8 %neg_or_zero
 }
 
 define <2 x i8> @ashr_bitwidth_mask_vec_commute(<2 x i8> %x, <2 x i8> %py) {
 ; CHECK-LABEL: @ashr_bitwidth_mask_vec_commute(
 ; CHECK-NEXT:    [[Y:%.*]] = mul <2 x i8> [[PY:%.*]], <i8 42, i8 2>
-; CHECK-NEXT:    [[SIGN:%.*]] = ashr <2 x i8> [[X:%.*]], <i8 7, i8 7>
-; CHECK-NEXT:    [[NEG_OR_ZERO:%.*]] = and <2 x i8> [[Y]], [[SIGN]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i8> [[X:%.*]], zeroinitializer
+; CHECK-NEXT:    [[NEG_OR_ZERO:%.*]] = select <2 x i1> [[ISNEG]], <2 x i8> [[Y]], <2 x i8> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i8> [[NEG_OR_ZERO]]
 ;
   %y = mul <2 x i8> %py, <i8 42, i8 2>      ; thwart complexity-based ordering
   %sign = ashr <2 x i8> %x, <i8 7, i8 7>
   %neg_or_zero = and <2 x i8> %y, %sign
   ret <2 x i8> %neg_or_zero
 }
 
+; negative test - extra use
+
 define i8 @ashr_bitwidth_mask_use(i8 %x, i8 %y) {
 ; CHECK-LABEL: @ashr_bitwidth_mask_use(
 ; CHECK-NEXT:    [[SIGN:%.*]] = ashr i8 [[X:%.*]], 7
 ; CHECK-NEXT:    call void @use8(i8 [[SIGN]])
 ; CHECK-NEXT:    [[R:%.*]] = and i8 [[SIGN]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %sign = ashr i8 %x, 7
   call void @use8(i8 %sign)
   %r = and i8 %sign, %y
   ret i8 %r
 }
 
+; negative test - wrong shift amount
+
 define i8 @ashr_not_bitwidth_mask(i8 %x, i8 %y) {
 ; CHECK-LABEL: @ashr_not_bitwidth_mask(
 ; CHECK-NEXT:    [[SIGN:%.*]] = ashr i8 [[X:%.*]], 6
 ; CHECK-NEXT:    [[R:%.*]] = and i8 [[SIGN]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %sign = ashr i8 %x, 6
   %r = and i8 %sign, %y
   ret i8 %r
 }
 
+; negative test - wrong shift opcode
+
 define i8 @lshr_bitwidth_mask(i8 %x, i8 %y) {
 ; CHECK-LABEL: @lshr_bitwidth_mask(
 ; CHECK-NEXT:    [[SIGN:%.*]] = lshr i8 [[X:%.*]], 7
 ; CHECK-NEXT:    [[R:%.*]] = and i8 [[SIGN]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %sign = lshr i8 %x, 7
   %r = and i8 %sign, %y
   ret i8 %r
 }

llvm/test/Transforms/InstCombine/icmp.ll

 ; CHECK-NEXT:    ret <2 x i1> undef
 ;
   %V = icmp eq <2 x i64> zeroinitializer, undef
   ret <2 x i1> %V
 }
 
 define i32 @test6(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test6(
-; CHECK-NEXT:    [[A_LOBIT_NEG:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[F:%.*]] = and i32 [[A_LOBIT_NEG]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[F:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[F]]
 ;
   %c = icmp sle i32 %a, -1
   %d = zext i1 %c to i32
   %e = sub i32 0, %d
   %f = and i32 %e, %b
   ret i32 %f
 }

llvm/test/Transforms/InstCombine/mul-inseltpoison.ll

;
   ret i32 %t
 }
 
 ; Use the sign-bit as a mask:
 ; (zext (A < 0)) * B --> (A >> 31) & B
 
 define i32 @test10(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test10(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %c = icmp slt i32 %a, 0
   %d = zext i1 %c to i32
   %e = mul i32 %d, %b
   ret i32 %e
 }
 
 define i32 @test11(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test11(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %c = icmp sle i32 %a, -1
   %d = zext i1 %c to i32
   %e = mul i32 %d, %b
   ret i32 %e
 }
 
 declare void @use32(i32)
 
 define i32 @test12(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test12(
 ; CHECK-NEXT:    [[A_LOBIT:%.*]] = lshr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    call void @use32(i32 [[A_LOBIT]])
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %c = icmp ugt i32 %a, 2147483647
   %d = zext i1 %c to i32
   %e = mul i32 %d, %b
   call void @use32(i32 %d)
   ret i32 %e
;
   %r = mul i32 %zy, %sx
   ret i32 %r
 }
 
 ; (A >>u 31) * B --> (A >>s 31) & B
 
 define i32 @signbit_mul(i32 %a, i32 %b) {
 ; CHECK-LABEL: @signbit_mul(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %d = lshr i32 %a, 31
   %e = mul i32 %d, %b
   ret i32 %e
 }
 
 define i32 @signbit_mul_commute_extra_use(i32 %a, i32 %b) {
 ; CHECK-LABEL: @signbit_mul_commute_extra_use(
 ; CHECK-NEXT:    [[D:%.*]] = lshr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    call void @use32(i32 [[D]])
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %d = lshr i32 %a, 31
   %e = mul i32 %b, %d
   call void @use32(i32 %d)
   ret i32 %e
 }
 
 ; (A >>u 31)) * B --> (A >>s 31) & B
 
 define <2 x i32> @signbit_mul_vec(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @signbit_mul_vec(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[A:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[E:%.*]] = and <2 x i32> [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i32> [[A:%.*]], zeroinitializer
+; CHECK-NEXT:    [[E:%.*]] = select <2 x i1> [[ISNEG]], <2 x i32> [[B:%.*]], <2 x i32> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i32> [[E]]
 ;
   %d = lshr <2 x i32> %a, <i32 31, i32 31>
   %e = mul <2 x i32> %d, %b
   ret <2 x i32> %e
 }
 
 define <2 x i32> @signbit_mul_vec_commute(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @signbit_mul_vec_commute(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[A:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[E:%.*]] = and <2 x i32> [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i32> [[A:%.*]], zeroinitializer
+; CHECK-NEXT:    [[E:%.*]] = select <2 x i1> [[ISNEG]], <2 x i32> [[B:%.*]], <2 x i32> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i32> [[E]]
 ;
   %d = lshr <2 x i32> %a, <i32 31, i32 31>
   %e = mul <2 x i32> %b, %d
   ret <2 x i32> %e
 }
 
 define i32 @test18(i32 %A, i32 %B) {

llvm/test/Transforms/InstCombine/mul.ll

;
   ret i32 %t
 }
 
 ; Use the sign-bit as a mask:
 ; (zext (A < 0)) * B --> (A >> 31) & B
 
 define i32 @test10(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test10(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %c = icmp slt i32 %a, 0
   %d = zext i1 %c to i32
   %e = mul i32 %d, %b
   ret i32 %e
 }
 
 define i32 @test11(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test11(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %c = icmp sle i32 %a, -1
   %d = zext i1 %c to i32
   %e = mul i32 %d, %b
   ret i32 %e
 }
 
 declare void @use32(i32)
 
 define i32 @test12(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test12(
 ; CHECK-NEXT:    [[A_LOBIT:%.*]] = lshr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    call void @use32(i32 [[A_LOBIT]])
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %c = icmp ugt i32 %a, 2147483647
   %d = zext i1 %c to i32
   %e = mul i32 %d, %b
   call void @use32(i32 %d)
   ret i32 %e
;
   %sx = sext i1 %x to i32
   call void @use32(i32 %sx)
   %zy = zext i1 %y to i32
   call void @use32(i32 %zy)
   %r = mul i32 %zy, %sx
   ret i32 %r
 }
 
-; (A >>u 31) * B --> (A >>s 31) & B
+; (A >>u 31) * B --> (A >>s 31) & B --> A < 0 ? B : 0
 
 define i32 @signbit_mul(i32 %a, i32 %b) {
 ; CHECK-LABEL: @signbit_mul(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %d = lshr i32 %a, 31
   %e = mul i32 %d, %b
   ret i32 %e
 }
 
 define i32 @signbit_mul_commute_extra_use(i32 %a, i32 %b) {
 ; CHECK-LABEL: @signbit_mul_commute_extra_use(
 ; CHECK-NEXT:    [[D:%.*]] = lshr i32 [[A:%.*]], 31
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[A]], 31
-; CHECK-NEXT:    [[E:%.*]] = and i32 [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[A]], 0
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[ISNEG]], i32 [[B:%.*]], i32 0
 ; CHECK-NEXT:    call void @use32(i32 [[D]])
 ; CHECK-NEXT:    ret i32 [[E]]
 ;
   %d = lshr i32 %a, 31
   %e = mul i32 %b, %d
   call void @use32(i32 %d)
   ret i32 %e
 }
 
-; (A >>u 31)) * B --> (A >>s 31) & B
+; (A >>u 31)) * B --> (A >>s 31) & B --> A < 0 ? B : 0
 
 define <2 x i32> @signbit_mul_vec(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @signbit_mul_vec(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[A:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[E:%.*]] = and <2 x i32> [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i32> [[A:%.*]], zeroinitializer
+; CHECK-NEXT:    [[E:%.*]] = select <2 x i1> [[ISNEG]], <2 x i32> [[B:%.*]], <2 x i32> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i32> [[E]]
 ;
   %d = lshr <2 x i32> %a, <i32 31, i32 31>
   %e = mul <2 x i32> %d, %b
   ret <2 x i32> %e
 }
 
 define <2 x i32> @signbit_mul_vec_commute(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @signbit_mul_vec_commute(
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[A:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[E:%.*]] = and <2 x i32> [[TMP1]], [[B:%.*]]
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i32> [[A:%.*]], zeroinitializer
+; CHECK-NEXT:    [[E:%.*]] = select <2 x i1> [[ISNEG]], <2 x i32> [[B:%.*]], <2 x i32> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i32> [[E]]
 ;
   %d = lshr <2 x i32> %a, <i32 31, i32 31>
   %e = mul <2 x i32> %b, %d
   ret <2 x i32> %e
 }
 
 define i32 @test18(i32 %A, i32 %B) {