[COMGR] Add VFS support for device library linking (llvm#2472)

Author: chinmaydd

amd/comgr/README.md

@@ -145,7 +145,7 @@ By default, the cache is enabled.
   termination. The string format aligns with [Clang's ThinLTO cache pruning policy](https://clang.llvm.org/docs/ThinLTO.html#cache-pruning).
   The default policy is set as: "prune_interval=1h:prune_expiration=0h:cache_size=75%:cache_size_bytes=30g:cache_size_files=0".
 
-Comgr also supports some environment variables to aid in debugging. These
+Comgr supports some environment variables to aid in debugging. These
 include:
 
 * `AMD_COMGR_SAVE_TEMPS`: If this is set, and is not "0", Comgr does not delete
@@ -163,6 +163,20 @@ include:
 * `AMD_COMGR_TIME_STATISTICS`: If this is set, and is not "0", logs will
   include additional Comgr-specific timing information for compilation actions.
 
+Comgr implements support for an in-memory, virtual filesystem (VFS) for storing
+temporaries generated during intermediate compilation steps. This is aimed at 
+improving performance by reducing on-disk file I/O. Currently, VFS is only supported 
+for the device library link step, but we aim to progressively add support for
+more actions.
+
+By default, VFS is turned on.
+
+* `AMD_COMGR_USE_VFS`: When set to "0", VFS support is turned off.
+* Users may use the API `amd_comgr_action_info_set_vfs` to disable VFS for individual actions
+  without having to modify system-wide environment variables.
+* If `AMD_COMGR_SAVE_TEMPS` is set and not "0", VFS support is turned off irrespective
+  of `AMD_COMGR_USE_VFS` or the use of `amd_comgr_action_info_set_vfs`.
+
 Versioning
 ----------
 

amd/comgr/docs/ReleaseNotes.md

@@ -29,12 +29,18 @@ behaviors:
 Extensions (the same license used by LLVM).
 - Added Image Support to Comgr's handling of ISA metadata. Support for images
 can now be queried with Comgr's metadata APIs.
+- Added support for linking device library files through the use of a Virtual 
+File System (VFS).
 
 Bug Fixes
 ---------
 
 New APIs
 --------
+- amd\_comgr\_info\_set\_vfs\_() (v3.1)
+    - By setting this ActionInfo property, users can explicitly dictate if
+    device libraries should be linked using the real file system or a
+    Virtual File System (VFS).
 
 Deprecated APIs
 ---------------

amd/comgr/include/amd_comgr.h.in

@@ -234,6 +234,18 @@ extern "C" {
  */
 #define AMD_COMGR_VERSION_2_9
 
+/**
+ * The function was introduced or changed in version 3.0 of the interface
+ * and has the symbol version string of ``"@amd_comgr_NAME@_3.0"``.
+ */
+#define AMD_COMGR_VERSION_3_0
+
+/**
+ * The function was introduced or changed in version 3.1 of the interface
+ * and has the symbol version string of ``"@amd_comgr_NAME@_3.1"``.
+ */
+#define AMD_COMGR_VERSION_3_1
+
 /** @} */
 
 /**
@@ -1421,6 +1433,41 @@ amd_comgr_action_info_get_bundle_entry_id(
   size_t *size,
   char *bundle_entry_id) AMD_COMGR_VERSION_2_8;
 
+/**
+ * @brief Set whether the specified action should use an 
+ * in-memory virtual file system (VFS).
+ *
+ * @warning Environment variable @p AMD_COMGR_SAVE_TEMPS may override options 
+ * set by this API and @p AMD_COMGR_USE_VFS. If @p AMD_COMGR_SAVE_TEMPS is set 
+ * to "1", all actions are performed using the real file system irrespective of 
+ * the value of @p should_use_vfs @p AMD_COMGR_USE_VFS;
+ *
+ * @warning Environment variable @p AMD_COMGR_USE_VFS may override options 
+ * set by this API. If @p AMD_COMGR_USE_VFS is set to "1", all actions 
+ * are performed using VFS. If @p AMD_COMGR_USE_VFS is set to "0", 
+ * none of the actions are performed using VFS.
+ *
+ * If @p AMD_COMGR_USE_VFS is unset, this API can be used to selectively
+ * turn VFS usage on/off for specified actions.
+ *
+ * @param[in] action_info A handle to the action info object to be
+ * updated.
+ *
+ * @param[in] should_use_vfs A boolean that directs the choice to
+ * use the VFS.
+ *
+ * @retval ::AMD_COMGR_STATUS_SUCCESS The function has
+ * been executed successfully.
+ *
+ * @retval ::AMD_COMGR_STATUS_ERROR_INVALID_ARGUMENT @p
+ * action_info is an invalid action info object.
+ *
+ */
+amd_comgr_status_t AMD_COMGR_API
+amd_comgr_action_info_set_vfs(
+  amd_comgr_action_info_t action_info,
+  bool should_use_vfs) AMD_COMGR_VERSION_3_1;
+
 /**
  * @brief Set the device library linking behavior of an action info object.
  *

amd/comgr/src/amdcomgr.def

@@ -48,3 +48,4 @@ amd_comgr_action_info_set_bundle_entry_ids
 amd_comgr_action_info_get_bundle_entry_id_count
 amd_comgr_action_info_get_bundle_entry_id
 amd_comgr_action_info_set_device_lib_linking
+amd_comgr_action_info_set_vfs

amd/comgr/src/comgr-compiler.cpp

@@ -553,6 +553,11 @@ SmallString<128> getFilePath(DataObject *Object, StringRef Dir) {
   return Path;
 }
 
+// TODO: Move inputFromFile and outputToFile within AMDGPUCompiler
+//
+// Currently, we only invoke these two methods in the context of AMDGPUCompiler.
+// Moreover, member functions that deal with file I/O should not worry whether
+// the underlying filesystem being used is virtual or real.
 amd_comgr_status_t inputFromFile(DataObject *Object, StringRef Path) {
   ProfilePoint Point("FileIO");
   auto BufOrError = MemoryBuffer::getFile(Path);
@@ -646,7 +651,7 @@ void logArgv(raw_ostream &OS, StringRef ProgramName,
 
 amd_comgr_status_t executeCommand(const Command &Job, raw_ostream &LogS,
                                   DiagnosticOptions &DiagOpts,
-                                  llvm::vfs::FileSystem &VFS) {
+                                  llvm::vfs::FileSystem &FS) {
   TextDiagnosticPrinter DiagClient(LogS, &DiagOpts);
   IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs);
   DiagnosticsEngine Diags(DiagID, &DiagOpts, &DiagClient, false);
@@ -672,17 +677,19 @@ amd_comgr_status_t executeCommand(const Command &Job, raw_ostream &LogS,
 
     std::unique_ptr<CompilerInstance> Clang(new CompilerInstance());
     Clang->setVerboseOutputStream(LogS);
+    Clang->setFileManager(new FileManager(Clang->getFileSystemOpts(), &FS));
     if (!Argv.back()) {
       Argv.pop_back();
     }
+
     if (!CompilerInvocation::CreateFromArgs(Clang->getInvocation(), Argv,
                                             Diags)) {
       return AMD_COMGR_STATUS_ERROR;
     }
     // Internally this call refers to the invocation created above, so at
     // this point the DiagnosticsEngine should accurately reflect all user
     // requested configuration from Argv.
-    Clang->createDiagnostics(VFS, &DiagClient, /* ShouldOwnClient */ false);
+    Clang->createDiagnostics(FS, &DiagClient, /* ShouldOwnClient */ false);
     if (!Clang->hasDiagnostics()) {
       return AMD_COMGR_STATUS_ERROR;
     }
@@ -755,12 +762,11 @@ AMDGPUCompiler::executeInProcessDriver(ArrayRef<const char *> Args) {
   IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs);
   DiagnosticsEngine Diags(DiagID, &*DiagOpts, DiagClient);
 
-  auto VFS = llvm::vfs::getRealFileSystem();
-  ProcessWarningOptions(Diags, *DiagOpts, *VFS, /*ReportDiags=*/false);
+  ProcessWarningOptions(Diags, *DiagOpts, *OverlayFS, /*ReportDiags=*/false);
 
   Driver TheDriver((Twine(env::getLLVMPath()) + "/bin/clang").str(),
                    llvm::sys::getDefaultTargetTriple(), Diags,
-                   "AMDGPU Code Object Manager", VFS);
+                   "AMDGPU Code Object Manager", OverlayFS);
   TheDriver.setCheckInputsExist(false);
 
   // Log arguments used to build compilation
@@ -782,7 +788,7 @@ AMDGPUCompiler::executeInProcessDriver(ArrayRef<const char *> Args) {
 
   auto Cache = CommandCache::get(LogS);
   for (auto &Job : C->getJobs()) {
-    ClangCommand C(Job, *DiagOpts, *VFS, executeCommand);
+    ClangCommand C(Job, *DiagOpts, *OverlayFS, executeCommand);
     if (Cache) {
       if (auto Status = Cache->execute(C, LogS)) {
         return Status;
@@ -1089,8 +1095,18 @@ amd_comgr_status_t AMDGPUCompiler::addDeviceLibraries() {
     for (auto DeviceLib : getDeviceLibraries()) {
       llvm::SmallString<128> DeviceLibPath = DeviceLibsDir;
       path::append(DeviceLibPath, std::get<0>(DeviceLib));
-      if (auto Status = outputToFile(std::get<1>(DeviceLib), DeviceLibPath)) {
-        return Status;
+      // TODO: We should abstract the logic of deciding whether to use the VFS
+      // or the real file system within inputFromFile and outputToFile.
+      if (UseVFS) {
+        if (!InMemoryFS->addFile(
+                DeviceLibPath, /* ModificationTime */ 0,
+                llvm::MemoryBuffer::getMemBuffer(std::get<1>(DeviceLib)))) {
+          return AMD_COMGR_STATUS_ERROR;
+        }
+      } else {
+        if (auto Status = outputToFile(std::get<1>(DeviceLib), DeviceLibPath)) {
+          return Status;
+        }
       }
     }
   }
@@ -1944,6 +1960,25 @@ AMDGPUCompiler::AMDGPUCompiler(DataAction *ActionInfo, DataSet *InSet,
     : ActionInfo(ActionInfo), InSet(InSet), OutSetT(DataSet::convert(OutSet)),
       LogS(LogS) {
   initializeCommandLineArgs(Args);
+
+  // Initialize OverlayFS with the real file system which helps redirect
+  // non-VFS reads and writes.
+  OverlayFS = new vfs::OverlayFileSystem(vfs::getRealFileSystem());
+
+  std::optional<bool> VFSStatus = env::shouldUseVFS();
+  if ((VFSStatus.has_value() && *VFSStatus) ||
+      (!VFSStatus.has_value() && ActionInfo->ShouldUseVFS)) {
+    if (env::shouldEmitVerboseLogs()) {
+      LogS << "    File System: VFS\n";
+    }
+    UseVFS = true;
+    InMemoryFS = new vfs::InMemoryFileSystem;
+    OverlayFS->pushOverlay(InMemoryFS);
+  } else {
+    if (env::shouldEmitVerboseLogs()) {
+      LogS << "    File System: Real\n";
+    }
+  }
 }
 
 AMDGPUCompiler::~AMDGPUCompiler() {

amd/comgr/src/comgr-compiler.h

@@ -41,6 +41,7 @@
 
 #include "comgr.h"
 #include "clang/Driver/Driver.h"
+#include "llvm/Support/VirtualFileSystem.h"
 
 namespace COMGR {
 
@@ -72,6 +73,10 @@ class AMDGPUCompiler {
   llvm::StringSaver Saver = Allocator;
   /// Whether we need to disable Clang's device-lib linking.
   bool NoGpuLib = true;
+  bool UseVFS = false;
+
+  llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> OverlayFS;
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> InMemoryFS;
 
   amd_comgr_status_t createTmpDirs();
   amd_comgr_status_t removeTmpDirs();

amd/comgr/src/comgr-env.cpp

@@ -47,6 +47,26 @@ bool shouldSaveTemps() {
   return SaveTemps && StringRef(SaveTemps) != "0";
 }
 
+bool shouldSaveLLVMTemps() {
+  static char *SaveTemps = getenv("AMD_COMGR_SAVE_LLVM_TEMPS");
+  return SaveTemps && StringRef(SaveTemps) != "0";
+}
+
+std::optional<bool> shouldUseVFS() {
+  if (shouldSaveTemps())
+    return false;
+
+  static char *UseVFS = getenv("AMD_COMGR_USE_VFS");
+  if (UseVFS) {
+    if (StringRef(UseVFS) == "0")
+      return false;
+    else if (StringRef(UseVFS) == "1")
+      return true;
+  }
+
+  return std::nullopt;
+}
+
 std::optional<StringRef> getRedirectLogs() {
   static char *RedirectLogs = getenv("AMD_COMGR_REDIRECT_LOGS");
   if (!RedirectLogs || StringRef(RedirectLogs) == "0") {

amd/comgr/src/comgr-env.h

@@ -43,6 +43,8 @@ namespace env {
 
 /// Return whether the environment requests temps be saved.
 bool shouldSaveTemps();
+bool shouldSaveLLVMTemps();
+std::optional<bool> shouldUseVFS();
 
 /// If the environment requests logs be redirected, return the string identifier
 /// of where to redirect. Otherwise return @p None.

amd/comgr/src/comgr.cpp

@@ -1186,6 +1186,22 @@ amd_comgr_status_t AMD_COMGR_API
   return ActionP->setBundleEntryIDs(ArrayRef<const char *>(EntryIDs, Count));
 }
 
+amd_comgr_status_t AMD_COMGR_API
+    // NOLINTNEXTLINE(readability-identifier-naming)
+    amd_comgr_action_info_set_vfs
+    //
+    (amd_comgr_action_info_t ActionInfo, bool ShouldUseVFS) {
+  DataAction *ActionP = DataAction::convert(ActionInfo);
+
+  if (!ActionP) {
+    return AMD_COMGR_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  ActionP->ShouldUseVFS = ShouldUseVFS;
+
+  return AMD_COMGR_STATUS_SUCCESS;
+}
+
 amd_comgr_status_t AMD_COMGR_API
     // NOLINTNEXTLINE(readability-identifier-naming)
     amd_comgr_action_info_set_device_lib_linking

amd/comgr/src/comgr.h

@@ -230,6 +230,7 @@ struct DataAction {
   amd_comgr_language_t Language;
   bool Logging;
   bool ShouldLinkDeviceLibs = false;
+  bool ShouldUseVFS = true;
 
   std::vector<std::string> BundleEntryIDs;