Merge pull request #39656 from aschwaighofer/delay_import_prespecializations

aschwaighofer · web-flow · commit 8866a2dcb1d7 · 2021-10-12T12:20:16.000-07:00
Delay import of prespecialized decls until generic specialization
diff --git a/include/swift/SIL/SILModule.h b/include/swift/SIL/SILModule.h
@@ -109,6 +109,7 @@ class SILUndef;
 class SourceFile;
 class SerializedSILLoader;
 class SILFunctionBuilder;
+class SILOptFunctionBuilder;
 class SILRemarkStreamer;
 
 namespace Lowering {
@@ -353,6 +354,8 @@ class SILModule {
   /// This gets set in OwnershipModelEliminator pass.
   bool regDeserializationNotificationHandlerForAllFuncOME;
 
+  bool prespecializedFunctionDeclsImported;
+
   /// Action to be executed for serializing the SILModule.
   ActionCallback SerializeSILAction;
 
@@ -894,6 +897,10 @@ class SILModule {
   /// declaration context ought to be serialized as part of this module.
   bool
   shouldSerializeEntitiesAssociatedWithDeclContext(const DeclContext *DC) const;
+
+  /// Gather prespecialized from extensions.
+  void performOnceForPrespecializedImportedExtensions(
+      llvm::function_ref<void(AbstractFunctionDecl *)> action);
 };
 
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const SILModule &M){
diff --git a/lib/SIL/IR/SILFunctionBuilder.cpp b/lib/SIL/IR/SILFunctionBuilder.cpp
@@ -184,10 +184,13 @@ SILFunction *SILFunctionBuilder::getOrCreateFunction(
     // assert.
     assert(mod.getStage() == SILStage::Raw ||
            fn->getLoweredFunctionType() == constantType);
+    auto linkageForDef = constant.getLinkage(ForDefinition_t::ForDefinition);
+    auto fnLinkage = fn->getLinkage();
     assert(mod.getStage() == SILStage::Raw || fn->getLinkage() == linkage ||
            (forDefinition == ForDefinition_t::NotForDefinition &&
-            fn->getLinkage() ==
-                constant.getLinkage(ForDefinition_t::ForDefinition)));
+            (fnLinkage == linkageForDef ||
+             (linkageForDef == SILLinkage::PublicNonABI &&
+              fnLinkage == SILLinkage::SharedExternal))));
     if (forDefinition) {
       // In all the cases where getConstantLinkage returns something
       // different for ForDefinition, it returns an available-externally
diff --git a/lib/SIL/IR/SILModule.cpp b/lib/SIL/IR/SILModule.cpp
@@ -98,7 +98,8 @@ SILModule::SILModule(llvm::PointerUnion<FileUnit *, ModuleDecl *> context,
       Options(Options), serialized(false),
       regDeserializationNotificationHandlerForNonTransparentFuncOME(false),
       regDeserializationNotificationHandlerForAllFuncOME(false),
-      SerializeSILAction(), Types(TC) {
+      prespecializedFunctionDeclsImported(false), SerializeSILAction(),
+      Types(TC) {
   assert(!context.isNull());
   if (auto *file = context.dyn_cast<FileUnit *>()) {
     AssociatedDeclContext = file;
@@ -874,6 +875,42 @@ void SILModule::installSILRemarkStreamer() {
 bool SILModule::isStdlibModule() const {
   return TheSwiftModule->isStdlibModule();
 }
+void SILModule::performOnceForPrespecializedImportedExtensions(
+    llvm::function_ref<void(AbstractFunctionDecl *)> action) {
+  if (prespecializedFunctionDeclsImported)
+    return;
+
+  SmallVector<ModuleDecl *, 8> importedModules;
+  // Add the Swift module.
+  if (!isStdlibModule()) {
+    auto *SwiftStdlib = getASTContext().getStdlibModule(true);
+    if (SwiftStdlib)
+      importedModules.push_back(SwiftStdlib);
+  }
+
+  // Add explicitly imported modules.
+  SmallVector<Decl *, 32> topLevelDecls;
+  getSwiftModule()->getTopLevelDecls(topLevelDecls);
+  for (const Decl *D : topLevelDecls) {
+    if (auto importDecl = dyn_cast<ImportDecl>(D)) {
+      if (!importDecl->getModule() ||
+          importDecl->getModule()->isNonSwiftModule())
+        continue;
+      importedModules.push_back(importDecl->getModule());
+    }
+  }
+
+  for (auto *module : importedModules) {
+    SmallVector<Decl *, 16> prespecializations;
+    module->getExportedPrespecializations(prespecializations);
+    for (auto *p : prespecializations) {
+      if (auto *vd = dyn_cast<AbstractFunctionDecl>(p)) {
+        action(vd);
+      }
+    }
+  }
+  prespecializedFunctionDeclsImported = true;
+}
 
 SILProperty *SILProperty::create(SILModule &M,
                                  bool Serialized,
diff --git a/lib/SILGen/SILGen.cpp b/lib/SILGen/SILGen.cpp
@@ -2203,52 +2203,3 @@ swift::performASTLowering(FileUnit &sf, Lowering::TypeConverter &tc,
   auto desc = ASTLoweringDescriptor::forFile(sf, tc, options);
   return llvm::cantFail(sf.getASTContext().evaluator(ASTLoweringRequest{desc}));
 }
-
-static void transferSpecializeAttributeTargets(SILGenModule &SGM, SILModule &M,
-                                               Decl *d) {
-  auto *vd = cast<AbstractFunctionDecl>(d);
-  for (auto *A : vd->getAttrs().getAttributes<SpecializeAttr>()) {
-    auto *SA = cast<SpecializeAttr>(A);
-    // Filter _spi.
-    auto spiGroups = SA->getSPIGroups();
-    auto hasSPIGroup = !spiGroups.empty();
-    if (hasSPIGroup) {
-      if (vd->getModuleContext() != M.getSwiftModule() &&
-          !M.getSwiftModule()->isImportedAsSPI(SA, vd)) {
-        continue;
-      }
-    }
-    if (auto *targetFunctionDecl = SA->getTargetFunctionDecl(vd)) {
-      auto target = SILDeclRef(targetFunctionDecl);
-      auto targetSILFunction = SGM.getFunction(target, NotForDefinition);
-      auto kind = SA->getSpecializationKind() ==
-                          SpecializeAttr::SpecializationKind::Full
-                      ? SILSpecializeAttr::SpecializationKind::Full
-                      : SILSpecializeAttr::SpecializationKind::Partial;
-      Identifier spiGroupIdent;
-      if (hasSPIGroup) {
-        spiGroupIdent = spiGroups[0];
-      }
-      auto availability =
-        AvailabilityInference::annotatedAvailableRangeForAttr(SA,
-                                                              M.getSwiftModule()->getASTContext());
-      targetSILFunction->addSpecializeAttr(SILSpecializeAttr::create(
-          M, SA->getSpecializedSignature(), SA->isExported(), kind, nullptr,
-          spiGroupIdent, vd->getModuleContext(), availability));
-    }
-  }
-}
-
-void SILGenModule::visitImportDecl(ImportDecl *import) {
-  auto *module = import->getModule();
-  if (module->isNonSwiftModule())
-    return;
-
-  SmallVector<Decl*, 16> prespecializations;
-  module->getExportedPrespecializations(prespecializations);
-  for (auto *p : prespecializations) {
-    if (auto *vd = dyn_cast<AbstractFunctionDecl>(p)) {
-      transferSpecializeAttributeTargets(*this, M, vd);
-    }
-  }
-}
diff --git a/lib/SILGen/SILGen.h b/lib/SILGen/SILGen.h
@@ -265,7 +265,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
 
   // These are either not allowed at global scope or don't require
   // code emission.
-  void visitImportDecl(ImportDecl *d);
+  void visitImportDecl(ImportDecl *d) {}
   void visitEnumCaseDecl(EnumCaseDecl *d) {}
   void visitEnumElementDecl(EnumElementDecl *d) {}
   void visitOperatorDecl(OperatorDecl *d) {}
diff --git a/lib/SILOptimizer/Utils/Generics.cpp b/lib/SILOptimizer/Utils/Generics.cpp
@@ -13,6 +13,7 @@
 #define DEBUG_TYPE "generic-specializer"
 
 #include "swift/SILOptimizer/Utils/Generics.h"
+#include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/GenericEnvironment.h"
@@ -2520,6 +2521,45 @@ usePrespecialized(SILOptFunctionBuilder &funcBuilder, ApplySite apply,
   return nullptr;
 }
 
+static void transferSpecializeAttributeTargets(SILModule &M,
+                                               SILOptFunctionBuilder &builder,
+                                               Decl *d) {
+  auto *vd = cast<AbstractFunctionDecl>(d);
+  for (auto *A : vd->getAttrs().getAttributes<SpecializeAttr>()) {
+    auto *SA = cast<SpecializeAttr>(A);
+    // Filter _spi.
+    auto spiGroups = SA->getSPIGroups();
+    auto hasSPIGroup = !spiGroups.empty();
+    if (hasSPIGroup) {
+      if (vd->getModuleContext() != M.getSwiftModule() &&
+          !M.getSwiftModule()->isImportedAsSPI(SA, vd)) {
+        continue;
+      }
+    }
+    if (auto *targetFunctionDecl = SA->getTargetFunctionDecl(vd)) {
+      auto target = SILDeclRef(targetFunctionDecl);
+      auto targetSILFunction = builder.getOrCreateFunction(
+          SILLocation(vd), target, NotForDefinition,
+          [&builder](SILLocation loc, SILDeclRef constant) -> SILFunction * {
+            return builder.getOrCreateFunction(loc, constant, NotForDefinition);
+          });
+      auto kind = SA->getSpecializationKind() ==
+                          SpecializeAttr::SpecializationKind::Full
+                      ? SILSpecializeAttr::SpecializationKind::Full
+                      : SILSpecializeAttr::SpecializationKind::Partial;
+      Identifier spiGroupIdent;
+      if (hasSPIGroup) {
+        spiGroupIdent = spiGroups[0];
+      }
+      auto availability = AvailabilityInference::annotatedAvailableRangeForAttr(
+          SA, M.getSwiftModule()->getASTContext());
+      targetSILFunction->addSpecializeAttr(SILSpecializeAttr::create(
+          M, SA->getSpecializedSignature(), SA->isExported(), kind, nullptr,
+          spiGroupIdent, vd->getModuleContext(), availability));
+    }
+  }
+}
+
 void swift::trySpecializeApplyOfGeneric(
     SILOptFunctionBuilder &FuncBuilder,
     ApplySite Apply, DeadInstructionSet &DeadApplies,
@@ -2573,6 +2613,13 @@ void swift::trySpecializeApplyOfGeneric(
   // Check if there is a pre-specialization available in a library.
   SILFunction *prespecializedF = nullptr;
   ReabstractionInfo prespecializedReInfo;
+
+  FuncBuilder.getModule().performOnceForPrespecializedImportedExtensions(
+      [&FuncBuilder](AbstractFunctionDecl *pre) {
+        transferSpecializeAttributeTargets(FuncBuilder.getModule(), FuncBuilder,
+                                           pre);
+      });
+
   if ((prespecializedF = usePrespecialized(FuncBuilder, Apply, RefF, ReInfo,
                                            prespecializedReInfo))) {
     ReInfo = prespecializedReInfo;
diff --git a/test/IRGen/dllimport.swift b/test/IRGen/dllimport.swift
@@ -1,5 +1,5 @@
-// RUN: %swift -target thumbv7--windows-itanium -emit-ir -parse-as-library -disable-legacy-type-info -parse-stdlib -module-name dllimport %s -o - -enable-source-import -I %S | %FileCheck %s -check-prefix CHECK -check-prefix CHECK-NO-OPT
-// RUN: %swift -target thumbv7--windows-itanium -O -emit-ir -parse-as-library -disable-legacy-type-info -parse-stdlib -module-name dllimport -primary-file %s -o - -enable-source-import -I %S | %FileCheck %s -check-prefix CHECK -check-prefix CHECK-OPT
+// RUN: %swift -Xllvm -sil-disable-pass=GenericSpecializer -target thumbv7--windows-itanium -emit-ir -parse-as-library -disable-legacy-type-info -parse-stdlib -module-name dllimport %s -o - -enable-source-import -I %S | %FileCheck %s -check-prefix CHECK -check-prefix CHECK-NO-OPT
+// RUN: %swift -Xllvm -sil-disable-pass=GenericSpecializer -target thumbv7--windows-itanium -O -emit-ir -parse-as-library -disable-legacy-type-info -parse-stdlib -module-name dllimport -primary-file %s -o - -enable-source-import -I %S | %FileCheck %s -check-prefix CHECK -check-prefix CHECK-OPT
 
 // REQUIRES: CODEGENERATOR=ARM
 
diff --git a/test/SILGen/specialize_attr.swift b/test/SILGen/specialize_attr.swift
@@ -3,7 +3,7 @@
 // RUN: %target-swift-frontend -module-name A -emit-module-path %t/A.swiftmodule -enable-library-evolution -swift-version 5 %S/Inputs/specialize_attr_module.swift
 // RUN: %target-swift-frontend -I %t -module-name B -emit-module-path %t/B.swiftmodule -enable-library-evolution -swift-version 5 %S/Inputs/specialize_attr_module2.swift
 // RUN: %target-swift-emit-silgen -I %t -module-name specialize_attr -emit-verbose-sil %s -swift-version 5 | %FileCheck %s -check-prefix=CHECK -check-prefix=CHECK-%target-os
-// RUN: %target-swift-emit-sil -I %t -sil-verify-all -O -module-name specialize_attr -emit-verbose-sil %s | %FileCheck -check-prefix=CHECK-OPT -check-prefix=CHECK-OPT-EVO %s
+// RUN: %target-swift-emit-sil -I %t -sil-verify-all -O -module-name specialize_attr -emit-verbose-sil %s | %FileCheck -check-prefix=CHECK-OPT -check-prefix=CHECK-OPT-EVO -check-prefix=CHECK-OPT-%target-os %s
 
 // Test .swiftinterface
 // RUN: %empty-directory(%t)
@@ -124,15 +124,15 @@ public struct CC2<T : PP> {
 }
 
 // Import from module A and B.
-// CHECK-macosx-LABEL: sil [serialized] [_specialize exported: true, kind: full, where T == Double] [_specialize exported: true, kind: full, available: 10.50, where T == Int16] [_specialize exported: true, kind: full, where T == Int] @$s1A11PublicThingV11doStuffWithyyxF : $@convention(method) <τ_0_0> (@in_guaranteed τ_0_0, PublicThing<τ_0_0>) -> ()
-// CHECK-linux-gnu-LABEL: sil [serialized] [_specialize exported: true, kind: full, where T == Double] [_specialize exported: true, kind: full, where T == Int16] [_specialize exported: true, kind: full, where T == Int] @$s1A11PublicThingV11doStuffWithyyxF : $@convention(method) <τ_0_0> (@in_guaranteed τ_0_0, PublicThing<τ_0_0>) -> ()
-// CHECK-LABEL: sil [serialized] [_specialize exported: true, kind: full, where T == Double] [_specialize exported: true, kind: full, where T == Int] @$s1A13InternalThingV11doStuffWith5boxedyAA05BoxedB0VyxG_tF : $@convention(method) <τ_0_0> (BoxedThing<τ_0_0>, InternalThing<τ_0_0>) -> ()
+// CHECK-OPT-macosx-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Double] [_specialize exported: true, kind: full, available: 10.50, where T == Int16] [_specialize exported: true, kind: full, where T == Int] @$s1A11PublicThingV11doStuffWithyyxF : $@convention(method) <T> (@in_guaranteed T, PublicThing<T>) -> ()
+// CHECK-OPT-linux-gnu-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Double] [_specialize exported: true, kind: full, where T == Int16] [_specialize exported: true, kind: full, where T == Int] @$s1A11PublicThingV11doStuffWithyyxF : $@convention(method) <T> (@in_guaranteed T, PublicThing<T>) -> ()
+// CHECK-OPT-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Double] [_specialize exported: true, kind: full, where T == Int] @$s1A13InternalThingV11doStuffWith5boxedyAA05BoxedB0VyxG_tF : $@convention(method) <T> (BoxedThing<T>, InternalThing<T>) -> ()
 
-// CHECK-DAG: sil [serialized] [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2V9computedZxvM : $@yield_once @convention(method) <τ_0_0> (@inout InternalThing2<τ_0_0>) -> @yields @inout τ_0_0
-// CHECK-DAG: sil [serialized] [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2V9computedZxvr : $@yield_once @convention(method) <τ_0_0> (@in_guaranteed InternalThing2<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+// CHECK-OPT-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2V9computedZxvM : $@yield_once @convention(method) <T> (@inout InternalThing2<T>) -> @yields @inout T
+// CHECK-OPT-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2V9computedZxvr : $@yield_once @convention(method) <T> (@in_guaranteed InternalThing2<T>) -> @yields @in_guaranteed T
 
-// CHECK-DAG: sil [serialized] [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2VyxSicig : $@convention(method) <τ_0_0> (Int, @in_guaranteed InternalThing2<τ_0_0>) -> @out τ_0_0
-// CHECK-DAG: sil [serialized] [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2VyxSicis : $@convention(method) <τ_0_0> (@in τ_0_0, Int, @inout InternalThing2<τ_0_0>) -> ()
+// CHECK-OPT-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2VyxSicig : $@convention(method) <T> (Int, @in_guaranteed InternalThing2<T>) -> @out T
+// CHECK-OPT-DAG: sil {{.*}} [_specialize exported: true, kind: full, where T == Int] @$s1A14InternalThing2VyxSicis : $@convention(method) <T> (@in T, Int, @inout InternalThing2<T>) -> ()
 
 // CHECK-LABEL: sil [_specialize exported: false, kind: full, where T == Klass1, U == FakeString] [_specialize exported: false, kind: full, where T == Int, U == Float] [ossa] @$s15specialize_attr0A4This_1uyx_q_tr0_lF : $@convention(thin) <T, U> (@in_guaranteed T, @in_guaranteed U) -> () {
 
diff --git a/test/SILOptimizer/pre_specialize-macos.swift b/test/SILOptimizer/pre_specialize-macos.swift
@@ -9,9 +9,6 @@
 import pre_specialized_module
 import pre_specialized_module2
 
-// OPT: sil [available 10.50] [noinline] @$s22pre_specialized_module21publicPrespecialized2yyxlFAA8SomeDataV_Ts5 : $@convention(thin) (SomeData) -> ()
-// OPT: sil [available 10.50] [noinline] @$s22pre_specialized_module21publicPrespecialized2yyxlF0a1_B8_module213SomeOtherDataV_Ts5 : $@convention(thin) (SomeOtherData) -> ()
-
 // CHECK: sil @$s4main28usePrespecializedEntryPointsyyF : $@convention(thin) () -> () {
 // OPT:  [[F1:%.*]] = function_ref @$s22pre_specialized_module21publicPrespecialized2yyxlFAA8SomeDataV_Ts5 : $@convention(thin) (SomeData) -> ()
 // OPT:  apply [[F1]](
@@ -28,3 +25,6 @@ public func usePrespecializedEntryPoints() {
   publicPrespecialized2(SomeData())
   publicPrespecialized2(SomeOtherData())
 }
+
+// OPT: sil [available 10.50] [noinline] @$s22pre_specialized_module21publicPrespecialized2yyxlFAA8SomeDataV_Ts5 : $@convention(thin) (SomeData) -> ()
+// OPT: sil [available 10.50] [noinline] @$s22pre_specialized_module21publicPrespecialized2yyxlF0a1_B8_module213SomeOtherDataV_Ts5 : $@convention(thin) (SomeOtherData) -> ()
diff --git a/test/SILOptimizer/specialize_checked_cast_branch.swift b/test/SILOptimizer/specialize_checked_cast_branch.swift
@@ -338,4 +338,5 @@ _ = ExistentialToArchetypeCast(o: o, t: b)
 // CHECK-NOT: checked_cast_br %
 // CHECK: return %0 : $AnyObject
 // CHECK-NOT: checked_cast_br %
+// CHECK: } // end sil function '$s30specialize_checked_cast_branch26ExistentialToArchetypeCast1o1txyXl_xtlFyXl_Tg5'
 _ = ExistentialToArchetypeCast(o: o, t: o)
