[InstCombine] Fold Select with binary op

[xbolva00]

Summary

Fold

%A = icmp eq i8 %x, 0
%B = xor i8 %x, %z
%C = select i1 %A, i8 %B, i8 %y

To

%C = select i1 %A, i8 %z, i8 %y

Fixes https://bugs.llvm.org/show_bug.cgi?id=38345
Proof: https://rise4fun.com/Alive/43J

[xbolva00]

test/Transforms/InstCombine/select-xor-icmp.ll

55 ↗

(On Diff #157841)

Hm I will remove this test, since other transformation matches it.

[xbolva00]

test/Transforms/InstCombine/select-xor-icmp.ll

55 ↗

(On Diff #157841)

Oh, well, just another variant of this fold.

https://rise4fun.com/Alive/bKj

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineSelect.cpp

68

Constants are always canonicalized to RHS, so no commutativity in here.
https://rise4fun.com/Alive/r7Q

1740

I'm not sure there is much point in not passing the SelectInst &SI itself.
It's certainly a bit harder to read foldSelectInstWithBinaryOp() itself if it does lot's of matching on a number of arguments
(as opposed to one huge match() on one single function argument)

test/Transforms/InstCombine/select-xor-icmp.ll

24 ↗

(On Diff #157841)

Since there is a constant, please also add a test with <i8 0, i8 undef, i8 0>

36 ↗

(On Diff #157841)

This never happens

62 ↗

(On Diff #157841)

Please also add one test with non-zero

[xbolva00]

• Tests

[spatel]

This was noted as an inline comment on the test file: this transform is not xor-specific. Example:
https://rise4fun.com/Alive/R4P

I think it works for any binop with an "identity constant", so I'd prefer to see this patch generalized to use ConstantExpr::getBinOpIdentity().
That said, these cmp+select transforms are getting into the gray area of equivalence substitutions that are better handled by GVN.
At the least, I'd like to know if we can get any stats on how often this happens (collect transform stats using test-suite or some other well-known benchmarks).

[xbolva00]

I picked it from the list of missed optimizations from here: https://blog.regehr.org/archives/1192

[spatel]

In D49954#1179378, @xbolva00 wrote:

I picked it from the list of missed optimizations from here: https://blog.regehr.org/archives/1192

Ok, that's probably still a good-enough motivation. But let's generalize it, so we won't have to deal with every binop and cmp variation individually.

[spatel]

test/Transforms/InstCombine/select-binop-icmp.ll

18–26

That's a miscompile:
https://rise4fun.com/Alive/9SC

The identity constant needs to match the compare constant.

[xbolva00]

• Fixed miscompilation, added tests

[xbolva00]

Can I precommit tests?

[spatel]

In D49954#1179406, @xbolva00 wrote:

Can I precommit tests?

Yes, but please make sure we have at least 1 positive test for each potential opcode that is folded. If you want to leave the non-commutative and FP opcodes out of the initial patch, that's fine, but it should be marked with a TODO comment (and we'll eventually need tests for those). We'll also need tests for the 'ne' predicate.

test/Transforms/InstCombine/select-binop-icmp.ll

4

Unnecessary - the transform is target-independent.

[xbolva00]

• Rebased

[spatel]

lib/Transforms/InstCombine/InstCombineSelect.cpp

64

Add another TODO for the 'ne' predicate.

66

We don't need the builder here. In fact, we don't even need to create a new instruction. Just replace the true value of the select with 'Z' and return &Sel to indicate that it was morphed.

This preserves metadata that this patch would unknowingly drop (please add a test for that). There are examples of tests with select transforms that have metadata in:
test/Transforms/InstCombine/select_meta.ll

71

'C' should be declared and matched as a Constant type.

test/Transforms/InstCombine/select-binop-icmp.ll

174–175

This case should fold, right? Please add a TODO comment here and in the source if that's correct.

[xbolva00]

test/Transforms/InstCombine/select-binop-icmp.ll

174–175

With undef probably not.

[xbolva00]

• Added support for ne variant
• Added additional tests

[xbolva00]

Is this ok now?

[spatel]

lib/Transforms/InstCombine/InstCombineSelect.cpp

66

foldSelectBinOpIdentity() would be a more descriptive name.

74–75

I'd like to avoid the code duplication if it doesn't make it too messy. How about exiting if the compare doesn't match what we need:

if (!match(Cond, m_ICmp(Pred, m_Value(X), m_Constant(C))) ||
    !ICmpInst::isEquality(Pred))
 return nullptr;

Then choose the select operand based on whether:
bool IsEq = Pred == ICmpInst::ICMP_EQ;

test/Transforms/InstCombine/select-binop-icmp.ll

75–76

Is the icmp predicate getting canonicalized to 'eq' before we reach this transform?

Add an extra use of the icmp with something like:
call @use(i1 %A)
...to make sure we have coverage for the 'ne' side of the transform.
(and commit the test diffs before this patch, no pre-commit review needed on more tests)

[spatel]

lib/Transforms/InstCombine/InstCombineSelect.cpp

66

Remove unused Builder param.

[xbolva00]

• Addressed notes

[spatel]

lib/Transforms/InstCombine/InstCombineSelect.cpp

78–82

Can we eliminate this duplication now? So something like:

auto *BO = dyn_cast<BinaryOperator>(
    IsEq ? Sel.getTrueValue() : Sel.getFalseValue());
if (BO && match(BO, ...) && 
    ConstantExpr::getBinOpIdentity(BO, Ty) == C) {
  Sel.setOperand(IsEq ? 1 : 2, Z);
  return &Sel;
}

[xbolva00]

lib/Transforms/InstCombine/InstCombineSelect.cpp

78–82

Yes, good idea

[spatel]

test/Transforms/InstCombine/select-binop-icmp.ll

174–175

Please explain why not.
This is the only remaining question I have for this patch
cc @lebedev.ri in case there's any other feedback.

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineSelect.cpp

79

I suspect the cases with undef do not fold only because of this.
I can't tell if it is intentional or not, but i'd make it more explicit, at least with a comment.

[xbolva00]

• Added todo comment
• Ran check-llvm, updated SimplifyCFG test

[xbolva00]

lib/Transforms/InstCombine/InstCombineSelect.cpp

79

Yes, a helper method would be needed to extend it and allow undefs.
Added a comment to test file.

[xbolva00]

• Added TODO comment to code about undefs

[spatel]

LGTM - are you planning to work on any of the TODO items?

[xbolva00]

In D49954#1180956, @spatel wrote:

LGTM - are you planning to work on any of the TODO items?

Yes.

Any advice what should be prioritized from TODO items? Support for undefs?

[spatel]

In D49954#1180973, @xbolva00 wrote:

In D49954#1180956, @spatel wrote:

LGTM - are you planning to work on any of the TODO items?

Yes.

Any advice what should be prioritized from TODO items? Support for undefs?

I'd like to have the folds for non-commutative and FP opcodes first. The undef support would be nice, but I think it's much less likely.

[xbolva00]

In D49954#1180981, @spatel wrote:

In D49954#1180973, @xbolva00 wrote:

In D49954#1180956, @spatel wrote:

LGTM - are you planning to work on any of the TODO items?

Yes.

Any advice what should be prioritized from TODO items? Support for undefs?

I'd like to have the folds for non-commutative and FP opcodes first. The undef support would be nice, but I think it's much less likely.

Ok, I am going to complete the second fold patch and I would return here again.

lib/Transforms/InstCombine/InstCombineSelect.cpp

 
 static Value *createMinMax(InstCombiner::BuilderTy &Builder,
                            SelectPatternFlavor SPF, Value *A, Value *B) {
   CmpInst::Predicate Pred = getMinMaxPred(SPF);
   assert(CmpInst::isIntPredicate(Pred) && "Expected integer predicate");
   return Builder.CreateSelect(Builder.CreateICmp(Pred, A, B), A, B);
 }
 
+/// Fold
+///   %A = icmp eq i8 %x, 0
+///   %B = op i8 %x, %z
+///   %C = select i1 %A, i8 %B, i8 %y
+/// To
+///   %C = select i1 %A, i8 %z, i8 %y
+/// OP: binop with an identity constant
+/// TODO: support for non-commutative and FP opcodes
+static Value *foldSelectInstWithBinaryOp(SelectInst &Sel,
+                                         InstCombiner::BuilderTy &Builder) {
+  Value *Cond = Sel.getCondition();
+  Value *TrueVal = Sel.getTrueValue();
+  Value *FalseVal = Sel.getFalseValue();
+
+  Value *X, *Z, *C;
+  CmpInst::Predicate Pred;
+  if (match(Cond, m_ICmp(Pred, m_Value(X), m_Value(C))) &&
+      Pred == ICmpInst::ICMP_EQ)
+    if (match(TrueVal, m_c_BinOp(m_Specific(X), m_Value(Z))) &&
+        ConstantExpr::getBinOpIdentity(cast<Instruction>(TrueVal)->getOpcode(),
+                                       X->getType()) == C)
+      return Builder.CreateSelect(Cond, Z, FalseVal);
+  return nullptr;
+}
+
 /// This folds:
 ///  select (icmp eq (and X, C1)), TC, FC
 ///    iff C1 is a power 2 and the difference between TC and FC is a power-of-2.
 /// To something like:
 ///  (shr (and (X, C1)), (log2(C1) - log2(TC-FC))) + FC
 /// Or:
 ///  (shl (and (X, C1)), (log2(TC-FC) - log2(C1))) + FC
 /// With some variations depending if FC is larger than TC, or the shift
 /// isn't needed, or the bit widths don't match.
Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   }
 
   // See if we are selecting two values based on a comparison of the two values.
   if (ICmpInst *ICI = dyn_cast<ICmpInst>(CondVal))
     if (Instruction *Result = foldSelectInstWithICmp(SI, ICI))
       return Result;
 
   if (Instruction *Add = foldAddSubSelect(SI, Builder))
     return Add;
 
   // Turn (select C, (op X, Y), (op X, Z)) -> (op X, (select C, Y, Z))
   auto *TI = dyn_cast<Instruction>(TrueVal);
   auto *FI = dyn_cast<Instruction>(FalseVal);
   if (TI && FI && TI->getOpcode() == FI->getOpcode())
     if (Instruction *IV = foldSelectOpOp(SI, TI, FI))
       return IV;
 
Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
 
   if (Instruction *BitCastSel = foldSelectCmpBitcasts(SI, Builder))
     return BitCastSel;
 
   // Simplify selects that test the returned flag of cmpxchg instructions.
   if (Instruction *Select = foldSelectCmpXchg(SI))
     return Select;
 
+  if (Value *V = foldSelectInstWithBinaryOp(SI, Builder))
+    return replaceInstUsesWith(SI, V);
+
   return nullptr;
 }

test/Transforms/InstCombine/select-binop-icmp.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define i32 @select_xor_icmp(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @select_xor_icmp(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = xor i32 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[A]], i32 [[Z:%.*]], i32 [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %A = icmp eq i32 %x, 0
   %B = xor i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
 define i32 @select_mul_icmp(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @select_mul_icmp(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 1
-; CHECK-NEXT:    [[B:%.*]] = mul i32 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[A]], i32 [[Z:%.*]], i32 [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %A = icmp eq i32 %x, 1
   %B = mul i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
 define i32 @select_add_icmp(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @select_add_icmp(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = add i32 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[A]], i32 [[Z:%.*]], i32 [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %A = icmp eq i32 %x, 0
   %B = add i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
 define i32 @select_or_icmp(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @select_or_icmp(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = or i32 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[A]], i32 [[Z:%.*]], i32 [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %A = icmp eq i32 %x, 0
   %B = or i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
 define i32 @select_and_icmp(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @select_and_icmp(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], -1
-; CHECK-NEXT:    [[B:%.*]] = and i32 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[A]], i32 [[Z:%.*]], i32 [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %A = icmp eq i32 %x, -1
   %B = and i32 %x, %z
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
 define <2 x i8> @select_xor_icmp_vec(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
 ; CHECK-LABEL: @select_xor_icmp_vec(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq <2 x i8> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[B:%.*]] = xor <2 x i8> [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[B]], <2 x i8> [[Y:%.*]]
-; CHECK-NEXT:    ret <2 x i8> [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select <2 x i1> [[A]], <2 x i8> [[Z:%.*]], <2 x i8> [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i8> [[TMP1]]
 ;
   %A = icmp eq <2 x i8>  %x, <i8 0, i8 0>
   %B = xor <2 x i8>  %x, %z
   %C = select <2 x i1>  %A, <2 x i8>  %B, <2 x i8>  %y
   ret <2 x i8>  %C
 }
 
-define <2 x i8> @select_xor_icmp_vec2(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
-; CHECK-LABEL: @select_xor_icmp_vec2(
-; CHECK-NEXT:    [[A:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 0, i8 undef>
-; CHECK-NEXT:    [[B:%.*]] = xor <2 x i8> [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[B]], <2 x i8> [[Y:%.*]]
-; CHECK-NEXT:    ret <2 x i8> [[C]]
-;
-  %A = icmp eq <2 x i8>  %x, <i8 0, i8 undef>
-  %B = xor <2 x i8>  %x, %z
-  %C = select <2 x i1>  %A, <2 x i8>  %B, <2 x i8>  %y
-  ret <2 x i8>  %C
-}
-
 define i32 @select_xor_inv_icmp(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @select_xor_inv_icmp(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = xor i32 [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[A]], i32 [[Z:%.*]], i32 [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %A = icmp eq i32 %x, 0
   %B = xor i32 %z, %x
   %C = select i1 %A, i32 %B, i32 %y
   ret i32 %C
 }
 
 ; Negative tests

 ; CHECK-NEXT:    ret <2 x i8> [[C]]
 ;
   %A = icmp eq <2 x i8>  %x, <i8 5, i8 3>
   %B = xor <2 x i8>  %x, %z
   %C = select <2 x i1>  %A, <2 x i8>  %B, <2 x i8>  %y
   ret <2 x i8>  %C
 }
 
+define <2 x i8> @select_xor_icmp_vec_bad_2(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
+; CHECK-LABEL: @select_xor_icmp_vec_bad_2(
+; CHECK-NEXT:    [[A:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 0, i8 undef>
+; CHECK-NEXT:    [[B:%.*]] = xor <2 x i8> [[X]], [[Z:%.*]]
+; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[B]], <2 x i8> [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i8> [[C]]
+;
+  %A = icmp eq <2 x i8>  %x, <i8 0, i8 undef>
+  %B = xor <2 x i8>  %x, %z
+  %C = select <2 x i1>  %A, <2 x i8>  %B, <2 x i8>  %y
+  ret <2 x i8>  %C
+}
+
 define i32 @select_mul_icmp_bad(i32 %x, i32 %y, i32 %z, i32 %k) {
 ; CHECK-LABEL: @select_mul_icmp_bad(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq i32 [[X:%.*]], 3
 ; CHECK-NEXT:    [[B:%.*]] = mul i32 [[X]], [[Z:%.*]]
 ; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i32 [[B]], i32 [[Y:%.*]]
 ; CHECK-NEXT:    ret i32 [[C]]
 ;
   %A = icmp eq i32 %x, 3