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

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.

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

Author: preames

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

@@ -90,13 +90,13 @@ void StatepointLoweringState::startNewStatepoint(SelectionDAGBuilder &Builder) {
   // 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");
 }
@@ -691,8 +691,9 @@ lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
 
       // 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));
       }
@@ -857,15 +858,31 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
 
   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.
@@ -1101,6 +1118,22 @@ void SelectionDAGBuilder::visitGCRelocate(const GCRelocateInst &Relocate) {
 #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) {
@@ -1110,17 +1143,6 @@ void SelectionDAGBuilder::visitGCRelocate(const GCRelocateInst &Relocate) {
     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");

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.h

@@ -103,9 +103,9 @@ class StatepointLoweringState {
     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)

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

@@ -310,10 +310,11 @@ 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