mirrors
/
codon
mirror of https://github.com/exaloop/codon.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

JIT fixes for Apple Silicon (#575)

pull/579/head
A. R. Shajii 2024-08-04 23:24:57 -04:00 committed by GitHub
parent 11d281d1b3
commit fa2904c15b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 126 additions and 146 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
codon/cir/llvm/llvisitor.cpp
View File

@ -2941,8 +2941,9 @@ void LLVMVisitor::visit(const TryCatchFlow *x) {
0)); 0));
// check for foreign exceptions // check for foreign exceptions
B->CreateCondBr(B->CreateICmpEQ(unwindExceptionClass, B->getInt64(seq_exc_class())), B->CreateCondBr(
tc.exceptionRouteBlock, externalExcBlock); B->CreateICmpEQ(unwindExceptionClass, B->getInt64(SEQ_EXCEPTION_CLASS)),
tc.exceptionRouteBlock, externalExcBlock);
// external exception (currently assumed to be unreachable) // external exception (currently assumed to be unreachable)
B->SetInsertPoint(externalExcBlock); B->SetInsertPoint(externalExcBlock);

110
codon/compiler/engine.cpp
View File

@ -8,86 +8,62 @@
namespace codon { namespace codon {
namespace jit { namespace jit {
void Engine::handleLazyCallThroughError() { Engine::Engine() : jit(), debug(nullptr) {
llvm::errs() << "LazyCallThrough error: Could not find function body"; auto eb = llvm::EngineBuilder();
exit(1); eb.setMArch(llvm::codegen::getMArch());
} eb.setMCPU(llvm::codegen::getCPUStr());
eb.setMAttrs(llvm::codegen::getFeatureList());
llvm::Expected<llvm::orc::ThreadSafeModule> auto target = eb.selectTarget();
Engine::optimizeModule(llvm::orc::ThreadSafeModule module, auto layout = target->createDataLayout();
const llvm::orc::MaterializationResponsibility &R) { auto epc = llvm::cantFail(llvm::orc::SelfExecutorProcessControl::Create(
module.withModuleDo([](llvm::Module &module) { std::make_shared<llvm::orc::SymbolStringPool>()));
ir::optimize(&module, /*debug=*/false, /*jit=*/true);
});
return std::move(module);
}
Engine::Engine(std::unique_ptr<llvm::orc::ExecutionSession> sess, llvm::orc::LLJITBuilder builder;
std::unique_ptr<llvm::orc::EPCIndirectionUtils> epciu, builder.setDataLayout(layout);
llvm::orc::JITTargetMachineBuilder jtmb, llvm::DataLayout layout) builder.setObjectLinkingLayerCreator(
: sess(std::move(sess)), epciu(std::move(epciu)), layout(std::move(layout)), [&](llvm::orc::ExecutionSession &es, const llvm::Triple &triple)
mangle(*this->sess, this->layout), -> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> {
objectLayer(*this->sess, auto L = std::make_unique<llvm::orc::ObjectLinkingLayer>(
[]() { return std::make_unique<BoehmGCMemoryManager>(); }), es, llvm::cantFail(BoehmGCJITLinkMemoryManager::Create()));
compileLayer(*this->sess, objectLayer, L->addPlugin(std::make_unique<llvm::orc::EHFrameRegistrationPlugin>(
std::make_unique<llvm::orc::ConcurrentIRCompiler>(std::move(jtmb))), es, llvm::cantFail(llvm::orc::EPCEHFrameRegistrar::Create(es))));
optimizeLayer(*this->sess, compileLayer, optimizeModule), L->addPlugin(std::make_unique<llvm::orc::DebugObjectManagerPlugin>(
codLayer(*this->sess, optimizeLayer, this->epciu->getLazyCallThroughManager(), es, llvm::cantFail(llvm::orc::createJITLoaderGDBRegistrar(es))));
[this] { return this->epciu->createIndirectStubsManager(); }), auto dbPlugin = std::make_unique<DebugPlugin>();
mainJD(this->sess->createBareJITDylib("<main>")), this->debug = dbPlugin.get();
dbListener(std::make_unique<DebugListener>()) { L->addPlugin(std::move(dbPlugin));
mainJD.addGenerator( L->setAutoClaimResponsibilityForObjectSymbols(true);
return L;
});
builder.setJITTargetMachineBuilder(
llvm::orc::JITTargetMachineBuilder(target->getTargetTriple()));
jit = llvm::cantFail(builder.create());
jit->getMainJITDylib().addGenerator(
llvm::cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess( llvm::cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
layout.getGlobalPrefix()))); layout.getGlobalPrefix())));
objectLayer.setAutoClaimResponsibilityForObjectSymbols(true);
objectLayer.registerJITEventListener(*dbListener);
}
Engine::~Engine() { jit->getIRTransformLayer().setTransform(
if (auto err = sess->endSession()) [&](llvm::orc::ThreadSafeModule module,
sess->reportError(std::move(err)); const llvm::orc::MaterializationResponsibility &R) {
if (auto err = epciu->cleanup()) module.withModuleDo([](llvm::Module &module) {
sess->reportError(std::move(err)); ir::optimize(&module, /*debug=*/false, /*jit=*/true);
} });
return std::move(module);
llvm::Expected<std::unique_ptr<Engine>> Engine::create() { });
auto epc = llvm::orc::SelfExecutorProcessControl::Create();
if (!epc)
return epc.takeError();
auto sess = std::make_unique<llvm::orc::ExecutionSession>(std::move(*epc));
auto epciu = llvm::orc::EPCIndirectionUtils::Create(*sess);
if (!epciu)
return epciu.takeError();
(*epciu)->createLazyCallThroughManager(
*sess, llvm::orc::ExecutorAddr::fromPtr(&handleLazyCallThroughError));
if (auto err = llvm::orc::setUpInProcessLCTMReentryViaEPCIU(**epciu))
return std::move(err);
llvm::orc::JITTargetMachineBuilder jtmb(
sess->getExecutorProcessControl().getTargetTriple());
auto layout = jtmb.getDefaultDataLayoutForTarget();
if (!layout)
return layout.takeError();
return std::make_unique<Engine>(std::move(sess), std::move(*epciu), std::move(jtmb),
std::move(*layout));
} }
llvm::Error Engine::addModule(llvm::orc::ThreadSafeModule module, llvm::Error Engine::addModule(llvm::orc::ThreadSafeModule module,
llvm::orc::ResourceTrackerSP rt) { llvm::orc::ResourceTrackerSP rt) {
if (!rt) if (!rt)
rt = mainJD.getDefaultResourceTracker(); rt = jit->getMainJITDylib().getDefaultResourceTracker();
return optimizeLayer.add(rt, std::move(module)); return jit->addIRModule(rt, std::move(module));
} }
llvm::Expected<llvm::orc::ExecutorSymbolDef> Engine::lookup(llvm::StringRef name) { llvm::Expected<llvm::orc::ExecutorAddr> Engine::lookup(llvm::StringRef name) {
return sess->lookup({&mainJD}, mangle(name.str())); return jit->lookup(name);
} }
} // namespace jit } // namespace jit

