[Mlir-commits] [mlir] [MLIR] Add callback functions for ModuleToObject (PR #116007)

[Mehdi Amini]

================
@@ -156,3 +156,74 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
     ASSERT_TRUE(!object->empty());
   }
 }
+
+// Test callback functions invoked with LLVM IR and ISA.
+TEST_F(MLIRTargetLLVMNVVM,
+       SKIP_WITHOUT_NVPTX(CallbackInvokedWithLLVMIRAndISA)) {
+  if (!hasPtxas())
+    GTEST_SKIP() << "PTXAS compiler not found, skipping test.";
+
+  MLIRContext context(registry);
+
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+
+  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
+
+  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
+  ASSERT_TRUE(!!serializer);
+
+  std::string initialLLVMIR;
+  auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
+    std::ostringstream oss;
+    llvm::raw_os_ostream ros(oss);
+    ros.flush();
+    module.print(ros, nullptr);
+    initialLLVMIR = oss.str();
+  };
+
+  std::string linkedLLVMIR;
+  auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
+    std::ostringstream oss;
+    llvm::raw_os_ostream ros(oss);
+    ros.flush();
+    module.print(ros, nullptr);
+    linkedLLVMIR = oss.str();
+  };
+
+  std::string optimizedLLVMIR;
+  auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {
+    std::ostringstream oss;
+    llvm::raw_os_ostream ros(oss);
+    ros.flush();
+    module.print(ros, nullptr);
+    optimizedLLVMIR = oss.str();
----------------
joker-eph wrote:

The `ros.flush()` is inserted before the `module.print()`, that does not seem right...

This should work instead: 

```suggestion
    llvm::raw_string_ostream rso(optimizedLLVMIR);
    module.print(rso, nullptr);
```

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