[Statepoints] Change statepoint machine instr format to better suit VReg lowering.

Current Statepoint MI format is this:

   STATEPOINT
   <id>, <num patch bytes >, <num call arguments>, <call target>,
   [call arguments...],
   <StackMaps::ConstantOp>, <calling convention>,
   <StackMaps::ConstantOp>, <statepoint flags>,
   <StackMaps::ConstantOp>, <num deopt args>, [deopt args...],
   <gc base/derived pairs...> <gc allocas...>

Note that GC pointers are listed in pairs <base,derived>.
This causes base pointers to appear many times (at least twice) in
instruction, which is bad for us when VReg lowering is ON.
The problem is that machine operand tiedness is 1-1 relation, so
it might look like this:

  %vr2 = STATEPOINT ... %vr1, %vr1(tied-def0)

Since only one instance of %vr1 is tied, that may lead to incorrect
codegen (see PR46917 for more details), so we have to always spill
base pointers. This mostly defeats new VReg lowering scheme.

This patch changes statepoint instruction format so that every
gc pointer appears only once in operand list. That way they all can
be tied. Additional set of operands is added to preserve base-derived
relation required to build stackmap.
New statepoint has following format:

  STATEPOINT
  <id>, <num patch bytes>, <num call arguments>, <call target>,
  [call arguments...],
  <StackMaps::ConstantOp>, <calling convention>,
  <StackMaps::ConstantOp>, <statepoint flags>,
  <StackMaps::ConstantOp>, <num deopt args>, [deopt args...],
  <StackMaps::ConstantOp>, <num gc pointers>, [gc pointers...],
  <StackMaps::ConstantOp>, <num gc allocas>,  [gc allocas...]
  <StackMaps::ConstantOp>, <num entries in gc map>, [base/derived indices...]

Changes are:
  - every gc pointer is listed only once in a flat length-prefixed list;
  - alloca list is prefixed with its length too;
  - following alloca list is length-prefixed list of base-derived
    indices of pointers from gc pointer list. Note that indices are
    logical (number of pointer), not absolute (index of machine operand).

Differential Revision: https://reviews.llvm.org/D87154

Author: dantrushin

llvm/include/llvm/CodeGen/StackMaps.h