38
codon/compiler/engine.h
View File

@ -13,46 +13,22 @@ namespace jit {
class Engine { class Engine {
private: private:
std::unique_ptr<llvm::orc::ExecutionSession> sess; std::unique_ptr<llvm::orc::LLJIT> jit;
std::unique_ptr<llvm::orc::EPCIndirectionUtils> epciu; DebugPlugin *debug;
llvm::DataLayout layout;
llvm::orc::MangleAndInterner mangle;
llvm::orc::RTDyldObjectLinkingLayer objectLayer;
llvm::orc::IRCompileLayer compileLayer;
llvm::orc::IRTransformLayer optimizeLayer;
llvm::orc::CompileOnDemandLayer codLayer;
llvm::orc::JITDylib &mainJD;
std::unique_ptr<DebugListener> dbListener;
static void handleLazyCallThroughError();
static llvm::Expected<llvm::orc::ThreadSafeModule>
optimizeModule(llvm::orc::ThreadSafeModule module,
const llvm::orc::MaterializationResponsibility &R);
public: public:
Engine(std::unique_ptr<llvm::orc::ExecutionSession> sess, Engine();
std::unique_ptr<llvm::orc::EPCIndirectionUtils> epciu,
llvm::orc::JITTargetMachineBuilder jtmb, llvm::DataLayout layout);
~Engine(); const llvm::DataLayout &getDataLayout() const { return jit->getDataLayout(); }
static llvm::Expected<std::unique_ptr<Engine>> create(); llvm::orc::JITDylib &getMainJITDylib() { return jit->getMainJITDylib(); }
const llvm::DataLayout &getDataLayout() const { return layout; } DebugPlugin *getDebugListener() const { return debug; }
llvm::orc::JITDylib &getMainJITDylib() { return mainJD; }
DebugListener *getDebugListener() const { return dbListener.get(); }
llvm::Error addModule(llvm::orc::ThreadSafeModule module, llvm::Error addModule(llvm::orc::ThreadSafeModule module,
llvm::orc::ResourceTrackerSP rt = nullptr); llvm::orc::ResourceTrackerSP rt = nullptr);
llvm::Expected<llvm::orc::ExecutorSymbolDef> lookup(llvm::StringRef name); llvm::Expected<llvm::orc::ExecutorAddr> lookup(llvm::StringRef name);
}; };
} // namespace jit } // namespace jit

