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

[mstorsjo]

Summary

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.

[paquette]

LGTM

[rnk]

lgtm

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

 }
 
 bool IRTranslator::translateCall(const User &U, MachineIRBuilder &MIRBuilder) {
   const CallInst &CI = cast<CallInst>(U);
   auto TII = MF->getTarget().getIntrinsicInfo();
   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.
   if (CI.countOperandBundlesOfType(LLVMContext::OB_cfguardtarget))
     return false;
 
   if (CI.isInlineAsm())
     return translateInlineAsm(CI, MIRBuilder);

llvm/lib/Target/AArch64/AArch64FastISel.cpp

for (auto *Val : CLI.OutVals) {
 
     OutVTs.push_back(VT);
   }
 
   Address Addr;
   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))
     return false;
 
   const AArch64RegisterInfo *RegInfo = Subtarget->getRegisterInfo();
   if (RegInfo->isAnyArgRegReserved(*MF))
     RegInfo->emitReservedArgRegCallError(*MF);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (auto &RegToPass : RegsToPass) {
     InFlag = Chain.getValue(1);
   }
 
   // If the callee is a GlobalAddress/ExternalSymbol node (quite common, every
   // direct call is) turn it into a TargetGlobalAddress/TargetExternalSymbol
   // 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);
     }
   } else if (auto *S = dyn_cast<ExternalSymbolSDNode>(Callee)) {
     if (getTargetMachine().getCodeModel() == CodeModel::Large &&
         Subtarget->isTargetMachO()) {
       const char *Sym = S->getSymbol();

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

if (TM.getCodeModel() == CodeModel::Large && isTargetMachO() &&
     return AArch64II::MO_GOT;
 
   // NonLazyBind goes via GOT unless we know it's available locally.
   auto *F = dyn_cast<Function>(GV);
   if (UseNonLazyBind && F && F->hasFnAttribute(Attribute::NonLazyBind) &&
       !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;
 }
 
 void AArch64Subtarget::overrideSchedPolicy(MachineSchedPolicy &Policy,
                                            unsigned NumRegionInstrs) const {
   // LNT run (at least on Cyclone) showed reasonably significant gains for
   // bi-directional scheduling. 253.perlbmk.
   Policy.OnlyTopDown = false;

llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call-weak.ll

+; 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
+}

llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll

define void @test_inreg(i8* inreg %ptr) {
   ; COMMON:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
   ; COMMON:   BL @simple_fn, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
   ; COMMON:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
   ; COMMON:   RET_ReallyLR
   tail call void @simple_fn()
   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
   ; COMMON: bb.1 (%ir-block.0):
   ; COMMON:   TCRETURNdi @fast_fn, 0, csr_aarch64_aapcs, implicit $sp
   tail call fastcc void @fast_fn()
   ret void
 }

llvm/test/CodeGen/AArch64/windows-extern-weak.ll

+; 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()