[InstCombine] Fold x + (x | -x) to x & (x - 1)

[marcauberer]

Summary

Fixes #57531

This transformation may be particularly useful on x86-64, because x & (x - 1) can be performed by a single blsr instruction.

[RKSimon]

Sorry for the short/blunt comments - running out of time today!

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1429

Use m_Neg()? Also, you need to use m_Deferred() for those later m_Value()

llvm/test/Transforms/InstCombine/add_or_sub.ll

26

vector cases? multiuse cases? negative cases?

[marcauberer]

Improve, format and add tests

[marcauberer]

@RKSimon thanks for the review. I've uploaded improvements. Just one question left ...

llvm/test/Transforms/InstCombine/add_or_sub.ll

26

I just added vector and multi-use cases. How do I do negative cases?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1434–1437

This shows the expected set of commuted variants of the general pattern, but I don't think the match is handling all of them.

You'll want to add tests for each of these variations, and that's going to require some bonus instructions to make the patterns stick until we reach this point in visitAdd().

Here are some tests that I recently added that can serve as templates - note especially the lines with "thwart complexity-based canoinicalization":
d4a4004c0f9d3070

[spatel]

llvm/test/Transforms/InstCombine/add_or_sub.ll

26

For each positive condition in the match statement(s), create a test that does not satisfy the requirement.

So for example, we're looking for a negate (sub) of a specific operand in the 'or', then create a test with a negate of a different value:

%sub = sub i8 0, %y  <-- mismatch, so this should not transform
%or = or i8 %sub, %x
%add = add i8 %or, %x

Or maybe the 'sub' isn't a negate; it might be: (1 - x). Change the 'or' to an 'xor', etc.

[marcauberer]

Test all combinations and add negative tests

[marcauberer]

@spatel @RKSimon I now added some more extensive tests, which should cover all relevant cases.

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1434–1437

As a matter of fact, it seems like the match is covering all four cases. The first four test cases in the test ref file should cover those. I didn't even need the "thwart complexity based canonicalization instructions"

llvm/test/Transforms/InstCombine/add_or_sub.ll

26

Thanks for the explanation! I've added a few negative tests. Please let me know if something is missing.

[marcauberer]

Rebase

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1434–1437

The patch doesn't transform this example:

define i32 @add_or_sub_comb_i32_variant1(i32 %p) {
  %x = mul i32 %p, %p ; thwart complexity-based canonicalization
  %sub = sub i32 0, %x
  %or = or i32 %x, %sub
  %add = add i32 %or, %x
  ret i32 %add
}

It will be easier to see what's happening if we commit the baseline tests first (and we'll do that anyway), but you should probably create the expected tests for commuted patterns first.

[marcauberer]

Fix bug and fine-tune tests

[marcauberer]

Re-format

[marcauberer]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1434–1437

I think I got now what you mean. I added the "thwart complexity-based canonicalization" to every positive test and fixed the bug in the code. Could you take another look?

The baseline tests are now live here: https://reviews.llvm.org/D133449

[marcauberer]

Rebase and update test file

[spatel]

llvm/test/Transforms/InstCombine/add_or_sub.ll

83

Oops - I missed this detail in the pre-commit. Shouldn't the vector constant have zeros rather than -1 to exercise this patch?

106

We need to bail out on this pattern, so we don't have extra instructions (ie, this will be a negative test). See matcher for "m_OneUse()".

124

We need to bail out on this pattern, so we don't have extra instructions (ie, this will be a negative test). See matcher for "m_OneUse()".

154

Can we propagate the no-wrap flags to the new add (ie, is it correct based on Alive2 proofs)?

[marcauberer]

llvm/test/Transforms/InstCombine/add_or_sub.ll

83

Oh yes, you are right! Is a separate change required for this?

[spatel]

llvm/test/Transforms/InstCombine/add_or_sub.ll

83

No - it's minor and obvious in hindsight. Just fix it up in this patch.

[marcauberer]

Add oneUse to or inst and fix test mistake

[marcauberer]

llvm/test/Transforms/InstCombine/add_or_sub.ll

154

Do we need to propagate the no-wrap flags?

Without propagating the flags: https://alive2.llvm.org/ce/z/3rM4KF
With propagating the flags: https://alive2.llvm.org/ce/z/tWyFm7

[spatel]

llvm/test/Transforms/InstCombine/add_or_sub.ll

154

If we can propagate/generate flags during the transform, that is best. It saves some potentially expensive (in compile-time) analysis that tries to regenerate the flags or prevents losing information completely.

Note that "nuw" on these patterns is a special-case:
https://alive2.llvm.org/ce/z/h1763a

I have no idea if that pattern occurs in the real world, but we don't seem to get the optimal result with or without this patch.

