diff --git a/mlir/include/mlir/Target/LLVM/ModuleToObject.h b/mlir/include/mlir/Target/LLVM/ModuleToObject.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Target/LLVM/ModuleToObject.h @@ -0,0 +1,121 @@ +//===- 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_TARGET_LLVM_MODULETOOBJECT_H +#define MLIR_TARGET_LLVM_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 operation 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::Module &module, llvm::TargetMachine &targetMachine) { + return SmallVector>(); + } + + /// Hook for performing additional actions on a loaded bitcode file. + virtual LogicalResult handleBitcodeFile(llvm::Module &module, + llvm::TargetMachine &targetMachine) { + return success(); + } + + /// 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, + llvm::TargetMachine &targetMachine, StringRef path); + + /// Loads multiple bitcode files. + LogicalResult loadBitcodeFilesFromList( + llvm::LLVMContext &context, llvm::TargetMachine &targetMachine, + 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_TARGET_LLVM_MODULETOOBJECT_H diff --git a/mlir/lib/Target/CMakeLists.txt b/mlir/lib/Target/CMakeLists.txt --- a/mlir/lib/Target/CMakeLists.txt +++ b/mlir/lib/Target/CMakeLists.txt @@ -1,3 +1,4 @@ add_subdirectory(Cpp) add_subdirectory(SPIRV) add_subdirectory(LLVMIR) +add_subdirectory(LLVM) diff --git a/mlir/lib/Target/LLVM/CMakeLists.txt b/mlir/lib/Target/LLVM/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Target/LLVM/CMakeLists.txt @@ -0,0 +1,20 @@ +add_mlir_library(MLIRTargetLLVM + ModuleToObject.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Target/LLVM + + DEPENDS + intrinsics_gen + + LINK_COMPONENTS + Core + IPO + nativecodegen + Passes + Support + Target + TargetParser + LINK_LIBS PUBLIC + MLIRExecutionEngineUtils +) \ No newline at end of file diff --git a/mlir/lib/Target/LLVM/ModuleToObject.cpp b/mlir/lib/Target/LLVM/ModuleToObject.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Target/LLVM/ModuleToObject.cpp @@ -0,0 +1,227 @@ +//===- 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/Target/LLVM/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, + llvm::TargetMachine &targetMachine, + 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; + } + if (failed(handleBitcodeFile(*library, targetMachine))) { + return nullptr; + } + return library; +} + +LogicalResult ModuleToObject::loadBitcodeFilesFromList( + llvm::LLVMContext &context, llvm::TargetMachine &targetMachine, + 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, targetMachine, 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(*llvmModule, *targetMachine); + 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); +} diff --git a/mlir/unittests/CMakeLists.txt b/mlir/unittests/CMakeLists.txt --- a/mlir/unittests/CMakeLists.txt +++ b/mlir/unittests/CMakeLists.txt @@ -16,6 +16,7 @@ add_subdirectory(Support) add_subdirectory(Rewrite) add_subdirectory(TableGen) +add_subdirectory(Target) add_subdirectory(Transforms) if(MLIR_ENABLE_EXECUTION_ENGINE) diff --git a/mlir/unittests/Target/CMakeLists.txt b/mlir/unittests/Target/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/unittests/Target/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(LLVM) diff --git a/mlir/unittests/Target/LLVM/CMakeLists.txt b/mlir/unittests/Target/LLVM/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/unittests/Target/LLVM/CMakeLists.txt @@ -0,0 +1,10 @@ +add_mlir_unittest(MLIRTargetLLVMTests + Serialize.cpp +) +get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) + +target_link_libraries(MLIRTargetLLVMTests + PRIVATE + MLIRTargetLLVM + ${dialect_libs} +) diff --git a/mlir/unittests/Target/LLVM/Serialize.cpp b/mlir/unittests/Target/LLVM/Serialize.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Target/LLVM/Serialize.cpp @@ -0,0 +1,68 @@ +//===- Serialize.cpp --------------------------------------------*- C++ -*-===// +// +// This file is licensed 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/IR/MLIRContext.h" +#include "mlir/InitAllDialects.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Target/LLVM/ModuleToObject.h" +#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" + +#include "llvm/IRReader/IRReader.h" +#include "llvm/Support/MemoryBufferRef.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/TargetParser/Host.h" + +#include "gmock/gmock.h" + +using namespace mlir; + +static struct LLVMInitializer { + LLVMInitializer() { + llvm::InitializeNativeTarget(); + llvm::InitializeNativeTargetAsmPrinter(); + } +} initializer; + +TEST(MLIRTargetLLVM, SerializeModule) { + std::string moduleStr = R"mlir( + llvm.func @foo(%arg0 : i32) { + llvm.return + } + )mlir"; + + DialectRegistry registry; + registerAllDialects(registry); + registerBuiltinDialectTranslation(registry); + registerLLVMDialectTranslation(registry); + MLIRContext context(registry); + + OwningOpRef module = + parseSourceString(moduleStr, &context); + ASSERT_TRUE(!!module); + + // Serialize the module. + std::string targetTriple = llvm::sys::getDefaultTargetTriple(); + ModuleToObject serializer(*(module->getOperation()), targetTriple, "", ""); + std::optional> serializedModule = serializer.run(); + ASSERT_TRUE(!!serializedModule); + ASSERT_TRUE(serializedModule->size() > 0); + + // Read the serialized module. + llvm::MemoryBufferRef buffer( + StringRef(serializedModule->data(), serializedModule->size()), "module"); + llvm::LLVMContext llvmContext; + llvm::Expected> llvmModule = + llvm::getLazyBitcodeModule(buffer, llvmContext); + ASSERT_TRUE(!!llvmModule); + ASSERT_TRUE(!!*llvmModule); + + // Check that it has a function named `foo`. + ASSERT_TRUE((*llvmModule)->getFunction("foo") != nullptr); +}