[RFC] [AArch64] Add a target feature flag for restoring X18 after foreign function calls

[mstorsjo]

Summary

This is a different, more comprehensive solution to the same issue as D61892 (https://bugs.winehq.org/show_bug.cgi?id=38780).

This is an incomplete PoC patch; in -O0 mode, the function calls are lowered with GlobalISel which doesn't use the modified functions in AArch64ISelLowering.cpp. When GlobalISel lowers a call in IRTranslator::translateCall in lib/CodeGen/GlobalISel/IRTranslator.cpp, it doesn't have target specific knowledge where it could check for this AArch64 specific feature flag.

Currently this is implemented as target feature flag (used via clang as -Xclang -target-feature -Xclang +protect-x18). This is not ideal though, because it's hard to check for this feature flag in e.g. a configure script; on clang versions that don't have the flag, test compiling with this flag does work but emits a warning (which -Werror doesn't make into an error either).

lib/Target/AArch64/AArch64.td

def FeatureStrictAlign : SubtargetFeature<"strict-align",
                                           "Disallow all unaligned memory "
                                           "access">;
 
 foreach i = {1-7,9-15,18,20-28} in
     def FeatureReserveX#i : SubtargetFeature<"reserve-x"#i, "ReserveXRegister["#i#"]", "true",
                                              "Reserve X"#i#", making it unavailable "
                                              "as a GPR">;
 
+def FeatureProtectX18 : SubtargetFeature<"protect-x18", "ProtectX18", "true",
+                                    "Guard against called code changing x18">;
+
 foreach i = {8-15,18} in
     def FeatureCallSavedX#i : SubtargetFeature<"call-saved-x"#i,
          "CustomCallSavedXRegs["#i#"]", "true", "Make X"#i#" callee saved.">;
 
 def FeatureUseAA : SubtargetFeature<"use-aa", "UseAA", "true",
                                     "Use alias analysis during codegen">;
 
 def FeatureBalanceFPOps : SubtargetFeature<"balance-fp-ops", "BalanceFPOps",

lib/Target/AArch64/AArch64ISelLowering.cpp

 
         SDValue Copy = DAG.getCopyToReg(DAG.getEntryNode(), DL, Reg, InVals[I]);
         Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, Copy, Chain);
         break;
       }
     }
   }
 
