[Comgr][Cache] Cache unbundle of compressed bundles (#526)

Author: jmmartinez

amd/comgr/CMakeLists.txt

@@ -85,6 +85,7 @@ set(SOURCES
   src/comgr-spirv-command.cpp
   src/comgr-symbol.cpp
   src/comgr-symbolizer.cpp
+  src/comgr-unbundle-command.cpp
   src/time-stat/time-stat.cpp)
 
 if(COMGR_BUILD_SHARED_LIBS)

amd/comgr/src/comgr-cache-command.cpp

@@ -98,4 +98,39 @@ CachedCommandAdaptor::getIdentifier() const {
   toHex(H.final(), true, Id);
   return Id;
 }
+
+llvm::Error
+CachedCommandAdaptor::writeSingleOutputFile(StringRef OutputFilename,
+                                            StringRef CachedBuffer) {
+  std::error_code EC;
+  raw_fd_ostream Out(OutputFilename, EC);
+  if (EC) {
+    Error E = createStringError(EC, Twine("Failed to open ") + OutputFilename +
+                                        " : " + EC.message() + "\n");
+    return E;
+  }
+
+  Out.write(CachedBuffer.data(), CachedBuffer.size());
+  Out.close();
+  if (Out.has_error()) {
+    Error E = createStringError(EC, Twine("Failed to write ") + OutputFilename +
+                                        " : " + EC.message() + "\n");
+    return E;
+  }
+
+  return Error::success();
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+CachedCommandAdaptor::readSingleOutputFile(StringRef OutputFilename) {
+  ErrorOr<std::unique_ptr<MemoryBuffer>> MBOrErr =
+      MemoryBuffer::getFile(OutputFilename);
+  if (!MBOrErr) {
+    std::error_code EC = MBOrErr.getError();
+    return createStringError(EC, Twine("Failed to open ") + OutputFilename +
+                                     " : " + EC.message() + "\n");
+  }
+
+  return std::move(*MBOrErr);
+}
 } // namespace COMGR

amd/comgr/src/comgr-cache-command.h

@@ -69,6 +69,12 @@ class CachedCommandAdaptor {
   static void addString(HashAlgorithm &H, llvm::StringRef S);
   static std::optional<size_t> searchComgrTmpModel(llvm::StringRef S);
 
+  // helper since several command types just write to a single output file
+  static llvm::Error writeSingleOutputFile(llvm::StringRef OutputFilename,
+                                           llvm::StringRef CachedBuffer);
+  static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
+  readSingleOutputFile(llvm::StringRef OutputFilename);
+
 protected:
   virtual ActionClass getClass() const = 0;
   virtual void addOptionsIdentifier(HashAlgorithm &) const = 0;

amd/comgr/src/comgr-clang-command.cpp

@@ -158,35 +158,16 @@ bool ClangCommand::canCache() const {
 
 Error ClangCommand::writeExecuteOutput(StringRef CachedBuffer) {
   StringRef OutputFilename = Command.getOutputFilenames().front();
-  std::error_code EC;
-  raw_fd_ostream Out(OutputFilename, EC);
-  if (EC) {
-    Error E = createStringError(EC, Twine("Failed to open ") + OutputFilename +
-                                        " : " + EC.message() + "\n");
-    return E;
-  }
-
-  Out.write(CachedBuffer.data(), CachedBuffer.size());
-  Out.close();
-  if (Out.has_error()) {
-    Error E = createStringError(EC, Twine("Failed to write ") + OutputFilename +
-                                        " : " + EC.message() + "\n");
-    return E;
-  }
-
-  return Error::success();
+  return CachedCommandAdaptor::writeSingleOutputFile(OutputFilename,
+                                                     CachedBuffer);
 }
 
 Expected<StringRef> ClangCommand::readExecuteOutput() {
-  StringRef OutputFilename = Command.getOutputFilenames().front();
-  ErrorOr<std::unique_ptr<MemoryBuffer>> MBOrErr =
-      MemoryBuffer::getFile(OutputFilename);
-  if (!MBOrErr) {
-    std::error_code EC = MBOrErr.getError();
-    return createStringError(EC, Twine("Failed to open ") + OutputFilename +
-                                     " : " + EC.message() + "\n");
-  }
-  Output = std::move(*MBOrErr);
+  auto MaybeBuffer = CachedCommandAdaptor::readSingleOutputFile(
+      Command.getOutputFilenames().front());
+  if (!MaybeBuffer)
+    return MaybeBuffer.takeError();
+  Output = std::move(*MaybeBuffer);
   return Output->getBuffer();
 }
 

amd/comgr/src/comgr-compiler.cpp

@@ -43,6 +43,7 @@
 #include "comgr-diagnostic-handler.h"
 #include "comgr-env.h"
 #include "comgr-spirv-command.h"
+#include "comgr-unbundle-command.h"
 #include "lld/Common/CommonLinkerContext.h"
 #include "lld/Common/Driver.h"
 #include "clang/CodeGen/CodeGenAction.h"
@@ -1255,9 +1256,10 @@ amd_comgr_status_t AMDGPUCompiler::unbundle() {
   }
 
   // Collect bitcode memory buffers from bitcodes, bundles, and archives
+  auto Cache = CommandCache::get(LogS);
   for (auto *Input : InSet->DataObjects) {
 
-    std::string FileExtension;
+    const char *FileExtension;
     amd_comgr_data_kind_t UnbundledDataKind;
     switch (Input->DataKind) {
     case AMD_COMGR_DATA_KIND_BC_BUNDLE:
@@ -1288,22 +1290,22 @@ amd_comgr_status_t AMDGPUCompiler::unbundle() {
       const size_t BufSize = sizeof(char) * 30;
       char *Buf = (char *)malloc(BufSize);
       snprintf(Buf, BufSize, "comgr-bundle-%d.%s", std::rand() % 10000,
-               FileExtension.c_str());
+               FileExtension);
       Input->Name = Buf;
     }
 
     // Write input file system so that OffloadBundler API can process
     // TODO: Switch write to VFS
-    std::string InputFilePath = getFilePath(Input, InputDir).str().str();
+    SmallString<128> InputFilePath = getFilePath(Input, InputDir);
     if (auto Status = outputToFile(Input, InputFilePath)) {
       return Status;
     }
 
     // Bundler input name
-    BundlerConfig.InputFileNames.push_back(InputFilePath);
+    BundlerConfig.InputFileNames.emplace_back(InputFilePath);
 
     // Generate prefix for output files
-    std::string OutputPrefix = std::string(Input->Name);
+    StringRef OutputPrefix = Input->Name;
     size_t Index = OutputPrefix.find_last_of(".");
     OutputPrefix = OutputPrefix.substr(0, Index);
 
@@ -1314,53 +1316,35 @@ amd_comgr_status_t AMDGPUCompiler::unbundle() {
            << "   Unbundled Files Extension: ." << FileExtension << "\n";
     }
 
-    for (size_t I = 0; I < ActionInfo->BundleEntryIDs.size(); I++) {
-      auto Entry = ActionInfo->BundleEntryIDs[I];
-      BundlerConfig.TargetNames.push_back(Entry);
-
+    for (StringRef Entry : ActionInfo->BundleEntryIDs) {
       // Add an output file for each target
-      std::string OutputFileName =
-          OutputPrefix + '-' + Entry + "." + FileExtension;
+      SmallString<128> OutputFilePath = OutputDir;
+      sys::path::append(OutputFilePath,
+                        OutputPrefix + "-" + Entry + "." + FileExtension);
 
-      // TODO: Switch this to LLVM path APIs
-      std::string OutputFilePath = OutputDir.str().str() + "/" + OutputFileName;
-      BundlerConfig.OutputFileNames.push_back(OutputFilePath);
+      BundlerConfig.TargetNames.emplace_back(Entry);
+      BundlerConfig.OutputFileNames.emplace_back(OutputFilePath);
 
       if (env::shouldEmitVerboseLogs()) {
-        LogS << "\tBundle Entry ID: " << BundlerConfig.TargetNames[I] << "\n"
-             << "\tOutput Filename: " << BundlerConfig.OutputFileNames[I]
-             << "\n";
+        LogS << "\tBundle Entry ID: " << Entry << "\n"
+             << "\tOutput Filename: " << OutputFilePath << "\n";
         LogS.flush();
       }
     }
 
-    OffloadBundler Bundler(BundlerConfig);
-
-    switch (Input->DataKind) {
-    case AMD_COMGR_DATA_KIND_BC_BUNDLE: {
-      llvm::Error Err = Bundler.UnbundleFiles();
-      llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(),
-                                  "Unbundle Bitcodes Error: ");
-      break;
-    }
-    case AMD_COMGR_DATA_KIND_AR_BUNDLE: {
-      llvm::Error Err = Bundler.UnbundleArchive();
-      llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(),
-                                  "Unbundle Archives Error: ");
-      break;
-    }
-    case AMD_COMGR_DATA_KIND_OBJ_BUNDLE: {
-      llvm::Error Err = Bundler.UnbundleFiles();
-      llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(),
-                                  "Unbundle Objects Error: ");
-      break;
-    }
-    default:
-      llvm_unreachable("invalid bundle type");
+    UnbundleCommand Unbundle(Input->DataKind, BundlerConfig);
+    if (Cache) {
+      if (auto Status = Cache->execute(Unbundle, LogS)) {
+        return Status;
+      }
+    } else {
+      if (auto Status = Unbundle.execute(LogS)) {
+        return Status;
+      }
     }
 
     // Add new bitcodes to OutSetT
-    for (auto OutputFilePath : BundlerConfig.OutputFileNames) {
+    for (StringRef OutputFilePath : BundlerConfig.OutputFileNames) {
 
       amd_comgr_data_t ResultT;
 
@@ -1371,22 +1355,15 @@ amd_comgr_status_t AMDGPUCompiler::unbundle() {
       ScopedDataObjectReleaser SDOR(ResultT);
 
       DataObject *Result = DataObject::convert(ResultT);
-      if (auto Status = inputFromFile(Result, StringRef(OutputFilePath)))
+      if (auto Status = inputFromFile(Result, OutputFilePath))
         return Status;
 
-      StringRef OutputFileName =
-          llvm::sys::path::filename(StringRef(OutputFilePath));
+      StringRef OutputFileName = sys::path::filename(OutputFilePath);
       Result->setName(OutputFileName);
 
       if (auto Status = amd_comgr_data_set_add(OutSetT, ResultT)) {
         return Status;
       }
-
-      // Remove input and output file after reading back into Comgr data
-      if (!env::shouldEmitVerboseLogs()) {
-        sys::fs::remove(InputFilePath);
-        sys::fs::remove(OutputFilePath);
-      }
     }
   }