[ValueTracking] Consider assumptions in isKnownNonEqual

[marcauberer]

Summary

Consider icmp ne and icmp eq assumption intrinsic calls when checking
for known non-equality.

[goldstein.w.n]

Can you add some tests that show this in action?

[jdoerfert]

llvm/lib/Analysis/ValueTracking.cpp

3150

Equality check and Operand checks are probably faster than valid assume check.

[marcauberer]

Add tests and change order of checks

[marcauberer]

@jdoerfert @goldstein.w.n ping

[goldstein.w.n]

No real compile time regressions from this:
https://llvm-compile-time-tracker.com/compare.php?from=facaede5ae43806d3123c18823cc173a4b3970ec&to=f8c93e123c527366c18f3a0b2b02b8cd629ce757&stat=instructions:uP

llvm/lib/Analysis/ValueTracking.cpp

3164

In generaly I'm not sure how valuable this is. Does it affect any of the codes in LLVM test suite or any common patterns?

llvm/test/Analysis/ValueTracking/known-non-equal.ll

1857

Can you pre-commit tests.
Also can you add some negative tests where the assume is 1) out of scope 2) doesn't actually imply equality, and 3) a case like if(X) { assume(X); ...(X) } else { ...(X) }.

[marcauberer]

llvm/lib/Analysis/ValueTracking.cpp

3164

The original use case vanished due to switching to a calculateKnownBits-approach.
So I think the relevant use cases are covered with calculateKnownBits, which is also refered by isKnownNonEqual for int and int vector types.

llvm/include/llvm/Analysis/ValueTracking.h

 /// Returns true if the given value is known be negative (i.e. non-positive
 /// and non-zero).
 bool isKnownNegative(const Value *V, const DataLayout &DL, unsigned Depth = 0,
                      AssumptionCache *AC = nullptr,
                      const Instruction *CxtI = nullptr,
                      const DominatorTree *DT = nullptr,
                      bool UseInstrInfo = true);
 
+/// Check if an assume intrinsic covers the non-equality information
+bool getKnownNonEqualFromAssume(const Value *V1, const Value *V2,
+                                const SimplifyQuery &Q);
+
 /// Return true if the given values are known to be non-equal when defined.
 /// Supports scalar integer types only.
 bool isKnownNonEqual(const Value *V1, const Value *V2, const DataLayout &DL,
                      AssumptionCache *AC = nullptr,
                      const Instruction *CxtI = nullptr,
                      const DominatorTree *DT = nullptr,
                      bool UseInstrInfo = true);
 

llvm/lib/Analysis/ValueTracking.cpp

for (const BasicBlock *IncomBB : PN1->blocks()) {
     RecQ.CxtI = IncomBB->getTerminator();
     if (!isKnownNonEqual(IV1, IV2, Depth + 1, RecQ))
       return false;
     UsedFullRecursion = true;
   }
   return true;
 }
 
