Changing the support for loading LLVM libraries to not rely on path to the filesystem

Author: joker-eph

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -51,7 +51,7 @@ class TargetOptions {
   /// obtaining the parent symbol table. The default compilation target is
   /// `Fatbin`.
   TargetOptions(
-      StringRef toolkitPath = {}, ArrayRef<std::string> linkFiles = {},
+      StringRef toolkitPath = {}, ArrayRef<Attribute> librariesToLink = {},
       StringRef cmdOptions = {}, StringRef elfSection = {},
       CompilationTarget compilationTarget = getDefaultCompilationTarget(),
       function_ref<SymbolTable *()> getSymbolTableCallback = {},
@@ -66,8 +66,8 @@ class TargetOptions {
   /// Returns the toolkit path.
   StringRef getToolkitPath() const;
 
-  /// Returns the files to link to.
-  ArrayRef<std::string> getLinkFiles() const;
+  /// Returns the LLVM libraries to link to.
+  ArrayRef<Attribute> getLibrariesToLink() const;
 
   /// Returns the command line options.
   StringRef getCmdOptions() const;
@@ -113,7 +113,7 @@ class TargetOptions {
   /// appropiate value: ie. `TargetOptions(TypeID::get<DerivedClass>())`.
   TargetOptions(
       TypeID typeID, StringRef toolkitPath = {},
-      ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
+      ArrayRef<Attribute> librariesToLink = {}, StringRef cmdOptions = {},
       StringRef elfSection = {},
       CompilationTarget compilationTarget = getDefaultCompilationTarget(),
       function_ref<SymbolTable *()> getSymbolTableCallback = {},
@@ -126,7 +126,7 @@ class TargetOptions {
   std::string toolkitPath;
 
   /// List of files to link with the LLVM module.
-  SmallVector<std::string> linkFiles;
+  SmallVector<Attribute> librariesToLink;
 
   /// An optional set of command line options to be used by the compilation
   /// process.

mlir/include/mlir/Target/LLVM/ModuleToObject.h

@@ -83,7 +83,7 @@ class ModuleToObject {
 
   /// Loads multiple bitcode files.
   LogicalResult loadBitcodeFilesFromList(
-      llvm::LLVMContext &context, ArrayRef<std::string> fileList,
+      llvm::LLVMContext &context, ArrayRef<Attribute> librariesToLink,
       SmallVector<std::unique_ptr<llvm::Module>> &llvmModules,
       bool failureOnError = true);
 

mlir/include/mlir/Target/LLVM/NVVM/Utils.h

@@ -46,14 +46,15 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   /// Returns the CUDA toolkit path.
   StringRef getToolkitPath() const;
 
-  /// Returns the bitcode files to be loaded.
-  ArrayRef<std::string> getFileList() const;
+  /// Returns the bitcode libraries to be linked into the gpu module after
+  /// translation to LLVM IR.
+  ArrayRef<Attribute> getLibrariesToLink() const;
 
-  /// Appends `nvvm/libdevice.bc` into `fileList`. Returns failure if the
+  /// Appends `nvvm/libdevice.bc` into `librariesToLink`. Returns failure if the
   /// library couldn't be found.
   LogicalResult appendStandardLibs();
 
-  /// Loads the bitcode files in `fileList`.
+  /// Loads the bitcode files in `librariesToLink`.
   virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
   loadBitcodeFiles(llvm::Module &module) override;
 
@@ -64,8 +65,10 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   /// CUDA toolkit path.
   std::string toolkitPath;
 
-  /// List of LLVM bitcode files to link to.
-  SmallVector<std::string> fileList;
+  /// List of LLVM bitcode to link into after translation to LLVM IR.
+  /// The attributes can be StringAttr pointing to a file path, or
+  /// a Resource blob pointing to the LLVM bitcode in-memory.
+  SmallVector<Attribute> librariesToLink;
 };
 } // namespace NVVM
 } // namespace mlir

mlir/include/mlir/Target/LLVM/ROCDL/Utils.h

@@ -16,6 +16,7 @@
 #include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVM/ModuleToObject.h"
 
@@ -61,8 +62,8 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   /// Returns the ROCM toolkit path.
   StringRef getToolkitPath() const;
 
-  /// Returns the bitcode files to be loaded.
-  ArrayRef<std::string> getFileList() const;
+  /// Returns the LLVM bitcode libraries to be linked.
+  ArrayRef<Attribute> getLibrariesToLink() const;
 
   /// Appends standard ROCm device libraries to `fileList`.
   LogicalResult appendStandardLibs(AMDGCNLibraries libs);
@@ -107,7 +108,7 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   std::string toolkitPath;
 
   /// List of LLVM bitcode files to link to.
-  SmallVector<std::string> fileList;
+  SmallVector<Attribute> librariesToLink;
 
   /// AMD GCN libraries to use when linking, the default is using none.
   AMDGCNLibraries deviceLibs = AMDGCNLibraries::None;

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -2483,30 +2483,30 @@ KernelMetadataAttr KernelTableAttr::lookup(StringAttr key) const {
 //===----------------------------------------------------------------------===//
 
 TargetOptions::TargetOptions(
-    StringRef toolkitPath, ArrayRef<std::string> linkFiles,
+    StringRef toolkitPath, ArrayRef<Attribute> librariesToLink,
     StringRef cmdOptions, StringRef elfSection,
     CompilationTarget compilationTarget,
     function_ref<SymbolTable *()> getSymbolTableCallback,
     function_ref<void(llvm::Module &)> initialLlvmIRCallback,
     function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
     function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
     function_ref<void(StringRef)> isaCallback)
-    : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,
+    : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, librariesToLink,
                     cmdOptions, elfSection, compilationTarget,
                     getSymbolTableCallback, initialLlvmIRCallback,
                     linkedLlvmIRCallback, optimizedLlvmIRCallback,
                     isaCallback) {}
 
 TargetOptions::TargetOptions(
-    TypeID typeID, StringRef toolkitPath, ArrayRef<std::string> linkFiles,
+    TypeID typeID, StringRef toolkitPath, ArrayRef<Attribute> librariesToLink,
     StringRef cmdOptions, StringRef elfSection,
     CompilationTarget compilationTarget,
     function_ref<SymbolTable *()> getSymbolTableCallback,
     function_ref<void(llvm::Module &)> initialLlvmIRCallback,
     function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
     function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
     function_ref<void(StringRef)> isaCallback)
-    : toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
+    : toolkitPath(toolkitPath.str()), librariesToLink(librariesToLink),
       cmdOptions(cmdOptions.str()), elfSection(elfSection.str()),
       compilationTarget(compilationTarget),
       getSymbolTableCallback(getSymbolTableCallback),
@@ -2519,7 +2519,9 @@ TypeID TargetOptions::getTypeID() const { return typeID; }
 
 StringRef TargetOptions::getToolkitPath() const { return toolkitPath; }
 
-ArrayRef<std::string> TargetOptions::getLinkFiles() const { return linkFiles; }
+ArrayRef<Attribute> TargetOptions::getLibrariesToLink() const {
+  return librariesToLink;
+}
 
 StringRef TargetOptions::getCmdOptions() const { return cmdOptions; }
 

mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp

@@ -68,8 +68,10 @@ void GpuModuleToBinaryPass::runOnOperation() {
     }
     return &parentTable.value();
   };
-
-  TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, elfSection,
+  SmallVector<Attribute> librariesToLink;
+  for (const std::string &path : linkFiles)
+    librariesToLink.push_back(StringAttr::get(&getContext(), path));
+  TargetOptions targetOptions(toolkitPath, librariesToLink, cmdOptions, elfSection,
                               *targetFormat, lazyTableBuilder);
   if (failed(transformGpuModulesToBinaries(
           getOperation(), OffloadingLLVMTranslationAttrInterface(nullptr),

mlir/lib/Target/LLVM/ModuleToObject.cpp

@@ -89,22 +89,29 @@ ModuleToObject::loadBitcodeFile(llvm::LLVMContext &context, StringRef path) {
 }
 
 LogicalResult ModuleToObject::loadBitcodeFilesFromList(
-    llvm::LLVMContext &context, ArrayRef<std::string> fileList,
+    llvm::LLVMContext &context, ArrayRef<Attribute> librariesToLink,
     SmallVector<std::unique_ptr<llvm::Module>> &llvmModules,
     bool failureOnError) {
-  for (const std::string &str : fileList) {
-    // Test if the path exists, if it doesn't abort.
-    StringRef pathRef = StringRef(str.data(), str.size());
-    if (!llvm::sys::fs::is_regular_file(pathRef)) {
+  for (Attribute linkLib : librariesToLink) {
+    if (auto filePath = dyn_cast<StringAttr>(linkLib)) {
+      // Test if the path exists, if it doesn't abort.
+      if (!llvm::sys::fs::is_regular_file(filePath.strref())) {
+        getOperation().emitError()
+            << "File path: " << filePath << " does not exist or is not a file.";
+        return failure();
+      }
+      // Load the file or abort on error.
+      if (auto bcFile = loadBitcodeFile(context, filePath))
+        llvmModules.push_back(std::move(bcFile));
+      else if (failureOnError)
+        return failure();
+      continue;
+    }
+    if (failureOnError) {
       getOperation().emitError()
-          << "File path: " << pathRef << " does not exist or is not a file.\n";
+          << "Unknown attribute describing LLVM library to load: " << linkLib;
       return failure();
     }
-    // Load the file or abort on error.
-    if (auto bcFile = loadBitcodeFile(context, pathRef))
-      llvmModules.push_back(std::move(bcFile));
-    else if (failureOnError)
-      return failure();
   }
   return success();
 }

mlir/lib/Target/LLVM/NVVM/Target.cpp

@@ -93,17 +93,14 @@ SerializeGPUModuleBase::SerializeGPUModuleBase(
                      targetOptions.getOptimizedLlvmIRCallback(),
                      targetOptions.getISACallback()),
       target(target), toolkitPath(targetOptions.getToolkitPath()),
-      fileList(targetOptions.getLinkFiles()) {
+      librariesToLink(targetOptions.getLibrariesToLink()) {
 
   // If `targetOptions` have an empty toolkitPath use `getCUDAToolkitPath`
   if (toolkitPath.empty())
     toolkitPath = getCUDAToolkitPath();
 
   // Append the files in the target attribute.
-  if (ArrayAttr files = target.getLink())
-    for (Attribute attr : files.getValue())
-      if (auto file = dyn_cast<StringAttr>(attr))
-        fileList.push_back(file.str());
+  librariesToLink.append(target.getLink().begin(), target.getLink().end());
 
   // Append libdevice to the files to be loaded.
   (void)appendStandardLibs();
@@ -126,8 +123,8 @@ NVVMTargetAttr SerializeGPUModuleBase::getTarget() const { return target; }
 
 StringRef SerializeGPUModuleBase::getToolkitPath() const { return toolkitPath; }
 
-ArrayRef<std::string> SerializeGPUModuleBase::getFileList() const {
-  return fileList;
+ArrayRef<Attribute> SerializeGPUModuleBase::getLibrariesToLink() const {
+  return librariesToLink;
 }
 
 // Try to append `libdevice` from a CUDA toolkit installation.
@@ -149,16 +146,16 @@ LogicalResult SerializeGPUModuleBase::appendStandardLibs() {
                                  << " does not exist or is not a file.\n";
       return failure();
     }
-    fileList.push_back(pathRef.str());
+    librariesToLink.push_back(StringAttr::get(target.getContext(), pathRef));
   }
   return success();
 }
 
 std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
 SerializeGPUModuleBase::loadBitcodeFiles(llvm::Module &module) {
   SmallVector<std::unique_ptr<llvm::Module>> bcFiles;
-  if (failed(loadBitcodeFilesFromList(module.getContext(), fileList, bcFiles,
-                                      true)))
+  if (failed(loadBitcodeFilesFromList(module.getContext(), librariesToLink,
+                                      bcFiles, true)))
     return std::nullopt;
   return std::move(bcFiles);
 }

mlir/lib/Target/LLVM/ROCDL/Target.cpp

@@ -97,17 +97,14 @@ SerializeGPUModuleBase::SerializeGPUModuleBase(
     : ModuleToObject(module, target.getTriple(), target.getChip(),
                      target.getFeatures(), target.getO()),
       target(target), toolkitPath(targetOptions.getToolkitPath()),
-      fileList(targetOptions.getLinkFiles()) {
+      librariesToLink(targetOptions.getLibrariesToLink()) {
 
   // If `targetOptions` has an empty toolkitPath use `getROCMPath`
   if (toolkitPath.empty())
     toolkitPath = getROCMPath();
 
   // Append the files in the target attribute.
-  if (ArrayAttr files = target.getLink())
-    for (Attribute attr : files.getValue())
-      if (auto file = dyn_cast<StringAttr>(attr))
-        fileList.push_back(file.str());
+  librariesToLink.append(target.getLink().begin(), target.getLink().end());
 }
 
 void SerializeGPUModuleBase::init() {
@@ -128,8 +125,8 @@ ROCDLTargetAttr SerializeGPUModuleBase::getTarget() const { return target; }
 
 StringRef SerializeGPUModuleBase::getToolkitPath() const { return toolkitPath; }
 
-ArrayRef<std::string> SerializeGPUModuleBase::getFileList() const {
-  return fileList;
+ArrayRef<Attribute> SerializeGPUModuleBase::getLibrariesToLink() const {
+  return librariesToLink;
 }
 
 LogicalResult SerializeGPUModuleBase::appendStandardLibs(AMDGCNLibraries libs) {
@@ -160,7 +157,7 @@ LogicalResult SerializeGPUModuleBase::appendStandardLibs(AMDGCNLibraries libs) {
                                   << " does not exist or is not a file";
       return true;
     }
-    fileList.push_back(pathRef.str());
+    librariesToLink.push_back(StringAttr::get(target.getContext(), pathRef));
     path.truncate(baseSize);
     return false;
   };
@@ -178,13 +175,13 @@ LogicalResult SerializeGPUModuleBase::appendStandardLibs(AMDGCNLibraries libs) {
 std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
 SerializeGPUModuleBase::loadBitcodeFiles(llvm::Module &module) {
   // Return if there are no libs to load.
-  if (deviceLibs == AMDGCNLibraries::None && fileList.empty())
+  if (deviceLibs == AMDGCNLibraries::None && librariesToLink.empty())
     return SmallVector<std::unique_ptr<llvm::Module>>();
   if (failed(appendStandardLibs(deviceLibs)))
     return std::nullopt;
   SmallVector<std::unique_ptr<llvm::Module>> bcFiles;
-  if (failed(loadBitcodeFilesFromList(module.getContext(), fileList, bcFiles,
-                                      true)))
+  if (failed(loadBitcodeFilesFromList(module.getContext(), librariesToLink,
+                                      bcFiles, true)))
     return std::nullopt;
   return std::move(bcFiles);
 }