[Statepoint] Consolidate relocation type tracking [NFC]

Change the way we track how a particular pointer was relocated at a statepoint in selection dag.  Previously, we used an optional<location> for the spill lowering, and a block local Register for the newly introduced vreg lowering.  Combine all three lowerings (norelocate, spill, and vreg) into a single helper class, and keep a single copy of the information.

This is submitted separately as it really does make the code more readible on it's own, but the indirect motivation is to move vreg tracking from StatepointLowering to FunctionLoweringInfo.  This is the last piece needed to support cross block relocations with vregs; that will follow in a separate (non-NFC) patch.

llvm/include/llvm/CodeGen/FunctionLoweringInfo.h

@@ -91,13 +91,33 @@ public:
   /// Track virtual registers created for exception pointers.
   DenseMap<const Value *, Register> CatchPadExceptionPointers;
 
-  /// Keep track of frame indices allocated for statepoints as they could be
-  /// used across basic block boundaries (e.g. for an invoke).  For each
-  /// gc.statepoint instruction, maps uniqued llvm IR values to the slots they
-  /// were spilled in.  If a value is mapped to None it means we visited the
-  /// value but didn't spill it (because it was a constant, for instance).
-  using StatepointSpillMapTy = DenseMap<const Value *, Optional<int>>;
-  DenseMap<const Instruction *, StatepointSpillMapTy> StatepointSpillMaps;
+  /// Helper object to track which of three possible relocation mechanisms are
+  /// used for a particular value being relocated over a statepoint.
+  struct StatepointRelocationRecord {
+    enum RelocType {
+      // Value did not need to be relocated and can be used directly.
+      NoRelocate,
+      // Value was spilled to stack and needs filled at the gc.relocate.
+      Spill,
+      // Value was lowered to tied def and gc.relocate should be replaced with
+      // copy from vreg.
+      VReg,
+    } type = NoRelocate;
+    // Payload contains either frame index of the stack slot in which the value
+    // was spilled, or virtual register which contains the re-definition.
+    union payload_t {
+      payload_t() : FI(-1) {}
+      int FI;
+      Register Reg;
+    } payload;
+  };
+
+  /// Keep track of each value which was relocated and the strategy used to
+  /// relocate that value.  This information is required when visiting
+  /// gc.relocates which may appear in following blocks.
+  using StatepointSpillMapTy =
+    DenseMap<const Value *, StatepointRelocationRecord>;
+  DenseMap<const Instruction *, StatepointSpillMapTy> StatepointRelocationMaps;
 
   /// StaticAllocaMap - Keep track of frame indices for fixed sized allocas in
   /// the entry block.  This allows the allocas to be efficiently referenced

llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp

@@ -360,7 +360,7 @@ void FunctionLoweringInfo::clear() {
   RegFixups.clear();
   RegsWithFixups.clear();
   StatepointStackSlots.clear();
-  StatepointSpillMaps.clear();
+  StatepointRelocationMaps.clear();
   PreferredExtendType.clear();
 }
 

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

@@ -71,6 +71,8 @@ cl::opt<unsigned> MaxRegistersForGCPointers(
     "max-registers-for-gc-values", cl::Hidden, cl::init(0),
     cl::desc("Max number of VRegs allowed to pass GC pointer meta args in"));
 
