[aarch64][x86][win] Add compiler support for MSVC's /funcoverride flag (Windows kernel loader replaceable functions)

clang/include/clang/Basic/CodeGenOptions.h

@@ -499,6 +499,9 @@ class CodeGenOptions : public CodeGenOptionsBase {
   /// The name of a file to use with \c .secure_log_unique directives.
   std::string AsSecureLogFile;
 
+  /// A list of functions that are replacable by the loader.
+  std::vector<std::string> LoaderReplaceableFunctionNames;
+
 public:
   // Define accessors/mutators for code generation options of enumeration type.
 #define CODEGENOPT(Name, Bits, Default)
@@ -571,6 +574,12 @@ class CodeGenOptions : public CodeGenOptionsBase {
   /// Reset all of the options that are not considered when building a
   /// module.
   void resetNonModularOptions(StringRef ModuleFormat);
+
+  // Is the given function name one of the functions that can be replaced by the
+  // loader?
+  bool isLoaderReplaceableFunctionName(StringRef FuncName) const {
+    return llvm::is_contained(LoaderReplaceableFunctionNames, FuncName);
+  }
 };
 
 }  // end namespace clang

clang/include/clang/Driver/Options.td

@@ -7826,6 +7826,9 @@ def import_call_optimization : Flag<["-"], "import-call-optimization">,
              "by the Windows kernel to enable import call optimization">,
     MarshallingInfoFlag<CodeGenOpts<"ImportCallOptimization">>;
 
+def replaceable_function: Joined<["-"], "loader-replaceable-function=">,
+  MarshallingInfoStringVector<CodeGenOpts<"LoaderReplaceableFunctionNames">>;
+
 } // let Visibility = [CC1Option]
 
 //===----------------------------------------------------------------------===//
@@ -9088,6 +9091,10 @@ def _SLASH_Gregcall : CLFlag<"Gregcall">,
 def _SLASH_Gregcall4 : CLFlag<"Gregcall4">,
   HelpText<"Set __regcall4 as a default calling convention to respect __regcall ABI v.4">;
 
+def _SLASH_funcoverride : CLCompileJoined<"funcoverride:">,
+  HelpText<"Mark <function> as being replaceable by the Windows kernel loader">,
+  MetaVarName<"<function>">;
+
 // GNU Driver aliases
 
 def : Separate<["-"], "Xmicrosoft-visualc-tools-root">, Alias<_SLASH_vctoolsdir>;

clang/lib/CodeGen/CGCall.cpp

@@ -2641,6 +2641,10 @@ void CodeGenModule::ConstructAttributeList(StringRef Name,
     GetCPUAndFeaturesAttributes(CalleeInfo.getCalleeDecl(), FuncAttrs);
   }
 
+  // Mark functions that are replaceable by the loader.
+  if (CodeGenOpts.isLoaderReplaceableFunctionName(Name))
+    FuncAttrs.addAttribute("loader-replaceable");
+
   // Collect attributes from arguments and return values.
   ClangToLLVMArgMapping IRFunctionArgs(getContext(), FI);
 

clang/lib/Driver/ToolChains/Clang.cpp

@@ -8578,6 +8578,12 @@ void Clang::AddClangCLArgs(const ArgList &Args, types::ID InputType,
     }
     A->claim();
   }
