[mlir] Add existing callback functions to GpuModuleToBinary pass

[MikaOvO]

After this PR, there are callback functions for LLVM IR and ISA. Downstream users usually use the GpuModuleToBinary pass to complete ModuleToObject.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Zichen Lu (MikaOvO)

Changes

After this PR, there are callback functions for LLVM IR and ISA. Downstream users usually use the GpuModuleToBinary pass to complete ModuleToObject.

---

Full diff: https://github.com/llvm/llvm-project/pull/117694.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/GPU/Transforms/Passes.td (+13-1)
• (modified) mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp (+4-2)

diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
index 4a9ddafdd177d2..313c58715ed511 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
@@ -95,7 +95,19 @@ def GpuModuleToBinaryPass
     Option<"cmdOptions", "opts", "std::string", [{""}],
            "Command line options to pass to the tools.">,
     Option<"compilationTarget", "format", "std::string", [{"fatbin"}],
-           "The target representation of the compilation process.">
+           "The target representation of the compilation process.">,
+    Option<"initialLlvmIRCallback", "initialLlvmIRCallback", 
+           "llvm::function_ref<void(llvm::Module &)>", "nullptr",
+           "Callback invoked with the initial LLVM IR for the device module.">,
+    Option<"linkedLlvmIRCallback", "linkedLlvmIRCallback", 
+           "llvm::function_ref<void(llvm::Module &)>", "nullptr",
+           "Callback invoked with LLVM IR for the device module after linking the device libraries.">,
+    Option<"optimizedLlvmIRCallback", "optimizedLlvmIRCallback", 
+           "llvm::function_ref<void(llvm::Module &)>", "nullptr",
+           "Callback invoked with LLVM IR for the device module after LLVM optimizations but before codegen.">,
+    Option<"isaCallback", "isaCallback", 
+           "llvm::function_ref<void(llvm::StringRef)>", "nullptr",
+           "Callback invoked with the target ISA for the device, for example PTX assembly.">
   ];
 }
 
diff --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
index 86a3b4780e88ce..ce399067c1b04e 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
@@ -69,8 +69,10 @@ void GpuModuleToBinaryPass::runOnOperation() {
     return &parentTable.value();
   };
 
-  TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, *targetFormat,
-                              lazyTableBuilder);
+  TargetOptions targetOptions(
+      toolkitPath, linkFiles, cmdOptions, *targetFormat, lazyTableBuilder,
+      initialLlvmIRCallback.getValue(), linkedLlvmIRCallback.getValue(),
+      optimizedLlvmIRCallback.getValue(), isaCallback.getValue());
   if (failed(transformGpuModulesToBinaries(
           getOperation(), OffloadingLLVMTranslationAttrInterface(nullptr),
           targetOptions)))

[MikaOvO]

I have tested it locally, is there a unittest here for me to modify? I didn't find it.

[joker-eph]

That seems problematic to me, because passes options are supposed to round-trip to text.
What is printed in this case if we print the pass-manager?

[MikaOvO]

Hi, do you mean:

  std::string ss;
  llvm::raw_string_ostream rso(ss);
  pm.printAsTextualPipeline(rso);
  llvm::outs() << ss << "\n";

I got:

===unknown data value for option
UNREACHABLE executed at ./llvm-solid/mlir/include/mlir/Pass/PassOptions.h:168!

[joker-eph]

Yes that was indeed my concern.
So in general we just don't expose that kind of things as pass options. Instead we make it so that the pass body is extracted as a utility, and the pass is a very thin wrapper. That way people can build their own pass downstream.

What we could do upstream though is add a "isaCallbackDumpPath" option and use this to let users provide a way to save to a file. That may not be enough for what you want to do, but you can adjust downstream with your own pass.

[MikaOvO]

Got it. Thanks!
So if I understand correctly, it's difficult to pass pointers or references in options, is it right?

If right, I think it's better to use serializer directly, and I will persuade others to believe it. (because the main reason others want me to use pass but not serializer now is to eliminate code duplication, but maybe it's difficult to write a new pass and achieve this.)

[joker-eph]

There is no reason that you can't use the pass I believe, ping me offline if you need to.

[MikaOvO]

OK, thanks!