[InstSimplify] enhance simplifyWithOpReplaced() to allow more 'select' removal

[spatel]

Summary

This is a generalization of a suggestion from issue #60799 that allows removing a redundant guard of an input value via icmp+select. It should also solve issue #60801.

This only comes into play for a select with an equality condition where we are trying to substitute a constant into the false arm of a select. (A 'true' select arm substitution allows "refinement", so it is not on this code path.)

The constant must be the same in the compare and the select, and it must be a "binop absorber" (X op C = C). That query currently includes 'or', 'and', and 'mul', so there are tests for all of those opcodes.

We then use "impliesPoison" on the false arm binop and the original "Op" to be replaced to ensure that the select is not actually blocking poison from leaking. That could be potentially expensive as we recursively test each operand, but it is currently limited to a depth of 2. That's enough to catch our motivating cases, but probably nothing more complicated (although that seems unlikely).

I don't know how to generalize a proof for Alive2 for this, but here's a positive and negative test example to help illustrate the subtle logic differences of poison propagation:
https://alive2.llvm.org/ce/z/Sz5K-c

[nikic]

I like the approach, but I'm not sure you're checking the right poison implications. My first thought would be that we need impliesPoison(V, Op). We are replacing the select with V, so what we want to establish is that if V is poison, the select would be poison anyway because Op is also implied poison.

[spatel]

In D144493#4141817, @nikic wrote:

I like the approach, but I'm not sure you're checking the right poison implications. My first thought would be that we need impliesPoison(V, Op). We are replacing the select with V, so what we want to establish is that if V is poison, the select would be poison anyway because Op is also implied poison.

Yes, that is clearer and seems to work. I'm not finding a test diff given the other constraints in place...
Also, I initially misread the recursive depth check in impliesPoison() - we can look past the binop itself and each of its operands.

[spatel]

Patch updated:
Reduced the impliesPoison() logic as suggested.

[nikic]

LGTM

llvm/lib/Analysis/InstructionSimplify.cpp

if (auto *BO = dyn_cast<BinaryOperator>(I)) {
       if (NewOps[1] == ConstantExpr::getBinOpIdentity(Opcode, I->getType(),
                                                       /* RHS */ true))
         return NewOps[0];
 
       // x & x -> x, x | x -> x
       if ((Opcode == Instruction::And || Opcode == Instruction::Or) &&
           NewOps[0] == NewOps[1])
         return NewOps[0];
+
+      // If we are substituting an absorber constant into a binop and extra
+      // poison can't leak if we remove the select -- because both operands of
+      // the binop are based on the same value -- then it may be safe to replace
+      // the value with the absorber constant. Examples:
+      // (Op == 0) ? 0 : (Op & -Op)            --> Op & -Op
+      // (Op == 0) ? 0 : (Op * (binop Op, C))  --> Op * (binop Op, C)
+      // (Op == -1) ? -1 : (Op | (binop C, Op) --> Op | (binop C, Op)
+      if (RepOp == ConstantExpr::getBinOpAbsorber(Opcode, I->getType()) &&
+          impliesPoison(BO, Op))
+        return RepOp;
     }
 
     if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
       // getelementptr x, 0 -> x
       if (NewOps.size() == 2 && match(NewOps[1], m_Zero()) &&
           !GEP->isInBounds())
         return NewOps[0];
     }

llvm/test/Transforms/InstSimplify/select.ll

 define <2 x i32> @poison4(<2 x i1> %cond, <2 x i32> %x) {
 ; CHECK-LABEL: @poison4(
 ; CHECK-NEXT:    ret <2 x i32> [[X:%.*]]
 ;
   %v = select <2 x i1> %cond, <2 x i32> %x, <2 x i32> poison
   ret <2 x i32> %v
 }
 
+; 0 is the absorber constant for 'and'.
+; The 'select' can't block extra poison because both sides of 'and' have 'x' operand.
+
 define i8 @replace_false_op_eq_neg_and(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_eq_neg_and(
-; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
+; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[NEG]], [[X]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[AND]]
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[AND]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %neg = sub i8 0, %x
   %and = and i8 %neg, %x
   %sel = select i1 %eq0, i8 0, i8 %and
   ret i8 %sel
 }
 
+; same as above, but commute 'and'
+
 define i8 @replace_false_op_eq_neg_and_commute(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_eq_neg_and_commute(
-; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
+; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], [[NEG]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[AND]]
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[AND]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %neg = sub i8 0, %x
   %and = and i8 %x, %neg
   %sel = select i1 %eq0, i8 0, i8 %and
   ret i8 %sel
 }
 
+; same as above, but swap 'select'
+
 define i8 @replace_false_op_ne_neg_and(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_ne_neg_and(
-; CHECK-NEXT:    [[NE0:%.*]] = icmp ne i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
+; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[NEG]], [[X]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[NE0]], i8 [[AND]], i8 0
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[AND]]
 ;
   %ne0 = icmp ne i8 %x, 0
   %neg = sub i8 0, %x
   %and = and i8 %neg, %x
   %sel = select i1 %ne0, i8 %and, i8 0
   ret i8 %sel
 }
 
+; same as above, but commute 'and' and swap 'select'
+
 define i8 @replace_false_op_ne_neg_and_commute(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_ne_neg_and_commute(
-; CHECK-NEXT:    [[NE0:%.*]] = icmp ne i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
+; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], [[NEG]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[NE0]], i8 [[AND]], i8 0
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[AND]]
 ;
   %ne0 = icmp ne i8 %x, 0
   %neg = sub i8 0, %x
   %and = and i8 %x, %neg
   %sel = select i1 %ne0, i8 %and, i8 0
   ret i8 %sel
 }
 
