[Statepoints] Properly handle const base pointer.

[dantrushin]

Summary

Current code in InstEmitter assumes all GC pointers are either
VRegs or stack slots - hence, taking only one operand.
But it is possible to have constant base, in which case it
occupies two machine operands.

Add a convinience function to StackMaps to get index of next
meta argument and use it in InsrEmitter to properly advance to
the next statepoint meta operand.

[reames]

LGTM. Good catch, and nicely structured cleanup in the process.

llvm/include/llvm/CodeGen/StackMaps.h

public:
 
   // OpTypes are used to encode information about the following logical
   // operand (which may consist of several MachineOperands) for the
   // OpParser.
   using OpType = enum { DirectMemRefOp, IndirectMemRefOp, ConstantOp };
 
   StackMaps(AsmPrinter &AP);
 
+  /// Get index of next meta operand.
+  /// Similar to parseOperand, but does not actually parses operand meaning.
+  static unsigned getNextMetaArgIdx(MachineInstr *MI, unsigned CurIdx);
+
   void reset() {
     CSInfos.clear();
     ConstPool.clear();
     FnInfos.clear();
   }
 
   using LocationVec = SmallVector<Location, 8>;
   using LiveOutVec = SmallVector<LiveOutReg, 8>;

llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp

 
 /// 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();
-  // At this point stack references has not been lowered yet, so they
-  // take single operand.
   ++OperIdx;
