[Mlir-commits] [mlir] [mlir][fix] Add callback functions for ModuleToObject (PR #116916)

[Mehdi Amini]

================
@@ -153,3 +155,91 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(TargetAttrAPI)) {
   // `serializeToObject`.
   ASSERT_TRUE(properties.contains("serialize_attr"));
 }
+
+// Test callback function invoked with initial LLVM IR
+TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithInitialLLVMIR)) {
+  MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  Builder builder(&context);
+  IntegerAttr target = builder.getI32IntegerAttr(0);
+  auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
+
+  std::string initialLLVMIR;
+  auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
+    llvm::raw_string_ostream ros(initialLLVMIR);
+    module.print(ros, nullptr);
+  };
+
+  gpu::TargetOptions opts(
+      {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {},
+      initialCallback);
+  std::optional<SmallVector<char, 0>> serializedBinary =
+      targetAttr.serializeToObject(*module, opts);
+
+  ASSERT_TRUE(serializedBinary != std::nullopt);
+  ASSERT_TRUE(!serializedBinary->empty());
+  ASSERT_TRUE(!initialLLVMIR.empty());
+}
+
+// Test callback function invoked with linked LLVM IR
+TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithLinkedLLVMIR)) {
+  MLIRContext context(registry);
+  context.loadAllAvailableDialects();
----------------
joker-eph wrote:

```suggestion
```

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