Merge pull request #28707 from eeckstein/cmo-fixes

Two fixes for cross module optimization

Author: eeckstein

include/swift/SILOptimizer/PassManager/Passes.def

@@ -311,6 +311,8 @@ PASS(SimplifyUnreachableContainingBlocks, "simplify-unreachable-containing-block
      "Utility pass. Removes all non-term insts from blocks with unreachable terms")
 PASS(SerializeSILPass, "serialize-sil",
      "Utility pass. Serializes the current SILModule")
+PASS(CMOSerializeSILPass, "cmo-serialize-sil",
+     "Utility pass. Serializes the current SILModule for cross-module-optimization")
 PASS(NonInlinableFunctionSkippingChecker, "check-non-inlinable-function-skipping",
      "Utility pass to ensure -experimental-skip-non-inlinable-function-bodies "
      "skips everything it should")

lib/SILOptimizer/IPO/CrossModuleSerializationSetup.cpp

@@ -49,7 +49,11 @@ class CrossModuleSerializationSetup {
     }
   }
 
-  bool setUpForSerialization(SILFunction *F);
+  bool canUseFromInline(SILFunction *F, bool lookIntoThunks);
+
+  bool canSerialize(SILFunction *F, bool lookIntoThunks);
+
+  void setUpForSerialization(SILFunction *F);
 
   void prepareInstructionForSerialization(SILInstruction *inst);
 
@@ -67,27 +71,6 @@ class CrossModuleSerializationSetup {
   void scanModule();
 };
 
-static bool canUseFromInline(SILFunction *F) {
-  if (!F)
-    return false;
-
-  switch (F->getLinkage()) {
-  case SILLinkage::PublicNonABI:
-  case SILLinkage::Shared:
-    return F->isSerialized() != IsNotSerialized;
-  case SILLinkage::Public:
-  case SILLinkage::Hidden:
-  case SILLinkage::Private:
-  case SILLinkage::PublicExternal:
-  case SILLinkage::SharedExternal:
-  case SILLinkage::PrivateExternal:
-  case SILLinkage::HiddenExternal:
-    break;
-  }
-
-  return true;
-}
-
 /// Visitor for making used types of an intruction inlinable.
 ///
 /// We use the SILCloner for visiting types, though it sucks that we allocate