+bool llvm::getKnownNonEqualFromAssume(const Value *V1, const Value *V2,
+                                      const SimplifyQuery &Q) {
+  // Use of assumptions is context-sensitive. If we don't have a context, we
+  // cannot use them!
+  if (!Q.AC || !Q.CxtI)
+    return false;
+
+  // Note that the patterns below need to be kept in sync with the code
+  // in AssumptionCache::updateAffectedValues.
+
+  for (auto &AssumeVH : Q.AC->assumptionsFor(V1)) {
+    if (!AssumeVH)
+      continue;
+
+    CallInst *I = cast<CallInst>(AssumeVH);
+    assert(I->getParent()->getParent() == Q.CxtI->getParent()->getParent() &&
+           "Got assumption for the wrong function!");
+
+    // Warning: This loop can end up being somewhat performance sensitive.
+    // We're running this loop for 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 *Arg = I->getArgOperand(0);
+    ICmpInst *Cmp = dyn_cast<ICmpInst>(Arg);
+    if (!Cmp || !Cmp->isEquality())
+      continue;
+
+    // If we have a matching assumption:
+    // - EQ: return false
+    // - NE: return true
+    Value *Op0 = Cmp->getOperand(0);
+    Value *Op1 = Cmp->getOperand(1);
+    if (!(Op0 == V1 && Op1 == V2) && !(Op0 == V2 && Op1 == V1))
+      continue;
+
+    if (isValidAssumeForContext(I, Q.CxtI, Q.DT))
+      return Cmp->getPredicate() == CmpInst::ICMP_NE;
+  }
+
+  return false;
+}
+
 /// Return true if it is known that V1 != V2.
 static bool isKnownNonEqual(const Value *V1, const Value *V2, unsigned Depth,
                             const SimplifyQuery &Q) {
   if (V1 == V2)
     return false;
   if (V1->getType() != V2->getType())
     // We can't look through casts yet.
     return false;
if (V1->getType()->isIntOrIntVectorTy()) {
     // has a known zero where V2 has a known one, they must not be equal.
     KnownBits Known1 = computeKnownBits(V1, Depth, Q);
     KnownBits Known2 = computeKnownBits(V2, Depth, Q);
 
     if (Known1.Zero.intersects(Known2.One) ||
         Known2.Zero.intersects(Known1.One))
       return true;
   }
-  return false;
+
+  // Check whether a nearby assume intrinsic can determine non-equality.
+  return getKnownNonEqualFromAssume(V1, V2, Q);
 }
 
 /// Return true if 'V & Mask' is known to be zero.  We use this predicate to
 /// simplify operations downstream. Mask is known to be zero for bits that V
 /// cannot have.
 ///
 /// This function is defined on values with integer type, values with pointer
 /// type, and vectors of integers.  In the case

llvm/test/Analysis/ValueTracking/known-non-equal.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=instsimplify < %s -S | FileCheck %s
 
+declare void @llvm.assume(i1)
+declare void @use.i1(i1)
+
 define i1 @test(ptr %pq, i8 %B) {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:    ret i1 false
 ;
   %q = load i8, ptr %pq, !range !0 ; %q is known nonzero; no known bits
   %A = add nsw i8 %B, %q
   %cmp = icmp eq i8 %A, %B
   ret i1 %cmp
loop:
   %B.iv.next = sub i8 %A.iv, %B.iv
   %cmp = icmp ne i64 %iv.next, 10
   br i1 %cmp, label %loop, label %exit
 exit:
   %res = icmp eq i8 %A.iv, %B.iv
   ret i1 %res
 }
 
+; Take the information about non-equality from a nearby assumption intrinsic call - ne
+; This test tries to take advantage of the following simplification: icmp eq|ne X, Y -> false|true if X != Y
+define i1 @from_assumption_ne(i16 %x) {
+; CHECK-LABEL: @from_assumption_ne(
+; CHECK-NEXT:    [[NE:%.*]] = icmp ne i16 [[X:%.*]], 56
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NE]])
+; CHECK-NEXT:    call void @use.i1(i1 false)
+; CHECK-NEXT:    call void @use.i1(i1 true)
+; CHECK-NEXT:    ret i1 true
+;
+  %ne = icmp ne i16 %x, 56
+  call void @llvm.assume(i1 %ne)
+  %res1 = icmp eq i16 %x, 56
+  call void @use.i1(i1 %res1)
+  %res2 = icmp ne i16 56, %x
+  call void @use.i1(i1 %res2)
+  %res = or i1 %res1, %res2
+  ret i1 %res
+}
+
+; Take the information about non-equality from a nearby assumption intrinsic call - eq
+; This test tries to take advantage of the following simplification: icmp eq|ne X, Y -> false|true if X != Y
+define i1 @from_assumption_eq(i64 %x, i64 %y) {
+; CHECK-LABEL: @from_assumption_eq(
+; CHECK-NEXT:    [[EQ:%.*]] = icmp eq i64 [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[EQ]])
+; CHECK-NEXT:    call void @use.i1(i1 true)
+; CHECK-NEXT:    call void @use.i1(i1 false)
+; CHECK-NEXT:    ret i1 true
+;
+  %eq = icmp eq i64 %y, %x
+  call void @llvm.assume(i1 %eq)
+  %res1 = icmp eq i64 %x, %y
+  call void @use.i1(i1 %res1)
+  %res2 = icmp ne i64 %y, %x
+  call void @use.i1(i1 %res2)
+  %res = or i1 %res1, %res2
+  ret i1 %res
+}
 
 !0 = !{ i8 1, i8 5 }