[MLIR][GPU] Run generic LLVM optimizations when serializing (on AMD)

[krzysz00]

Summary

• Adds hooks that allow SerializeTo* passes to arbitrarily transform

the produced LLVM Module before it is passed to the code generation
passes.

• Uses these hooks within the SerializeToHsaco pass in order to run

LLVM optimizations and to set the optimization level on the
TargetMachine.

• Adds an optLevel parameter to SerializeToHsaco

Future work may include moving much of what's been added to
SerializeToHsaco to SerializeToBlob, but that would require
confirmation from the NVVM backend maintainers that it would be
appropriate to do so.

Depends on D114107

[mehdi_amini]

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

59

Will this method be overridden / specialized?

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

41

Nit: remove trivial braces

76

(same: remove trivial braces)

77

You have an extra copy here, take a reference (or std::move)?

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

174

Please use diagnostics

[krzysz00]

• Adress review comments

[krzysz00]

Commends addressed

While addressing the comments, I found this file was never being compiled, and in fact couldn't be compiled with how CMake was written. So, we now have a dependency on D114184

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

59

Not by anything we've got, and if anyone wants to specialize it in the future they could just make it virtual in the next revision. Good catch.

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

76

Is not having trivial braces the local style? I've removed them because the style guide is the style guide, but I'd like to be vaguely unhappy about this decision.

[krzysz00]

• Add comment, remove trivial braces

[krzysz00]

• Remove commits sitting in a different revision

[krzysz00]

• Adress review comments
• Add comment, remove trivial braces

[mehdi_amini]

mlir/lib/Dialect/GPU/CMakeLists.txt

165

I just pushed https://github.com/llvm/llvm-project/commit/b743850b736e4a89378be8bed61c1b3489b56d19 because this does not seem to do what we need, can you look into this? What's the difference between this and my commit?

[mehdi_amini]

mlir/lib/Dialect/GPU/CMakeLists.txt

165

Actually forget it: my patch does not fix it either!

[krzysz00]

mlir/lib/Dialect/GPU/CMakeLists.txt

165

Can I get some context as to the problem?

[mehdi_amini]

mlir/lib/Dialect/GPU/CMakeLists.txt

165

See: https://github.com/llvm/llvm-project/issues/54242

There is some sort of layering issue here as there is a dependency from a transform on the ExecutionEngine.

This is also somehow not necessary as you're just reusing a convenient wrapper to setup a LLVM pipeline here.

[krzysz00]

mlir/lib/Dialect/GPU/CMakeLists.txt

165

