[MLIR][GPU] Add debug output to enable dumping GPU assembly

[krzysz00]

Summary

• Set the DEBUG_TYPE of SerializeToBlob to serialize-to-blob
• Add debug output to print the assembly or PTX for GPU modules before they are assembled and linked

Note that, as SerializeToBlob is a superclass of SerializeToCubin and
SerializeToHsaco, --debug-only=serialize-to-blom will dump the
intermediate compiler result for both of these passes.

In addition, if LLVM options such as --stop-after are used to control
the GPU kernel compilation process, the debug output will contain the
appropriate intermediate IR.

[krzysz00]

(Looks like this also pulled in createSerializeToHsacoPass which wasn't in the upstream codebase, apparently)

[mehdi_amini]

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

31

In general we rely on people to disable threading for debugging. A global mutex for debugging purpose seems overkill to me.

90

Can we just use the usual Debugging facilities?
LLVM_DEBUG({

  llvm::dbgs() << targetISA << "\n";
  llvm::dbgs().flush();
}

);

[krzysz00]

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

31

Fair enough, though this lock came from

1. The fact that, when I was debugging what LLVM was doing, I couldn't for the life of me get threading off in the interests of usable --print-after-all output
2. That, without a lock and with manually patching in an llvm::outs() << isa << "\n"; (though in a slightly different spot) I'd get the output of the assembled program interleaved with the assembly

This might just be a case of me being overly paranoid wrt locks - I can remove.

90

Can we just use the usual Debugging facilities?

Unless there's a way to special-case these even further so that even --debug doesn't turn on the output ... probably not without killing the usability of --debug output for everything else, since you'll have many screenfuls of raw assembly in the middle of the usually much tamer debug logs.

(On that note, the Tensorflow folks also use debug options for this type of print

[herhut]

Thanks for adding this. I have wanted this a couple of times, too, but never went as far as creating a diff for it.

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

90

You could set the LLVM_DEBUG_TYPE to something meaningful, e.g., 'serialize-to-blob' and then '--debug-only=serialize-to-blob' would do what you want. Using the general --debug is not helpful IMHO because it already creates too much output.

[bondhugula]

This will be quite useful - thanks!

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

92

Slight rephrase to:

Dump final generated instructions ... for a kernel to the debug stream

?

111–112

Terminate all comments with a full stop.

[bondhugula]

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

90

I too think -debug-only=serialize-to-blob should be fine. dump-asm isn't really a pass functionality option but a debugging one. It's useful either way though.

[mehdi_amini]

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

90

Alternatively, you can make it an option that would attach the ASM as an attribute to the op instead of printing it.

[bondhugula]

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

90

That would mess up the first line and also make it hard to navigate/find, etc. It's better in a way if the MLIR and the ASM are separate -- but for multiple cubins, it'll be good to have an easy way to identify which asm is for which.

[krzysz00]

• Adress review feedback, will reword commit message later

[mehdi_amini]

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

90

The tradeoff to me is more about whether you need this in production or not.

[krzysz00]

I agree about the LLVM_DEBUG point, and have edited the code accordingly.

This definitely isn't something I'd see coming around in production, but it is pretty helpful when debugging weird compiler bugs.

[krzysz00]

I agree with the LLVM_DEBUG arguments and have modified the pass accordingly.

I definitely don't see this being something that gets used in production, but it is pretty helpful for debugging weird compiler failures and the like.

[herhut]

Thanks.

[herhut]

Please fix the commit message though. I just noticed this now.

[krzysz00]

Change commit message

[krzysz00]

Add tags to commit message

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

protected:
   Option<std::string> chip{*this, "chip",
                            ::llvm::cl::desc("Target architecture")};
   Option<std::string> features{*this, "features",
                                ::llvm::cl::desc("Target features")};
   Option<std::string> gpuBinaryAnnotation{
       *this, "gpu-binary-annotation",
       llvm::cl::desc("Annotation attribute string for GPU binary"),
       llvm::cl::init(getDefaultGpuBinaryAnnotation())};
+  Option<bool> dumpAsm{
+      *this, "dump-asm",
+      llvm::cl::desc("Whether the final generated instructions or intermediate "
+                     "IR (if stopping early) for a kernel should be dumped to "
+                     "the debug stream"),
+      llvm::cl::init(false)};
 };
 } // namespace gpu
 
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Register pass to serialize GPU kernel functions to a CUBIN binary
 /// annotation.
 void registerGpuSerializeToCubinPass();
 
 /// Register pass to serialize GPU kernel functions to a HSAco binary
 /// annotation.
 void registerGpuSerializeToHsacoPass();
 
+/// Create an instance of the GPU kernel function to HSAco binary serialization
+/// pass
+std::unique_ptr<Pass> createGpuSerializeToHsacoPass(StringRef triple,
+                                                    StringRef arch,
+                                                    StringRef features,
+                                                    int optLevel,
+                                                    bool dumpAsm = false);
+
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/GPU/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_GPU_PASSES_H_

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

 #include "mlir/Pass/Pass.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 
+#include <mutex>
+#include <string>
+
 using namespace mlir;
 