17
codon/compiler/jit.cpp
View File

@ -23,14 +23,9 @@ const std::string JIT_FILENAME = "<jit>";
} // namespace } // namespace
JIT::JIT(const std::string &argv0, const std::string &mode) JIT::JIT(const std::string &argv0, const std::string &mode)
: compiler(std::make_unique<Compiler>(argv0, Compiler::Mode::JIT)), engine(), : compiler(std::make_unique<Compiler>(argv0, Compiler::Mode::JIT)),
pydata(std::make_unique<PythonData>()), mode(mode) { engine(std::make_unique<Engine>()), pydata(std::make_unique<PythonData>()),
if (auto e = Engine::create()) { mode(mode) {
engine = std::move(e.get());
} else {
engine = {};
seqassertn(false, "JIT engine creation error");
}
compiler->getLLVMVisitor()->setJIT(true); compiler->getLLVMVisitor()->setJIT(true);
} }
@ -60,7 +55,7 @@ llvm::Error JIT::init() {
if (auto err = func.takeError()) if (auto err = func.takeError())
return err; return err;
auto *main = func->getAddress().toPtr<MainFunc>(); auto *main = func->toPtr<MainFunc>();
(*main)(0, nullptr); (*main)(0, nullptr);
return llvm::Error::success(); return llvm::Error::success();
} }
@ -174,7 +169,7 @@ llvm::Expected<void *> JIT::address(const ir::Func *input) {
if (auto err = func.takeError()) if (auto err = func.takeError())
return std::move(err); return std::move(err);
return (void *)func->getAddress().getValue(); return (void *)func->getValue();
} }
llvm::Expected<std::string> JIT::run(const ir::Func *input) { llvm::Expected<std::string> JIT::run(const ir::Func *input) {
@ -292,7 +287,7 @@ JITResult JIT::executePython(const std::string &name,
auto *wrapper = it->second; auto *wrapper = it->second;
const std::string name = ir::LLVMVisitor::getNameForFunction(wrapper); const std::string name = ir::LLVMVisitor::getNameForFunction(wrapper);
auto func = llvm::cantFail(engine->lookup(name)); auto func = llvm::cantFail(engine->lookup(name));
wrap = func.getAddress().toPtr<PyWrapperFunc>(); wrap = func.toPtr<PyWrapperFunc>();
} else { } else {
static int idx = 0; static int idx = 0;
auto wrapname = "__codon_wrapped__" + name + "_" + std::to_string(idx++); auto wrapname = "__codon_wrapped__" + name + "_" + std::to_string(idx++);

95
codon/runtime/exc.cpp
View File

@ -14,6 +14,52 @@
#include <string> #include <string>
#include <vector> #include <vector>
#ifdef __APPLE__
#include <TargetConditionals.h>
#if TARGET_OS_MAC && __arm64__
#define APPLE_SILICON
#endif
#endif
#ifdef APPLE_SILICON
#include "llvm/BinaryFormat/MachO.h"
#include <dlfcn.h>
// https://github.com/llvm/llvm-project/issues/49036
// Define a minimal mach header for JIT'd code.
static llvm::MachO::mach_header_64 fake_mach_header = {
.magic = llvm::MachO::MH_MAGIC_64,
.cputype = llvm::MachO::CPU_TYPE_ARM64,
.cpusubtype = llvm::MachO::CPU_SUBTYPE_ARM64_ALL,
.filetype = llvm::MachO::MH_DYLIB,
.ncmds = 0,
.sizeofcmds = 0,
.flags = 0,
.reserved = 0};
// Declare libunwind SPI types and functions.
struct unw_dynamic_unwind_sections {
uintptr_t dso_base;
uintptr_t dwarf_section;
size_t dwarf_section_length;
uintptr_t compact_unwind_section;
size_t compact_unwind_section_length;
};
int find_dynamic_unwind_sections(uintptr_t addr, unw_dynamic_unwind_sections *info) {
info->dso_base = (uintptr_t)&fake_mach_header;
info->dwarf_section = 0;
info->dwarf_section_length = 0;
info->compact_unwind_section = 0;
info->compact_unwind_section_length = 0;
return 1;
}
// Typedef for callback above.
typedef int (*unw_find_dynamic_unwind_sections)(
uintptr_t addr, struct unw_dynamic_unwind_sections *info);
#endif
struct BacktraceFrame { struct BacktraceFrame {
char *function; char *function;
char *filename; char *filename;
@ -90,24 +136,6 @@ template <typename Type_> static uintptr_t ReadType(const uint8_t *&p) {
} }
} // namespace } // namespace
static int64_t ourBaseFromUnwindOffset;
static const unsigned char ourBaseExcpClassChars[] = {'o', 'b', 'j', '\0',
's', 'e', 'q', '\0'};
static uint64_t genClass(const unsigned char classChars[], size_t classCharsSize) {
uint64_t ret = classChars[0];
for (unsigned i = 1; i < classCharsSize; i++) {
ret <<= 8;
ret += classChars[i];
}
return ret;
}
static uint64_t ourBaseExceptionClass = 0;
struct OurExceptionType_t { struct OurExceptionType_t {
int type; int type;
}; };
@ -131,15 +159,22 @@ struct SeqExcHeader_t {
void *python_type; void *python_type;
}; };
void seq_exc_init() { void seq_exc_init(int flags) {
ourBaseFromUnwindOffset = seq_exc_offset(); #ifdef APPLE_SILICON
ourBaseExceptionClass = seq_exc_class(); if (!(flags & SEQ_FLAG_STANDALONE)) {
if (auto *unw_add_find_dynamic_unwind_sections =
(int (*)(unw_find_dynamic_unwind_sections find_dynamic_unwind_sections))
dlsym(RTLD_DEFAULT, "__unw_add_find_dynamic_unwind_sections")) {
unw_add_find_dynamic_unwind_sections(find_dynamic_unwind_sections);
}
}
#endif
} }
static void seq_delete_exc(_Unwind_Exception *expToDelete) { static void seq_delete_exc(_Unwind_Exception *expToDelete) {
if (!expToDelete || expToDelete->exception_class != ourBaseExceptionClass) if (!expToDelete || expToDelete->exception_class != SEQ_EXCEPTION_CLASS)
return; return;
auto *exc = (OurException *)((char *)expToDelete + ourBaseFromUnwindOffset); auto *exc = (OurException *)((char *)expToDelete + seq_exc_offset());
if (seq_flags & SEQ_FLAG_DEBUG) { if (seq_flags & SEQ_FLAG_DEBUG) {
exc->bt.free(); exc->bt.free();
} }
@ -160,7 +195,7 @@ SEQ_FUNC void *seq_alloc_exc(int type, void *obj) {
assert(e); assert(e);
e->type.type = type; e->type.type = type;
e->obj = obj; e->obj = obj;
e->unwindException.exception_class = ourBaseExceptionClass; e->unwindException.exception_class = SEQ_EXCEPTION_CLASS;
e->unwindException.exception_cleanup = seq_delete_unwind_exc; e->unwindException.exception_cleanup = seq_delete_unwind_exc;
if (seq_flags & SEQ_FLAG_DEBUG) { if (seq_flags & SEQ_FLAG_DEBUG) {
e->bt.frames = nullptr; e->bt.frames = nullptr;
@ -420,11 +455,11 @@ static bool handleActionValue(int64_t *resultAction, uint8_t TTypeEncoding,
_Unwind_Exception *exceptionObject) { _Unwind_Exception *exceptionObject) {
bool ret = false; bool ret = false;
if (!resultAction || !exceptionObject || (exceptionClass != ourBaseExceptionClass)) if (!resultAction || !exceptionObject || (exceptionClass != SEQ_EXCEPTION_CLASS))
return ret; return ret;
auto *excp = (struct OurBaseException_t *)(((char *)exceptionObject) + auto *excp =
ourBaseFromUnwindOffset); (struct OurBaseException_t *)(((char *)exceptionObject) + seq_exc_offset());
OurExceptionType_t *excpType = &(excp->type); OurExceptionType_t *excpType = &(excp->type);
seq_int_t type = excpType->type; seq_int_t type = excpType->type;
@ -523,7 +558,7 @@ static _Unwind_Reason_Code handleLsda(int version, const uint8_t *lsda,
// Note: Action value // Note: Action value
uintptr_t actionEntry = readULEB128(&callSitePtr); uintptr_t actionEntry = readULEB128(&callSitePtr);
if (exceptionClass != ourBaseExceptionClass) { if (exceptionClass != SEQ_EXCEPTION_CLASS) {
// We have been notified of a foreign exception being thrown, // We have been notified of a foreign exception being thrown,
// and we therefore need to execute cleanup landing pads // and we therefore need to execute cleanup landing pads
actionEntry = 0; actionEntry = 0;
@ -596,10 +631,6 @@ SEQ_FUNC int64_t seq_exc_offset() {
return (int64_t)((uintptr_t)&dummy - (uintptr_t) & (dummy.unwindException)); return (int64_t)((uintptr_t)&dummy - (uintptr_t) & (dummy.unwindException));
} }
SEQ_FUNC uint64_t seq_exc_class() {
return genClass(ourBaseExcpClassChars, sizeof(ourBaseExcpClassChars));
}
std::string codon::runtime::makeBacktraceFrameString(uintptr_t pc, std::string codon::runtime::makeBacktraceFrameString(uintptr_t pc,
const std::string &func, const std::string &func,
const std::string &file, int line, const std::string &file, int line,

4
codon/runtime/lib.cpp
View File

@ -41,7 +41,7 @@ extern "C" void __kmpc_set_gc_callbacks(gc_setup_callback get_stack_base,
gc_roots_callback add_roots, gc_roots_callback add_roots,
gc_roots_callback del_roots); gc_roots_callback del_roots);
void seq_exc_init(); void seq_exc_init(int flags);
#ifdef CODON_GPU #ifdef CODON_GPU
void seq_nvptx_init(); void seq_nvptx_init();
@ -55,7 +55,7 @@ SEQ_FUNC void seq_init(int flags) {
GC_allow_register_threads(); GC_allow_register_threads();
__kmpc_set_gc_callbacks(GC_get_stack_base, (gc_setup_callback)GC_register_my_thread, __kmpc_set_gc_callbacks(GC_get_stack_base, (gc_setup_callback)GC_register_my_thread,
GC_add_roots, GC_remove_roots); GC_add_roots, GC_remove_roots);
seq_exc_init(); seq_exc_init(flags);
#ifdef CODON_GPU #ifdef CODON_GPU
seq_nvptx_init(); seq_nvptx_init();
#endif #endif

3
codon/runtime/lib.h
View File

@ -17,6 +17,8 @@
#define SEQ_FLAG_CAPTURE_OUTPUT (1 << 1) // capture writes to stdout/stderr #define SEQ_FLAG_CAPTURE_OUTPUT (1 << 1) // capture writes to stdout/stderr
#define SEQ_FLAG_STANDALONE (1 << 2) // compiled as a standalone object/binary #define SEQ_FLAG_STANDALONE (1 << 2) // compiled as a standalone object/binary
#define SEQ_EXCEPTION_CLASS 0x6f626a0073657100
#define SEQ_FUNC extern "C" #define SEQ_FUNC extern "C"
typedef int64_t seq_int_t; typedef int64_t seq_int_t;
@ -74,7 +76,6 @@ SEQ_FUNC _Unwind_Reason_Code seq_personality(int version, _Unwind_Action actions
_Unwind_Exception *exceptionObject, _Unwind_Exception *exceptionObject,
_Unwind_Context *context); _Unwind_Context *context);
SEQ_FUNC int64_t seq_exc_offset(); SEQ_FUNC int64_t seq_exc_offset();
SEQ_FUNC uint64_t seq_exc_class();
SEQ_FUNC seq_str_t seq_str_int(seq_int_t n, seq_str_t format, bool *error); SEQ_FUNC seq_str_t seq_str_int(seq_int_t n, seq_str_t format, bool *error);
SEQ_FUNC seq_str_t seq_str_uint(seq_int_t n, seq_str_t format, bool *error); SEQ_FUNC seq_str_t seq_str_uint(seq_int_t n, seq_str_t format, bool *error);