Make optimize llvm common to both gpu-to-hsaco/cubin

[vinayaka-bandishti]

Summary

Before serializing, optimizations on llvm were only called on path to
hsaco, and not cubin. Define opt-level for gpu-to-cubin pass as well,
and move call to optimize llvm to a common place.

[vinayaka-bandishti]

mlir/lib/Dialect/GPU/CMakeLists.txt

127

I am a little skeptical about adding this. Please let me know if there is a better way to do this.

[bondhugula]

The update looks good but the layering needs to be fixed.

mlir/lib/Dialect/GPU/CMakeLists.txt

127

This can't be right. GPUTransforms can't depend on the ExecutionEngine. You'll need to move the OptUtils to its own library.

[bondhugula]

mlir/lib/Dialect/GPU/CMakeLists.txt

127

You can create an MLIRLLVMOptUtils with the needed functions shared between SerializeToBlob and ExecutionEngine.

[bondhugula]

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

323

Since this was removed, a header include would need to be updated/trimmed?

[vinayaka-bandishti]

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

323

The declaration was in the same file. I have removed it.

[vinayaka-bandishti]

Fix layering

[bondhugula]

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

136

... with optLevel (default level 2).

mlir/lib/Dialect/GPU/CMakeLists.txt

127

Missing dep on MLIRExecutionEngineUtils - please see below.

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

16

This is again an issue I believe. OptUtils.h/cpp is part of MLIRExecutionEngineUtils. If you include this file, you'll need to depend on MLIRExecutionEngineUtils.

Without it and in this form, the build/link will simply fail with shared libs turned on.

[vinayaka-bandishti]

Address comments, layering fix

[vinayaka-bandishti]

mlir/lib/Dialect/GPU/CMakeLists.txt

12

I think, we can remove this IPO dep. Perhaps in a different change.

127

This dep is already present.

[ftynse]

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

115

I wouldn't have the default level be 2. This is surprising for both "classical compiler" users that would expect it to be 0 (as in clang/gcc) and "ml performance-oriented compiler" that would expect it ot be 3 or the maximum level. I see that this seems to have been the case for HSACO before, can we track down the commit that introduced that and ask the author for their rationale?

140

Ditto. Especially dangerous because CUBIN defaults to 2 and HSACO doesn't have a default.

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

109

Ultra-nit: don't start diagnostic messages from a letter.

[vinayaka-bandishti]

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

115

Sure.
Commit introduced default opt level = 2 for hsaco is https://github.com/llvm/llvm-project/commit/bd22554af06e1f16dc9ff12eac8987f0ceebe8c1,
reviewed at https://reviews.llvm.org/D114113

@krzysz00 Can you please check this comment on rationale behind opt level?

[vinayaka-bandishti]

Address comment on diagnostics style

[vinayaka-bandishti]

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

109

You meant capital letters?

[bondhugula]

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

114

This method is in the MLIRExecutionEngineUtils library. The dependency on LLVMTarget itself won't help.

[bondhugula]

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

115

As for HSACO, if we want to change the default opt level, it should be done in another revision. This revision is on adding the optimizer on the serialize-to-cubin path. The commit shouldn't surprise users by changing the existing default.

[bondhugula]

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

109

Diagnostic messages shouldn't end with a new line. It is appropriately suffixed with more things. You are also missing a space after "level". The message should start in lowercase, I believe.

[vinayaka-bandishti]

Resolve comments on diagnostics

[vinayaka-bandishti]

A note about the default opt level.

Before:
HSACO path had default opt level 2. CUBIN path had no(dummy) call to optimizeLlvm - effectively opt level was 0.

Since we make the default opt level common to both paths in this patch,
if we keep it at 2, cubin users will see the change/difference.
if we make it 0, hsaco users will see the change/difference.

[bondhugula]

In D151554#4394485, @vinayaka-bandishti wrote:

A note about the default opt level.

Before:
HSACO path had default opt level 2. CUBIN path had no(dummy) call to optimizeLlvm - effectively opt level was 0.

Since we make the default opt level common to both paths in this patch,
if we keep it at 2, cubin users will see the change/difference.
if we make it 0, hsaco users will see the change/difference.

The commit title and purpose already captures the fact that the cubin path will now have optimization. So it's fine this way since HSACO path's behavior is unchanged. We can change both to O3 in a subsequent commit if that's natural for the amd/hsaco path as well.

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

 
 protected:
   Option<std::string> triple{*this, "triple",
                              ::llvm::cl::desc("Target triple")};
   Option<std::string> chip{*this, "chip",
                            ::llvm::cl::desc("Target architecture")};
   Option<std::string> features{*this, "features",
                                ::llvm::cl::desc("Target features")};
+  Option<int> optLevel{*this, "opt-level",
+                       llvm::cl::desc("Optimization level for compilation"),
+                       llvm::cl::init(2)};
   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 CUBIN binary serialization
-/// pass.
+/// pass with optLevel (default level 2).
 std::unique_ptr<Pass> createGpuSerializeToCubinPass(StringRef triple,
                                                     StringRef chip,
-                                                    StringRef features);
+                                                    StringRef features,
+                                                    int optLevel = 2);
 
 /// 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/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_GPU_TRANSFORMS_PASSES_H_

mlir/lib/Dialect/GPU/CMakeLists.txt

 if (MLIR_ENABLE_CUDA_CONVERSIONS)
   set(NVPTX_LIBS
     NVPTXCodeGen
     NVPTXDesc
     NVPTXInfo
   )
 endif()
 
 if (MLIR_ENABLE_ROCM_CONVERSIONS)
   set(AMDGPU_LIBS
     IRReader
     IPO
     linker
     MCParser
     AMDGPUAsmParser
     AMDGPUCodeGen
     AMDGPUDesc
     AMDGPUInfo
     target
   )
