[InstCombine] Fold `(-1 + A) & B` into `A ? 0 : B` where A is effectively a bool

[dtcxzyw]

Summary

Solves issue https://github.com/llvm/llvm-project/issues/63321.

This patch explicitly folds (-1 + A) & B into A ? 0 : B. Additional trunc will be created when A is neither i1 nor <N x i1>.

https://alive2.llvm.org/ce/z/pWv9jJ

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2616

This comment and the title are missing an important detail that the value but me 0/1 (either via load range as is your alive2 link or from i1)

2618

The add doesn't need to be nsw

2624

Instead of duplicating the logic for op1/op2 use 'match(&I, m_c_And(...))`

2626

Comment that 0 refers to depth.

2626

Dont cast op1 cxti, make a simplifyquery at I (sorry not computer at the moment, so can't lookup exact syntax)

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2626

I think using knownbits:::ule(1) would ne a lot clearer.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2618

Why does this need a zext?

[dtcxzyw]

fix comment and rebase

[dtcxzyw]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2618

It needs a zext to avoid creating redundant trunc (zext i1).

[dtcxzyw]

Polite ping :)

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2605

m_c_Add is unnecessary due to constant on RHS.

2605

You need a one use restriction on the add and a multiuse test.

2610

Use countMaxActiveBits() <= 1 instead of KnownBits::ule.

2612

This gets canonicalized to an icmp eq 0, so you should directly emit that.

2618

You are always creating the trunc, so I don't understand how explicitly matching the zext avoids a redundant trunc. The trunc will get folded away later in either case.

[dtcxzyw]

• Rebase
• Address comments

[dtcxzyw]

Ping.

[goldstein.w.n]

llvm/test/Transforms/InstCombine/binop-cast.ll

309

Can you:

1. Split the tests to a seperate patch (create a series so that patch1 is the tests, patch2 is the impl). This allows us to see exact what changes the impl causes
2. Can you add a negative test where the load range is out of range (maybe 0-3)
3. Can you add a test where %x is the same as %y and the 0/1 range is created using and.

[dtcxzyw]

Rebase on D154187

[goldstein.w.n]

llvm/test/Transforms/InstCombine/binop-cast.ll

307

I'm not sure this isn't a regression.
Generally I'd think add; and is as easier to analyze than icmp; select and likely to get better codegen. I think maybe this should only be done if we can actuall save inst count (so above require the zext).

[dtcxzyw]

llvm/test/Transforms/InstCombine/binop-cast.ll

307

icmp; select seems to get better codegen on x86-64 and aarch64.

x86-64: https://godbolt.org/z/jv9v9xTMT
riscv64: https://godbolt.org/z/zTbbnn9x4
aarch64: https://godbolt.org/z/eK6xWKnrb

[dtcxzyw]

Ping.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2609

computeKnownBits(A, /* Depth */ 0, &I).countMaxActiveBits() <= 1 -> isKnownToBeAPowerOfTwo(A,..., /*OrZero*/true, ...)

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2601

Please land this code cleanup as a separate NFC commit.

2609

That's not the same thing though? The current check only allows 0 and 1, which result in -1 and 0 after the add. If it's just any power of two we have 0 and 2^x instead, which become -1 and 2^x-1.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2609

