Move and strip the jit engine to the v1

This commit is contained in:
Sameer Rahmani 2022-06-14 15:56:57 +01:00
parent cf71fcde38
commit c2fb6bf8cb
7 changed files with 535 additions and 5 deletions

View File

@ -0,0 +1,18 @@
#+TITLE: Serene JIT
#+AUTHOR: Sameer Rahmani
#+STARTUP: logdrawer logdone logreschedule indent content align constSI entitiespretty nolatexpreview
#+OPTIONS: tex:t
* SereneJIT
** Development Notes
*** Steps to take
You usually want a custom =MaterializationUnit= for your program representation, and a custom =Layer=. The Layer will have two
operations: =add= and =emit=. The =add= operation takes an instance of your program representation, builds one of your custom
=MaterializationUnits= to hold it, then adds it to a =JITDylib=. The =emit= operation takes a =MaterializationResponsibility=
object and an instance of your program representation and materializes it, usually by compiling it and handing the resulting
object off to an =ObjectLinkingLayer=.
Your custom =MaterializationUnit= will have two operations: =materialize= and =discard=. The =materialize= function will be
called for you when any symbol provided by the unit is looked up, and it should just call the =emit= function on your layer,
passing in the given =MaterializationResponsibility= and the wrapped program representation. The =discard= function will be
called if some weak symbol provided by your unit is not needed (because the JIT found an overriding definition).
You can use this to drop your definition early, or just ignore it and let the linker drop the definition later.

View File

@ -0,0 +1,164 @@
/* -*- C++ -*-
* Serene Programming Language
*
* Copyright (c) 2019-2022 Sameer Rahmani <lxsameer@gnu.org>
*
* This program is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, version 2.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
/**
* Commentary:
This is the first working attempt on building a JIT engine for Serene
and named after Edmond Halley.
- It supports both ASTs and Namespaces
- Every Namespace might have one or more JITDylibs. Depends on the method
of the compilation.
- It operates in lazy (for REPL) and non-lazy mode and wraps LLJIT
and LLLazyJIT
- It uses an object cache layer to cache module (not NSs) objects.
*/
#ifndef SERENE_JIT_HALLEY_H
#define SERENE_JIT_HALLEY_H
#include "serene/export.h" // for SERENE...
#include <llvm/ADT/SmallVector.h> // for SmallV...
#include <llvm/ADT/StringMap.h> // for StringMap
#include <llvm/ADT/StringRef.h> // for StringRef
#include <llvm/ExecutionEngine/ObjectCache.h> // for Object...
#include <llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h> // for JITTar...
#include <llvm/ExecutionEngine/Orc/LLJIT.h> // for LLJIT
#include <llvm/Support/Debug.h> // for dbgs
#include <llvm/Support/Error.h> // for Expected
#include <llvm/Support/MemoryBuffer.h> // for Memory...
#include <llvm/Support/MemoryBufferRef.h> // for Memory...
#include <llvm/Support/raw_ostream.h> // for raw_os...
#include <memory> // for unique...
#define HALLEY_LOG(...) \
DEBUG_WITH_TYPE("halley", llvm::dbgs() \
<< "[HALLEY]: " << __VA_ARGS__ << "\n");
namespace llvm {
class DataLayout;
class JITEventListener;
class Module;
} // namespace llvm
namespace serene {
class SereneContext;
namespace jit {
class Halley;
using MaybeEngine = llvm::Expected<std::unique_ptr<Halley>>;
using MaybeEnginePtr = llvm::Expected<void (*)(void **)>;
/// A simple object cache following Lang's LLJITWithObjectCache example and
/// MLIR's SimpelObjectCache.
class ObjectCache : public llvm::ObjectCache {
public:
/// Cache the given `objBuffer` for the given module `m`. The buffer contains
/// the combiled objects of the module
void notifyObjectCompiled(const llvm::Module *m,
llvm::MemoryBufferRef objBuffer) override;
// Lookup the cache for the given module `m` or returen a nullptr.
std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *m) override;
/// Dump cached object to output file `filename`.
void dumpToObjectFile(llvm::StringRef filename);
private:
llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> cachedObjects;
};
class SERENE_EXPORT Halley {
// TODO: Replace this with a variant of LLJIT and LLLazyJIT
std::unique_ptr<llvm::orc::LLJIT> engine;
std::unique_ptr<ObjectCache> cache;
/// GDB notification listener.
llvm::JITEventListener *gdbListener;
/// Perf notification listener.
llvm::JITEventListener *perfListener;
llvm::orc::JITTargetMachineBuilder jtmb;
llvm::DataLayout &dl;
SereneContext &ctx;
bool isLazy = false;
public:
Halley(serene::SereneContext &ctx, llvm::orc::JITTargetMachineBuilder &&jtmb,
llvm::DataLayout &&dl);
static MaybeEngine make(serene::SereneContext &ctx,
llvm::orc::JITTargetMachineBuilder &&jtmb);
void setEngine(std::unique_ptr<llvm::orc::LLJIT> e, bool isLazy);
/// Looks up a packed-argument function with the given sym name and returns a
/// pointer to it. Propagates errors in case of failure.
// MaybeJITPtr lookup(exprs::Symbol &sym) const;
/// Invokes the function with the given name passing it the list of opaque
/// pointers to the actual arguments.
// llvm::Error
// invokePacked(llvm::StringRef name,
// llvm::MutableArrayRef<void *> args = llvm::None) const;
/// Trait that defines how a given type is passed to the JIT code. This
/// defaults to passing the address but can be specialized.
template <typename T>
struct Argument {
static void pack(llvm::SmallVectorImpl<void *> &args, T &val) {
args.push_back(&val);
}
};
/// Tag to wrap an output parameter when invoking a jitted function.
template <typename T>
struct FnResult {
explicit FnResult(T &result) : value(result) {}
T &value;
};
/// Helper function to wrap an output operand when using
/// ExecutionEngine::invoke.
template <typename T>
static FnResult<T> result(T &t) {
return FnResult<T>(t);
}
// Specialization for output parameter: their address is forwarded directly to
// the native code.
template <typename T>
struct Argument<FnResult<T>> {
static void pack(llvm::SmallVectorImpl<void *> &args, FnResult<T> &result) {
args.push_back(&result.value);
}
};
void dumpToObjectFile(llvm::StringRef filename);
};
MaybeEngine makeHalleyJIT(SereneContext &ctx);
} // namespace jit
} // namespace serene
#endif

