[MLIR] Add a BlobAttr interface for attribute to wrap arbitrary content and use it as linkLibs for ModuleToObject

This change allows to expose through an interface attributes wrapping content as external
resources, and the usage inside the ModuleToObject show how we will be able to provide
runtime libraries without relying on the filesystem.

Author: joker-eph

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td

@@ -34,6 +34,24 @@ def TypedAttrInterface : AttrInterface<"TypedAttr"> {
   >];
 }
 
+//===----------------------------------------------------------------------===//
+// BlobAttrInterface
+//===----------------------------------------------------------------------===//
+
+def BlobAttrInterface : AttrInterface<"BlobAttr"> {
+  let cppNamespace = "::mlir";
+  let description = [{
+    This interface allows an attribute to expose a blob of data without more
+    information. The data must be stored so that it can be accessed as a
+    contiguous ArrayRef.
+  }];
+
+  let methods = [InterfaceMethod<
+    "Get the attribute's data",
+    "::llvm::ArrayRef<char>", "getData"
+  >];
+}
+
 //===----------------------------------------------------------------------===//
 // ElementsAttrInterface
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/BuiltinAttributes.td

@@ -153,7 +153,8 @@ def Builtin_DenseArrayRawDataParameter : ArrayRefParameter<
   }];
 }
 