Yup ignore my incorrect comment. I misread as a population count check :(

[nikic]

llvm/test/Transforms/InstCombine/binop-cast.ll

302

This isn't a negative test?

[dtcxzyw]

llvm/test/Transforms/InstCombine/binop-cast.ll

302

It was added as a negative test since @goldstein.w.n thought it should be a regression.
However, I tested the inst seq on some backends and found that icmp; select seems to get better codegen.

x86-64: https://godbolt.org/z/jv9v9xTMT
riscv64: https://godbolt.org/z/zTbbnn9x4
riscv64 (sifive-u74 with SFB opt): https://godbolt.org/z/aKjEs5n1P (icmp; select only costs 2 cycles with SFB)
aarch64: https://godbolt.org/z/eK6xWKnrb

I think the pattern select i1 %x, i32 %y, i32 0 is more likely to get better codegen.

[dtcxzyw]

Ping.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2601

This comment is outstanding.

2604

i1/!range {0,2} isn't really the necessary condition. The necessary condition is A is 0 or 1 which can be proven in a lot more cases.

Please also update the title/summary with this precondition of A being 0/1.

[dtcxzyw]

Rebase.

[goldstein.w.n]

LGTM.

[dtcxzyw]

llvm/test/Transforms/InstCombine/binop-cast.ll

302

@goldstein.w.n Any more comments here? I think it should not be a negative test.

[dtcxzyw]

Ping.

[goldstein.w.n]

In D153148#4647805, @dtcxzyw wrote:

Ping.

That was been approved.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Instruction *InstCombinerImpl::visitAnd(BinaryOperator &I) {
   if (Instruction *Sel = foldBinopOfSextBoolToSelect(I))
     return Sel;
 
   // and(sext(A), B) / and(B, sext(A)) --> A ? B : 0, where A is i1 or <N x i1>.
   // TODO: Move this into foldBinopOfSextBoolToSelect as a more generalized fold
   //       with binop identity constant. But creating a select with non-constant
   //       arm may not be reversible due to poison semantics. Is that a good
   //       canonicalization?
-  Value *A;
-  if (match(Op0, m_OneUse(m_SExt(m_Value(A)))) &&
-      A->getType()->isIntOrIntVectorTy(1))
-    return SelectInst::Create(A, Op1, Constant::getNullValue(Ty));
-  if (match(Op1, m_OneUse(m_SExt(m_Value(A)))) &&
+  Value *A, *B;
+  if (match(&I, m_c_And(m_OneUse(m_SExt(m_Value(A))), m_Value(B))) &&
       A->getType()->isIntOrIntVectorTy(1))
-    return SelectInst::Create(A, Op0, Constant::getNullValue(Ty));
+    return SelectInst::Create(A, B, Constant::getNullValue(Ty));
 
   // Similarly, a 'not' of the bool translates to a swap of the select arms:
-  // ~sext(A) & Op1 --> A ? 0 : Op1
-  // Op0 & ~sext(A) --> A ? 0 : Op0
-  if (match(Op0, m_Not(m_SExt(m_Value(A)))) &&
+  // ~sext(A) & B / B & ~sext(A) --> A ? 0 : B
+  if (match(&I, m_c_And(m_Not(m_SExt(m_Value(A))), m_Value(B))) &&
       A->getType()->isIntOrIntVectorTy(1))
-    return SelectInst::Create(A, Constant::getNullValue(Ty), Op1);
-  if (match(Op1, m_Not(m_SExt(m_Value(A)))) &&
-      A->getType()->isIntOrIntVectorTy(1))
-    return SelectInst::Create(A, Constant::getNullValue(Ty), Op0);
+    return SelectInst::Create(A, Constant::getNullValue(Ty), B);
+
+  // (-1 + A) & B --> A ? 0 : B where A is effectively a bool (zext of i1/<N x
+  // i1>, or inst with !range {0, 2}).
+  if (match(&I, m_c_And(m_OneUse(m_Add(m_ZExtOrSelf(m_Value(A)), m_AllOnes())),
+                        m_Value(B)))) {
+    if (A->getType()->isIntOrIntVectorTy(1))
+      return SelectInst::Create(A, Constant::getNullValue(Ty), B);
+    if (computeKnownBits(A, /* Depth */ 0, &I).countMaxActiveBits() <= 1) {
+      return SelectInst::Create(
+          Builder.CreateICmpEQ(A, Constant::getNullValue(A->getType())), B,
+          Constant::getNullValue(Ty));
+    }
+  }
 
   // (iN X s>> (N-1)) & Y --> (X s< 0) ? Y : 0 -- with optional sext
   if (match(&I, m_c_And(m_OneUse(m_SExtOrSelf(
                             m_AShr(m_Value(X), m_APIntAllowUndef(C)))),
                         m_Value(Y))) &&
       *C == X->getType()->getScalarSizeInBits() - 1) {
     Value *IsNeg = Builder.CreateIsNeg(X, "isneg");
     return SelectInst::Create(IsNeg, Y, ConstantInt::getNullValue(Ty));
   }
   // If there's a 'not' of the shifted value, swap the select operands:
   // ~(iN X s>> (N-1)) & Y --> (X s< 0) ? 0 : Y -- with optional sext
   if (match(&I, m_c_And(m_OneUse(m_SExtOrSelf(
                             m_Not(m_AShr(m_Value(X), m_APIntAllowUndef(C))))),
                         m_Value(Y))) &&
       *C == X->getType()->getScalarSizeInBits() - 1) {
     Value *IsNeg = Builder.CreateIsNeg(X, "isneg");
     return SelectInst::Create(IsNeg, ConstantInt::getNullValue(Ty), Y);
   }
 
   // (~x) & y  -->  ~(x | (~y))  iff that gets rid of inversions

llvm/test/Transforms/InstCombine/binop-cast.ll

;
   call void @use(i32 %sext)
   %r = xor i32 %sext, 42
   ret i32 %r
 }
 
 define i64 @PR63321(ptr %ptr, i64 %c) {
 ; CHECK-LABEL: @PR63321(
 ; CHECK-NEXT:    [[VAL:%.*]] = load i8, ptr [[PTR:%.*]], align 1, !range [[RNG0:![0-9]+]]
-; CHECK-NEXT:    [[RHS:%.*]] = zext i8 [[VAL]] to i64
-; CHECK-NEXT:    [[MASK:%.*]] = add nsw i64 [[RHS]], -1
-; CHECK-NEXT:    [[RES:%.*]] = and i64 [[MASK]], [[C:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[VAL]], 0
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[TMP1]], i64 [[C:%.*]], i64 0
 ; CHECK-NEXT:    ret i64 [[RES]]
 ;
   %val = load i8, ptr %ptr, align 1, !range !{i8 0, i8 2}
   %rhs = zext i8 %val to i64
   %mask = add i64 -1, %rhs
   %res = and i64 %mask, %c
   ret i64 %res
 }

 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %val = zext i1 %x to i32
   %mask = add i32 -1, %val
   %res = and i32 %mask, %y
   ret i32 %res
 }
 
 ; Negative test of and_add_bool_to_select
 define i32 @and_add_bool_no_fold(i32 %y) {
 ; CHECK-LABEL: @and_add_bool_no_fold(
 ; CHECK-NEXT:    [[X:%.*]] = and i32 [[Y:%.*]], 1
-; CHECK-NEXT:    [[MASK:%.*]] = add nsw i32 [[X]], -1
-; CHECK-NEXT:    [[RES:%.*]] = and i32 [[MASK]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i32 [[X]], 0
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[TMP1]], i32 [[Y]], i32 0
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %x = and i32 %y, 1
   %mask = add i32 -1, %x
   %res = and i32 %mask, %y
   ret i32 %res
 }
 
 define <2 x i32> @and_add_bool_vec_to_select(<2 x i1> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @and_add_bool_vec_to_select(