IRGen: Co-locate metadata instatiation/completions/accessor functions in a special section

For spatial locality on startup.

Hide collocating metadata functions in a separate section behind a flag.
The default is not to collocate functions.

rdar://101593202

Author: aschwaighofer

include/swift/AST/IRGenOptions.h

@@ -419,6 +419,9 @@ class IRGenOptions {
 
   unsigned DisableReadonlyStaticObjects : 1;
 
+  /// Collocate metadata functions in their own section.
+  unsigned CollocatedMetadataFunctions : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -489,7 +492,8 @@ class IRGenOptions {
         WitnessMethodElimination(false), ConditionalRuntimeRecords(false),
         InternalizeAtLink(false), InternalizeSymbols(false),
         EmitGenericRODatas(false), NoPreallocatedInstantiationCaches(false),
-        DisableReadonlyStaticObjects(false), CmdArgs(),
+        DisableReadonlyStaticObjects(false),
+        CollocatedMetadataFunctions(false), CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All) {
 #ifndef NDEBUG

include/swift/Option/FrontendOptions.td

@@ -1105,6 +1105,13 @@ def enable_move_inout_stack_protector :
   Flag<["-"], "enable-move-inout-stack-protector">,
   HelpText<"Enable the stack protector by moving values to temporaries">;
 
+def enable_collocate_metadata_functions :
+  Flag<["-"], "enable-collocate-metadata-functions">,
+  HelpText<"Enable collocate metadata functions">;
+def disable_collocate_metadata_functions :
+  Flag<["-"], "disable-collocate-metadata-functions">,
+  HelpText<"Disable collocate metadata functions">;
+
 def enable_new_llvm_pass_manager :
   Flag<["-"], "enable-new-llvm-pass-manager">,
   HelpText<"Enable the new llvm pass manager">;

lib/Frontend/CompilerInvocation.cpp

@@ -2472,7 +2472,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
       Args.hasFlag(OPT_disable_new_llvm_pass_manager,
                    OPT_enable_new_llvm_pass_manager,
                    Opts.LegacyPassManager);
-
+  Opts.CollocatedMetadataFunctions =
+      Args.hasFlag(OPT_enable_collocate_metadata_functions,
+                   OPT_disable_collocate_metadata_functions,
+                   Opts.CollocatedMetadataFunctions);
   return false;
 }
 

lib/IRGen/GenClass.cpp

@@ -2511,6 +2511,7 @@ static llvm::Function *emitObjCMetadataUpdateFunction(IRGenModule &IGM,
   llvm::Function *f =
     IGM.getAddrOfObjCMetadataUpdateFunction(D, ForDefinition);
   f->setAttributes(IGM.constructInitialAttributes());
+  IGM.setColocateMetadataSection(f);
 
   IRGenFunction IGF(IGM, f);
   if (IGM.DebugInfo)

lib/IRGen/GenDecl.cpp

@@ -1979,6 +1979,7 @@ void IRGenerator::emitEagerClassInitialization() {
     IGM->DebugInfo->emitArtificialFunction(RegisterIGF, RegisterIGF.CurFn);
   RegisterFn->setAttributes(IGM->constructInitialAttributes());
   RegisterFn->setCallingConv(IGM->DefaultCC);
+  IGM->setColocateMetadataSection(RegisterFn);
 
   for (ClassDecl *CD : ClassesForEagerInitialization) {
     auto Ty = CD->getDeclaredType()->getCanonicalType();
@@ -2022,6 +2023,7 @@ void IRGenerator::emitObjCActorsNeedingSuperclassSwizzle() {
     IGM->DebugInfo->emitArtificialFunction(RegisterIGF, RegisterIGF.CurFn);
   RegisterFn->setAttributes(IGM->constructInitialAttributes());
   RegisterFn->setCallingConv(IGM->DefaultCC);
+  IGM->setColocateMetadataSection(RegisterFn);
 
   // Look up the SwiftNativeNSObject class.
   auto swiftNativeNSObjectName =
@@ -5858,3 +5860,23 @@ IRGenModule::getOrCreateHelperFunction(StringRef fnName, llvm::Type *resultTy,
 
   return fn;
 }
+
+void IRGenModule::setColocateMetadataSection(llvm::Function *f) {
+  if (!IRGen.Opts.CollocatedMetadataFunctions)
+    return;
+
+  switch (TargetInfo.OutputObjectFormat) {
+  case llvm::Triple::MachO:
+    f->setSection("__TEXT, __textg_swiftm, regular, pure_instructions");
+    break;
+  case llvm::Triple::DXContainer:
+  case llvm::Triple::GOFF:
+  case llvm::Triple::SPIRV:
+  case llvm::Triple::UnknownObjectFormat:
+  case llvm::Triple::Wasm:
+  case llvm::Triple::ELF:
+  case llvm::Triple::XCOFF:
+  case llvm::Triple::COFF:
+    break;
+  }
+}

lib/IRGen/GenMeta.cpp

@@ -206,6 +206,7 @@ static void emitMetadataCompletionFunction(IRGenModule &IGM,
   f->setAttributes(IGM.constructInitialAttributes());
   f->setDoesNotThrow();
   IGM.setHasNoFramePointer(f);
+  IGM.setColocateMetadataSection(f);
 
   IRGenFunction IGF(IGM, f);
 
@@ -2969,6 +2970,7 @@ namespace {
       f->setAttributes(IGM.constructInitialAttributes());
       f->setDoesNotThrow();
       IGM.setHasNoFramePointer(f);
+      IGM.setColocateMetadataSection(f);
 
       IRGenFunction IGF(IGM, f);
 

lib/IRGen/IRGenModule.h

@@ -1572,6 +1572,8 @@ private:                                                                       \
   void appendLLVMUsedConditionalEntry(llvm::GlobalVariable *var,
                                       const ProtocolConformance *conformance);
 
+  void setColocateMetadataSection(llvm::Function *f);
+
   llvm::Constant *
   getAddrOfTypeMetadata(CanType concreteType,
                         TypeMetadataCanonicality canonicality =

lib/IRGen/MetadataRequest.cpp

@@ -2052,6 +2052,7 @@ void irgen::emitCacheAccessFunction(IRGenModule &IGM, llvm::Function *accessor,
   accessor->addFnAttr(llvm::Attribute::NoInline);
   // Accessor functions don't need frame pointers.
   IGM.setHasNoFramePointer(accessor);
+  IGM.setColocateMetadataSection(accessor);
 
   // This function is logically 'readnone': the caller does not need
   // to reason about any side effects or stores it might perform.
@@ -2404,6 +2405,7 @@ MetadataResponse irgen::emitGenericTypeMetadataAccessFunction(
           generateThunkFn,
           /*noinline*/ true));
     }
+    IGM.setColocateMetadataSection(thunkFn);
 
     // Call out to the helper.
     auto arg0 = numArguments >= 1
@@ -2674,6 +2676,8 @@ irgen::getGenericTypeMetadataAccessFunction(IRGenModule &IGM,
   llvm::Function *accessor =
     IGM.getAddrOfGenericTypeMetadataAccessFunction(
         nominal, genericArgs.Types, shouldDefine);
+  if (shouldDefine)
+    IGM.setColocateMetadataSection(accessor);
 
   // If we're not supposed to define the accessor, or if we already
   // have defined it, just return the pointer.
@@ -2917,6 +2921,7 @@ emitMetadataAccessByMangledName(IRGenFunction &IGF, CanType type,
     IGM.setHasNoFramePointer(subIGF.CurFn);
     if (IGM.DebugInfo)
       IGM.DebugInfo->emitArtificialFunction(subIGF, subIGF.CurFn);
+    IGM.setColocateMetadataSection(subIGF.CurFn);
 
     auto params = subIGF.collectParameters();
     auto cache = params.claimNext();

test/IRGen/collocate_metadata_functions.swift

@@ -0,0 +1,9 @@
+// RUN: %target-swift-frontend -primary-file %s -emit-ir -enable-collocate-metadata-functions | %FileCheck %s
+
+// REQUIRES: OS=macosx
+
+// CHECK: define{{.*}} swiftcc %swift.metadata_response @"$s28collocate_metadata_functions13GenericStructVMr"({{.*}} section "__TEXT, __textg_swiftm, regular, pure_instructions"
+
+public struct GenericStruct<T> {
+    var field: T?
+}

test/IRGen/generic_classes.sil

@@ -415,5 +415,5 @@ entry(%c : $RootGeneric<Int32>):
 // CHECK:   ret %swift.metadata_response [[T1]]
 // CHECK: }
 
-// OSIZE: define hidden swiftcc %swift.metadata_response @"$s15generic_classes11RootGenericCMa"(i64 %0, %swift.type* {{.*}}) [[ATTRS:#[0-9]+]] {
+// OSIZE: define hidden swiftcc %swift.metadata_response @"$s15generic_classes11RootGenericCMa"(i64 %0, %swift.type* {{.*}}) [[ATTRS:#[0-9]+]] {{(section)?.*}}{
 // OSIZE: [[ATTRS]] = {{{.*}}noinline