[MLIR][NVVM] Add dumpISA and dumpMachineISA Flags

Currently, dumping virtual and machine-level ISA is restricted to debug builds, making it unavailable in release builds. However, the ability to view these ISAs can be valuable even in release builds or with production compilers. For instance, `nvcc` provides similar functionality, making it a useful reference.

This PR introduces `dumpISA` and `dumpMachineISA` flags to the `GpuModuleToBinaryPass`.

Additionally, it adds `dump-ptx` and `dump-sass` flags to the `GPUToNVVMPipelineOptions`.

Author: grypp

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -52,7 +52,8 @@ class TargetOptions {
       StringRef toolkitPath = {}, ArrayRef<std::string> linkFiles = {},
       StringRef cmdOptions = {},
       CompilationTarget compilationTarget = getDefaultCompilationTarget(),
-      function_ref<SymbolTable *()> getSymbolTableCallback = {});
+      function_ref<SymbolTable *()> getSymbolTableCallback = {},
+      bool dumpISA = false, bool dumpMachineISA = false);
 
   /// Returns the typeID.
   TypeID getTypeID() const;
@@ -66,6 +67,12 @@ class TargetOptions {
   /// Returns the command line options.
   StringRef getCmdOptions() const;
 
+  /// Returns the dump-isa command line options.
+  bool getDumpISA() const;
+
+  /// Returns the dump-machine-isa command line options.
+  bool getDumpMachineISA() const;
+
   /// Returns a tokenization of the command line options.
   std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
   tokenizeCmdOptions() const;
@@ -90,7 +97,8 @@ class TargetOptions {
       TypeID typeID, StringRef toolkitPath = {},
       ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
       CompilationTarget compilationTarget = getDefaultCompilationTarget(),
-      function_ref<SymbolTable *()> getSymbolTableCallback = {});
+      function_ref<SymbolTable *()> getSymbolTableCallback = {},
+      bool dumpISA = false, bool dumpMachineISA = false);
 
   /// Path to the target toolkit.
   std::string toolkitPath;
@@ -102,6 +110,12 @@ class TargetOptions {
   /// process.
   std::string cmdOptions;
 
+  /// An optional flag to dump generated ISA.
+  bool dumpISA = false;
+
+  /// An optional flag to dump generated and disassembled machine ISA.
+  bool dumpMachineISA = false;
+
   /// Compilation process target format.
   CompilationTarget compilationTarget;
 

mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h

@@ -53,6 +53,12 @@ struct GPUToNVVMPipelineOptions
           "Whether to use the bareptr calling convention on the host (warning "
           "this should be false until the GPU layering is fixed)"),
       llvm::cl::init(false)};
+  PassOptions::Option<bool> dumpPtx{
+      *this, "dump-ptx", llvm::cl::desc("Dumps PTX code to the error output"),
+      llvm::cl::init(false)};
+  PassOptions::Option<bool> dumpSass{
+      *this, "dump-sass", llvm::cl::desc("Dumps SASS code to the error output"),
+      llvm::cl::init(false)};
 };
 
 //===----------------------------------------------------------------------===//

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

@@ -95,7 +95,15 @@ 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<"dumpISA", "dump-isa", "bool",
+           /*default=*/"false",
+           "Dumps generated ISA to the error output.">,
+    Option<"dumpMachineISA", "dump-machine-isa", "bool",
+           /*default=*/"false",
+           "Dumps the generated machine-level ISA to the error output. "
+           "If the generated ISA is virtual, it instead dumps the"
+           "machine-level equivalent.">
   ];
 }
 

mlir/include/mlir/Target/LLVM/NVVM/Utils.h

@@ -54,7 +54,7 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   LogicalResult appendStandardLibs();
 
   /// Loads the bitcode files in `fileList`.
-  virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
+  std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
   loadBitcodeFiles(llvm::Module &module) override;
 
 protected:

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -2302,18 +2302,26 @@ KernelMetadataAttr KernelTableAttr::lookup(StringAttr key) const {
 TargetOptions::TargetOptions(
     StringRef toolkitPath, ArrayRef<std::string> linkFiles,
     StringRef cmdOptions, CompilationTarget compilationTarget,
-    function_ref<SymbolTable *()> getSymbolTableCallback)
+    function_ref<SymbolTable *()> getSymbolTableCallback, bool dumpISA,
+    bool dumpMachineISA)
     : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,
-                    cmdOptions, compilationTarget, getSymbolTableCallback) {}
+                    cmdOptions, compilationTarget, getSymbolTableCallback,
+                    dumpISA, dumpMachineISA) {}
 
 TargetOptions::TargetOptions(
     TypeID typeID, StringRef toolkitPath, ArrayRef<std::string> linkFiles,
     StringRef cmdOptions, CompilationTarget compilationTarget,
-    function_ref<SymbolTable *()> getSymbolTableCallback)
+    function_ref<SymbolTable *()> getSymbolTableCallback, bool dumpISA,
+    bool dumpMachineISA)
     : toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
-      cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget),
+      cmdOptions(cmdOptions.str()), dumpISA(dumpISA),
+      dumpMachineISA(dumpMachineISA), compilationTarget(compilationTarget),
       getSymbolTableCallback(getSymbolTableCallback), typeID(typeID) {}
 