+// Ensure multiple threads don't try to simultaneously dump the assembly for
+// separate modules
+static std::mutex dumpAsmLock;
+
 std::string gpu::getDefaultGpuBinaryAnnotation() { return "gpu.binary"; }
 
 gpu::SerializeToBlobPass::SerializeToBlobPass(TypeID passID)
     : OperationPass<gpu::GPUModuleOp>(passID) {}
 
 gpu::SerializeToBlobPass::SerializeToBlobPass(const SerializeToBlobPass &other)
     : OperationPass<gpu::GPUModuleOp>(other) {}
 
void gpu::SerializeToBlobPass::runOnOperation() {
   Optional<std::string> maybeTargetISA =
       translateToISA(*llvmModule, *targetMachine);
 
   if (!maybeTargetISA.hasValue())
     return signalPassFailure();
 
   std::string targetISA = std::move(maybeTargetISA.getValue());
 
+  if (dumpAsm.getValue()) {
+    const std::lock_guard<std::mutex> lock(dumpAsmLock);
+    llvm::dbgs() << targetISA << "\n";
+    llvm::dbgs().flush();
+  }
+
   // Serialize the target ISA.
   std::unique_ptr<std::vector<char>> blob = serializeISA(targetISA);
   if (!blob)
     return signalPassFailure();
 
   // Add the blob as module attribute.
   auto attr =
       StringAttr::get(&getContext(), StringRef(blob->data(), blob->size()));

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

 
 using namespace mlir;
 
 namespace {
 class SerializeToHsacoPass
     : public PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass> {
 public:
   SerializeToHsacoPass(StringRef triple, StringRef arch, StringRef features,
-                       int optLevel);
+                       int optLevel, bool dumpAsm = false);
   SerializeToHsacoPass(const SerializeToHsacoPass &other);
   StringRef getArgument() const override { return "gpu-to-hsaco"; }
   StringRef getDescription() const override {
     return "Lower GPU kernel function to HSACO binary annotations";
   }
 
 protected:
   Option<int> optLevel{

 // Sets the 'option' to 'value' unless it already has a value.
 static void maybeSetOption(Pass::Option<std::string> &option,
                            function_ref<std::string()> getValue) {
   if (!option.hasValue())
     option = getValue();
 }
 
 SerializeToHsacoPass::SerializeToHsacoPass(StringRef triple, StringRef arch,
-                                           StringRef features, int optLevel) {
+                                           StringRef features, int optLevel,
+                                           bool dumpAsm) {
   maybeSetOption(this->triple, [&triple] { return triple.str(); });
   maybeSetOption(this->chip, [&arch] { return arch.str(); });
   maybeSetOption(this->features, [&features] { return features.str(); });
   if (this->optLevel.getNumOccurrences() == 0)
     this->optLevel.setValue(optLevel);
+  if (this->dumpAsm.getNumOccurrences() == 0 && dumpAsm)
+    this->dumpAsm.setValue(dumpAsm);
 }
 
 void SerializeToHsacoPass::getDependentDialects(
     DialectRegistry &registry) const {
   registerROCDLDialectTranslation(registry);
   gpu::SerializeToBlobPass::getDependentDialects(registry);
 }
 
PassRegistration<SerializeToHsacoPass> registerSerializeToHSACO(
         LLVMInitializeAMDGPUTarget();
         LLVMInitializeAMDGPUTargetInfo();
         LLVMInitializeAMDGPUTargetMC();
 
         return std::make_unique<SerializeToHsacoPass>("amdgcn-amd-amdhsa", "",
                                                       "", 2);
       });
 }
+
+/// Create an instance of the GPU kernel function to HSAco binary serialization
+/// pass
+std::unique_ptr<Pass> mlir::createGpuSerializeToHsacoPass(StringRef triple,
+                                                          StringRef arch,
+                                                          StringRef features,
+                                                          int optLevel,
+                                                          bool dumpASM) {
+  return std::make_unique<SerializeToHsacoPass>(triple, arch, features,
+                                                optLevel, dumpASM);
+}
+
 #else  // MLIR_GPU_TO_HSACO_PASS_ENABLE
 void mlir::registerGpuSerializeToHsacoPass() {}
 #endif // MLIR_GPU_TO_HSACO_PASS_ENABLE

mlir/test/Conversion/GPUToROCm/lower-rocdl-kernel-to-hsaco.mlir

 // RUN: mlir-opt %s --test-gpu-to-hsaco | FileCheck %s
+// RUN: mlir-opt %s --test-gpu-to-hsaco=dump-asm=true 2>&1 |\
+// RUN:   FileCheck %s --check-prefix=CHECK-ASM
 
 // CHECK: gpu.module @foo attributes {gpu.binary = "HSACO"}
+// CHECK-ASM: .globl kernel
 gpu.module @foo {
   llvm.func @kernel(%arg0 : f32, %arg1 : !llvm.ptr<f32>)
     // CHECK: attributes  {gpu.kernel}
     attributes  { gpu.kernel } {
     llvm.return
   }
 }
 
 // CHECK: gpu.module @bar attributes {gpu.binary = "HSACO"}
 gpu.module @bar {
   // CHECK: func @kernel_a
   llvm.func @kernel_a()
     attributes  { gpu.kernel } {
     llvm.return
   }
 
   // CHECK: func @kernel_b
   llvm.func @kernel_b()
     attributes  { gpu.kernel } {
     llvm.return
   }
 }
+// CHECK-ASM: amdhsa.target: