Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JIT fixes for Apple Silicon #575

Merged
merged 1 commit into from
Aug 5, 2024
Merged

JIT fixes for Apple Silicon #575

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions codon/cir/llvm/llvisitor.cpp
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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