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

Add option to specify dir to dump inetrmediate files during gpu binaries generation for debug.

Also fix ROCDL lowering bug where callbacks weren't propagated.

Author: Hardcode84

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/Debug.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) {
+    LLVM_DEBUG(llvm::dbgs() << "Failed to create file '" << path
+                            << "': " << ec.message() << "\n");
+    return;
+  }
+
+  writeContent(output.os());
+  output.keep();
+  LLVM_DEBUG(llvm::dbgs() << "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 &module) {
+    dumpToFile(dumpIntermediates, module.getName() + ".initial.ll",
+               [&](llvm::raw_ostream &os) { module.print(os, nullptr); });
+  };
+
+  auto linkedIRCallback = [&](llvm::Module &module) {
+    dumpToFile(dumpIntermediates, module.getName() + ".linked.ll",
+               [&](llvm::raw_ostream &os) { module.print(os, nullptr); });
+  };
+
+  auto optimizedIRCallback = [&](llvm::Module &module) {
+    dumpToFile(dumpIntermediates, module.getName() + ".opt.ll",
+               [&](llvm::raw_ostream &os) { module.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,11 @@ LogicalResult ModuleToObject::loadBitcodeFilesFromList(
 
 std::unique_ptr<llvm::Module>
 ModuleToObject::translateToLLVMIR(llvm::LLVMContext &llvmContext) {
-  return translateModuleToLLVMIR(&getOperation(), llvmContext);
+  Operation &op = getOperation();
+  // Try to get nicer name from the operation.
+  auto nameAttr = op.getAttrOfType<StringAttr>("sym_name");
+  StringRef name = nameAttr ? nameAttr.getValue() : "LLVMDialectModule";
+  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 || true
+// 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 || true
+// 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
+    }
+  }
+}