diff --git a/mlir/include/mlir/ExecutionEngine/ModuleToObject.h b/mlir/include/mlir/ExecutionEngine/ModuleToObject.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/ExecutionEngine/ModuleToObject.h @@ -0,0 +1,117 @@ +//===- ModuleToObject.h - Module to object base class -----------*- C++ -*-===// +// +// 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 declares the base class for transforming Operations into binary +// objects. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_EXECUTIONENGINE_MODULETOOBJECT_H +#define MLIR_EXECUTIONENGINE_MODULETOOBJECT_H + +#include "mlir/IR/Operation.h" +#include "llvm/IR/Module.h" + +namespace llvm { +class TargetMachine; +} // namespace llvm + +namespace mlir { +namespace LLVM { +class ModuleTranslation; +} +/// Utility base class for transforming Operations into binary objects, by +/// default it returns the serialized bitcode for the module. +class ModuleToObject { +public: + ModuleToObject(Operation &module, StringRef triple, StringRef chip, + StringRef features = {}, int optLevel = 3); + virtual ~ModuleToObject() = default; + + /// Returns the gpu.module being serialized. + Operation &getOperation(); + + /// Runs the serialization pipeline, returning `std::nullopt` on error. + virtual std::optional> run(); + +protected: + // Hooks to be implemented by derived classes. + + /// Hook for loading bitcode files, returns std::nullopt on failure. + virtual std::optional>> + loadBitcodeFiles(llvm::LLVMContext &context, llvm::Module &module) { + return SmallVector>(); + } + + /// Hook for performing additional actions on a loaded bitcode file. + virtual void handleBitcodeFile(llvm::Module &module, + llvm::TargetMachine &targetMachine) {} + + /// Hook for performing additional actions on the llvmModule pre linking. + virtual void handleModulePreLink(llvm::Module &module, + llvm::TargetMachine &targetMachine) {} + + /// Hook for performing additional actions on the llvmModule post linking. + virtual void handleModulePostLink(llvm::Module &module, + llvm::TargetMachine &targetMachine) {} + + /// Serializes the LLVM IR bitcode to an object file, by default it serializes + /// to LLVM bitcode. + virtual std::optional> + moduleToObject(llvm::Module &llvmModule, llvm::TargetMachine &targetMachine); + +protected: + /// Create the target machine based on the target triple and chip. + std::unique_ptr createTargetMachine(); + + /// Loads a bitcode file from path. + std::unique_ptr loadBitcodeFile(llvm::LLVMContext &context, + StringRef path); + + /// Loads multiple bitcode files. + LogicalResult loadBitcodeFilesFromList( + llvm::LLVMContext &context, ArrayRef fileList, + SmallVector> &llvmModules, + bool failureOnError = true); + + /// Translates the gpu.module to LLVM IR. + std::unique_ptr + translateToLLVMIR(llvm::LLVMContext &llvmContext); + + /// Link the llvmModule to other bitcode file. + LogicalResult linkFiles(llvm::Module &module, + SmallVector> &&libs); + + /// Optimize the module. + LogicalResult optimizeModule(llvm::Module &module, + llvm::TargetMachine &targetMachine, int optL); + + /// Utility function for translating to ISA, returns `std::nullopt` on + /// failure. + static std::optional + translateToISA(llvm::Module &llvmModule, llvm::TargetMachine &targetMachine); + +protected: + /// Module to transform to a binary object. + Operation &module; + + /// Target triple. + StringRef triple; + + /// Target chip. + StringRef chip; + + /// Target features. + StringRef features; + + /// Optimization level. + int optLevel; +}; +} // namespace mlir + +#endif // MLIR_EXECUTIONENGINE_MODULETOOBJECT_H diff --git a/mlir/lib/ExecutionEngine/CMakeLists.txt b/mlir/lib/ExecutionEngine/CMakeLists.txt --- a/mlir/lib/ExecutionEngine/CMakeLists.txt +++ b/mlir/lib/ExecutionEngine/CMakeLists.txt @@ -19,6 +19,7 @@ # libMLIR.so. add_mlir_library(MLIRExecutionEngineUtils OptUtils.cpp + ModuleToObject.cpp ADDITIONAL_HEADER_DIRS ${MLIR_MAIN_INCLUDE_DIR}/mlir/ExecutionEngine diff --git a/mlir/lib/ExecutionEngine/ModuleToObject.cpp b/mlir/lib/ExecutionEngine/ModuleToObject.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/ExecutionEngine/ModuleToObject.cpp @@ -0,0 +1,221 @@ +//===- ModuleToObject.cpp - Module to object base class ---------*- C++ -*-===// +// +// 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 the base class for transforming Operations into binary +// objects. +// +//===----------------------------------------------------------------------===// + +#include "mlir/ExecutionEngine/ModuleToObject.h" + +#include "mlir/ExecutionEngine/OptUtils.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Export.h" +#include "mlir/Target/LLVMIR/ModuleTranslation.h" + +#include "llvm/Bitcode/BitcodeWriter.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IRReader/IRReader.h" +#include "llvm/Linker/Linker.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/TargetParser/TargetParser.h" +#include "llvm/Transforms/IPO/Internalize.h" + +using namespace mlir; + +ModuleToObject::ModuleToObject(Operation &module, StringRef triple, + StringRef chip, StringRef features, int optLevel) + : module(module), triple(triple), chip(chip), features(features), + optLevel(optLevel) {} + +Operation &ModuleToObject::getOperation() { return module; } + +std::unique_ptr ModuleToObject::createTargetMachine() { + std::string error; + // Load the target. + const llvm::Target *target = + llvm::TargetRegistry::lookupTarget(triple, error); + if (!target) { + getOperation().emitError() << "Failed to lookup target: " << error; + return {}; + } + + // Create the target machine using the target. + llvm::TargetMachine *machine = + target->createTargetMachine(triple, chip, features, {}, {}); + if (!machine) { + getOperation().emitError() << "Failed to create the target machine."; + return {}; + } + return std::unique_ptr{machine}; +} + +std::unique_ptr +ModuleToObject::loadBitcodeFile(llvm::LLVMContext &context, StringRef path) { + llvm::SMDiagnostic error; + std::unique_ptr library = + llvm::getLazyIRFileModule(path, error, context); + if (!library) { + getOperation().emitError() << "Failed loading file from " << path + << ", error: " << error.getMessage(); + return nullptr; + } + return library; +} + +LogicalResult ModuleToObject::loadBitcodeFilesFromList( + llvm::LLVMContext &context, ArrayRef fileList, + SmallVector> &llvmModules, + bool failureOnError) { + for (const std::string &str : fileList) { + // Test if the path exists, if it doesn't abort. + StringRef pathRef = StringRef(str.data(), str.size()); + if (!llvm::sys::fs::is_regular_file(pathRef)) { + getOperation().emitError() + << "File path: " << pathRef << " does not exist or is not a file.\n"; + return failure(); + } + // Load the file or abort on error. + if (auto bcFile = loadBitcodeFile(context, pathRef)) + llvmModules.push_back(std::move(bcFile)); + else if (failureOnError) + return failure(); + } + return success(); +} + +std::unique_ptr +ModuleToObject::translateToLLVMIR(llvm::LLVMContext &llvmContext) { + return translateModuleToLLVMIR(&getOperation(), llvmContext); +} + +LogicalResult +ModuleToObject::linkFiles(llvm::Module &module, + SmallVector> &&libs) { + if (libs.empty()) + return success(); + llvm::Linker linker(module); + for (std::unique_ptr &libModule : libs) { + // This bitcode linking imports the library functions into the module, + // allowing LLVM optimization passes (which must run after linking) to + // optimize across the libraries and the module's code. We also only import + // symbols if they are referenced by the module or a previous library since + // there will be no other source of references to those symbols in this + // compilation and since we don't want to bloat the resulting code object. + bool err = linker.linkInModule( + std::move(libModule), llvm::Linker::Flags::LinkOnlyNeeded, + [](llvm::Module &m, const StringSet<> &gvs) { + llvm::internalizeModule(m, [&gvs](const llvm::GlobalValue &gv) { + return !gv.hasName() || (gvs.count(gv.getName()) == 0); + }); + }); + // True is linker failure + if (err) { + getOperation().emitError("Unrecoverable failure during bitcode linking."); + // We have no guaranties about the state of `ret`, so bail + return failure(); + } + } + return success(); +} + +LogicalResult ModuleToObject::optimizeModule(llvm::Module &module, + llvm::TargetMachine &targetMachine, + int optLevel) { + if (optLevel < 0 || optLevel > 3) + return getOperation().emitError() + << "Invalid optimization level: " << optLevel << "."; + + targetMachine.setOptLevel(static_cast(optLevel)); + + auto transformer = + makeOptimizingTransformer(optLevel, /*sizeLevel=*/0, &targetMachine); + auto error = transformer(&module); + 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::optional +ModuleToObject::translateToISA(llvm::Module &llvmModule, + llvm::TargetMachine &targetMachine) { + std::string targetISA; + llvm::raw_string_ostream stream(targetISA); + + { // Drop pstream after this to prevent the ISA from being stuck buffering + llvm::buffer_ostream pstream(stream); + llvm::legacy::PassManager codegenPasses; + + if (targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr, + llvm::CGFT_AssemblyFile)) + return std::nullopt; + + codegenPasses.run(llvmModule); + } + return stream.str(); +} + +std::optional> +ModuleToObject::moduleToObject(llvm::Module &llvmModule, + llvm::TargetMachine &targetMachine) { + SmallVector binaryData; + // Write the LLVM module bitcode to a buffer. + llvm::raw_svector_ostream outputStream(binaryData); + llvm::WriteBitcodeToFile(llvmModule, outputStream); + return binaryData; +} + +std::optional> ModuleToObject::run() { + // Translate the module to LLVM IR. + llvm::LLVMContext llvmContext; + std::unique_ptr llvmModule = translateToLLVMIR(llvmContext); + if (!llvmModule) { + getOperation().emitError() << "Failed creating the llvm::Module."; + return std::nullopt; + } + + // Create the target machine. + std::unique_ptr targetMachine = createTargetMachine(); + if (!targetMachine) + return std::nullopt; + + // Set the data layout and target triple of the module. + llvmModule->setDataLayout(targetMachine->createDataLayout()); + llvmModule->setTargetTriple(targetMachine->getTargetTriple().getTriple()); + + // Link bitcode files. + handleModulePreLink(*llvmModule, *targetMachine); + { + auto libs = loadBitcodeFiles(llvmContext, *llvmModule); + if (!libs) + return std::nullopt; + if (libs->size()) + if (failed(linkFiles(*llvmModule, std::move(*libs)))) + return std::nullopt; + handleModulePostLink(*llvmModule, *targetMachine); + } + + // Optimize the module. + if (failed(optimizeModule(*llvmModule, *targetMachine, optLevel))) + return std::nullopt; + + // Return the serialized object. + return moduleToObject(*llvmModule, *targetMachine); +}