ARM64: Avoid jump stubs where possible

src/coreclr/jit/lower.cpp

@@ -3704,11 +3704,17 @@ GenTree* Lowering::LowerDirectCall(GenTreeCall* call)
     switch (accessType)
     {
         case IAT_VALUE:
+        {
+            bool targetIsNear = true;
+#ifdef TARGET_ARM64
+            targetIsNear = comp->eeGetRelocTypeHint(addr) == IMAGE_REL_ARM64_BRANCH26;
+#endif
+
             // Non-virtual direct call to known address.
             // For JIT helper based tailcall (only used on x86) the target
             // address is passed as an arg to the helper so we want a node for
             // it.
-            if (!IsCallTargetInRange(addr) || call->IsTailCallViaJitHelper())
+            if (!targetIsNear || !IsCallTargetInRange(addr) || call->IsTailCallViaJitHelper())
             {
                 result = AddrGen(addr);
             }
@@ -3719,6 +3725,7 @@ GenTree* Lowering::LowerDirectCall(GenTreeCall* call)
                 call->gtDirectCallAddress = addr;
             }
             break;
+        }
 
         case IAT_PVALUE:
         {

src/coreclr/vm/codeman.h

@@ -444,6 +444,8 @@ class CodeHeap
     // Space for header is reserved immediately before. It is not included in size.
     virtual void* AllocMemForCode_NoThrow(size_t header, size_t size, DWORD alignment, size_t reserveForJumpStubs) = 0;
 
+    virtual PTR_BYTE GetAllocPtr() = 0;
+
 #ifdef DACCESS_COMPILE
     virtual void EnumMemoryRegions(CLRDataEnumMemoryFlags flags) = 0;
 #endif
@@ -538,6 +540,11 @@ class LoaderCodeHeap : CodeHeap
         m_LoaderHeap.EnumMemoryRegions(flags);
     }
 #endif
+
+    virtual PTR_BYTE GetAllocPtr()
+    {
+        return m_LoaderHeap.GetAllocPtr();
+    }
 };
 
 #if defined(HOST_64BIT)

src/coreclr/vm/dynamicmethod.h

@@ -298,6 +298,8 @@ class HostCodeHeap : CodeHeap
 
     LoaderAllocator* GetAllocator() { return m_pAllocator; }
 
+    virtual PTR_BYTE GetAllocPtr() { return 0; }
+
 #ifdef DACCESS_COMPILE
     virtual void EnumMemoryRegions(CLRDataEnumMemoryFlags flags);
 #endif

src/coreclr/vm/jitinterface.cpp

@@ -11333,6 +11333,27 @@ WORD CEEJitInfo::getRelocTypeHint(void * target)
     }
 #endif // TARGET_AMD64
 
+#ifdef TARGET_ARM64
+    // Unlike x86/x86_64 we can't rely on IsPreferredExecutableRange here, instead we'll try to
+    // guess whether the native instruction (e.g. br) will be able to reach this target or not
+    // within 128Mb (1Mb is subtracted because we don't know the exact code size at this point)
+    // NOTE: it's still just a hint
+    if (m_pMethodBeingCompiled)
+    {
+        PTR_LoaderAllocator alloc = m_pMethodBeingCompiled->GetLoaderAllocator();
+        auto heapList = (HeapList*)alloc->m_pLastUsedCodeHeap;
+        if (heapList != nullptr)
+        {
+            BYTE* codeHeap = heapList->pHeap->GetAllocPtr();
+            INT64 delta = (INT64)((BYTE*)target - codeHeap);
+            if (FitsInRel28(delta))
+            {
+                return IMAGE_REL_ARM64_BRANCH26;
+            }
+        }
+    }
+#endif
+
     // No hints
     return (WORD)-1;
 }