Skip to content

Commit

Permalink
JIT fixes for Apple Silicon (exaloop#575)
Browse files Browse the repository at this point in the history
  • Loading branch information
@arshajii
arshajii authored Aug 5, 2024
1 parent 11d281d commit fa2904c
Show file tree
Hide file tree
Showing 7 changed files with 128 additions and 148 deletions.
5 changes: 3 additions & 2 deletions codon/cir/llvm/llvisitor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2941,8 +2941,9 @@ void LLVMVisitor::visit(const TryCatchFlow *x) {
0));

// check for foreign exceptions
B->CreateCondBr(B->CreateICmpEQ(unwindExceptionClass, B->getInt64(seq_exc_class())),
tc.exceptionRouteBlock, externalExcBlock);
B->CreateCondBr(
B->CreateICmpEQ(unwindExceptionClass, B->getInt64(SEQ_EXCEPTION_CLASS)),
tc.exceptionRouteBlock, externalExcBlock);

// external exception (currently assumed to be unreachable)
B->SetInsertPoint(externalExcBlock);
Expand Down
114 changes: 45 additions & 69 deletions codon/compiler/engine.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,86 +8,62 @@
namespace codon {
namespace jit {

void Engine::handleLazyCallThroughError() {
llvm::errs() << "LazyCallThrough error: Could not find function body";
exit(1);
}

llvm::Expected<llvm::orc::ThreadSafeModule>
Engine::optimizeModule(llvm::orc::ThreadSafeModule module,
const llvm::orc::MaterializationResponsibility &R) {
module.withModuleDo([](llvm::Module &module) {
ir::optimize(&module, /*debug=*/false, /*jit=*/true);
});
return std::move(module);
}

Engine::Engine(std::unique_ptr<llvm::orc::ExecutionSession> sess,
std::unique_ptr<llvm::orc::EPCIndirectionUtils> epciu,
llvm::orc::JITTargetMachineBuilder jtmb, llvm::DataLayout layout)
: sess(std::move(sess)), epciu(std::move(epciu)), layout(std::move(layout)),
mangle(*this->sess, this->layout),
objectLayer(*this->sess,
[]() { return std::make_unique<BoehmGCMemoryManager>(); }),
compileLayer(*this->sess, objectLayer,
std::make_unique<llvm::orc::ConcurrentIRCompiler>(std::move(jtmb))),
optimizeLayer(*this->sess, compileLayer, optimizeModule),
codLayer(*this->sess, optimizeLayer, this->epciu->getLazyCallThroughManager(),
[this] { return this->epciu->createIndirectStubsManager(); }),
mainJD(this->sess->createBareJITDylib("<main>")),
dbListener(std::make_unique<DebugListener>()) {
mainJD.addGenerator(
Engine::Engine() : jit(), debug(nullptr) {
auto eb = llvm::EngineBuilder();
eb.setMArch(llvm::codegen::getMArch());
eb.setMCPU(llvm::codegen::getCPUStr());
eb.setMAttrs(llvm::codegen::getFeatureList());

auto target = eb.selectTarget();
auto layout = target->createDataLayout();
auto epc = llvm::cantFail(llvm::orc::SelfExecutorProcessControl::Create(
std::make_shared<llvm::orc::SymbolStringPool>()));

llvm::orc::LLJITBuilder builder;
builder.setDataLayout(layout);
builder.setObjectLinkingLayerCreator(
[&](llvm::orc::ExecutionSession &es, const llvm::Triple &triple)
-> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> {
auto L = std::make_unique<llvm::orc::ObjectLinkingLayer>(
es, llvm::cantFail(BoehmGCJITLinkMemoryManager::Create()));
L->addPlugin(std::make_unique<llvm::orc::EHFrameRegistrationPlugin>(
es, llvm::cantFail(llvm::orc::EPCEHFrameRegistrar::Create(es))));
L->addPlugin(std::make_unique<llvm::orc::DebugObjectManagerPlugin>(
es, llvm::cantFail(llvm::orc::createJITLoaderGDBRegistrar(es))));
auto dbPlugin = std::make_unique<DebugPlugin>();
this->debug = dbPlugin.get();
L->addPlugin(std::move(dbPlugin));
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(
layout.getGlobalPrefix())));
objectLayer.setAutoClaimResponsibilityForObjectSymbols(true);
objectLayer.registerJITEventListener(*dbListener);
}

Engine::~Engine() {
if (auto err = sess->endSession())
sess->reportError(std::move(err));
if (auto err = epciu->cleanup())
sess->reportError(std::move(err));
}

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));
jit->getIRTransformLayer().setTransform(
[&](llvm::orc::ThreadSafeModule module,
const llvm::orc::MaterializationResponsibility &R) {
module.withModuleDo([](llvm::Module &module) {
ir::optimize(&module, /*debug=*/false, /*jit=*/true);
});
return std::move(module);
});
}

llvm::Error Engine::addModule(llvm::orc::ThreadSafeModule module,
llvm::orc::ResourceTrackerSP 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) {
return sess->lookup({&mainJD}, mangle(name.str()));
llvm::Expected<llvm::orc::ExecutorAddr> Engine::lookup(llvm::StringRef name) {
return jit->lookup(name);
}

} // namespace jit
Expand Down
38 changes: 7 additions & 31 deletions codon/compiler/engine.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,46 +13,22 @@ namespace jit {

class Engine {
private:
std::unique_ptr<llvm::orc::ExecutionSession> sess;
std::unique_ptr<llvm::orc::EPCIndirectionUtils> epciu;

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);
std::unique_ptr<llvm::orc::LLJIT> jit;
DebugPlugin *debug;

public:
Engine(std::unique_ptr<llvm::orc::ExecutionSession> sess,
std::unique_ptr<llvm::orc::EPCIndirectionUtils> epciu,
llvm::orc::JITTargetMachineBuilder jtmb, llvm::DataLayout layout);

~Engine();

static llvm::Expected<std::unique_ptr<Engine>> create();
Engine();

const llvm::DataLayout &getDataLayout() const { return layout; }
const llvm::DataLayout &getDataLayout() const { return jit->getDataLayout(); }

llvm::orc::JITDylib &getMainJITDylib() { return mainJD; }
llvm::orc::JITDylib &getMainJITDylib() { return jit->getMainJITDylib(); }

DebugListener *getDebugListener() const { return dbListener.get(); }
DebugPlugin *getDebugListener() const { return debug; }

llvm::Error addModule(llvm::orc::ThreadSafeModule module,
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
Expand Down
17 changes: 6 additions & 11 deletions codon/compiler/jit.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,9 @@ const std::string JIT_FILENAME = "<jit>";
} // namespace

JIT::JIT(const std::string &argv0, const std::string &mode)
: compiler(std::make_unique<Compiler>(argv0, Compiler::Mode::JIT)), engine(),
pydata(std::make_unique<PythonData>()), mode(mode) {
if (auto e = Engine::create()) {
engine = std::move(e.get());
} else {
engine = {};
seqassertn(false, "JIT engine creation error");
}
: compiler(std::make_unique<Compiler>(argv0, Compiler::Mode::JIT)),
engine(std::make_unique<Engine>()), pydata(std::make_unique<PythonData>()),
mode(mode) {
compiler->getLLVMVisitor()->setJIT(true);
}

Expand Down Expand Up @@ -60,7 +55,7 @@ llvm::Error JIT::init() {
if (auto err = func.takeError())
return err;

auto *main = func->getAddress().toPtr<MainFunc>();
auto *main = func->toPtr<MainFunc>();
(*main)(0, nullptr);
return llvm::Error::success();
}
Expand Down Expand Up @@ -174,7 +169,7 @@ llvm::Expected<void *> JIT::address(const ir::Func *input) {
if (auto err = func.takeError())
return std::move(err);

return (void *)func->getAddress().getValue();
return (void *)func->getValue();
}

llvm::Expected<std::string> JIT::run(const ir::Func *input) {
Expand Down Expand Up @@ -292,7 +287,7 @@ JITResult JIT::executePython(const std::string &name,
auto *wrapper = it->second;
const std::string name = ir::LLVMVisitor::getNameForFunction(wrapper);
auto func = llvm::cantFail(engine->lookup(name));
wrap = func.getAddress().toPtr<PyWrapperFunc>();
wrap = func.toPtr<PyWrapperFunc>();
} else {
static int idx = 0;
auto wrapname = "__codon_wrapped__" + name + "_" + std::to_string(idx++);
Expand Down
95 changes: 63 additions & 32 deletions codon/runtime/exc.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,52 @@
#include <string>
#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 {
char *function;
char *filename;
Expand Down Expand Up @@ -90,24 +136,6 @@ template <typename Type_> static uintptr_t ReadType(const uint8_t *&p) {
}
} // 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 {
int type;
};
Expand All @@ -131,15 +159,22 @@ struct SeqExcHeader_t {
void *python_type;
};

void seq_exc_init() {
ourBaseFromUnwindOffset = seq_exc_offset();
ourBaseExceptionClass = seq_exc_class();
void seq_exc_init(int flags) {
#ifdef APPLE_SILICON
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) {
if (!expToDelete || expToDelete->exception_class != ourBaseExceptionClass)
if (!expToDelete || expToDelete->exception_class != SEQ_EXCEPTION_CLASS)
return;
auto *exc = (OurException *)((char *)expToDelete + ourBaseFromUnwindOffset);
auto *exc = (OurException *)((char *)expToDelete + seq_exc_offset());
if (seq_flags & SEQ_FLAG_DEBUG) {
exc->bt.free();
}
Expand All @@ -160,7 +195,7 @@ SEQ_FUNC void *seq_alloc_exc(int type, void *obj) {
assert(e);
e->type.type = type;
e->obj = obj;
e->unwindException.exception_class = ourBaseExceptionClass;
e->unwindException.exception_class = SEQ_EXCEPTION_CLASS;
e->unwindException.exception_cleanup = seq_delete_unwind_exc;
if (seq_flags & SEQ_FLAG_DEBUG) {
e->bt.frames = nullptr;
Expand Down Expand Up @@ -420,11 +455,11 @@ static bool handleActionValue(int64_t *resultAction, uint8_t TTypeEncoding,
_Unwind_Exception *exceptionObject) {
bool ret = false;

if (!resultAction || !exceptionObject || (exceptionClass != ourBaseExceptionClass))
if (!resultAction || !exceptionObject || (exceptionClass != SEQ_EXCEPTION_CLASS))
return ret;

auto *excp = (struct OurBaseException_t *)(((char *)exceptionObject) +
ourBaseFromUnwindOffset);
auto *excp =
(struct OurBaseException_t *)(((char *)exceptionObject) + seq_exc_offset());
OurExceptionType_t *excpType = &(excp->type);
seq_int_t type = excpType->type;

Expand Down Expand Up @@ -523,7 +558,7 @@ static _Unwind_Reason_Code handleLsda(int version, const uint8_t *lsda,
// Note: Action value
uintptr_t actionEntry = readULEB128(&callSitePtr);

if (exceptionClass != ourBaseExceptionClass) {
if (exceptionClass != SEQ_EXCEPTION_CLASS) {
// We have been notified of a foreign exception being thrown,
// and we therefore need to execute cleanup landing pads
actionEntry = 0;
Expand Down Expand Up @@ -596,10 +631,6 @@ SEQ_FUNC int64_t seq_exc_offset() {
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,
const std::string &func,
const std::string &file, int line,
Expand Down
Loading

0 comments on commit fa2904c

Please sign in to comment.