[RS4GC] Make PointerToBase mapping be global. NFC.

[skatkov]

Summary

PointerToBase is a mapping between potentially derived pointer to its base.
As soon as we are in SSA form If there is a base of derived pointer and it
is available at def of derived pointer, the same base will be available at any
point where derived pointer is alive.

So the mapping of derived pointer to base pointer is a global property, not
a property of a call site.

[reames]

LGTM, though can you update the commit comment not to call this "global"? That generally makes me thing global variable (i.e. static at TU level), not global to the function.

Once this lands, I think there might be some further cleanup available. But that can be done lazily as needed.

I'm somewhat shaking my head that we never noticed the reuse potential before. Thanks!

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

 //
 // 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 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;
 
-  /// Mapping from live pointers to a base-defining-value
-  MapVector<Value *, Value *> PointerToBase;
-
   /// The *new* gc.statepoint instruction itself.  This produces the token
   /// that normal path gc.relocates and the gc.result are tied to.
   GCStatepointInst *StatepointToken;
 
   /// Instruction to which exceptional gc relocates are attached
   /// Makes it easier to iterate through them during relocationViaAlloca.
   Instruction *UnwindToken;
 

 // preconditions: live is a set of pointer type Values
 //
 // 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,
-                 MapVector<Value *, Value *> &PointerToBase,
-                 DominatorTree *DT, DefiningValueMapTy &DVCache) {
+static void findBasePointers(const StatepointLiveSetTy &live,
+                             PointerToBaseTy &PointerToBase, DominatorTree *DT,
+                             DefiningValueMapTy &DVCache) {
   for (Value *ptr : live) {
     Value *base = findBasePointer(ptr, DVCache);
     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) {
-  MapVector<Value *, Value *> PointerToBase;
+                             PartiallyConstructedSafepointRecord &result,
+                             PointerToBaseTy &PointerToBase) {
   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);
-
-  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);
-      errs() << "\n";;
-    }
-  }
-
-  result.PointerToBase = PointerToBase;
 }
 
 /// 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);
+                                  PartiallyConstructedSafepointRecord &result,
+                                  PointerToBaseTy &PointerToBase);
 
 static void recomputeLiveInValues(
     Function &F, DominatorTree &DT, ArrayRef<CallBase *> toUpdate,
-    MutableArrayRef<struct PartiallyConstructedSafepointRecord> records) {
+    MutableArrayRef<struct PartiallyConstructedSafepointRecord> records,
+    PointerToBaseTy &PointerToBase) {
   // TODO-PERF: reuse the original liveness, then simply run the dataflow
   // again.  The old values are still live and will help it stabilize quickly.
   GCPtrLivenessData RevisedLivenessData;
   computeLiveInValues(DT, F, RevisedLivenessData);
   for (size_t i = 0; i < records.size(); i++) {
     struct PartiallyConstructedSafepointRecord &info = records[i];
-    recomputeLiveInValues(RevisedLivenessData, toUpdate[i], info);
+    recomputeLiveInValues(RevisedLivenessData, toUpdate[i], info,
+                          PointerToBase);
   }
 }
 
 // When inserting gc.relocate and gc.result calls, we need to ensure there are
 // no uses of the original value / return value between the gc.statepoint and
 // the gc.relocate / gc.result call.  One case which can arise is a phi node
 // starting one of the successor blocks.  We also need to be able to insert the
 // gc.relocates only on the path which goes through the statepoint.  We might
