brintos

brintos / llvm-project-archived public Read only

0
0
Text · 38.4 KiB · d0c4b1b Raw
1029 lines · cpp
1//===-- AArch64Arm64ECCallLowering.cpp - Lower Arm64EC calls ----*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8///9/// \file10/// This file contains the IR transform to lower external or indirect calls for11/// the ARM64EC calling convention. Such calls must go through the runtime, so12/// we can translate the calling convention for calls into the emulator.13///14/// This subsumes Control Flow Guard handling.15///16//===----------------------------------------------------------------------===//17 18#include "AArch64.h"19#include "llvm/ADT/SetVector.h"20#include "llvm/ADT/SmallString.h"21#include "llvm/ADT/SmallVector.h"22#include "llvm/ADT/Statistic.h"23#include "llvm/IR/CallingConv.h"24#include "llvm/IR/GlobalAlias.h"25#include "llvm/IR/IRBuilder.h"26#include "llvm/IR/Instruction.h"27#include "llvm/IR/Mangler.h"28#include "llvm/IR/Module.h"29#include "llvm/Object/COFF.h"30#include "llvm/Pass.h"31#include "llvm/Support/CommandLine.h"32#include "llvm/TargetParser/Triple.h"33 34using namespace llvm;35using namespace llvm::COFF;36 37using OperandBundleDef = OperandBundleDefT<Value *>;38 39#define DEBUG_TYPE "arm64eccalllowering"40 41STATISTIC(Arm64ECCallsLowered, "Number of Arm64EC calls lowered");42 43static cl::opt<bool> LowerDirectToIndirect("arm64ec-lower-direct-to-indirect",44                                           cl::Hidden, cl::init(true));45static cl::opt<bool> GenerateThunks("arm64ec-generate-thunks", cl::Hidden,46                                    cl::init(true));47 48namespace {49 50enum ThunkArgTranslation : uint8_t {51  Direct,52  Bitcast,53  PointerIndirection,54};55 56struct ThunkArgInfo {57  Type *Arm64Ty;58  Type *X64Ty;59  ThunkArgTranslation Translation;60};61 62class AArch64Arm64ECCallLowering : public ModulePass {63public:64  static char ID;65  AArch64Arm64ECCallLowering() : ModulePass(ID) {}66 67  Function *buildExitThunk(FunctionType *FnTy, AttributeList Attrs);68  Function *buildEntryThunk(Function *F);69  void lowerCall(CallBase *CB);70  Function *buildGuestExitThunk(Function *F);71  Function *buildPatchableThunk(GlobalAlias *UnmangledAlias,72                                GlobalAlias *MangledAlias);73  bool processFunction(Function &F, SetVector<GlobalValue *> &DirectCalledFns,74                       DenseMap<GlobalAlias *, GlobalAlias *> &FnsMap);75  bool runOnModule(Module &M) override;76 77private:78  int cfguard_module_flag = 0;79  FunctionType *GuardFnType = nullptr;80  FunctionType *DispatchFnType = nullptr;81  Constant *GuardFnCFGlobal = nullptr;82  Constant *GuardFnGlobal = nullptr;83  Constant *DispatchFnGlobal = nullptr;84  Module *M = nullptr;85 86  Type *PtrTy;87  Type *I64Ty;88  Type *VoidTy;89 90  void getThunkType(FunctionType *FT, AttributeList AttrList,91                    Arm64ECThunkType TT, raw_ostream &Out,92                    FunctionType *&Arm64Ty, FunctionType *&X64Ty,93                    SmallVector<ThunkArgTranslation> &ArgTranslations);94  void getThunkRetType(FunctionType *FT, AttributeList AttrList,95                       raw_ostream &Out, Type *&Arm64RetTy, Type *&X64RetTy,96                       SmallVectorImpl<Type *> &Arm64ArgTypes,97                       SmallVectorImpl<Type *> &X64ArgTypes,98                       SmallVector<ThunkArgTranslation> &ArgTranslations,99                       bool &HasSretPtr);100  void getThunkArgTypes(FunctionType *FT, AttributeList AttrList,101                        Arm64ECThunkType TT, raw_ostream &Out,102                        SmallVectorImpl<Type *> &Arm64ArgTypes,103                        SmallVectorImpl<Type *> &X64ArgTypes,104                        SmallVectorImpl<ThunkArgTranslation> &ArgTranslations,105                        bool HasSretPtr);106  ThunkArgInfo canonicalizeThunkType(Type *T, Align Alignment, bool Ret,107                                     uint64_t ArgSizeBytes, raw_ostream &Out);108};109 110} // end anonymous namespace111 112void AArch64Arm64ECCallLowering::getThunkType(113    FunctionType *FT, AttributeList AttrList, Arm64ECThunkType TT,114    raw_ostream &Out, FunctionType *&Arm64Ty, FunctionType *&X64Ty,115    SmallVector<ThunkArgTranslation> &ArgTranslations) {116  Out << (TT == Arm64ECThunkType::Entry ? "$ientry_thunk$cdecl$"117                                        : "$iexit_thunk$cdecl$");118 119  Type *Arm64RetTy;120  Type *X64RetTy;121 122  SmallVector<Type *> Arm64ArgTypes;123  SmallVector<Type *> X64ArgTypes;124 125  // The first argument to a thunk is the called function, stored in x9.126  // For exit thunks, we pass the called function down to the emulator;127  // for entry/guest exit thunks, we just call the Arm64 function directly.128  if (TT == Arm64ECThunkType::Exit)129    Arm64ArgTypes.push_back(PtrTy);130  X64ArgTypes.push_back(PtrTy);131 132  bool HasSretPtr = false;133  getThunkRetType(FT, AttrList, Out, Arm64RetTy, X64RetTy, Arm64ArgTypes,134                  X64ArgTypes, ArgTranslations, HasSretPtr);135 136  getThunkArgTypes(FT, AttrList, TT, Out, Arm64ArgTypes, X64ArgTypes,137                   ArgTranslations, HasSretPtr);138 139  Arm64Ty = FunctionType::get(Arm64RetTy, Arm64ArgTypes, false);140 141  X64Ty = FunctionType::get(X64RetTy, X64ArgTypes, false);142}143 144void AArch64Arm64ECCallLowering::getThunkArgTypes(145    FunctionType *FT, AttributeList AttrList, Arm64ECThunkType TT,146    raw_ostream &Out, SmallVectorImpl<Type *> &Arm64ArgTypes,147    SmallVectorImpl<Type *> &X64ArgTypes,148    SmallVectorImpl<ThunkArgTranslation> &ArgTranslations, bool HasSretPtr) {149 150  Out << "$";151  if (FT->isVarArg()) {152    // We treat the variadic function's thunk as a normal function153    // with the following type on the ARM side:154    //   rettype exitthunk(155    //     ptr x9, ptr x0, i64 x1, i64 x2, i64 x3, ptr x4, i64 x5)156    //157    // that can coverage all types of variadic function.158    // x9 is similar to normal exit thunk, store the called function.159    // x0-x3 is the arguments be stored in registers.160    // x4 is the address of the arguments on the stack.161    // x5 is the size of the arguments on the stack.162    //163    // On the x64 side, it's the same except that x5 isn't set.164    //165    // If both the ARM and X64 sides are sret, there are only three166    // arguments in registers.167    //168    // If the X64 side is sret, but the ARM side isn't, we pass an extra value169    // to/from the X64 side, and let SelectionDAG transform it into a memory170    // location.171    Out << "varargs";172 173    // x0-x3174    for (int i = HasSretPtr ? 1 : 0; i < 4; i++) {175      Arm64ArgTypes.push_back(I64Ty);176      X64ArgTypes.push_back(I64Ty);177      ArgTranslations.push_back(ThunkArgTranslation::Direct);178    }179 180    // x4181    Arm64ArgTypes.push_back(PtrTy);182    X64ArgTypes.push_back(PtrTy);183    ArgTranslations.push_back(ThunkArgTranslation::Direct);184    // x5185    Arm64ArgTypes.push_back(I64Ty);186    if (TT != Arm64ECThunkType::Entry) {187      // FIXME: x5 isn't actually used by the x64 side; revisit once we188      // have proper isel for varargs189      X64ArgTypes.push_back(I64Ty);190      ArgTranslations.push_back(ThunkArgTranslation::Direct);191    }192    return;193  }194 195  unsigned I = 0;196  if (HasSretPtr)197    I++;198 199  if (I == FT->getNumParams()) {200    Out << "v";201    return;202  }203 204  for (unsigned E = FT->getNumParams(); I != E; ++I) {205#if 0206    // FIXME: Need more information about argument size; see207    // https://reviews.llvm.org/D132926208    uint64_t ArgSizeBytes = AttrList.getParamArm64ECArgSizeBytes(I);209    Align ParamAlign = AttrList.getParamAlignment(I).valueOrOne();210#else211    uint64_t ArgSizeBytes = 0;212    Align ParamAlign = Align();213#endif214    auto [Arm64Ty, X64Ty, ArgTranslation] =215        canonicalizeThunkType(FT->getParamType(I), ParamAlign,216                              /*Ret*/ false, ArgSizeBytes, Out);217    Arm64ArgTypes.push_back(Arm64Ty);218    X64ArgTypes.push_back(X64Ty);219    ArgTranslations.push_back(ArgTranslation);220  }221}222 223void AArch64Arm64ECCallLowering::getThunkRetType(224    FunctionType *FT, AttributeList AttrList, raw_ostream &Out,225    Type *&Arm64RetTy, Type *&X64RetTy, SmallVectorImpl<Type *> &Arm64ArgTypes,226    SmallVectorImpl<Type *> &X64ArgTypes,227    SmallVector<ThunkArgTranslation> &ArgTranslations, bool &HasSretPtr) {228  Type *T = FT->getReturnType();229#if 0230  // FIXME: Need more information about argument size; see231  // https://reviews.llvm.org/D132926232  uint64_t ArgSizeBytes = AttrList.getRetArm64ECArgSizeBytes();233#else234  int64_t ArgSizeBytes = 0;235#endif236  if (T->isVoidTy()) {237    if (FT->getNumParams()) {238      Attribute SRetAttr0 = AttrList.getParamAttr(0, Attribute::StructRet);239      Attribute InRegAttr0 = AttrList.getParamAttr(0, Attribute::InReg);240      Attribute SRetAttr1, InRegAttr1;241      if (FT->getNumParams() > 1) {242        // Also check the second parameter (for class methods, the first243        // parameter is "this", and the second parameter is the sret pointer.)244        // It doesn't matter which one is sret.245        SRetAttr1 = AttrList.getParamAttr(1, Attribute::StructRet);246        InRegAttr1 = AttrList.getParamAttr(1, Attribute::InReg);247      }248      if ((SRetAttr0.isValid() && InRegAttr0.isValid()) ||249          (SRetAttr1.isValid() && InRegAttr1.isValid())) {250        // sret+inreg indicates a call that returns a C++ class value. This is251        // actually equivalent to just passing and returning a void* pointer252        // as the first or second argument. Translate it that way, instead of253        // trying to model "inreg" in the thunk's calling convention; this254        // simplfies the rest of the code, and matches MSVC mangling.255        Out << "i8";256        Arm64RetTy = I64Ty;257        X64RetTy = I64Ty;258        return;259      }260      if (SRetAttr0.isValid()) {261        // FIXME: Sanity-check the sret type; if it's an integer or pointer,262        // we'll get screwy mangling/codegen.263        // FIXME: For large struct types, mangle as an integer argument and264        // integer return, so we can reuse more thunks, instead of "m" syntax.265        // (MSVC mangles this case as an integer return with no argument, but266        // that's a miscompile.)267        Type *SRetType = SRetAttr0.getValueAsType();268        Align SRetAlign = AttrList.getParamAlignment(0).valueOrOne();269        canonicalizeThunkType(SRetType, SRetAlign, /*Ret*/ true, ArgSizeBytes,270                              Out);271        Arm64RetTy = VoidTy;272        X64RetTy = VoidTy;273        Arm64ArgTypes.push_back(FT->getParamType(0));274        X64ArgTypes.push_back(FT->getParamType(0));275        ArgTranslations.push_back(ThunkArgTranslation::Direct);276        HasSretPtr = true;277        return;278      }279    }280 281    Out << "v";282    Arm64RetTy = VoidTy;283    X64RetTy = VoidTy;284    return;285  }286 287  auto info =288      canonicalizeThunkType(T, Align(), /*Ret*/ true, ArgSizeBytes, Out);289  Arm64RetTy = info.Arm64Ty;290  X64RetTy = info.X64Ty;291  if (X64RetTy->isPointerTy()) {292    // If the X64 type is canonicalized to a pointer, that means it's293    // passed/returned indirectly. For a return value, that means it's an294    // sret pointer.295    X64ArgTypes.push_back(X64RetTy);296    X64RetTy = VoidTy;297  }298}299 300ThunkArgInfo AArch64Arm64ECCallLowering::canonicalizeThunkType(301    Type *T, Align Alignment, bool Ret, uint64_t ArgSizeBytes,302    raw_ostream &Out) {303 304  auto direct = [](Type *T) {305    return ThunkArgInfo{T, T, ThunkArgTranslation::Direct};306  };307 308  auto bitcast = [this](Type *Arm64Ty, uint64_t SizeInBytes) {309    return ThunkArgInfo{Arm64Ty,310                        llvm::Type::getIntNTy(M->getContext(), SizeInBytes * 8),311                        ThunkArgTranslation::Bitcast};312  };313 314  auto pointerIndirection = [this](Type *Arm64Ty) {315    return ThunkArgInfo{Arm64Ty, PtrTy,316                        ThunkArgTranslation::PointerIndirection};317  };318 319  if (T->isHalfTy()) {320    // Prefix with `llvm` since MSVC doesn't specify `_Float16`321    Out << "__llvm_h__";322    return direct(T);323  }324 325  if (T->isFloatTy()) {326    Out << "f";327    return direct(T);328  }329 330  if (T->isDoubleTy()) {331    Out << "d";332    return direct(T);333  }334 335  if (T->isFloatingPointTy()) {336    report_fatal_error("Only 16, 32, and 64 bit floating points are supported "337                       "for ARM64EC thunks");338  }339 340  auto &DL = M->getDataLayout();341 342  if (auto *StructTy = dyn_cast<StructType>(T))343    if (StructTy->getNumElements() == 1)344      T = StructTy->getElementType(0);345 346  if (T->isArrayTy()) {347    Type *ElementTy = T->getArrayElementType();348    uint64_t ElementCnt = T->getArrayNumElements();349    uint64_t ElementSizePerBytes = DL.getTypeSizeInBits(ElementTy) / 8;350    uint64_t TotalSizeBytes = ElementCnt * ElementSizePerBytes;351    if (ElementTy->isHalfTy() || ElementTy->isFloatTy() ||352        ElementTy->isDoubleTy()) {353      if (ElementTy->isHalfTy())354        // Prefix with `llvm` since MSVC doesn't specify `_Float16`355        Out << "__llvm_H__";356      else if (ElementTy->isFloatTy())357        Out << "F";358      else if (ElementTy->isDoubleTy())359        Out << "D";360      Out << TotalSizeBytes;361      if (Alignment.value() >= 16 && !Ret)362        Out << "a" << Alignment.value();363      if (TotalSizeBytes <= 8) {364        // Arm64 returns small structs of float/double in float registers;365        // X64 uses RAX.366        return bitcast(T, TotalSizeBytes);367      } else {368        // Struct is passed directly on Arm64, but indirectly on X64.369        return pointerIndirection(T);370      }371    } else if (T->isFloatingPointTy()) {372      report_fatal_error(373          "Only 16, 32, and 64 bit floating points are supported "374          "for ARM64EC thunks");375    }376  }377 378  if ((T->isIntegerTy() || T->isPointerTy()) && DL.getTypeSizeInBits(T) <= 64) {379    Out << "i8";380    return direct(I64Ty);381  }382 383  unsigned TypeSize = ArgSizeBytes;384  if (TypeSize == 0)385    TypeSize = DL.getTypeSizeInBits(T) / 8;386  Out << "m";387  if (TypeSize != 4)388    Out << TypeSize;389  if (Alignment.value() >= 16 && !Ret)390    Out << "a" << Alignment.value();391  // FIXME: Try to canonicalize Arm64Ty more thoroughly?392  if (TypeSize == 1 || TypeSize == 2 || TypeSize == 4 || TypeSize == 8) {393    // Pass directly in an integer register394    return bitcast(T, TypeSize);395  } else {396    // Passed directly on Arm64, but indirectly on X64.397    return pointerIndirection(T);398  }399}400 401// This function builds the "exit thunk", a function which translates402// arguments and return values when calling x64 code from AArch64 code.403Function *AArch64Arm64ECCallLowering::buildExitThunk(FunctionType *FT,404                                                     AttributeList Attrs) {405  SmallString<256> ExitThunkName;406  llvm::raw_svector_ostream ExitThunkStream(ExitThunkName);407  FunctionType *Arm64Ty, *X64Ty;408  SmallVector<ThunkArgTranslation> ArgTranslations;409  getThunkType(FT, Attrs, Arm64ECThunkType::Exit, ExitThunkStream, Arm64Ty,410               X64Ty, ArgTranslations);411  if (Function *F = M->getFunction(ExitThunkName))412    return F;413 414  Function *F = Function::Create(Arm64Ty, GlobalValue::LinkOnceODRLinkage, 0,415                                 ExitThunkName, M);416  F->setCallingConv(CallingConv::ARM64EC_Thunk_Native);417  F->setSection(".wowthk$aa");418  F->setComdat(M->getOrInsertComdat(ExitThunkName));419  // Copy MSVC, and always set up a frame pointer. (Maybe this isn't necessary.)420  F->addFnAttr("frame-pointer", "all");421  // Only copy sret from the first argument. For C++ instance methods, clang can422  // stick an sret marking on a later argument, but it doesn't actually affect423  // the ABI, so we can omit it. This avoids triggering a verifier assertion.424  if (FT->getNumParams()) {425    auto SRet = Attrs.getParamAttr(0, Attribute::StructRet);426    auto InReg = Attrs.getParamAttr(0, Attribute::InReg);427    if (SRet.isValid() && !InReg.isValid())428      F->addParamAttr(1, SRet);429  }430  // FIXME: Copy anything other than sret?  Shouldn't be necessary for normal431  // C ABI, but might show up in other cases.432  BasicBlock *BB = BasicBlock::Create(M->getContext(), "", F);433  IRBuilder<> IRB(BB);434  Value *CalleePtr =435      M->getOrInsertGlobal("__os_arm64x_dispatch_call_no_redirect", PtrTy);436  Value *Callee = IRB.CreateLoad(PtrTy, CalleePtr);437  auto &DL = M->getDataLayout();438  SmallVector<Value *> Args;439 440  // Pass the called function in x9.441  auto X64TyOffset = 1;442  Args.push_back(F->arg_begin());443 444  Type *RetTy = Arm64Ty->getReturnType();445  if (RetTy != X64Ty->getReturnType()) {446    // If the return type is an array or struct, translate it. Values of size447    // 8 or less go into RAX; bigger values go into memory, and we pass a448    // pointer.449    if (DL.getTypeStoreSize(RetTy) > 8) {450      Args.push_back(IRB.CreateAlloca(RetTy));451      X64TyOffset++;452    }453  }454 455  for (auto [Arg, X64ArgType, ArgTranslation] : llvm::zip_equal(456           make_range(F->arg_begin() + 1, F->arg_end()),457           make_range(X64Ty->param_begin() + X64TyOffset, X64Ty->param_end()),458           ArgTranslations)) {459    // Translate arguments from AArch64 calling convention to x86 calling460    // convention.461    //462    // For simple types, we don't need to do any translation: they're463    // represented the same way. (Implicit sign extension is not part of464    // either convention.)465    //466    // The big thing we have to worry about is struct types... but467    // fortunately AArch64 clang is pretty friendly here: the cases that need468    // translation are always passed as a struct or array. (If we run into469    // some cases where this doesn't work, we can teach clang to mark it up470    // with an attribute.)471    //472    // The first argument is the called function, stored in x9.473    if (ArgTranslation != ThunkArgTranslation::Direct) {474      Value *Mem = IRB.CreateAlloca(Arg.getType());475      IRB.CreateStore(&Arg, Mem);476      if (ArgTranslation == ThunkArgTranslation::Bitcast) {477        Type *IntTy = IRB.getIntNTy(DL.getTypeStoreSizeInBits(Arg.getType()));478        Args.push_back(IRB.CreateLoad(IntTy, Mem));479      } else {480        assert(ArgTranslation == ThunkArgTranslation::PointerIndirection);481        Args.push_back(Mem);482      }483    } else {484      Args.push_back(&Arg);485    }486    assert(Args.back()->getType() == X64ArgType);487  }488  // FIXME: Transfer necessary attributes? sret? anything else?489 490  CallInst *Call = IRB.CreateCall(X64Ty, Callee, Args);491  Call->setCallingConv(CallingConv::ARM64EC_Thunk_X64);492 493  Value *RetVal = Call;494  if (RetTy != X64Ty->getReturnType()) {495    // If we rewrote the return type earlier, convert the return value to496    // the proper type.497    if (DL.getTypeStoreSize(RetTy) > 8) {498      RetVal = IRB.CreateLoad(RetTy, Args[1]);499    } else {500      Value *CastAlloca = IRB.CreateAlloca(RetTy);501      IRB.CreateStore(Call, CastAlloca);502      RetVal = IRB.CreateLoad(RetTy, CastAlloca);503    }504  }505 506  if (RetTy->isVoidTy())507    IRB.CreateRetVoid();508  else509    IRB.CreateRet(RetVal);510  return F;511}512 513// This function builds the "entry thunk", a function which translates514// arguments and return values when calling AArch64 code from x64 code.515Function *AArch64Arm64ECCallLowering::buildEntryThunk(Function *F) {516  SmallString<256> EntryThunkName;517  llvm::raw_svector_ostream EntryThunkStream(EntryThunkName);518  FunctionType *Arm64Ty, *X64Ty;519  SmallVector<ThunkArgTranslation> ArgTranslations;520  getThunkType(F->getFunctionType(), F->getAttributes(),521               Arm64ECThunkType::Entry, EntryThunkStream, Arm64Ty, X64Ty,522               ArgTranslations);523  if (Function *F = M->getFunction(EntryThunkName))524    return F;525 526  Function *Thunk = Function::Create(X64Ty, GlobalValue::LinkOnceODRLinkage, 0,527                                     EntryThunkName, M);528  Thunk->setCallingConv(CallingConv::ARM64EC_Thunk_X64);529  Thunk->setSection(".wowthk$aa");530  Thunk->setComdat(M->getOrInsertComdat(EntryThunkName));531  // Copy MSVC, and always set up a frame pointer. (Maybe this isn't necessary.)532  Thunk->addFnAttr("frame-pointer", "all");533 534  BasicBlock *BB = BasicBlock::Create(M->getContext(), "", Thunk);535  IRBuilder<> IRB(BB);536 537  Type *RetTy = Arm64Ty->getReturnType();538  Type *X64RetType = X64Ty->getReturnType();539 540  bool TransformDirectToSRet = X64RetType->isVoidTy() && !RetTy->isVoidTy();541  unsigned ThunkArgOffset = TransformDirectToSRet ? 2 : 1;542  unsigned PassthroughArgSize =543      (F->isVarArg() ? 5 : Thunk->arg_size()) - ThunkArgOffset;544  assert(ArgTranslations.size() == (F->isVarArg() ? 5 : PassthroughArgSize));545 546  // Translate arguments to call.547  SmallVector<Value *> Args;548  for (unsigned i = 0; i != PassthroughArgSize; ++i) {549    Value *Arg = Thunk->getArg(i + ThunkArgOffset);550    Type *ArgTy = Arm64Ty->getParamType(i);551    ThunkArgTranslation ArgTranslation = ArgTranslations[i];552    if (ArgTranslation != ThunkArgTranslation::Direct) {553      // Translate array/struct arguments to the expected type.554      if (ArgTranslation == ThunkArgTranslation::Bitcast) {555        Value *CastAlloca = IRB.CreateAlloca(ArgTy);556        IRB.CreateStore(Arg, CastAlloca);557        Arg = IRB.CreateLoad(ArgTy, CastAlloca);558      } else {559        assert(ArgTranslation == ThunkArgTranslation::PointerIndirection);560        Arg = IRB.CreateLoad(ArgTy, Arg);561      }562    }563    assert(Arg->getType() == ArgTy);564    Args.push_back(Arg);565  }566 567  if (F->isVarArg()) {568    // The 5th argument to variadic entry thunks is used to model the x64 sp569    // which is passed to the thunk in x4, this can be passed to the callee as570    // the variadic argument start address after skipping over the 32 byte571    // shadow store.572 573    // The EC thunk CC will assign any argument marked as InReg to x4.574    Thunk->addParamAttr(5, Attribute::InReg);575    Value *Arg = Thunk->getArg(5);576    Arg = IRB.CreatePtrAdd(Arg, IRB.getInt64(0x20));577    Args.push_back(Arg);578 579    // Pass in a zero variadic argument size (in x5).580    Args.push_back(IRB.getInt64(0));581  }582 583  // Call the function passed to the thunk.584  Value *Callee = Thunk->getArg(0);585  CallInst *Call = IRB.CreateCall(Arm64Ty, Callee, Args);586 587  auto SRetAttr = F->getAttributes().getParamAttr(0, Attribute::StructRet);588  auto InRegAttr = F->getAttributes().getParamAttr(0, Attribute::InReg);589  if (SRetAttr.isValid() && !InRegAttr.isValid()) {590    Thunk->addParamAttr(1, SRetAttr);591    Call->addParamAttr(0, SRetAttr);592  }593 594  Value *RetVal = Call;595  if (TransformDirectToSRet) {596    IRB.CreateStore(RetVal, Thunk->getArg(1));597  } else if (X64RetType != RetTy) {598    Value *CastAlloca = IRB.CreateAlloca(X64RetType);599    IRB.CreateStore(Call, CastAlloca);600    RetVal = IRB.CreateLoad(X64RetType, CastAlloca);601  }602 603  // Return to the caller.  Note that the isel has code to translate this604  // "ret" to a tail call to __os_arm64x_dispatch_ret.  (Alternatively, we605  // could emit a tail call here, but that would require a dedicated calling606  // convention, which seems more complicated overall.)607  if (X64RetType->isVoidTy())608    IRB.CreateRetVoid();609  else610    IRB.CreateRet(RetVal);611 612  return Thunk;613}614 615std::optional<std::string> getArm64ECMangledFunctionName(GlobalValue &GV) {616  if (!GV.hasName()) {617    GV.setName("__unnamed");618  }619 620  return llvm::getArm64ECMangledFunctionName(GV.getName());621}622 623// Builds the "guest exit thunk", a helper to call a function which may or may624// not be an exit thunk. (We optimistically assume non-dllimport function625// declarations refer to functions defined in AArch64 code; if the linker626// can't prove that, we use this routine instead.)627Function *AArch64Arm64ECCallLowering::buildGuestExitThunk(Function *F) {628  llvm::raw_null_ostream NullThunkName;629  FunctionType *Arm64Ty, *X64Ty;630  SmallVector<ThunkArgTranslation> ArgTranslations;631  getThunkType(F->getFunctionType(), F->getAttributes(),632               Arm64ECThunkType::GuestExit, NullThunkName, Arm64Ty, X64Ty,633               ArgTranslations);634  auto MangledName = getArm64ECMangledFunctionName(*F);635  assert(MangledName && "Can't guest exit to function that's already native");636  std::string ThunkName = *MangledName;637  if (ThunkName[0] == '?' && ThunkName.find("@") != std::string::npos) {638    ThunkName.insert(ThunkName.find("@"), "$exit_thunk");639  } else {640    ThunkName.append("$exit_thunk");641  }642  Function *GuestExit =643      Function::Create(Arm64Ty, GlobalValue::WeakODRLinkage, 0, ThunkName, M);644  GuestExit->setComdat(M->getOrInsertComdat(ThunkName));645  GuestExit->setSection(".wowthk$aa");646  GuestExit->addMetadata(647      "arm64ec_unmangled_name",648      *MDNode::get(M->getContext(),649                   MDString::get(M->getContext(), F->getName())));650  GuestExit->setMetadata(651      "arm64ec_ecmangled_name",652      MDNode::get(M->getContext(),653                  MDString::get(M->getContext(), *MangledName)));654  F->setMetadata("arm64ec_hasguestexit", MDNode::get(M->getContext(), {}));655  BasicBlock *BB = BasicBlock::Create(M->getContext(), "", GuestExit);656  IRBuilder<> B(BB);657 658  // Create new call instruction. The call check should always be a call,659  // even if the original CallBase is an Invoke or CallBr instructio.660  // This is treated as a direct call, so do not use GuardFnCFGlobal.661  LoadInst *GuardCheckLoad = B.CreateLoad(PtrTy, GuardFnGlobal);662  Function *Thunk = buildExitThunk(F->getFunctionType(), F->getAttributes());663  CallInst *GuardCheck = B.CreateCall(664      GuardFnType, GuardCheckLoad, {F, Thunk});665  Value *GuardCheckDest = B.CreateExtractValue(GuardCheck, 0);666  Value *GuardFinalDest = B.CreateExtractValue(GuardCheck, 1);667 668  // Ensure that the first argument is passed in the correct register.669  GuardCheck->setCallingConv(CallingConv::CFGuard_Check);670 671  SmallVector<Value *> Args(llvm::make_pointer_range(GuestExit->args()));672  OperandBundleDef OB("cfguardtarget", GuardFinalDest);673  CallInst *Call = B.CreateCall(Arm64Ty, GuardCheckDest, Args, OB);674  Call->setTailCallKind(llvm::CallInst::TCK_MustTail);675 676  if (Call->getType()->isVoidTy())677    B.CreateRetVoid();678  else679    B.CreateRet(Call);680 681  auto SRetAttr = F->getAttributes().getParamAttr(0, Attribute::StructRet);682  auto InRegAttr = F->getAttributes().getParamAttr(0, Attribute::InReg);683  if (SRetAttr.isValid() && !InRegAttr.isValid()) {684    GuestExit->addParamAttr(0, SRetAttr);685    Call->addParamAttr(0, SRetAttr);686  }687 688  return GuestExit;689}690 691Function *692AArch64Arm64ECCallLowering::buildPatchableThunk(GlobalAlias *UnmangledAlias,693                                                GlobalAlias *MangledAlias) {694  llvm::raw_null_ostream NullThunkName;695  FunctionType *Arm64Ty, *X64Ty;696  Function *F = cast<Function>(MangledAlias->getAliasee());697  SmallVector<ThunkArgTranslation> ArgTranslations;698  getThunkType(F->getFunctionType(), F->getAttributes(),699               Arm64ECThunkType::GuestExit, NullThunkName, Arm64Ty, X64Ty,700               ArgTranslations);701  std::string ThunkName(MangledAlias->getName());702  if (ThunkName[0] == '?' && ThunkName.find("@") != std::string::npos) {703    ThunkName.insert(ThunkName.find("@"), "$hybpatch_thunk");704  } else {705    ThunkName.append("$hybpatch_thunk");706  }707 708  Function *GuestExit =709      Function::Create(Arm64Ty, GlobalValue::WeakODRLinkage, 0, ThunkName, M);710  GuestExit->setComdat(M->getOrInsertComdat(ThunkName));711  GuestExit->setSection(".wowthk$aa");712  BasicBlock *BB = BasicBlock::Create(M->getContext(), "", GuestExit);713  IRBuilder<> B(BB);714 715  // Load the global symbol as a pointer to the check function.716  LoadInst *DispatchLoad = B.CreateLoad(PtrTy, DispatchFnGlobal);717 718  // Create new dispatch call instruction.719  Function *ExitThunk =720      buildExitThunk(F->getFunctionType(), F->getAttributes());721  CallInst *Dispatch =722      B.CreateCall(DispatchFnType, DispatchLoad,723                   {UnmangledAlias, ExitThunk, UnmangledAlias->getAliasee()});724 725  // Ensure that the first arguments are passed in the correct registers.726  Dispatch->setCallingConv(CallingConv::CFGuard_Check);727 728  SmallVector<Value *> Args(llvm::make_pointer_range(GuestExit->args()));729  CallInst *Call = B.CreateCall(Arm64Ty, Dispatch, Args);730  Call->setTailCallKind(llvm::CallInst::TCK_MustTail);731 732  if (Call->getType()->isVoidTy())733    B.CreateRetVoid();734  else735    B.CreateRet(Call);736 737  auto SRetAttr = F->getAttributes().getParamAttr(0, Attribute::StructRet);738  auto InRegAttr = F->getAttributes().getParamAttr(0, Attribute::InReg);739  if (SRetAttr.isValid() && !InRegAttr.isValid()) {740    GuestExit->addParamAttr(0, SRetAttr);741    Call->addParamAttr(0, SRetAttr);742  }743 744  MangledAlias->setAliasee(GuestExit);745  return GuestExit;746}747 748// Lower an indirect call with inline code.749void AArch64Arm64ECCallLowering::lowerCall(CallBase *CB) {750  IRBuilder<> B(CB);751  Value *CalledOperand = CB->getCalledOperand();752 753  // If the indirect call is called within catchpad or cleanuppad,754  // we need to copy "funclet" bundle of the call.755  SmallVector<llvm::OperandBundleDef, 1> Bundles;756  if (auto Bundle = CB->getOperandBundle(LLVMContext::OB_funclet))757    Bundles.push_back(OperandBundleDef(*Bundle));758 759  // Load the global symbol as a pointer to the check function.760  Value *GuardFn;761  if (cfguard_module_flag == 2 && !CB->hasFnAttr("guard_nocf"))762    GuardFn = GuardFnCFGlobal;763  else764    GuardFn = GuardFnGlobal;765  LoadInst *GuardCheckLoad = B.CreateLoad(PtrTy, GuardFn);766 767  // Create new call instruction. The CFGuard check should always be a call,768  // even if the original CallBase is an Invoke or CallBr instruction.769  Function *Thunk = buildExitThunk(CB->getFunctionType(), CB->getAttributes());770  CallInst *GuardCheck =771      B.CreateCall(GuardFnType, GuardCheckLoad, {CalledOperand, Thunk},772                   Bundles);773  Value *GuardCheckDest = B.CreateExtractValue(GuardCheck, 0);774  Value *GuardFinalDest = B.CreateExtractValue(GuardCheck, 1);775 776  // Ensure that the first argument is passed in the correct register.777  GuardCheck->setCallingConv(CallingConv::CFGuard_Check);778 779  // Update the call: set the callee, and add a bundle with the final780  // destination,781  CB->setCalledOperand(GuardCheckDest);782  OperandBundleDef OB("cfguardtarget", GuardFinalDest);783  auto *NewCall = CallBase::addOperandBundle(CB, LLVMContext::OB_cfguardtarget,784                                             OB, CB->getIterator());785  NewCall->copyMetadata(*CB);786  CB->replaceAllUsesWith(NewCall);787  CB->eraseFromParent();788}789 790bool AArch64Arm64ECCallLowering::runOnModule(Module &Mod) {791  if (!GenerateThunks)792    return false;793 794  M = &Mod;795 796  // Check if this module has the cfguard flag and read its value.797  if (auto *MD =798          mdconst::extract_or_null<ConstantInt>(M->getModuleFlag("cfguard")))799    cfguard_module_flag = MD->getZExtValue();800 801  PtrTy = PointerType::getUnqual(M->getContext());802  I64Ty = Type::getInt64Ty(M->getContext());803  VoidTy = Type::getVoidTy(M->getContext());804 805  GuardFnType =806      FunctionType::get(StructType::get(PtrTy, PtrTy), {PtrTy, PtrTy}, false);807  DispatchFnType = FunctionType::get(PtrTy, {PtrTy, PtrTy, PtrTy}, false);808  GuardFnCFGlobal = M->getOrInsertGlobal("__os_arm64x_check_icall_cfg", PtrTy);809  GuardFnGlobal = M->getOrInsertGlobal("__os_arm64x_check_icall", PtrTy);810  DispatchFnGlobal = M->getOrInsertGlobal("__os_arm64x_dispatch_call", PtrTy);811 812  // Mangle names of function aliases and add the alias name to813  // arm64ec_unmangled_name metadata to ensure a weak anti-dependency symbol is814  // emitted for the alias as well. Do this early, before handling815  // hybrid_patchable functions, to avoid mangling their aliases.816  for (GlobalAlias &A : Mod.aliases()) {817    auto F = dyn_cast_or_null<Function>(A.getAliaseeObject());818    if (!F)819      continue;820    if (std::optional<std::string> MangledName =821            getArm64ECMangledFunctionName(A)) {822      F->addMetadata("arm64ec_unmangled_name",823                     *MDNode::get(M->getContext(),824                                  MDString::get(M->getContext(), A.getName())));825      A.setName(MangledName.value());826    }827  }828 829  DenseMap<GlobalAlias *, GlobalAlias *> FnsMap;830  SetVector<GlobalAlias *> PatchableFns;831 832  for (Function &F : Mod) {833    if (F.hasPersonalityFn()) {834      GlobalValue *PersFn =835          cast<GlobalValue>(F.getPersonalityFn()->stripPointerCasts());836      if (PersFn->getValueType() && PersFn->getValueType()->isFunctionTy()) {837        if (std::optional<std::string> MangledName =838                getArm64ECMangledFunctionName(*PersFn)) {839          PersFn->setName(MangledName.value());840        }841      }842    }843 844    if (!F.hasFnAttribute(Attribute::HybridPatchable) ||845        F.isDeclarationForLinker() || F.hasLocalLinkage() ||846        F.getName().ends_with(HybridPatchableTargetSuffix))847      continue;848 849    // Rename hybrid patchable functions and change callers to use a global850    // alias instead.851    if (std::optional<std::string> MangledName =852            getArm64ECMangledFunctionName(F)) {853      std::string OrigName(F.getName());854      F.setName(MangledName.value() + HybridPatchableTargetSuffix);855 856      // The unmangled symbol is a weak alias to an undefined symbol with the857      // "EXP+" prefix. This undefined symbol is resolved by the linker by858      // creating an x86 thunk that jumps back to the actual EC target. Since we859      // can't represent that in IR, we create an alias to the target instead.860      // The "EXP+" symbol is set as metadata, which is then used by861      // emitGlobalAlias to emit the right alias.862      auto *A =863          GlobalAlias::create(GlobalValue::LinkOnceODRLinkage, OrigName, &F);864      auto *AM = GlobalAlias::create(GlobalValue::LinkOnceODRLinkage,865                                     MangledName.value(), &F);866      F.replaceUsesWithIf(AM,867                          [](Use &U) { return isa<GlobalAlias>(U.getUser()); });868      F.replaceAllUsesWith(A);869      F.setMetadata("arm64ec_exp_name",870                    MDNode::get(M->getContext(),871                                MDString::get(M->getContext(),872                                              "EXP+" + MangledName.value())));873      A->setAliasee(&F);874      AM->setAliasee(&F);875 876      if (F.hasDLLExportStorageClass()) {877        A->setDLLStorageClass(GlobalValue::DLLExportStorageClass);878        F.setDLLStorageClass(GlobalValue::DefaultStorageClass);879      }880 881      FnsMap[A] = AM;882      PatchableFns.insert(A);883    }884  }885 886  SetVector<GlobalValue *> DirectCalledFns;887  for (Function &F : Mod)888    if (!F.isDeclarationForLinker() &&889        F.getCallingConv() != CallingConv::ARM64EC_Thunk_Native &&890        F.getCallingConv() != CallingConv::ARM64EC_Thunk_X64)891      processFunction(F, DirectCalledFns, FnsMap);892 893  struct ThunkInfo {894    Constant *Src;895    Constant *Dst;896    Arm64ECThunkType Kind;897  };898  SmallVector<ThunkInfo> ThunkMapping;899  for (Function &F : Mod) {900    if (!F.isDeclarationForLinker() &&901        (!F.hasLocalLinkage() || F.hasAddressTaken()) &&902        F.getCallingConv() != CallingConv::ARM64EC_Thunk_Native &&903        F.getCallingConv() != CallingConv::ARM64EC_Thunk_X64) {904      if (!F.hasComdat())905        F.setComdat(Mod.getOrInsertComdat(F.getName()));906      ThunkMapping.push_back(907          {&F, buildEntryThunk(&F), Arm64ECThunkType::Entry});908    }909  }910  for (GlobalValue *O : DirectCalledFns) {911    auto GA = dyn_cast<GlobalAlias>(O);912    auto F = dyn_cast<Function>(GA ? GA->getAliasee() : O);913    ThunkMapping.push_back(914        {O, buildExitThunk(F->getFunctionType(), F->getAttributes()),915         Arm64ECThunkType::Exit});916    if (!GA && !F->hasDLLImportStorageClass())917      ThunkMapping.push_back(918          {buildGuestExitThunk(F), F, Arm64ECThunkType::GuestExit});919  }920  for (GlobalAlias *A : PatchableFns) {921    Function *Thunk = buildPatchableThunk(A, FnsMap[A]);922    ThunkMapping.push_back({Thunk, A, Arm64ECThunkType::GuestExit});923  }924 925  if (!ThunkMapping.empty()) {926    SmallVector<Constant *> ThunkMappingArrayElems;927    for (ThunkInfo &Thunk : ThunkMapping) {928      ThunkMappingArrayElems.push_back(ConstantStruct::getAnon(929          {Thunk.Src, Thunk.Dst,930           ConstantInt::get(M->getContext(), APInt(32, uint8_t(Thunk.Kind)))}));931    }932    Constant *ThunkMappingArray = ConstantArray::get(933        llvm::ArrayType::get(ThunkMappingArrayElems[0]->getType(),934                             ThunkMappingArrayElems.size()),935        ThunkMappingArrayElems);936    new GlobalVariable(Mod, ThunkMappingArray->getType(), /*isConstant*/ false,937                       GlobalValue::ExternalLinkage, ThunkMappingArray,938                       "llvm.arm64ec.symbolmap");939  }940 941  return true;942}943 944bool AArch64Arm64ECCallLowering::processFunction(945    Function &F, SetVector<GlobalValue *> &DirectCalledFns,946    DenseMap<GlobalAlias *, GlobalAlias *> &FnsMap) {947  SmallVector<CallBase *, 8> IndirectCalls;948 949  // For ARM64EC targets, a function definition's name is mangled differently950  // from the normal symbol. We currently have no representation of this sort951  // of symbol in IR, so we change the name to the mangled name, then store952  // the unmangled name as metadata.  Later passes that need the unmangled953  // name (emitting the definition) can grab it from the metadata.954  //955  // FIXME: Handle functions with weak linkage?956  if (!F.hasLocalLinkage() || F.hasAddressTaken()) {957    if (std::optional<std::string> MangledName =958            getArm64ECMangledFunctionName(F)) {959      F.addMetadata("arm64ec_unmangled_name",960                    *MDNode::get(M->getContext(),961                                 MDString::get(M->getContext(), F.getName())));962      if (F.hasComdat() && F.getComdat()->getName() == F.getName()) {963        Comdat *MangledComdat = M->getOrInsertComdat(MangledName.value());964        SmallVector<GlobalObject *> ComdatUsers =965            to_vector(F.getComdat()->getUsers());966        for (GlobalObject *User : ComdatUsers)967          User->setComdat(MangledComdat);968      }969      F.setName(MangledName.value());970    }971  }972 973  // Iterate over the instructions to find all indirect call/invoke/callbr974  // instructions. Make a separate list of pointers to indirect975  // call/invoke/callbr instructions because the original instructions will be976  // deleted as the checks are added.977  for (BasicBlock &BB : F) {978    for (Instruction &I : BB) {979      auto *CB = dyn_cast<CallBase>(&I);980      if (!CB || CB->getCallingConv() == CallingConv::ARM64EC_Thunk_X64 ||981          CB->isInlineAsm())982        continue;983 984      // We need to instrument any call that isn't directly calling an985      // ARM64 function.986      //987      // FIXME: getCalledFunction() fails if there's a bitcast (e.g.988      // unprototyped functions in C)989      if (Function *F = CB->getCalledFunction()) {990        if (!LowerDirectToIndirect || F->hasLocalLinkage() ||991            F->isIntrinsic() || !F->isDeclarationForLinker())992          continue;993 994        DirectCalledFns.insert(F);995        continue;996      }997 998      // Use mangled global alias for direct calls to patchable functions.999      if (GlobalAlias *A = dyn_cast<GlobalAlias>(CB->getCalledOperand())) {1000        auto I = FnsMap.find(A);1001        if (I != FnsMap.end()) {1002          CB->setCalledOperand(I->second);1003          DirectCalledFns.insert(I->first);1004          continue;1005        }1006      }1007 1008      IndirectCalls.push_back(CB);1009      ++Arm64ECCallsLowered;1010    }1011  }1012 1013  if (IndirectCalls.empty())1014    return false;1015 1016  for (CallBase *CB : IndirectCalls)1017    lowerCall(CB);1018 1019  return true;1020}1021 1022char AArch64Arm64ECCallLowering::ID = 0;1023INITIALIZE_PASS(AArch64Arm64ECCallLowering, "Arm64ECCallLowering",1024                "AArch64Arm64ECCallLowering", false, false)1025 1026ModulePass *llvm::createAArch64Arm64ECCallLoweringPass() {1027  return new AArch64Arm64ECCallLowering;1028}1029