[RS4GC] Cache BDVs and bases alogn with IsKnownBase flag (NFC)

[dmakogon]

Summary

This refactors RS4GC to cache results returned by findBaseDefiningValue. That function performs a recursive search for a base or at least a BDV for a given value.
Imagine we have the following IR:

%b = call i8* @foo()
%d1 = getelementptr %b, 1
%d2 = getelementptr %d1, 1

and imagine we call findBaseDefiningValue on %d2. The function would make a recursive call on gep operand which is %d1 and then another one on %b, then it'll return %b as base value for %d2. If later we call it on %d1, we would do the same recursive call on %b to again infer that it is a base value.
With this change findBaseDefiningValue caches returned result on each recursive call, so in the example above when being called on %d1 we'd already know that it's base is %b from the cache and wouldn't make a recursive call.

This patch also gets rid of BaseDefiningValueResult by caching the IsKnownBase flags for BDVs and bases.

[mkazantsev]

General questions: there are multiple places where we call eraseFromParent. Is there a guarantee that after this the cache is valid? Any way to verify this?

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

291

Remove?

[dmakogon]

Removed comment with debug prints

[dmakogon]

In D125000#3502849, @mkazantsev wrote:

General questions: there are multiple places where we call eraseFromParent. Is there a guarantee that after this the cache is valid? Any way to verify this?

Actually we never erase any instructions that may happen to be in the cache. In some places where eraseFromParent is called we do an explicit assert that the value to be erased is not in the cache.
In other places the values being erased are definitely not present in the cache: there's a stage when some dummy calls are inserted and then erased shortly after. And last place where the values are erased is when we replace the existing statepoints with new ones with new live variables.

[mkazantsev]

It looks like all you want BaseDefiningValueResult for is simply keeping one boolean for each Value you know the answer for. To me it looks like you simply need a map Value -> bool, and isKnownBaseResult should just return result from this map. After that, we don't need the structure BaseDefiningValueResult at all.

What I don't like in current approach is that we are not guarded against situation when we have two BaseDefiningValueResult instances with same value but different booleans, and it's not what we want in reality. What we want is map.

Am I missing something here?

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

128

nit: ///, \p V.

250

if (MustBeBase) assert(IsKnownBase)? I think it's more understandable.

289

Why not just return LHS == RHS?

[mkazantsev]

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

260–261

I guess this comment above would need an update.

[dmakogon]

Removed BaseDefiningValueResult struct. Instead we now have a cache which stores whether a base or BDV is a known base.

[mkazantsev]

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

476–481

Should it be findBaseDefiningValueCached? Otherwise you don't benefit from caching in recursion.

543–546

nit: tab

[mkazantsev]

LGTM with nit.

[mkazantsev]

Pls assert that you don't overwrite values in your cache.

[mkazantsev]

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

741

const IsKnownBase ... ?

742

Pls add assert with descriptive message what happens if the value is missing

[mkazantsev]

Withdrawing LGTM until asserts are added. I also think you should shortcut through cached method, unless there are reasons not to.

[dmakogon]

Implemented caching in findBaseDefiningValue. Added a weak assert in setKnownBase - it checks if the key is already present in the map and if so fails the assert if trying to push a different value from the cached one.

[dmakogon]

const qualifier in isKnownBase

[mkazantsev]

That looks much safer now! Thanks. LGTM.

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

442–446

