diff --git a/mlir/include/mlir/Dialect/GPU/Passes.h b/mlir/include/mlir/Dialect/GPU/Passes.h index a207c6b2279e..0ce874fc4c3b 100644 --- a/mlir/include/mlir/Dialect/GPU/Passes.h +++ b/mlir/include/mlir/Dialect/GPU/Passes.h @@ -54,14 +54,23 @@ public: protected: void getDependentDialects(DialectRegistry ®istry) const override; -private: - /// Creates the LLVM target machine to generate the ISA. - std::unique_ptr 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 translateToLLVMIR(llvm::LLVMContext &llvmContext); +private: + /// Creates the LLVM target machine to generate the ISA. + std::unique_ptr createTargetMachine(); + + /// Translates the module to ISA + Optional translateToISA(llvm::Module &llvmModule, + llvm::TargetMachine &targetMachine); + /// Serializes the target ISA to binary form. virtual std::unique_ptr> serializeISA(const std::string &isa) = 0; diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt index 5fe79070737c..1243f4e540ab 100644 --- a/mlir/lib/Dialect/GPU/CMakeLists.txt +++ b/mlir/lib/Dialect/GPU/CMakeLists.txt @@ -162,6 +162,7 @@ if(MLIR_ENABLE_ROCM_RUNNER) target_link_libraries(MLIRGPUOps PRIVATE lldELF + MLIRExecutionEngine MLIRROCDLToLLVMIRTranslation ) diff --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp index 07c138e380e2..f8be171d1f5d 100644 --- a/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp @@ -31,18 +31,28 @@ gpu::SerializeToBlobPass::SerializeToBlobPass(TypeID passID) gpu::SerializeToBlobPass::SerializeToBlobPass(const SerializeToBlobPass &other) : OperationPass(other) {} -static std::string translateToISA(llvm::Module &llvmModule, - llvm::TargetMachine &targetMachine) { +Optional +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); - codegenPasses.run(llvmModule); - return targetISA; + + { // 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 stream.str(); } void gpu::SerializeToBlobPass::runOnOperation() { @@ -58,7 +68,13 @@ void gpu::SerializeToBlobPass::runOnOperation() { if (!targetMachine) return signalPassFailure(); - std::string targetISA = translateToISA(*llvmModule, *targetMachine); + Optional maybeTargetISA = + translateToISA(*llvmModule, *targetMachine); + + if (!maybeTargetISA.hasValue()) + return signalPassFailure(); + + std::string targetISA = std::move(maybeTargetISA.getValue()); // Serialize the target ISA. std::unique_ptr> blob = serializeISA(targetISA); @@ -71,6 +87,14 @@ void gpu::SerializeToBlobPass::runOnOperation() { 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 ®istry) const { registerLLVMDialectTranslation(registry); diff --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp index 29edd535e3cf..16949a70a810 100644 --- a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp @@ -15,6 +15,7 @@ #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" @@ -53,12 +54,24 @@ namespace { class SerializeToHsacoPass : public PassWrapper { 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 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 ®istry) const override; @@ -72,6 +85,8 @@ private: }; } // namespace +SerializeToHsacoPass::SerializeToHsacoPass(const SerializeToHsacoPass &other) + : PassWrapper(other) {} static std::string getDefaultChip() { const char kDefaultChip[] = "gfx900"; @@ -137,10 +152,12 @@ static void maybeSetOption(Pass::Option &option, } 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( @@ -149,6 +166,30 @@ void SerializeToHsacoPass::getDependentDialects( 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(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> SerializeToHsacoPass::assembleIsa(const std::string &isa) { auto loc = getOperation().getLoc(); @@ -286,7 +327,7 @@ void mlir::registerGpuSerializeToHsacoPass() { LLVMInitializeAMDGPUTargetMC(); return std::make_unique("amdgcn-amd-amdhsa", "", - ""); + "", 2); }); } #else // MLIR_GPU_TO_HSACO_PASS_ENABLE