add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/SerializeToHsaco.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
 
   LINK_COMPONENTS
   Core
   MC
+  Target
   ${NVPTX_LIBS}
   ${AMDGPU_LIBS}
 
   DEPENDS
   MLIRGPUPassIncGen
   MLIRParallelLoopMapperEnumsGen
 
   LINK_LIBS PUBLIC
if(MLIR_ENABLE_CUDA_RUNNER)
   target_include_directories(obj.MLIRGPUTransforms
     PRIVATE
     ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
   )
 
   find_library(CUDA_DRIVER_LIBRARY cuda HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
   target_link_libraries(MLIRGPUTransforms
     PRIVATE
     MLIRNVVMToLLVMIRTranslation
     ${CUDA_DRIVER_LIBRARY}
   )
 
 endif()
 
 if(MLIR_ENABLE_ROCM_CONVERSIONS)
   if (NOT ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD))
     message(SEND_ERROR

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

 //===- SerializeToBlob.cpp - MLIR GPU lowering pass -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a base class for a pass to serialize a gpu module
 // into a binary blob that can be executed on a GPU. The binary blob is added
 // as a string attribute to the gpu module.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.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"
void gpu::SerializeToBlobPass::runOnOperation() {
   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
+  int optLevel = this->optLevel.getValue();
+  if (optLevel < 0 || optLevel > 3)
+    return getOperation().emitError()
+           << "invalid optimization level " << optLevel;
+
+  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();
+        });
+    return mlirError;
+  }
   return success();
 }
 
 void gpu::SerializeToBlobPass::getDependentDialects(
     DialectRegistry &registry) const {
   registerGPUDialectTranslation(registry);
   registerLLVMDialectTranslation(registry);
   OperationPass<gpu::GPUModuleOp>::getDependentDialects(registry);

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

 
 namespace {
 class SerializeToCubinPass
     : public PassWrapper<SerializeToCubinPass, gpu::SerializeToBlobPass> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SerializeToCubinPass)
 
   SerializeToCubinPass(StringRef triple = "nvptx64-nvidia-cuda",
-                       StringRef chip = "sm_35", StringRef features = "+ptx60");
+                       StringRef chip = "sm_35", StringRef features = "+ptx60",
+                       int optLevel = 2);
 
   StringRef getArgument() const override { return "gpu-to-cubin"; }
   StringRef getDescription() const override {
     return "Lower GPU kernel function to CUBIN binary annotations";
   }
 
 private:
   void getDependentDialects(DialectRegistry &registry) const override;
 
   // Serializes PTX to CUBIN.
   std::unique_ptr<std::vector<char>>
   serializeISA(const std::string &isa) override;
 };
 } // namespace
 
 // Sets the 'option' to 'value' unless it already has a value.
 static void maybeSetOption(Pass::Option<std::string> &option, StringRef value) {
   if (!option.hasValue())
     option = value.str();
 }
 
 SerializeToCubinPass::SerializeToCubinPass(StringRef triple, StringRef chip,
-                                           StringRef features) {
+                                           StringRef features, int optLevel) {
   maybeSetOption(this->triple, triple);
   maybeSetOption(this->chip, chip);
   maybeSetOption(this->features, features);
+  if (this->optLevel.getNumOccurrences() == 0)
+    this->optLevel.setValue(optLevel);
 }
 
 void SerializeToCubinPass::getDependentDialects(
     DialectRegistry &registry) const {
   registerNVVMDialectTranslation(registry);
   gpu::SerializeToBlobPass::getDependentDialects(registry);
 }
 
PassRegistration<SerializeToCubinPass> registerSerializeToCubin([] {
     LLVMInitializeNVPTXAsmPrinter();
 
     return std::make_unique<SerializeToCubinPass>();
   });
 }
 
 std::unique_ptr<Pass> mlir::createGpuSerializeToCubinPass(StringRef triple,
                                                           StringRef arch,
-                                                          StringRef features) {
-  return std::make_unique<SerializeToCubinPass>(triple, arch, features);
+                                                          StringRef features,
+                                                          int optLevel) {
+  return std::make_unique<SerializeToCubinPass>(triple, arch, features,
+                                                optLevel);
 }
 
 #else  // MLIR_GPU_TO_CUBIN_PASS_ENABLE
 void mlir::registerGpuSerializeToCubinPass() {}
 #endif // MLIR_GPU_TO_CUBIN_PASS_ENABLE

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

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)};
-
   Option<std::string> rocmPath{*this, "rocm-path",
                                llvm::cl::desc("Path to ROCm install")};
 
   // Overload to allow linking in device libs
   std::unique_ptr<llvm::Module>
   translateToLLVMIR(llvm::LLVMContext &llvmContext) override;
 
-  /// Adds LLVM optimization passes
-  LogicalResult optimizeLlvm(llvm::Module &llvmModule,
-                             llvm::TargetMachine &targetMachine) override;
-
 private:
   void getDependentDialects(DialectRegistry &registry) const override;
 
   // Loads LLVM bitcode libraries
   std::optional<SmallVector<std::unique_ptr<llvm::Module>, 3>>
   loadLibraries(SmallVectorImpl<char> &path,
                 SmallVectorImpl<StringRef> &libraries,
                 llvm::LLVMContext &context);
if (err) {
       // We have no guaranties about the state of `ret`, so bail
       return nullptr;
     }
   }
 
   return ret;
 }
 
-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));