+typedef FunctionLoweringInfo::StatepointRelocationRecord RecordType;
+
 static void pushStackMapConstant(SmallVectorImpl<SDValue>& Ops,
                                  SelectionDAGBuilder &Builder, uint64_t Value) {
   SDLoc L = Builder.getCurSDLoc();
@@ -90,13 +92,11 @@ void StatepointLoweringState::startNewStatepoint(SelectionDAGBuilder &Builder) {
   // FunctionLoweringInfo.  Also need to ensure used bits get cleared.
   AllocatedStackSlots.clear();
   AllocatedStackSlots.resize(Builder.FuncInfo.StatepointStackSlots.size());
-  VirtRegs.clear();
 }
 
 void StatepointLoweringState::clear() {
   Locations.clear();
   AllocatedStackSlots.clear();
-  VirtRegs.clear();
   assert(PendingGCRelocateCalls.empty() &&
          "cleared before statepoint sequence completed");
 }
@@ -162,14 +162,18 @@ static Optional<int> findPreviousSpillSlot(const Value *Val,
 
   // Spill location is known for gc relocates
   if (const auto *Relocate = dyn_cast<GCRelocateInst>(Val)) {
-    const auto &SpillMap =
-        Builder.FuncInfo.StatepointSpillMaps[Relocate->getStatepoint()];
+    const auto &RelocationMap =
+        Builder.FuncInfo.StatepointRelocationMaps[Relocate->getStatepoint()];
 
-    auto It = SpillMap.find(Relocate->getDerivedPtr());
-    if (It == SpillMap.end())
+    auto It = RelocationMap.find(Relocate->getDerivedPtr());
+    if (It == RelocationMap.end())
       return None;
 
-    return It->second;
+    auto &Record = It->second;
+    if (Record.type != RecordType::Spill)
+      return None;
+
+    return Record.payload.FI;
   }
 
   // Look through bitcast instructions.
@@ -666,39 +670,6 @@ lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
       MemRefs.push_back(MMO);
     }
   }
-
-  // Record computed locations for all lowered values.
-  // This can not be embedded in lowering loops as we need to record *all*
-  // values, while previous loops account only values with unique SDValues.
-  const Instruction *StatepointInstr = SI.StatepointInstr;
-  auto &SpillMap = Builder.FuncInfo.StatepointSpillMaps[StatepointInstr];
-
-  for (const GCRelocateInst *Relocate : SI.GCRelocates) {
-    const Value *V = Relocate->getDerivedPtr();
-    SDValue SDV = Builder.getValue(V);
-    SDValue Loc = Builder.StatepointLowering.getLocation(SDV);
-
-    if (Loc.getNode()) {
-      // If this is a value we spilled, remember where for when we visit the
-      // gc.relocate corresponding to this gc.statepoint
-      SpillMap[V] = cast<FrameIndexSDNode>(Loc)->getIndex();
-    } else {
-      // If we didn't spill the value - allocas, constants, and values lowered
-      // as tied vregs - mark them as visited, but not spilled.  Marking them
-      // visited (as opposed to simply missing in the map), allows tighter
-      // assertion checking.  
-      SpillMap[V] = None;
-
-      // Conservatively export all values used by gc.relocates outside this
-      // block.  This is currently only needed for expressions which don't need
-      // relocation (such as constants and allocas).
-      if (!LowerAsVReg.count(SDV) &&
-          Relocate->getParent() != StatepointInstr->getParent()) {
-        Builder.ExportFromCurrentBlock(V);
-        assert(!LowerAsVReg.count(SDV));
-      }
-    }
-  }
 }
 
 SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
@@ -856,12 +827,10 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
     DAG.getMachineNode(TargetOpcode::STATEPOINT, getCurSDLoc(), NodeTys, Ops);
   DAG.setNodeMemRefs(StatepointMCNode, MemRefs);
 
-  SDNode *SinkNode = StatepointMCNode;
 
-  // Remember the mapping between values relocated in registers and the virtual
-  // register holding the relocation.  Note that for simplicity, we *always*
-  // create a vreg even within a single block.
-  auto &VirtRegs = StatepointLowering.VirtRegs;
+  // For values lowered to tied-defs, create the virtual registers.  Note that
+  // for simplicity, we *always* create a vreg even within a single block. 
+  DenseMap<const Value *, Register> VirtRegs;
   for (const auto *Relocate : SI.GCRelocates) {
     Value *Derived = Relocate->getDerivedPtr();
     SDValue SD = getValue(Derived);
@@ -885,6 +854,39 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
     VirtRegs[Derived] = Reg; 
   }
 
