[Statepoint] When using the tied def lowering, unconditionally use vregs [almost NFC]

[reames]

Summary

This builds on 3da1a96 on the path towards supporting invokes and cross block relocations. The actual change attempts to be NFC, but does fail in one corner-case explained below.

The change itself is fairly mechanical. Rather than remember SDValues - which are inherently block local - immediately produce a virtual register copy and remember that.

Once this lands, we'll update the FunctionLoweringInfo::StatepointSpillMap map to allow register based lowerings, delete VirtRegs from StatepointLowering, and drop the restriction against cross block relocations. I deliberately separate the semantic part into it's own change for easy of understanding and fault isolation.

The corner-case which isn't quite NFC is that the old implementation implicitly CSEd gc.relocates of the same SDValue regardless of type. The new implementation still only relocates once, but it produces distinct vregs for the bitcast and it's source, whereas SelectionDAG's generic CSE was able to remove the bitcast in the old implementation. Note that the final assembly doesn't change (at least in the test), as our MI level optimizations catch the duplication.

I assert that this is an uninteresting corner-case. It's functionally correct, and if we find a case where this influences performance, we should really be canonicalizing types to i8* at the IR level.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

878

Register

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.h

108

Should use Register

[reames]

Address Register related review comments

[skatkov]

lgtm

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

1120

do you really need these empty lines?

1122

redundant dot.

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.h

public:
   }
 
   bool isStackSlotAllocated(int Offset) {
     assert(Offset >= 0 && Offset < (int)AllocatedStackSlots.size() &&
            "out of bounds");
     return AllocatedStackSlots.test(Offset);
   }
 
-  /// For each statepoint keep mapping from original derived pointer to
-  /// the statepoint node result defining its new value.
-  DenseMap<const Value *, SDValue> DerivedPtrMap;
+  /// 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)
   DenseMap<SDValue, SDValue> Locations;
 
   /// A boolean indicator for each slot listed in the FunctionInfo as to
   /// whether it has been used in the current statepoint.  Since we try to

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

assert(PendingGCRelocateCalls.empty() &&
          "Trying to visit statepoint before finished processing previous one");
   Locations.clear();
   NextSlotToAllocate = 0;
   // Need to resize this on each safepoint - we need the two to stay in sync and
   // the clear patterns of a SelectionDAGBuilder have no relation to
   // FunctionLoweringInfo.  Also need to ensure used bits get cleared.
   AllocatedStackSlots.clear();
   AllocatedStackSlots.resize(Builder.FuncInfo.StatepointStackSlots.size());
-  DerivedPtrMap.clear();
+  VirtRegs.clear();
 }
 
 void StatepointLoweringState::clear() {
   Locations.clear();
   AllocatedStackSlots.clear();
-  DerivedPtrMap.clear();
+  VirtRegs.clear();
   assert(PendingGCRelocateCalls.empty() &&
          "cleared before statepoint sequence completed");
 }
 
 SDValue
 StatepointLoweringState::allocateStackSlot(EVT ValueType,
                                            SelectionDAGBuilder &Builder) {
   NumSlotsAllocatedForStatepoints++;
if (Loc.getNode()) {
       // 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, but will likely be extended for vreg case shortly.
-      if (Relocate->getParent() != StatepointInstr->getParent()) {
+      // relocation (such as constants and allocas).
+      if (!LowerAsVReg.count(SDV) &&
+          Relocate->getParent() != StatepointInstr->getParent()) {
         Builder.ExportFromCurrentBlock(V);
         assert(!LowerAsVReg.count(SDV));
       }
     }
   }
 }
 
 SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
#endif
 
   unsigned NumResults = NodeTys.size();
   MachineSDNode *StatepointMCNode =
     DAG.getMachineNode(TargetOpcode::STATEPOINT, getCurSDLoc(), NodeTys, Ops);
   DAG.setNodeMemRefs(StatepointMCNode, MemRefs);
 
   SDNode *SinkNode = StatepointMCNode;
 
-  // Fill mapping from derived pointer to statepoint result denoting its
-  // relocated value.
-  auto &DPtrMap = StatepointLowering.DerivedPtrMap;
+  // 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 (const auto *Relocate : SI.GCRelocates) {
     Value *Derived = Relocate->getDerivedPtr();
     SDValue SD = getValue(Derived);
     if (!LowerAsVReg.count(SD))
       continue;