@@ -253,11 +236,27 @@ prepareInstructionForSerialization(SILInstruction *inst) {
 void CrossModuleSerializationSetup::handleReferencedFunction(SILFunction *func) {
   if (!func->isDefinition() || func->isAvailableExternally())
     return;
+  if (func->getLinkage() == SILLinkage::Shared) {
+    assert(func->isThunk() != IsNotThunk &&
+      "only thunks are accepted to have shared linkage");
+    assert(canSerialize(func, /*lookIntoThunks*/ false) &&
+      "we should already have checked that the thunk is serializable");
+    
+    if (func->isSerialized() == IsSerialized)
+      return;
+
+    // We cannot make shared functions "usableFromInline", i.e. make them Public
+    // because this could result in duplicate-symbol errors. Instead we make
+    // them "@alwaysEmitIntoClient"
+    setUpForSerialization(func);
+    return;
+  }
   if (shouldSerialize(func)) {
     addToWorklistIfNotHandled(func);
-  } else {
-    makeFunctionUsableFromInline(func);
+    return;
   }
+  makeFunctionUsableFromInline(func);
+  return;
 }
 
 void CrossModuleSerializationSetup::handleReferencedMethod(SILDeclRef method) {
@@ -268,23 +267,24 @@ void CrossModuleSerializationSetup::handleReferencedMethod(SILDeclRef method) {
   M.addExternallyVisibleDecl(getBaseMethod(methodDecl));
 }
 
-/// Setup the function \p param F for serialization and put callees onto the
-/// worklist for further processing.
+/// Check if the function \p F can be serialized.
 ///
-/// Returns false in case this is not possible for some reason.
-bool CrossModuleSerializationSetup::setUpForSerialization(SILFunction *F) {
+/// If \p lookIntoThunks is true, function_ref instructions of shared
+/// thunks are also accepted.
+bool CrossModuleSerializationSetup::canSerialize(SILFunction *F,
+                                                 bool lookIntoThunks) {
   // First step: check if serializing F is even possible.
   for (SILBasicBlock &block : *F) {
     for (SILInstruction &inst : block) {
       if (auto *FRI = dyn_cast<FunctionRefBaseInst>(&inst)) {
         SILFunction *callee = FRI->getReferencedFunctionOrNull();
-        if (!canUseFromInline(callee))
+        if (!canUseFromInline(callee, lookIntoThunks))
           return false;
       } else if (auto *KPI = dyn_cast<KeyPathInst>(&inst)) {
         bool canUse = true;
         KPI->getPattern()->visitReferencedFunctionsAndMethods(
             [&](SILFunction *func) {
-              if (!canUseFromInline(func))
+              if (!canUseFromInline(func, lookIntoThunks))
                 canUse = false;
             },
             [](SILDeclRef method) { });
@@ -293,6 +293,45 @@ bool CrossModuleSerializationSetup::setUpForSerialization(SILFunction *F) {
       }
     }
   }
+  return true;
+}
+
+/// Returns true if the function \p func can be used from a serialized function.
+///
+/// If \p lookIntoThunks is true, serializable shared thunks are also accepted.
+bool CrossModuleSerializationSetup::canUseFromInline(SILFunction *func,
+                                                     bool lookIntoThunks) {
+  if (!func)
+    return false;
+
+  switch (func->getLinkage()) {
+  case SILLinkage::PublicNonABI:
+    return func->isSerialized() != IsNotSerialized;
+  case SILLinkage::Shared:
+    if (func->isThunk() != IsNotThunk && lookIntoThunks &&
+        // Don't recursively lookIntoThunks to avoid infinite loops.
+        canSerialize(func, /*lookIntoThunks*/ false)) {
+      return true;
+    }
+    return false;
+  case SILLinkage::Public:
+  case SILLinkage::Hidden:
+  case SILLinkage::Private:
+  case SILLinkage::PublicExternal:
+  case SILLinkage::SharedExternal:
+  case SILLinkage::PrivateExternal:
+  case SILLinkage::HiddenExternal:
+    break;
+  }
+  return true;
+}
+
+/// Setup the function \p param F for serialization and put callees onto the
+/// worklist for further processing.
+///
+/// Returns false in case this is not possible for some reason.
+void CrossModuleSerializationSetup::setUpForSerialization(SILFunction *F) {
+  assert(F->isSerialized() != IsSerialized);
 
   // Second step: go through all instructions and prepare them for
   // for serialization.
@@ -301,7 +340,14 @@ bool CrossModuleSerializationSetup::setUpForSerialization(SILFunction *F) {
       prepareInstructionForSerialization(&inst);
     }
   }
-  return true;
+  F->setSerialized(IsSerialized);
+
+  // As a code size optimization, make serialized functions
+  // @alwaysEmitIntoClient.
+  // Also, for shared thunks it's required to make them @alwaysEmitIntoClient.
+  // SILLinkage::Public would not work for shared functions, because it could
+  // result in duplicate-symbol linker errors.
+  F->setLinkage(SILLinkage::PublicNonABI);
 }
 
 /// Select functions in the module which should be serialized.
@@ -319,12 +365,8 @@ void CrossModuleSerializationSetup::scanModule() {
     // Decide whether we want to serialize the function.
     if (shouldSerialize(F)) {
       // Try to serialize.
-      if (setUpForSerialization(F)) {
-        F->setSerialized(IsSerialized);
-
-        // As a code size optimization, make serialized functions
-        // @alwaysEmitIntoClient.
-        F->setLinkage(SILLinkage::PublicNonABI);
+      if (canSerialize(F, /*lookIntoThunks*/ true)) {
+        setUpForSerialization(F);
       } else {
         // If for some reason the function cannot be serialized, we mark it as
         // usable-from-inline.

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -413,6 +413,12 @@ static void addPerfEarlyModulePassPipeline(SILPassPipelinePlan &P) {
   P.addOutliner();
 
   P.addCrossModuleSerializationSetup();
+  
+  // In case of cross-module-optimization, we need to serialize right after
+  // CrossModuleSerializationSetup. Eventually we want to serialize early
+  // anyway, but for now keep the SerializeSILPass at the later stage of the
+  // pipeline in case cross-module-optimization is not enabled.
+  P.addCMOSerializeSILPass();
 }
 
 static void addHighLevelEarlyLoopOptPipeline(SILPassPipelinePlan &P) {

lib/SILOptimizer/UtilityPasses/SerializeSILPass.cpp

@@ -391,6 +391,9 @@ void updateOpaqueArchetypes(SILFunction &F) {
 /// A utility pass to serialize a SILModule at any place inside the optimization
 /// pipeline.
 class SerializeSILPass : public SILModuleTransform {
+    
+  bool onlyForCrossModuleOptimization;
+    
   /// Removes [serialized] from all functions. This allows for more
   /// optimizations and for a better dead function elimination.
   void removeSerializedFlagFromAllFunctions(SILModule &M) {
@@ -422,12 +425,19 @@ class SerializeSILPass : public SILModuleTransform {
   }
 
 public:
-  SerializeSILPass() {}
+  SerializeSILPass(bool onlyForCrossModuleOptimization)
+    : onlyForCrossModuleOptimization(onlyForCrossModuleOptimization)
+  { }
+  
   void run() override {
     auto &M = *getModule();
     // Nothing to do if the module was serialized already.
     if (M.isSerialized())
       return;
+    
+    if (onlyForCrossModuleOptimization &&
+        !M.getOptions().CrossModuleOptimization)
+      return;
 
     // Mark all reachable functions as "anchors" so that they are not
     // removed later by the dead function elimination pass. This
@@ -449,5 +459,9 @@ class SerializeSILPass : public SILModuleTransform {
 };
 
 SILTransform *swift::createSerializeSILPass() {
-  return new SerializeSILPass();
+  return new SerializeSILPass(/* onlyForCrossModuleOptimization */ false);
+}
+
+SILTransform *swift::createCMOSerializeSILPass() {
+  return new SerializeSILPass(/* onlyForCrossModuleOptimization */ true);
 }

test/SILOptimizer/Inputs/cross-module.swift

@@ -236,3 +236,14 @@ public func useClassKeypath<T>(_ t: T) -> Int {
   return c[keyPath: getClassKeypath(t)]
 }
 
+@inline(never)
+func unrelated<U>(_ u: U) {
+  print(u)
+}
+
+@inline(never)
+public func callUnrelated<T>(_ t: T) -> T {
+  unrelated(43)
+  return t
+}
+

test/SILOptimizer/cross-module-optimization.swift

@@ -113,10 +113,18 @@ func testKeypath() {
   print(useClassKeypath(0))
 }
 
+func testMisc() {
+  // CHECK-OUTPUT: 43
+  // CHECK-OUTPUT: 42
+  // CHECK-SIL-DAG: sil shared {{.*}} @$s4Test13callUnrelatedyxxlFSi_Tg5
+  print(callUnrelated(42))
+}
+
 testNestedTypes()
 testClass()
 testError()
 testProtocolsAndClasses()
 testSubModule()
 testClosures()
 testKeypath()
+testMisc()