Merge pull request #58723 from aschwaighofer/transfer_prespecialize_targetFunction_earlier

GenericSpecialization: Move once initialization of pre-specializations earlier before the first query for pre-specialziations

Author: aschwaighofer

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -33,6 +33,44 @@
 using namespace swift;
 
 namespace {
+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));
+    }
+  }
+}
 
 static bool specializeAppliesInFunction(SILFunction &F,
                                         SILTransform *transform,
@@ -60,6 +98,13 @@ static bool specializeAppliesInFunction(SILFunction &F,
       auto *Callee = Apply.getReferencedFunctionOrNull();
       if (!Callee)
         continue;
+
+      FunctionBuilder.getModule().performOnceForPrespecializedImportedExtensions(
+        [&FunctionBuilder](AbstractFunctionDecl *pre) {
+        transferSpecializeAttributeTargets(FunctionBuilder.getModule(), FunctionBuilder,
+                                           pre);
+        });
+
       if (!Callee->isDefinition() && !Callee->hasPrespecialization()) {
         ORE.emit([&]() {
           using namespace OptRemark;

lib/SILOptimizer/Utils/Generics.cpp

@@ -2552,45 +2552,6 @@ 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,
@@ -2646,12 +2607,6 @@ void swift::trySpecializeApplyOfGeneric(
   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;

test/SILOptimizer/Inputs/prespecialize_import_module.swift

@@ -0,0 +1,7 @@
+public func someFunc<T>(_ t: T) {
+  print(t)
+}
+
+@_specialize(exported: true, target: someFunc(_:), where T == Int)
+@usableFromInline
+func __specialize_someFunc<T>(_: T) {}

test/SILOptimizer/prespecialize_import.swift

@@ -0,0 +1,11 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -module-name A -emit-module-path %t/A.swiftmodule %S/Inputs/prespecialize_import_module.swift
+// RUN: %target-swift-frontend -O -emit-sil -module-name B  -I %t %s | %FileCheck %s
+import A
+
+// CHECK-LABEL: sil{{.*}} @$s1B4testyyF
+public func test() {
+  // CHECK: s1A8someFuncyyxlFSi_Ts5
+  someFunc(5)
+}
+// CHECK: end sil function '$s1B4testyyF'