+bool TargetOptions::getDumpISA() const { return dumpISA; }
+
+bool TargetOptions::getDumpMachineISA() const { return dumpMachineISA; }
+
 TypeID TargetOptions::getTypeID() const { return typeID; }
 
 StringRef TargetOptions::getToolkitPath() const { return toolkitPath; }

mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp

@@ -95,6 +95,8 @@ void buildHostPostPipeline(OpPassManager &pm,
 
   GpuModuleToBinaryPassOptions gpuModuleToBinaryPassOptions;
   gpuModuleToBinaryPassOptions.compilationTarget = options.cubinFormat;
+  gpuModuleToBinaryPassOptions.dumpISA = options.dumpPtx;
+  gpuModuleToBinaryPassOptions.dumpMachineISA = options.dumpSass;
   pm.addPass(createGpuModuleToBinaryPass(gpuModuleToBinaryPassOptions));
   pm.addPass(createConvertMathToLLVMPass());
   pm.addPass(createCanonicalizerPass());

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

@@ -70,7 +70,7 @@ void GpuModuleToBinaryPass::runOnOperation() {
   };
 
   TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, *targetFormat,
-                              lazyTableBuilder);
+                              lazyTableBuilder, dumpISA, dumpMachineISA);
   if (failed(transformGpuModulesToBinaries(
           getOperation(), OffloadingLLVMTranslationAttrInterface(nullptr),
           targetOptions)))

mlir/lib/Target/LLVM/NVVM/Target.cpp

@@ -292,8 +292,8 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
     return std::nullopt;
   TmpFile cubinFile;
   if (createFatbin) {
-    Twine cubinFilename = ptxFile->first + ".cubin";
-    cubinFile = TmpFile(cubinFilename.str(), llvm::FileRemover(cubinFilename));
+    std::string cubinFilename = (ptxFile->first + ".cubin").str();
+    cubinFile = TmpFile(cubinFilename, llvm::FileRemover(cubinFilename));
   } else {
     cubinFile.first = binaryFile->first;
   }
@@ -402,8 +402,8 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
                                 /*MemoryLimit=*/0,
                                 /*ErrMsg=*/&message))
     return emitLogError("`ptxas`");
-#define DEBUG_TYPE "dump-sass"
-  LLVM_DEBUG({
+
+  if (targetOptions.getDumpMachineISA()) {
     std::optional<std::string> nvdisasm = findTool("nvdisasm");
     SmallVector<StringRef> nvdisasmArgs(
         {StringRef("nvdisasm"), StringRef(cubinFile.first)});
@@ -417,11 +417,10 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> logBuffer =
         llvm::MemoryBuffer::getFile(logFile->first);
     if (logBuffer && !(*logBuffer)->getBuffer().empty()) {
-      llvm::dbgs() << "Output:\n" << (*logBuffer)->getBuffer() << "\n";
-      llvm::dbgs().flush();
+      llvm::errs() << "Output:\n" << (*logBuffer)->getBuffer() << "\n";
+      llvm::errs().flush();
     }
-  });
-#undef DEBUG_TYPE
+  }
 
   // Invoke `fatbin`.
   message.clear();
@@ -572,12 +571,13 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
     getOperation().emitError() << "Failed translating the module to ISA.";
     return std::nullopt;
   }
-#define DEBUG_TYPE "serialize-to-isa"
-  LLVM_DEBUG({
-    llvm::dbgs() << "PTX for module: " << getOperation().getNameAttr() << "\n";
-    llvm::dbgs() << *serializedISA << "\n";
-    llvm::dbgs().flush();
-  });
+  if (targetOptions.getDumpISA()) {
+    llvm::errs() << "// Generated PTX for module: "
+                 << getOperation().getNameAttr() << "\n";
+    llvm::errs() << *serializedISA << "\n";
+    llvm::errs().flush();
+  }
+
 #undef DEBUG_TYPE
 
   // Return PTX if the compilation target is `assembly`.

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

@@ -430,13 +430,12 @@ std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
     getOperation().emitError() << "failed translating the module to ISA";
     return std::nullopt;
   }
-#define DEBUG_TYPE "serialize-to-isa"
-  LLVM_DEBUG({
+  if (targetOptions.getDumpISA()) {
     llvm::dbgs() << "ISA for module: "
                  << cast<gpu::GPUModuleOp>(getOperation()).getNameAttr() << "\n"
                  << *serializedISA << "\n";
-  });
-#undef DEBUG_TYPE
+  }
+
   // Return ISA assembly code if the compilation target is assembly.
   if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Assembly)
     return SmallVector<char, 0>(serializedISA->begin(), serializedISA->end());

mlir/test/Integration/GPU/CUDA/dump-ptx.mlir

@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s \
-// RUN:  | mlir-opt -gpu-lower-to-nvvm-pipeline -debug-only=serialize-to-isa \
+// RUN:  | mlir-opt -gpu-lower-to-nvvm-pipeline="dump-ptx=true" \
 // RUN:  2>&1 | FileCheck %s
 
 // CHECK: Generated by LLVM NVPTX Back-End