b24de9f684
The enableObjectCache option was added in https://reviews.llvm.org/rG06e8101034e, defaulting to false. However, the init code added there got its logic reversed (cache(enableObjectCache ? nullptr : new SimpleObjectCache()), which was fixed in https://reviews.llvm.org/rGd1186fcb04 by setting the default to true, thereby preserving the existing behavior even if it was unintentional. Default now the object cache to false as it was originally intended. While at it, mention in enableObjectCache's documentation how the cache can be dumped. Reviewed-by: mehdi_amini Differential Revision: https://reviews.llvm.org/D121291
381 lines
14 KiB
C++
381 lines
14 KiB
C++
//===- ExecutionEngine.cpp - MLIR Execution engine and utils --------------===//
|
|
//
|
|
// 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 execution engine for MLIR modules based on LLVM Orc
|
|
// JIT engine.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
#include "mlir/ExecutionEngine/ExecutionEngine.h"
|
|
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/Support/FileUtilities.h"
|
|
#include "mlir/Target/LLVMIR/Export.h"
|
|
|
|
#include "llvm/ExecutionEngine/JITEventListener.h"
|
|
#include "llvm/ExecutionEngine/ObjectCache.h"
|
|
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
|
|
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
|
|
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
|
|
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
|
|
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
|
|
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
|
|
#include "llvm/IR/IRBuilder.h"
|
|
#include "llvm/MC/SubtargetFeature.h"
|
|
#include "llvm/MC/TargetRegistry.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include "llvm/Support/Error.h"
|
|
#include "llvm/Support/Host.h"
|
|
#include "llvm/Support/ToolOutputFile.h"
|
|
|
|
#define DEBUG_TYPE "execution-engine"
|
|
|
|
using namespace mlir;
|
|
using llvm::dbgs;
|
|
using llvm::Error;
|
|
using llvm::errs;
|
|
using llvm::Expected;
|
|
using llvm::LLVMContext;
|
|
using llvm::MemoryBuffer;
|
|
using llvm::MemoryBufferRef;
|
|
using llvm::Module;
|
|
using llvm::SectionMemoryManager;
|
|
using llvm::StringError;
|
|
using llvm::Triple;
|
|
using llvm::orc::DynamicLibrarySearchGenerator;
|
|
using llvm::orc::ExecutionSession;
|
|
using llvm::orc::IRCompileLayer;
|
|
using llvm::orc::JITTargetMachineBuilder;
|
|
using llvm::orc::MangleAndInterner;
|
|
using llvm::orc::RTDyldObjectLinkingLayer;
|
|
using llvm::orc::SymbolMap;
|
|
using llvm::orc::ThreadSafeModule;
|
|
using llvm::orc::TMOwningSimpleCompiler;
|
|
|
|
/// Wrap a string into an llvm::StringError.
|
|
static Error makeStringError(const Twine &message) {
|
|
return llvm::make_error<StringError>(message.str(),
|
|
llvm::inconvertibleErrorCode());
|
|
}
|
|
|
|
void SimpleObjectCache::notifyObjectCompiled(const Module *m,
|
|
MemoryBufferRef objBuffer) {
|
|
cachedObjects[m->getModuleIdentifier()] = MemoryBuffer::getMemBufferCopy(
|
|
objBuffer.getBuffer(), objBuffer.getBufferIdentifier());
|
|
}
|
|
|
|
std::unique_ptr<MemoryBuffer> SimpleObjectCache::getObject(const Module *m) {
|
|
auto i = cachedObjects.find(m->getModuleIdentifier());
|
|
if (i == cachedObjects.end()) {
|
|
LLVM_DEBUG(dbgs() << "No object for " << m->getModuleIdentifier()
|
|
<< " in cache. Compiling.\n");
|
|
return nullptr;
|
|
}
|
|
LLVM_DEBUG(dbgs() << "Object for " << m->getModuleIdentifier()
|
|
<< " loaded from cache.\n");
|
|
return MemoryBuffer::getMemBuffer(i->second->getMemBufferRef());
|
|
}
|
|
|
|
void SimpleObjectCache::dumpToObjectFile(StringRef outputFilename) {
|
|
// Set up the output file.
|
|
std::string errorMessage;
|
|
auto file = openOutputFile(outputFilename, &errorMessage);
|
|
if (!file) {
|
|
llvm::errs() << errorMessage << "\n";
|
|
return;
|
|
}
|
|
|
|
// Dump the object generated for a single module to the output file.
|
|
assert(cachedObjects.size() == 1 && "Expected only one object entry.");
|
|
auto &cachedObject = cachedObjects.begin()->second;
|
|
file->os() << cachedObject->getBuffer();
|
|
file->keep();
|
|
}
|
|
|
|
void ExecutionEngine::dumpToObjectFile(StringRef filename) {
|
|
if (cache == nullptr) {
|
|
llvm::errs() << "cannot dump ExecutionEngine object code to file: "
|
|
"object cache is disabled\n";
|
|
return;
|
|
}
|
|
cache->dumpToObjectFile(filename);
|
|
}
|
|
|
|
void ExecutionEngine::registerSymbols(
|
|
llvm::function_ref<SymbolMap(MangleAndInterner)> symbolMap) {
|
|
auto &mainJitDylib = jit->getMainJITDylib();
|
|
cantFail(mainJitDylib.define(
|
|
absoluteSymbols(symbolMap(llvm::orc::MangleAndInterner(
|
|
mainJitDylib.getExecutionSession(), jit->getDataLayout())))));
|
|
}
|
|
|
|
// Setup LLVM target triple from the current machine.
|
|
bool ExecutionEngine::setupTargetTriple(Module *llvmModule) {
|
|
// Setup the machine properties from the current architecture.
|
|
auto targetTriple = llvm::sys::getDefaultTargetTriple();
|
|
std::string errorMessage;
|
|
const auto *target =
|
|
llvm::TargetRegistry::lookupTarget(targetTriple, errorMessage);
|
|
if (!target) {
|
|
errs() << "NO target: " << errorMessage << "\n";
|
|
return true;
|
|
}
|
|
|
|
std::string cpu(llvm::sys::getHostCPUName());
|
|
llvm::SubtargetFeatures features;
|
|
llvm::StringMap<bool> hostFeatures;
|
|
|
|
if (llvm::sys::getHostCPUFeatures(hostFeatures))
|
|
for (auto &f : hostFeatures)
|
|
features.AddFeature(f.first(), f.second);
|
|
|
|
std::unique_ptr<llvm::TargetMachine> machine(target->createTargetMachine(
|
|
targetTriple, cpu, features.getString(), {}, {}));
|
|
if (!machine) {
|
|
errs() << "Unable to create target machine\n";
|
|
return true;
|
|
}
|
|
llvmModule->setDataLayout(machine->createDataLayout());
|
|
llvmModule->setTargetTriple(targetTriple);
|
|
return false;
|
|
}
|
|
|
|
static std::string makePackedFunctionName(StringRef name) {
|
|
return "_mlir_" + name.str();
|
|
}
|
|
|
|
// For each function in the LLVM module, define an interface function that wraps
|
|
// all the arguments of the original function and all its results into an i8**
|
|
// pointer to provide a unified invocation interface.
|
|
static void packFunctionArguments(Module *module) {
|
|
auto &ctx = module->getContext();
|
|
llvm::IRBuilder<> builder(ctx);
|
|
DenseSet<llvm::Function *> interfaceFunctions;
|
|
for (auto &func : module->getFunctionList()) {
|
|
if (func.isDeclaration()) {
|
|
continue;
|
|
}
|
|
if (interfaceFunctions.count(&func)) {
|
|
continue;
|
|
}
|
|
|
|
// Given a function `foo(<...>)`, define the interface function
|
|
// `mlir_foo(i8**)`.
|
|
auto *newType = llvm::FunctionType::get(
|
|
builder.getVoidTy(), builder.getInt8PtrTy()->getPointerTo(),
|
|
/*isVarArg=*/false);
|
|
auto newName = makePackedFunctionName(func.getName());
|
|
auto funcCst = module->getOrInsertFunction(newName, newType);
|
|
llvm::Function *interfaceFunc = cast<llvm::Function>(funcCst.getCallee());
|
|
interfaceFunctions.insert(interfaceFunc);
|
|
|
|
// Extract the arguments from the type-erased argument list and cast them to
|
|
// the proper types.
|
|
auto *bb = llvm::BasicBlock::Create(ctx);
|
|
bb->insertInto(interfaceFunc);
|
|
builder.SetInsertPoint(bb);
|
|
llvm::Value *argList = interfaceFunc->arg_begin();
|
|
SmallVector<llvm::Value *, 8> args;
|
|
args.reserve(llvm::size(func.args()));
|
|
for (auto &indexedArg : llvm::enumerate(func.args())) {
|
|
llvm::Value *argIndex = llvm::Constant::getIntegerValue(
|
|
builder.getInt64Ty(), APInt(64, indexedArg.index()));
|
|
llvm::Value *argPtrPtr =
|
|
builder.CreateGEP(builder.getInt8PtrTy(), argList, argIndex);
|
|
llvm::Value *argPtr =
|
|
builder.CreateLoad(builder.getInt8PtrTy(), argPtrPtr);
|
|
llvm::Type *argTy = indexedArg.value().getType();
|
|
argPtr = builder.CreateBitCast(argPtr, argTy->getPointerTo());
|
|
llvm::Value *arg = builder.CreateLoad(argTy, argPtr);
|
|
args.push_back(arg);
|
|
}
|
|
|
|
// Call the implementation function with the extracted arguments.
|
|
llvm::Value *result = builder.CreateCall(&func, args);
|
|
|
|
// Assuming the result is one value, potentially of type `void`.
|
|
if (!result->getType()->isVoidTy()) {
|
|
llvm::Value *retIndex = llvm::Constant::getIntegerValue(
|
|
builder.getInt64Ty(), APInt(64, llvm::size(func.args())));
|
|
llvm::Value *retPtrPtr =
|
|
builder.CreateGEP(builder.getInt8PtrTy(), argList, retIndex);
|
|
llvm::Value *retPtr =
|
|
builder.CreateLoad(builder.getInt8PtrTy(), retPtrPtr);
|
|
retPtr = builder.CreateBitCast(retPtr, result->getType()->getPointerTo());
|
|
builder.CreateStore(result, retPtr);
|
|
}
|
|
|
|
// The interface function returns void.
|
|
builder.CreateRetVoid();
|
|
}
|
|
}
|
|
|
|
ExecutionEngine::ExecutionEngine(bool enableObjectCache,
|
|
bool enableGDBNotificationListener,
|
|
bool enablePerfNotificationListener)
|
|
: cache(enableObjectCache ? new SimpleObjectCache() : nullptr),
|
|
gdbListener(enableGDBNotificationListener
|
|
? llvm::JITEventListener::createGDBRegistrationListener()
|
|
: nullptr),
|
|
perfListener(nullptr) {
|
|
if (enablePerfNotificationListener) {
|
|
if (auto *listener = llvm::JITEventListener::createPerfJITEventListener())
|
|
perfListener = listener;
|
|
else if (auto *listener =
|
|
llvm::JITEventListener::createIntelJITEventListener())
|
|
perfListener = listener;
|
|
}
|
|
}
|
|
|
|
Expected<std::unique_ptr<ExecutionEngine>>
|
|
ExecutionEngine::create(ModuleOp m, const ExecutionEngineOptions &options) {
|
|
auto engine = std::make_unique<ExecutionEngine>(
|
|
options.enableObjectCache, options.enableGDBNotificationListener,
|
|
options.enablePerfNotificationListener);
|
|
|
|
std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
|
|
auto llvmModule = options.llvmModuleBuilder
|
|
? options.llvmModuleBuilder(m, *ctx)
|
|
: translateModuleToLLVMIR(m, *ctx);
|
|
if (!llvmModule)
|
|
return makeStringError("could not convert to LLVM IR");
|
|
// FIXME: the triple should be passed to the translation or dialect conversion
|
|
// instead of this. Currently, the LLVM module created above has no triple
|
|
// associated with it.
|
|
setupTargetTriple(llvmModule.get());
|
|
packFunctionArguments(llvmModule.get());
|
|
|
|
auto dataLayout = llvmModule->getDataLayout();
|
|
|
|
// Callback to create the object layer with symbol resolution to current
|
|
// process and dynamically linked libraries.
|
|
auto objectLinkingLayerCreator = [&](ExecutionSession &session,
|
|
const Triple &tt) {
|
|
auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
|
|
session, [sectionMemoryMapper = options.sectionMemoryMapper]() {
|
|
return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
|
|
});
|
|
|
|
// Register JIT event listeners if they are enabled.
|
|
if (engine->gdbListener)
|
|
objectLayer->registerJITEventListener(*engine->gdbListener);
|
|
if (engine->perfListener)
|
|
objectLayer->registerJITEventListener(*engine->perfListener);
|
|
|
|
// COFF format binaries (Windows) need special handling to deal with
|
|
// exported symbol visibility.
|
|
// cf llvm/lib/ExecutionEngine/Orc/LLJIT.cpp LLJIT::createObjectLinkingLayer
|
|
llvm::Triple targetTriple(llvm::Twine(llvmModule->getTargetTriple()));
|
|
if (targetTriple.isOSBinFormatCOFF()) {
|
|
objectLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
|
|
objectLayer->setAutoClaimResponsibilityForObjectSymbols(true);
|
|
}
|
|
|
|
// Resolve symbols from shared libraries.
|
|
for (auto libPath : options.sharedLibPaths) {
|
|
auto mb = llvm::MemoryBuffer::getFile(libPath);
|
|
if (!mb) {
|
|
errs() << "Failed to create MemoryBuffer for: " << libPath
|
|
<< "\nError: " << mb.getError().message() << "\n";
|
|
continue;
|
|
}
|
|
auto &jd = session.createBareJITDylib(std::string(libPath));
|
|
auto loaded = DynamicLibrarySearchGenerator::Load(
|
|
libPath.data(), dataLayout.getGlobalPrefix());
|
|
if (!loaded) {
|
|
errs() << "Could not load " << libPath << ":\n " << loaded.takeError()
|
|
<< "\n";
|
|
continue;
|
|
}
|
|
jd.addGenerator(std::move(*loaded));
|
|
cantFail(objectLayer->add(jd, std::move(mb.get())));
|
|
}
|
|
|
|
return objectLayer;
|
|
};
|
|
|
|
// Callback to inspect the cache and recompile on demand. This follows Lang's
|
|
// LLJITWithObjectCache example.
|
|
auto compileFunctionCreator = [&](JITTargetMachineBuilder jtmb)
|
|
-> Expected<std::unique_ptr<IRCompileLayer::IRCompiler>> {
|
|
if (options.jitCodeGenOptLevel)
|
|
jtmb.setCodeGenOptLevel(options.jitCodeGenOptLevel.getValue());
|
|
auto tm = jtmb.createTargetMachine();
|
|
if (!tm)
|
|
return tm.takeError();
|
|
return std::make_unique<TMOwningSimpleCompiler>(std::move(*tm),
|
|
engine->cache.get());
|
|
};
|
|
|
|
// Create the LLJIT by calling the LLJITBuilder with 2 callbacks.
|
|
auto jit =
|
|
cantFail(llvm::orc::LLJITBuilder()
|
|
.setCompileFunctionCreator(compileFunctionCreator)
|
|
.setObjectLinkingLayerCreator(objectLinkingLayerCreator)
|
|
.create());
|
|
|
|
// Add a ThreadSafemodule to the engine and return.
|
|
ThreadSafeModule tsm(std::move(llvmModule), std::move(ctx));
|
|
if (options.transformer)
|
|
cantFail(tsm.withModuleDo(
|
|
[&](llvm::Module &module) { return options.transformer(&module); }));
|
|
cantFail(jit->addIRModule(std::move(tsm)));
|
|
engine->jit = std::move(jit);
|
|
|
|
// Resolve symbols that are statically linked in the current process.
|
|
llvm::orc::JITDylib &mainJD = engine->jit->getMainJITDylib();
|
|
mainJD.addGenerator(
|
|
cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
|
|
dataLayout.getGlobalPrefix())));
|
|
|
|
return std::move(engine);
|
|
}
|
|
|
|
Expected<void (*)(void **)>
|
|
ExecutionEngine::lookupPacked(StringRef name) const {
|
|
auto result = lookup(makePackedFunctionName(name));
|
|
if (!result)
|
|
return result.takeError();
|
|
return reinterpret_cast<void (*)(void **)>(result.get());
|
|
}
|
|
|
|
Expected<void *> ExecutionEngine::lookup(StringRef name) const {
|
|
auto expectedSymbol = jit->lookup(name);
|
|
|
|
// JIT lookup may return an Error referring to strings stored internally by
|
|
// the JIT. If the Error outlives the ExecutionEngine, it would want have a
|
|
// dangling reference, which is currently caught by an assertion inside JIT
|
|
// thanks to hand-rolled reference counting. Rewrap the error message into a
|
|
// string before returning. Alternatively, ORC JIT should consider copying
|
|
// the string into the error message.
|
|
if (!expectedSymbol) {
|
|
std::string errorMessage;
|
|
llvm::raw_string_ostream os(errorMessage);
|
|
llvm::handleAllErrors(expectedSymbol.takeError(),
|
|
[&os](llvm::ErrorInfoBase &ei) { ei.log(os); });
|
|
return makeStringError(os.str());
|
|
}
|
|
|
|
auto rawFPtr = expectedSymbol->getAddress();
|
|
auto *fptr = reinterpret_cast<void *>(rawFPtr);
|
|
if (!fptr)
|
|
return makeStringError("looked up function is null");
|
|
return fptr;
|
|
}
|
|
|
|
Error ExecutionEngine::invokePacked(StringRef name,
|
|
MutableArrayRef<void *> args) {
|
|
auto expectedFPtr = lookupPacked(name);
|
|
if (!expectedFPtr)
|
|
return expectedFPtr.takeError();
|
|
auto fptr = *expectedFPtr;
|
|
|
|
(*fptr)(args.data());
|
|
|
|
return Error::success();
|
|
}
|