Constants are so trivial that we might not cache them at all. But it's up to you.

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

 /// attributes and metadata. gc.statepoint can touch the entire heap including
 /// noalias objects.
 /// Apart from attributes and metadata, we also remove instructions that imply
 /// constant physical memory: llvm.invariant.start.
 static void stripNonValidData(Module &M);
 
 static bool shouldRewriteStatepointsIn(Function &F);
 
 PreservedAnalyses RewriteStatepointsForGC::run(Module &M,
                                                ModuleAnalysisManager &AM) {
   bool Changed = false;
   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   for (Function &F : M) {
     // Nothing to do for declarations.
     if (F.isDeclaration() || F.empty())
       continue;
 
struct GCPtrLivenessData {
   /// instruction in this basic block or ones reachable from here)
   MapVector<BasicBlock *, SetVector<Value *>> LiveIn;
 
   /// Values live out of this basic block (i.e. live into
   /// any successor block)
   MapVector<BasicBlock *, SetVector<Value *>> LiveOut;
 };
 
 // The type of the internal cache used inside the findBasePointers family
 // of functions.  From the callers perspective, this is an opaque type and
 // should not be inspected.
 //
 // In the actual implementation this caches two relations:
 // - The base relation itself (i.e. this pointer is based on that one)
 // - The base defining value relation (i.e. before base_phi insertion)
 // Generally, after the execution of a full findBasePointer call, only the
 // base relation will remain.  Internally, we add a mixture of the two
 // types, then update all the second type to the first type
 using DefiningValueMapTy = MapVector<Value *, Value *>;
+using IsKnownBaseMapTy = MapVector<Value *, bool>;
 using PointerToBaseTy = MapVector<Value *, Value *>;
 using StatepointLiveSetTy = SetVector<Value *>;
 using RematerializedValueMapTy =
     MapVector<AssertingVH<Instruction>, AssertingVH<Value>>;
 
 struct PartiallyConstructedSafepointRecord {
   /// The set of values known to be live across this safepoint
   StatepointLiveSetTy LiveSet;
struct PartiallyConstructedSafepointRecord {
   /// Maps rematerialized copy to it's original value.
   RematerializedValueMapTy RematerializedValues;
 };
 
 struct RematerizlizationCandidateRecord {
   // Chain from derived pointer to base.
   SmallVector<Instruction *, 3> ChainToBase;
   // Original base.
   Value *RootOfChain;
   // Cost of chain.
   InstructionCost Cost;
 };
 using RematCandTy = MapVector<Value *, RematerizlizationCandidateRecord>;
 
 } // end anonymous namespace
 
 static ArrayRef<Use> GetDeoptBundleOperands(const CallBase *Call) {
   Optional<OperandBundleUse> DeoptBundle =
       Call->getOperandBundle(LLVMContext::OB_deopt);
static void analyzeParsePointLiveness(
   }
   if (PrintLiveSetSize) {
     dbgs() << "Safepoint For: " << Call->getCalledOperand()->getName() << "\n";
     dbgs() << "Number live values: " << LiveSet.size() << "\n";
   }
   Result.LiveSet = LiveSet;
 }
 
-// Returns true is V is a knownBaseResult.
-static bool isKnownBaseResult(Value *V);
-
 // Returns true if V is a BaseResult that already exists in the IR, i.e. it is
 // not created by the findBasePointers algorithm.
 static bool isOriginalBaseResult(Value *V);
 
-namespace {
+/// Returns true if V is a known base.
+static bool isKnownBase(Value *V, const IsKnownBaseMapTy &KnownBases);
 
-/// A single base defining value - An immediate base defining value for an
-/// instruction 'Def' is an input to 'Def' whose base is also a base of 'Def'.
-/// For instructions which have multiple pointer [vector] inputs or that
-/// transition between vector and scalar types, there is no immediate base
-/// defining value.  The 'base defining value' for 'Def' is the transitive
-/// closure of this relation stopping at the first instruction which has no
-/// immediate base defining value.  The b.d.v. might itself be a base pointer,
-/// but it can also be an arbitrary derived pointer.
-struct BaseDefiningValueResult {
-  /// Contains the value which is the base defining value.
-  Value * const BDV;
-
-  /// True if the base defining value is also known to be an actual base
-  /// pointer.
-  const bool IsKnownBase;
-
-  BaseDefiningValueResult(Value *BDV, bool IsKnownBase)
-    : BDV(BDV), IsKnownBase(IsKnownBase) {
-#ifndef NDEBUG
-    // Check consistency between new and old means of checking whether a BDV is
-    // a base.
-    bool MustBeBase = isKnownBaseResult(BDV);
-    assert(!MustBeBase || MustBeBase == IsKnownBase);
-#endif
-  }
-};
+/// Caches the IsKnownBase flag for a value and asserts that it wasn't present
+/// in the cache before.
+static void setKnownBase(Value *V, bool IsKnownBase,
+                         IsKnownBaseMapTy &KnownBases);
 
-} // end anonymous namespace
-
-static BaseDefiningValueResult findBaseDefiningValue(Value *I);
+static Value *findBaseDefiningValue(Value *I, DefiningValueMapTy &Cache,
+                                    IsKnownBaseMapTy &KnownBases);
 
 /// Return a base defining value for the 'Index' element of the given vector
 /// instruction 'I'.  If Index is null, returns a BDV for the entire vector
 /// 'I'.  As an optimization, this method will try to determine when the
 /// element is known to already be a base pointer.  If this can be established,
 /// the second value in the returned pair will be true.  Note that either a
 /// vector or a pointer typed value can be returned.  For the former, the
 /// vector returned is a BDV (and possibly a base) of the entire vector 'I'.
 /// If the later, the return pointer is a BDV (or possibly a base) for the
 /// particular element in 'I'.
-static BaseDefiningValueResult
-findBaseDefiningValueOfVector(Value *I) {
+static Value *findBaseDefiningValueOfVector(Value *I, DefiningValueMapTy &Cache,
+                                            IsKnownBaseMapTy &KnownBases) {
   // Each case parallels findBaseDefiningValue below, see that code for
   // detailed motivation.
 
-  if (isa<Argument>(I))
+  auto Cached = Cache.find(I);
+  if (Cached != Cache.end())
+    return Cached->second;
+
+  if (isa<Argument>(I)) {
     // An incoming argument to the function is a base pointer
-    return BaseDefiningValueResult(I, true);
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
-  if (isa<Constant>(I))
+  if (isa<Constant>(I)) {
     // Base of constant vector consists only of constant null pointers.
     // For reasoning see similar case inside 'findBaseDefiningValue' function.
-    return BaseDefiningValueResult(ConstantAggregateZero::get(I->getType()),
-                                   true);
+    auto *CAZ = ConstantAggregateZero::get(I->getType());
+    Cache[I] = CAZ;
+    setKnownBase(CAZ, /* IsKnownBase */true, KnownBases);
+    return CAZ;
+  }
 
-  if (isa<LoadInst>(I))
-    return BaseDefiningValueResult(I, true);
+  if (isa<LoadInst>(I)) {
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
-  if (isa<InsertElementInst>(I))
+  if (isa<InsertElementInst>(I)) {
     // We don't know whether this vector contains entirely base pointers or
     // not.  To be conservatively correct, we treat it as a BDV and will
     // duplicate code as needed to construct a parallel vector of bases.
-    return BaseDefiningValueResult(I, false);
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */false, KnownBases);
+    return I;
+  }
 
-  if (isa<ShuffleVectorInst>(I))
+  if (isa<ShuffleVectorInst>(I)) {
     // We don't know whether this vector contains entirely base pointers or
     // not.  To be conservatively correct, we treat it as a BDV and will
     // duplicate code as needed to construct a parallel vector of bases.
     // TODO: There a number of local optimizations which could be applied here
     // for particular sufflevector patterns.
-    return BaseDefiningValueResult(I, false);
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */false, KnownBases);
+    return I;
+  }
 
   // The behavior of getelementptr instructions is the same for vector and
   // non-vector data types.
-  if (auto *GEP = dyn_cast<GetElementPtrInst>(I))
-    return findBaseDefiningValue(GEP->getPointerOperand());
+  if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
+    auto *BDV =
+        findBaseDefiningValue(GEP->getPointerOperand(), Cache, KnownBases);
+    Cache[GEP] = BDV;
+    return BDV;
+  }
 
   // If the pointer comes through a bitcast of a vector of pointers to
   // a vector of another type of pointer, then look through the bitcast
-  if (auto *BC = dyn_cast<BitCastInst>(I))
-    return findBaseDefiningValue(BC->getOperand(0));
+  if (auto *BC = dyn_cast<BitCastInst>(I)) {
+    auto *BDV = findBaseDefiningValue(BC->getOperand(0), Cache, KnownBases);
+    Cache[BC] = BDV;
+    return BDV;
+  }
 
   // We assume that functions in the source language only return base
   // pointers.  This should probably be generalized via attributes to support
   // both source language and internal functions.
-  if (isa<CallInst>(I) || isa<InvokeInst>(I))
-    return BaseDefiningValueResult(I, true);
+  if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
   // A PHI or Select is a base defining value.  The outer findBasePointer
   // algorithm is responsible for constructing a base value for this BDV.
   assert((isa<SelectInst>(I) || isa<PHINode>(I)) &&
          "unknown vector instruction - no base found for vector element");
-  return BaseDefiningValueResult(I, false);
+  Cache[I] = I;
+  setKnownBase(I, /* IsKnownBase */false, KnownBases);
+  return I;
 }
 
 /// Helper function for findBasePointer - Will return a value which either a)
 /// defines the base pointer for the input, b) blocks the simple search
 /// (i.e. a PHI or Select of two derived pointers), or c) involves a change
 /// from pointer to vector type or back.
-static BaseDefiningValueResult findBaseDefiningValue(Value *I) {
+static Value *findBaseDefiningValue(Value *I, DefiningValueMapTy &Cache,
+                                    IsKnownBaseMapTy &KnownBases) {
   assert(I->getType()->isPtrOrPtrVectorTy() &&
          "Illegal to ask for the base pointer of a non-pointer type");
+  auto Cached = Cache.find(I);
+  if (Cached != Cache.end())
+    return Cached->second;
 
   if (I->getType()->isVectorTy())
-    return findBaseDefiningValueOfVector(I);
+    return findBaseDefiningValueOfVector(I, Cache, KnownBases);
 
-  if (isa<Argument>(I))
+  if (isa<Argument>(I)) {
     // An incoming argument to the function is a base pointer
     // We should have never reached here if this argument isn't an gc value
-    return BaseDefiningValueResult(I, true);
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
   if (isa<Constant>(I)) {
     // We assume that objects with a constant base (e.g. a global) can't move
     // and don't need to be reported to the collector because they are always
     // live. Besides global references, all kinds of constants (e.g. undef,
     // constant expressions, null pointers) can be introduced by the inliner or
     // the optimizer, especially on dynamically dead paths.
     // Here we treat all of them as having single null base. By doing this we
     // trying to avoid problems reporting various conflicts in a form of
     // "phi (const1, const2)" or "phi (const, regular gc ptr)".
     // See constant.ll file for relevant test cases.
 
-    return BaseDefiningValueResult(
-        ConstantPointerNull::get(cast<PointerType>(I->getType())), true);
+    auto *CPN = ConstantPointerNull::get(cast<PointerType>(I->getType()));
+    Cache[I] = CPN;
+    setKnownBase(CPN, /* IsKnownBase */true, KnownBases);
+    return CPN;
   }
 
   // inttoptrs in an integral address space are currently ill-defined.  We
   // treat them as defining base pointers here for consistency with the
   // constant rule above and because we don't really have a better semantic
   // to give them.  Note that the optimizer is always free to insert undefined
   // behavior on dynamically dead paths as well.
-  if (isa<IntToPtrInst>(I))
-    return BaseDefiningValueResult(I, true);
+  if (isa<IntToPtrInst>(I)) {
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
   if (CastInst *CI = dyn_cast<CastInst>(I)) {
     Value *Def = CI->stripPointerCasts();
     // If stripping pointer casts changes the address space there is an
     // addrspacecast in between.
     assert(cast<PointerType>(Def->getType())->getAddressSpace() ==
                cast<PointerType>(CI->getType())->getAddressSpace() &&
            "unsupported addrspacecast");
     // If we find a cast instruction here, it means we've found a cast which is
     // not simply a pointer cast (i.e. an inttoptr).  We don't know how to
     // handle int->ptr conversion.
     assert(!isa<CastInst>(Def) && "shouldn't find another cast here");
-    return findBaseDefiningValue(Def);
+    auto *BDV = findBaseDefiningValue(Def, Cache, KnownBases);
+    Cache[CI] = BDV;
+    return BDV;
   }
 
-  if (isa<LoadInst>(I))
+  if (isa<LoadInst>(I)) {
     // The value loaded is an gc base itself
-    return BaseDefiningValueResult(I, true);
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
-  if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(I))
+  if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(I)) {
     // The base of this GEP is the base
-    return findBaseDefiningValue(GEP->getPointerOperand());
+    auto *BDV =
+        findBaseDefiningValue(GEP->getPointerOperand(), Cache, KnownBases);
+    Cache[GEP] = BDV;
+    return BDV;
+  }
 
-  if (auto *Freeze = dyn_cast<FreezeInst>(I))
-    return findBaseDefiningValue(Freeze->getOperand(0));
+  if (auto *Freeze = dyn_cast<FreezeInst>(I)) {
+    auto *BDV = findBaseDefiningValue(Freeze->getOperand(0), Cache, KnownBases);
+    Cache[Freeze] = BDV;
+    return BDV;
+  }
 
   if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     switch (II->getIntrinsicID()) {
     default:
       // fall through to general call handling
       break;
     case Intrinsic::experimental_gc_statepoint:
       llvm_unreachable("statepoints don't produce pointers");
     case Intrinsic::experimental_gc_relocate:
       // Rerunning safepoint insertion after safepoints are already
       // inserted is not supported.  It could probably be made to work,
       // but why are you doing this?  There's no good reason.
       llvm_unreachable("repeat safepoint insertion is not supported");
     case Intrinsic::gcroot:
       // Currently, this mechanism hasn't been extended to work with gcroot.
       // There's no reason it couldn't be, but I haven't thought about the
       // implications much.
       llvm_unreachable(
           "interaction with the gcroot mechanism is not supported");
     case Intrinsic::experimental_gc_get_pointer_base:
-      return findBaseDefiningValue(II->getOperand(0));
+      auto *BDV = findBaseDefiningValue(II->getOperand(0), Cache, KnownBases);
+      Cache[II] = BDV;
+      return BDV;
     }
   }
   // We assume that functions in the source language only return base
   // pointers.  This should probably be generalized via attributes to support
   // both source language and internal functions.
-  if (isa<CallInst>(I) || isa<InvokeInst>(I))
-    return BaseDefiningValueResult(I, true);
+  if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
   // TODO: I have absolutely no idea how to implement this part yet.  It's not
   // necessarily hard, I just haven't really looked at it yet.
   assert(!isa<LandingPadInst>(I) && "Landing Pad is unimplemented");
 
-  if (isa<AtomicCmpXchgInst>(I))
+  if (isa<AtomicCmpXchgInst>(I)) {
     // A CAS is effectively a atomic store and load combined under a
     // predicate.  From the perspective of base pointers, we just treat it
     // like a load.
-    return BaseDefiningValueResult(I, true);
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
   assert(!isa<AtomicRMWInst>(I) && "Xchg handled above, all others are "
                                    "binary ops which don't apply to pointers");
 
   // The aggregate ops.  Aggregates can either be in the heap or on the
   // stack, but in either case, this is simply a field load.  As a result,
   // this is a defining definition of the base just like a load is.
-  if (isa<ExtractValueInst>(I))
-    return BaseDefiningValueResult(I, true);
+  if (isa<ExtractValueInst>(I)) {
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */true, KnownBases);
+    return I;
+  }
 
   // We should never see an insert vector since that would require we be
   // tracing back a struct value not a pointer value.
   assert(!isa<InsertValueInst>(I) &&
          "Base pointer for a struct is meaningless");
 
   // This value might have been generated by findBasePointer() called when
   // substituting gc.get.pointer.base() intrinsic.
   bool IsKnownBase =
       isa<Instruction>(I) && cast<Instruction>(I)->getMetadata("is_base_value");
+  setKnownBase(I, /* IsKnownBase */IsKnownBase, KnownBases);
+  Cache[I] = I;
 
   // An extractelement produces a base result exactly when it's input does.
   // We may need to insert a parallel instruction to extract the appropriate
   // element out of the base vector corresponding to the input. Given this,
   // it's analogous to the phi and select case even though it's not a merge.
   if (isa<ExtractElementInst>(I))
     // Note: There a lot of obvious peephole cases here.  This are deliberately
     // handled after the main base pointer inference algorithm to make writing
     // test cases to exercise that code easier.
-    return BaseDefiningValueResult(I, IsKnownBase);
+    return I;
 
   // The last two cases here don't return a base pointer.  Instead, they
   // return a value which dynamically selects from among several base
   // derived pointers (each with it's own base potentially).  It's the job of
   // the caller to resolve these.
   assert((isa<SelectInst>(I) || isa<PHINode>(I)) &&
          "missing instruction case in findBaseDefiningValue");
-  return BaseDefiningValueResult(I, IsKnownBase);
+  return I;
 }
 
 /// Returns the base defining value for this value.
-static Value *findBaseDefiningValueCached(Value *I, DefiningValueMapTy &Cache) {
-  Value *&Cached = Cache[I];
-  if (!Cached) {
-    Cached = findBaseDefiningValue(I).BDV;
+static Value *findBaseDefiningValueCached(Value *I, DefiningValueMapTy &Cache,
+                                          IsKnownBaseMapTy &KnownBases) {
+  if (Cache.find(I) == Cache.end()) {
+    auto *BDV = findBaseDefiningValue(I, Cache, KnownBases);
+    Cache[I] = BDV;
     LLVM_DEBUG(dbgs() << "fBDV-cached: " << I->getName() << " -> "
-                      << Cached->getName() << "\n");
+                      << Cache[I]->getName() << ", is known base = "
+                      << KnownBases[I] << "\n");
   }
   assert(Cache[I] != nullptr);
-  return Cached;
+  assert(KnownBases.find(Cache[I]) != KnownBases.end() &&
+         "Cached value must be present in known bases map");
+  return Cache[I];
 }
 
 /// Return a base pointer for this value if known.  Otherwise, return it's
 /// base defining value.
-static Value *findBaseOrBDV(Value *I, DefiningValueMapTy &Cache) {
-  Value *Def = findBaseDefiningValueCached(I, Cache);
+static Value *findBaseOrBDV(Value *I, DefiningValueMapTy &Cache,
+                            IsKnownBaseMapTy &KnownBases) {
+  Value *Def = findBaseDefiningValueCached(I, Cache, KnownBases);
   auto Found = Cache.find(Def);
   if (Found != Cache.end()) {
     // Either a base-of relation, or a self reference.  Caller must check.
     return Found->second;
   }
   // Only a BDV available
   return Def;
 }
 
 /// This value is a base pointer that is not generated by RS4GC, i.e. it already
 /// exists in the code.
 static bool isOriginalBaseResult(Value *V) {
   // no recursion possible
   return !isa<PHINode>(V) && !isa<SelectInst>(V) &&
          !isa<ExtractElementInst>(V) && !isa<InsertElementInst>(V) &&
          !isa<ShuffleVectorInst>(V);
 }
 
-/// Given the result of a call to findBaseDefiningValue, or findBaseOrBDV,
-/// is it known to be a base pointer?  Or do we need to continue searching.
-static bool isKnownBaseResult(Value *V) {
-  if (isOriginalBaseResult(V))
-    return true;
-  if (isa<Instruction>(V) &&
-      cast<Instruction>(V)->getMetadata("is_base_value")) {
-    // This is a previously inserted base phi or select.  We know
-    // that this is a base value.
-    return true;
-  }
+static bool isKnownBase(Value *V, const IsKnownBaseMapTy &KnownBases) {
+  auto It = KnownBases.find(V);
+  assert(It != KnownBases.end() && "Value not present in the map");
+  return It->second;
 
-  // We need to keep searching
-  return false;
+static void setKnownBase(Value *V, bool IsKnownBase,
+                         IsKnownBaseMapTy &KnownBases) {
+#ifndef NDEBUG
+  auto It = KnownBases.find(V);
+  if (It != KnownBases.end())
+    assert(It->second == IsKnownBase && "Changing already present value");
+#endif
+  KnownBases[V] = IsKnownBase;
 }
 
 // Returns true if First and Second values are both scalar or both vector.
 static bool areBothVectorOrScalar(Value *First, Value *Second) {
   return isa<VectorType>(First->getType()) ==
          isa<VectorType>(Second->getType());
 }
 
 namespace {
 
 /// Models the state of a single base defining value in the findBasePointer
 /// algorithm for determining where a new instruction is needed to propagate
static raw_ostream &operator<<(raw_ostream &OS, const BDVState &State) {
   return OS;
 }
 #endif
 
 /// For a given value or instruction, figure out what base ptr its derived from.
 /// For gc objects, this is simply itself.  On success, returns a value which is
 /// the base pointer.  (This is reliable and can be used for relocation.)  On
 /// failure, returns nullptr.
-static Value *findBasePointer(Value *I, DefiningValueMapTy &Cache) {
-  Value *Def = findBaseOrBDV(I, Cache);
+static Value *findBasePointer(Value *I, DefiningValueMapTy &Cache,
+                              IsKnownBaseMapTy &KnownBases) {
+  Value *Def = findBaseOrBDV(I, Cache, KnownBases);
 
-  if (isKnownBaseResult(Def) && areBothVectorOrScalar(Def, I))
+  if (isKnownBase(Def, KnownBases) && areBothVectorOrScalar(Def, I))
     return Def;
 
   // Here's the rough algorithm:
   // - For every SSA value, construct a mapping to either an actual base
   //   pointer or a PHI which obscures the base pointer.
   // - Construct a mapping from PHI to unknown TOP state.  Use an
   //   optimistic algorithm to propagate base pointer information.  Lattice
   //   looks like:
/* scope */ {
     SmallVector<Value*, 16> Worklist;
     Worklist.push_back(Def);
     States.insert({Def, BDVState(Def)});
     while (!Worklist.empty()) {
       Value *Current = Worklist.pop_back_val();
       assert(!isOriginalBaseResult(Current) && "why did it get added?");
 
       auto visitIncomingValue = [&](Value *InVal) {
-        Value *Base = findBaseOrBDV(InVal, Cache);
-        if (isKnownBaseResult(Base) && areBothVectorOrScalar(Base, InVal))
+        Value *Base = findBaseOrBDV(InVal, Cache, KnownBases);
+        if (isKnownBase(Base, KnownBases) && areBothVectorOrScalar(Base, InVal))
           // Known bases won't need new instructions introduced and can be
           // ignored safely. However, this can only be done when InVal and Base
           // are both scalar or both vector. Otherwise, we need to find a
           // correct BDV for InVal, by creating an entry in the lattice
           // (States).
           return;
         assert(isExpectedBDVType(Base) && "the only non-base values "
                "we see should be base defining values");
do {
     ToRemove.clear();
     for (auto Pair : States) {
       Value *BDV = Pair.first;
       auto canPruneInput = [&](Value *V) {
         // If the input of the BDV is the BDV itself we can prune it. This is
         // only possible if the BDV is a PHI node.
         if (V->stripPointerCasts() == BDV)
           return true;
-        Value *VBDV = findBaseOrBDV(V, Cache);
+        Value *VBDV = findBaseOrBDV(V, Cache, KnownBases);
         if (V->stripPointerCasts() != VBDV)
           return false;
         // The assumption is that anything not in the state list is
         // propagates a base pointer.
         return States.count(VBDV) == 0;
       };
 
       bool CanPrune = true;
#endif
     // Since this is computing a fixed point, the order of visit does not
     // effect the result.  TODO: We could use a worklist here and make this run
     // much faster.
     for (auto Pair : States) {
       Value *BDV = Pair.first;
       // Only values that do not have known bases or those that have differing
       // type (scalar versus vector) from a possible known base should be in the
       // lattice.
-      assert((!isKnownBaseResult(BDV) ||
+      assert((!isKnownBase(BDV, KnownBases) ||
              !areBothVectorOrScalar(BDV, Pair.second.getBaseValue())) &&
                  "why did it get added?");
 
       BDVState NewState(BDV);
       visitBDVOperands(BDV, [&](Value *Op) {
-        Value *BDV = findBaseOrBDV(Op, Cache);
+        Value *BDV = findBaseOrBDV(Op, Cache, KnownBases);
         auto OpState = GetStateForBDV(BDV, Op);
         NewState.meet(OpState);
       });
 
       BDVState OldState = States[BDV];
       if (OldState != NewState) {
         Progress = true;
         States[BDV] = NewState;
#endif
   // is of scalar type.
   for (auto Pair : States) {
     Instruction *I = cast<Instruction>(Pair.first);
     BDVState State = Pair.second;
     auto *BaseValue = State.getBaseValue();
     // Only values that do not have known bases or those that have differing
     // type (scalar versus vector) from a possible known base should be in the
     // lattice.
-    assert((!isKnownBaseResult(I) || !areBothVectorOrScalar(I, BaseValue)) &&
-           "why did it get added?");
+    assert(
+        (!isKnownBase(I, KnownBases) || !areBothVectorOrScalar(I, BaseValue)) &&
     assert(!State.isUnknown() && "Optimistic algorithm didn't complete!");
 
     if (!State.isBase() || !isa<VectorType>(BaseValue->getType()))
       continue;
     // extractelement instructions are a bit special in that we may need to
     // insert an extract even when we know an exact base for the instruction.
     // The problem is that we need to convert from a vector base to a scalar
     // base for the particular indice we're interested in.
     if (isa<ExtractElementInst>(I)) {
       auto *EE = cast<ExtractElementInst>(I);
       // TODO: In many cases, the new instruction is just EE itself.  We should
       // exploit this, but can't do it here since it would break the invariant
       // about the BDV not being known to be a base.
       auto *BaseInst = ExtractElementInst::Create(
           State.getBaseValue(), EE->getIndexOperand(), "base_ee", EE);
       BaseInst->setMetadata("is_base_value", MDNode::get(I->getContext(), {}));
       States[I] = BDVState(I, BDVState::Base, BaseInst);
+      setKnownBase(BaseInst, /* IsKnownBase */true, KnownBases);
     } else if (!isa<VectorType>(I->getType())) {
       // We need to handle cases that have a vector base but the instruction is
       // a scalar type (these could be phis or selects or any instruction that
       // are of scalar type, but the base can be a vector type).  We
       // conservatively set this as conflict.  Setting the base value for these
       // conflicts is handled in the next loop which traverses States.
       States[I] = BDVState(I, BDVState::Conflict);
     }
   }
 
 #ifndef NDEBUG
   VerifyStates();
 #endif
 
   // Insert Phis for all conflicts
   // TODO: adjust naming patterns to avoid this order of iteration dependency
   for (auto Pair : States) {
     Instruction *I = cast<Instruction>(Pair.first);
     BDVState State = Pair.second;
     // Only values that do not have known bases or those that have differing
     // type (scalar versus vector) from a possible known base should be in the
     // lattice.
-    assert((!isKnownBaseResult(I) || !areBothVectorOrScalar(I, State.getBaseValue())) &&
+    assert((!isKnownBase(I, KnownBases) ||
+            !areBothVectorOrScalar(I, State.getBaseValue())) &&
            "why did it get added?");
     assert(!State.isUnknown() && "Optimistic algorithm didn't complete!");
 
     // Since we're joining a vector and scalar base, they can never be the
     // same.  As a result, we should always see insert element having reached
     // the conflict state.
     assert(!isa<InsertElementInst>(I) || State.isConflict());
 
for (auto Pair : States) {
     };
 
     Instruction *BaseInst = I->clone();
     BaseInst->insertBefore(I);
     BaseInst->setName(getMangledName(I));
     // Add metadata marking this as a base value
     BaseInst->setMetadata("is_base_value", MDNode::get(I->getContext(), {}));
     States[I] = BDVState(I, BDVState::Conflict, BaseInst);
+    setKnownBase(BaseInst, /* IsKnownBase */true, KnownBases);
   }
 
 #ifndef NDEBUG
   VerifyStates();
 #endif
 
   // Returns a instruction which produces the base pointer for a given
   // instruction.  The instruction is assumed to be an input to one of the BDVs
   // seen in the inference algorithm above.  As such, we must either already
   // know it's base defining value is a base, or have inserted a new
   // instruction to propagate the base of it's BDV and have entered that newly
   // introduced instruction into the state table.  In either case, we are
   // assured to be able to determine an instruction which produces it's base
   // pointer.
   auto getBaseForInput = [&](Value *Input, Instruction *InsertPt) {
-    Value *BDV = findBaseOrBDV(Input, Cache);
+    Value *BDV = findBaseOrBDV(Input, Cache, KnownBases);
     Value *Base = nullptr;
     if (!States.count(BDV)) {
       assert(areBothVectorOrScalar(BDV, Input));
       Base = BDV;
     } else {
       // Either conflict or base.
       assert(States.count(BDV));
       Base = States[BDV].getBaseValue();
#endif
   // instructions.
   for (auto Pair : States) {
     Instruction *BDV = cast<Instruction>(Pair.first);
     BDVState State = Pair.second;
 
     // Only values that do not have known bases or those that have differing
     // type (scalar versus vector) from a possible known base should be in the
     // lattice.
-    assert((!isKnownBaseResult(BDV) ||
+    assert((!isKnownBase(BDV, KnownBases) ||
             !areBothVectorOrScalar(BDV, State.getBaseValue())) &&
            "why did it get added?");
     assert(!State.isUnknown() && "Optimistic algorithm didn't complete!");
     if (!State.isConflict())
       continue;
 
     if (PHINode *BasePHI = dyn_cast<PHINode>(State.getBaseValue())) {
       PHINode *PN = cast<PHINode>(BDV);
#endif
   // relation and one of the base pointer relation!  FIXME
   for (auto Pair : States) {
     auto *BDV = Pair.first;
     Value *Base = Pair.second.getBaseValue();
     assert(BDV && Base);
     // Only values that do not have known bases or those that have differing
     // type (scalar versus vector) from a possible known base should be in the
     // lattice.
-    assert((!isKnownBaseResult(BDV) || !areBothVectorOrScalar(BDV, Base)) &&
-           "why did it get added?");
+    assert(
+        (!isKnownBase(BDV, KnownBases) || !areBothVectorOrScalar(BDV, Base)) &&
 
     LLVM_DEBUG(
         dbgs() << "Updating base value cache"
                << " for: " << BDV->getName() << " from: "
                << (Cache.count(BDV) ? Cache[BDV]->getName().str() : "none")
                << " to: " << Base->getName() << "\n");
 
     Cache[BDV] = Base;

 // side effects: may insert PHI nodes into the existing CFG, will preserve
 // CFG, will not remove or mutate any existing nodes
 //
 // post condition: PointerToBase contains one (derived, base) pair for every
 // pointer in live.  Note that derived can be equal to base if the original
 // pointer was a base pointer.
 static void findBasePointers(const StatepointLiveSetTy &live,
                              PointerToBaseTy &PointerToBase, DominatorTree *DT,
-                             DefiningValueMapTy &DVCache) {
+                             DefiningValueMapTy &DVCache,
+                             IsKnownBaseMapTy &KnownBases) {
   for (Value *ptr : live) {
-    Value *base = findBasePointer(ptr, DVCache);
+    Value *base = findBasePointer(ptr, DVCache, KnownBases);
     assert(base && "failed to find base pointer");
     PointerToBase[ptr] = base;
     assert((!isa<Instruction>(base) || !isa<Instruction>(ptr) ||
             DT->dominates(cast<Instruction>(base)->getParent(),
                           cast<Instruction>(ptr)->getParent())) &&
            "The base we found better dominate the derived pointer");
   }
 }
 
 /// Find the required based pointers (and adjust the live set) for the given
 /// parse point.
 static void findBasePointers(DominatorTree &DT, DefiningValueMapTy &DVCache,
                              CallBase *Call,
                              PartiallyConstructedSafepointRecord &result,
-                             PointerToBaseTy &PointerToBase) {
+                             PointerToBaseTy &PointerToBase,
+                             IsKnownBaseMapTy &KnownBases) {
   StatepointLiveSetTy PotentiallyDerivedPointers = result.LiveSet;
   // We assume that all pointers passed to deopt are base pointers; as an
   // optimization, we can use this to avoid seperately materializing the base
   // pointer graph.  This is only relevant since we're very conservative about
   // generating new conflict nodes during base pointer insertion.  If we were
   // smarter there, this would be irrelevant.
   if (auto Opt = Call->getOperandBundle(LLVMContext::OB_deopt))
     for (Value *V : Opt->Inputs) {
       if (!PotentiallyDerivedPointers.count(V))
         continue;
       PotentiallyDerivedPointers.remove(V);
       PointerToBase[V] = V;
     }
-  findBasePointers(PotentiallyDerivedPointers, PointerToBase, &DT, DVCache);
+  findBasePointers(PotentiallyDerivedPointers, PointerToBase, &DT, DVCache,
+                   KnownBases);
 }
 
 /// Given an updated version of the dataflow liveness results, update the
 /// liveset and base pointer maps for the call site CS.
 static void recomputeLiveInValues(GCPtrLivenessData &RevisedLivenessData,
                                   CallBase *Call,
                                   PartiallyConstructedSafepointRecord &result,
                                   PointerToBaseTy &PointerToBase);
#endif
   // Remove rematerializaed values from the live set
   for (auto LiveValue: LiveValuesToBeDeleted) {
     Info.LiveSet.remove(LiveValue);
   }
 }
 
 static bool inlineGetBaseAndOffset(Function &F,
                                    SmallVectorImpl<CallInst *> &Intrinsics,
-                                   DefiningValueMapTy &DVCache) {
+                                   DefiningValueMapTy &DVCache,
+                                   IsKnownBaseMapTy &KnownBases) {
   auto &Context = F.getContext();
   auto &DL = F.getParent()->getDataLayout();
   bool Changed = false;
 
   for (auto *Callsite : Intrinsics)
     switch (Callsite->getIntrinsicID()) {
     case Intrinsic::experimental_gc_get_pointer_base: {
       Changed = true;
-      Value *Base = findBasePointer(Callsite->getOperand(0), DVCache);
+      Value *Base =
+          findBasePointer(Callsite->getOperand(0), DVCache, KnownBases);
       assert(!DVCache.count(Callsite));
       auto *BaseBC = IRBuilder<>(Callsite).CreateBitCast(
           Base, Callsite->getType(), suffixed_name_or(Base, ".cast", ""));
       if (BaseBC != Base)
         DVCache[BaseBC] = Base;
       Callsite->replaceAllUsesWith(BaseBC);
       if (!BaseBC->hasName())
         BaseBC->takeName(Callsite);
       Callsite->eraseFromParent();
       break;
     }
     case Intrinsic::experimental_gc_get_pointer_offset: {
       Changed = true;
       Value *Derived = Callsite->getOperand(0);
-      Value *Base = findBasePointer(Derived, DVCache);
+      Value *Base = findBasePointer(Derived, DVCache, KnownBases);
       assert(!DVCache.count(Callsite));
       unsigned AddressSpace = Derived->getType()->getPointerAddressSpace();
       unsigned IntPtrSize = DL.getPointerSizeInBits(AddressSpace);
       IRBuilder<> Builder(Callsite);
       Value *BaseInt =
           Builder.CreatePtrToInt(Base, Type::getIntNTy(Context, IntPtrSize),
                                  suffixed_name_or(Base, ".int", ""));
       Value *DerivedInt =
for (auto *Callsite : Intrinsics)
     }
 
   return Changed;
 }
 
 static bool insertParsePoints(Function &F, DominatorTree &DT,
                               TargetTransformInfo &TTI,
                               SmallVectorImpl<CallBase *> &ToUpdate,
-                              DefiningValueMapTy &DVCache) {
+                              DefiningValueMapTy &DVCache,
+                              IsKnownBaseMapTy &KnownBases) {
 #ifndef NDEBUG
   // Validate the input
   std::set<CallBase *> Uniqued;
   Uniqued.insert(ToUpdate.begin(), ToUpdate.end());
   assert(Uniqued.size() == ToUpdate.size() && "no duplicates please!");
 
   for (CallBase *Call : ToUpdate)
     assert(Call->getFunction() == &F);
#endif
   findLiveReferences(F, DT, ToUpdate, Records);
 
   /// Global mapping from live pointers to a base-defining-value.
   PointerToBaseTy PointerToBase;
 
   // B) Find the base pointers for each live pointer
   for (size_t i = 0; i < Records.size(); i++) {
     PartiallyConstructedSafepointRecord &info = Records[i];
-    findBasePointers(DT, DVCache, ToUpdate[i], info, PointerToBase);
+    findBasePointers(DT, DVCache, ToUpdate[i], info, PointerToBase, KnownBases);
   }
   if (PrintBasePointers) {
     errs() << "Base Pairs (w/o Relocation):\n";
     for (auto &Pair : PointerToBase) {
       errs() << " derived ";
       Pair.first->printAsOperand(errs(), false);
       errs() << " base ";
       Pair.second->printAsOperand(errs(), false);
bool RewriteStatepointsForGC::runOnFunction(Function &F, DominatorTree &DT,
   }
 
   // Cache the 'defining value' relation used in the computation and
   // insertion of base phis and selects.  This ensures that we don't insert
   // large numbers of duplicate base_phis. Use one cache for both
   // inlineGetBaseAndOffset() and insertParsePoints().
   DefiningValueMapTy DVCache;
 
+  // Mapping between a base values and a flag indicating whether it's a known
+  // base or not.
+  IsKnownBaseMapTy KnownBases;
+
   if (!Intrinsics.empty())
     // Inline @gc.get.pointer.base() and @gc.get.pointer.offset() before finding
     // live references.
-    MadeChange |= inlineGetBaseAndOffset(F, Intrinsics, DVCache);
+    MadeChange |= inlineGetBaseAndOffset(F, Intrinsics, DVCache, KnownBases);
 
   if (!ParsePointNeeded.empty())
-    MadeChange |= insertParsePoints(F, DT, TTI, ParsePointNeeded, DVCache);
+    MadeChange |=
+        insertParsePoints(F, DT, TTI, ParsePointNeeded, DVCache, KnownBases);
 
   return MadeChange;
 }
 
 // liveness computation via standard dataflow
 // -------------------------------------------------------------------
 
 // TODO: Consider using bitvectors for liveness, the set of potentially