Add a flag to disable compile-time preallocated instantiation caches (for type metadata and protocol conformances) (#41148)

Author: kubamracek

include/swift/AST/IRGenOptions.h

@@ -370,6 +370,10 @@ class IRGenOptions {
 
   unsigned InternalizeAtLink : 1;
 
+  /// Whether to avoid emitting zerofill globals as preallocated type metadata
+  /// and prototol conformance caches.
+  unsigned NoPreallocatedInstantiationCaches : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -433,6 +437,7 @@ class IRGenOptions {
         EnableGlobalISel(false), VirtualFunctionElimination(false),
         WitnessMethodElimination(false), ConditionalRuntimeRecords(false),
         InternalizeAtLink(false),
+        NoPreallocatedInstantiationCaches(false),
         CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All) {}

include/swift/Option/FrontendOptions.td

@@ -579,6 +579,10 @@ def conditional_runtime_records : Flag<["-"], "conditional-runtime-records">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
   HelpText<"Allow removal of runtime metadata records (public types, protocol conformances) based on whether they're used or unused">;
 
+def disable_preallocated_instantiation_caches : Flag<["-"], "disable-preallocated-instantiation-caches">,
+  Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
+  HelpText<"Avoid preallocating metadata instantiation caches in globals">;
+
 def disable_previous_implementation_calls_in_dynamic_replacements :
   Flag<["-"], "disable-previous-implementation-calls-in-dynamic-replacements">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,

lib/Frontend/CompilerInvocation.cpp

@@ -2182,6 +2182,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     Opts.InternalizeAtLink = true;
   }
 
+  if (Args.hasArg(OPT_disable_preallocated_instantiation_caches)) {
+    Opts.NoPreallocatedInstantiationCaches = true;
+  }
+
   // Default to disabling swift async extended frame info on anything but
   // darwin. Other platforms are unlikely to have support for extended frame
   // pointer information.

lib/IRGen/GenMeta.cpp

@@ -1267,7 +1267,7 @@ namespace {
     }
 
     void addMetadataInstantiationCache() {
-      if (!HasMetadata) {
+      if (!HasMetadata || IGM.getOptions().NoPreallocatedInstantiationCaches) {
         B.addInt32(0);
         return;
       }
@@ -2569,6 +2569,8 @@ namespace {
 
     /// Emit the instantiation cache variable for the template.
     void emitInstantiationCache() {
+      if (IGM.IRGen.Opts.NoPreallocatedInstantiationCaches) return;
+      
       auto cache = cast<llvm::GlobalVariable>(
         IGM.getAddrOfTypeMetadataInstantiationCache(Target, ForDefinition));
       auto init =

lib/IRGen/GenProto.cpp

@@ -1978,8 +1978,11 @@ namespace {
                 Description.requiresSpecialization);
       // Instantiation function
       B.addRelativeAddressOrNull(Description.instantiationFn);
+
       // Private data
-      {
+      if (IGM.IRGen.Opts.NoPreallocatedInstantiationCaches) {
+        B.addInt32(0);
+      } else {
         auto privateDataTy =
           llvm::ArrayType::get(IGM.Int8PtrTy,
                                swift::NumGenericMetadataPrivateDataWords);

stdlib/public/runtime/Metadata.cpp

@@ -389,7 +389,35 @@ namespace {
   };
 
   using LazyGenericMetadataCache = Lazy<GenericMetadataCache>;
-}
+
+  class GlobalMetadataCacheEntry {
+  public:
+    const TypeContextDescriptor *Description;
+    GenericMetadataCache Cache;
+
+    GlobalMetadataCacheEntry(const TypeContextDescriptor *description)
+        : Description(description), Cache(*description->getGenericContext()) {}
+
+    intptr_t getKeyIntValueForDump() {
+      return reinterpret_cast<intptr_t>(Description);
+    }
+    bool matchesKey(const TypeContextDescriptor *description) const {
+      return description == Description;
+    }
+    friend llvm::hash_code hash_value(const GlobalMetadataCacheEntry &value) {
+      return llvm::hash_value(value.Description);
+    }
+    static size_t
+    getExtraAllocationSize(const TypeContextDescriptor *description) {
+      return 0;
+    }
+    size_t getExtraAllocationSize() const { return 0; }
+  };
+
+  static SimpleGlobalCache<GlobalMetadataCacheEntry, GlobalMetadataCacheTag>
+      GlobalMetadataCache;
+
+} // end anonymous namespace
 
 /// Fetch the metadata cache for a generic metadata structure.
 static GenericMetadataCache &getCache(
@@ -401,6 +429,11 @@ static GenericMetadataCache &getCache(
                 sizeof(GenericMetadataInstantiationCache::PrivateData),
                 "metadata cache is larger than the allowed space");
 
+  auto *cacheStorage = generics.getInstantiationCache();
+  if (cacheStorage == nullptr) {
+    return GlobalMetadataCache.getOrInsert(&description).first->Cache;
+  }
+
   auto lazyCache =
     reinterpret_cast<LazyGenericMetadataCache*>(
       generics.getInstantiationCache()->PrivateData);
@@ -4606,19 +4639,50 @@ class WitnessTableCacheEntry :
                          const void * const *instantiationArgs);
 };
 
-} // end anonymous namespace
-
 using GenericWitnessTableCache =
   MetadataCache<WitnessTableCacheEntry, GenericWitnessTableCacheTag>;
 using LazyGenericWitnessTableCache = Lazy<GenericWitnessTableCache>;
 
+class GlobalWitnessTableCacheEntry {
+public:
+  const GenericWitnessTable *Gen;
+  GenericWitnessTableCache Cache;
+
+  GlobalWitnessTableCacheEntry(const GenericWitnessTable *gen)
+      : Gen(gen), Cache() {}
+
+  intptr_t getKeyIntValueForDump() {
+    return reinterpret_cast<intptr_t>(Gen);
+  }
+  bool matchesKey(const GenericWitnessTable *gen) const {
+    return gen == Gen;
+  }
+  friend llvm::hash_code hash_value(const GlobalWitnessTableCacheEntry &value) {
+    return llvm::hash_value(value.Gen);
+  }
+  static size_t
+  getExtraAllocationSize(const GenericWitnessTable *gen) {
+    return 0;
+  }
+  size_t getExtraAllocationSize() const { return 0; }
+};
+
+static SimpleGlobalCache<GlobalWitnessTableCacheEntry, GlobalWitnessTableCacheTag>
+    GlobalWitnessTableCache;
+
+} // end anonymous namespace
+
 /// Fetch the cache for a generic witness-table structure.
 static GenericWitnessTableCache &getCache(const GenericWitnessTable *gen) {
   // Keep this assert even if you change the representation above.
   static_assert(sizeof(LazyGenericWitnessTableCache) <=
                 sizeof(GenericWitnessTable::PrivateDataType),
                 "metadata cache is larger than the allowed space");
 
+  if (gen->PrivateData == nullptr) {
+    return GlobalWitnessTableCache.getOrInsert(gen).first->Cache;
+  }
+
   auto lazyCache =
     reinterpret_cast<LazyGenericWitnessTableCache*>(gen->PrivateData.get());
   return lazyCache->get();

stdlib/public/runtime/MetadataAllocatorTags.def

@@ -42,5 +42,7 @@ TAG(GenericWitnessTableCache, 16)
 TAG(GenericClassMetadata, 17)
 TAG(GenericValueMetadata, 18)
 TAG(SingletonGenericWitnessTableCache, 19)
+TAG(GlobalMetadataCache, 20)
+TAG(GlobalWitnessTableCache, 21)
 
 #undef TAG

test/IRGen/disable-instantiation-cache-exec.swift

@@ -0,0 +1,23 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -Xfrontend -disable-preallocated-instantiation-caches -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out
+
+// REQUIRES: executable_test
+
+public class Generic<T> {
+  public func m1(t: T) -> T { return t }
+  public func m2(t: T) -> T { return t }
+}
+
+protocol MyProtocol {
+  associatedtype T
+  func foo() -> T
+}
+
+public struct MyStruct<T>: MyProtocol {
+  func foo() -> T { fatalError() }
+}
+
+print(Generic<Int>())
+print(MyStruct<Int>())

test/IRGen/disable-instantiation-cache.swift

@@ -0,0 +1,32 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -module-name main %s -emit-ir | %FileCheck %s --check-prefix=CHECK-CACHE
+// RUN: %target-swift-frontend -module-name main %s -emit-ir -disable-preallocated-instantiation-caches | %FileCheck %s --check-prefix=CHECK-NOCACHE
+
+public class Generic<T> {
+  public func m1(t: T) -> T { return t }
+  public func m2(t: T) -> T { return t }
+}
+
+protocol MyProtocol {
+  associatedtype T
+  func foo() -> T
+}
+
+public struct MyStruct<T>: MyProtocol {
+  func foo() -> T { fatalError() }
+}
+
+// "metadata instantiation cache for protocol conformance descriptor for main.MyStruct<A> : main.MyProtocol in main"
+// CHECK-CACHE: @"$s4main8MyStructVyxGAA0B8ProtocolAAMcMK" = internal global [{{.*}} x i8*] zeroinitializer
+// CHECK-CACHE: @"$s4main8MyStructVyxGAA0B8ProtocolAAMc" = {{.*}} @"$s4main8MyStructVyxGAA0B8ProtocolAAMcMK" {{.*}}
+// CHECK-NOCACHE-NOT: @"$s4main8MyStructVyxGAA0B8ProtocolAAMcMK"
+
+// "type metadata instantiation cache for main.Generic"
+// CHECK-CACHE: @"$s4main7GenericCMI" = internal global [{{.*}} x i8*] zeroinitializer
+// CHECK-CACHE: @"$s4main7GenericCMn" = {{.*}} @"$s4main7GenericCMI" {{.*}}
+// CHECK-NOCACHE-NOT: @"$s4main7GenericCMI"
+
+// "type metadata instantiation cache for main.MyStruct"
+// CHECK-CACHE: @"$s4main8MyStructVMI" = internal global [{{.*}} x i8*] zeroinitializer
+// CHECK-CACHE: @"$s4main8MyStructVMn" = {{.*}} @"$s4main8MyStructVMI" {{.*}} 
+// CHECK-NOCACHE-NOT: @"$s4main8MyStructVMI"