static StringRef getDeoptLowering(CallBase *Call) {
   return "live-through";
 }
 
 static void
 makeStatepointExplicitImpl(CallBase *Call, /* to replace */
                            const SmallVectorImpl<Value *> &BasePtrs,
                            const SmallVectorImpl<Value *> &LiveVariables,
                            PartiallyConstructedSafepointRecord &Result,
-                           std::vector<DeferredReplacement> &Replacements) {
+                           std::vector<DeferredReplacement> &Replacements,
+                           const PointerToBaseTy &PointerToBase) {
   assert(BasePtrs.size() == LiveVariables.size());
 
   // Then go ahead and use the builder do actually do the inserts.  We insert
   // immediately before the previous instruction under the assumption that all
   // arguments will be available here.  We can't insert afterwards since we may
   // be replacing a terminator.
   IRBuilder<> Builder(Call);
 
if (IID == Intrinsic::experimental_deoptimize) {
       // intrinsic should know the corresponding base pointers.
       //
       // To make the base pointers available pass them explicitly as arguments:
       //   memcpy(dest_derived, source_derived, ...) =>
       //   memcpy(dest_base, dest_offset, source_base, source_offset, ...)
       auto &Context = Call->getContext();
       auto &DL = Call->getModule()->getDataLayout();
       auto GetBaseAndOffset = [&](Value *Derived) {
-        assert(Result.PointerToBase.count(Derived));
+        assert(PointerToBase.count(Derived));
         unsigned AddressSpace = Derived->getType()->getPointerAddressSpace();
         unsigned IntPtrSize = DL.getPointerSizeInBits(AddressSpace);
-        Value *Base = Result.PointerToBase.find(Derived)->second;
+        Value *Base = PointerToBase.find(Derived)->second;
         Value *Base_int = Builder.CreatePtrToInt(
             Base, Type::getIntNTy(Context, IntPtrSize));
         Value *Derived_int = Builder.CreatePtrToInt(
             Derived, Type::getIntNTy(Context, IntPtrSize));
         return std::make_pair(Base, Builder.CreateSub(Derived_int, Base_int));
       };
 
       auto *Dest = CallArgs[0];

 // Replace an existing gc.statepoint with a new one and a set of gc.relocates
 // which make the relocations happening at this safepoint explicit.
 //
 // WARNING: Does not do any fixup to adjust users of the original live
 // values.  That's the callers responsibility.
 static void
 makeStatepointExplicit(DominatorTree &DT, CallBase *Call,
                        PartiallyConstructedSafepointRecord &Result,
-                       std::vector<DeferredReplacement> &Replacements) {
+                       std::vector<DeferredReplacement> &Replacements,
+                       const PointerToBaseTy &PointerToBase) {
   const auto &LiveSet = Result.LiveSet;
-  const auto &PointerToBase = Result.PointerToBase;
 
   // Convert to vector for efficient cross referencing.
   SmallVector<Value *, 64> BaseVec, LiveVec;
   LiveVec.reserve(LiveSet.size());
   BaseVec.reserve(LiveSet.size());
   for (Value *L : LiveSet) {
     LiveVec.push_back(L);
     assert(PointerToBase.count(L));
     Value *Base = PointerToBase.find(L)->second;
     BaseVec.push_back(Base);
   }
   assert(LiveVec.size() == BaseVec.size());
 
   // Do the actual rewriting and delete the old statepoint
-  makeStatepointExplicitImpl(Call, BaseVec, LiveVec, Result, Replacements);
+  makeStatepointExplicitImpl(Call, BaseVec, LiveVec, Result, Replacements,
+                             PointerToBase);
 }
 
 // Helper function for the relocationViaAlloca.
 //
 // It receives iterator to the statepoint gc relocates and emits a store to the
 // assigned location (via allocaMap) for the each one of them.  It adds the
 // visited values into the visitedLiveValues set, which we will later use them
 // for validation checking.

 }
 
 // From the statepoint live set pick values that are cheaper to recompute then
 // to relocate. Remove this values from the live set, rematerialize them after
 // statepoint and record them in "Info" structure. Note that similar to
 // relocated values we don't do any user adjustments here.
 static void rematerializeLiveValues(CallBase *Call,
                                     PartiallyConstructedSafepointRecord &Info,
+                                    PointerToBaseTy &PointerToBase,
                                     TargetTransformInfo &TTI) {
   const unsigned int ChainLengthThreshold = 10;
 
   // Record values we are going to delete from this statepoint live set.
   // We can not di this in following loop due to iterator invalidation.
   SmallVector<Value *, 32> LiveValuesToBeDeleted;
 
   for (Value *LiveValue: Info.LiveSet) {
     // For each live pointer find its defining chain
     SmallVector<Instruction *, 3> ChainToBase;
-    assert(Info.PointerToBase.count(LiveValue));
+    assert(PointerToBase.count(LiveValue));
     Value *RootOfChain =
       findRematerializableChainToBasePointer(ChainToBase,
                                              LiveValue);
 
     // Nothing to do, or chain is too long
     if ( ChainToBase.size() == 0 ||
         ChainToBase.size() > ChainLengthThreshold)
       continue;
 
     // Handle the scenario where the RootOfChain is not equal to the
     // Base Value, but they are essentially the same phi values.
-    if (RootOfChain != Info.PointerToBase[LiveValue]) {
+    if (RootOfChain != PointerToBase[LiveValue]) {
       PHINode *OrigRootPhi = dyn_cast<PHINode>(RootOfChain);
-      PHINode *AlternateRootPhi = dyn_cast<PHINode>(Info.PointerToBase[LiveValue]);
+      PHINode *AlternateRootPhi = dyn_cast<PHINode>(PointerToBase[LiveValue]);
       if (!OrigRootPhi || !AlternateRootPhi)
         continue;
       // PHI nodes that have the same incoming values, and belonging to the same
       // basic blocks are essentially the same SSA value.  When the original phi
       // has incoming values with different base pointers, the original phi is
       // marked as conflict, and an additional `AlternateRootPhi` with the same
       // incoming values get generated by the findBasePointer function. We need
       // to identify the newly generated AlternateRootPhi (.base version of phi)
#endif
     };
 
     // Different cases for calls and invokes. For invokes we need to clone
     // instructions both on normal and unwind path.
     if (isa<CallInst>(Call)) {
       Instruction *InsertBefore = Call->getNextNode();
       assert(InsertBefore);
       Instruction *RematerializedValue = rematerializeChain(
-          InsertBefore, RootOfChain, Info.PointerToBase[LiveValue]);
+          InsertBefore, RootOfChain, PointerToBase[LiveValue]);
       Info.RematerializedValues[RematerializedValue] = LiveValue;
     } else {
       auto *Invoke = cast<InvokeInst>(Call);
 
       Instruction *NormalInsertBefore =
           &*Invoke->getNormalDest()->getFirstInsertionPt();
       Instruction *UnwindInsertBefore =
           &*Invoke->getUnwindDest()->getFirstInsertionPt();
 
       Instruction *NormalRematerializedValue = rematerializeChain(
-          NormalInsertBefore, RootOfChain, Info.PointerToBase[LiveValue]);
+          NormalInsertBefore, RootOfChain, PointerToBase[LiveValue]);
       Instruction *UnwindRematerializedValue = rematerializeChain(
-          UnwindInsertBefore, RootOfChain, Info.PointerToBase[LiveValue]);
+          UnwindInsertBefore, RootOfChain, PointerToBase[LiveValue]);
 
       Info.RematerializedValues[NormalRematerializedValue] = LiveValue;
       Info.RematerializedValues[UnwindRematerializedValue] = LiveValue;
     }
   }
 
   // Remove rematerializaed values from the live set
   for (auto LiveValue: LiveValuesToBeDeleted) {
#endif
   }
 
   SmallVector<PartiallyConstructedSafepointRecord, 64> Records(ToUpdate.size());
 
   // A) Identify all gc pointers which are statically live at the given call
   // site.
   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);