@@ -148,9 +148,13 @@ class PatchPointOpers {
 ///   <StackMaps::ConstantOp>, <calling convention>,
 ///   <StackMaps::ConstantOp>, <statepoint flags>,
 ///   <StackMaps::ConstantOp>, <num deopt args>, [deopt args...],
-///   <gc base/derived pairs...> <gc allocas...>
-/// Note that the last two sets of arguments are not currently length
-///   prefixed.
+///   <StackMaps::ConstantOp>, <num gc pointer args>, [gc pointer args...],
+///   <StackMaps::ConstantOp>, <num gc allocas>, [gc allocas args...],
+///   <StackMaps::ConstantOp>, <num  entries in gc map>, [base/derived pairs]
+///   base/derived pairs in gc map are logical indices into <gc pointer args>
+///   section.
+///   All gc pointers assigned to VRegs produce new value (in form of MI Def
+///   operand) and are tied to it.
 class StatepointOpers {
   // TODO:: we should change the STATEPOINT representation so that CC and
   // Flags should be part of meta operands, with args and deopt operands, and
@@ -217,6 +221,19 @@ class StatepointOpers {
   /// Return the statepoint flags.
   uint64_t getFlags() const { return MI->getOperand(getFlagsIdx()).getImm(); }
 
+  uint64_t getNumDeoptArgs() const {
+    return MI->getOperand(getNumDeoptArgsIdx()).getImm();
+  }
+
+  /// Get index of first GC pointer operand of -1 if there are none.
+  int getFirstGCPtrIdx();
+
+  /// Get vector of base/derived pairs from statepoint.
+  /// Elements are indices into GC Pointer operand list (logical).
+  /// Returns number of elements in GCMap.
+  unsigned
+  getGCPointerMap(SmallVectorImpl<std::pair<unsigned, unsigned>> &GCMap);
+
 private:
   const MachineInstr *MI;
   unsigned NumDefs;
@@ -263,7 +280,7 @@ class StackMaps {
 
   /// Get index of next meta operand.
   /// Similar to parseOperand, but does not actually parses operand meaning.
-  static unsigned getNextMetaArgIdx(MachineInstr *MI, unsigned CurIdx);
+  static unsigned getNextMetaArgIdx(const MachineInstr *MI, unsigned CurIdx);
 
   void reset() {
     CSInfos.clear();
@@ -337,6 +354,13 @@ class StackMaps {
                MachineInstr::const_mop_iterator MOE, LocationVec &Locs,
                LiveOutVec &LiveOuts) const;
 
+  /// Specialized parser of statepoint operands.
+  /// They do not directly correspond to StackMap record entries.
+  void parseStatepointOpers(const MachineInstr &MI,
+                            MachineInstr::const_mop_iterator MOI,
+                            MachineInstr::const_mop_iterator MOE,
+                            LocationVec &Locations, LiveOutVec &LiveOuts);
+
   /// Create a live-out register record for the given register @p Reg.
   LiveOutReg createLiveOutReg(unsigned Reg,
                               const TargetRegisterInfo *TRI) const;

llvm/lib/CodeGen/FixupStatepointCallerSaved.cpp

@@ -98,48 +98,6 @@ static unsigned getRegisterSize(const TargetRegisterInfo &TRI, Register Reg) {
   return TRI.getSpillSize(*RC);
 }
 
-// Advance iterator to the next stack map entry
-static MachineInstr::const_mop_iterator
-advanceToNextStackMapElt(MachineInstr::const_mop_iterator MOI) {
-  if (MOI->isImm()) {
-    switch (MOI->getImm()) {
-    default:
-      llvm_unreachable("Unrecognized operand type.");
-    case StackMaps::DirectMemRefOp:
-      MOI += 2; // <Reg>, <Imm>
-      break;
-    case StackMaps::IndirectMemRefOp:
-      MOI += 3; // <Size>, <Reg>, <Imm>
-      break;
-    case StackMaps::ConstantOp:
-      MOI += 1;
-      break;
-    }
-  }
-  return ++MOI;
-}
-
-// Return statepoint GC args as a set
-static SmallSet<Register, 8> collectGCRegs(MachineInstr &MI) {
-  StatepointOpers SO(&MI);
-  unsigned NumDeoptIdx = SO.getNumDeoptArgsIdx();
-  unsigned NumDeoptArgs = MI.getOperand(NumDeoptIdx).getImm();
-  MachineInstr::const_mop_iterator MOI(MI.operands_begin() + NumDeoptIdx + 1),
-      MOE(MI.operands_end());
-
-  // Skip deopt args
-  while (NumDeoptArgs--)
-    MOI = advanceToNextStackMapElt(MOI);
-
-  SmallSet<Register, 8> Result;
-  while (MOI != MOE) {
-    if (MOI->isReg() && !MOI->isImplicit())
-      Result.insert(MOI->getReg());
-    MOI = advanceToNextStackMapElt(MOI);
-  }
-  return Result;
-}
-
 // Try to eliminate redundant copy to register which we're going to
 // spill, i.e. try to change:
 //    X = COPY Y
@@ -411,8 +369,13 @@ class StatepointState {
   // Also cache the size of found registers.
   // Returns true if caller save registers found.
   bool findRegistersToSpill() {
+    SmallSet<Register, 8> GCRegs;
+    // All GC pointer operands assigned to registers produce new value.
+    // Since they're tied to their defs, it is enough to collect def registers.
+    for (const auto &Def : MI.defs())
+      GCRegs.insert(Def.getReg());
+
     SmallSet<Register, 8> VisitedRegs;
-    SmallSet<Register, 8> GCRegs = collectGCRegs(MI);
     for (unsigned Idx = StatepointOpers(&MI).getVarIdx(),
                   EndIdx = MI.getNumOperands();
          Idx < EndIdx; ++Idx) {

llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp

@@ -82,19 +82,6 @@ static unsigned countOperands(SDNode *Node, unsigned NumExpUses,
   return N;
 }
 
-/// Return starting index of GC operand list.
-// FIXME: need a better place for this. Put it in StackMaps?
-static unsigned getStatepointGCArgStartIdx(MachineInstr *MI) {
-  assert(MI->getOpcode() == TargetOpcode::STATEPOINT &&
-         "STATEPOINT node expected");
-  unsigned OperIdx = StatepointOpers(MI).getNumDeoptArgsIdx();
-  unsigned NumDeopts = MI->getOperand(OperIdx).getImm();
-  ++OperIdx;
-  while (NumDeopts--)
-    OperIdx = StackMaps::getNextMetaArgIdx(MI, OperIdx);
-  return OperIdx;
-}
-
 /// EmitCopyFromReg - Generate machine code for an CopyFromReg node or an
 /// implicit physical register output.
 void InstrEmitter::
@@ -993,14 +980,13 @@ EmitMachineNode(SDNode *Node, bool IsClone, bool IsCloned,
     assert(!HasPhysRegOuts && "STATEPOINT mishandled");
     MachineInstr *MI = MIB;
     unsigned Def = 0;
-    unsigned Use = getStatepointGCArgStartIdx(MI);
-    Use = StackMaps::getNextMetaArgIdx(MI, Use); // first derived
-    assert(Use < MI->getNumOperands());
+    int First = StatepointOpers(MI).getFirstGCPtrIdx();
+    assert(First > 0 && "Statepoint has Defs but no GC ptr list");
+    unsigned Use = (unsigned)First;
     while (Def < NumDefs) {
       if (MI->getOperand(Use).isReg())
         MI->tieOperands(Def++, Use);
-      Use = StackMaps::getNextMetaArgIdx(MI, Use); // next base
-      Use = StackMaps::getNextMetaArgIdx(MI, Use); // next derived
+      Use = StackMaps::getNextMetaArgIdx(MI, Use);
     }
   }
 

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

@@ -495,6 +495,7 @@ lowerIncomingStatepointValue(SDValue Incoming, bool RequireSpillSlot,
 static void
 lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
                         SmallVectorImpl<MachineMemOperand *> &MemRefs,
+                        SmallVectorImpl<SDValue> &GCPtrs,
                         DenseMap<SDValue, int> &LowerAsVReg,
                         SelectionDAGBuilder::StatepointLoweringInfo &SI,
                         SelectionDAGBuilder &Builder) {
@@ -547,21 +548,39 @@ lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
   unsigned MaxVRegPtrs =
       std::min(MaxTiedRegs, MaxRegistersForGCPointers.getValue());
 
-  LLVM_DEBUG(dbgs() << "Desiding how to lower GC Pointers:\n");
+  LLVM_DEBUG(dbgs() << "Deciding how to lower GC Pointers:\n");
+
+  // List of unique lowered GC Pointer values.
+  SmallSetVector<SDValue, 16> LoweredGCPtrs;
+  // Map lowered GC Pointer value to the index in above vector
+  DenseMap<SDValue, unsigned> GCPtrIndexMap;
+
   unsigned CurNumVRegs = 0;
-  for (const Value *P : SI.Ptrs) {
+
+  auto processGCPtr = [&](const Value *V) {
+    SDValue PtrSD = Builder.getValue(V);
+    if (!LoweredGCPtrs.insert(PtrSD))
+      return; // skip duplicates
+    GCPtrIndexMap[PtrSD] = LoweredGCPtrs.size() - 1;
+
+    assert(!LowerAsVReg.count(PtrSD) && "must not have been seen");
     if (LowerAsVReg.size() == MaxVRegPtrs)
-      break;
-    SDValue PtrSD = Builder.getValue(P);
-    if (willLowerDirectly(PtrSD) || P->getType()->isVectorTy()) {
+      return;
+    if (willLowerDirectly(PtrSD) || V->getType()->isVectorTy()) {
       LLVM_DEBUG(dbgs() << "direct/spill "; PtrSD.dump(&Builder.DAG));
-      continue;
+      return;
     }
     LLVM_DEBUG(dbgs() << "vreg "; PtrSD.dump(&Builder.DAG));
     LowerAsVReg[PtrSD] = CurNumVRegs++;
-  }
-  LLVM_DEBUG(dbgs() << LowerAsVReg.size()
-                    << " derived pointers will go in vregs\n");
+  };
+
+  // Process derived pointers first to give them more chance to go on VReg.
+  for (const Value *V : SI.Ptrs)
+    processGCPtr(V);
+  for (const Value *V : SI.Bases)
+    processGCPtr(V);
+
+  LLVM_DEBUG(dbgs() << LowerAsVReg.size() << " pointers will go in vregs\n");
 
   auto isGCValue = [&](const Value *V) {
     auto *Ty = V->getType();
@@ -589,13 +608,16 @@ lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
       reservePreviousStackSlotForValue(V, Builder);
   }
 
-  for (unsigned i = 0; i < SI.Bases.size(); ++i) {
-    SDValue SDV = Builder.getValue(SI.Bases[i]);
-    if (AlwaysSpillBase || !LowerAsVReg.count(SDV))
-      reservePreviousStackSlotForValue(SI.Bases[i], Builder);
-    SDV = Builder.getValue(SI.Ptrs[i]);
+  for (const Value *V : SI.Ptrs) {
+    SDValue SDV = Builder.getValue(V);
+    if (!LowerAsVReg.count(SDV))
+      reservePreviousStackSlotForValue(V, Builder);
+  }
+
+  for (const Value *V : SI.Bases) {
+    SDValue SDV = Builder.getValue(V);
     if (!LowerAsVReg.count(SDV))
-      reservePreviousStackSlotForValue(SI.Ptrs[i], Builder);
+      reservePreviousStackSlotForValue(V, Builder);
   }
 
   // First, prefix the list with the number of unique values to be
@@ -624,43 +646,51 @@ lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
                                  Builder);
   }
 
-  // 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], ...)
-  for (unsigned i = 0; i < SI.Bases.size(); ++i) {
-    bool RequireSpillSlot;
-    SDValue Base = Builder.getValue(SI.Bases[i]);
-    RequireSpillSlot = AlwaysSpillBase || !LowerAsVReg.count(Base);
-    lowerIncomingStatepointValue(Base, RequireSpillSlot, Ops, MemRefs,
+  // Finally, go ahead and lower all the gc arguments.
+  pushStackMapConstant(Ops, Builder, LoweredGCPtrs.size());
+  for (SDValue SDV : LoweredGCPtrs)
+    lowerIncomingStatepointValue(SDV, !LowerAsVReg.count(SDV), Ops, MemRefs,
                                  Builder);
 
-    SDValue Derived = Builder.getValue(SI.Ptrs[i]);
-    RequireSpillSlot = !LowerAsVReg.count(Derived);
-    lowerIncomingStatepointValue(Derived, RequireSpillSlot, Ops, MemRefs,
-                                 Builder);
-  }
+  // Copy to out vector. LoweredGCPtrs will be empty after this point.
+  GCPtrs = LoweredGCPtrs.takeVector();
 
   // 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
+  SmallVector<SDValue, 4> Allocas;
   for (Value *V : SI.GCArgs) {
     SDValue Incoming = Builder.getValue(V);
     if (FrameIndexSDNode *FI = dyn_cast<FrameIndexSDNode>(Incoming)) {
       // This handles allocas as arguments to the statepoint
       assert(Incoming.getValueType() == Builder.getFrameIndexTy() &&
              "Incoming value is a frame index!");
-      Ops.push_back(Builder.DAG.getTargetFrameIndex(FI->getIndex(),
-                                                    Builder.getFrameIndexTy()));
+      Allocas.push_back(Builder.DAG.getTargetFrameIndex(
+          FI->getIndex(), Builder.getFrameIndexTy()));
 
       auto &MF = Builder.DAG.getMachineFunction();
       auto *MMO = getMachineMemOperand(MF, *FI);
       MemRefs.push_back(MMO);
     }
   }
+  pushStackMapConstant(Ops, Builder, Allocas.size());
+  Ops.append(Allocas.begin(), Allocas.end());
+
+  // Now construct GC base/derived map;
+  pushStackMapConstant(Ops, Builder, SI.Ptrs.size());
+  SDLoc L = Builder.getCurSDLoc();
+  for (unsigned i = 0; i < SI.Ptrs.size(); ++i) {
+    SDValue Base = Builder.getValue(SI.Bases[i]);
+    assert(GCPtrIndexMap.count(Base) && "base not found in index map");
+    Ops.push_back(
+        Builder.DAG.getTargetConstant(GCPtrIndexMap[Base], L, MVT::i64));
+    SDValue Derived = Builder.getValue(SI.Ptrs[i]);
+    assert(GCPtrIndexMap.count(Derived) && "derived not found in index map");
+    Ops.push_back(
+        Builder.DAG.getTargetConstant(GCPtrIndexMap[Derived], L, MVT::i64));
+  }
 }
 
 SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
@@ -683,11 +713,16 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
 #endif
 
   // Lower statepoint vmstate and gcstate arguments
+
+  // All lowered meta args.
   SmallVector<SDValue, 10> LoweredMetaArgs;
+  // Lowered GC pointers (subset of above).
+  SmallVector<SDValue, 16> LoweredGCArgs;
   SmallVector<MachineMemOperand*, 16> MemRefs;
   // Maps derived pointer SDValue to statepoint result of relocated pointer.
   DenseMap<SDValue, int> LowerAsVReg;
-  lowerStatepointMetaArgs(LoweredMetaArgs, MemRefs, LowerAsVReg, SI, *this);
+  lowerStatepointMetaArgs(LoweredMetaArgs, MemRefs, LoweredGCArgs, LowerAsVReg,
+                          SI, *this);
 
   // Now that we've emitted the spills, we need to update the root so that the
   // call sequence is ordered correctly.
@@ -802,8 +837,7 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
   // Compute return values.  Provide a glue output since we consume one as
   // input.  This allows someone else to chain off us as needed.
   SmallVector<EVT, 8> NodeTys;
-  for (auto &Ptr : SI.Ptrs) {
-    SDValue SD = getValue(Ptr);
+  for (auto SD : LoweredGCArgs) {
     if (!LowerAsVReg.count(SD))
       continue;
     NodeTys.push_back(SD.getValueType());