Fix tailcall regression with compiled F# (#41206)

* Fix tailcall regression with compiled F#

This change skips instantiating stubs for direct tailcalls and instead passes the inst argument directly to the target method.

Fixes #40864

* Add regression test

Co-authored-by: Jakob Botsch Nielsen <Jakob.botsch.nielsen@gmail.com>

src/coreclr/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeHelpers.CoreCLR.cs

@@ -286,9 +286,6 @@ public static IntPtr AllocateTypeAssociatedMemory(Type type, int size)
  [MethodImpl(MethodImplOptions.InternalCall)]
  private static extern IntPtr AllocTailCallArgBuffer(int size, IntPtr gcDesc);
 
- [MethodImpl(MethodImplOptions.InternalCall)]
- private static extern void FreeTailCallArgBuffer();
-
  [MethodImpl(MethodImplOptions.InternalCall)]
  private static unsafe extern TailCallTls* GetTailCallInfo(IntPtr retAddrSlot, IntPtr* retAddr);
 
@@ -323,6 +320,12 @@ private static unsafe void DispatchTailCalls(
  finally
  {
  tls->Frame = prevFrame;
+
+ // If the arg buffer is reporting inst argument, it is safe to abandon it now
+ if (tls->ArgBuffer != IntPtr.Zero && *(int*)tls->ArgBuffer == 1 /* TAILCALLARGBUFFER_INSTARG_ONLY */)
+ {
+ *(int*)tls->ArgBuffer = 2 /* TAILCALLARGBUFFER_ABANDONED */;
+ }
  }
  }
 
@@ -481,9 +484,6 @@ internal unsafe struct TailCallTls
  {
  public PortableTailCallFrame* Frame;
  public IntPtr ArgBuffer;
- private IntPtr _argBufferSize;
- private IntPtr _argBufferGCDesc;
- private fixed byte _argBufferInline[64];
  }
 
 }

src/coreclr/src/jit/compiler.h

@@ -5579,7 +5579,6 @@ class Compiler
  GenTree* fgCreateCallDispatcherAndGetResult(GenTreeCall* origCall,
  CORINFO_METHOD_HANDLE callTargetStubHnd,
  CORINFO_METHOD_HANDLE dispatcherHnd);