-    DPtrMap[Derived] = SDValue(StatepointMCNode, LowerAsVReg[SD]);
+
+    // Handle multiple gc.relocates of the same input efficiently.
+    if (VirtRegs.count(Derived))
+      continue;
+
+    SDValue Relocated = SDValue(StatepointMCNode, LowerAsVReg[SD]);
+
+    auto *RetTy = Relocate->getType();
+    Register Reg = FuncInfo.CreateRegs(RetTy);
+    RegsForValue RFV(*DAG.getContext(), DAG.getTargetLoweringInfo(),
+                     DAG.getDataLayout(), Reg, RetTy, None);
+    SDValue Chain = DAG.getEntryNode();
+    RFV.getCopyToRegs(Relocated, DAG, getCurSDLoc(), Chain, nullptr);
+    PendingExports.push_back(Chain);
+    
+    VirtRegs[Derived] = Reg; 
   }
 
   // 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;
#ifndef NDEBUG
   if (Relocate.getStatepoint()->getParent() == 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();
+
+  // If relocation was done via virtual register..
+  auto &VirtRegs = StatepointLowering.VirtRegs;
+  auto It = VirtRegs.find(DerivedPtr);
+  if (It != VirtRegs.end()) {
+    Register InReg = It->second;
+    RegsForValue RFV(*DAG.getContext(), DAG.getTargetLoweringInfo(),
+                     DAG.getDataLayout(), InReg, Relocate.getType(),
+                     None); // This is not an ABI copy.
+    SDValue Chain = DAG.getEntryNode();
+    SDValue Relocation = RFV.getCopyFromRegs(DAG, FuncInfo, getCurSDLoc(),
+                                             Chain, nullptr, nullptr);
+    setValue(&Relocate, Relocation);
+    return;
+  }
+  
   SDValue SD = getValue(DerivedPtr);
 
   if (SD.isUndef() && SD.getValueType().getSizeInBits() <= 64) {
     // Lowering relocate(undef) as arbitrary constant. Current constant value
     // is chosen such that it's unlikely to be a valid pointer.
     setValue(&Relocate, DAG.getTargetConstant(0xFEFEFEFE, SDLoc(SD), MVT::i64));
     return;
   }
 
-  // Relocate is local to statepoint block and its pointer was assigned
-  // to VReg. Use corresponding statepoint result.
-  auto &DPtrMap = StatepointLowering.DerivedPtrMap;
-  auto It = DPtrMap.find(DerivedPtr);
-  if (It != DPtrMap.end()) {
-    setValue(&Relocate, It->second);
-    assert(Relocate.getParent() == Relocate.getStatepoint()->getParent() &&
-           "unexpected DPtrMap entry");
-    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) {

llvm/test/CodeGen/X86/statepoint-vreg.ll

; CHECK-ASM-NEXT: retq
   ret void
 }
 
 ; Two gc.relocates of a bitcasted pointer should only require a single spill/fill
 define void @test_gcptr_uniqueing(i32 addrspace(1)* %ptr) gc "statepoint-example" {
 ; CHECK-VREG-LABEL: name:            test_gcptr_uniqueing
 ; CHECK-VREG:    %0:gr64 = COPY $rdi
 ; CHECK-VREG:    ADJCALLSTACKDOWN64 0, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
-; CHECK-VREG:    %1:gr64 = STATEPOINT 0, 0, 0, @func, 2, 0, 2, 0, 2, 2, %0, 2, 4278124286, %0, %0(tied-def 0), csr_64, implicit-def $rsp, implicit-def $ssp
+; CHECK-VREG:    %2:gr64 = STATEPOINT 0, 0, 0, @func, 2, 0, 2, 0, 2, 2, %0, 2, 4278124286, %0, %0(tied-def 0), csr_64, implicit-def $rsp, implicit-def $ssp
 ; CHECK-VREG:    ADJCALLSTACKUP64 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+; CHECK-VREG:    %1:gr64 = COPY %2
 ; CHECK-VREG:    ADJCALLSTACKDOWN64 0, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
-; CHECK-VREG:    $rdi = COPY %1
+; CHECK-VREG:    $rdi = COPY %2
 ; CHECK-VREG:    $rsi = COPY %1
 ; CHECK-VREG:    CALL64pcrel32 @use1, csr_64, implicit $rsp, implicit $ssp, implicit $rdi, implicit $rsi, implicit-def $rsp, implicit-def $ssp
 
 ; CHECK-PREG-LABEL: name:            test_gcptr_uniqueing
 ; CHECK-PREG:    renamable $rbx = COPY $rdi
 ; CHECK-PREG:    renamable $rbx = STATEPOINT 0, 0, 0, @func, 2, 0, 2, 0, 2, 2, killed renamable $rbx, 2, 4278124286, renamable $rbx, renamable $rbx(tied-def 0), csr_64, implicit-def $rsp, implicit-def $ssp
 ; CHECK-PREG:    $rdi = COPY renamable $rbx
 ; CHECK-PREG:    $rsi = COPY killed renamable $rbx