[InstCombine] fold add+negate through select into sub

[spatel]

Summary

This transform came up as a potential DAGCombine in D133282, so I wanted to see how it escaped in IR too.

We do general folds in InstCombiner::SimplifySelectsFeedingBinaryOp() by checking if either arm of a select simplifies when the trailing binop is threaded into the select.

So as long as one side simplifies, it's a good fold to combine a negate and add into 1 subtract.

This is an example with a zero arm in the select:
https://alive2.llvm.org/ce/z/Hgu_Tj

And this models the tests with a cancelling 'not' op:
https://alive2.llvm.org/ce/z/BuzVV_

[craig.topper]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

882

Are we wanting to ensure that True and False are different or that only one of them is non-null?

[spatel]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

882

We want exactly one to be non-null. If they're both non-null (both arms simplified), then it'll get handled by the existing code path. If both are null (not simplifiable), then the transform won't work.

[craig.topper]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

882

That’s not exactly what that check does though. Do we need to convert to book first before comparing?

[spatel]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

882

Yes - sorry, clicked 'submit' before finishing the comment. I'll change that check.

[spatel]

Patch updated:
Changed the check for simplified operands to do what the code comment says. I couldn't think of a test to show a failure given the other constraints.

[craig.topper]

LGTM

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (isa<FPMathOperator>(&I)) {
     FMF = I.getFastMathFlags();
     Builder.setFastMathFlags(FMF);
   }
 
   Instruction::BinaryOps Opcode = I.getOpcode();
   SimplifyQuery Q = SQ.getWithInstruction(&I);
 
   Value *Cond, *True = nullptr, *False = nullptr;
+
+  // Special-case for add/negate combination. Replace the zero in the negation
+  // with the trailing add operand:
+  // (Cond ? TVal : -N) + Z --> Cond ? True : (Z - N)
+  // (Cond ? -N : FVal) + Z --> Cond ? (Z - N) : False
+  auto foldAddNegate = [&](Value *TVal, Value *FVal, Value *Z) -> Value * {
+    // We need an 'add' and exactly 1 arm of the select to have been simplified.
+    if (Opcode != Instruction::Add || (!True && !False) || (True && False))
+      return nullptr;
+
+    Value *N;
+    if (True && match(FVal, m_Neg(m_Value(N)))) {
+      Value *Sub = Builder.CreateSub(Z, N);
+      return Builder.CreateSelect(Cond, True, Sub, I.getName());
+    }
+    if (False && match(TVal, m_Neg(m_Value(N)))) {
+      Value *Sub = Builder.CreateSub(Z, N);
+      return Builder.CreateSelect(Cond, Sub, False, I.getName());
+    }
+    return nullptr;
+  };
+
   if (LHSIsSelect && RHSIsSelect && A == D) {
     // (A ? B : C) op (A ? E : F) -> A ? (B op E) : (C op F)
     Cond = A;
     True = simplifyBinOp(Opcode, B, E, FMF, Q);
     False = simplifyBinOp(Opcode, C, F, FMF, Q);
 
     if (LHS->hasOneUse() && RHS->hasOneUse()) {
       if (False && !True)
         True = Builder.CreateBinOp(Opcode, B, E);
       else if (True && !False)
         False = Builder.CreateBinOp(Opcode, C, F);
     }
   } else if (LHSIsSelect && LHS->hasOneUse()) {
     // (A ? B : C) op Y -> A ? (B op Y) : (C op Y)
     Cond = A;
     True = simplifyBinOp(Opcode, B, RHS, FMF, Q);
     False = simplifyBinOp(Opcode, C, RHS, FMF, Q);
+    if (Value *NewSel = foldAddNegate(B, C, RHS))
+      return NewSel;
   } else if (RHSIsSelect && RHS->hasOneUse()) {
     // X op (D ? E : F) -> D ? (X op E) : (X op F)
     Cond = D;
     True = simplifyBinOp(Opcode, LHS, E, FMF, Q);
     False = simplifyBinOp(Opcode, LHS, F, FMF, Q);
+    if (Value *NewSel = foldAddNegate(E, F, LHS))
+      return NewSel;
   }
 
   if (!True || !False)
     return nullptr;
 
   Value *SI = Builder.CreateSelect(Cond, True, False);
   SI->takeName(&I);
   return SI;

