[AArch64] [Windows] Use COFF stubs for calls to extern_weak functions

As the extern_weak target might be missing, resolving to the absolute
address zero, we can't use the normal direct PC-relative branch
instructions (as that would result in relocations out of range).

Improve the classifyGlobalFunctionReference method to set
MO_DLLIMPORT/MO_COFFSTUB, and simplify the existing code in
AArch64TargetLowering::LowerCall to use the return value from
classifyGlobalFunctionReference for these cases.

Add code in both AArch64FastISel and GlobalISel/IRTranslator to
bail out for function calls to extern weak functions on windows,
to let SelectionDAG handle them.

This matches what was done for X86 in 6bf108d77a3c.

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

7 files changed, 60 insertions, 25 deletions

diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 223ed954854..d71382868e4 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -1593,7 +1593,9 @@ bool IRTranslator::translateCall(const User &U, MachineIRBuilder &MIRBuilder) {
   const Function *F = CI.getCalledFunction();
 
   // FIXME: support Windows dllimport function calls.
-  if (F && F->hasDLLImportStorageClass())
+  if (F && (F->hasDLLImportStorageClass() ||
+            (MF->getTarget().getTargetTriple().isOSWindows() &&
+             F->hasExternalWeakLinkage())))
     return false;
 
   // FIXME: support control flow guard targets.
diff --git a/llvm/lib/Target/AArch64/AArch64FastISel.cpp b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
index 7a8243ad7c9..7e9c68f2bb3 100644
--- a/llvm/lib/Target/AArch64/AArch64FastISel.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
@@ -3253,6 +3253,13 @@ bool AArch64FastISel::fastLowerCall(CallLoweringInfo &CLI) {
   if (Callee && !computeCallAddress(Callee, Addr))
     return false;
 
+  // The weak function target may be zero; in that case we must use indirect
+  // addressing via a stub on windows as it may be out of range for a
+  // PC-relative jump.
+  if (Subtarget->isTargetWindows() && Addr.getGlobalValue() &&
+      Addr.getGlobalValue()->hasExternalWeakLinkage())
+    return false;
+
   // Handle the arguments now that we've gotten them.
   unsigned NumBytes;
   if (!processCallArgs(CLI, OutVTs, NumBytes))
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index d01979d8be8..9ae84f51407 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -4153,14 +4153,11 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   // node so that legalize doesn't hack it.
   if (auto *G = dyn_cast<GlobalAddressSDNode>(Callee)) {
     auto GV = G->getGlobal();
-    if (Subtarget->classifyGlobalFunctionReference(GV, getTargetMachine()) ==
-        AArch64II::MO_GOT) {
-      Callee = DAG.getTargetGlobalAddress(GV, DL, PtrVT, 0, AArch64II::MO_GOT);
+    unsigned OpFlags =
+        Subtarget->classifyGlobalFunctionReference(GV, getTargetMachine());
+    if (OpFlags & AArch64II::MO_GOT) {
+      Callee = DAG.getTargetGlobalAddress(GV, DL, PtrVT, 0, OpFlags);
       Callee = DAG.getNode(AArch64ISD::LOADgot, DL, PtrVT, Callee);
-    } else if (Subtarget->isTargetCOFF() && GV->hasDLLImportStorageClass()) {
-      assert(Subtarget->isTargetWindows() &&
-             "Windows is the only supported COFF target");
-      Callee = getGOT(G, DAG, AArch64II::MO_DLLIMPORT);
     } else {
       const GlobalValue *GV = G->getGlobal();
       Callee = DAG.getTargetGlobalAddress(GV, DL, PtrVT, 0, 0);
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
index 8e131b0caa6..931f0c0ffb8 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
@@ -251,6 +251,10 @@ unsigned AArch64Subtarget::classifyGlobalFunctionReference(
       !TM.shouldAssumeDSOLocal(*GV->getParent(), GV))
     return AArch64II::MO_GOT;
 
+  // Use ClassifyGlobalReference for setting MO_DLLIMPORT/MO_COFFSTUB.
+  if (getTargetTriple().isOSWindows())
+    return ClassifyGlobalReference(GV, TM);
+
   return AArch64II::MO_NO_FLAG;
 }
 
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call-weak.ll b/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call-weak.ll
new file mode 100644
index 00000000000..bea5961efff
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call-weak.ll
@@ -0,0 +1,15 @@
+; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+; RUN: llc %s -stop-after=irtranslator -verify-machineinstrs -mtriple aarch64-apple-darwin -global-isel -o - 2>&1 | FileCheck %s --check-prefixes=DARWIN,COMMON
+
+; Shouldn't tail call when the OS doesn't support it.
+declare extern_weak void @extern_weak_fn()
+define void @test_extern_weak() {
+  ; DARWIN-LABEL: name: test_extern_weak
+  ; DARWIN: bb.1 (%ir-block.0):
+  ; DARWIN:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
+  ; DARWIN:   BL @extern_weak_fn, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
+  ; DARWIN:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
+  ; DARWIN:   RET_ReallyLR
+  tail call void @extern_weak_fn()
+  ret void
+}
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll b/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll
index b49307ca216..174c6ecb899 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll
@@ -193,23 +193,6 @@ define void @test_inreg(i8* inreg %ptr) {
   ret void
 }
 
-; Shouldn't tail call when the OS doesn't support it. Windows supports this,
-; so we should be able to tail call there.
-declare extern_weak void @extern_weak_fn()
-define void @test_extern_weak() {
-  ; DARWIN-LABEL: name: test_extern_weak
-  ; DARWIN: bb.1 (%ir-block.0):
-  ; DARWIN:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
-  ; DARWIN:   BL @extern_weak_fn, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
-  ; DARWIN:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
-  ; DARWIN:   RET_ReallyLR
-  ; WINDOWS-LABEL: name: test_extern_weak
-  ; WINDOWS: bb.1 (%ir-block.0):
-  ; WINDOWS:   TCRETURNdi @extern_weak_fn, 0, csr_aarch64_aapcs, implicit $sp
-  tail call void @extern_weak_fn()
-  ret void
-}
-
 declare fastcc void @fast_fn()
 define void @test_mismatched_caller() {
   ; COMMON-LABEL: name: test_mismatched_caller
diff --git a/llvm/test/CodeGen/AArch64/windows-extern-weak.ll b/llvm/test/CodeGen/AArch64/windows-extern-weak.ll
new file mode 100644
index 00000000000..e1e6de49604
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/windows-extern-weak.ll
@@ -0,0 +1,27 @@
+; RUN: llc -mtriple aarch64-windows -filetype asm -o - < %s | FileCheck %s
+; RUN: llc -mtriple aarch64-windows -filetype asm -o - -fast-isel %s | FileCheck %s
+; RUN: llc -mtriple aarch64-windows -filetype asm -o - -global-isel -global-isel-abort=0 %s | FileCheck %s
+
+define void @func() {
+; CHECK-LABEL: func:
+; CHECK:      str x30, [sp, #-16]!
+; CHECK-NEXT: adrp x8, .refptr.weakfunc
+; CHECK-NEXT: ldr x8, [x8, .refptr.weakfunc]
+; CHECK-NEXT: cbz     x8, .LBB0_2
+; CHECK-NEXT: ; %bb.1:
+; CHECK-NEXT: blr     x8
+; CHECK-NEXT: .LBB0_2:
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: ret
+
+  br i1 icmp ne (void ()* @weakfunc, void ()* null), label %1, label %2
+
+1:
+  call void @weakfunc()
+  br label %2
+
+2:
+  ret void
+}
+
+declare extern_weak void @weakfunc()