[Mlir-commits] [mlir] [mlir] Add existing callback functions to GpuModuleToBinary pass (PR #117694)

[Zichen Lu]

https://github.com/MikaOvO created https://github.com/llvm/llvm-project/pull/117694

After [this PR](https://github.com/llvm/llvm-project/pull/116916), there are callback functions for LLVM IR and ISA. Downstream users usually use the `GpuModuleToBinary` pass to complete `ModuleToObject`. 

>From 4fb3b44840385fdb300cc8fe239c041dbd09f1c7 Mon Sep 17 00:00:00 2001
From: Zichen Lu <mikaovo2000 at gmail.com>
Date: Tue, 26 Nov 2024 17:58:19 +0800
Subject: [PATCH] [mlir] Add existing callback functions to GpuModuleToBinary
 pass

---
 mlir/include/mlir/Dialect/GPU/Transforms/Passes.td | 14 +++++++++++++-
 mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp |  6 ++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

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)))