-  while (NumDeopts--) {
-    MachineOperand &MO = MI->getOperand(OperIdx);
-    if (MO.isImm() && MO.getImm() == StackMaps::ConstantOp) {
-      ++OperIdx;
-      assert(MI->getOperand(OperIdx).isImm() &&
-             "Unexpected statepoint operand");
-    }
-    ++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::
 EmitCopyFromReg(SDNode *Node, unsigned ResNo, bool IsClone, bool IsCloned,
                 Register SrcReg, DenseMap<SDValue, Register> &VRBaseMap) {
if (!UsedRegs.empty() || II.getImplicitDefs() || II.hasOptionalDef())
     MIB->setPhysRegsDeadExcept(UsedRegs, *TRI);
 
   // STATEPOINT is too 'dynamic' to have meaningful machine description.
   // We have to manually tie operands.
   if (Opc == TargetOpcode::STATEPOINT && NumDefs > 0) {
     assert(!HasPhysRegOuts && "STATEPOINT mishandled");
     MachineInstr *MI = MIB;
     unsigned Def = 0;
-    unsigned Use = getStatepointGCArgStartIdx(MI) + 1;
+    unsigned Use = getStatepointGCArgStartIdx(MI);
+    Use = StackMaps::getNextMetaArgIdx(MI, Use); // first derived
+    assert(Use < MI->getNumOperands());
     while (Def < NumDefs) {
       if (MI->getOperand(Use).isReg())
         MI->tieOperands(Def++, Use);
-      Use += 2;
+      Use = StackMaps::getNextMetaArgIdx(MI, Use); // next base
+      Use = StackMaps::getNextMetaArgIdx(MI, Use); // next derived
     }
   }
 
   // Run post-isel target hook to adjust this instruction if needed.
   if (II.hasPostISelHook())
     TLI->AdjustInstrPostInstrSelection(*MIB, Node);
 }
 

llvm/lib/CodeGen/StackMaps.cpp

unsigned PatchPointOpers::getNextScratchIdx(unsigned StartIdx) const {
   return ScratchIdx;
 }
 
 StackMaps::StackMaps(AsmPrinter &AP) : AP(AP) {
   if (StackMapVersion != 3)
     llvm_unreachable("Unsupported stackmap version!");
 }
 
+unsigned StackMaps::getNextMetaArgIdx(MachineInstr *MI, unsigned CurIdx) {
+  assert(CurIdx < MI->getNumOperands() && "Bad meta arg index");
+  const auto &MO = MI->getOperand(CurIdx);
+  if (MO.isImm()) {
+    switch (MO.getImm()) {
+    default:
+      llvm_unreachable("Unrecognized operand type.");
+    case StackMaps::DirectMemRefOp:
+      CurIdx += 2;
+      break;
+    case StackMaps::IndirectMemRefOp:
+      CurIdx += 3;
+      break;
+    case StackMaps::ConstantOp:
+      ++CurIdx;
+      break;
+    }
+  }
+  ++CurIdx;
+  assert(CurIdx < MI->getNumOperands() && "points past operand list");
+  return CurIdx;
+}
+
 /// Go up the super-register chain until we hit a valid dwarf register number.
 static unsigned getDwarfRegNum(unsigned Reg, const TargetRegisterInfo *TRI) {
   int RegNum = TRI->getDwarfRegNum(Reg, false);
   for (MCSuperRegIterator SR(Reg, TRI); SR.isValid() && RegNum < 0; ++SR)
     RegNum = TRI->getDwarfRegNum(*SR, false);
 
   assert(RegNum >= 0 && "Invalid Dwarf register number.");
   return (unsigned)RegNum;

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

 ; CHECK-NEXT:    retq
 entry:
   %safepoint_token = tail call token (i64, i32, i1 ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_i1f(i64 0, i32 0, i1 ()* @return_i1, i32 0, i32 0, i32 0, i32 0) ["gc-live" (i32 addrspace(1)* %a)]
   %rel1 = call i32 addrspace(1)* @llvm.experimental.gc.relocate.p1i32(token %safepoint_token,  i32 0, i32 0)
   %res1 = call zeroext i1 @llvm.experimental.gc.result.i1(token %safepoint_token)
   call void @consume(i32 addrspace(1)* %rel1)
   ret i1 %res1
 }
+
 ; test pointer variables intermixed with pointer constants
 define void @test_mixed(i32 addrspace(1)* %a, i32 addrspace(1)* %b, i32 addrspace(1)* %c) gc "statepoint-example" {
 ; CHECK-LABEL: test_mixed:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    pushq %r15
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    pushq %r14
 ; CHECK-NEXT:    .cfi_def_cfa_offset 24

 
 exceptional_return.right:
   %landing_pad1 = landingpad token
           cleanup
   %val.relocated3 = call coldcc i64 addrspace(1)* @llvm.experimental.gc.relocate.p1i64(token %landing_pad1, i32 0, i32 0)
   ret i64 addrspace(1)* %val.relocated3
 }
 
+; test ISEL for constant base pointer - must properly tie operands
+define void @test_const_base(i32 addrspace(1)* %a) gc "statepoint-example" {
+; CHECK-LABEL: test_const_base:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset %rbx, -16
+; CHECK-NEXT:    movq %rdi, %rbx
+; CHECK-NEXT:    callq func
+; CHECK-NEXT:  .Ltmp24:
+; CHECK-NEXT:    movq %rbx, %rdi
+; CHECK-NEXT:    callq consume
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 8
+; CHECK-NEXT:    retq
+entry:
+  %token1 = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @func, i32 0, i32 0, i32 0, i32 0) ["deopt" (i32 0, i32 1, i32 7, i32 addrspace(1)* null, i32 9), "gc-live" (i32 addrspace(1)* null, i32 addrspace(1)* %a)]
+  %rel = call i32 addrspace(1)* @llvm.experimental.gc.relocate.p1i32(token %token1,  i32 0, i32 1)
+  call void @consume(i32 addrspace(1)* %rel)
+  ret void
+}
+
 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 token @llvm.experimental.gc.statepoint.p0f_isVoidp1i64f(i64, i32, void (i64 addrspace(1)*)*, i32, i32, ...)
 declare i32 addrspace(1)* @llvm.experimental.gc.relocate.p1i32(token, i32, i32)
 declare i64 addrspace(1)* @llvm.experimental.gc.relocate.p1i64(token, i32, i32)
 declare 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 i1 @llvm.experimental.gc.result.i1(token)