+    findBasePointers(DT, DVCache, ToUpdate[i], info, PointerToBase);
+  }
+  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);
+      errs() << "\n";
+      ;
+    }
   }
 
   // The base phi insertion logic (for any safepoint) may have inserted new
   // instructions which are now live at some safepoint.  The simplest such
   // example is:
   // loop:
   //   phi a  <-- will be a new base_phi here
   //   safepoint 1 <-- that needs to be live here
   //   gep a + 1
   //   safepoint 2
   //   br loop
   // We insert some dummy calls after each safepoint to definitely hold live
   // the base pointers which were identified for that safepoint.  We'll then
   // ask liveness for _every_ base inserted to see what is now live.  Then we
   // remove the dummy calls.
   Holders.reserve(Holders.size() + Records.size());
   for (size_t i = 0; i < Records.size(); i++) {
     PartiallyConstructedSafepointRecord &Info = Records[i];
 
     SmallVector<Value *, 128> Bases;
-    for (auto Pair : Info.PointerToBase)
-      Bases.push_back(Pair.second);
+    for (auto *Derived : Info.LiveSet) {
+      assert(PointerToBase.count(Derived) && "Missed base for derived pointer");
+      Bases.push_back(PointerToBase[Derived]);
+    }
 
     insertUseHolderAfter(ToUpdate[i], Bases, Holders);
   }
 
   // By selecting base pointers, we've effectively inserted new uses. Thus, we
   // need to rerun liveness.  We may *also* have inserted new defs, but that's
   // not the key issue.
