[InstCombine] Add transforms for `(or/and (icmp eq/ne X,0),(icmp eq/ne X,Pow2OrZero))`

[goldstein.w.n]

Summary

(or (icmp eq X, 0), (icmp eq X, Pow2OrZero))

--> `(icmp eq (and X, Pow2OrZero), X)`

(and (icmp ne X, 0), (icmp ne X, Pow2OrZero))

--> `(icmp ne (and X, Pow2OrZero), X)`

Proofs: https://alive2.llvm.org/ce/z/nPo2BN

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

764

LHS and RHS are commutative . Do we need to consider it?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

764

Above we swap LHS/RHS so that LHS is the comparison against zero: L757.

If you are refering to not using m_c_ICmp here then its also no issue b.c of canonicalization.

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

743

nit: Pow2 -> Pow2OrZero

764

oh, I missed the swap, thanks.

[nikic]

Not sure we should do this transform in IR. The original form seems to be a lot simpler to reason about than the new one -- e.g. if you have that as condition, it's obvious that %x can't be zero (as there's an explicit check for it). Afterwards, that would be fairly hard to determine.

[goldstein.w.n]

In D157312#4616620, @nikic wrote:

Not sure we should do this transform in IR. The original form seems to be a lot simpler to reason about than the new one -- e.g. if you have that as condition, it's obvious that %x can't be zero (as there's an explicit check for it). Afterwards, that would be fairly hard to determine.

Really? We have a fair amount of logic for X & Y eq/ne ... and we do this fold already for constants (and similiar folds for non-constants).
Also since its an equivilency and only relies on Pow2OrZero being a pow2orzero, theres no information loss.

[nikic]

In D157312#4618198, @goldstein.w.n wrote:

In D157312#4616620, @nikic wrote:

Not sure we should do this transform in IR. The original form seems to be a lot simpler to reason about than the new one -- e.g. if you have that as condition, it's obvious that %x can't be zero (as there's an explicit check for it). Afterwards, that would be fairly hard to determine.

Really? We have a fair amount of logic for X & Y eq/ne ... and we do this fold already for constants (and similiar folds for non-constants).
Also since its an equivilency and only relies on Pow2OrZero being a pow2orzero, theres no information loss.

For constants this is turned into an X & ~C eq 0 or X & ~C ne 0 pattern. That is well understood due to the obvious KnownBits implications. X & Y == X not so much I think.

[goldstein.w.n]

In D157312#4618244, @nikic wrote:

In D157312#4618198, @goldstein.w.n wrote:

In D157312#4616620, @nikic wrote:

Not sure we should do this transform in IR. The original form seems to be a lot simpler to reason about than the new one -- e.g. if you have that as condition, it's obvious that %x can't be zero (as there's an explicit check for it). Afterwards, that would be fairly hard to determine.

Really? We have a fair amount of logic for X & Y eq/ne ... and we do this fold already for constants (and similiar folds for non-constants).
Also since its an equivilency and only relies on Pow2OrZero being a pow2orzero, theres no information loss.

For constants this is turned into an X & ~C eq 0 or X & ~C ne 0 pattern. That is well understood due to the obvious KnownBits implications. X & Y == X not so much I think.

How about if we get D145424, D145425, and D145426 in as well?
Otherwise I'll move to the DAG, but since we split branches in SelectionDAGLowering folds like this lose alot of there effectiveness.

[nikic]

Okay, let's give this a try.

llvm/test/Transforms/InstCombine/and-or-icmps.ll

3135

Please also add tests with logical and/or that don't fold (with pow2 on the right).

I thin kwe'd also missing a test where the comparison is not against zero but another constant.

[goldstein.w.n]

Rebase (ontop of a lot of new stuff and AndXX Support)

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

751

out op?

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[goldstein.w.n]

Abandoning and resubmitting on GH

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

if (!Known.isNonNegative())
     return nullptr;
 
   if (Inverted)
     NewPred = ICmpInst::getInversePredicate(NewPred);
 
   return Builder.CreateICmp(NewPred, Input, RangeEnd);
 }
 
+// (or (icmp eq X, 0), (icmp eq X, Pow2OrZero))
+//      -> (icmp eq (and X, Pow2), X)
+// (and (icmp ne X, 0), (icmp ne X, Pow2OrZero))
+//      -> (icmp ne (and X, Pow2OrZero), X)
+static Value *
+foldAndOrOfICmpsWithPow2AndWithZero(InstCombiner::BuilderTy &Builder,
+                                    ICmpInst *LHS, ICmpInst *RHS, bool IsAnd,
+                                    const SimplifyQuery &Q) {
+  CmpInst::Predicate Pred = IsAnd ? CmpInst::ICMP_NE : CmpInst::ICMP_EQ;
+  // Make sure we have right compares for out op.
+  if (LHS->getPredicate() != Pred || RHS->getPredicate() != Pred)
+    return nullptr;
+
+  // Make it so we can match LHS against the (icmp eq/ne X, 0) just for
+  // simplicity.
+  if (match(RHS->getOperand(1), m_Zero()))
+    std::swap(LHS, RHS);
+
+  Value *Pow2, *Op;
+  // Match the desired pattern:
+  // LHS: (icmp eq/ne X, 0)
+  // RHS: (icmp eq/ne X, Pow2OrZero)
+  if (!match(LHS, m_ICmp(Pred, m_Value(Op), m_Zero())) ||
+      !match(RHS, m_c_ICmp(Pred, m_Specific(Op), m_Value(Pow2))) ||
+      !isKnownToBeAPowerOfTwo(Pow2, Q.DL, /*OrZero*/ true, /*Depth*/ 0, Q.AC,
+                              Q.CxtI, Q.DT))
+    return nullptr;
+
+  Value *And = Builder.CreateAnd(Op, Pow2);
+  return Builder.CreateICmp(Pred, And, Op);
+}
+
 // Fold (iszero(A & K1) | iszero(A & K2)) -> (A & (K1 | K2)) != (K1 | K2)
 // Fold (!iszero(A & K1) & !iszero(A & K2)) -> (A & (K1 | K2)) == (K1 | K2)
 Value *InstCombinerImpl::foldAndOrOfICmpsOfAndWithPow2(ICmpInst *LHS,
                                                        ICmpInst *RHS,
                                                        Instruction *CxtI,
                                                        bool IsAnd,
                                                        bool IsLogical) {
   CmpInst::Predicate Pred = IsAnd ? CmpInst::ICMP_NE : CmpInst::ICMP_EQ;
Value *InstCombinerImpl::foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
   const SimplifyQuery Q = SQ.getWithInstruction(&I);
 
   // Fold (iszero(A & K1) | iszero(A & K2)) ->  (A & (K1 | K2)) != (K1 | K2)
   // Fold (!iszero(A & K1) & !iszero(A & K2)) ->  (A & (K1 | K2)) == (K1 | K2)
   // if K1 and K2 are a one-bit mask.
   if (Value *V = foldAndOrOfICmpsOfAndWithPow2(LHS, RHS, &I, IsAnd, IsLogical))
     return V;
 
+  if (!IsLogical)
+    if (Value *V =
+            foldAndOrOfICmpsWithPow2AndWithZero(Builder, LHS, RHS, IsAnd, Q))
+      return V;
+
   ICmpInst::Predicate PredL = LHS->getPredicate(), PredR = RHS->getPredicate();
   Value *LHS0 = LHS->getOperand(0), *RHS0 = RHS->getOperand(0);
   Value *LHS1 = LHS->getOperand(1), *RHS1 = RHS->getOperand(1);
   const APInt *LHSC = nullptr, *RHSC = nullptr;
   match(LHS1, m_APInt(LHSC));
   match(RHS1, m_APInt(RHSC));
 
   // (icmp1 A, B) | (icmp2 A, B) --> (icmp3 A, B)

llvm/test/Transforms/InstCombine/and-or-icmps.ll

;
   %D = or i32 %C, %B
   ret i32 %D
 }
 
 define i1 @icmp_eq_or_z_or_pow2orz(i8 %x, i8 %y) {
 ; CHECK-LABEL: @icmp_eq_or_z_or_pow2orz(
 ; CHECK-NEXT:    [[NY:%.*]] = sub i8 0, [[Y:%.*]]
 ; CHECK-NEXT:    [[POW2ORZ:%.*]] = and i8 [[NY]], [[Y]]
-; CHECK-NEXT:    [[C0:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[CP2:%.*]] = icmp eq i8 [[POW2ORZ]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = or i1 [[C0]], [[CP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[POW2ORZ]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i8 [[TMP1]], [[X]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %ny = sub i8 0, %y
   %pow2orz = and i8 %ny, %y
 
   %c0 = icmp eq i8 %x, 0
   %cp2 = icmp eq i8 %x, %pow2orz
   %r = or i1 %c0, %cp2
   ret i1 %r
 }
 
 define <2 x i1> @icmp_ne_and_z_and_pow2orz(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @icmp_ne_and_z_and_pow2orz(
 ; CHECK-NEXT:    [[NY:%.*]] = sub <2 x i8> zeroinitializer, [[Y:%.*]]
 ; CHECK-NEXT:    [[POW2ORZ:%.*]] = and <2 x i8> [[NY]], [[Y]]
-; CHECK-NEXT:    [[C0:%.*]] = icmp ne <2 x i8> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[CP2:%.*]] = icmp ne <2 x i8> [[POW2ORZ]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = and <2 x i1> [[C0]], [[CP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i8> [[POW2ORZ]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = icmp ne <2 x i8> [[TMP1]], [[X]]
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %ny = sub <2 x i8> zeroinitializer, %y
   %pow2orz = and <2 x i8> %ny, %y
 
   %c0 = icmp ne <2 x i8> %x, zeroinitializer
   %cp2 = icmp ne <2 x i8> %x, %pow2orz
   %r = and <2 x i1> %c0, %cp2

 ;
   %ny = sub i8 0, %y
   %pow2orz = and i8 %ny, %y
 
   %c0 = icmp sle i8 %x, 0
   %cp2 = icmp sle i8 %x, %pow2orz
   %r = or i1 %c0, %cp2
   ret i1 %r
 }