Could we (I can try to do it tomorrow before I leave town on Friday, but maybe it makes more sense for y'all to do it) factor out that wrapper into a separate component that both this transformation and the execution engine can depend on?

The reason we're using this wrapper is so that we don't have to maintain our own copy of how to set up opt -On, which could bitrot.

[mehdi_amini]

mlir/lib/Dialect/GPU/CMakeLists.txt

165

To me " how to set up opt -On" isn't something that belongs to MLIR in the first place. If such a helper is desirable: it should be in LLVM.

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

public:
   SerializeToBlobPass(TypeID passID);
   SerializeToBlobPass(const SerializeToBlobPass &other);
 
   void runOnOperation() final;
 
 protected:
   void getDependentDialects(DialectRegistry &registry) const override;
 
-private:
-  /// Creates the LLVM target machine to generate the ISA.
-  std::unique_ptr<llvm::TargetMachine> createTargetMachine();
+  /// Hook allowing the application of optimizations before codegen
+  /// By default, does nothing
+  virtual LogicalResult optimizeLlvm(llvm::Module &llvmModule,
+                                     llvm::TargetMachine &targetMachine);
 
   /// Translates the 'getOperation()' result to an LLVM module.
   virtual std::unique_ptr<llvm::Module>
   translateToLLVMIR(llvm::LLVMContext &llvmContext);
 
+private:
+  /// Creates the LLVM target machine to generate the ISA.
+  std::unique_ptr<llvm::TargetMachine> createTargetMachine();
+
+  /// Translates the module to ISA
+  Optional<std::string> translateToISA(llvm::Module &llvmModule,
+                                       llvm::TargetMachine &targetMachine);
+
   /// Serializes the target ISA to binary form.
   virtual std::unique_ptr<std::vector<char>>
   serializeISA(const std::string &isa) = 0;
 
 protected:
   Option<std::string> triple{*this, "triple",
                              ::llvm::cl::desc("Target triple")};
   Option<std::string> chip{*this, "chip",

mlir/lib/Dialect/GPU/CMakeLists.txt

target_include_directories(obj.MLIRGPUTransforms
     ${MLIR_SOURCE_DIR}/../lld/include
     ${HIP_PATH}/include
     ${ROCM_PATH}/include
   )
 
   target_link_libraries(MLIRGPUOps
     PRIVATE
     lldELF
+    MLIRExecutionEngine
     MLIRROCDLToLLVMIRTranslation
   )
 
   # Link lldELF also to libmlir.so. Create an alias that starts with LLVM
   # because LINK_COMPONENTS elements are implicitly prefixed with LLVM.
   add_library(LLVMAliasTolldELF ALIAS lldELF)
   set_property(GLOBAL APPEND PROPERTY MLIR_LLVM_LINK_COMPONENTS AliasTolldELF)
 
 endif()

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

 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) {}
 
-static std::string translateToISA(llvm::Module &llvmModule,
+Optional<std::string>
+gpu::SerializeToBlobPass::translateToISA(llvm::Module &llvmModule,
                                          llvm::TargetMachine &targetMachine) {
   llvmModule.setDataLayout(targetMachine.createDataLayout());
 
+  if (failed(optimizeLlvm(llvmModule, targetMachine)))
+    return llvm::None;
+
   std::string targetISA;
   llvm::raw_string_ostream stream(targetISA);
-  llvm::buffer_ostream pstream(stream);
+
   llvm::legacy::PassManager codegenPasses;
-  targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
-                                    llvm::CGFT_AssemblyFile);
+
+  { // Drop pstream after this to prevent the ISA from being stuck buffering
+    llvm::buffer_ostream pstream(stream);
+    if (targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
+                                          llvm::CGFT_AssemblyFile))
+      return llvm::None;
+
     codegenPasses.run(llvmModule);
-  return targetISA;
+  }
+  return stream.str();
 }
 
 void gpu::SerializeToBlobPass::runOnOperation() {
   // Lower the module to an LLVM IR module using a separate context to enable
   // multi-threaded processing.
   llvm::LLVMContext llvmContext;
   std::unique_ptr<llvm::Module> llvmModule = translateToLLVMIR(llvmContext);
   if (!llvmModule)
     return signalPassFailure();
 
   // Lower the LLVM IR module to target ISA.
   std::unique_ptr<llvm::TargetMachine> targetMachine = createTargetMachine();
   if (!targetMachine)
     return signalPassFailure();
 
-  std::string targetISA = translateToISA(*llvmModule, *targetMachine);
+  Optional<std::string> maybeTargetISA =
+      translateToISA(*llvmModule, *targetMachine);
+
+  if (!maybeTargetISA.hasValue())
+    return signalPassFailure();
+
+  std::string targetISA = std::move(maybeTargetISA.getValue());
 
   // 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()));
   getOperation()->setAttr(gpuBinaryAnnotation, attr);
 }
 