-  recomputeLiveInValues(F, DT, ToUpdate, Records);
+  recomputeLiveInValues(F, DT, ToUpdate, Records, PointerToBase);
 
   if (PrintBasePointers) {
-    for (auto &Info : Records) {
-      errs() << "Base Pairs: (w/Relocation)\n";
-      for (auto Pair : Info.PointerToBase) {
-        errs() << " derived ";
-        Pair.first->printAsOperand(errs(), false);
-        errs() << " base ";
-        Pair.second->printAsOperand(errs(), false);
-        errs() << "\n";
-      }
-    }
-  }
+    for (auto Pair : PointerToBase) {
 
   // It is possible that non-constant live variables have a constant base.  For
   // example, a GEP with a variable offset from a global.  In this case we can
   // remove it from the liveset.  We already don't add constants to the liveset
   // because we assume they won't move at runtime and the GC doesn't need to be
   // informed about them.  The same reasoning applies if the base is constant.
   // Note that the relocation placement code relies on this filtering for
   // correctness as it expects the base to be in the liveset, which isn't true
   // if the base is constant.
-  for (auto &Info : Records)
-    for (auto &BasePair : Info.PointerToBase)
-      if (isa<Constant>(BasePair.second))
-        Info.LiveSet.remove(BasePair.first);
+  for (auto &Info : Records) {
+    Info.LiveSet.remove_if([&](Value *LiveV) {
+      assert(PointerToBase.count(LiveV) && "Missed base for derived pointer");
+      return isa<Constant>(PointerToBase[LiveV]);
+    });
+  }
 
   for (CallInst *CI : Holders)
     CI->eraseFromParent();
 
   Holders.clear();
 
   // In order to reduce live set of statepoint we might choose to rematerialize
   // some values instead of relocating them. This is purely an optimization and
   // does not influence correctness.
   for (size_t i = 0; i < Records.size(); i++)
-    rematerializeLiveValues(ToUpdate[i], Records[i], TTI);
+    rematerializeLiveValues(ToUpdate[i], Records[i], PointerToBase, TTI);
 
   // We need this to safely RAUW and delete call or invoke return values that
   // may themselves be live over a statepoint.  For details, please see usage in
   // makeStatepointExplicitImpl.
   std::vector<DeferredReplacement> Replacements;
 
   // Now run through and replace the existing statepoints with new ones with
   // the live variables listed.  We do not yet update uses of the values being
   // relocated. We have references to live variables that need to
   // survive to the last iteration of this loop.  (By construction, the
   // previous statepoint can not be a live variable, thus we can and remove
   // the old statepoint calls as we go.)
   for (size_t i = 0; i < Records.size(); i++)
-    makeStatepointExplicit(DT, ToUpdate[i], Records[i], Replacements);
+    makeStatepointExplicit(DT, ToUpdate[i], Records[i], Replacements,
+                           PointerToBase);
 
   ToUpdate.clear(); // prevent accident use of invalid calls.
 
   for (auto &PR : Replacements)
     PR.doReplacement();
 
   Replacements.clear();
 
   for (auto &Info : Records) {
     // These live sets may contain state Value pointers, since we replaced calls
     // with operand bundles with calls wrapped in gc.statepoint, and some of
     // those calls may have been def'ing live gc pointers.  Clear these out to
     // avoid accidentally using them.
     //
     // TODO: We should create a separate data structure that does not contain
     // these live sets, and migrate to using that data structure from this point
     // onward.
     Info.LiveSet.clear();
-    Info.PointerToBase.clear();
   }
+  PointerToBase.clear();
 
   // Do all the fixups of the original live variables to their relocated selves
   SmallVector<Value *, 128> Live;
   for (size_t i = 0; i < Records.size(); i++) {
     PartiallyConstructedSafepointRecord &Info = Records[i];
 
     // We can't simply save the live set from the original insertion.  One of
     // the live values might be the result of a call which needs a safepoint.
static void findLiveSetAtInst(Instruction *Inst, GCPtrLivenessData &Data,
   computeLiveInValues(BB->rbegin(), ++Inst->getIterator().getReverse(),
                       LiveOut);
   LiveOut.remove(Inst);
   Out.insert(LiveOut.begin(), LiveOut.end());
 }
 
 static void recomputeLiveInValues(GCPtrLivenessData &RevisedLivenessData,
                                   CallBase *Call,
-                                  PartiallyConstructedSafepointRecord &Info) {
+                                  PartiallyConstructedSafepointRecord &Info,
+                                  PointerToBaseTy &PointerToBase) {
   StatepointLiveSetTy Updated;
   findLiveSetAtInst(Call, RevisedLivenessData, Updated);
 
   // We may have base pointers which are now live that weren't before.  We need
   // to update the PointerToBase structure to reflect this.
   for (auto V : Updated)
-    Info.PointerToBase.insert({V, V});
-
-#ifndef NDEBUG
-  for (auto V : Updated)
-    assert(Info.PointerToBase.count(V) &&
-           "Must be able to find base for live value!");
-#endif
-
-  // Remove any stale base mappings - this can happen since our liveness is
-  // more precise then the one inherent in the base pointer analysis.
-  DenseSet<Value *> ToErase;
-  for (auto KVPair : Info.PointerToBase)
-    if (!Updated.count(KVPair.first))
-      ToErase.insert(KVPair.first);
-
-  for (auto *V : ToErase)
-    Info.PointerToBase.erase(V);
-
-#ifndef NDEBUG
-  for (auto KVPair : Info.PointerToBase)
-    assert(Updated.count(KVPair.first) && "record for non-live value");
-#endif
+    PointerToBase.insert({ V, V });
 
   Info.LiveSet = Updated;
 }