-def Builtin_DenseArray : Builtin_Attr<"DenseArray", "dense_array"> {
+def Builtin_DenseArray : Builtin_Attr<"DenseArray", "dense_array", 
+    [BlobAttrInterface]> {
   let summary = "A dense array of integer or floating point elements.";
   let description = [{
     A dense array attribute is an attribute that represents a dense array of
@@ -211,6 +212,10 @@ def Builtin_DenseArray : Builtin_Attr<"DenseArray", "dense_array"> {
     int64_t size() const { return getSize(); }
     /// Return true if there are no elements in the dense array.
     bool empty() const { return !size(); }
+    /// BlobAttrInterface method.
+    ArrayRef<char> getData() {
+      return getRawData();
+    }
   }];
 }
 
@@ -431,7 +436,7 @@ def Builtin_DenseStringElementsAttr : Builtin_Attr<
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseResourceElementsAttr : Builtin_Attr<"DenseResourceElements",
-    "dense_resource_elements", [ElementsAttrInterface]> {
+    "dense_resource_elements", [ElementsAttrInterface, BlobAttrInterface]> {
   let summary = "An Attribute containing a dense multi-dimensional array "
                 "backed by a resource";
   let description = [{
@@ -485,6 +490,10 @@ def Builtin_DenseResourceElementsAttr : Builtin_Attr<"DenseResourceElements",
       "ShapedType":$type, "StringRef":$blobName, "AsmResourceBlob":$blob
     )>
   ];
+  let extraClassDeclaration = [{
+    /// BlobAttrInterface method.
+    ArrayRef<char> getData();
+  }];
 
   let skipDefaultBuilders = 1;
 }

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

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

mlir/lib/IR/BuiltinAttributes.cpp

@@ -1544,6 +1544,12 @@ DenseResourceElementsAttr DenseResourceElementsAttr::get(ShapedType type,
   return get(type, manager.insert(blobName, std::move(blob)));
 }
 
+ArrayRef<char> DenseResourceElementsAttr::getData() {
+  if (AsmResourceBlob *blob = this->getRawHandle().getBlob())
+    return blob->getDataAs<char>();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // DenseResourceElementsAttrBase
 

mlir/lib/Target/LLVM/ModuleToObject.cpp

@@ -14,6 +14,8 @@
 #include "mlir/Target/LLVM/ModuleToObject.h"
 
 #include "mlir/ExecutionEngine/OptUtils.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
@@ -25,6 +27,7 @@
 #include "llvm/Linker/Linker.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
@@ -93,6 +96,9 @@ LogicalResult ModuleToObject::loadBitcodeFilesFromList(
     SmallVector<std::unique_ptr<llvm::Module>> &llvmModules,
     bool failureOnError) {
   for (Attribute linkLib : librariesToLink) {
+    // Attributes in this list can be either list of file paths using
+    // StringAttr, or a resource attribute pointing to the LLVM bitcode in
+    // memory.
     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())) {
@@ -107,6 +113,27 @@ LogicalResult ModuleToObject::loadBitcodeFilesFromList(
         return failure();
       continue;
     }
+    if (auto blobAttr = dyn_cast<BlobAttr>(linkLib)) {
+      // Load the file or abort on error.
+      llvm::SMDiagnostic error;
+      ArrayRef<char> data = blobAttr.getData();
+      std::unique_ptr<llvm::MemoryBuffer> buffer =
+          llvm::MemoryBuffer::getMemBuffer(StringRef(data.data(), data.size()),
+                                           "blobLinkedLib",
+                                           /*RequiresNullTerminator=*/false);
+      std::unique_ptr<llvm::Module> mod =
+          getLazyIRModule(std::move(buffer), error, context);
+      if (mod) {
+        if (failed(handleBitcodeFile(*mod)))
+          return failure();
+        llvmModules.push_back(std::move(mod));
+      } else if (failureOnError) {
+        getOperation().emitError()
+            << "Couldn't load LLVM library for linking: " << error.getMessage();
+        return failure();
+      }
+      continue;
+    }
     if (failureOnError) {
       getOperation().emitError()
           << "Unknown attribute describing LLVM library to load: " << linkLib;

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

@@ -100,7 +100,8 @@ SerializeGPUModuleBase::SerializeGPUModuleBase(
     toolkitPath = getCUDAToolkitPath();
 
   // Append the files in the target attribute.
-  librariesToLink.append(target.getLink().begin(), target.getLink().end());
+  if (target.getLink())
+    librariesToLink.append(target.getLink().begin(), target.getLink().end());
 
   // Append libdevice to the files to be loaded.
   (void)appendStandardLibs();

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

@@ -104,7 +104,8 @@ SerializeGPUModuleBase::SerializeGPUModuleBase(
     toolkitPath = getROCMPath();
 
   // Append the files in the target attribute.
-  librariesToLink.append(target.getLink().begin(), target.getLink().end());
+  if (target.getLink())
+    librariesToLink.append(target.getLink().begin(), target.getLink().end());
 }
 
 void SerializeGPUModuleBase::init() {

mlir/unittests/Target/LLVM/CMakeLists.txt

@@ -1,9 +1,13 @@
 set(LLVM_LINK_COMPONENTS nativecodegen)
 
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+
 add_mlir_unittest(MLIRTargetLLVMTests
   SerializeNVVMTarget.cpp
   SerializeROCDLTarget.cpp
   SerializeToLLVMBitcode.cpp
+DEPENDS
+  ${dialect_libs}
 )
 
 mlir_target_link_libraries(MLIRTargetLLVMTests

mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp

@@ -18,15 +18,18 @@
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 
+#include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/Config/llvm-config.h" // for LLVM_HAS_NVPTX_TARGET
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/Process.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Host.h"
 
 #include "gmock/gmock.h"
+#include <cstdint>
 
 using namespace mlir;
 
@@ -215,3 +218,81 @@ TEST_F(MLIRTargetLLVMNVVM,
     isaResult.clear();
   }
 }
+
+// Test linking LLVM IR from a resource attribute.
+TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(LinkedLLVMIRResource)) {
+  MLIRContext context(registry);
+  std::string moduleStr = R"mlir(
+    gpu.module @nvvm_test {
+      llvm.func @bar()
+      llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {
+        llvm.call @bar() : () -> ()
+        llvm.return
+      }
+    }
+  )mlir";
+  // Provide the library to link as a serialized bitcode blob.
+  SmallVector<char> bitcodeToLink;
+  {
+    std::string linkedLib = R"llvm(
+  define void @bar() {
+    ret void
+  }
+  )llvm";
+    llvm::SMDiagnostic err;
+    llvm::MemoryBufferRef buffer(linkedLib, "linkedLib");
+    llvm::LLVMContext llvmCtx;
+    std::unique_ptr<llvm::Module> module = llvm::parseIR(buffer, err, llvmCtx);
+    ASSERT_TRUE(module) << " Can't parse IR: " << err.getMessage();
+    {
+      llvm::raw_svector_ostream os(bitcodeToLink);
+      WriteBitcodeToFile(*module, os);
+    }
+  }
+
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  Builder builder(&context);
+
+  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
+  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
+
+  // Hook to intercept the LLVM IR after linking external libs.
+  std::string linkedLLVMIR;
+  auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
+    llvm::raw_string_ostream ros(linkedLLVMIR);
+    module.print(ros, nullptr);
+  };
+
+  // Store the bitcode as a DenseI8ArrayAttr.
+  SmallVector<Attribute> librariesToLink;
+  librariesToLink.push_back(DenseI8ArrayAttr::get(
+      &context,
+      ArrayRef<int8_t>((int8_t *)bitcodeToLink.data(), bitcodeToLink.size())));
+  gpu::TargetOptions options({}, librariesToLink, {}, {},
+                             gpu::CompilationTarget::Assembly, {}, {},
+                             linkedCallback);
+  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
+    std::optional<SmallVector<char, 0>> object =
+        serializer.serializeToObject(gpuModule, options);
+
+    // Verify that we correctly linked in the library: the external call is
+    // replaced by the definition.
+    ASSERT_TRUE(!linkedLLVMIR.empty());
+    {
+      llvm::SMDiagnostic err;
+      llvm::MemoryBufferRef buffer(linkedLLVMIR, "linkedLLVMIR");
+      llvm::LLVMContext llvmCtx;
+      std::unique_ptr<llvm::Module> module =
+          llvm::parseIR(buffer, err, llvmCtx);
+      ASSERT_TRUE(module) << " Can't parse linkedLLVMIR: " << err.getMessage()
+                          << " IR: \n\b" << linkedLLVMIR;
+      llvm::Function *bar = module->getFunction("bar");
+      ASSERT_TRUE(bar);
+      ASSERT_FALSE(bar->empty());
+    }
+    ASSERT_TRUE(object != std::nullopt);
+    ASSERT_TRUE(!object->empty());
+  }
+}