It's a small change to the Builder.CreateAdd() call to propagate flags from "I", so we should do that. Please add tests with "nsw" alone and "nuw" alone to verify that we do the right thing in all cases.

[spatel]

llvm/test/Transforms/InstCombine/add_or_sub.ll

154

On 2nd thought, just add no-wrap flags to some of the adds in the existing tests, and that will more efficiently verify that we are behaving correctly. We don't really need explicit tests just to show flag propagation. You can add a code comment above the test or modify the test function name to explain the intent.

[marcauberer]

Better handling for nuw and nsw flags in tests

[marcauberer]

Code cleanup

[marcauberer]

llvm/test/Transforms/InstCombine/add_or_sub.ll

154

Done, but the positive tests including nuw behave differently, like you mentioned. Do we leave it like this?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1441

Nit: use CamelCase for variable names, so HasNuw or HasNUW (similar for the others)...or just roll those values anonymously into the CreateAdd() line and let clang-format wrap the long line as it likes.

llvm/test/Transforms/InstCombine/add_or_sub.ll

154

I think it's ok since we have plenty of tests covering all of the code paths. The tests with 'nuw' demonstrate that we can chain together 2 or more transforms. If we ever implement the ideal nuw-specific transform after this patch, those tests can confirm if the folds are in the right order to get the best result.

[marcauberer]

Code cleanup

[marcauberer]

llvm/test/Transforms/InstCombine/add_or_sub.ll

154

Okay, perfect.

[spatel]

Thanks for iterating on this! LGTM

[marcauberer]

Update test ref

[marcauberer]

@spatel sorry, forgot to refresh the test ref. Now the change is ready for the final review ...

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
   // (add (and A, B) (or A, B)) --> (add A, B)
   if (match(&I, m_c_BinOp(m_Or(m_Value(A), m_Value(B)),
                           m_c_And(m_Deferred(A), m_Deferred(B))))) {
     // Replacing operands in-place to preserve nuw/nsw flags.
     replaceOperand(I, 0, A);
     replaceOperand(I, 1, B);
     return &I;
   }
   
+  // (add (or (sub 0, A) A) A) --> (and (add A, -1) A)
+  // (add (or A (sub 0, A)) A) --> (and (add A, -1) A)
+  // (add A (or (sub 0, A) A)) --> (and (add A, -1) A)
+  // (add A (or A (sub 0, A))) --> (and (add A, -1) A)
+  if (match(&I, m_c_BinOp(m_Or(m_Sub(m_ZeroInt(), m_Value(A)), m_Value(A)), m_Value(A)))) {
+    Value *AddInst = Builder.CreateAdd(A, Constant::getAllOnesValue(A->getType()));
+    return BinaryOperator::CreateAnd(AddInst, A);
+  }
+
   // Canonicalize ((A & -A) - 1) --> ((A - 1) & ~A)
   // Forms all commutable operations, and simplifies ctpop -> cttz folds.
   if (match(&I,
             m_Add(m_OneUse(m_c_And(m_Value(A), m_OneUse(m_Neg(m_Deferred(A))))),
                   m_AllOnes()))) {
     Constant *AllOnes = ConstantInt::getAllOnesValue(RHS->getType());
     Value *Dec = Builder.CreateAdd(A, AllOnes);
     Value *Not = Builder.CreateXor(A, AllOnes);
     return BinaryOperator::CreateAnd(Dec, Not);
   }
 
   // Disguised reassociation/factorization:
   // ~(A * C1) + A
   // ((A * -C1) - 1) + A
   // ((A * -C1) + A) - 1
   // (A * (1 - C1)) - 1

llvm/test/Transforms/InstCombine/add_or_sub.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i32 @add_or_sub_i32_comb(i32 %x) {
+; CHECK-LABEL: @add_or_sub_i32_comb(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[X:%.*]], -1
+; CHECK-NEXT:    [[ADD:%.*]] = and i32 [[TMP1]], [[X]]
+; CHECK-NEXT:    ret i32 [[ADD]]
+;
+  %sub = sub i32 0, %x
+  %or = or i32 %sub, %x
+  %add = add i32 %or, %x
+  ret i32 %add
+}
+
+define i8 @add_or_sub_i8_comb(i8 %x) {
+; CHECK-LABEL: @add_or_sub_i8_comb(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X:%.*]], -1
+; CHECK-NEXT:    [[ADD:%.*]] = and i8 [[TMP1]], [[X]]
+; CHECK-NEXT:    ret i8 [[ADD]]
+;
+  %sub = sub i8 0, %x
+  %or = or i8 %sub, %x
+  %add = add i8 %or, %x
+  ret i8 %add
+}