[mlir][gpu] gpu-module-to-binary: add option to dump intermediate files

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

@@ -113,7 +113,9 @@ def GpuModuleToBinaryPass
     Option<"compilationTarget", "format", "std::string", [{"fatbin"}],
            "The target representation of the compilation process.">,
     Option<"elfSection", "section", "std::string", [{""}],
-           "ELF section where binary is to be located.">
+           "ELF section where binary is to be located.">,
+    Option<"dumpIntermediates", "dump-intermediates", "std::string", [{""}],
+           "Directory to dump intermediate artifacts (LLVM IR, device assembly).">
   ];
 }
 

mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp

@@ -17,6 +17,12 @@
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/DebugLog.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/ToolOutputFile.h"
+
+#define DEBUG_TYPE "gpu-module-to-binary"
 
 using namespace mlir;
 using namespace mlir::gpu;
@@ -26,6 +32,27 @@ namespace mlir {
 #include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
 } // namespace mlir
 
+static void dumpToFile(StringRef dumpDir, const llvm::Twine &filename,
+                       function_ref<void(llvm::raw_ostream &)> writeContent) {
+  if (dumpDir.empty())
+    return;
+
+  llvm::SmallString<128> path(dumpDir);
+  llvm::sys::path::append(path, filename);
+
+  std::error_code ec;
+  llvm::ToolOutputFile output(path, ec, llvm::sys::fs::OF_None);
+  if (ec) {
+    LDBG() << "Failed to create file '" << path << "': " << ec.message()
+           << "\n";
+    return;
+  }
+
+  writeContent(output.os());
+  output.keep();
+  LDBG() << "Dumped intermediate to: " << path << "\n";
+}
+
 namespace {
 class GpuModuleToBinaryPass
     : public impl::GpuModuleToBinaryPassBase<GpuModuleToBinaryPass> {
@@ -64,8 +91,42 @@ void GpuModuleToBinaryPass::runOnOperation() {
   SmallVector<Attribute> librariesToLink;
   for (const std::string &path : linkFiles)
     librariesToLink.push_back(StringAttr::get(&getContext(), path));
+
+  // Create dump directory if specified.
+  if (!dumpIntermediates.empty()) {
+    if (std::error_code ec =
+            llvm::sys::fs::create_directories(dumpIntermediates)) {
+      getOperation()->emitError() << "Failed to create dump directory '"
+                                  << dumpIntermediates << "': " << ec.message();
+      return signalPassFailure();
+    }
+  }
+
+  // Create callbacks for dumping intermediate artifacts if requested.
+  auto initialIRCallback = [&](llvm::Module &mod) {
+    dumpToFile(dumpIntermediates, mod.getName() + ".initial.ll",
+               [&](llvm::raw_ostream &os) { mod.print(os, nullptr); });
+  };
+
+  auto linkedIRCallback = [&](llvm::Module &mod) {
+    dumpToFile(dumpIntermediates, mod.getName() + ".linked.ll",
+               [&](llvm::raw_ostream &os) { mod.print(os, nullptr); });
+  };
+
+  auto optimizedIRCallback = [&](llvm::Module &mod) {
+    dumpToFile(dumpIntermediates, mod.getName() + ".opt.ll",
+               [&](llvm::raw_ostream &os) { mod.print(os, nullptr); });
+  };
+
+  auto isaCallback = [&](StringRef isa) {
+    dumpToFile(dumpIntermediates, "kernel.isa",
+               [&](llvm::raw_ostream &os) { os << isa; });
+  };
+
   TargetOptions targetOptions(toolkitPath, librariesToLink, cmdOptions,
-                              elfSection, *targetFormat, lazyTableBuilder);
+                              elfSection, *targetFormat, lazyTableBuilder,
+                              initialIRCallback, linkedIRCallback,
+                              optimizedIRCallback, isaCallback);
   if (failed(transformGpuModulesToBinaries(
           getOperation(), OffloadingLLVMTranslationAttrInterface(nullptr),
           targetOptions)))

mlir/lib/Target/LLVM/ModuleToObject.cpp

@@ -143,7 +143,13 @@ LogicalResult ModuleToObject::loadBitcodeFilesFromList(
 
 std::unique_ptr<llvm::Module>
 ModuleToObject::translateToLLVMIR(llvm::LLVMContext &llvmContext) {
-  return translateModuleToLLVMIR(&getOperation(), llvmContext);
+  Operation &op = getOperation();
+  StringRef name = "LLVMDialectModule";
+  // Try to get nicer name from the operation.
+  if (auto symOp = dyn_cast<SymbolOpInterface>(op);
+      symOp && symOp.getNameAttr())
+    name = symOp.getNameAttr().getValue();
+  return translateModuleToLLVMIR(&op, llvmContext, name);
 }
 
 LogicalResult

mlir/lib/Target/LLVM/ROCDL/Target.cpp

@@ -95,7 +95,11 @@ SerializeGPUModuleBase::SerializeGPUModuleBase(
     Operation &module, ROCDLTargetAttr target,
     const gpu::TargetOptions &targetOptions)
     : ModuleToObject(module, target.getTriple(), target.getChip(),
-                     target.getFeatures(), target.getO()),
+                     target.getFeatures(), target.getO(),
+                     targetOptions.getInitialLlvmIRCallback(),
+                     targetOptions.getLinkedLlvmIRCallback(),
+                     targetOptions.getOptimizedLlvmIRCallback(),
+                     targetOptions.getISACallback()),
       target(target), toolkitPath(targetOptions.getToolkitPath()),
       librariesToLink(targetOptions.getLibrariesToLink()) {
 
@@ -428,6 +432,10 @@ std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
     getOperation().emitError() << "failed translating the module to ISA";
     return std::nullopt;
   }
+
+  if (isaCallback)
+    isaCallback(serializedISA.value());
+
 #define DEBUG_TYPE "serialize-to-isa"
   LLVM_DEBUG({
     llvm::dbgs() << "ISA for module: "

mlir/test/Dialect/GPU/module-to-binary-nvvm-intermediates.mlir

@@ -0,0 +1,17 @@
+// REQUIRES: host-supports-nvptx
+// RUN: rm -rf %t
+// RUN: mlir-opt %s --gpu-module-to-binary='format=isa dump-intermediates=%t' | FileCheck %s
+// RUN: test -f %t/kernel_module.initial.ll
+// RUN: test -f %t/kernel_module.linked.ll
+// RUN: test -f %t/kernel_module.opt.ll
+// RUN: test -f %t/kernel.isa
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: gpu.binary @kernel_module
+
+  gpu.module @kernel_module [#nvvm.target<chip = "sm_70">] {
+    llvm.func @kernel(%arg0: f32) {
+      llvm.return
+    }
+  }
+}

mlir/test/Dialect/GPU/module-to-binary-rocdl-intermediates.mlir

@@ -0,0 +1,17 @@
+// REQUIRES: host-supports-amdgpu
+// RUN: rm -rf %t
+// RUN: mlir-opt %s --gpu-module-to-binary='format=isa dump-intermediates=%t' | FileCheck %s
+// RUN: test -f %t/kernel_module.initial.ll
+// RUN: test -f %t/kernel_module.linked.ll
+// RUN: test -f %t/kernel_module.opt.ll
+// RUN: test -f %t/kernel.isa
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: gpu.binary @kernel_module
+
+  gpu.module @kernel_module [#rocdl.target<chip = "gfx942">] {
+    llvm.func @kernel(%arg0: f32) {
+      llvm.return
+    }
+  }
+}