+LogicalResult
+gpu::SerializeToBlobPass::optimizeLlvm(llvm::Module &llvmModule,
+                                       llvm::TargetMachine &targetMachine) {
+  // TODO: If serializeToCubin ends up defining optimizations, factor them
+  // into here from SerializeToHsaco
+  return success();
+}
+
 void gpu::SerializeToBlobPass::getDependentDialects(
     DialectRegistry &registry) const {
   registerLLVMDialectTranslation(registry);
   OperationPass<gpu::GPUModuleOp>::getDependentDialects(registry);
 }
 
 std::unique_ptr<llvm::TargetMachine>
 gpu::SerializeToBlobPass::createTargetMachine() {

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

 // adds that blob as a string attribute of the module.
 //
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 
 #if MLIR_GPU_TO_HSACO_PASS_ENABLE
+#include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCCodeEmitter.h"

 #include <mutex>
 
 using namespace mlir;
 
 namespace {
 class SerializeToHsacoPass
     : public PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass> {
 public:
-  SerializeToHsacoPass(StringRef triple, StringRef arch, StringRef features);
+  SerializeToHsacoPass(StringRef triple, StringRef arch, StringRef features,
+                       int optLevel);
+  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{
+      *this, "opt-level",
+      llvm::cl::desc("Optimization level for HSACO compilation"),
+      llvm::cl::init(2)};
+
+  /// Adds LLVM optimization passes
+  LogicalResult optimizeLlvm(llvm::Module &llvmModule,
+                             llvm::TargetMachine &targetMachine) override;
+
 private:
   void getDependentDialects(DialectRegistry &registry) const override;
 
   // Serializes ROCDL to HSACO.
   std::unique_ptr<std::vector<char>>
   serializeISA(const std::string &isa) override;
 
   std::unique_ptr<SmallVectorImpl<char>> assembleIsa(const std::string &isa);
   std::unique_ptr<std::vector<char>>
   createHsaco(const SmallVectorImpl<char> &isaBinary);
 };
 } // namespace
 
+SerializeToHsacoPass::SerializeToHsacoPass(const SerializeToHsacoPass &other)
+    : PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass>(other) {}
 static std::string getDefaultChip() {
   const char kDefaultChip[] = "gfx900";
 
   // Locate rocm_agent_enumerator.
   const char kRocmAgentEnumerator[] = "rocm_agent_enumerator";
   llvm::ErrorOr<std::string> rocmAgentEnumerator = llvm::sys::findProgramByName(
       kRocmAgentEnumerator, {__ROCM_PATH__ "/bin"});
   if (!rocmAgentEnumerator) {

 // 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) {
+                                           StringRef features, int optLevel) {
   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);
 }
 
 void SerializeToHsacoPass::getDependentDialects(
     DialectRegistry &registry) const {
   registerROCDLDialectTranslation(registry);
   gpu::SerializeToBlobPass::getDependentDialects(registry);
 }
 
+LogicalResult
+SerializeToHsacoPass::optimizeLlvm(llvm::Module &llvmModule,
+                                   llvm::TargetMachine &targetMachine) {
+  int optLevel = this->optLevel.getValue();
+  if (optLevel < 0 || optLevel > 3)
+    return getOperation().emitError()
+           << "Invalid HSA optimization level" << optLevel << "\n";
+
+  targetMachine.setOptLevel(static_cast<llvm::CodeGenOpt::Level>(optLevel));
+
+  auto transformer =
+      makeOptimizingTransformer(optLevel, /*sizeLevel=*/0, &targetMachine);
+  auto error = transformer(&llvmModule);
+  if (error) {
+    InFlightDiagnostic mlirError = getOperation()->emitError();
+    llvm::handleAllErrors(
+        std::move(error), [&mlirError](const llvm::ErrorInfoBase &ei) {
+          mlirError << "Could not optimize LLVM IR: " << ei.message() << "\n";
+        });
+    return mlirError;
+  }
+  return success();
+}
+
 std::unique_ptr<SmallVectorImpl<char>>
 SerializeToHsacoPass::assembleIsa(const std::string &isa) {
   auto loc = getOperation().getLoc();
 
   SmallVector<char, 0> result;
   llvm::raw_svector_ostream os(result);
 
   llvm::Triple triple(llvm::Triple::normalize(this->triple));
PassRegistration<SerializeToHsacoPass> registerSerializeToHSACO(
         // Initialize LLVM AMDGPU backend.
         LLVMInitializeAMDGPUAsmParser();
         LLVMInitializeAMDGPUAsmPrinter();
         LLVMInitializeAMDGPUTarget();
         LLVMInitializeAMDGPUTargetInfo();
         LLVMInitializeAMDGPUTargetMC();
 
         return std::make_unique<SerializeToHsacoPass>("amdgcn-amd-amdhsa", "",
-                                                      "");
+                                                      "", 2);
       });
 }
 #else  // MLIR_GPU_TO_HSACO_PASS_ENABLE
 void mlir::registerGpuSerializeToHsacoPass() {}
 #endif // MLIR_GPU_TO_HSACO_PASS_ENABLE