[StatepointLowering] Properly handle local and non-local relocates of the same value.

[dantrushin]

Summary

 FunctionLoweringInfo::StatepointRelocationMaps map is used to pass GC pointer
 lowering information from statepoint to gc.relocate  which may appear in the
 different block.
 D124444 introduced different lowering for local and non-local relocates.
 Local relocates use statepoint SDValue and non-local relocates use value
exported to VReg.
 But I overlooked the fact that StatepointRelocationMap is indexed not by
 GCRelocate instruction, but by derived pointer. This works incorrectly when
 we have two relocates (one local and another non-local) of the same value,
 because they need different relocation records.
 
 This patch fixes the problem by recording relocation information per relocate
 instruction, not per derived pointer. This way, each gc.relocate can be lowered
 differently.

[skatkov]

Good catch.

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

static Optional<int> findPreviousSpillSlot(const Value *Val,
   if (LookUpDepth <= 0)
     return None;
 
   // Spill location is known for gc relocates
   if (const auto *Relocate = dyn_cast<GCRelocateInst>(Val)) {
     const auto &RelocationMap =
         Builder.FuncInfo.StatepointRelocationMaps[Relocate->getStatepoint()];
 
-    auto It = RelocationMap.find(Relocate->getDerivedPtr());
+    auto It = RelocationMap.find(Relocate);
     if (It == RelocationMap.end())
       return None;
 
     auto &Record = It->second;
     if (Record.type != RecordType::Spill)
       return None;
 
     return Record.payload.FI;
if (IsLocal && LowerAsVReg.count(SDV)) {
       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;
+    RelocationMap[Relocate] = Record;
   }
 
   
 
   SDNode *SinkNode = StatepointMCNode;
 
   // Build the GC_TRANSITION_END node if necessary.
   //
#ifndef NDEBUG
   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();
   auto &RelocationMap =
     FuncInfo.StatepointRelocationMaps[Relocate.getStatepoint()];
-  auto SlotIt = RelocationMap.find(DerivedPtr);
+  auto SlotIt = RelocationMap.find(&Relocate);
   assert(SlotIt != RelocationMap.end() && "Relocating not lowered gc value");
   const RecordType &Record = SlotIt->second;
 
   // If relocation was done via virtual register..
   if (Record.type == RecordType::SDValueNode) {
     assert(Relocate.getStatepoint()->getParent() == Relocate.getParent() &&
            "Nonlocal gc.relocate mapped via SDValue");
     SDValue SDV = StatepointLowering.getLocation(getValue(DerivedPtr));

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

 
 declare dso_local i1 @return_i1()
 declare dso_local void @func()
 declare dso_local void @consume(i32 addrspace(1)*)
 declare dso_local void @consume2(i32 addrspace(1)*, i32 addrspace(1)*)
 declare dso_local void @consume5(i32 addrspace(1)*, i32 addrspace(1)*, i32 addrspace(1)*, i32 addrspace(1)*, i32 addrspace(1)*)
 declare dso_local void @use1(i32 addrspace(1)*, i8 addrspace(1)*)
 declare dso_local void @bar(i8 addrspace(1)*, i8 addrspace(1)*)
+declare i8 addrspace(1)* @dummy(i32)
 
 ; test most simple relocate
 define i1 @test_relocate(i32 addrspace(1)* %a) gc "statepoint-example" {
 ; CHECK-VREG-LABEL: name:            test_relocate
 ; CHECK-VREG:    %0:gr64 = COPY $rdi
 ; CHECK-VREG:    %1:gr64 = STATEPOINT 0, 0, 0, @return_i1, 2, 0, 2, 0, 2, 0, 2, 1, %0(tied-def 0), 2, 0, 2, 1, 0, 0, csr_64, implicit-def $rsp, implicit-def $ssp, implicit-def $al
 ; CHECK-VREG:    %2:gr8 = COPY $al
 ; CHECK-VREG:    $rdi = COPY %1
left:
   %call2 = call zeroext i1 @llvm.experimental.gc.result.i1(token %safepoint_token)
   call void @consume(i32 addrspace(1)* %call1)
   ret i1 %call2
 
 right:
   ret i1 true
 }
 
+; Local and non-local relocates of the same value
+; CHECK-VREG-LABEL: name:            test_local_non_local_reloc
+; CHECK-VREG:  bb.0.entry:
+; CHECK-VREG:    %2:gr64 = COPY $rsi
+; CHECK-VREG:    %1:gr32 = COPY $edi
+; CHECK-VREG:    %4:gr8 = COPY %1.sub_8bit
+; CHECK-VREG:    ADJCALLSTACKDOWN64 0, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+; CHECK-VREG:    %5:gr32 = MOV32r0 implicit-def dead $eflags
+; CHECK-VREG:    $edi = COPY %5
+; CHECK-VREG:    %6:gr64 = IMPLICIT_DEF
+; CHECK-VREG:    %0:gr64 = STATEPOINT 2, 5, 1, killed %6, $edi, 2, 0, 2, 0, 2, 0, 2, 1, %2(tied-def 0), 2, 0, 2, 1, 0, 0, csr_64, implicit-def $rsp, implicit-def $ssp, implicit-def $rax
+; CHECK-VREG:    ADJCALLSTACKUP64 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+; CHECK-VREG:    %7:gr64 = COPY $rax
+; CHECK-VREG:    %3:gr64 = COPY %0
+; CHECK-VREG:    TEST8ri killed %4, 1, implicit-def $eflags
+; CHECK-VREG:    JCC_1 %bb.2, 5, implicit $eflags
+; CHECK-VREG:    JMP_1 %bb.1
+; CHECK-VREG:  bb.1.left:
+; CHECK-VREG:    $rax = COPY %3
+; CHECK-VREG:    RET 0, $rax
+; CHECK-VREG:  bb.2.right:
+; CHECK-VREG:    $rax = COPY %0
+; CHECK-VREG:    RET 0, $rax
+define i8 addrspace(1)* @test_local_non_local_reloc(i1 %c, i8 addrspace(1)* %p) gc "statepoint-example" {
+entry:
+  %statepoint = call token (i64, i32, i8 addrspace(1)* (i32)*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_p1i8i32f(i64 2, i32 5, i8 addrspace(1)* (i32)* nonnull elementtype(i8 addrspace(1)* (i32)) @dummy, i32 1, i32 0, i32 0, i32 0, i32 0) [ "deopt"(), "gc-live"(i8 addrspace(1)* %p) ]
+  %p.relocated = call coldcc i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token %statepoint, i32 0, i32 0) ; (%p, %p)
+  br i1 %c, label %right, label %left
+
+left:
+  %p.relocated.2 = call coldcc i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token %statepoint, i32 0, i32 0) ; (%p, %p)
+  ret i8 addrspace(1)* %p.relocated.2
+
+right:
+  ret i8 addrspace(1)* %p.relocated
+}
+
 ; No need to check post-regalloc output as it is the same
 define i1 @duplicate_reloc() gc "statepoint-example" {
 ; CHECK-VREG-LABEL: name:            duplicate_reloc
 ; CHECK-VREG:  bb.0.entry:
 ; CHECK-VREG:    STATEPOINT 0, 0, 0, @func, 2, 0, 2, 0, 2, 0, 2, 1, 2, 0, 2, 0, 2, 1, 0, 0, csr_64, implicit-def $rsp, implicit-def $ssp
 ; CHECK-VREG:    STATEPOINT 0, 0, 0, @func, 2, 0, 2, 0, 2, 0, 2, 1, 2, 0, 2, 0, 2, 1, 0, 0, csr_64, implicit-def $rsp, implicit-def $ssp
 ; CHECK-VREG:    %0:gr8 = MOV8ri 1
 ; CHECK-VREG:    $al = COPY %0

 
 declare token @llvm.experimental.gc.statepoint.p0f_i1f(i64, i32, i1 ()*, i32, i32, ...)
 declare token @llvm.experimental.gc.statepoint.p0f_isVoidf(i64, i32, void ()*, i32, i32, ...)
 declare dso_local i32 addrspace(1)* @llvm.experimental.gc.relocate.p1i32(token, i32, i32)
 declare dso_local i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token, i32, i32)
 declare <2 x i8 addrspace(1)*> @llvm.experimental.gc.relocate.v2p1i8(token, i32, i32)
 declare dso_local i1 @llvm.experimental.gc.result.i1(token)
 declare token @llvm.experimental.gc.statepoint.p0f_isVoidp1i8p1i8f(i64 immarg, i32 immarg, void (i8 addrspace(1)*, i8 addrspace(1)*)*, i32 immarg, i32 immarg, ...)
+declare token @llvm.experimental.gc.statepoint.p0f_p1i8i32f(i64, i32, i8 addrspace(1)* (i32)*, i32, i32, ...)