[Mlir-commits] [mlir] [MLIR] Add a BlobAttr interface for attribute to wrap arbitrary content and use it as linkLibs for ModuleToObject (PR #120116)

[Mehdi Amini]

================
@@ -215,3 +218,72 @@ 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, {},
+      mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {}, {},
+      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());
+    ASSERT_TRUE(
+        StringRef(linkedLLVMIR).contains("define internal void @bar() {"));
----------------
joker-eph wrote:

True, but isn't this on the same level as all the FileCheck CHECK in LLVM?

https://github.com/llvm/llvm-project/pull/120116