+  // Record for later use how each relocation was lowered.  This is needed to
+  // allow later gc.relocates to mirror the lowering chosen.
+  const Instruction *StatepointInstr = SI.StatepointInstr;
+  auto &RelocationMap = FuncInfo.StatepointRelocationMaps[StatepointInstr];
+  for (const GCRelocateInst *Relocate : SI.GCRelocates) {
+    const Value *V = Relocate->getDerivedPtr();
+    SDValue SDV = getValue(V);
+    SDValue Loc = StatepointLowering.getLocation(SDV);
+
+    RecordType Record;
+    if (Loc.getNode()) {
+      Record.type = RecordType::Spill;
+      Record.payload.FI = cast<FrameIndexSDNode>(Loc)->getIndex();
+    } else if (LowerAsVReg.count(SDV)) {
+      Record.type = RecordType::VReg;
+      assert(VirtRegs.count(V));
+      Record.payload.Reg = VirtRegs[V];
+    } else {
+      Record.type = RecordType::NoRelocate;
+      // If we didn't relocate a value, we'll essentialy end up inserting an
+      // additional use of the original value when lowering the gc.relocate.
+      // We need to make sure the value is available at the new use, which
+      // might be in another block.
+      if (Relocate->getParent() != StatepointInstr->getParent())
+        ExportFromCurrentBlock(V);
+    }
+    RelocationMap[V] = Record;
+  }
+
+  
+
+  SDNode *SinkNode = StatepointMCNode;
+
   // Build the GC_TRANSITION_END node if necessary.
   //
   // See the comment above regarding GC_TRANSITION_START for the layout of
@@ -1118,12 +1120,15 @@ void SelectionDAGBuilder::visitGCRelocate(const GCRelocateInst &Relocate) {
 #endif
 
   const Value *DerivedPtr = Relocate.getDerivedPtr();
+  auto &RelocationMap =
+    FuncInfo.StatepointRelocationMaps[Relocate.getStatepoint()];
+  auto SlotIt = RelocationMap.find(DerivedPtr);
+  assert(SlotIt != RelocationMap.end() && "Relocating not lowered gc value");
+  const RecordType &Record = SlotIt->second;
 
   // If relocation was done via virtual register..
-  auto &VirtRegs = StatepointLowering.VirtRegs;
-  auto It = VirtRegs.find(DerivedPtr);
-  if (It != VirtRegs.end()) {
-    Register InReg = It->second;
+  if (Record.type == RecordType::VReg) {
+    Register InReg = Record.payload.Reg;
     RegsForValue RFV(*DAG.getContext(), DAG.getTargetLoweringInfo(),
                      DAG.getDataLayout(), InReg, Relocate.getType(),
                      None); // This is not an ABI copy.
@@ -1143,19 +1148,17 @@ void SelectionDAGBuilder::visitGCRelocate(const GCRelocateInst &Relocate) {
     return;
   }
 
-  auto &SpillMap = FuncInfo.StatepointSpillMaps[Relocate.getStatepoint()];
-  auto SlotIt = SpillMap.find(DerivedPtr);
-  assert(SlotIt != SpillMap.end() && "Relocating not lowered gc value");
-  Optional<int> DerivedPtrLocation = SlotIt->second;
 
   // We didn't need to spill these special cases (constants and allocas).
   // See the handling in spillIncomingValueForStatepoint for detail.
-  if (!DerivedPtrLocation) {
+  if (Record.type == RecordType::NoRelocate) {
     setValue(&Relocate, SD);
     return;
   }
 
-  unsigned Index = *DerivedPtrLocation;
+  assert(Record.type == RecordType::Spill);
+
+  unsigned Index = Record.payload.FI;;
   SDValue SpillSlot = DAG.getTargetFrameIndex(Index, getFrameIndexTy());
 
   // All the reloads are independent and are reading memory only modified by

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.h

@@ -103,10 +103,6 @@ public:
     return AllocatedStackSlots.test(Offset);
   }
 
-  /// Maps pre-relocated value to virtual register holding it's relocation if
-  /// vreg lowering was used.
-  DenseMap<const Value *, Register> VirtRegs;
-
 private:
   /// Maps pre-relocation value (gc pointer directly incoming into statepoint)
   /// into it's location (currently only stack slots)