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

chinmaydd · web-flow · commit 3332d18c2507 · 2025-04-16T09:00:30.000-07:00
This commit adds support for linking device library files through the use of a Virtual File System (VFS).

Change-Id: I27f1dccfe772ebe9d11f4a570a05d0b9d55a705d
diff --git a/amd/comgr/README.md b/amd/comgr/README.md
@@ -142,7 +142,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
@@ -162,6 +162,22 @@ 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 off. Set the environment variable `AMD_COMGR_USE_VFS=1` to 
+enable it.
+
+* `AMD_COMGR_USE_VFS`: When unset or "0", VFS support is turned off.
+* Users may use the API `amd_comgr_action_info_set_vfs` to enable VFS for individual actions
+  without having to modify system-wide environment variables. Note that API is effective only 
+  if `AMD_COMGR_USE_VFS` is unset.
+* 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
 ----------
 
diff --git a/amd/comgr/docs/ReleaseNotes.md b/amd/comgr/docs/ReleaseNotes.md
@@ -29,12 +29,17 @@ 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, a user can explicitly dictate if
+    device libraries should be linked using a Virtual File System (VFS).
 
 Deprecated APIs
 ---------------
diff --git a/amd/comgr/include/amd_comgr.h.in b/amd/comgr/include/amd_comgr.h.in
@@ -225,6 +225,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
+
 /** @} */
 
 /**
@@ -1420,6 +1432,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.
  *
diff --git a/amd/comgr/src/amdcomgr.def b/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
diff --git a/amd/comgr/src/comgr-compiler.cpp b/amd/comgr/src/comgr-compiler.cpp
@@ -532,6 +532,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);
@@ -625,7 +630,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);
@@ -651,17 +656,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;
     }
@@ -734,12 +741,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
@@ -761,7 +767,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;
@@ -1075,8 +1081,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;
+        }
       }
     }
   }
@@ -2074,6 +2090,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() {
diff --git a/amd/comgr/src/comgr-compiler.h b/amd/comgr/src/comgr-compiler.h
@@ -11,6 +11,7 @@
 
 #include "comgr.h"
 #include "clang/Driver/Driver.h"
+#include "llvm/Support/VirtualFileSystem.h"
 
 namespace COMGR {
 
@@ -36,6 +37,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();
diff --git a/amd/comgr/src/comgr-env.cpp b/amd/comgr/src/comgr-env.cpp
@@ -31,6 +31,21 @@ bool shouldSaveLLVMTemps() {
   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") {
diff --git a/amd/comgr/src/comgr-env.h b/amd/comgr/src/comgr-env.h
@@ -17,6 +17,7 @@ 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.
diff --git a/amd/comgr/src/comgr.cpp b/amd/comgr/src/comgr.cpp
@@ -1169,6 +1169,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
diff --git a/amd/comgr/src/comgr.h b/amd/comgr/src/comgr.h
@@ -204,6 +204,7 @@ struct DataAction {
   amd_comgr_language_t Language;
   bool Logging;
   bool ShouldLinkDeviceLibs = false;
+  bool ShouldUseVFS = false;
 
   std::vector<std::string> BundleEntryIDs;
 
diff --git a/amd/comgr/src/exportmap.in b/amd/comgr/src/exportmap.in
@@ -91,3 +91,7 @@ global: amd_comgr_action_info_set_bundle_entry_ids;
 @amd_comgr_NAME@_2.9 {
 global: amd_comgr_action_info_set_device_lib_linking;
 } @amd_comgr_NAME@_2.8;
+
+@amd_comgr_NAME@3.1 {
+global: amd_comgr_action_info_set_vfs;
+} @amd_comgr_NAME@_2.9;
diff --git a/amd/comgr/test-lit/comgr-sources/source-to-bc-with-dev-libs.c b/amd/comgr/test-lit/comgr-sources/source-to-bc-with-dev-libs.c
@@ -22,9 +22,9 @@ int main(int argc, char *argv[]) {
                                   "-amdgpu-prelink"};
   size_t CodeGenOptionsCount =
       sizeof(CodeGenOptions) / sizeof(CodeGenOptions[0]);
-  if (argc != 4) {
-    fprintf(stderr,
-            "Usage: source-to-bc-with-device-libs file.cl -o file.bc\n");
+  if (argc < 4 || argc > 5) {
+    fprintf(stderr, "Usage: source-to-bc-with-device-libs file.cl "
+                    "[--vfs|--novfs] -o file.bc\n");
     exit(1);
   }
 
@@ -42,6 +42,12 @@ int main(int argc, char *argv[]) {
   amd_comgr_(action_info_set_isa_name(DataAction, "amdgcn-amd-amdhsa--gfx900"));
   amd_comgr_(create_data_set(&DataSetPch));
 
+  if (!strncmp(argv[2], "--vfs", 5)) {
+    amd_comgr_(action_info_set_vfs(DataAction, true));
+  } else if (!strncmp(argv[2], "--novfs", 7)) {
+    amd_comgr_(action_info_set_vfs(DataAction, false));
+  }
+
   amd_comgr_(do_action(AMD_COMGR_ACTION_ADD_PRECOMPILED_HEADERS, DataAction,
                        DataSetIn, DataSetPch));
 
@@ -74,7 +80,7 @@ int main(int argc, char *argv[]) {
   amd_comgr_data_t DataBc;
   amd_comgr_(
       action_data_get_data(DataSetBc, AMD_COMGR_DATA_KIND_BC, 0, &DataBc));
-  dumpData(DataBc, argv[3]);
+  dumpData(DataBc, argv[argc - 1]);
 
   amd_comgr_(release_data(DataSource));
   amd_comgr_(release_data(DataBc));
diff --git a/amd/comgr/test-lit/vfs-tests/lit.local.cfg b/amd/comgr/test-lit/vfs-tests/lit.local.cfg
@@ -0,0 +1,2 @@
+config.environment['AMD_COMGR_EMIT_VERBOSE_LOGS'] = "1"
+config.environment['AMD_COMGR_REDIRECT_LOGS'] = "stdout"
diff --git a/amd/comgr/test-lit/vfs-tests/vfs-tests.cl b/amd/comgr/test-lit/vfs-tests/vfs-tests.cl
diff --git a/amd/comgr/test/CMakeLists.txt b/amd/comgr/test/CMakeLists.txt
diff --git a/amd/comgr/test/compile_source_with_device_libs_to_bc_with_vfs_test.c b/amd/comgr/test/compile_source_with_device_libs_to_bc_with_vfs_test.c
