[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

32

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

87

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

32

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.

87

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

87

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

?

105–106

Terminate all comments with a full stop.

[bondhugula]

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

87

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

87

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

87

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

87

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> 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())};
 };
 } // 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);
+
 /// 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 <string>
+
+#define DEBUG_TYPE "serialize-to-blob"
+
 using namespace mlir;
 
 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) {}
 
 Optional<std::string>
 gpu::SerializeToBlobPass::translateToISA(llvm::Module &llvmModule,
                                          llvm::TargetMachine &targetMachine) {
void gpu::SerializeToBlobPass::runOnOperation() {
   Optional<std::string> maybeTargetISA =
       translateToISA(*llvmModule, *targetMachine);
 
   if (!maybeTargetISA.hasValue())
     return signalPassFailure();
 
   std::string targetISA = std::move(maybeTargetISA.getValue());
 
+  LLVM_DEBUG({
+    llvm::dbgs() << "ISA for module: " << getOperation().getNameAttr() << "\n";
+    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

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) {
+  return std::make_unique<SerializeToHsacoPass>(triple, arch, features,
+                                                optLevel);
+}
+
 #else  // MLIR_GPU_TO_HSACO_PASS_ENABLE
 void mlir::registerGpuSerializeToHsacoPass() {}
 #endif // MLIR_GPU_TO_HSACO_PASS_ENABLE