Add an optional cache to computeKnownBits.

[jlebar]

Summary

BasicAA (via MemoryDependenceResults, via GVN) may make many calls to
computeKnownBits. On my testcase (not able to share, sorry), this cache
reduces the time spent in computeKnownBits from 2s to 270ms, a 7.5x
speedup.

[nikic]

llvm/include/llvm/Analysis/AliasAnalysis.h

41

Incorrect include paths

llvm/include/llvm/Analysis/ValueTracking.h

52

This is incorrect, the context instruction is used to check applicability of assumes for example, see isValidAssumeForContext().

[foad]

from 2 to 270ms

Makes it sound like a 135x slow down.

[foad]

llvm/include/llvm/Analysis/ValueTracking.h

49

If you want to reformat this file can you do it first as a separate NFC patch please?

[jlebar]

llvm/include/llvm/Analysis/ValueTracking.h

52

isGuaranteedToTransferExecutionToSuccessor, wow, TIL. Thanks.

But, very sad: With the change to look at the real CtxI rather than its BB, this no longer provides a significant speedup.

This all seems so silly, because in 99.9% of cases the BB terminator really *is* sufficient.

Any thoughts? Or do you think I should give up on this?

[jlebar]

llvm/include/llvm/Analysis/ValueTracking.h

49

I didn't really *want* to reformat the file, but arc was doing it for me because I touched these lines... :-/

Anyway, I pushed 8cc3bfd13f3135985e5b15ee65f2fc43239fb9fe.

[nikic]

BasicAA (via MemoryDependenceResults, via GVN) may make many calls to
computeKnownBits. On my testcase (not able to share, sorry), this cache
reduces the time spent in computeKnownBits from 2s to 270ms, a 7.5x
speedup.

MDA has a known problem with overly large recursion cutoffs. I'd be interested to know whether passing -memdep-block-number-limit=100 would improve compile-time for your case.

llvm/include/llvm/Analysis/ValueTracking.h

52

I don't have any great ideas on how to do more effective caching without affecting result quality. I guess we could decouple the context-independent from the context-dependent parts somehow, but given that they are currently interleaved, this doesn't sound easy.

[jlebar]

MDA has a known problem with overly large recursion cutoffs. I'd be interested to know whether passing -memdep-block-number-limit=100 would improve compile-time for your case.

Turns out that tweaking the memdep limits did help in my case.

So I guess I'm happy dropping this patch. Though I *do* feel like there's still something here.

llvm/include/llvm/Analysis/AliasAnalysis.h

 //    MustAlias at the same time. The current API can only return one result,
 //    though this is rarely a problem in practice.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_ALIASANALYSIS_H
 #define LLVM_ANALYSIS_ALIASANALYSIS_H
 
+#include "third_party/llvm/llvm-project/llvm/include/llvm/Analysis/ValueTracking.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/Support/KnownBits.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include <cstdint>
public:
   /// How many active NoAlias assumption uses there are.
   int NumAssumptionUses = 0;
 
   /// Location pairs for which an assumption based result is currently stored.
   /// Used to remove all potentially incorrect results from the cache if an
   /// assumption is disproven.
   SmallVector<AAQueryInfo::LocPair, 4> AssumptionBasedResults;
 
+  KnownBitsCache KnownBitsCache;
+
   AAQueryInfo(CaptureInfo *CI) : CI(CI) {}
 
   /// Create a new AAQueryInfo based on this one, but with the cache cleared.
   /// This is used for recursive queries across phis, where cache results may
   /// not be valid.
   AAQueryInfo withEmptyCache() {
     AAQueryInfo NewAAQI(CI);
     NewAAQI.Depth = Depth;

llvm/include/llvm/Analysis/ValueTracking.h

 class MDNode;
 class OptimizationRemarkEmitter;
 class StringRef;
 class TargetLibraryInfo;
 class Value;
 
 constexpr unsigned MaxAnalysisRecursionDepth = 6;
 
-  /// Determine which bits of V are known to be either zero or one and return
-  /// them in the KnownZero/KnownOne bit sets.
-  ///
-  /// This function is defined on values with integer type, values with pointer
-  /// type, and vectors of integers.  In the case
-  /// where V is a vector, the known zero and known one values are the
-  /// same width as the vector element, and the bit is set only if it is true
-  /// for all of the elements in the vector.
-  void computeKnownBits(const Value *V, KnownBits &Known,
-                        const DataLayout &DL, unsigned Depth = 0,
-                        AssumptionCache *AC = nullptr,
-                        const Instruction *CxtI = nullptr,
-                        const DominatorTree *DT = nullptr,
-                        OptimizationRemarkEmitter *ORE = nullptr,
-                        bool UseInstrInfo = true);
+// Optional cache used by KnownBits.  Populated as we walk the use/def graph.
+//
+// We can key lookups on the basic block that contains CxtI, rather than on CxtI
+// itself, because the context instr itself is not significant, only its BB.
+using KnownBitsCache = DenseMap<
+    std::pair<const Value *, const BasicBlock * /*BB containing CxtI*/>,
+    KnownBits>;
+
+void computeKnownBits(const Value *V, KnownBits &Known, const DataLayout &DL,
+                      unsigned Depth = 0, AssumptionCache *AC = nullptr,
+                      bool UseInstrInfo = true,
+                      KnownBitsCache *Cache = nullptr);
 
-  /// Determine which bits of V are known to be either zero or one and return
-  /// them in the KnownZero/KnownOne bit sets.
-  ///
-  /// This function is defined on values with integer type, values with pointer
-  /// type, and vectors of integers.  In the case
-  /// where V is a vector, the known zero and known one values are the
-  /// same width as the vector element, and the bit is set only if it is true
-  /// for all of the demanded elements in the vector.
-  void computeKnownBits(const Value *V, const APInt &DemandedElts,
-                        KnownBits &Known, const DataLayout &DL,
-                        unsigned Depth = 0, AssumptionCache *AC = nullptr,
-                        const Instruction *CxtI = nullptr,
-                        const DominatorTree *DT = nullptr,
-                        OptimizationRemarkEmitter *ORE = nullptr,
-                        bool UseInstrInfo = true);
+void computeKnownBits(
+    const Value *V, const APInt &DemandedElts, KnownBits &Known,
+    const DataLayout &DL, unsigned Depth = 0, AssumptionCache *AC = nullptr,
+    const Instruction *CxtI = nullptr, const DominatorTree *DT = nullptr,
+    OptimizationRemarkEmitter *ORE = nullptr, bool UseInstrInfo = true,
+    KnownBitsCache *Cache = nullptr);
 
-  /// Returns the known bits rather than passing by reference.
-  KnownBits computeKnownBits(const Value *V, const DataLayout &DL,
-                             unsigned Depth = 0, AssumptionCache *AC = nullptr,
-                             const Instruction *CxtI = nullptr,
-                             const DominatorTree *DT = nullptr,
-                             OptimizationRemarkEmitter *ORE = nullptr,
-                             bool UseInstrInfo = true);
+KnownBits computeKnownBits(
+    const Value *V, const DataLayout &DL, unsigned Depth = 0,
+    AssumptionCache *AC = nullptr, const Instruction *CxtI = nullptr,
+    const DominatorTree *DT = nullptr, OptimizationRemarkEmitter *ORE = nullptr,
+    bool UseInstrInfo = true,
+    KnownBitsCache *Cache = nullptr);
 
-  /// Returns the known bits rather than passing by reference.
-  KnownBits computeKnownBits(const Value *V, const APInt &DemandedElts,
-                             const DataLayout &DL, unsigned Depth = 0,
-                             AssumptionCache *AC = nullptr,
-                             const Instruction *CxtI = nullptr,
-                             const DominatorTree *DT = nullptr,
-                             OptimizationRemarkEmitter *ORE = nullptr,
-                             bool UseInstrInfo = true);
+KnownBits computeKnownBits(
+    const Value *V, const APInt &DemandedElts, const DataLayout &DL,
+    unsigned Depth = 0, AssumptionCache *AC = nullptr,
+    const Instruction *CxtI = nullptr, const DominatorTree *DT = nullptr,
+    OptimizationRemarkEmitter *ORE = nullptr, bool UseInstrInfo = true,
+    KnownBitsCache *Cache = nullptr);
 
-  /// Compute known bits from the range metadata.
-  /// \p KnownZero the set of bits that are known to be zero
-  /// \p KnownOne the set of bits that are known to be one
-  void computeKnownBitsFromRangeMetadata(const MDNode &Ranges,
-                                         KnownBits &Known);
+void computeKnownBitsFromRangeMetadata(const MDNode &Ranges, KnownBits &Known);
 
-  /// Return true if LHS and RHS have no common bits set.
-  bool haveNoCommonBitsSet(const Value *LHS, const Value *RHS,
-                           const DataLayout &DL,
-                           AssumptionCache *AC = nullptr,
-                           const Instruction *CxtI = nullptr,
-                           const DominatorTree *DT = nullptr,
-                           bool UseInstrInfo = true);
+                         const DataLayout &DL, AssumptionCache *AC = nullptr,
 
-  /// Return true if the given value is known to have exactly one bit set when
-  /// defined. For vectors return true if every element is known to be a power
-  /// of two when defined. Supports values with integer or pointer type and
-  /// vectors of integers. If 'OrZero' is set, then return true if the given
-  /// value is either a power of two or zero.
-  bool isKnownToBeAPowerOfTwo(const Value *V, const DataLayout &DL,
-                              bool OrZero = false, unsigned Depth = 0,
-                              AssumptionCache *AC = nullptr,
-                              const Instruction *CxtI = nullptr,
-                              const DominatorTree *DT = nullptr,
-                              bool UseInstrInfo = true);
 
-  bool isOnlyUsedInZeroEqualityComparison(const Instruction *CxtI);
 
-  /// Return true if the given value is known to be non-zero when defined. For
-  /// vectors, return true if every element is known to be non-zero when
-  /// defined. For pointers, if the context instruction and dominator tree are
-  /// specified, perform context-sensitive analysis and return true if the
-  /// pointer couldn't possibly be null at the specified instruction.
-  /// Supports values with integer or pointer type and vectors of integers.
-  bool isKnownNonZero(const Value *V, const DataLayout &DL, unsigned Depth = 0,
-                      AssumptionCache *AC = nullptr,
-                      const Instruction *CxtI = nullptr,
-                      const DominatorTree *DT = nullptr,
-                      bool UseInstrInfo = true);
 
-  /// Return true if the two given values are negation.
-  /// Currently can recoginze Value pair:
-  /// 1: <X, Y> if X = sub (0, Y) or Y = sub (0, X)
-  /// 2: <X, Y> if X = sub (A, B) and Y = sub (B, A)
-  bool isKnownNegation(const Value *X, const Value *Y, bool NeedNSW = false);
 
-  /// Returns true if the give value is known to be non-negative.
-  bool isKnownNonNegative(const Value *V, const DataLayout &DL,
-                          unsigned Depth = 0,
-                          AssumptionCache *AC = nullptr,
-                          const Instruction *CxtI = nullptr,
-                          const DominatorTree *DT = nullptr,
-                          bool UseInstrInfo = true);
+                        unsigned Depth = 0, AssumptionCache *AC = nullptr,
 
-  /// Returns true if the given value is known be positive (i.e. non-negative
-  /// and non-zero).
-  bool isKnownPositive(const Value *V, const DataLayout &DL, unsigned Depth = 0,
-                       AssumptionCache *AC = nullptr,
-                       const Instruction *CxtI = nullptr,
-                       const DominatorTree *DT = nullptr,
-                       bool UseInstrInfo = true);
 
-  /// 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);
 
-  /// 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);
 
-  /// 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
-  /// where V is a vector, the mask, known zero, and known one values are the
-  /// same width as the vector element, and the bit is set only if it is true
-  /// for all of the elements in the vector.
-  bool MaskedValueIsZero(const Value *V, const APInt &Mask,
-                         const DataLayout &DL,
-                         unsigned Depth = 0, AssumptionCache *AC = nullptr,
-                         const Instruction *CxtI = nullptr,
-                         const DominatorTree *DT = nullptr,
-                         bool UseInstrInfo = true);
+bool MaskedValueIsZero(const Value *V, const APInt &Mask, const DataLayout &DL,
 
-  /// Return the number of times the sign bit of the register is replicated into
-  /// the other bits. We know that at least 1 bit is always equal to the sign
-  /// bit (itself), but other cases can give us information. For example,
-  /// immediately after an "ashr X, 2", we know that the top 3 bits are all
-  /// equal to each other, so we return 3. For vectors, return the number of
-  /// sign bits for the vector element with the mininum number of known sign
-  /// bits.
-  unsigned ComputeNumSignBits(const Value *Op, const DataLayout &DL,
-                              unsigned Depth = 0, AssumptionCache *AC = nullptr,
-                              const Instruction *CxtI = nullptr,
-                              const DominatorTree *DT = nullptr,
-                              bool UseInstrInfo = true);
 
-  /// Get the upper bound on bit size for this Value \p Op as a signed integer.
-  /// i.e.  x == sext(trunc(x to MaxSignificantBits) to bitwidth(x)).
-  /// Similar to the APInt::getSignificantBits function.
-  unsigned ComputeMaxSignificantBits(const Value *Op, const DataLayout &DL,
-                                     unsigned Depth = 0,
-                                     AssumptionCache *AC = nullptr,
-                                     const Instruction *CxtI = nullptr,
-                                     const DominatorTree *DT = nullptr);
 
-  /// Map a call instruction to an intrinsic ID.  Libcalls which have equivalent
-  /// intrinsics are treated as-if they were intrinsics.
-  Intrinsic::ID getIntrinsicForCallSite(const CallBase &CB,
-                                        const TargetLibraryInfo *TLI);
 
-  /// Return true if we can prove that the specified FP value is never equal to
-  /// -0.0.
-  bool CannotBeNegativeZero(const Value *V, const TargetLibraryInfo *TLI,
-                            unsigned Depth = 0);
 
-  /// Return true if we can prove that the specified FP value is either NaN or
-  /// never less than -0.0.
-  ///
-  ///      NaN --> true
-  ///       +0 --> true
-  ///       -0 --> true
-  ///   x > +0 --> true
-  ///   x < -0 --> false
-  bool CannotBeOrderedLessThanZero(const Value *V, const TargetLibraryInfo *TLI);
 
-  /// Return true if the floating-point scalar value is not an infinity or if
-  /// the floating-point vector value has no infinities. Return false if a value
-  /// could ever be infinity.
-  bool isKnownNeverInfinity(const Value *V, const TargetLibraryInfo *TLI,
-                            unsigned Depth = 0);
 
-  /// Return true if the floating-point scalar value is not a NaN or if the
-  /// floating-point vector value has no NaN elements. Return false if a value
-  /// could ever be NaN.
-  bool isKnownNeverNaN(const Value *V, const TargetLibraryInfo *TLI,
-                       unsigned Depth = 0);
 
-  /// Return true if we can prove that the specified FP value's sign bit is 0.
-  ///
-  ///      NaN --> true/false (depending on the NaN's sign bit)
-  ///       +0 --> true
-  ///       -0 --> false
-  ///   x > +0 --> true
-  ///   x < -0 --> false
-  bool SignBitMustBeZero(const Value *V, const TargetLibraryInfo *TLI);
 
-  /// If the specified value can be set by repeating the same byte in memory,
-  /// return the i8 value that it is represented with. This is true for all i8
-  /// values obviously, but is also true for i32 0, i32 -1, i16 0xF0F0, double
-  /// 0.0 etc. If the value can't be handled with a repeated byte store (e.g.
-  /// i16 0x1234), return null. If the value is entirely undef and padding,
-  /// return undef.
-  Value *isBytewiseValue(Value *V, const DataLayout &DL);
 
-  /// Given an aggregate and an sequence of indices, see if the scalar value
-  /// indexed is already around as a register, for example if it were inserted
-  /// directly into the aggregate.
-  ///
-  /// If InsertBefore is not null, this function will duplicate (modified)
-  /// insertvalues when a part of a nested struct is extracted.
-  Value *FindInsertedValue(Value *V,
-                           ArrayRef<unsigned> idx_range,
-                           Instruction *InsertBefore = nullptr);
+Value *FindInsertedValue(Value *V, ArrayRef<unsigned> idx_range,
 
-  /// Analyze the specified pointer to see if it can be expressed as a base
-  /// pointer plus a constant offset. Return the base and offset to the caller.
-  ///
-  /// This is a wrapper around Value::stripAndAccumulateConstantOffsets that
-  /// creates and later unpacks the required APInt.
-  inline Value *GetPointerBaseWithConstantOffset(Value *Ptr, int64_t &Offset,
-                                                 const DataLayout &DL,
-                                                 bool AllowNonInbounds = true) {
-    APInt OffsetAPInt(DL.getIndexTypeSizeInBits(Ptr->getType()), 0);
-    Value *Base =
-        Ptr->stripAndAccumulateConstantOffsets(DL, OffsetAPInt, AllowNonInbounds);
 
-    Offset = OffsetAPInt.getSExtValue();
-    return Base;
   }
   inline const Value *
   GetPointerBaseWithConstantOffset(const Value *Ptr, int64_t &Offset,
                                    const DataLayout &DL,
                                    bool AllowNonInbounds = true) {
     return GetPointerBaseWithConstantOffset(const_cast<Value *>(Ptr), Offset, DL,
                                             AllowNonInbounds);
   }

llvm/lib/Analysis/BasicAliasAnalysis.cpp

for (unsigned i = 0, e = DecompGEP1.VarIndices.size(); i != e; ++i) {
 
     if (i == 0)
       GCD = ScaleForGCD.abs();
     else
       GCD = APIntOps::GreatestCommonDivisor(GCD, ScaleForGCD.abs());
 
     ConstantRange CR = computeConstantRange(Index.Val.V, /* ForSigned */ false,
                                             true, &AC, Index.CxtI);
-    KnownBits Known =
-        computeKnownBits(Index.Val.V, DL, 0, &AC, Index.CxtI, DT);
+    KnownBits Known = computeKnownBits(Index.Val.V, DL, 0, &AC, Index.CxtI, DT,
+                                       /*ORE=*/nullptr, /*UseInstrInfo=*/true,
+                                       &AAQI.KnownBitsCache);
     CR = CR.intersectWith(
         ConstantRange::fromKnownBits(Known, /* Signed */ true),
         ConstantRange::Signed);
     CR = Index.Val.evaluateWith(CR).sextOrTrunc(OffsetRange.getBitWidth());
 
     assert(OffsetRange.getBitWidth() == Scale.getBitWidth() &&
            "Bit widths are normalized to MaxIndexSize");
     if (Index.IsNSW)

llvm/lib/Analysis/ValueTracking.cpp

 //===- ValueTracking.cpp - Walk computations to compute properties --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains routines that help analyze properties that chains of
 // computations have.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/ValueTracking.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <utility>
+
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Cleanup.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"

 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
-#include <algorithm>
-#include <cassert>
-#include <cstdint>
-#include <utility>
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 // Controls the number of uses of the value searched for possible
 // dominating comparisons.
 static cl::opt<unsigned> DomConditionsMaxUses("dom-conditions-max-uses",
                                               cl::Hidden, cl::init(20));
struct Query {
   AssumptionCache *AC;
   const Instruction *CxtI;
   const DominatorTree *DT;
 
   // Unlike the other analyses, this may be a nullptr because not all clients
   // provide it currently.
   OptimizationRemarkEmitter *ORE;
 
+  // Optional cache.
+  KnownBitsCache *Cache = nullptr;
+
   /// If true, it is safe to use metadata during simplification.
   InstrInfoQuery IIQ;
 
   Query(const DataLayout &DL, AssumptionCache *AC, const Instruction *CxtI,
         const DominatorTree *DT, bool UseInstrInfo,
-        OptimizationRemarkEmitter *ORE = nullptr)
-      : DL(DL), AC(AC), CxtI(CxtI), DT(DT), ORE(ORE), IIQ(UseInstrInfo) {}
+        OptimizationRemarkEmitter *ORE = nullptr,
+        KnownBitsCache *Cache = nullptr)
+      : DL(DL),
+        AC(AC),
+        CxtI(CxtI),
+        DT(DT),
+        ORE(ORE),
+        Cache(Cache),
+        IIQ(UseInstrInfo) {}
 };
 
 } // end anonymous namespace
 
 // Given the provided Value and, potentially, a context instruction, return
 // the preferred context instruction (if any).
 static const Instruction *safeCxtI(const Value *V, const Instruction *CxtI) {
   // If we've been provided with a context instruction, then use that (provided
APInt DemandedElts =
       FVTy ? APInt::getAllOnes(FVTy->getNumElements()) : APInt(1, 1);
   computeKnownBits(V, DemandedElts, Known, Depth, Q);
 }
 
 void llvm::computeKnownBits(const Value *V, KnownBits &Known,
                             const DataLayout &DL, unsigned Depth,
                             AssumptionCache *AC, const Instruction *CxtI,
                             const DominatorTree *DT,
-                            OptimizationRemarkEmitter *ORE, bool UseInstrInfo) {
+                            OptimizationRemarkEmitter *ORE, bool UseInstrInfo,
+                            KnownBitsCache *Cache) {
   ::computeKnownBits(V, Known, Depth,
-                     Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE));
+                     Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE, Cache));
 }
 
 void llvm::computeKnownBits(const Value *V, const APInt &DemandedElts,
                             KnownBits &Known, const DataLayout &DL,
                             unsigned Depth, AssumptionCache *AC,
                             const Instruction *CxtI, const DominatorTree *DT,
-                            OptimizationRemarkEmitter *ORE, bool UseInstrInfo) {
+                            OptimizationRemarkEmitter *ORE, bool UseInstrInfo,
+                            KnownBitsCache *Cache) {
   ::computeKnownBits(V, DemandedElts, Known, Depth,
-                     Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE));
+                     Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE, Cache));
 }
 
 static KnownBits computeKnownBits(const Value *V, const APInt &DemandedElts,
                                   unsigned Depth, const Query &Q);
 
 static KnownBits computeKnownBits(const Value *V, unsigned Depth,
                                   const Query &Q);
 
 KnownBits llvm::computeKnownBits(const Value *V, const DataLayout &DL,
                                  unsigned Depth, AssumptionCache *AC,
                                  const Instruction *CxtI,
                                  const DominatorTree *DT,
                                  OptimizationRemarkEmitter *ORE,
-                                 bool UseInstrInfo) {
+                                 bool UseInstrInfo, KnownBitsCache *Cache) {
   return ::computeKnownBits(
-      V, Depth, Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE));
+      V, Depth, Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE, Cache));
 }
 
 KnownBits llvm::computeKnownBits(const Value *V, const APInt &DemandedElts,
                                  const DataLayout &DL, unsigned Depth,
                                  AssumptionCache *AC, const Instruction *CxtI,
                                  const DominatorTree *DT,
                                  OptimizationRemarkEmitter *ORE,
-                                 bool UseInstrInfo) {
+                                 bool UseInstrInfo, KnownBitsCache* Cache) {
   return ::computeKnownBits(
       V, DemandedElts, Depth,
-      Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE));
+      Query(DL, AC, safeCxtI(V, CxtI), DT, UseInstrInfo, ORE, Cache));
 }
 
 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");
if (ScalarTy->isPointerTy()) {
     assert(BitWidth == Q.DL.getPointerTypeSizeInBits(ScalarTy) &&
            "V and Known should have same BitWidth");
   } else {
     assert(BitWidth == Q.DL.getTypeSizeInBits(ScalarTy) &&
            "V and Known should have same BitWidth");
   }
 #endif
 
+  bool UseCache = Q.Cache && DemandedElts.isAllOnes();
+  auto CacheKey = std::make_pair(V, Q.CxtI ? Q.CxtI->getParent() : nullptr);
+  if (UseCache) {
+    auto it = Q.Cache->find(CacheKey);
+    if (it != Q.Cache->end()) {
+      Known = it->second;
+      return;
+    }
+  }
+
+  // Insert `Known` into the cache when this call to computeKnownBits returns.
+  Cleanup InsertIntoCacheRAII = [&] {
+    if (UseCache) {
+      // We can't reuse the `it` that we looked up at the beginning of this
+      // function, because Q.Cache is a DenseMap, which does not guarantee
+      // iterator stability.
+      auto [_, inserted] = Q.Cache->insert({CacheKey, Known});
+      (void) inserted;
+      assert(inserted);
+    }
+  };
+
   const APInt *C;
   if (match(V, m_APInt(C))) {
     // We know all of the bits for a scalar constant or a splat vector constant!
     Known = KnownBits::makeConstant(*C);
     return;
   }
   // Null and aggregate-zero are all-zeros.
   if (isa<ConstantPointerNull>(V) || isa<ConstantAggregateZero>(V)) {