Merge pull request #40589 from eeckstein/new-cmo

rework cross-module-optimization

Author: eeckstein

include/swift/AST/TBDGenRequests.h

@@ -50,8 +50,13 @@ class TBDGenDescriptor final {
   FileOrModule Input;
   TBDGenOptions Opts;
 
-  TBDGenDescriptor(FileOrModule input, const TBDGenOptions &opts)
-      : Input(input), Opts(opts) {
+  /// Symbols (e.g. function names) which are made public by the
+  /// CrossModuleOptimization pass and therefore must be included in the TBD file.
+  TBDSymbolSetPtr publicCMOSymbols;
+
+  TBDGenDescriptor(FileOrModule input, const TBDGenOptions &opts,
+                   TBDSymbolSetPtr publicCMOSymbols)
+      : Input(input), Opts(opts), publicCMOSymbols(publicCMOSymbols) {
     assert(input);
   }
 
@@ -73,17 +78,21 @@ class TBDGenDescriptor final {
   const StringRef getDataLayoutString() const;
   const llvm::Triple &getTarget() const;
 
+  TBDSymbolSetPtr getPublicCMOSymbols() const { return publicCMOSymbols; }
+
   bool operator==(const TBDGenDescriptor &other) const;
   bool operator!=(const TBDGenDescriptor &other) const {
     return !(*this == other);
   }
 
-  static TBDGenDescriptor forFile(FileUnit *file, const TBDGenOptions &opts) {
-    return TBDGenDescriptor(file, opts);
+  static TBDGenDescriptor forFile(FileUnit *file, const TBDGenOptions &opts,
+                                  TBDSymbolSetPtr publicCMOSymbols) {
+    return TBDGenDescriptor(file, opts, publicCMOSymbols);
   }
 
-  static TBDGenDescriptor forModule(ModuleDecl *M, const TBDGenOptions &opts) {
-    return TBDGenDescriptor(M, opts);
+  static TBDGenDescriptor forModule(ModuleDecl *M, const TBDGenOptions &opts,
+                                    TBDSymbolSetPtr publicCMOSymbols) {
+    return TBDGenDescriptor(M, opts, publicCMOSymbols);
   }
 };
 
@@ -154,6 +163,9 @@ class SymbolSource {
     /// A symbol used to customize linker behavior, introduced by TBDGen.
     LinkerDirective,
 
+    /// A symbol which was made public by the CrossModuleOptimization pass.
+    CrossModuleOptimization,
+
     /// A symbol with an unknown origin.
     // FIXME: This should be eliminated.
     Unknown
@@ -173,7 +185,8 @@ class SymbolSource {
     irEntity = entity;
   }
   explicit SymbolSource(Kind kind) : kind(kind) {
-    assert(kind == Kind::LinkerDirective || kind == Kind::Unknown);
+    assert(kind == Kind::LinkerDirective || kind == Kind::Unknown ||
+           kind == Kind::CrossModuleOptimization);
   }
 
 public:
@@ -186,6 +199,9 @@ class SymbolSource {
   static SymbolSource forLinkerDirective() {
     return SymbolSource{Kind::LinkerDirective};
   }
+  static SymbolSource forCrossModuleOptimization() {
+    return SymbolSource{Kind::CrossModuleOptimization};
+  }
   static SymbolSource forUnknown() {
     return SymbolSource{Kind::Unknown};
   }
@@ -194,6 +210,10 @@ class SymbolSource {
     return kind == Kind::LinkerDirective;
   }
 
+  bool isFromCrossModuleOptimization() const {
+    return kind == Kind::CrossModuleOptimization;
+  }
+
   SILDeclRef getSILDeclRef() const {
     assert(kind == Kind::SIL);
     return silDeclRef;

include/swift/Frontend/Frontend.h

@@ -432,6 +432,7 @@ class CompilerInstance {
   std::unique_ptr<ASTContext> Context;
   std::unique_ptr<Lowering::TypeConverter> TheSILTypes;
   std::unique_ptr<DiagnosticVerifier> DiagVerifier;
+  TBDSymbolSetPtr publicCMOSymbols;
 
   /// A cache describing the set of inter-module dependencies that have been queried.
   /// Null if not present.
@@ -583,6 +584,10 @@ class CompilerInstance {
   /// file.
   SourceFile *getCodeCompletionFile() const;
 
+  /// Return the symbols (e.g. function names) which are made public by the
+  /// CrossModuleOptimization pass and therefore must be included in the TBD file.
+  TBDSymbolSetPtr getPublicCMOSymbols() const { return publicCMOSymbols; }
+
 private:
   /// Set up the file system by loading and validating all VFS overlay YAML
   /// files. If the process of validating VFS files failed, or the overlay

include/swift/SIL/SILModule.h

@@ -38,6 +38,7 @@
 #include "swift/SIL/SILVTable.h"
 #include "swift/SIL/SILWitnessTable.h"
 #include "swift/SIL/TypeLowering.h"
+#include "swift/TBDGen/TBDGen.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/ADT/MapVector.h"
@@ -367,6 +368,10 @@ class SILModule {
 
   /// Folding set for key path patterns.
   llvm::FoldingSet<KeyPathPattern> KeyPathPatterns;
+  
+  /// Symbols (e.g. function names) which are made public by the
+  /// CrossModuleOptimization pass and therefore must be included in the TBD file.
+  TBDSymbolSetPtr publicCMOSymbols;
 
 public:
   ~SILModule();
@@ -506,6 +511,12 @@ class SILModule {
 
   const SILOptions &getOptions() const { return Options; }
 
+  /// Return the symbols (e.g. function names) which are made public by the
+  /// CrossModuleOptimization pass and therefore must be included in the TBD file.
+  TBDSymbolSetPtr getPublicCMOSymbols() { return publicCMOSymbols; }
+
+  void addPublicCMOSymbol(StringRef symbol);
+
   using iterator = FunctionListType::iterator;
   using const_iterator = FunctionListType::const_iterator;
   FunctionListType &getFunctionList() { return functions; }

include/swift/SILOptimizer/PassManager/Passes.def

@@ -110,8 +110,8 @@ PASS(AccessEnforcementSelection, "access-enforcement-selection",
      "Access Enforcement Selection")
 PASS(AccessEnforcementWMO, "access-enforcement-wmo",
      "Access Enforcement Whole Module Optimization")
-PASS(CrossModuleSerializationSetup, "cross-module-serialization-setup",
-     "Setup serialization flags for cross-module optimization")
+PASS(CrossModuleOptimization, "cmo",
+     "Perform cross-module optimization")
 PASS(AccessSummaryDumper, "access-summary-dump",
      "Dump Address Parameter Access Summary")
 PASS(AccessStorageAnalysisDumper, "access-storage-analysis-dump",
@@ -405,8 +405,6 @@ 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(YieldOnceCheck, "yield-once-check",
     "Check correct usage of yields in yield-once coroutines")
 PASS(OSLogOptimization, "os-log-optimization", "Optimize os log calls")

include/swift/TBDGen/TBDGen.h

@@ -16,6 +16,9 @@
 #include "llvm/ADT/StringSet.h"
 #include "swift/Basic/Version.h"
 #include <vector>
+#include <set>
+#include <string>
+#include <memory>
 
 namespace llvm {
 class raw_ostream;
@@ -104,12 +107,26 @@ struct TBDGenOptions {
   }
 };
 
+/// Used for symbols which are made public by the CrossModuleOptimization pass
+/// and therefore must be included in the TBD file.
+///
+/// We cannot use llvm::StringSet, because we need deterministic iteration order.
+using TBDSymbolSet = std::set<std::string>;
+
+/// A pointer to TBDSymbolSet.
+///
+/// Do reference counting for memory management.
+/// This set is created in the optimizer and primarily stored in the SILModule.
+/// But then they need to be kept alive beyond the lifetime of the SILModule.
+using TBDSymbolSetPtr = std::shared_ptr<TBDSymbolSet>;
+
 std::vector<std::string> getPublicSymbols(TBDGenDescriptor desc);
 
 void writeTBDFile(ModuleDecl *M, llvm::raw_ostream &os,
-                  const TBDGenOptions &opts);
+                  const TBDGenOptions &opts, TBDSymbolSetPtr publicCMOSymbols);
 
-void writeAPIJSONFile(ModuleDecl *M, llvm::raw_ostream &os, bool PrettyPrint);
+void writeAPIJSONFile(ModuleDecl *M, llvm::raw_ostream &os, bool PrettyPrint,
+                      TBDSymbolSetPtr publicCMOSymbols);
 
 } // end namespace swift
 

lib/AST/Module.cpp

@@ -2374,17 +2374,18 @@ canBeUsedForCrossModuleOptimization(DeclContext *ctxt) const {
   // See if context is imported in a "regular" way, i.e. not with
   // @_implementationOnly or @_spi.
   ModuleDecl::ImportFilter filter = {
-    ModuleDecl::ImportFilterKind::Exported,
-    ModuleDecl::ImportFilterKind::Default};
+    ModuleDecl::ImportFilterKind::ImplementationOnly,
+    ModuleDecl::ImportFilterKind::SPIAccessControl
+  };
   SmallVector<ImportedModule, 4> results;
   getImportedModules(results, filter);
 
   auto &imports = getASTContext().getImportCache();
   for (auto &desc : results) {
     if (imports.isImportedBy(moduleOfCtxt, desc.importedModule))
-      return true;
+      return false;
   }
-  return false;
+  return true;
 }
 
 void SourceFile::lookupImportedSPIGroups(

lib/DriverTool/swift_api_extract_main.cpp

@@ -215,7 +215,7 @@ class SwiftAPIExtractInvocation {
       return 1;
 
     if (OutputFilename == "-") {
-      writeAPIJSONFile(M, llvm::outs(), PrettyPrint);
+      writeAPIJSONFile(M, llvm::outs(), PrettyPrint, nullptr);
       return 0;
     }
 
@@ -227,7 +227,7 @@ class SwiftAPIExtractInvocation {
       return 1;
     }
 
-    writeAPIJSONFile(M, OS, PrettyPrint);
+    writeAPIJSONFile(M, OS, PrettyPrint, nullptr);
     return 0;
   }
 };

lib/Frontend/Frontend.cpp

@@ -1336,6 +1336,8 @@ bool CompilerInstance::performSILProcessing(SILModule *silModule) {
 
   performSILOptimizations(Invocation, silModule);
 
+  publicCMOSymbols = silModule->getPublicCMOSymbols();
+
   if (auto *stats = getStatsReporter())
     countStatsPostSILOpt(*stats, *silModule);
 

lib/FrontendTool/FrontendTool.cpp

@@ -692,15 +692,10 @@ static bool writeTBDIfNeeded(CompilerInstance &Instance) {
     return false;
   }
 
-  if (Invocation.getSILOptions().CrossModuleOptimization) {
-    Instance.getDiags().diagnose(SourceLoc(),
-                                 diag::tbd_not_supported_with_cmo);
-    return false;
-  }
-
   const std::string &TBDPath = Invocation.getTBDPathForWholeModule();
 
-  return writeTBD(Instance.getMainModule(), TBDPath, tbdOpts);
+  return writeTBD(Instance.getMainModule(), TBDPath, tbdOpts,
+                  Instance.getPublicCMOSymbols());
 }
 
 static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
@@ -1353,19 +1348,19 @@ static bool processCommandLineAndRunImmediately(CompilerInstance &Instance,
 
 static bool validateTBDIfNeeded(const CompilerInvocation &Invocation,
                                 ModuleOrSourceFile MSF,
-                                const llvm::Module &IRModule) {
-  const auto mode = Invocation.getFrontendOptions().ValidateTBDAgainstIR;
+                                const llvm::Module &IRModule,
+                                TBDSymbolSetPtr publicCMOSymbols) {
+  auto mode = Invocation.getFrontendOptions().ValidateTBDAgainstIR;
+  if (mode == FrontendOptions::TBDValidationMode::All &&
+      Invocation.getSILOptions().CrossModuleOptimization)
+    mode = FrontendOptions::TBDValidationMode::MissingFromTBD;
+
   const bool canPerformTBDValidation = [&]() {
     // If the user has requested we skip validation, honor it.
     if (mode == FrontendOptions::TBDValidationMode::None) {
       return false;
     }
 
-    // Cross-module optimization does not support TBD.
-    if (Invocation.getSILOptions().CrossModuleOptimization) {
-      return false;
-    }
-
     // If we can't validate the given input file, bail early. This covers cases
     // like passing raw SIL as a primary file.
     const auto &IO = Invocation.getFrontendOptions().InputsAndOutputs;
@@ -1426,9 +1421,10 @@ static bool validateTBDIfNeeded(const CompilerInvocation &Invocation,
   // noise from e.g. statically-linked libraries.
   Opts.embedSymbolsFromModules.clear();
   if (auto *SF = MSF.dyn_cast<SourceFile *>()) {
-    return validateTBD(SF, IRModule, Opts, diagnoseExtraSymbolsInTBD);
+    return validateTBD(SF, IRModule, Opts, publicCMOSymbols,
+                       diagnoseExtraSymbolsInTBD);
   } else {
-    return validateTBD(MSF.get<ModuleDecl *>(), IRModule, Opts,
+    return validateTBD(MSF.get<ModuleDecl *>(), IRModule, Opts, publicCMOSymbols,
                        diagnoseExtraSymbolsInTBD);
   }
 }
@@ -1621,6 +1617,8 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
     return processCommandLineAndRunImmediately(
         Instance, std::move(SM), MSF, observer, ReturnValue);
 
+  TBDSymbolSetPtr publicCMOSymbols = SM->getPublicCMOSymbols();
+
   StringRef OutputFilename = PSPs.OutputFilename;
   std::vector<std::string> ParallelOutputFilenames =
       opts.InputsAndOutputs.copyOutputFilenames();
@@ -1635,7 +1633,8 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
   if (!IRModule)
     return Instance.getDiags().hadAnyError();
 
-  if (validateTBDIfNeeded(Invocation, MSF, *IRModule.getModule()))
+  if (validateTBDIfNeeded(Invocation, MSF, *IRModule.getModule(),
+                          publicCMOSymbols))
     return true;
 
   return generateCode(Instance, OutputFilename, IRModule.getModule(),

lib/FrontendTool/TBD.cpp

@@ -42,7 +42,7 @@ static std::vector<StringRef> sortSymbols(llvm::StringSet<> &symbols) {
 }
 
 bool swift::writeTBD(ModuleDecl *M, StringRef OutputFilename,
-                     const TBDGenOptions &Opts) {
+                     const TBDGenOptions &Opts, TBDSymbolSetPtr publicCMOSymbols) {
   std::error_code EC;
   llvm::raw_fd_ostream OS(OutputFilename, EC, llvm::sys::fs::OF_None);
   if (EC) {
@@ -51,7 +51,7 @@ bool swift::writeTBD(ModuleDecl *M, StringRef OutputFilename,
     return true;
   }
 
-  writeTBDFile(M, OS, Opts);
+  writeTBDFile(M, OS, Opts, publicCMOSymbols);
 
   return false;
 }
@@ -126,17 +126,21 @@ static bool validateSymbols(DiagnosticEngine &diags,
 bool swift::validateTBD(ModuleDecl *M,
                         const llvm::Module &IRModule,
                         const TBDGenOptions &opts,
+                        TBDSymbolSetPtr publicCMOSymbols,
                         bool diagnoseExtraSymbolsInTBD) {
-  auto symbols = getPublicSymbols(TBDGenDescriptor::forModule(M, opts));
+  auto symbols = getPublicSymbols(TBDGenDescriptor::forModule(M, opts,
+                                  publicCMOSymbols));
   return validateSymbols(M->getASTContext().Diags, symbols, IRModule,
                          diagnoseExtraSymbolsInTBD);
 }
 
 bool swift::validateTBD(FileUnit *file,
                         const llvm::Module &IRModule,
                         const TBDGenOptions &opts,
+                        TBDSymbolSetPtr publicCMOSymbols,
                         bool diagnoseExtraSymbolsInTBD) {
-  auto symbols = getPublicSymbols(TBDGenDescriptor::forFile(file, opts));
+  auto symbols = getPublicSymbols(TBDGenDescriptor::forFile(file, opts,
+                                  publicCMOSymbols));
   return validateSymbols(file->getParentModule()->getASTContext().Diags,
                          symbols, IRModule, diagnoseExtraSymbolsInTBD);
 }

lib/FrontendTool/TBD.h

@@ -14,6 +14,7 @@
 #define SWIFT_FRONTENDTOOL_TBD_H
 
 #include "swift/Frontend/FrontendOptions.h"
+#include "swift/TBDGen/TBDGen.h"
 
 namespace llvm {
 class StringRef;
@@ -26,14 +27,16 @@ class FrontendOptions;
 struct TBDGenOptions;
 
 bool writeTBD(ModuleDecl *M, StringRef OutputFilename,
-              const TBDGenOptions &Opts);
+              const TBDGenOptions &Opts, TBDSymbolSetPtr publicCMOSymbols);
 bool validateTBD(ModuleDecl *M,
                  const llvm::Module &IRModule,
                  const TBDGenOptions &opts,
+                 TBDSymbolSetPtr publicCMOSymbols,
                  bool diagnoseExtraSymbolsInTBD);
 bool validateTBD(FileUnit *M,
                  const llvm::Module &IRModule,
                  const TBDGenOptions &opts,
+                 TBDSymbolSetPtr publicCMOSymbols,
                  bool diagnoseExtraSymbolsInTBD);
 }
 

lib/IRGen/IRGen.cpp

@@ -1032,6 +1032,7 @@ getSymbolSourcesToEmit(const IRGenDescriptor &desc) {
       irEntitiesToEmit.push_back(source->getIRLinkEntity());
       break;
     case SymbolSource::Kind::LinkerDirective:
+    case SymbolSource::Kind::CrossModuleOptimization:
     case SymbolSource::Kind::Unknown:
       llvm_unreachable("Not supported");
     }