+
+  for (const auto &FuncOverride :
+       Args.getAllArgValues(options::OPT__SLASH_funcoverride)) {
+    CmdArgs.push_back(Args.MakeArgString(
+        Twine("-loader-replaceable-function=") + FuncOverride));
+  }
 }
 
 const char *Clang::getBaseInputName(const ArgList &Args,

clang/test/CodeGen/loader-replaceable-function.cpp

@@ -0,0 +1,17 @@
+// RUN: %clang_cc1 -triple=x86_64-pc-windows-msvc -loader-replaceable-function=override_me -loader-replaceable-function="?override_me_cpp@@YAXXZ" -emit-llvm -o - %s | FileCheck %s
+
+// CHECK: define dso_local void @override_me() #0
+extern "C" void override_me() {}
+
+// CHECK: define dso_local void @"?override_me_cpp@@YAXXZ"() #0
+void override_me_cpp() {}
+
+// CHECK: define dso_local void @dont_override_me() #1
+extern "C" void dont_override_me() {}
+
+// CHECK: attributes #0 = { 
+// CHECK-SAME: loader-replaceable
+
+// CHECK: attributes #1 = { 
+// CHECK-NOT: loader-replaceable
+// CHECK-SAME: }

clang/test/Driver/cl-options.c

@@ -823,4 +823,8 @@
 // RUN: %clang_cl /d2epilogunwind /c -### -- %s 2>&1 | FileCheck %s --check-prefix=EPILOGUNWIND
 // EPILOGUNWIND: -fwinx64-eh-unwindv2
 
+// RUN: %clang_cl /funcoverride:override_me1 /funcoverride:override_me2 /c -### -- %s 2>&1 | FileCheck %s --check-prefix=FUNCOVERRIDE
+// FUNCOVERRIDE: -loader-replaceable-function=override_me1
+// FUNCOVERRIDE-SAME: -loader-replaceable-function=override_me2
+
 void f(void) { }

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -812,6 +812,17 @@ class AsmPrinter : public MachineFunctionPass {
   getCodeViewJumpTableInfo(int JTI, const MachineInstr *BranchInstr,
                            const MCSymbol *BranchLabel) const;
 
+  //===------------------------------------------------------------------===//
+  // COFF Helper Routines
+  //===------------------------------------------------------------------===//
+
+  /// Emits symbols and data to allow functions marked with the
+  /// loader-replaceable attribute to be replaceable.
+  void emitCOFFReplaceableFunctionData(Module &M);
+
+  /// Emits the @feat.00 symbol indicating the features enabled in this module.
+  void emitCOFFFeatureSymbol(Module &M);
+
   //===------------------------------------------------------------------===//
   // Inline Asm Support
   //===------------------------------------------------------------------===//

llvm/include/llvm/IR/Attributes.td

@@ -400,6 +400,7 @@ def NoJumpTables : StrBoolAttr<"no-jump-tables">;
 def NoInlineLineTables : StrBoolAttr<"no-inline-line-tables">;
 def ProfileSampleAccurate : StrBoolAttr<"profile-sample-accurate">;
 def UseSampleProfile : StrBoolAttr<"use-sample-profile">;
+def LoaderReplaceable : StrBoolAttr<"loader-replaceable">;
 
 def DenormalFPMath : ComplexStrAttr<"denormal-fp-math", [FnAttr]>;
 def DenormalFPMathF32 : ComplexStrAttr<"denormal-fp-math-f32", [FnAttr]>;

llvm/include/llvm/IR/Mangler.h

@@ -25,6 +25,8 @@ class Triple;
 class Twine;
 class raw_ostream;
 
+constexpr std::string_view HybridPatchableTargetSuffix = "$hp_target";
+
 class Mangler {
   /// We need to give global values the same name every time they are mangled.
   /// This keeps track of the number we give to anonymous ones.

llvm/lib/Analysis/InlineCost.cpp

@@ -3174,6 +3174,10 @@ std::optional<InlineResult> llvm::getAttributeBasedInliningDecision(
   if (Call.isNoInline())
     return InlineResult::failure("noinline call site attribute");
 
+  // Don't inline functions that are loader replaceable.
+  if (Callee->hasFnAttribute("loader-replaceable"))
+    return InlineResult::failure("loader replaceable function attribute");
+
   return std::nullopt;
 }
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -4723,3 +4723,109 @@ AsmPrinter::getCodeViewJumpTableInfo(int JTI, const MachineInstr *BranchInstr,
   return std::make_tuple(Base, 0, BranchLabel,
                          codeview::JumpTableEntrySize::Int32);
 }
+
+void AsmPrinter::emitCOFFReplaceableFunctionData(Module &M) {
+  const Triple &TT = TM.getTargetTriple();
+  assert(TT.isOSBinFormatCOFF());
+
+  bool IsTargetArm64EC = TT.isWindowsArm64EC();
+  SmallVector<char> Buf;
+  SmallVector<MCSymbol *> FuncOverrideDefaultSymbols;
+  bool SwitchedToDirectiveSection = false;
+  for (const Function &F : M.functions()) {
+    if (F.hasFnAttribute("loader-replaceable")) {
+      if (!SwitchedToDirectiveSection) {
+        OutStreamer->switchSection(
+            OutContext.getObjectFileInfo()->getDrectveSection());
+        SwitchedToDirectiveSection = true;
+      }
+
+      StringRef Name = F.getName();
+
+      // For hybrid-patchable targets, strip the prefix so that we can mark
+      // the real function as replaceable.
+      if (IsTargetArm64EC && Name.ends_with(HybridPatchableTargetSuffix)) {
+        Name = Name.drop_back(HybridPatchableTargetSuffix.size());
+      }
+
+      MCSymbol *FuncOverrideSymbol =
+          MMI->getContext().getOrCreateSymbol(Name + "_$fo$");
+      OutStreamer->beginCOFFSymbolDef(FuncOverrideSymbol);
+      OutStreamer->emitCOFFSymbolStorageClass(COFF::IMAGE_SYM_CLASS_EXTERNAL);
+      OutStreamer->emitCOFFSymbolType(COFF::IMAGE_SYM_DTYPE_NULL);
+      OutStreamer->endCOFFSymbolDef();
+
+      MCSymbol *FuncOverrideDefaultSymbol =
+          MMI->getContext().getOrCreateSymbol(Name + "_$fo_default$");
+      OutStreamer->beginCOFFSymbolDef(FuncOverrideDefaultSymbol);
+      OutStreamer->emitCOFFSymbolStorageClass(COFF::IMAGE_SYM_CLASS_EXTERNAL);
+      OutStreamer->emitCOFFSymbolType(COFF::IMAGE_SYM_DTYPE_NULL);
+      OutStreamer->endCOFFSymbolDef();
+      FuncOverrideDefaultSymbols.push_back(FuncOverrideDefaultSymbol);
+
+      OutStreamer->emitBytes((Twine(" /ALTERNATENAME:") +
+                              FuncOverrideSymbol->getName() + "=" +
+                              FuncOverrideDefaultSymbol->getName())
+                                 .toStringRef(Buf));
+      Buf.clear();
+    }
+  }
+
+  if (SwitchedToDirectiveSection)
+    OutStreamer->popSection();
+
+  if (FuncOverrideDefaultSymbols.empty())
+    return;
+
+  // MSVC emits the symbols for the default variables pointing at the start of
+  // the .data section, but doesn't actually allocate any space for them. LLVM
+  // can't do this, so have all of the variables pointing at a single byte
+  // instead.
+  OutStreamer->switchSection(OutContext.getObjectFileInfo()->getDataSection());
+  for (MCSymbol *Symbol : FuncOverrideDefaultSymbols) {
+    OutStreamer->emitLabel(Symbol);
+  }
+  OutStreamer->emitZeros(1);
+  OutStreamer->popSection();
+}
+
+void AsmPrinter::emitCOFFFeatureSymbol(Module &M) {
+  const Triple &TT = TM.getTargetTriple();
+  assert(TT.isOSBinFormatCOFF());
+
+  // Emit an absolute @feat.00 symbol.
+  MCSymbol *S = MMI->getContext().getOrCreateSymbol(StringRef("@feat.00"));
+  OutStreamer->beginCOFFSymbolDef(S);
+  OutStreamer->emitCOFFSymbolStorageClass(COFF::IMAGE_SYM_CLASS_STATIC);
+  OutStreamer->emitCOFFSymbolType(COFF::IMAGE_SYM_DTYPE_NULL);
+  OutStreamer->endCOFFSymbolDef();
+  int64_t Feat00Value = 0;
+
+  if (TT.getArch() == Triple::x86) {
+    // According to the PE-COFF spec, the LSB of this value marks the object
+    // for "registered SEH".  This means that all SEH handler entry points
+    // must be registered in .sxdata.  Use of any unregistered handlers will
+    // cause the process to terminate immediately.  LLVM does not know how to
+    // register any SEH handlers, so its object files should be safe.
+    Feat00Value |= COFF::Feat00Flags::SafeSEH;
+  }
+
+  if (M.getModuleFlag("cfguard")) {
+    // Object is CFG-aware.
+    Feat00Value |= COFF::Feat00Flags::GuardCF;
+  }
+
+  if (M.getModuleFlag("ehcontguard")) {
+    // Object also has EHCont.
+    Feat00Value |= COFF::Feat00Flags::GuardEHCont;
+  }
+
+  if (M.getModuleFlag("ms-kernel")) {
+    // Object is compiled with /kernel.
+    Feat00Value |= COFF::Feat00Flags::Kernel;
+  }
+
+  OutStreamer->emitSymbolAttribute(S, MCSA_Global);
+  OutStreamer->emitAssignment(
+      S, MCConstantExpr::create(Feat00Value, MMI->getContext()));
+}

llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp

@@ -815,15 +815,16 @@ bool AArch64Arm64ECCallLowering::runOnModule(Module &Mod) {
     }
 
     if (!F.hasFnAttribute(Attribute::HybridPatchable) || F.isDeclaration() ||
-        F.hasLocalLinkage() || F.getName().ends_with("$hp_target"))
+        F.hasLocalLinkage() ||
+        F.getName().ends_with(HybridPatchableTargetSuffix))
       continue;
 
     // Rename hybrid patchable functions and change callers to use a global
     // alias instead.
     if (std::optional<std::string> MangledName =
             getArm64ECMangledFunctionName(F.getName().str())) {
       std::string OrigName(F.getName());
-      F.setName(MangledName.value() + "$hp_target");
+      F.setName(MangledName.value() + HybridPatchableTargetSuffix);
 
       // The unmangled symbol is a weak alias to an undefined symbol with the
       // "EXP+" prefix. This undefined symbol is resolved by the linker by

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -329,32 +329,8 @@ void AArch64AsmPrinter::emitStartOfAsmFile(Module &M) {
   const Triple &TT = TM.getTargetTriple();
 
   if (TT.isOSBinFormatCOFF()) {
-    // Emit an absolute @feat.00 symbol
-    MCSymbol *S = MMI->getContext().getOrCreateSymbol(StringRef("@feat.00"));
-    OutStreamer->beginCOFFSymbolDef(S);
-    OutStreamer->emitCOFFSymbolStorageClass(COFF::IMAGE_SYM_CLASS_STATIC);
-    OutStreamer->emitCOFFSymbolType(COFF::IMAGE_SYM_DTYPE_NULL);
-    OutStreamer->endCOFFSymbolDef();
-    int64_t Feat00Value = 0;
-
-    if (M.getModuleFlag("cfguard")) {
-      // Object is CFG-aware.
-      Feat00Value |= COFF::Feat00Flags::GuardCF;
-    }
-
-    if (M.getModuleFlag("ehcontguard")) {
-      // Object also has EHCont.
-      Feat00Value |= COFF::Feat00Flags::GuardEHCont;
-    }
-
-    if (M.getModuleFlag("ms-kernel")) {
-      // Object is compiled with /kernel.
-      Feat00Value |= COFF::Feat00Flags::Kernel;
-    }
-
-    OutStreamer->emitSymbolAttribute(S, MCSA_Global);
-    OutStreamer->emitAssignment(
-        S, MCConstantExpr::create(Feat00Value, MMI->getContext()));
+    emitCOFFFeatureSymbol(M);
+    emitCOFFReplaceableFunctionData(M);
 
     if (M.getModuleFlag("import-call-optimization"))
       EnableImportCallOptimization = true;

llvm/lib/Target/X86/X86AsmPrinter.cpp

@@ -910,41 +910,8 @@ void X86AsmPrinter::emitStartOfAsmFile(Module &M) {
     OutStreamer->switchSection(getObjFileLowering().getTextSection());
 
   if (TT.isOSBinFormatCOFF()) {
-    // Emit an absolute @feat.00 symbol.
-    MCSymbol *S = MMI->getContext().getOrCreateSymbol(StringRef("@feat.00"));
-    OutStreamer->beginCOFFSymbolDef(S);
-    OutStreamer->emitCOFFSymbolStorageClass(COFF::IMAGE_SYM_CLASS_STATIC);
-    OutStreamer->emitCOFFSymbolType(COFF::IMAGE_SYM_DTYPE_NULL);
-    OutStreamer->endCOFFSymbolDef();
-    int64_t Feat00Value = 0;
-
-    if (TT.getArch() == Triple::x86) {
-      // According to the PE-COFF spec, the LSB of this value marks the object
-      // for "registered SEH".  This means that all SEH handler entry points
-      // must be registered in .sxdata.  Use of any unregistered handlers will
-      // cause the process to terminate immediately.  LLVM does not know how to
-      // register any SEH handlers, so its object files should be safe.
-      Feat00Value |= COFF::Feat00Flags::SafeSEH;
-    }
-
-    if (M.getModuleFlag("cfguard")) {
-      // Object is CFG-aware.
-      Feat00Value |= COFF::Feat00Flags::GuardCF;
-    }
-
-    if (M.getModuleFlag("ehcontguard")) {
-      // Object also has EHCont.
-      Feat00Value |= COFF::Feat00Flags::GuardEHCont;
-    }
-
-    if (M.getModuleFlag("ms-kernel")) {
-      // Object is compiled with /kernel.
-      Feat00Value |= COFF::Feat00Flags::Kernel;
-    }
-
-    OutStreamer->emitSymbolAttribute(S, MCSA_Global);
-    OutStreamer->emitAssignment(
-        S, MCConstantExpr::create(Feat00Value, MMI->getContext()));
+    emitCOFFFeatureSymbol(M);
+    emitCOFFReplaceableFunctionData(M);
   }
   OutStreamer->emitSyntaxDirective();
 

llvm/test/CodeGen/AArch64/win-loader-replaceable-function.ll

@@ -0,0 +1,40 @@
+; RUN: llc -mtriple=aarch64-pc-windows-msvc < %s | FileCheck %s
+
+define dso_local i32 @override_me1() "loader-replaceable" {
+entry:
+  ret i32 1
+}
+
+define dso_local i32 @override_me2() "loader-replaceable" {
+entry:
+  ret i32 2
+}
+
+define dso_local i32 @dont_override_me() {
+entry:
+  ret i32 3
+}
+
+; CHECK:              .section        .drectve,"yni"
+; CHECK-NEXT:         .def    override_me1_$fo$;
+; CHECK-NEXT:         .scl    2;
+; CHECK-NEXT:         .type   0;
+; CHECK-NEXT:         .endef
+; CHECK-NEXT:         .def    override_me1_$fo_default$;
+; CHECK-NEXT:         .scl    2;
+; CHECK-NEXT:         .type   0;
+; CHECK-NEXT:         .endef
+; CHECK-NEXT:         .ascii  " /ALTERNATENAME:override_me1_$fo$=override_me1_$fo_default$"
+; CHECK-NEXT:         .def    override_me2_$fo$;
+; CHECK-NEXT:         .scl    2;
+; CHECK-NEXT:         .type   0;
+; CHECK-NEXT:         .endef
+; CHECK-NEXT:         .def    override_me2_$fo_default$;
+; CHECK-NEXT:         .scl    2;
+; CHECK-NEXT:         .type   0;
+; CHECK-NEXT:         .endef
+; CHECK-NEXT:         .ascii  " /ALTERNATENAME:override_me2_$fo$=override_me2_$fo_default$"
+; CHECK-NEXT:         .data
+; CHECK-NEXT: override_me1_$fo_default$:
+; CHECK-NEXT: override_me2_$fo_default$:
+; CHECK-NEXT:         .zero   1