MIR Statepoint refactoring: pass GC pointers in VRegs. Part 3/5.

[dantrushin]

Summary

Change MIR representation of STATEPOINT instruction from explicit
spill/reload of GC pointer arguments around statepoint to simply
assigning them to virtual registers. Then they can be spilled as
necessary. Relocated pointers are represented as STATEPOINT's Def
operands (tied to corresponding derived pointers from GC args list):

rel1,rel2,... = STATEPOINT ..., derived1<tied-def0>, derived2<tied-def1>, ...

This patch implements actual lowering with new statepoint representation.
Due to storage limitation of SDNode, first N base and derived pointers
pairs from GC Args list are assigned to virtual registers (of course,
only those which need relocations for derived pointers).
Upper value of N is bounded by the maximal amount of tied registers
MachineInstr can manage (currently, 15). The rest, if any, is passed
on stack slots, as in the old implementation.

Depends On: D80192

[dantrushin]

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

230

@reames : Philip, is there reason why ISD::UNDEF is not handled just like constants in upstream LLVM?
LLVM spills these UNDEFs and they are eliminated by later passes (if unused).

[dantrushin]

Abandoned in favor of D81648

llvm/include/llvm/CodeGen/FunctionLoweringInfo.h

public:
   /// 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;
 
+  /// For each statepoint keep mapping from original derived pointer to
+  /// the index of StatepointSDNode result defining its new value.
+  using DerivedPtrMapTy = DenseMap<const Value *, unsigned>;
+  DenseMap<const Instruction *, DerivedPtrMapTy> DerivedPtrMap;
+
+  /// For each statepoint keep virtual registers its result values has
+  /// been exported to.
+  DenseMap<const Instruction *, SmallVector<unsigned, 8>> StatepointRegs;
+
   /// StaticAllocaMap - Keep track of frame indices for fixed sized allocas in
   /// the entry block.  This allows the allocas to be efficiently referenced
   /// anywhere in the function.
   DenseMap<const AllocaInst*, int> StaticAllocaMap;
 
   /// ByValArgFrameIndexMap - Keep track of frame indices for byval arguments.
   DenseMap<const Argument*, int> ByValArgFrameIndexMap;
 

llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp

if (RegisterSDNode *RN = dyn_cast<RegisterSDNode>(Node->getOperand(I - 1)))
         continue;
     NumImpUses = N - I;
     break;
   }
 
   return N;
 }
 