llvm/test/Transforms/InstCombine/add.ll

 ;
   %a = ashr i8 %x, 7
   %r = add i8 %a, 2
   ret i8 %r
 }
 
 define i8 @select_negate_or_zero(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_negate_or_zero(
-; CHECK-NEXT:    [[NEGX:%.*]] = sub i8 0, [[X:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i8 0, i8 [[NEGX]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[SEL]], [[Y:%.*]]
-; CHECK-NEXT:    ret i8 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[B:%.*]], i8 0, i8 [[X:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = sub i8 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[ADD1]]
 ;
   %negx = sub i8 0, %x
   %sel = select i1 %b, i8 0, i8 %negx
   %add = add i8 %sel, %y
   ret i8 %add
 }
 
+; commuted add operands - same result
+
 define <2 x i8> @select_negate_or_zero_commute(<2 x i1> %b, <2 x i8> %x, <2 x i8> %p) {
 ; CHECK-LABEL: @select_negate_or_zero_commute(
 ; CHECK-NEXT:    [[Y:%.*]] = mul <2 x i8> [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[NEGX:%.*]] = sub <2 x i8> <i8 poison, i8 0>, [[X:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[B:%.*]], <2 x i8> <i8 poison, i8 0>, <2 x i8> [[NEGX]]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[Y]], [[SEL]]
-; CHECK-NEXT:    ret <2 x i8> [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select <2 x i1> [[B:%.*]], <2 x i8> zeroinitializer, <2 x i8> [[X:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = sub <2 x i8> [[Y]], [[TMP1]]
+; CHECK-NEXT:    ret <2 x i8> [[ADD1]]
 ;
   %y = mul <2 x i8> %p, %p ; thwart complexity-based canonicalization
   %negx = sub <2 x i8> <i8 poison, i8 0>, %x
   %sel = select <2 x i1> %b, <2 x i8> <i8 poison, i8 0>, <2 x i8> %negx
   %add = add <2 x i8> %y, %sel
   ret <2 x i8> %add
 }
 
+; swapped select operands and extra use are ok
+
 define i8 @select_negate_or_zero_swap(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_negate_or_zero_swap(
 ; CHECK-NEXT:    [[NEGX:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[NEGX]])
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i8 [[NEGX]], i8 0
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[SEL]], [[Y:%.*]]
-; CHECK-NEXT:    ret i8 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[B:%.*]], i8 [[X]], i8 0
+; CHECK-NEXT:    [[ADD1:%.*]] = sub i8 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[ADD1]]
 ;
   %negx = sub i8 0, %x
   call void @use(i8 %negx)
   %sel = select i1 %b, i8 %negx, i8 0
   %add = add i8 %sel, %y
   ret i8 %add
 }
 
+; commuted add operands - same result
+
 define i8 @select_negate_or_zero_swap_commute(i1 %b, i8 %x, i8 %p) {
 ; CHECK-LABEL: @select_negate_or_zero_swap_commute(
 ; CHECK-NEXT:    [[Y:%.*]] = mul i8 [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[NEGX:%.*]] = sub i8 0, [[X:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i8 [[NEGX]], i8 0
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[Y]], [[SEL]]
-; CHECK-NEXT:    ret i8 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[B:%.*]], i8 [[X:%.*]], i8 0
+; CHECK-NEXT:    [[ADD1:%.*]] = sub i8 [[Y]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[ADD1]]
 ;
   %y = mul i8 %p, %p ; thwart complexity-based canonicalization
   %negx = sub i8 0, %x
   %sel = select i1 %b, i8 %negx, i8 0
   %add = add i8 %y, %sel
   ret i8 %add
 }
 
+; negative test - one arm of the select must simplify
+
 define i8 @select_negate_or_nonzero(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_negate_or_nonzero(
 ; CHECK-NEXT:    [[NEGX:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i8 42, i8 [[NEGX]]
 ; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[SEL]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[ADD]]
 ;
   %negx = sub i8 0, %x
   %sel = select i1 %b, i8 42, i8 %negx
   %add = add i8 %sel, %y
   ret i8 %add
 }
 
+; negative test - must have a negate, not any subtract
+
 define i8 @select_nonnegate_or_zero(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_nonnegate_or_zero(
 ; CHECK-NEXT:    [[NEGX:%.*]] = sub i8 42, [[X:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i8 0, i8 [[NEGX]]
 ; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[SEL]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[ADD]]
 ;
   %negx = sub i8 42, %x
   %sel = select i1 %b, i8 0, i8 %negx
   %add = add i8 %sel, %y
   ret i8 %add
 }
 
+; negative test - don't create an extra instruction
+
 define i8 @select_negate_or_nonzero_use(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_negate_or_nonzero_use(
 ; CHECK-NEXT:    [[NEGX:%.*]] = sub i8 0, [[X:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i8 0, i8 [[NEGX]]
 ; CHECK-NEXT:    call void @use(i8 [[SEL]])
 ; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[SEL]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i8 [[ADD]]
 ;
   %negx = sub i8 0, %x
   %sel = select i1 %b, i8 0, i8 %negx
   call void @use(i8 %sel)
   %add = add i8 %sel, %y
   ret i8 %add
 }
 
+; extra reduction because y + ~y -> -1
+
 define i5 @select_negate_not(i1 %b, i5 %x, i5 %y) {
 ; CHECK-LABEL: @select_negate_not(
-; CHECK-NEXT:    [[NEGX:%.*]] = sub i5 0, [[X:%.*]]
-; CHECK-NEXT:    [[NOTY:%.*]] = xor i5 [[Y:%.*]], -1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i5 [[NOTY]], i5 [[NEGX]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i5 [[SEL]], [[Y]]
-; CHECK-NEXT:    ret i5 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i5 [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = select i1 [[B:%.*]], i5 -1, i5 [[TMP1]]
+; CHECK-NEXT:    ret i5 [[ADD1]]
 ;
   %negx = sub i5 0, %x
   %noty = xor i5 %y, -1
   %sel = select i1 %b, i5 %noty, i5 %negx
   %add = add i5 %sel, %y
   ret i5 %add
 }
 
 define i5 @select_negate_not_commute(i1 %b, i5 %x, i5 %p) {
 ; CHECK-LABEL: @select_negate_not_commute(
 ; CHECK-NEXT:    [[Y:%.*]] = mul i5 [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[NEGX:%.*]] = sub i5 0, [[X:%.*]]
-; CHECK-NEXT:    [[NOTY:%.*]] = xor i5 [[Y]], -1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i5 [[NOTY]], i5 [[NEGX]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i5 [[Y]], [[SEL]]
-; CHECK-NEXT:    ret i5 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i5 [[Y]], [[X:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = select i1 [[B:%.*]], i5 -1, i5 [[TMP1]]
+; CHECK-NEXT:    ret i5 [[ADD1]]
 ;
   %y = mul i5 %p, %p ; thwart complexity-based canonicalization
   %negx = sub i5 0, %x
   %noty = xor i5 %y, -1
   %sel = select i1 %b, i5 %noty, i5 %negx
   %add = add i5 %y, %sel
   ret i5 %add
 }
 
 define i5 @select_negate_not_swap(i1 %b, i5 %x, i5 %y) {
 ; CHECK-LABEL: @select_negate_not_swap(
-; CHECK-NEXT:    [[NEGX:%.*]] = sub i5 0, [[X:%.*]]
-; CHECK-NEXT:    [[NOTY:%.*]] = xor i5 [[Y:%.*]], -1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i5 [[NEGX]], i5 [[NOTY]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i5 [[SEL]], [[Y]]
-; CHECK-NEXT:    ret i5 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i5 [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = select i1 [[B:%.*]], i5 [[TMP1]], i5 -1
+; CHECK-NEXT:    ret i5 [[ADD1]]
 ;
   %negx = sub i5 0, %x
   %noty = xor i5 %y, -1
   %sel = select i1 %b, i5 %negx, i5 %noty
   %add = add i5 %sel, %y
   ret i5 %add
 }
 
 define i5 @select_negate_not_swap_commute(i1 %b, i5 %x, i5 %p) {
 ; CHECK-LABEL: @select_negate_not_swap_commute(
 ; CHECK-NEXT:    [[Y:%.*]] = mul i5 [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[NEGX:%.*]] = sub i5 0, [[X:%.*]]
-; CHECK-NEXT:    [[NOTY:%.*]] = xor i5 [[Y]], -1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[B:%.*]], i5 [[NEGX]], i5 [[NOTY]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i5 [[Y]], [[SEL]]
-; CHECK-NEXT:    ret i5 [[ADD]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i5 [[Y]], [[X:%.*]]
+; CHECK-NEXT:    [[ADD1:%.*]] = select i1 [[B:%.*]], i5 [[TMP1]], i5 -1
+; CHECK-NEXT:    ret i5 [[ADD1]]
 ;
   %y = mul i5 %p, %p ; thwart complexity-based canonicalization
   %negx = sub i5 0, %x
   %noty = xor i5 %y, -1
   %sel = select i1 %b, i5 %negx, i5 %noty
   %add = add i5 %y, %sel
   ret i5 %add
 }