[ValueTracking] haveNoCommonBitsSet - add assumption cache handling (PR58624)

[RKSimon]

Summary

As detailed on Issue #58624, we fail to make use of assume calls when attempting to determine if two values haveNoCommonBitsSet

I've added a basic iterator across the assumption cache, looking for 'no common bits' comparisons - both the basic (lhs & rhs) == 0 and a boolean variant !(lhs & rhs) as comparisons of booleans will get folded.

I had hoped to support vectors as well but it looks like we'd have problems with poison if we try to match through reduction intrinsics: https://alive2.llvm.org/ce/z/Hn3Tnn

If people are happy with this approach, I'll get the test coverage pre-committed to simplify the patch

[nikic]

Not sure this is useful. The test case from the issue is over-reduced, the actual motivating case would have to handle a dominating condition, not an assume.

llvm/lib/Analysis/ValueTracking.cpp

313

Shouldn't we be lookup up the assumptions for either LHS or RHS here, rather than all of them?

339

You're missing a context validity check. Currently you're using assumes that don't dominate the context instruction.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

3800

Is this useful independently of the assume case?

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

3800

I guess the answer is no.
The code in llvm::haveNoCommonBitsSet check already exist at:

line 3776: if (SimplifyDemandedInstructionBits(I))
line 3791: if (match(&I, m_c_Xor(m_c_And(m_Not(m_Value(M)), m_Value())

If we add llvm::haveNoCommonBitsSet maybe we can move the code before line 3776 and remove line 3791.

llvm/lib/Analysis/ValueTracking.cpp

 bool llvm::haveNoCommonBitsSet(const Value *LHS, const Value *RHS,
                                const DataLayout &DL, AssumptionCache *AC,
                                const Instruction *CxtI, const DominatorTree *DT,
                                bool UseInstrInfo) {
   assert(LHS->getType() == RHS->getType() &&
          "LHS and RHS should have the same type");
   assert(LHS->getType()->isIntOrIntVectorTy() &&
          "LHS and RHS should be integers");
+
   // Look for an inverted mask: (X & ~M) op (Y & M).
   {
     Value *M;
     if (match(LHS, m_c_And(m_Not(m_Value(M)), m_Value())) &&
         match(RHS, m_c_And(m_Specific(M), m_Value())))
       return true;
     if (match(RHS, m_c_And(m_Not(m_Value(M)), m_Value())) &&
         match(LHS, m_c_And(m_Specific(M), m_Value())))
// Look for: (A & B) op ~(A | B)
     Value *A, *B;
     if (match(LHS, m_And(m_Value(A), m_Value(B))) &&
         match(RHS, m_Not(m_c_Or(m_Specific(A), m_Specific(B)))))
       return true;
     if (match(RHS, m_And(m_Value(A), m_Value(B))) &&
         match(LHS, m_Not(m_c_Or(m_Specific(A), m_Specific(B)))))
       return true;
   }
-  IntegerType *IT = cast<IntegerType>(LHS->getType()->getScalarType());
+
+  auto *IT = cast<IntegerType>(LHS->getType()->getScalarType());
+  if (AC && CxtI) {
+    for (auto &AssumeVH : AC->assumptions()) {
+      if (!AssumeVH)
+        continue;
+      CallInst *I = cast<CallInst>(AssumeVH);
+      assert(I->getFunction() == CxtI->getFunction() &&
+             "Got assumption for the wrong function!");
+
+      // Warning: This loop can end up being somewhat performance sensitive.
+      // We're running this loop once for each value queried resulting in a
+      // runtime of ~O(#assumes * #values).
+
+      assert(I->getCalledFunction()->getIntrinsicID() == Intrinsic::assume &&
+             "must be an assume intrinsic");
+
+      Value *Src = I->getArgOperand(0);
+
+      // Check for assume(and(lhs, rhs) == 0).
+      CmpInst::Predicate Pred;
+      if (match(Src, m_c_ICmp(Pred, m_c_And(m_Specific(LHS), m_Specific(RHS)),
+                              m_Zero())) &&
+          CmpInst::ICMP_EQ == Pred)
+        return true;
+
+      // Check for boolean assume(!and(lhs, rhs)).
+      if (IT->getBitWidth() == 1 &&
+          match(Src, m_Not(m_c_And(m_Specific(LHS), m_Specific(RHS)))))
+        return true;
+    }
+  }
+
   KnownBits LHSKnown(IT->getBitWidth());
   KnownBits RHSKnown(IT->getBitWidth());
   computeKnownBits(LHS, LHSKnown, DL, 0, AC, CxtI, DT, nullptr, UseInstrInfo);
   computeKnownBits(RHS, RHSKnown, DL, 0, AC, CxtI, DT, nullptr, UseInstrInfo);
   return KnownBits::haveNoCommonBitsSet(LHSKnown, RHSKnown);
 }
 
 bool llvm::isOnlyUsedInZeroEqualityComparison(const Instruction *I) {

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Instruction *InstCombinerImpl::visitXor(BinaryOperator &I) {
   Value *M;
   if (match(&I, m_c_Xor(m_c_And(m_Not(m_Value(M)), m_Value()),
                         m_c_And(m_Deferred(M), m_Value()))))
     return BinaryOperator::CreateOr(Op0, Op1);
 
   if (Instruction *Xor = visitMaskedMerge(I, Builder))
     return Xor;
 
+  // A^B --> A|B iff A and B have no bits set in common.
+  if (llvm::haveNoCommonBitsSet(Op0, Op1, DL, &AC, &I, &DT))
+    return BinaryOperator::CreateOr(Op0, Op1);
+
   Value *X, *Y;
   Constant *C1;
   if (match(Op1, m_Constant(C1))) {
     Constant *C2;
 
     if (match(Op0, m_OneUse(m_Or(m_Value(X), m_ImmConstant(C2)))) &&
         match(C1, m_ImmConstant())) {
       // (X | C2) ^ C1 --> (X & ~C2) ^ (C1^C2)

llvm/test/Transforms/InstCombine/xor.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 @G1 = global i32 0
 @G2 = global i32 0
 
 declare i32 @llvm.ctlz.i32(i32, i1)
 declare <2 x i8> @llvm.cttz.v2i8(<2 x i8>, i1)
+declare void @llvm.assume(i1)
 declare void @use(i8)
 
 define i1 @test0(i1 %A) {
 ; CHECK-LABEL: @test0(
 ; CHECK-NEXT:    ret i1 [[A:%.*]]
 ;
   %B = xor i1 %A, false
   ret i1 %B
;
   %r = xor i8 %a, -2
   ret i8 %r
 }
 
 ; (~A & B) ^ A --> A | B
 
 define <2 x i32> @xor_andn_commute1(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @xor_andn_commute1(
-; CHECK-NEXT:    [[Z:%.*]] = or <2 x i32> [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[Z:%.*]] = or <2 x i32> [[B:%.*]], [[A:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[Z]]
 ;
   %nota = xor <2 x i32> %a, <i32 -1, i32 -1>
   %r = and <2 x i32> %nota, %b
   %z = xor <2 x i32> %r, %a
   ret <2 x i32> %z
 }
 

 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %n = sub i32 0, %x
   %a = and i32 %n, %x
   %z = call i32 @llvm.ctlz.i32(i32 %a, i1 true) ; 0 is poison
   %r = xor i32 %z, 30
   ret i32 %r
 }
+
+;
+; (xor x, y) -> (or x, y) iff x and y have no common bits
+;
+
+define i32 @test_nocommonbits_i32(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_nocommonbits_i32(
+; CHECK-NEXT:    [[COMMONBITS:%.*]] = and i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[NOCOMMONBITS:%.*]] = icmp eq i32 [[COMMONBITS]], 0
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[NOCOMMONBITS]])
+; CHECK-NEXT:    [[R:%.*]] = or i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %commonbits = and i32 %x, %y
+  %nocommonbits = icmp eq i32 %commonbits, 0
+  tail call void @llvm.assume(i1 %nocommonbits)
+  %r = xor i32 %x, %y
+  ret i32 %r
+}
+
+; handle icmpeq -> not case for bool types
+define i1 @test_nocommonbits_i1(i1 %x, i1 %y) {
+; CHECK-LABEL: @test_nocommonbits_i1(
+; CHECK-NEXT:    [[COMMONBITS:%.*]] = and i1 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[NOCOMMONBITS:%.*]] = xor i1 [[COMMONBITS]], true
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[NOCOMMONBITS]])
+; CHECK-NEXT:    [[R:%.*]] = or i1 [[X]], [[Y]]
+; CHECK-NEXT:    ret i1 [[R]]
+;
+  %commonbits = and i1 %x, %y
+  %nocommonbits = xor i1 %commonbits, 1
+  tail call void @llvm.assume(i1 %nocommonbits)
+  %r = xor i1 %x, %y
+  ret i1 %r
+}
+
+; negative test - wrong predicate
+define i32 @test_nocommonbits_i32_ne(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_nocommonbits_i32_ne(
+; CHECK-NEXT:    [[COMMONBITS:%.*]] = and i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[NOCOMMONBITS:%.*]] = icmp ne i32 [[COMMONBITS]], 0
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[NOCOMMONBITS]])
+; CHECK-NEXT:    [[R:%.*]] = xor i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %commonbits = and i32 %x, %y
+  %nocommonbits = icmp ne i32 %commonbits, 0
+  tail call void @llvm.assume(i1 %nocommonbits)
+  %r = xor i32 %x, %y
+  ret i32 %r
+}
+
+; negative test - wrong binop
+define i32 @test_nocommonbits_i32_add(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_nocommonbits_i32_add(
+; CHECK-NEXT:    [[COMMONBITS:%.*]] = mul i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[NOCOMMONBITS:%.*]] = icmp eq i32 [[COMMONBITS]], 0
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[NOCOMMONBITS]])
+; CHECK-NEXT:    [[R:%.*]] = xor i32 [[X]], [[Y]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %commonbits = mul i32 %x, %y
+  %nocommonbits = icmp eq i32 %commonbits, 0
+  tail call void @llvm.assume(i1 %nocommonbits)
+  %r = xor i32 %x, %y
+  ret i32 %r
+}