+/// Collect used physical registers up the glue chain.
+static void collectPhysRegsFromGlueChain(const SDNode *Node,
+                                         const TargetInstrInfo *TII,
+                                         SmallVectorImpl<Register> &UsedRegs) {
+  if (Node->getValueType(Node->getNumValues() - 1) != MVT::Glue)
+    return;
+  for (SDNode *F = Node->getGluedUser(); F; F = F->getGluedUser()) {
+    if (F->getOpcode() == ISD::CopyFromReg) {
+      UsedRegs.push_back(cast<RegisterSDNode>(F->getOperand(1))->getReg());
+      continue;
+    } else if (F->getOpcode() == ISD::CopyToReg) {
+      // Skip CopyToReg nodes that are internal to the glue chain.
+      continue;
+    }
+    // Collect declared implicit uses.
+    if (F->isMachineOpcode()) {
+      const MCInstrDesc &MCID = TII->get(F->getMachineOpcode());
+      UsedRegs.append(MCID.getImplicitUses(),
+                      MCID.getImplicitUses() + MCID.getNumImplicitUses());
+    }
+    // In addition to declared implicit uses, we must also check for
+    // direct RegisterSDNode operands.
+    for (unsigned i = 0, e = F->getNumOperands(); i != e; ++i)
+      if (RegisterSDNode *R = dyn_cast<RegisterSDNode>(F->getOperand(i))) {
+        Register Reg = R->getReg();
+        if (Reg.isPhysical())
+          UsedRegs.push_back(Reg);
+      }
+  }
+}
+
 /// EmitCopyFromReg - Generate machine code for an CopyFromReg node or an
 /// implicit physical register output.
 void InstrEmitter::
 EmitCopyFromReg(SDNode *Node, unsigned ResNo, bool IsClone, bool IsCloned,
                 Register SrcReg, DenseMap<SDValue, Register> &VRBaseMap) {
   Register VRBase;
   if (SrcReg.isVirtual()) {
     // Just use the input register directly!
for (unsigned i = NumDefs; i < NumResults; ++i) {
       if (!Node->hasAnyUseOfValue(i))
         continue;
       // This implicitly defined physreg has a use.
       UsedRegs.push_back(Reg);
       EmitCopyFromReg(Node, i, IsClone, IsCloned, Reg, VRBaseMap);
     }
   }
 
-  // Scan the glue chain for any used physregs.
-  if (Node->getValueType(Node->getNumValues()-1) == MVT::Glue) {
-    for (SDNode *F = Node->getGluedUser(); F; F = F->getGluedUser()) {
-      if (F->getOpcode() == ISD::CopyFromReg) {
-        UsedRegs.push_back(cast<RegisterSDNode>(F->getOperand(1))->getReg());
-        continue;
-      } else if (F->getOpcode() == ISD::CopyToReg) {
-        // Skip CopyToReg nodes that are internal to the glue chain.
-        continue;
-      }
-      // Collect declared implicit uses.
-      const MCInstrDesc &MCID = TII->get(F->getMachineOpcode());
-      UsedRegs.append(MCID.getImplicitUses(),
-                      MCID.getImplicitUses() + MCID.getNumImplicitUses());
-      // In addition to declared implicit uses, we must also check for
-      // direct RegisterSDNode operands.
-      for (unsigned i = 0, e = F->getNumOperands(); i != e; ++i)
-        if (RegisterSDNode *R = dyn_cast<RegisterSDNode>(F->getOperand(i))) {
-          Register Reg = R->getReg();
-          if (Reg.isPhysical())
-            UsedRegs.push_back(Reg);
-        }
-    }
-  }
+  collectPhysRegsFromGlueChain(Node, TII, UsedRegs);
 
   // Finally mark unused registers as dead.
   if (!UsedRegs.empty() || II.getImplicitDefs() || II.hasOptionalDef())
     MIB->setPhysRegsDeadExcept(UsedRegs, *TRI);
 
   // Run post-isel target hook to adjust this instruction if needed.
   if (II.hasPostISelHook())
     TLI->AdjustInstrPostInstrSelection(*MIB, Node);
case ISD::LIFETIME_END: {
     TargetOpcode::LIFETIME_START : TargetOpcode::LIFETIME_END;
 
     FrameIndexSDNode *FI = dyn_cast<FrameIndexSDNode>(Node->getOperand(1));
     BuildMI(*MBB, InsertPos, Node->getDebugLoc(), TII->get(TarOp))
     .addFrameIndex(FI->getIndex());
     break;
   }
 
+  case ISD::GC_STATEPOINT: {
+    StatepointSDNode *SN = cast<StatepointSDNode>(Node);
+    unsigned GCArgStart = SN->getGCArgStart();
+    unsigned NumValues = CountResults(Node);
+    MachineInstrBuilder MIB =
+        BuildMI(*MF, Node->getDebugLoc(), TII->get(TargetOpcode::STATEPOINT));
+
+    for (unsigned i = 0; i < NumValues; ++i) {
+      SDValue Def(SN, i);
+      SDValue Use = SN->getOperand(GCArgStart + 1 + i * 2);
+      unsigned UseReg = getVR(Use, VRBaseMap);
+      unsigned DefReg = MRI->cloneVirtualRegister(UseReg);
+      MIB = MIB.addDef(DefReg);
+      VRBaseMap[Def] = DefReg;
+    }
+
+    unsigned NumOperands = Node->getNumOperands();
+    for (unsigned i = 0; i < NumOperands; ++i) {
+      const SDValue &O = Node->getOperand(i);
+      if (O.getValueType() == MVT::Other || O.getValueType() == MVT::Glue)
+        continue;
+      AddOperand(MIB, O, 0, nullptr, VRBaseMap, false, false, false);
+    }
+
+    unsigned Use = NumValues + GCArgStart + 1;
+    for (unsigned Def = 0; Def < NumValues; ++Def, Use += 2) {
+      MIB->tieOperands(Def, Use);
+    }
+
+    MBB->insert(InsertPos, MIB);
+
+    SmallVector<Register, 8> UsedRegs;
+    collectPhysRegsFromGlueChain(Node, TII, UsedRegs);
+    if (!UsedRegs.empty())
+      MIB->setPhysRegsDeadExcept(UsedRegs, *TRI);
+
+    break;
+  }
   case ISD::INLINEASM:
   case ISD::INLINEASM_BR: {
     unsigned NumOps = Node->getNumOperands();
     if (Node->getOperand(NumOps-1).getValueType() == MVT::Glue)
       --NumOps;  // Ignore the glue operand.
 
     // Create the inline asm machine instruction.
     unsigned TgtOpc = Node->getOpcode() == ISD::INLINEASM_BR

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

void SelectionDAGISel::SelectCodeCommon(SDNode *NodeToMatch,
   case ISD::TargetGlobalAddress:
   case ISD::TokenFactor:
   case ISD::CopyFromReg:
   case ISD::CopyToReg:
   case ISD::EH_LABEL:
   case ISD::ANNOTATION_LABEL:
   case ISD::LIFETIME_START:
   case ISD::LIFETIME_END:
+  case ISD::GC_STATEPOINT:
     NodeToMatch->setNodeId(-1); // Mark selected.
     return;
   case ISD::AssertSext:
   case ISD::AssertZext:
     ReplaceUses(SDValue(NodeToMatch, 0), NodeToMatch->getOperand(0));
     CurDAG->RemoveDeadNode(NodeToMatch);
     return;
   case ISD::INLINEASM:

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

 STATISTIC(NumOfStatepoints, "Number of statepoint nodes encountered");
 STATISTIC(StatepointMaxSlotsRequired,
           "Maximum number of stack slots required for a singe statepoint");
 
 cl::opt<bool> UseRegistersForDeoptValues(
     "use-registers-for-deopt-values", cl::Hidden, cl::init(false),
     cl::desc("Allow using registers for non pointer deopt args"));
 
+cl::opt<bool> UseRegistersForGCPointers(
+    "use-registers-for-gcptrs", cl::Hidden, cl::init(false),
+    cl::desc("Allow using registers for GC pointer meta args"));
+
 static void pushStackMapConstant(SmallVectorImpl<SDValue>& Ops,
                                  SelectionDAGBuilder &Builder, uint64_t Value) {
   SDLoc L = Builder.getCurSDLoc();
   Ops.push_back(Builder.DAG.getTargetConstant(StackMaps::ConstantOp, L,
                                               MVT::i64));
   Ops.push_back(Builder.DAG.getTargetConstant(Value, L, MVT::i64));
 }
 
static Optional<int> findPreviousSpillSlot(const Value *Val,
   // put handling of simple modifications in this function like it's done
   // for bitcasts we might end up reserving i's slot for 'i+1' because order in
   // which we visit values is unspecified.
 
   // Don't know any information about this instruction
   return None;
 }
 
+// Return true if V is a GC pointer which need not to be relocated.
+static bool isNonRelocatablePtr(SDValue V) {
+  return (isa<ConstantSDNode>(V) || isa<FrameIndexSDNode>(V));
+}
+
 /// Try to find existing copies of the incoming values in stack slots used for
 /// statepoint spilling.  If we can find a spill slot for the incoming value,
 /// mark that slot as allocated, and reuse the same slot for this safepoint.
 /// This helps to avoid series of loads and stores that only serve to reshuffle
 /// values on the stack between calls.
 static void reservePreviousStackSlotForValue(const Value *IncomingValue,
                                              SelectionDAGBuilder &Builder) {
   SDValue Incoming = Builder.getValue(IncomingValue);
 
-  if (isa<ConstantSDNode>(Incoming) || isa<FrameIndexSDNode>(Incoming)) {
-    // We won't need to spill this, so no need to check for previously
-    // allocated stack slots
+  if (isNonRelocatablePtr(Incoming))
     return;
-  }
 
   SDValue OldLocation = Builder.StatepointLowering.getLocation(Incoming);
   if (OldLocation.getNode())
     // Duplicates in input
     return;
 
   const int LookUpDepth = 6;
   Optional<int> Index =
static void reservePreviousStackSlotForValue(const Value *IncomingValue,
 
   // Cache this slot so we find it when going through the normal
   // assignment loop.
   SDValue Loc =
       Builder.DAG.getTargetFrameIndex(*Index, Builder.getFrameIndexTy());
   Builder.StatepointLowering.setLocation(Incoming, Loc);
 }
 
+/// Sort Ptrs vector so that pointers which need no relocation (constants and
+/// allocas) are placed at the end and those which need relocation are
+/// contiguously occupy beginning of the vector.
+/// Synchroniously update Bases and Relocs vectors.
+static unsigned sortGCPtrs(SmallVectorImpl<const Value *> &Bases,
+                           SmallVectorImpl<const Value *> &Ptrs,
+                           SmallVectorImpl<const GCRelocateInst *> &Relocs,
+                           SelectionDAGBuilder &Builder) {
+  unsigned curPos = 0;
+  for (unsigned i = 0, e = Ptrs.size(); i < e; ++i) {
+    SDValue SDV = Builder.getValue(Ptrs[i]);
+    if (isNonRelocatablePtr(SDV) || SDV.getOpcode() == ISD::UNDEF ||
+        SDV.getValueType().getSizeInBits() > 64) {
+      continue;
+    }
+    if (curPos < i) {
+      std::swap(Bases[curPos], Bases[i]);
+      std::swap(Ptrs[curPos], Ptrs[i]);
+      std::swap(Relocs[curPos], Relocs[i]);
+    }
+    ++curPos;
+  }
+  return curPos;
+}
+
 /// Extract call from statepoint, lower it and return pointer to the
 /// call node. Also update NodeMap so that getValue(statepoint) will
 /// reference lowered call result
 static std::pair<SDValue, SDNode *> lowerCallFromStatepointLoweringInfo(
     SelectionDAGBuilder::StatepointLoweringInfo &SI,
     SelectionDAGBuilder &Builder, SmallVectorImpl<SDValue> &PendingExports) {
   SDValue ReturnValue, CallEndVal;
   std::tie(ReturnValue, CallEndVal) =
if (!Loc.getNode()) {
     auto PtrInfo = MachinePointerInfo::getFixedStack(MF, Index);
     auto *StoreMMO = MF.getMachineMemOperand(
         PtrInfo, MachineMemOperand::MOStore, MFI.getObjectSize(Index),
         MFI.getObjectAlign(Index));
     Chain = Builder.DAG.getStore(Chain, Builder.getCurSDLoc(), Incoming, Loc,
                                  StoreMMO);
 
     MMO = getMachineMemOperand(MF, *cast<FrameIndexSDNode>(Loc));
-    
     Builder.StatepointLowering.setLocation(Incoming, Loc);
   }
 
   assert(Loc.getNode());
   return std::make_tuple(Loc, Chain, MMO);
 }
 
 /// Lower a single value incoming to a statepoint node.  This value can be

 /// lowering is described in lowerIncomingStatepointValue.  This function is
 /// responsible for lowering everything in the right position and playing some
 /// tricks to avoid redundant stack manipulation where possible.  On
 /// completion, 'Ops' will contain ready to use operands for machine code
 /// statepoint. The chain nodes will have already been created and the DAG root
 /// will be set to the last value spilled (if any were).
 static void
 lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
-                        SmallVectorImpl<MachineMemOperand*> &MemRefs,                                    SelectionDAGBuilder::StatepointLoweringInfo &SI,
+                        SmallVectorImpl<MachineMemOperand *> &MemRefs,
+                        unsigned NumVRegGCArgs, unsigned &GCArgStart,
+                        SelectionDAGBuilder::StatepointLoweringInfo &SI,
                         SelectionDAGBuilder &Builder) {
   // Lower the deopt and gc arguments for this statepoint.  Layout will be:
   // deopt argument length, deopt arguments.., gc arguments...
 #ifndef NDEBUG
   if (auto *GFI = Builder.GFI) {
     // Check that each of the gc pointer and bases we've gotten out of the
     // safepoint is something the strategy thinks might be a pointer (or vector
     // of pointers) into the GC heap.  This is basically just here to help catch
if (!Ty->isPtrOrPtrVectorTy())
       return false;
     if (auto *GFI = Builder.GFI)
       if (auto IsManaged = GFI->getStrategy().isGCManagedPointer(Ty))
         return *IsManaged;
     return true; // conservative
   };
 
   auto requireSpillSlot = [&](const Value *V) {
-    return !(LiveInDeopt || UseRegistersForDeoptValues) || isGCValue(V);
+    if (isGCValue(V)) {
+      auto *Ty = V->getType();
+      return Ty->isVectorTy() || !UseRegistersForGCPointers;
+    }
+    return !(LiveInDeopt || UseRegistersForDeoptValues);
   };
 
   // Before we actually start lowering (and allocating spill slots for values),
   // reserve any stack slots which we judge to be profitable to reuse for a
   // particular value.  This is purely an optimization over the code below and
   // doesn't change semantics at all.  It is important for performance that we
   // reserve slots for both deopt and gc values before lowering either.
   for (const Value *V : SI.DeoptState) {
     if (requireSpillSlot(V))
       reservePreviousStackSlotForValue(V, Builder);
   }
-  for (unsigned i = 0; i < SI.Bases.size(); ++i) {
+  for (unsigned i = NumVRegGCArgs; i < SI.Bases.size(); ++i) {
     reservePreviousStackSlotForValue(SI.Bases[i], Builder);
     reservePreviousStackSlotForValue(SI.Ptrs[i], Builder);
   }
 
   // First, prefix the list with the number of unique values to be
   // lowered.  Note that this is the number of *Values* not the
   // number of SDValues required to lower them.
   const int NumVMSArgs = SI.DeoptState.size();
if (const Argument *Arg = dyn_cast<Argument>(V)) {
         Incoming = Builder.DAG.getFrameIndex(FI, Builder.getFrameIndexTy());
     }
     if (!Incoming.getNode())
       Incoming = Builder.getValue(V);
     lowerIncomingStatepointValue(Incoming, requireSpillSlot(V), Ops, MemRefs,
                                  Builder);
   }
 
+  GCArgStart = Ops.size();
   // Finally, go ahead and lower all the gc arguments.  There's no prefixed
   // length for this one.  After lowering, we'll have the base and pointer
   // arrays interwoven with each (lowered) base pointer immediately followed by
   // it's (lowered) derived pointer.  i.e
   // (base[0], ptr[0], base[1], ptr[1], ...)
+  // Lower first `NumVRegGCArgs` base AND derived pointers through VRegs.
+  // In future we might use more sophisticated strategy for choosing which
+  // pointers to pass via virtual registers, but for now this simple approach
+  // looks good enough. Take into account these facts:
+  //  - NumVRegGCArgs is limited by the  max number of tied registers in MI;
+  //  - We relocate (and so need tied defs for) only derived pointers;
+  //  - Quite often base and derived pointer are the same.
+  auto &SL = Builder.StatepointLowering;
   for (unsigned i = 0; i < SI.Bases.size(); ++i) {
+    bool RequireSpillSlot = (i >= NumVRegGCArgs);
     const Value *Base = SI.Bases[i];
-    lowerIncomingStatepointValue(Builder.getValue(Base),
-                                 /*RequireSpillSlot*/ true, Ops, MemRefs,
-                                 Builder);
+    lowerIncomingStatepointValue(Builder.getValue(Base), RequireSpillSlot, Ops,
+                                 MemRefs, Builder);
 
     const Value *Ptr = SI.Ptrs[i];
-    lowerIncomingStatepointValue(Builder.getValue(Ptr),
-                                 /*RequireSpillSlot*/ true, Ops, MemRefs,
-                                 Builder);
+    SDValue SDV = Builder.getValue(Ptr);
+    lowerIncomingStatepointValue(SDV, RequireSpillSlot, Ops, MemRefs, Builder);
+    if (!RequireSpillSlot && !SL.getLocation(SDV))
+      SL.setLocation(SDV, Builder.DAG.getConstant(i, SDLoc(), MVT::i64));
   }
 
   // If there are any explicit spill slots passed to the statepoint, record
   // them, but otherwise do not do anything special.  These are user provided
   // allocas and give control over placement to the consumer.  In this case,
   // it is the contents of the slot which may get updated, not the pointer to
   // the alloca
   for (Value *V : SI.GCArgs) {
for (Value *V : SI.GCArgs) {
     }
   }
 
   // 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];
+  auto &DPtrMap = Builder.FuncInfo.DerivedPtrMap[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()) {
-      SpillMap[V] = cast<FrameIndexSDNode>(Loc)->getIndex();
+      if (auto FI = dyn_cast<FrameIndexSDNode>(Loc))
+        SpillMap[V] = FI->getIndex();
+      else {
+        DPtrMap[V] = cast<ConstantSDNode>(Loc)->getZExtValue();
+        SpillMap[V] = None;
+      }
     } else {
       // Record value as visited, but not spilled. This is case for allocas
       // and constants. For this values we can avoid emitting spill load while
       // visiting corresponding gc_relocate.
       // Actually we do not need to record them in this map at all.
       // We do this only to check that we are not relocating any unvisited
       // value.
       SpillMap[V] = None;
assert(SI.Bases.size() == SI.Ptrs.size() &&
          SI.Ptrs.size() <= SI.GCRelocates.size());
 
 #ifndef NDEBUG
   for (auto *Reloc : SI.GCRelocates)
     if (Reloc->getParent() == SI.StatepointInstr->getParent())
       StatepointLowering.scheduleRelocCall(*Reloc);
 #endif
 
+  unsigned NumVRegs = 0;
+
+  if (UseRegistersForGCPointers) {
+    const unsigned MaxTiedRegs = 15U;
+
+    // Sort vectors so that elements which need relocation are laid out
+    // contiguously at the beginning of vectors.
+    // This is dictated by the StatepointSDNode implementation: due to size limit,
+    // one cannot put vector into class derived from SDNode, we can only keep a
+    // number of relocations. So N results of StatepointSDNode map 1-1 to the
+    // first N derived pointers.
+    NumVRegs = sortGCPtrs(SI.Bases, SI.Ptrs, SI.GCRelocates, *this);
+    NumVRegs = std::min(NumVRegs, MaxTiedRegs);
+  }
+
   // Lower statepoint vmstate and gcstate arguments
   SmallVector<SDValue, 10> LoweredMetaArgs;
   SmallVector<MachineMemOperand*, 16> MemRefs;
-  lowerStatepointMetaArgs(LoweredMetaArgs, MemRefs, SI, *this);
+  unsigned GCArgStart;
+  lowerStatepointMetaArgs(LoweredMetaArgs, MemRefs, NumVRegs, GCArgStart, SI,
+                          *this);
 
   // Now that we've emitted the spills, we need to update the root so that the
   // call sequence is ordered correctly.
   SI.CLI.setChain(getRoot());
 
   // Get call node, we will replace it later with statepoint
   SDValue ReturnVal;
   SDNode *CallNode;
#endif
 
   // Add a constant argument for the flags
   uint64_t Flags = SI.StatepointFlags;
   assert(((Flags & ~(uint64_t)StatepointFlags::MaskAll) == 0) &&
          "Unknown flag used");
   pushStackMapConstant(Ops, *this, Flags);
 
   // Insert all vmstate and gcstate arguments
+  GCArgStart += Ops.size();
   Ops.insert(Ops.end(), LoweredMetaArgs.begin(), LoweredMetaArgs.end());
 
   // Add register mask from call node
   Ops.push_back(*RegMaskIt);
 
   // Add chain
   Ops.push_back(Chain);
 
   // Same for the glue, but we add it only if original call had it
   if (Glue.getNode())
     Ops.push_back(Glue);
 
-  // Compute return values.  Provide a glue output since we consume one as
-  // input.  This allows someone else to chain off us as needed.
-  SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
-
-  MachineSDNode *StatepointMCNode =
-    DAG.getMachineNode(TargetOpcode::STATEPOINT, getCurSDLoc(), NodeTys, Ops);
-  DAG.setNodeMemRefs(StatepointMCNode, MemRefs);
+  SDValue STV =
+      DAG.getStatepoint(getCurSDLoc(), ReturnVal, GCArgStart, NumVRegs, Ops);
 
-  SDNode *SinkNode = StatepointMCNode;
+  SDNode *SinkNode = STV.getNode();
 
   // Build the GC_TRANSITION_END node if necessary.
   //
   // See the comment above regarding GC_TRANSITION_START for the layout of
   // the operands to the GC_TRANSITION_END node.
   if (IsGCTransition) {
     SmallVector<SDValue, 8> TEOps;
 
     // Add chain
-    TEOps.push_back(SDValue(StatepointMCNode, 0));
+    TEOps.push_back(SDValue(STV.getNode(), STV->getNumValues() - 2));
 
     // Add GC transition arguments
     for (const Value *V : SI.GCTransitionArgs) {
       TEOps.push_back(getValue(V));
       if (V->getType()->isPointerTy())
         TEOps.push_back(DAG.getSrcValue(V));
     }
 
     // Add glue
-    TEOps.push_back(SDValue(StatepointMCNode, 1));
+    TEOps.push_back(SDValue(STV.getNode(), STV->getNumValues() - 1));
 
     SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
 
     SDValue GCTransitionStart =
         DAG.getNode(ISD::GC_TRANSITION_END, getCurSDLoc(), NodeTys, TEOps);
 
     SinkNode = GCTransitionStart.getNode();
   }
 
   // Replace original call
-  DAG.ReplaceAllUsesWith(CallNode, SinkNode); // This may update Root
+  // Call: ch,glue = CALL ...
+  // Statepoint: [gc relocates],ch,glue = GC_STATEPOINT ...
+  unsigned NumSinkValues = SinkNode->getNumValues();
+  SDValue StatepointValues[2] = {SDValue(SinkNode, NumSinkValues - 2),
+                                 SDValue(SinkNode, NumSinkValues - 1)};
+  DAG.ReplaceAllUsesWith(CallNode, StatepointValues);
   // Remove original call node
   DAG.DeleteNode(CallNode);
 
   // DON'T set the root - under the assumption that it's already set past the
   // inserted node we created.
 
   // TODO: A better future implementation would be to emit a single variable
   // argument, variable return value STATEPOINT node here and then hookup the
   // return value of each gc.relocate to the respective output of the
   // previously emitted STATEPOINT value.  Unfortunately, this doesn't appear
   // to actually be possible today.
 
-  return ReturnVal;
+  return STV;
 }
 
 void
 SelectionDAGBuilder::LowerStatepoint(ImmutableStatepoint ISP,
                                      const BasicBlock *EHPadBB /*= nullptr*/) {
   assert(ISP.getCall()->getCallingConv() != CallingConv::AnyReg &&
          "anyregcc is not supported on statepoints!");
 
#endif
   SI.GCTransitionArgs = ArrayRef<const Use>(ISP.gc_transition_args_begin(),
                                             ISP.gc_transition_args_end());
   SI.ID = ISP.getID();
   SI.DeoptState = ArrayRef<const Use>(ISP.deopt_begin(), ISP.deopt_end());
   SI.StatepointFlags = ISP.getFlags();
   SI.NumPatchBytes = ISP.getNumPatchBytes();
   SI.EHPadBB = EHPadBB;
 
-  SDValue ReturnValue = LowerAsSTATEPOINT(SI);
+  SDValue STV = LowerAsSTATEPOINT(SI);
 
   // Export the result value if needed
+  const BasicBlock *BB = ISP.getCall()->getParent();
+  std::vector<const GCRelocateInst *> RV = ISP.getRelocates();
+  bool NeedExport = llvm::any_of(
+      RV, [&BB](const GCRelocateInst *I) { return I->getParent() != BB; });
+
+  // If any of relocates or result value will be used in different basic
+  // block, we need to export them manually. Default exporting mechanism
+  // will not work here because it is based on IR Value types, and
+  // IR statepoint has different type than the actual call or relocates.
+  // It means that by default llvm will create export register of the wrong
+  // type (always i32 - TokenTy - in our case). So instead we need to create
+  // export registers manually.
+  // TODO: To eliminate this problem we can remove gc.result/gc.relocate
+  //       intrinsics completely and make statepoint call to return a tuple.
+  setValue(ISP.getInstruction(), STV);
+  if (NeedExport) {
+    LLVMContext *Context = DAG.getContext();
+    const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+    const DataLayout &DL = DAG.getDataLayout();
+    for (unsigned i = 0, e = STV->getNumValues() - 2; i < e; ++i) {
+      Value *DerivedPtr = SI.GCRelocates[i]->getDerivedPtr();
+      SDValue Res(STV.getNode(), i);
+      Type *Ty = DerivedPtr->getType();
+      unsigned Reg = FuncInfo.CreateRegs(Ty);
+      RegsForValue RFV(*Context, TLI, DL, Reg, Ty, None);
+      SDValue Chain = DAG.getEntryNode();
+
+      RFV.getCopyToRegs(Res, DAG, getCurSDLoc(), Chain, nullptr, DerivedPtr);
+      PendingExports.push_back(Chain);
+      FuncInfo.StatepointRegs[SI.StatepointInstr].push_back(Reg);
+    }
+  }
   const GCResultInst *GCResult = ISP.getGCResult();
   Type *RetTy = ISP.getActualReturnType();
   if (!RetTy->isVoidTy() && GCResult) {
     if (GCResult->getParent() != ISP.getCall()->getParent()) {
-      // Result value will be used in a different basic block so we need to
-      // export it now.  Default exporting mechanism will not work here because
-      // statepoint call has a different type than the actual call. It means
-      // that by default llvm will create export register of the wrong type
-      // (always i32 in our case). So instead we need to create export register
-      // with correct type manually.
-      // TODO: To eliminate this problem we can remove gc.result intrinsics
-      //       completely and make statepoint call to return a tuple.
+      SDValue ReturnValue = cast<StatepointSDNode>(STV)->getActualRetVal();
       unsigned Reg = FuncInfo.CreateRegs(RetTy);
       RegsForValue RFV(*DAG.getContext(), DAG.getTargetLoweringInfo(),
                        DAG.getDataLayout(), Reg, RetTy,
                        ISP.getCall()->getCallingConv());
       SDValue Chain = DAG.getEntryNode();
 
       RFV.getCopyToRegs(ReturnValue, DAG, getCurSDLoc(), Chain, nullptr);
       PendingExports.push_back(Chain);
       FuncInfo.ValueMap[ISP.getInstruction()] = Reg;
-    } else {
-      // Result value will be used in a same basic block. Don't export it or
-      // perform any explicit register copies.
-      // We'll replace the actuall call node shortly. gc_result will grab
-      // this value.
-      setValue(ISP.getInstruction(), ReturnValue);
     }
-  } else {
-    // The token value is never used from here on, just generate a poison value
-    setValue(ISP.getInstruction(), DAG.getIntPtrConstant(-1, getCurSDLoc()));
   }
 }
 
 void SelectionDAGBuilder::LowerCallSiteWithDeoptBundleImpl(
     const CallBase *Call, SDValue Callee, const BasicBlock *EHPadBB,
     bool VarArgDisallowed, bool ForceVoidReturnTy) {
   StatepointLoweringInfo SI(DAG);
   unsigned ArgBeginIndex = Call->arg_begin() - Call->op_begin();
void SelectionDAGBuilder::LowerCallSiteWithDeoptBundleImpl(
 
   SI.DeoptState =
       ArrayRef<const Use>(DeoptBundle.Inputs.begin(), DeoptBundle.Inputs.end());
   SI.StatepointFlags = static_cast<uint64_t>(StatepointFlags::None);
   SI.EHPadBB = EHPadBB;
 
   // NB! The GC arguments are deliberately left empty.
 
-  if (SDValue ReturnVal = LowerAsSTATEPOINT(SI)) {
+  SDValue STV = LowerAsSTATEPOINT(SI);
+  if (SDValue ReturnVal =
+          cast<StatepointSDNode>(STV.getNode())->getActualRetVal()) {
     ReturnVal = lowerRangeToAssertZExt(DAG, *Call, ReturnVal);
     setValue(Call, ReturnVal);
   }
 }
 
 void SelectionDAGBuilder::LowerCallSiteWithDeoptBundle(
     const CallBase *Call, SDValue Callee, const BasicBlock *EHPadBB) {
   LowerCallSiteWithDeoptBundleImpl(Call, Callee, EHPadBB,
PointerType *CalleeType = cast<PointerType>(
         ImmutableStatepoint(I).getCalledValue()->getType());
     Type *RetTy =
         cast<FunctionType>(CalleeType->getElementType())->getReturnType();
     SDValue CopyFromReg = getCopyFromRegs(I, RetTy);
 
     assert(CopyFromReg.getNode());
     setValue(&CI, CopyFromReg);
   } else {
-    setValue(&CI, getValue(I));
+    StatepointSDNode *STN = cast<StatepointSDNode>(getValue(I));
+    setValue(&CI, STN->getActualRetVal());
   }
 }
 
 void SelectionDAGBuilder::visitGCRelocate(const GCRelocateInst &Relocate) {
+  const BasicBlock *StatepointBB = Relocate.getStatepoint()->getParent();
 #ifndef NDEBUG
   // Consistency check
   // We skip this check for relocates not in the same basic block as their
   // statepoint. It would be too expensive to preserve validation info through
   // different basic blocks.
-  if (Relocate.getStatepoint()->getParent() == Relocate.getParent())
+  if (StatepointBB == Relocate.getParent())
     StatepointLowering.relocCallVisited(Relocate);
 
   auto *Ty = Relocate.getType()->getScalarType();
   if (auto IsManaged = GFI->getStrategy().isGCManagedPointer(Ty))
     assert(*IsManaged && "Non gc managed pointer relocated!");
 #endif
 
   const Value *DerivedPtr = Relocate.getDerivedPtr();
   SDValue SD = getValue(DerivedPtr);
 
   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;
 
+  auto &DPtrMap = FuncInfo.DerivedPtrMap[Relocate.getStatepoint()];
+  auto It = DPtrMap.find(Relocate.getDerivedPtr());
+  if (It != DPtrMap.end()) {
+    // This GC ptr is lowered through VReg.
+    unsigned Index = It->second;
+    SDValue Result;
+    auto &StatepointRegs = FuncInfo.StatepointRegs[Relocate.getStatepoint()];
+    if (StatepointBB != Relocate.getParent()) {
+      // Statepoint is in different basic block. Default getValue() mechanism
+      // does not work here, so we need create CopyFromRegs manually.
+      // See comment in LowerStatepoint for details.
+      assert(Index < StatepointRegs.size());
+      unsigned InReg = StatepointRegs[Index];
+      RegsForValue RFV(*DAG.getContext(), DAG.getTargetLoweringInfo(),
+                       DAG.getDataLayout(), InReg, DerivedPtr->getType(),
+                       None); // This is not an ABI copy.
+      SDValue Chain = DAG.getEntryNode();
+      Result = RFV.getCopyFromRegs(DAG, FuncInfo, getCurSDLoc(), Chain, nullptr,
+                                   DerivedPtr);
+    } else {
+      SDNode *Statepoint = getValue(Relocate.getStatepoint()).getNode();
+      Result = SDValue(Statepoint, Index);
+    }
+    setValue(&Relocate, Result);
+    return;
+  }
+
   // We didn't need to spill these special cases (constants and allocas).
   // See the handling in spillIncomingValueForStatepoint for detail.
   if (!DerivedPtrLocation) {
     setValue(&Relocate, SD);
     return;
   }
 
   unsigned Index = *DerivedPtrLocation;

llvm/lib/CodeGen/TargetLoweringBase.cpp

TargetLoweringBase::emitPatchPoint(MachineInstr &InitialMI,
   MachineFrameInfo &MFI = MF.getFrameInfo();
 
   // We're handling multiple types of operands here:
   // PATCHPOINT MetaArgs - live-in, read only, direct
   // STATEPOINT Deopt Spill - live-through, read only, indirect
   // STATEPOINT Deopt Alloca - live-through, read only, direct
   // (We're currently conservative and mark the deopt slots read/write in
   // practice.)
-  // STATEPOINT GC Spill - live-through, read/write, indirect
+  // STATEPOINT GC Spill - live-through, read/write, indirect or vreg
   // STATEPOINT GC Alloca - live-through, read/write, direct
-  // The live-in vs live-through is handled already (the live through ones are
-  // all stack slots), but we need to handle the different type of stackmap
-  // operands and memory effects here.
+  // The live-in vs live-through is handled already but we need to handle
+  // the different type of stackmap operands and memory effects here.
 
   // MI changes inside this loop as we grow operands.
   for(unsigned OperIdx = 0; OperIdx != MI->getNumOperands(); ++OperIdx) {
     MachineOperand &MO = MI->getOperand(OperIdx);
     if (!MO.isFI())
       continue;
 
     // foldMemoryOperand builds a new MI after replacing a single FI operand
     // with the canonical set of five x86 addressing-mode operands.
     int FI = MO.getIndex();
     MachineInstrBuilder MIB = BuildMI(MF, MI->getDebugLoc(), MI->getDesc());
 
+    auto transferOperand = [MIB, MI](unsigned i) {
+      const MachineOperand &O = MI->getOperand(i);
+      unsigned TiedTo = i;
+      if (O.isReg() && O.isTied())
+        TiedTo = MI->findTiedOperandIdx(i);
+      MIB.add(O);
+      if (TiedTo < i)
+        MIB->tieOperands(TiedTo, MIB->getNumOperands() - 1);
+    };
+
     // Copy operands before the frame-index.
     for (unsigned i = 0; i < OperIdx; ++i)
-      MIB.add(MI->getOperand(i));
+      transferOperand(i);
     // Add frame index operands recognized by stackmaps.cpp
     if (MFI.isStatepointSpillSlotObjectIndex(FI)) {
       // indirect-mem-ref tag, size, #FI, offset.
       // Used for spills inserted by StatepointLowering.  This codepath is not
       // used for patchpoints/stackmaps at all, for these spilling is done via
       // foldMemoryOperand callback only.
       assert(MI->getOpcode() == TargetOpcode::STATEPOINT && "sanity");
       MIB.addImm(StackMaps::IndirectMemRefOp);
       MIB.addImm(MFI.getObjectSize(FI));
       MIB.add(MI->getOperand(OperIdx));
       MIB.addImm(0);
     } else {
       // direct-mem-ref tag, #FI, offset.
       // Used by patchpoint, and direct alloca arguments to statepoints
       MIB.addImm(StackMaps::DirectMemRefOp);
       MIB.add(MI->getOperand(OperIdx));
       MIB.addImm(0);
     }
     // Copy the operands after the frame index.
     for (unsigned i = OperIdx + 1; i != MI->getNumOperands(); ++i)
-      MIB.add(MI->getOperand(i));
+      transferOperand(i);
 
     // Inherit previous memory operands.
     MIB.cloneMemRefs(*MI);
     assert(MIB->mayLoad() && "Folded a stackmap use to a non-load!");
 
     // Add a new memory operand for this FI.
     assert(MFI.getObjectOffset(FI) != -1);