Make processing @llvm.assume more efficient - Add affected values to the assumption cache

[hfinkel]

Summary

Here's my second try at making @llvm.assume processing more efficient. My previous attempt, which leveraged operand bundles, r289755, didn't end up working: it did make assume processing more efficient but eliminating the assumption cache made ephemeral value computation too expensive. This is a more-targeted change. We'll keep the assumption cache, but extend it to keep a map of affected values (i.e. values about which an assumption might provide some information) to the corresponding assumption intrinsics. This allows ValueTracking and LVI to find assumptions relevant to the value being queried without scanning all assumptions in the function. The fact that ValueTracking started doing O(number of assumptions in the function) work, for every known-bits query, has become prohibitively expensive in some cases.

[davide]

I'm going to try this on the testcases where your previous approach regressed and review again. In the meanwhile, some early comments.

include/llvm/Analysis/AssumptionCache.h

125–128 ↗

(On Diff #83570)

I prefer using the ternary operator here, but up to you.

lib/Analysis/AssumptionCache.cpp

39–42 ↗

(On Diff #83570)

I see why you need this, but, can we think of a way of sharing (at least partially) the code between this and ValueTracking?
I'm very worried the two implementations going out of sync quickly.

78–82 ↗

(On Diff #83570)

I would add a comment on this matchers (and the others).

[chandlerc]

Can you provide an explanation at a high level for why you're going with this approach rather than having the AC for ephemeral value finding and operand bundles for affected values? Essentially, why not just use both previous approaches rather than adding the affected values to the AC? (I'm not questioning what you've done here, just trying to understand the thought process behind it.)

[hfinkel]

In D28459#639387, @chandlerc wrote:

Can you provide an explanation at a high level for why you're going with this approach rather than having the AC for ephemeral value finding and operand bundles for affected values? Essentially, why not just use both previous approaches rather than adding the affected values to the AC? (I'm not questioning what you've done here, just trying to understand the thought process behind it.)

There are potential downsides to the operand-bundle approach. First, values in operand bundles are, by default, assumed to be captured, alias, etc. which could cause undesirable side effects for pointer-typed values. This could have been fixed, but there is another issue which is that adding all of the extra uses on intermediate values could inhibit simplification of more--complex conditions because of hasOneUse interference. Perhaps a minor problem, but a less tractable one. Also, depending on how you look at it, having the affected values "cached" in the IR, and having a dependence on InstCombine for them to work even for other passes, was suboptimal. As a result, I think that if we can't get rid of the assumption cache, we should use it.

[hfinkel]

lib/Analysis/AssumptionCache.cpp

39–42 ↗

(On Diff #83570)

I'm worried about this too. I tried to make the code here as general as possible, without casting too wide a net, to mitigate this as much as possible.

When I first started working on this, I tried to think of a way to share the matching patterns between ValueTracking and here (perhaps with some small generalization to pick up things that LVI needs that wouldn't otherwise be covered by ValueTracking's patterns. This was the best thing I could come up with. Part of the problem is that, in ValueTracking, we're matching against specific values. Here, we want to find all potentially-matching values. If you can think of a better way, I'l all ears :-)

78–82 ↗

(On Diff #83570)

You mean like this?

// A & B or A | B or A ^ B.

[davide]

lib/Analysis/AssumptionCache.cpp

39–42 ↗

(On Diff #83570)

Yeah, unfortunately, I don't think there's an easy way to generalize. The best we can do here is trying to be very careful when reviewing new changes (and I think your comment in ValueTracking is pretty much needed to make sure we don't forget about this in a month or two).

78–82 ↗

(On Diff #83570)

Yes, exactly.

[hfinkel]

Added some comments to the non-trivial matchers.

[Prazek]

lib/Analysis/AssumptionCache.cpp

34 ↗

(On Diff #83574)

emplace(first, second) or insert({first, second})

[hfinkel]

Use {} instead of make_pair.

[Prazek]

Are you planning to put it in 4.0 release? Sorry for only syntactic comments, didn't have time to go through code in details.

include/llvm/Analysis/AssumptionCache.h

55 ↗

(On Diff #83578)

using

57 ↗

(On Diff #83578)

explicit?

65 ↗

(On Diff #83578)

using?

[hfinkel]

In D28459#639952, @Prazek wrote:

Are you planning to put it in 4.0 release? Sorry for only syntactic comments, didn't have time to go through code in details.

I'd like to.

[Prazek]

lib/Analysis/AssumptionCache.cpp

52 ↗

(On Diff #83578)

auto *Op = I->getOperand(0)

reassigning V can be easy to miss

59–63 ↗

(On Diff #83578)

A and B are undefined, and they are not assigned in m_Value, so what is going on here?
I am just not familiar with this matcher library, but it looks very odd.

[hfinkel]

lib/Analysis/AssumptionCache.cpp

59–63 ↗

(On Diff #83578)

They are assigned in m_Value (as in the example at the very top of include/llvm/IR/PatternMatch.h). You might be thinking of m_Specific, which matches a specific provided value.

[hfinkel]

include/llvm/Analysis/AssumptionCache.h

57 ↗

(On Diff #83578)

No, DenseMap won't work if this is explicit (because we're using Value*'s KeyInfo).

[hfinkel]

Updated per review comments.

[hfinkel]

Any other comments? I'd like to get this in before we branch for 4.0 if possible.

[davide]

Sorry for the delay, Hal.
I just checked and this doesn't regress the cases your previous change regressed, so we should be good on that side.
I somehow share Dan's feeling that this could be solved with an infrastructural changes rather than caching, but I don't feel to be in a position to hinder progress without a concrete/implemented alternative.

[davide]

In D28459#642235, @davide wrote:

Sorry for the delay, Hal.
I just checked and this doesn't regress the cases your previous change regressed, so we should be good on that side.
I somehow share Dan's feeling that this could be solved with an infrastructural changes rather than caching, but I don't feel to be in a position to hinder progress without a concrete/implemented alternative.

For those wondering, I mean http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20170102/416193.html

[hfinkel]

In D28459#642239, @davide wrote:

In D28459#642235, @davide wrote:

Sorry for the delay, Hal.
I just checked and this doesn't regress the cases your previous change regressed, so we should be good on that side.
I somehow share Dan's feeling that this could be solved with an infrastructural changes rather than caching, but I don't feel to be in a position to hinder progress without a concrete/implemented alternative.

For those wondering, I mean http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20170102/416193.html

@dberlin , @chandlerc , et al. does anyone object to me committing this solution at this point? I'm obviously happy to help replace it later with something based on an extended SSA form once we figure out how that should work. In the mean time, this fixes the compile-time problems, which users are certainly hitting, in a fairly transparent way.

[hfinkel]

In D28459#642235, @davide wrote:

Sorry for the delay, Hal.
I just checked and this doesn't regress the cases your previous change regressed, so we should be good on that side.

Also, thanks for checking!

[chandlerc]

In D28459#642252, @hfinkel wrote:

In D28459#642239, @davide wrote:

In D28459#642235, @davide wrote:

Sorry for the delay, Hal.
I just checked and this doesn't regress the cases your previous change regressed, so we should be good on that side.
I somehow share Dan's feeling that this could be solved with an infrastructural changes rather than caching, but I don't feel to be in a position to hinder progress without a concrete/implemented alternative.

For those wondering, I mean http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20170102/416193.html

@dberlin , @chandlerc , et al. does anyone object to me committing this solution at this point? I'm obviously happy to help replace it later with something based on an extended SSA form once we figure out how that should work. In the mean time, this fixes the compile-time problems, which users are certainly hitting, in a fairly transparent way.

I think you should land this, at the very least for 4.0. But I'm not thrilled about it. It seems pragmatic and to address a very real set of problems, but as we've discussed several times here and elsewhere, I look forward to more principled or just less invasive approaches.

Anyways, LGTM.

[dyung]

Hi, I just wanted to say I applied your patch to the top of trunk, built it, and then used it to compile the example that previously resulted in a big compile time increase. I am happy to report that this patch did not meaningfully increase the compile time.

[hfinkel]

In D28459#642270, @dyung wrote:

Hi, I just wanted to say I applied your patch to the top of trunk, built it, and then used it to compile the example that previously resulted in a big compile time increase. I am happy to report that this patch did not meaningfully increase the compile time.

Thanks for checking!

[hfinkel]

In D28459#642265, @chandlerc wrote:

In D28459#642252, @hfinkel wrote:

In D28459#642239, @davide wrote:

In D28459#642235, @davide wrote:

Sorry for the delay, Hal.
I just checked and this doesn't regress the cases your previous change regressed, so we should be good on that side.
I somehow share Dan's feeling that this could be solved with an infrastructural changes rather than caching, but I don't feel to be in a position to hinder progress without a concrete/implemented alternative.

For those wondering, I mean http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20170102/416193.html

@dberlin , @chandlerc , et al. does anyone object to me committing this solution at this point? I'm obviously happy to help replace it later with something based on an extended SSA form once we figure out how that should work. In the mean time, this fixes the compile-time problems, which users are certainly hitting, in a fairly transparent way.

I think you should land this, at the very least for 4.0. But I'm not thrilled about it. It seems pragmatic and to address a very real set of problems, but as we've discussed several times here and elsewhere, I look forward to more principled or just less invasive approaches.

Anyways, LGTM.

Thanks. To be clear, we're all on the same page here.

llvm/trunk/include/llvm/Analysis/AssumptionCache.h

class AssumptionCache {
   ///
   /// We track this to lazily populate our assumptions.
   Function &F;
 
   /// \brief Vector of weak value handles to calls of the @llvm.assume
   /// intrinsic.
   SmallVector<WeakVH, 4> AssumeHandles;
 
+  class AffectedValueCallbackVH final : public CallbackVH {
+    AssumptionCache *AC;
+    void deleted() override;
+    void allUsesReplacedWith(Value *) override;
+
+  public:
+    using DMI = DenseMapInfo<Value *>;
+
+    AffectedValueCallbackVH(Value *V, AssumptionCache *AC = nullptr)
+        : CallbackVH(V), AC(AC) {}
+  };
+
+  friend AffectedValueCallbackVH;
+
+  /// \brief A map of values about which an assumption might be providing
+  /// information to the relevant set of assumptions.
+  using AffectedValuesMap =
+    DenseMap<AffectedValueCallbackVH, SmallVector<WeakVH, 1>,
+             AffectedValueCallbackVH::DMI>;
+  AffectedValuesMap AffectedValues;
+
+  /// Get the vector of assumptions which affect a value from the cache.
+  SmallVector<WeakVH, 1> &getAffectedValues(Value *V);
+
   /// \brief Flag tracking whether we have scanned the function yet.
   ///
   /// We want to be as lazy about this as possible, and so we scan the function
   /// at the last moment.
   bool Scanned;
 
   /// \brief Scan the function for assumptions and add them to the cache.
   void scanFunction();
 
 public:
   /// \brief Construct an AssumptionCache from a function by scanning all of
   /// its instructions.
   AssumptionCache(Function &F) : F(F), Scanned(false) {}
 
   /// \brief Add an @llvm.assume intrinsic to this function's cache.
   ///
   /// The call passed in must be an instruction within this function and must
   /// not already be in the cache.
   void registerAssumption(CallInst *CI);
 
+  /// \brief Update the cache of values being affected by this assumption (i.e.
+  /// the values about which this assumption provides information).
+  void updateAffectedValues(CallInst *CI);
+
   /// \brief Clear the cache of @llvm.assume intrinsics for a function.
   ///
   /// It will be re-scanned the next time it is requested.
   void clear() {
     AssumeHandles.clear();
+    AffectedValues.clear();
     Scanned = false;
   }
 
   /// \brief Access the list of assumption handles currently tracked for this
   /// function.
   ///
   /// Note that these produce weak handles that may be null. The caller must
   /// handle that case.
   /// FIXME: We should replace this with pointee_iterator<filter_iterator<...>>
   /// when we can write that to filter out the null values. Then caller code
   /// will become simpler.
   MutableArrayRef<WeakVH> assumptions() {
     if (!Scanned)
       scanFunction();
     return AssumeHandles;
   }
+
+  /// \brief Access the list of assumptions which affect this value.
+  MutableArrayRef<WeakVH> assumptionsFor(const Value *V) {
+    if (!Scanned)
+      scanFunction();
+
+    auto AVI = AffectedValues.find_as(const_cast<Value *>(V));
+    if (AVI == AffectedValues.end())
+      return MutableArrayRef<WeakVH>();
+
+    return AVI->second;
+  }
 };
 
 /// \brief A function analysis which provides an \c AssumptionCache.
 ///
 /// This analysis is intended for use with the new pass manager and will vend
 /// assumption caches for a given function.
 class AssumptionAnalysis : public AnalysisInfoMixin<AssumptionAnalysis> {
   friend AnalysisInfoMixin<AssumptionAnalysis>;

llvm/trunk/lib/Analysis/AssumptionCache.cpp

 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Debug.h"
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
+SmallVector<WeakVH, 1> &AssumptionCache::getAffectedValues(Value *V) {
+  // Try using find_as first to avoid creating extra value handles just for the
+  // purpose of doing the lookup.
+  auto AVI = AffectedValues.find_as(V);
+  if (AVI != AffectedValues.end())
+    return AVI->second;
+
+  auto AVIP = AffectedValues.insert({
+      AffectedValueCallbackVH(V, this), SmallVector<WeakVH, 1>()});
+  return AVIP.first->second;
+}
+
+void AssumptionCache::updateAffectedValues(CallInst *CI) {
+  // Note: This code must be kept in-sync with the code in
+  // computeKnownBitsFromAssume in ValueTracking.
+
+  SmallVector<Value *, 16> Affected;
+  auto AddAffected = [&Affected](Value *V) {
+    if (isa<Argument>(V)) {
+      Affected.push_back(V);
+    } else if (auto *I = dyn_cast<Instruction>(V)) {
+      Affected.push_back(I);
+
+      if (I->getOpcode() == Instruction::BitCast ||
+          I->getOpcode() == Instruction::PtrToInt) {
+        auto *Op = I->getOperand(0);
+        if (isa<Instruction>(Op) || isa<Argument>(Op))
+          Affected.push_back(Op);
+      }
+    }
+  };
+
+  Value *Cond = CI->getArgOperand(0), *A, *B;
+  AddAffected(Cond);
+
+  CmpInst::Predicate Pred;
+  if (match(Cond, m_ICmp(Pred, m_Value(A), m_Value(B)))) {
+    AddAffected(A);
+    AddAffected(B);
+
+    if (Pred == ICmpInst::ICMP_EQ) {
+      // For equality comparisons, we handle the case of bit inversion.
+      auto AddAffectedFromEq = [&AddAffected](Value *V) {
+        Value *A;
+        if (match(V, m_Not(m_Value(A)))) {
+          AddAffected(A);
+          V = A;
+        }
+
+        Value *B;
+        ConstantInt *C;
+        // (A & B) or (A | B) or (A ^ B).
+        if (match(V,
+                  m_CombineOr(m_And(m_Value(A), m_Value(B)),
+                    m_CombineOr(m_Or(m_Value(A), m_Value(B)),
+                                m_Xor(m_Value(A), m_Value(B)))))) {
+          AddAffected(A);
+          AddAffected(B);
+        // (A << C) or (A >>_s C) or (A >>_u C) where C is some constant.
+        } else if (match(V,
+                         m_CombineOr(m_Shl(m_Value(A), m_ConstantInt(C)),
+                           m_CombineOr(m_LShr(m_Value(A), m_ConstantInt(C)),
+                                       m_AShr(m_Value(A),
+                                              m_ConstantInt(C)))))) {
+          AddAffected(A);
+        }
+      };
+
+      AddAffectedFromEq(A);
+      AddAffectedFromEq(B);
+    }
+  }
+
+  for (auto &AV : Affected) {
+    auto &AVV = getAffectedValues(AV);
+    if (std::find(AVV.begin(), AVV.end(), CI) == AVV.end())
+      AVV.push_back(CI);
+  }
+}
+
+void AssumptionCache::AffectedValueCallbackVH::deleted() {
+  auto AVI = AC->AffectedValues.find(getValPtr());
+  if (AVI != AC->AffectedValues.end())
+    AC->AffectedValues.erase(AVI);
+  // 'this' now dangles!
+}
+
+void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) {
+  if (!isa<Instruction>(NV) && !isa<Argument>(NV))
+    return;
+
+  // Any assumptions that affected this value now affect the new value.
+
+  auto &NAVV = AC->getAffectedValues(NV);
+  auto AVI = AC->AffectedValues.find(getValPtr());
+  if (AVI == AC->AffectedValues.end())
+    return;
+
+  for (auto &A : AVI->second)
+    if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end())
+      NAVV.push_back(A);
+}
+
 void AssumptionCache::scanFunction() {
   assert(!Scanned && "Tried to scan the function twice!");
   assert(AssumeHandles.empty() && "Already have assumes when scanning!");
 
   // Go through all instructions in all blocks, add all calls to @llvm.assume
   // to this cache.
   for (BasicBlock &B : F)
     for (Instruction &II : B)
       if (match(&II, m_Intrinsic<Intrinsic::assume>()))
         AssumeHandles.push_back(&II);
 
   // Mark the scan as complete.
   Scanned = true;
+
+  // Update affected values.
+  for (auto &A : AssumeHandles)
+    updateAffectedValues(cast<CallInst>(A));
 }
 
 void AssumptionCache::registerAssumption(CallInst *CI) {
   assert(match(CI, m_Intrinsic<Intrinsic::assume>()) &&
          "Registered call does not call @llvm.assume");
 
   // If we haven't scanned the function yet, just drop this assumption. It will
   // be found when we scan later.
for (auto &VH : AssumeHandles) {
     assert(&F == cast<Instruction>(VH)->getParent()->getParent() &&
            "Cached assumption not inside this function!");
     assert(match(cast<CallInst>(VH), m_Intrinsic<Intrinsic::assume>()) &&
            "Cached something other than a call to @llvm.assume!");
     assert(AssumptionSet.insert(VH).second &&
            "Cache contains multiple copies of a call!");
   }
 #endif
+
+  updateAffectedValues(CI);
 }
 
 AnalysisKey AssumptionAnalysis::Key;
 
 PreservedAnalyses AssumptionPrinterPass::run(Function &F,
                                              FunctionAnalysisManager &AM) {
   AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
 

llvm/trunk/lib/Analysis/LazyValueInfo.cpp

 // If we can determine a constraint on the value given conditions assumed by
 // the program, intersect those constraints with BBLV
 void LazyValueInfoImpl::intersectAssumeOrGuardBlockValueConstantRange(
         Value *Val, LVILatticeVal &BBLV, Instruction *BBI) {
   BBI = BBI ? BBI : dyn_cast<Instruction>(Val);
   if (!BBI)
     return;
 
-  for (auto &AssumeVH : AC->assumptions()) {
+  for (auto &AssumeVH : AC->assumptionsFor(Val)) {
     if (!AssumeVH)
       continue;
     auto *I = cast<CallInst>(AssumeVH);
     if (!isValidAssumeForContext(I, BBI, DT))
       continue;
 
     BBLV = intersect(BBLV, getValueFromCondition(Val, I->getArgOperand(0)));
   }

llvm/trunk/lib/Analysis/ValueTracking.cpp

static void computeKnownBitsFromAssume(const Value *V, APInt &KnownZero,
                                        const Query &Q) {
   // Use of assumptions is context-sensitive. If we don't have a context, we
   // cannot use them!
   if (!Q.AC || !Q.CxtI)
     return;
 
   unsigned BitWidth = KnownZero.getBitWidth();
 
-  for (auto &AssumeVH : Q.AC->assumptions()) {
+  // Note that the patterns below need to be kept in sync with the code
+  // in AssumptionCache::updateAffectedValues.
+
+  for (auto &AssumeVH : Q.AC->assumptionsFor(V)) {
     if (!AssumeVH)
       continue;
     CallInst *I = cast<CallInst>(AssumeVH);
     assert(I->getParent()->getParent() == Q.CxtI->getParent()->getParent() &&
            "Got assumption for the wrong function!");
     if (Q.isExcluded(I))
       continue;
 

llvm/trunk/lib/Transforms/InstCombine/InstCombineCalls.cpp

case Intrinsic::assume: {
 
     // If there is a dominating assume with the same condition as this one,
     // then this one is redundant, and should be removed.
     APInt KnownZero(1, 0), KnownOne(1, 0);
     computeKnownBits(IIOperand, KnownZero, KnownOne, 0, II);
     if (KnownOne.isAllOnesValue())
       return eraseInstFromFunction(*II);
 
+    // Update the cache of affected values for this assumption (we might be
+    // here because we just simplified the condition).
+    AC.updateAffectedValues(II);
     break;
   }
   case Intrinsic::experimental_gc_relocate: {
     // Translate facts known about a pointer before relocating into
     // facts about the relocate value, while being careful to
     // preserve relocation semantics.
     Value *DerivedPtr = cast<GCRelocateInst>(II)->getDerivedPtr();