View File

@ -19,10 +19,12 @@
#ifndef SERENE_SERENE_H
#define SERENE_SERENE_H
#include "serene/export.h"
#include "serene/export.h" // for SERENE_EXPORT
#include "serene/jit/halley.h" // for MaybeEngine
namespace serene {
int SERENE_EXPORT makeEngine();
serene::jit::MaybeEngine SERENE_EXPORT makeEngine();
} // namespace serene
#endif

View File

@ -24,7 +24,9 @@ endif()
add_library(serene
serene.cpp
context.cpp)
context.cpp
jit/halley.cpp)
# Create an ALIAS target. This way if we mess up the name
# there will be an cmake error inseat of a linker error which is harder
@ -86,4 +88,5 @@ target_compile_definitions(
target_link_libraries(serene PRIVATE
LLVMOrcJIT
${llvm_libs})

View File

@ -21,6 +21,21 @@
#include <cstdlib> // for exit
namespace serene {
int SereneContext::getOptimizatioLevel() {
if (targetPhase <= CompilationPhase::NoOptimization) {
return 0;
}
if (targetPhase == CompilationPhase::O1) {
return 1;
}
if (targetPhase == CompilationPhase::O2) {
return 2;
}
return 3;
}
void terminate(SereneContext &ctx, int exitCode) {
(void)ctx;
// TODO: Since we are running in a single thread for now using exit is fine

View File

@ -0,0 +1,323 @@
/* -*- C++ -*-
* Serene Programming Language
*
* Copyright (c) 2019-2022 Sameer Rahmani <lxsameer@gnu.org>
*
* This program is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, version 2.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
#include "serene/jit/halley.h"
#include "serene/context.h" // for Seren...
#include <system_error> // for error...
#include <llvm/ADT/StringMapEntry.h> // for Strin...
#include <llvm/ADT/Triple.h> // for Triple
#include <llvm/ADT/iterator.h> // for itera...
#include <llvm/ExecutionEngine/JITEventListener.h> // for JITEv...
#include <llvm/ExecutionEngine/Orc/CompileUtils.h> // for TMOwn...
#include <llvm/ExecutionEngine/Orc/Core.h> // for Execu...
#include <llvm/ExecutionEngine/Orc/DebugUtils.h> // for opera...
#include <llvm/ExecutionEngine/Orc/ExecutionUtils.h> // for Dynam...
#include <llvm/ExecutionEngine/Orc/IRCompileLayer.h> // for IRCom...
#include <llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h> // for JITTa...
#include <llvm/ExecutionEngine/Orc/LLJIT.h> // for LLJIT...
#include <llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h> // for RTDyl...
#include <llvm/ExecutionEngine/Orc/ThreadSafeModule.h> // for Threa...
#include <llvm/ExecutionEngine/SectionMemoryManager.h> // for Secti...
#include <llvm/IR/DataLayout.h> // for DataL...
#include <llvm/IR/LLVMContext.h> // for LLVMC...
#include <llvm/IR/Module.h> // for Module
#include <llvm/Support/CodeGen.h> // for Level
#include <llvm/Support/FileSystem.h> // for OF_None
#include <llvm/Support/ToolOutputFile.h> // for ToolO...
#include <llvm/Support/raw_ostream.h> // for raw_o...
#include <algorithm> // for max
#include <assert.h> // for assert
#include <memory> // for uniqu...
#include <string> // for opera...
#include <utility> // for move
#define COMMON_ARGS_COUNT 8
namespace serene {
namespace jit {
void ObjectCache::notifyObjectCompiled(const llvm::Module *m,
llvm::MemoryBufferRef objBuffer) {
cachedObjects[m->getModuleIdentifier()] =
llvm::MemoryBuffer::getMemBufferCopy(objBuffer.getBuffer(),
objBuffer.getBufferIdentifier());
}
std::unique_ptr<llvm::MemoryBuffer>
ObjectCache::getObject(const llvm::Module *m) {
auto i = cachedObjects.find(m->getModuleIdentifier());
if (i == cachedObjects.end()) {
HALLEY_LOG("No object for " + m->getModuleIdentifier() +
" in cache. Compiling.");
return nullptr;
}
HALLEY_LOG("Object for " + m->getModuleIdentifier() + " loaded from cache.");
return llvm::MemoryBuffer::getMemBuffer(i->second->getMemBufferRef());
}
void ObjectCache::dumpToObjectFile(llvm::StringRef outputFilename) {
// Set up the output file.
std::error_code error;
auto file = std::make_unique<llvm::ToolOutputFile>(outputFilename, error,
llvm::sys::fs::OF_None);
if (error) {
llvm::errs() << "cannot open output file '" + outputFilename.str() +
"': " + error.message()
<< "\n";
return;
}
// Dump the object generated for a single module to the output file.
// TODO: Replace this with a runtime check
assert(cachedObjects.size() == 1 && "Expected only one object entry.");
auto &cachedObject = cachedObjects.begin()->second;
file->os() << cachedObject->getBuffer();
file->keep();
}
Halley::Halley(serene::SereneContext &ctx,
llvm::orc::JITTargetMachineBuilder &&jtmb, llvm::DataLayout &&dl)
: cache(ctx.opts.JITenableObjectCache ? new ObjectCache() : nullptr),
gdbListener(ctx.opts.JITenableGDBNotificationListener
? llvm::JITEventListener::createGDBRegistrationListener()
: nullptr),
perfListener(ctx.opts.JITenablePerfNotificationListener
? llvm::JITEventListener::createPerfJITEventListener()
: nullptr),
jtmb(jtmb), dl(dl), ctx(ctx){};
// MaybeJITPtr Halley::lookup(exprs::Symbol &sym) const {
// HALLEY_LOG("Looking up: " << sym.toString());
// auto *ns = ctx.getNS(sym.nsName);
// if (ns == nullptr) {
// return errors::makeError(ctx, errors::CantResolveSymbol, sym.location,
// "Can't find the namespace in the context: " +
// sym.nsName);
// }
// auto *dylib = ctx.getLatestJITDylib(*ns);
// //
// if (dylib == nullptr) {
// return errors::makeError(ctx, errors::CantResolveSymbol, sym.location,
// "Don't know about namespace: " + sym.nsName);
// }
// auto expectedSymbol =
// engine->lookup(*dylib, makePackedFunctionName(sym.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 errors::makeError(ctx, errors::CantResolveSymbol, sym.location,
// os.str());
// }
// auto rawFPtr = expectedSymbol->getValue();
// // NOLINTNEXTLINE(performance-no-int-to-ptr)
// auto fptr = reinterpret_cast<void (*)(void **)>(rawFPtr);
// if (fptr == nullptr) {
// return errors::makeError(ctx, errors::CantResolveSymbol, sym.location,
// "Lookup function is null!");
// }
// return fptr;
// };
void Halley::setEngine(std::unique_ptr<llvm::orc::LLJIT> e, bool isLazy) {
// Later on we might use different classes of JIT which might need some
// work for lazyness
(void)ctx;
engine = std::move(e);
this->isLazy = isLazy;
};
void Halley::dumpToObjectFile(llvm::StringRef filename) {
cache->dumpToObjectFile(filename);
};
MaybeEngine Halley::make(SereneContext &serene_ctx,
llvm::orc::JITTargetMachineBuilder &&jtmb) {
auto dl = jtmb.getDefaultDataLayoutForTarget();
if (!dl) {
return dl.takeError();
}
auto jitEngine =
std::make_unique<Halley>(serene_ctx, std::move(jtmb), std::move(*dl));
// Why not the llvmcontext from the SereneContext??
// Sice we're going to pass the ownership of this context to a thread
// safe module later on and we will only create the jit function wrappers
// with it, then it is fine to use a new context.
//
// What might go wrong?
// in a repl env when we have to create new modules on top of each other
// having two different contex might be a problem, but i think since we
// use the first context to generate the IR and the second one to just
// run it.
std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
// Callback to create the object layer with symbol resolution to current
// process and dynamically linked libraries.
auto objectLinkingLayerCreator = [&](llvm::orc::ExecutionSession &session,
const llvm::Triple &tt) {
(void)tt;
auto objectLayer =
std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(session, []() {
return std::make_unique<llvm::SectionMemoryManager>();
});
// Register JIT event listeners if they are enabled.
if (jitEngine->gdbListener != nullptr) {
objectLayer->registerJITEventListener(*jitEngine->gdbListener);
}
if (jitEngine->perfListener != nullptr) {
objectLayer->registerJITEventListener(*jitEngine->perfListener);
}
// COFF format binaries (Windows) need special handling to deal with
// exported symbol visibility.
// cf llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
// LLJIT::createObjectLinkingLayer
if (serene_ctx.triple.isOSBinFormatCOFF()) {
objectLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
objectLayer->setAutoClaimResponsibilityForObjectSymbols(true);
}
// Resolve symbols from shared libraries.
// for (auto libPath : sharedLibPaths) {
// auto mb = llvm::MemoryBuffer::getFile(libPath);
// if (!mb) {
// llvm::errs() << "Failed to create MemoryBuffer for: " << libPath
// << "\nError: " << mb.getError().message() << "\n";
// continue;
// }
// auto &JD = session.createBareJITDylib(std::string(libPath));
// auto loaded = llvm::orc::DynamicLibrarySearchGenerator::Load(
// libPath.data(), dataLayout.getGlobalPrefix());
// if (!loaded) {
// llvm::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 = [&](llvm::orc::JITTargetMachineBuilder JTMB)
-> llvm::Expected<
std::unique_ptr<llvm::orc::IRCompileLayer::IRCompiler>> {
llvm::CodeGenOpt::Level jitCodeGenOptLevel =
static_cast<llvm::CodeGenOpt::Level>(serene_ctx.getOptimizatioLevel());
JTMB.setCodeGenOptLevel(jitCodeGenOptLevel);
auto targetMachine = JTMB.createTargetMachine();
if (!targetMachine) {
return targetMachine.takeError();
}
return std::make_unique<llvm::orc::TMOwningSimpleCompiler>(
std::move(*targetMachine), jitEngine->cache.get());
};
if (serene_ctx.opts.JITLazy) {
// Setup a LLLazyJIT instance to the times that latency is important
// for example in a REPL. This way
auto jit =
cantFail(llvm::orc::LLLazyJITBuilder()
.setCompileFunctionCreator(compileFunctionCreator)
.setObjectLinkingLayerCreator(objectLinkingLayerCreator)
.create());
jitEngine->setEngine(std::move(jit), true);
} else {
// Setup a LLJIT instance for the times that performance is important
// and we want to compile everything as soon as possible. For instance
// when we run the JIT in the compiler
auto jit =
cantFail(llvm::orc::LLJITBuilder()
.setCompileFunctionCreator(compileFunctionCreator)
.setObjectLinkingLayerCreator(objectLinkingLayerCreator)
.create());
jitEngine->setEngine(std::move(jit), false);
}
jitEngine->engine->getIRCompileLayer().setNotifyCompiled(
[&](llvm::orc::MaterializationResponsibility &r,
llvm::orc::ThreadSafeModule tsm) {
auto syms = r.getRequestedSymbols();
tsm.withModuleDo([&](llvm::Module &m) {
HALLEY_LOG("Compiled "
<< syms << " for the module: " << m.getModuleIdentifier());
});
});
// Resolve symbols that are statically linked in the current process.
llvm::orc::JITDylib &mainJD = jitEngine->engine->getMainJITDylib();
mainJD.addGenerator(
cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
jitEngine->dl.getGlobalPrefix())));
return MaybeEngine(std::move(jitEngine));
};
MaybeEngine makeHalleyJIT(SereneContext &ctx) {
llvm::orc::JITTargetMachineBuilder jtmb(ctx.triple);
auto maybeJIT = Halley::make(ctx, std::move(jtmb));
if (!maybeJIT) {
return maybeJIT.takeError();
}
return maybeJIT;
};
} // namespace jit
} // namespace serene

View File

@ -18,8 +18,13 @@
#include "serene/serene.h"
namespace serene {
#include "serene/jit/halley.h" // for makeHalleyJIT, MaybeEngine
int makeEngine() { return 0; };
namespace serene {
class SereneContext;
serene::jit::MaybeEngine makeEngine(SereneContext &ctx) {
return serene::jit::makeHalleyJIT(ctx);
};
} // namespace serene