- GenTree* getMethodPointerTree(CORINFO_RESOLVED_TOKEN* pResolvedToken, CORINFO_CALL_INFO* pCallInfo);
  GenTree* getLookupTree(CORINFO_RESOLVED_TOKEN* pResolvedToken,
  CORINFO_LOOKUP* pLookup,
  unsigned handleFlags,

src/coreclr/src/jit/morph.cpp

@@ -7700,6 +7700,8 @@ GenTree* Compiler::fgMorphTailCallViaHelpers(GenTreeCall* call, CORINFO_TAILCALL
  assert(!call->IsImplicitTailCall());
  assert(!fgCanFastTailCall(call, nullptr));
 
+ bool virtualCall = call->IsVirtual();
+
  // If VSD then get rid of arg to VSD since we turn this into a direct call.
  // The extra arg will be the first arg so this needs to be done before we
  // handle the retbuf below.
@@ -7734,41 +7736,39 @@ GenTree* Compiler::fgMorphTailCallViaHelpers(GenTreeCall* call, CORINFO_TAILCALL
  // where we pass instantiating stub.
  if ((help.flags & CORINFO_TAILCALL_STORE_TARGET) != 0)
  {
- // If asked to store target and we have a type arg we will store
- // instantiating stub, so in that case we should not pass the type arg.
- if (call->tailCallInfo->GetSig()->hasTypeArg())
+ JITDUMP("Adding target since VM requested it\n");
+ GenTree* target;
+ if (!virtualCall)
  {
- JITDUMP("Removing type arg");
-
- assert(call->gtCallArgs != nullptr);
- if (Target::g_tgtArgOrder == Target::ARG_ORDER_R2L)
+ if (call->gtCallType == CT_INDIRECT)
  {
- // Generic context is first arg
- call->gtCallArgs = call->gtCallArgs->GetNext();
+ noway_assert(call->gtCallAddr != nullptr);
+ target = call->gtCallAddr;
  }
  else
  {
- // Generic context is last arg
- GenTreeCall::Use** lastArgSlot = &call->gtCallArgs;
- while ((*lastArgSlot)->GetNext() != nullptr)
+ CORINFO_CONST_LOOKUP addrInfo;
+ info.compCompHnd->getFunctionEntryPoint(call->gtCallMethHnd, &addrInfo);
+
+ CORINFO_GENERIC_HANDLE handle = nullptr;
+ void* pIndirection = nullptr;
+ assert(addrInfo.accessType != IAT_PPVALUE && addrInfo.accessType != IAT_RELPVALUE);
+
+ if (addrInfo.accessType == IAT_VALUE)
  {
- lastArgSlot = &(*lastArgSlot)->NextRef();
+ handle = addrInfo.handle;
  }
-
- *lastArgSlot = nullptr;
+ else if (addrInfo.accessType == IAT_PVALUE)
+ {
+ pIndirection = addrInfo.addr;
+ }
+ target = gtNewIconEmbHndNode(handle, pIndirection, GTF_ICON_FTN_ADDR, call->gtCallMethHnd);
  }
- call->fgArgInfo = nullptr;
- }
-
- JITDUMP("Adding target since VM requested it\n");
- GenTree* target;
- if (call->tailCallInfo->IsCalli())
- {
- noway_assert(call->gtCallType == CT_INDIRECT && call->gtCallAddr != nullptr);
- target = call->gtCallAddr;
  }
  else
  {
+ assert(!call->tailCallInfo->GetSig()->hasTypeArg());
+
  CORINFO_CALL_INFO callInfo;
  unsigned flags = CORINFO_CALLINFO_LDFTN;
  if (call->tailCallInfo->IsCallvirt())
@@ -7778,19 +7778,10 @@ GenTree* Compiler::fgMorphTailCallViaHelpers(GenTreeCall* call, CORINFO_TAILCALL
 
  eeGetCallInfo(call->tailCallInfo->GetToken(), nullptr, (CORINFO_CALLINFO_FLAGS)flags, &callInfo);
 
- if (!call->tailCallInfo->IsCallvirt() ||
- ((callInfo.methodFlags & (CORINFO_FLG_FINAL | CORINFO_FLG_STATIC)) != 0) ||
- ((callInfo.methodFlags & CORINFO_FLG_VIRTUAL) == 0))
- {
- target = getMethodPointerTree(call->tailCallInfo->GetToken(), &callInfo);
- }
- else
- {
- assert(call->gtCallThisArg != nullptr);
- // TODO: Proper cloning of the this pointer.
- target = getVirtMethodPointerTree(gtCloneExpr(call->gtCallThisArg->GetNode()),
- call->tailCallInfo->GetToken(), &callInfo);
- }
+ assert(call->gtCallThisArg != nullptr);
+ // TODO: Proper cloning of the this pointer.
+ target = getVirtMethodPointerTree(gtCloneExpr(call->gtCallThisArg->GetNode()),
+ call->tailCallInfo->GetToken(), &callInfo);
  }
 
  // Insert target as last arg
@@ -8044,30 +8035,6 @@ GenTree* Compiler::fgCreateCallDispatcherAndGetResult(GenTreeCall* orig
  return finalTree;
 }
 
-//------------------------------------------------------------------------
-// getMethodPointerTree: get a method pointer tree
-//
-// Arguments:
-// pResolvedToken - resolved token of the call
-// pCallInfo - the call info of the call
-//
-// Return Value:
-// A node representing the method pointer
-//
-GenTree* Compiler::getMethodPointerTree(CORINFO_RESOLVED_TOKEN* pResolvedToken, CORINFO_CALL_INFO* pCallInfo)
-{
- switch (pCallInfo->kind)
- {
- case CORINFO_CALL:
- return new (this, GT_FTN_ADDR) GenTreeFptrVal(TYP_I_IMPL, pCallInfo->hMethod);
- case CORINFO_CALL_CODE_POINTER:
- return getLookupTree(pResolvedToken, &pCallInfo->codePointerLookup, GTF_ICON_FTN_ADDR, pCallInfo->hMethod);
- default:
- noway_assert(!"unknown call kind");
- return nullptr;
- }
-}
-
 //------------------------------------------------------------------------
 // getLookupTree: get a lookup tree
 //

src/coreclr/src/vm/corelib.h

@@ -708,7 +708,6 @@ DEFINE_METHOD(RUNTIME_HELPERS, ENUM_EQUALS, EnumEquals, NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS, ENUM_COMPARE_TO, EnumCompareTo, NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS, ALLOC_TAILCALL_ARG_BUFFER, AllocTailCallArgBuffer, SM_Int_IntPtr_RetIntPtr)
 DEFINE_METHOD(RUNTIME_HELPERS, GET_TAILCALL_INFO, GetTailCallInfo, NoSig)
-DEFINE_METHOD(RUNTIME_HELPERS, FREE_TAILCALL_ARG_BUFFER, FreeTailCallArgBuffer, SM_RetVoid)
 DEFINE_METHOD(RUNTIME_HELPERS, DISPATCH_TAILCALLS, DispatchTailCalls, NoSig)
 
 DEFINE_CLASS(UNSAFE, InternalCompilerServices, Unsafe)
@@ -760,9 +759,6 @@ DEFINE_FIELD(PORTABLE_TAIL_CALL_FRAME, NEXT_CALL, NextCall)
 DEFINE_CLASS(TAIL_CALL_TLS, CompilerServices, TailCallTls)
 DEFINE_FIELD(TAIL_CALL_TLS, FRAME, Frame)
 DEFINE_FIELD(TAIL_CALL_TLS, ARG_BUFFER, ArgBuffer)
-DEFINE_FIELD(TAIL_CALL_TLS, ARG_BUFFER_SIZE, _argBufferSize)
-DEFINE_FIELD(TAIL_CALL_TLS, ARG_BUFFER_GC_DESC, _argBufferGCDesc)
-DEFINE_FIELD(TAIL_CALL_TLS, ARG_BUFFER_INLINE, _argBufferInline)
 
 DEFINE_CLASS_U(CompilerServices, PortableTailCallFrame, PortableTailCallFrame)
 DEFINE_FIELD_U(Prev, PortableTailCallFrame, Prev)
@@ -772,10 +768,6 @@ DEFINE_FIELD_U(NextCall, PortableTailCallFrame, NextCall)
 DEFINE_CLASS_U(CompilerServices, TailCallTls, TailCallTls)
 DEFINE_FIELD_U(Frame, TailCallTls, m_frame)
 DEFINE_FIELD_U(ArgBuffer, TailCallTls, m_argBuffer)
-DEFINE_FIELD_U(_argBufferSize, TailCallTls, m_argBufferSize)
-DEFINE_FIELD_U(_argBufferGCDesc, TailCallTls, m_argBufferGCDesc)
-DEFINE_FIELD_U(_argBufferInline, TailCallTls, m_argBufferInline)
-
 
 DEFINE_CLASS(RUNTIME_WRAPPED_EXCEPTION, CompilerServices, RuntimeWrappedException)
 DEFINE_METHOD(RUNTIME_WRAPPED_EXCEPTION, OBJ_CTOR, .ctor, IM_Obj_RetVoid)

src/coreclr/src/vm/ecalllist.h

@@ -881,7 +881,6 @@ FCFuncStart(gRuntimeHelpers)
  FCFuncElement("GetUninitializedObjectInternal", ReflectionSerialization::GetUninitializedObject)
  QCFuncElement("AllocateTypeAssociatedMemoryInternal", RuntimeTypeHandle::AllocateTypeAssociatedMemory)
  FCFuncElement("AllocTailCallArgBuffer", TailCallHelp::AllocTailCallArgBuffer)
- FCFuncElement("FreeTailCallArgBuffer", TailCallHelp::FreeTailCallArgBuffer)
  FCFuncElement("GetTailCallInfo", TailCallHelp::GetTailCallInfo)
  FCFuncElement("GetILBytesJitted", GetJittedBytes)
  FCFuncElement("GetMethodsJittedCount", GetJittedMethodsCount)

src/coreclr/src/vm/gcenv.ee.cpp

@@ -157,27 +157,33 @@ static void ScanStackRoots(Thread * pThread, promote_func* fn, ScanContext* sc)
 
 static void ScanTailCallArgBufferRoots(Thread* pThread, promote_func* fn, ScanContext* sc)
 {
- void* gcDesc;
- char* argBuffer = pThread->GetTailCallTls()->GetArgBuffer(&gcDesc);
- if (gcDesc == NULL)
+ TailCallArgBuffer* argBuffer = pThread->GetTailCallTls()->GetArgBuffer();
+ if (argBuffer == NULL || argBuffer->GCDesc == NULL)
  return;
 
- GCRefMapDecoder decoder(static_cast<PTR_BYTE>(gcDesc));
+ if (argBuffer->State == TAILCALLARGBUFFER_ABANDONED)
+ return;
+
+ bool instArgOnly = argBuffer->State == TAILCALLARGBUFFER_INSTARG_ONLY;
+
+ GCRefMapDecoder decoder(static_cast<PTR_BYTE>(argBuffer->GCDesc));
  while (!decoder.AtEnd())
  {
  int pos = decoder.CurrentPos();
  int token = decoder.ReadToken();
 
- PTR_TADDR ppObj = dac_cast<PTR_TADDR>(argBuffer + pos * sizeof(TADDR));
+ PTR_TADDR ppObj = dac_cast<PTR_TADDR>(((BYTE*)argBuffer->Args) + pos * sizeof(TADDR));
  switch (token)
  {
  case GCREFMAP_SKIP:
  break;
  case GCREFMAP_REF:
- fn(dac_cast<PTR_PTR_Object>(ppObj), sc, CHECK_APP_DOMAIN);
+ if (!instArgOnly)
+ fn(dac_cast<PTR_PTR_Object>(ppObj), sc, CHECK_APP_DOMAIN);
  break;
  case GCREFMAP_INTERIOR:
- PromoteCarefully(fn, dac_cast<PTR_PTR_Object>(ppObj), sc, GC_CALL_INTERIOR);
+ if (!instArgOnly)
+ PromoteCarefully(fn, dac_cast<PTR_PTR_Object>(ppObj), sc, GC_CALL_INTERIOR);
  break;
  case GCREFMAP_METHOD_PARAM:
  if (sc->promotion)

src/coreclr/src/vm/jitinterface.cpp

@@ -13952,7 +13952,7 @@ bool CEEInfo::getTailCallHelpersInternal(CORINFO_RESOLVED_TOKEN* callToken,
 
  TailCallHelp::CreateTailCallHelperStubs(
  m_pMethodBeingCompiled, pTargetMD,
- msig, isCallvirt, isThisArgByRef,
+ msig, isCallvirt, isThisArgByRef, sig->hasTypeArg(),
  &pStoreArgsMD, &needsTarget,
  &pCallTargetMD);