+; the first binop can be anything as long as it has the common operand
+
 define i8 @replace_false_op_eq_dec_and(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_eq_dec_and(
-; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[DEC:%.*]] = add i8 [[X]], -1
+; CHECK-NEXT:    [[DEC:%.*]] = add i8 [[X:%.*]], -1
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[DEC]], [[X]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[AND]]
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[AND]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %dec = add i8 %x, -1
   %and = and i8 %dec, %x
   %sel = select i1 %eq0, i8 0, i8 %and
   ret i8 %sel
 }
 
+; mul has the same absorber constant - "0"
+
 define i8 @replace_false_op_eq_add_mul(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_eq_add_mul(
-; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 42
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X:%.*]], 42
 ; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[ADD]], [[X]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[MUL]]
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[MUL]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %add = add i8 %x, 42
   %mul = mul i8 %add, %x
   %sel = select i1 %eq0, i8 0, i8 %mul
   ret i8 %sel
 }
 
+; or has a different absorber constant = "-1"
+
 define i8 @replace_false_op_eq_shl_or(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_eq_shl_or(
-; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], -1
-; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[X]], 3
+; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[X:%.*]], 3
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], [[SHL]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 -1, i8 [[OR]]
-; CHECK-NEXT:    ret i8 [[SEL]]
+; CHECK-NEXT:    ret i8 [[OR]]
 ;
   %eq0 = icmp eq i8 %x, -1
   %shl = shl i8 %x, 3
   %or = or i8 %x, %shl
   %sel = select i1 %eq0, i8 -1, i8 %or
   ret i8 %sel
 }
 
+; negative test - wrong cmp predicate
+
 define i8 @replace_false_op_sgt_neg_and(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_sgt_neg_and(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp sgt i8 [[X:%.*]], 0
 ; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[NEG]], [[X]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[AND]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq0 = icmp sgt i8 %x, 0
   %neg = sub i8 0, %x
   %and = and i8 %neg, %x
   %sel = select i1 %eq0, i8 0, i8 %and
   ret i8 %sel
 }
 
+; negative test - the binop must use a compare operand
+
 define i8 @replace_false_op_eq_shl_or_wrong_cmp_op(i8 %x, i8 %y) {
 ; CHECK-LABEL: @replace_false_op_eq_shl_or_wrong_cmp_op(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[Y:%.*]], -1
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[X:%.*]], 3
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], [[SHL]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 -1, i8 [[OR]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq0 = icmp eq i8 %y, -1
   %shl = shl i8 %x, 3
   %or = or i8 %x, %shl
   %sel = select i1 %eq0, i8 -1, i8 %or
   ret i8 %sel
 }
 
+; negative test - can't have extra source of potential poison
+
 define i8 @replace_false_op_eq_neg_and_leak1(i8 %x, i8 %y) {
 ; CHECK-LABEL: @replace_false_op_eq_neg_and_leak1(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
 ; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[Y:%.*]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[NEG]], [[X]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[AND]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %neg = sub i8 0, %y
   %and = and i8 %neg, %x
   %sel = select i1 %eq0, i8 0, i8 %and
   ret i8 %sel
 }
 
+; negative test - can't have extra source of potential poison
+
 define i8 @replace_false_op_eq_neg_and_leak2(i8 %x, i8 %y) {
 ; CHECK-LABEL: @replace_false_op_eq_neg_and_leak2(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
 ; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[NEG]], [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[AND]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %neg = sub i8 0, %x
   %and = and i8 %neg, %y
   %sel = select i1 %eq0, i8 0, i8 %and
   ret i8 %sel
 }
 
+; negative test - can't have extra source of potential poison
+
 define i8 @replace_false_op_eq_add_mul_leak3(i8 %x, i8 %y) {
 ; CHECK-LABEL: @replace_false_op_eq_add_mul_leak3(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], 0
 ; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], [[Y:%.*]]
 ; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[ADD]], [[X]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 0, i8 [[MUL]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq0 = icmp eq i8 %x, 0
   %add = add i8 %x, %y
   %mul = mul i8 %add, %x
   %sel = select i1 %eq0, i8 0, i8 %mul
   ret i8 %sel
 }
 
+; negative test - can't have extra source of potential poison
+
 define i8 @replace_false_op_eq_shl_or_leak4(i8 %x, i8 %y) {
 ; CHECK-LABEL: @replace_false_op_eq_shl_or_leak4(
 ; CHECK-NEXT:    [[EQ0:%.*]] = icmp eq i8 [[X:%.*]], -1
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[Y:%.*]], [[X]]
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], [[SHL]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ0]], i8 -1, i8 [[OR]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq0 = icmp eq i8 %x, -1
   %shl = shl i8 %y, %x
   %or = or i8 %x, %shl
   %sel = select i1 %eq0, i8 -1, i8 %or
   ret i8 %sel
 }
 
+; negative test - wrong cmp constant
+
 define i8 @replace_false_op_eq42_neg_and(i8 %x) {
 ; CHECK-LABEL: @replace_false_op_eq42_neg_and(
 ; CHECK-NEXT:    [[EQ42:%.*]] = icmp eq i8 [[X:%.*]], 42
 ; CHECK-NEXT:    [[NEG:%.*]] = sub i8 0, [[X]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[NEG]], [[X]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[EQ42]], i8 0, i8 [[AND]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %eq42 = icmp eq i8 %x, 42
   %neg = sub i8 0, %x
   %and = and i8 %neg, %x
   %sel = select i1 %eq42, i8 0, i8 %and
   ret i8 %sel
 }