[Caching] Fix multi-threaded WMO with MCCAS

include/swift/AST/IRGenRequests.h

@@ -33,6 +33,10 @@ class SILOptions;
 struct TBDGenOptions;
 class TBDGenDescriptor;
 
+namespace cas {
+  class SwiftCASOutputBackend;
+}
+
 namespace irgen {
   class IRGenModule;
 }
@@ -149,12 +153,15 @@ struct IRGenDescriptor {
 
   StringRef ModuleName;
   const PrimarySpecificPaths &PSPs;
+  std::shared_ptr<llvm::cas::ObjectStore> CAS;
   StringRef PrivateDiscriminator;
   ArrayRef<std::string> parallelOutputFilenames;
   ArrayRef<std::string> parallelIROutputFilenames;
   llvm::GlobalVariable **outModuleHash;
+  swift::cas::SwiftCASOutputBackend *casBackend = nullptr;
   llvm::raw_pwrite_stream *out = nullptr;
 
+
   friend llvm::hash_code hash_value(const IRGenDescriptor &owner) {
     return llvm::hash_combine(owner.Ctx, owner.SymbolsToEmit, owner.SILMod);
   }
@@ -176,8 +183,10 @@ struct IRGenDescriptor {
           const TBDGenOptions &TBDOpts, const SILOptions &SILOpts,
           Lowering::TypeConverter &Conv, std::unique_ptr<SILModule> &&SILMod,
           StringRef ModuleName, const PrimarySpecificPaths &PSPs,
+          std::shared_ptr<llvm::cas::ObjectStore> CAS,
           StringRef PrivateDiscriminator, SymsToEmit symsToEmit = std::nullopt,
-          llvm::GlobalVariable **outModuleHash = nullptr) {
+          llvm::GlobalVariable **outModuleHash = nullptr,
+          cas::SwiftCASOutputBackend *casBackend = nullptr) {
     return IRGenDescriptor{file,
                            symsToEmit,
                            Opts,
@@ -187,20 +196,26 @@ struct IRGenDescriptor {
                            SILMod.release(),
                            ModuleName,
                            PSPs,
+                           std::move(CAS),
                            PrivateDiscriminator,
                            {},
                            {},
-                           outModuleHash};
+                           outModuleHash,
+                           casBackend};
   }
 
-  static IRGenDescriptor forWholeModule(
-      ModuleDecl *M, const IRGenOptions &Opts, const TBDGenOptions &TBDOpts,
-      const SILOptions &SILOpts, Lowering::TypeConverter &Conv,
-      std::unique_ptr<SILModule> &&SILMod, StringRef ModuleName,
-      const PrimarySpecificPaths &PSPs, SymsToEmit symsToEmit = std::nullopt,
-      ArrayRef<std::string> parallelOutputFilenames = {},
-      ArrayRef<std::string> parallelIROutputFilenames = {},
-      llvm::GlobalVariable **outModuleHash = nullptr) {
+  static IRGenDescriptor
+  forWholeModule(ModuleDecl *M, const IRGenOptions &Opts,
+                 const TBDGenOptions &TBDOpts, const SILOptions &SILOpts,
+                 Lowering::TypeConverter &Conv,
+                 std::unique_ptr<SILModule> &&SILMod, StringRef ModuleName,
+                 const PrimarySpecificPaths &PSPs,
+                 std::shared_ptr<llvm::cas::ObjectStore> CAS,
+                 SymsToEmit symsToEmit = std::nullopt,
+                 ArrayRef<std::string> parallelOutputFilenames = {},
+                 ArrayRef<std::string> parallelIROutputFilenames = {},
+                 llvm::GlobalVariable **outModuleHash = nullptr,
+                 cas::SwiftCASOutputBackend *casBackend = nullptr) {
     return IRGenDescriptor{M,
                            symsToEmit,
                            Opts,
@@ -210,10 +225,12 @@ struct IRGenDescriptor {
                            SILMod.release(),
                            ModuleName,
                            PSPs,
+                           std::move(CAS),
                            "",
                            parallelOutputFilenames,
                            parallelIROutputFilenames,
-                           outModuleHash};
+                           outModuleHash,
+                           casBackend};
   }
 
   /// Retrieves the files to perform IR generation for. If the descriptor is

include/swift/Subsystems.h

@@ -80,6 +80,10 @@ namespace swift {
     class TypeConverter;
   }
 
+  namespace cas {
+    class SwiftCASOutputBackend;
+  }
+
   namespace fine_grained_dependencies {
     class SourceFileDepGraph;
   }
@@ -243,7 +247,8 @@ namespace swift {
   /// Get the CPU, subtarget feature options, and triple to use when emitting code.
   std::tuple<llvm::TargetOptions, std::string, std::vector<std::string>,
              std::string>
-  getIRTargetOptions(const IRGenOptions &Opts, ASTContext &Ctx);
+  getIRTargetOptions(const IRGenOptions &Opts, ASTContext &Ctx,
+                     std::shared_ptr<llvm::cas::ObjectStore> CAS = nullptr);
 
   /// Turn the given Swift module into LLVM IR and return the generated module.
   /// To compile and output the generated code, call \c performLLVM.
@@ -252,19 +257,23 @@ namespace swift {
                       const TBDGenOptions &TBDOpts,
                       std::unique_ptr<SILModule> SILMod, StringRef ModuleName,
                       const PrimarySpecificPaths &PSPs,
+                      std::shared_ptr<llvm::cas::ObjectStore> CAS,
                       ArrayRef<std::string> parallelOutputFilenames,
                       ArrayRef<std::string> parallelIROutputFilenames,
-                      llvm::GlobalVariable **outModuleHash = nullptr);
+                      llvm::GlobalVariable **outModuleHash = nullptr,
+                      cas::SwiftCASOutputBackend *casBackend = nullptr);
 
   /// Turn the given Swift file into LLVM IR and return the generated module.
   /// To compile and output the generated code, call \c performLLVM.
   GeneratedModule
-  performIRGeneration(FileUnit *file, const IRGenOptions &Opts, 
+  performIRGeneration(FileUnit *file, const IRGenOptions &Opts,
                       const TBDGenOptions &TBDOpts,
-                      std::unique_ptr<SILModule> SILMod,
-                      StringRef ModuleName, const PrimarySpecificPaths &PSPs,
+                      std::unique_ptr<SILModule> SILMod, StringRef ModuleName,
+                      const PrimarySpecificPaths &PSPs,
+                      std::shared_ptr<llvm::cas::ObjectStore> CAS,
                       StringRef PrivateDiscriminator,
-                      llvm::GlobalVariable **outModuleHash = nullptr);
+                      llvm::GlobalVariable **outModuleHash = nullptr,
+                      cas::SwiftCASOutputBackend *casBackend = nullptr);
 
   /// Given an already created LLVM module, construct a pass pipeline and run
   /// the Swift LLVM Pipeline upon it. This will include the emission of LLVM IR
@@ -330,7 +339,8 @@ namespace swift {
 
   /// Creates a TargetMachine from the IRGen opts and AST Context.
   std::unique_ptr<llvm::TargetMachine>
-  createTargetMachine(const IRGenOptions &Opts, ASTContext &Ctx);
+  createTargetMachine(const IRGenOptions &Opts, ASTContext &Ctx,
+                      std::shared_ptr<llvm::cas::ObjectStore> CAS);
 
   /// A convenience wrapper for Parser functionality.
   class ParserUnit {

lib/DriverTool/sil_llvm_gen_main.cpp

@@ -456,12 +456,13 @@ int sil_llvm_gen_main(ArrayRef<const char *> argv, void *MainAddr) {
     if (options.PerformWMO) {
       return IRGenDescriptor::forWholeModule(
           mod, Opts, TBDOpts, SILOpts, SILTypes,
-          /*SILMod*/ nullptr, moduleName, PSPs);
+          /*SILMod*/ nullptr, moduleName, PSPs, /*CAS=*/nullptr);
     }
 
-    return IRGenDescriptor::forFile(
-        mod->getFiles()[0], Opts, TBDOpts, SILOpts, SILTypes,
-        /*SILMod*/ nullptr, moduleName, PSPs, /*discriminator*/ "");
+    return IRGenDescriptor::forFile(mod->getFiles()[0], Opts, TBDOpts, SILOpts,
+                                    SILTypes,
+                                    /*SILMod*/ nullptr, moduleName, PSPs,
+                                    /*CAS=*/nullptr, /*discriminator*/ "");
   };
 
   auto &eval = CI.getASTContext().evaluator;

lib/FrontendTool/FrontendTool.cpp

@@ -1785,21 +1785,22 @@ static bool serializeModuleSummary(SILModule *SM,
 static GeneratedModule
 generateIR(const IRGenOptions &IRGenOpts, const TBDGenOptions &TBDOpts,
            std::unique_ptr<SILModule> SM, const PrimarySpecificPaths &PSPs,
+           std::shared_ptr<llvm::cas::ObjectStore> CAS,
+           cas::SwiftCASOutputBackend *casBackend,
            StringRef OutputFilename, ModuleOrSourceFile MSF,
            llvm::GlobalVariable *&HashGlobal,
            ArrayRef<std::string> parallelOutputFilenames,
            ArrayRef<std::string> parallelIROutputFilenames) {
-  if (auto *SF = MSF.dyn_cast<SourceFile *>()) {
-    return performIRGeneration(SF, IRGenOpts, TBDOpts,
-                               std::move(SM), OutputFilename, PSPs,
-                               SF->getPrivateDiscriminator().str(),
-                               &HashGlobal);
-  } else {
-    return performIRGeneration(cast<ModuleDecl *>(MSF), IRGenOpts, TBDOpts,
-                               std::move(SM), OutputFilename, PSPs,
-                               parallelOutputFilenames,
-                               parallelIROutputFilenames, &HashGlobal);
-  }
+  if (auto *SF = MSF.dyn_cast<SourceFile *>())
+    return performIRGeneration(SF, IRGenOpts, TBDOpts, std::move(SM),
+                               OutputFilename, PSPs, std::move(CAS),
+                               SF->getPrivateDiscriminator().str(), &HashGlobal,
+                               casBackend);
+
+  return performIRGeneration(
+      cast<ModuleDecl *>(MSF), IRGenOpts, TBDOpts, std::move(SM),
+      OutputFilename, PSPs, std::move(CAS), parallelOutputFilenames,
+      parallelIROutputFilenames, &HashGlobal, casBackend);
 }
 
 static bool processCommandLineAndRunImmediately(CompilerInstance &Instance,
@@ -1950,10 +1951,8 @@ static bool generateCode(CompilerInstance &Instance, StringRef OutputFilename,
                          llvm::Module *IRModule,
                          llvm::GlobalVariable *HashGlobal) {
   const auto &opts = Instance.getInvocation().getIRGenOptions();
-  std::unique_ptr<llvm::TargetMachine> TargetMachine =
-      createTargetMachine(opts, Instance.getASTContext());
-
-  TargetMachine->Options.MCOptions.CAS = Instance.getSharedCASInstance();
+  std::unique_ptr<llvm::TargetMachine> TargetMachine = createTargetMachine(
+      opts, Instance.getASTContext(), Instance.getSharedCASInstance());
 
   if (Instance.getInvocation().getCASOptions().EnableCaching &&
       opts.UseCASBackend)
@@ -2162,10 +2161,14 @@ static bool performCompileStepsPostSILGen(
                                       options::OPT_ir_output_path);
 
   llvm::GlobalVariable *HashGlobal;
-  auto IRModule =
-      generateIR(IRGenOpts, Invocation.getTBDGenOptions(), std::move(SM), PSPs,
-                 OutputFilename, MSF, HashGlobal, ParallelOutputFilenames,
-                 ParallelIROutputFilenames);
+  cas::SwiftCASOutputBackend *casBackend =
+      Invocation.getCASOptions().EnableCaching && IRGenOpts.UseCASBackend
+          ? &Instance.getCASOutputBackend()
+          : nullptr;
+  auto IRModule = generateIR(
+      IRGenOpts, Invocation.getTBDGenOptions(), std::move(SM), PSPs,
+      Instance.getSharedCASInstance(), casBackend, OutputFilename, MSF,
+      HashGlobal, ParallelOutputFilenames, ParallelIROutputFilenames);
 
   // Write extra LLVM IR output if requested
   if (IRModule && !PSPs.SupplementaryOutputs.LLVMIROutputPath.empty()) {

lib/IRGen/IRGen.cpp

@@ -129,7 +129,8 @@ static cl::opt<bool> AlignModuleToPageSize(
 
 std::tuple<llvm::TargetOptions, std::string, std::vector<std::string>,
            std::string>
-swift::getIRTargetOptions(const IRGenOptions &Opts, ASTContext &Ctx) {
+swift::getIRTargetOptions(const IRGenOptions &Opts, ASTContext &Ctx,
+                          std::shared_ptr<llvm::cas::ObjectStore> CAS) {
   // Things that maybe we should collect from the command line:
   //   - relocation model
   //   - code model
@@ -150,6 +151,9 @@ swift::getIRTargetOptions(const IRGenOptions &Opts, ASTContext &Ctx) {
   // Set option to select the CASBackendMode.
   TargetOpts.MCOptions.CASObjMode = Opts.CASObjMode;
 
+  // Set CAS and CASID callbacks.
+  TargetOpts.MCOptions.CAS = std::move(CAS);
+
   auto *Clang = static_cast<ClangImporter *>(Ctx.getClangModuleLoader());
 
   // Set UseInitArray appropriately.
@@ -1104,7 +1108,8 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
 }
 
 std::unique_ptr<llvm::TargetMachine>
-swift::createTargetMachine(const IRGenOptions &Opts, ASTContext &Ctx) {
+swift::createTargetMachine(const IRGenOptions &Opts, ASTContext &Ctx,
+                           std::shared_ptr<llvm::cas::ObjectStore> CAS) {
   CodeGenOptLevel OptLevel = Opts.shouldOptimize()
                                    ? CodeGenOptLevel::Default // -Os
                                    : CodeGenOptLevel::None;
@@ -1114,8 +1119,8 @@ swift::createTargetMachine(const IRGenOptions &Opts, ASTContext &Ctx) {
   std::string CPU;
   std::string EffectiveClangTriple;
   std::vector<std::string> targetFeaturesArray;
-  std::tie(TargetOpts, CPU, targetFeaturesArray, EffectiveClangTriple)
-    = getIRTargetOptions(Opts, Ctx);
+  std::tie(TargetOpts, CPU, targetFeaturesArray, EffectiveClangTriple) =
+      getIRTargetOptions(Opts, Ctx, std::move(CAS));
   const llvm::Triple &EffectiveTriple = llvm::Triple(EffectiveClangTriple);
   std::string targetFeatures;
   if (!targetFeaturesArray.empty()) {
@@ -1168,12 +1173,12 @@ swift::createTargetMachine(const IRGenOptions &Opts, ASTContext &Ctx) {
   return std::unique_ptr<llvm::TargetMachine>(TargetMachine);
 }
 
-IRGenerator::IRGenerator(const IRGenOptions &options, SILModule &module)
-  : Opts(options), SIL(module), QueueIndex(0) {
-}
+IRGenerator::IRGenerator(const IRGenOptions &options, SILModule &module,
+                         std::shared_ptr<llvm::cas::ObjectStore> CAS)
+    : Opts(options), SIL(module), CAS(std::move(CAS)), QueueIndex(0) {}
 
 std::unique_ptr<llvm::TargetMachine> IRGenerator::createTargetMachine() {
-  return ::createTargetMachine(Opts, SIL.getASTContext());
+  return ::createTargetMachine(Opts, SIL.getASTContext(), CAS);
 }
 
 // With -embed-bitcode, save a copy of the llvm IR as data in the
@@ -1479,7 +1484,7 @@ GeneratedModule IRGenRequest::evaluate(Evaluator &evaluator,
   auto *primaryFile =
       dyn_cast_or_null<SourceFile>(desc.Ctx.dyn_cast<FileUnit *>());
 
-  IRGenerator irgen(Opts, *SILMod);
+  IRGenerator irgen(Opts, *SILMod, desc.CAS);
 
   auto targetMachine = irgen.createTargetMachine();
   if (!targetMachine) return GeneratedModule::null();
@@ -1687,7 +1692,7 @@ static void performParallelIRGeneration(IRGenDescriptor desc) {
   auto SILMod = std::unique_ptr<SILModule>(desc.SILMod);
   auto *M = desc.getParentModule();
 
-  IRGenerator irgen(Opts, *SILMod);
+  IRGenerator irgen(Opts, *SILMod, desc.CAS);
 
   // Enter a cleanup to delete all the IGMs and their associated LLVMContexts
   // that have been associated with the IRGenerator.
@@ -1726,6 +1731,16 @@ static void performParallelIRGeneration(IRGenDescriptor desc) {
 
     // Create the IR emitter.
     auto outputName = *OutputIter++;
+    if (desc.casBackend) {
+      targetMachine->Options.MCOptions.ResultCallBack =
+          [=](const llvm::cas::CASID &id) -> llvm::Error {
+        if (auto Err = desc.casBackend->storeMCCASObjectID(outputName, id))
+          return Err;
+
+        return llvm::Error::success();
+      };
+    }
+
     IRGenModule *IGM = new IRGenModule(
         irgen, std::move(targetMachine), nextSF, desc.ModuleName, outputName,
         nextSF->getFilename(), nextSF->getPrivateDiscriminator().str());
@@ -1929,16 +1944,19 @@ GeneratedModule swift::performIRGeneration(
     swift::ModuleDecl *M, const IRGenOptions &Opts,
     const TBDGenOptions &TBDOpts, std::unique_ptr<SILModule> SILMod,
     StringRef ModuleName, const PrimarySpecificPaths &PSPs,
+    std::shared_ptr<llvm::cas::ObjectStore> CAS,
     ArrayRef<std::string> parallelOutputFilenames,
     ArrayRef<std::string> parallelIROutputFilenames,
-    llvm::GlobalVariable **outModuleHash) {
+    llvm::GlobalVariable **outModuleHash,
+    cas::SwiftCASOutputBackend *casBackend) {
   // Get a pointer to the SILModule to avoid a potential use-after-move.
   const auto *SILModPtr = SILMod.get();
   const auto &SILOpts = SILModPtr->getOptions();
   auto desc = IRGenDescriptor::forWholeModule(
       M, Opts, TBDOpts, SILOpts, SILModPtr->Types, std::move(SILMod),
-      ModuleName, PSPs, /*symsToEmit*/ std::nullopt, parallelOutputFilenames,
-      parallelIROutputFilenames, outModuleHash);
+      ModuleName, PSPs, std::move(CAS), /*symsToEmit*/ std::nullopt,
+      parallelOutputFilenames, parallelIROutputFilenames, outModuleHash,
+      casBackend);
 
   if (Opts.shouldPerformIRGenerationInParallel() &&
       !parallelOutputFilenames.empty() &&
@@ -1951,20 +1969,20 @@ GeneratedModule swift::performIRGeneration(
   return evaluateOrFatal(M->getASTContext().evaluator, IRGenRequest{desc});
 }
 
-GeneratedModule swift::
-performIRGeneration(FileUnit *file, const IRGenOptions &Opts,
-                    const TBDGenOptions &TBDOpts,
-                    std::unique_ptr<SILModule> SILMod,
-                    StringRef ModuleName, const PrimarySpecificPaths &PSPs,
-                    StringRef PrivateDiscriminator,
-                    llvm::GlobalVariable **outModuleHash) {
+GeneratedModule swift::performIRGeneration(
+    FileUnit *file, const IRGenOptions &Opts, const TBDGenOptions &TBDOpts,
+    std::unique_ptr<SILModule> SILMod, StringRef ModuleName,
+    const PrimarySpecificPaths &PSPs,
+    std::shared_ptr<llvm::cas::ObjectStore> CAS, StringRef PrivateDiscriminator,
+    llvm::GlobalVariable **outModuleHash,
+    cas::SwiftCASOutputBackend *casBackend) {
   // Get a pointer to the SILModule to avoid a potential use-after-move.
   const auto *SILModPtr = SILMod.get();
   const auto &SILOpts = SILModPtr->getOptions();
   auto desc = IRGenDescriptor::forFile(
       file, Opts, TBDOpts, SILOpts, SILModPtr->Types, std::move(SILMod),
-      ModuleName, PSPs, PrivateDiscriminator,
-      /*symsToEmit*/ std::nullopt, outModuleHash);
+      ModuleName, PSPs, std::move(CAS), PrivateDiscriminator,
+      /*symsToEmit*/ std::nullopt, outModuleHash, casBackend);
   return evaluateOrFatal(file->getASTContext().evaluator, IRGenRequest{desc});
 }
 
@@ -2038,7 +2056,7 @@ void swift::createSwiftModuleObjectFile(SILModule &SILMod, StringRef Buffer,
 bool swift::performLLVM(const IRGenOptions &Opts, ASTContext &Ctx,
                         llvm::Module *Module, StringRef OutputFilename) {
   // Build TargetMachine.
-  auto TargetMachine = createTargetMachine(Opts, Ctx);
+  auto TargetMachine = createTargetMachine(Opts, Ctx, /*CAS=*/nullptr);
   if (!TargetMachine)
     return true;