+  if (Subtarget->shouldProtectX18()) {
+    // If protecting X18, make a copy of the value of X18 on entry.
+    MVT PtrTy = getPointerTy(DAG.getDataLayout());
+    unsigned Reg = MF.getRegInfo().createVirtualRegister(getRegClassFor(PtrTy));
+    FuncInfo->setX18BackupReg(Reg);
+
+    SDValue X18Reg = DAG.getRegister(AArch64::X18, MVT::i64);
+    SDValue Copy = DAG.getCopyToReg(DAG.getEntryNode(), DL, Reg, X18Reg);
+    Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, Copy, Chain);
+  }
+
   unsigned StackArgSize = CCInfo.getNextStackOffset();
   bool TailCallOpt = MF.getTarget().Options.GuaranteedTailCallOpt;
   if (DoesCalleeRestoreStack(CallConv, TailCallOpt)) {
     // This is a non-standard ABI so by fiat I say we're allowed to make full
     // use of the stack area to be popped, which must be aligned to 16 bytes in
     // any case:
     StackArgSize = alignTo(StackArgSize, 16);
 

   if (!mayTailCallThisCC(CalleeCC))
     return false;
 
   MachineFunction &MF = DAG.getMachineFunction();
   const Function &CallerF = MF.getFunction();
   CallingConv::ID CallerCC = CallerF.getCallingConv();
   bool CCMatch = CallerCC == CalleeCC;
 
+  if (Subtarget->shouldProtectX18()) {
+    // If protecting X18, we can't do a tail call to an unknown function
+    // (functions from the same translation unit are deemed safe), as we'd
+    // need to restore X18 after the function.
+    bool SameTU = false;
+    if (auto *G = dyn_cast<GlobalAddressSDNode>(Callee))
+      SameTU = !G->getGlobal()->isDeclarationForLinker();
+    if (!SameTU)
+      return false;
+  }
+
   // Byval parameters hand the function a pointer directly into the stack area
   // we want to reuse during a tail call. Working around this *is* possible (see
   // X86) but less efficient and uglier in LowerCall.
   for (Function::const_arg_iterator i = CallerF.arg_begin(),
                                     e = CallerF.arg_end();
        i != e; ++i) {
     if (i->hasByValAttr())
       return false;

     Chain = DAG.getCopyToReg(Chain, DL, RegToPass.first,
                              RegToPass.second, InFlag);
     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.
+  bool SameTU = false;
   if (auto *G = dyn_cast<GlobalAddressSDNode>(Callee)) {
     auto GV = G->getGlobal();
+    SameTU = !GV->isDeclarationForLinker();
     if (Subtarget->classifyGlobalFunctionReference(GV, getTargetMachine()) ==
         AArch64II::MO_GOT) {
       Callee = DAG.getTargetGlobalAddress(GV, DL, PtrVT, 0, AArch64II::MO_GOT);
       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);

   Chain = DAG.getCALLSEQ_END(Chain, DAG.getIntPtrConstant(NumBytes, DL, true),
                              DAG.getIntPtrConstant(CalleePopBytes, DL, true),
                              InFlag, DL);
   if (!Ins.empty())
     InFlag = Chain.getValue(1);
 
   // Handle result values, copying them out of physregs into vregs that we
   // return.
-  return LowerCallResult(Chain, InFlag, CallConv, IsVarArg, Ins, DL, DAG,
-                         InVals, IsThisReturn,
-                         IsThisReturn ? OutVals[0] : SDValue());
+  Chain =
+      LowerCallResult(Chain, InFlag, CallConv, IsVarArg, Ins, DL, DAG, InVals,
+                      IsThisReturn, IsThisReturn ? OutVals[0] : SDValue());
+
+  // If protecting X18, restore X18 after a call to an unknown function
+  // (functions from the same translation unit are deemed safe).
+  if (Subtarget->shouldProtectX18() && !SameTU)
+    Chain = DAG.getCopyToReg(
+        Chain, DL, AArch64::X18,
+        DAG.getCopyFromReg(Chain, DL, FuncInfo->getX18BackupReg(), PtrVT));
+
+  return Chain;
 }
 
 bool AArch64TargetLowering::CanLowerReturn(
     CallingConv::ID CallConv, MachineFunction &MF, bool isVarArg,
     const SmallVectorImpl<ISD::OutputArg> &Outs, LLVMContext &Context) const {
   CCAssignFn *RetCC = CallConv == CallingConv::WebKit_JS
                           ? RetCC_AArch64_WebKit_JS
                           : RetCC_AArch64_AAPCS;

lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// other stack allocations.
   bool CalleeSaveStackHasFreeSpace = false;
 
   /// SRetReturnReg - sret lowering includes returning the value of the
   /// returned struct in a register. This field holds the virtual register into
   /// which the sret argument is passed.
   unsigned SRetReturnReg = 0;
 
+  unsigned X18BackupReg = 0;
+
   /// Has a value when it is known whether or not the function uses a
   /// redzone, and no value otherwise.
   /// Initialized during frame lowering, unless the function has the noredzone
   /// attribute, in which case it is set to false at construction.
   Optional<bool> HasRedZone;
 
   /// ForwardedMustTailRegParms - A list of virtual and physical registers
   /// that must be forwarded to every musttail call.
public:
   void setVarArgsFPRIndex(int Index) { VarArgsFPRIndex = Index; }
 
   unsigned getVarArgsFPRSize() const { return VarArgsFPRSize; }
   void setVarArgsFPRSize(unsigned Size) { VarArgsFPRSize = Size; }
 
   unsigned getSRetReturnReg() const { return SRetReturnReg; }
   void setSRetReturnReg(unsigned Reg) { SRetReturnReg = Reg; }
 
+  unsigned getX18BackupReg() const { return X18BackupReg; }
+  void setX18BackupReg(unsigned Reg) { X18BackupReg = Reg; }
+
   unsigned getJumpTableEntrySize(int Idx) const {
     auto It = JumpTableEntryInfo.find(Idx);
     if (It != JumpTableEntryInfo.end())
       return It->second.first;
     return 4;
   }
   MCSymbol *getJumpTableEntryPCRelSymbol(int Idx) const {
     return JumpTableEntryInfo.find(Idx)->second.second;

lib/Target/AArch64/AArch64Subtarget.h

protected:
   unsigned PrefFunctionAlignment = 0;
   unsigned PrefLoopAlignment = 0;
   unsigned MaxJumpTableSize = 0;
   unsigned WideningBaseCost = 0;
 
   // ReserveXRegister[i] - X#i is not available as a general purpose register.
   BitVector ReserveXRegister;
 
+  // ProtectX18 - try to guard against foreign code changing X18.
+  bool ProtectX18 = false;
+
   // CustomCallUsedXRegister[i] - X#i call saved.
   BitVector CustomCallSavedXRegs;
 
   bool IsLittle;
 
   /// TargetTriple - What processor and OS we're targeting.
   Triple TargetTriple;
 
public:
   bool isXRaySupported() const override { return true; }
 
   unsigned getMinVectorRegisterBitWidth() const {
     return MinVectorRegisterBitWidth;
   }
 
   bool isXRegisterReserved(size_t i) const { return ReserveXRegister[i]; }
   unsigned getNumXRegisterReserved() const { return ReserveXRegister.count(); }
+  bool shouldProtectX18() const { return ProtectX18; }
   bool isXRegCustomCalleeSaved(size_t i) const {
     return CustomCallSavedXRegs[i];
   }
   bool hasCustomCallingConv() const { return CustomCallSavedXRegs.any(); }
   bool hasFPARMv8() const { return HasFPARMv8; }
   bool hasNEON() const { return HasNEON; }
   bool hasCrypto() const { return HasCrypto; }
   bool hasDotProd() const { return HasDotProd; }

test/CodeGen/AArch64/protect-x18.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=aarch64-linux-gnu | FileCheck %s
+
+define dso_local void @localfunc() #0 {
+; CHECK-LABEL: localfunc:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    //APP
+; CHECK-NEXT:    //NO_APP
+; CHECK-NEXT:    ret
+entry:
+  tail call void asm sideeffect "", ""() #1
+  ret void
+}
+
+define dso_local i32 @func(i32 ()* nocapture %ptr) #0 {
+; CHECK-LABEL: func:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str x20, [sp, #-32]! // 8-byte Folded Spill
+; CHECK-NEXT:    stp x19, x30, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    mov x19, x0
+; CHECK-NEXT:    mov x20, x18
+; CHECK-NEXT:    bl localfunc
+; CHECK-NEXT:    blr x19
+; CHECK-NEXT:    mov w19, w0
+; CHECK-NEXT:    mov x18, x20
+; CHECK-NEXT:    bl other
+; CHECK-NEXT:    mov x18, x20
+; CHECK-NEXT:    bl localfunc
+; CHECK-NEXT:    mov w0, w19
+; CHECK-NEXT:    ldp x19, x30, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr x20, [sp], #32 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  tail call void @localfunc()
+  %call = tail call i32 %ptr()
+  tail call void @other()
+  tail call void @localfunc()
+  ret i32 %call
+}
+
+declare dso_local void @other()
+
+define dso_local void @tailcall_other() #0 {
+; CHECK-LABEL: tailcall_other:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    stp x19, x30, [sp, #-16]! // 16-byte Folded Spill
+; CHECK-NEXT:    mov x19, x18
+; CHECK-NEXT:    bl other
+; CHECK-NEXT:    mov x18, x19
+; CHECK-NEXT:    ldp x19, x30, [sp], #16 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  tail call void @other()
+  ret void
+}
+
+define dso_local void @tailcall_local() #0 {
+; CHECK-LABEL: tailcall_local:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    b localfunc
+entry:
+  tail call void @localfunc()
+  ret void
+}
+
+attributes #0 = { nounwind "target-features"="+reserve-x18,+protect-x18" }
+
+attributes #1 = { nounwind }