Expose the Julia JIT with a C API #49858

gbaraldi · 2023-05-17T16:27:31Z

This is a proof of concept to expose the Julia JIT to foreign users. The idea is to use this interface to solve #49056.

For now it only half works
Before

julia> call_delayed(foo, 3)
ERROR: heyy
Stacktrace:
 [1] macro expansion
   @ ./REPL[19]:337 [inlined]
 [2] abi_call
   @ ./REPL[19]:261 [inlined]
 [3] call_delayed(f::typeof(foo), args::Int64)
   @ Main.LazyCodegen ./REPL[19]:347
 [4] top-level scope
   @ REPL[58]:1

With this PR + the LLVM.jl changes JuliaLLVM/LLVM.jl#346

julia> call_delayed(foo, 3)
ERROR: heyy
Stacktrace:
 [1] foo
   @ ./REPL[16]:3
 [2] macro expansion
   @ Main.LazyCodegen ./REPL[14]:325 [inlined]
 [3] abi_call(f::Ptr{Nothing}, rt::Type{Int64}, tt::Type{Tuple{Int64}}, func::typeof(foo), args::Int64)
   @ Main.LazyCodegen ./REPL[14]:249 [inlined]
 [4] call_delayed(f::typeof(foo), args::Int64)
   @ Main.LazyCodegen ./REPL[14]:335
 [5] top-level scope
   @ REPL[17]:1

This still misses some traces, but it's a start.

I want to move those functions to a separate file and maybe some of them to LLVMExtra

Apparently it was missing debug info. Starting julia with -g2 lets us get the full trace (with no details because of no code/method_instances)

julia> declared_bar(3)
ERROR: heyehye
Stacktrace:
 [1] error
   @ ./error.jl:35
 [2] foo
   @ ./REPL[13]:3
 [3] bar
   @ ./REPL[22]:1
 [4] declared_bar(x::Int64)
   @ Main ./REPL[31]:2
 [5] top-level scope
   @ REPL[32]:1

src/codegen-stubs.c

pchintalapudi · 2023-05-17T21:18:34Z

Is there a reason to allow people to add LLVM IR to an IRCompileLayer instead of just requiring them to compile to an object file on their own?

gbaraldi · 2023-05-18T13:36:30Z

I added in order to have something to call emit on. It's probably not super necessary.

maleadt · 2023-05-22T09:33:59Z

src/llvm_api.cpp

+JL_DLLEXPORT_CODEGEN const char *
+LLVMExtraJLJITGetDataLayoutString_impl(JuliaOJITRef JIT)
+{
+    return unwrap(JIT)->getDataLayout().getStringRepresentation().c_str();


The C API supports DataLayout objects, so just return it directly?

Oh, I didn't know that, because the other functions (GetDataLayout,SetDatalayout etc) that move it around all use C strings

Ah, I see that the existing APIs use strings too (LLVMOrcLLJITGetDataLayoutStr), so I'm OK keeping it as strings for consistency.

maleadt · 2023-05-22T09:36:59Z

I'm not familiar with ORC or our IR compiler layer, so I'll defer to @vchuravy and @pchintalapudi.

vchuravy · 2023-05-24T19:14:21Z

src/jitlayers.cpp

@@ -1289,6 +1289,7 @@ JuliaOJIT::JuliaOJIT()
    ES(cantFail(orc::SelfExecutorProcessControl::Create())),
    GlobalJD(ES.createBareJITDylib("JuliaGlobals")),
    JD(ES.createBareJITDylib("JuliaOJIT")),
+    ExternalJD(ES.createBareJITDylib("JuliaExternal")),


Long-term we might want many ExternalJDs and was considering a design where we had one JD per world in Julia.

We can add as many JDs as we want with the ExecutionSession. The difference is that this one is linked in by default.

vtjnash · 2023-05-30T15:43:54Z

src/jitlayers.cpp

@@ -1385,6 +1388,9 @@ JuliaOJIT::JuliaOJIT()
    }

    JD.addToLinkOrder(GlobalJD, orc::JITDylibLookupFlags::MatchExportedSymbolsOnly);
+    JD.addToLinkOrder(ExternalJD, orc::JITDylibLookupFlags::MatchExportedSymbolsOnly);
+    ExternalJD.addToLinkOrder(GlobalJD, orc::JITDylibLookupFlags::MatchExportedSymbolsOnly);
+    ExternalJD.addToLinkOrder(JD, orc::JITDylibLookupFlags::MatchExportedSymbolsOnly);


There should not be any way for ExternalJD to access JD or vice versa. Neither of those should be accidentally breaking the content of the other.

I guess this one is kind of wrong, but we do want to keep the ability to llvmcall a function from an external JD to avoid doing the trampoline/function pointer dance we do right now.

That seems likely to go very badly for users. The llvmcall itself right now should not be able to access JD state either, but that is a different bug.

How would it go bad for the users? (llvmcall currently aborts if it doesn't find a symbol but that's separate)

JD is an implementation detail and not one that is stable, or one that we think is good right now. We could fix the llvmcall crash by checking that it does not contain declarations (except of intrinsic things) and otherwise reject it when we parsed it. That would do the right thing.

vtjnash · 2023-05-30T15:44:11Z

src/jitlayers.cpp

+    orc::JITDylib* SearchOrders[3] = {&JD, &GlobalJD, &ExternalJD};
+    ArrayRef<orc::JITDylib*> SearchOrder = makeArrayRef(&SearchOrders[0], ExportedSymbolsOnly ? 3 : 1);


This change should be reverted?

vtjnash · 2023-05-30T15:46:18Z

src/jitlayers.cpp

+Expected<JITEvaluatedSymbol> JuliaOJIT::findExternalJDSymbol(StringRef Name, bool ExternalJDOnly)
+{
+    orc::JITDylib* SearchOrders[3] = {&ExternalJD, &GlobalJD, &JD};


This should take some sort of handle that was returned by from addExternalModule / addObjectFile to do the lookup inside of. No global variables (esp. GlobalJD and JD) should be accessed here since their content is unreliable.

Expose the Julia JIT with a C API

c8e5305

gbaraldi mentioned this pull request May 17, 2023

Add Interface to julia OJIT JuliaLLVM/LLVM.jl#346

Merged

vchuravy reviewed May 17, 2023

View reviewed changes

src/codegen-stubs.c Show resolved Hide resolved

vchuravy requested review from pchintalapudi and maleadt May 17, 2023 18:50

Cleanup of code

9b418ed

gbaraldi force-pushed the ojit-llvm-c branch from 207b483 to 9b418ed Compare May 19, 2023 20:34

gbaraldi marked this pull request as ready for review May 19, 2023 20:34

maleadt reviewed May 22, 2023

View reviewed changes

gbaraldi mentioned this pull request May 22, 2023

Demonstrate backtraces using the juliaojit EnzymeAD/Enzyme.jl#863

Merged

Change name and add test

fe67d67

gbaraldi force-pushed the ojit-llvm-c branch from 7da1b67 to fe67d67 Compare May 24, 2023 18:12

Make analyse gc a bit happier

5ad6f4a

vchuravy reviewed May 24, 2023

View reviewed changes

vchuravy approved these changes May 24, 2023

View reviewed changes

gbaraldi force-pushed the ojit-llvm-c branch from 784e99d to 62f1dc9 Compare May 29, 2023 16:38

Add JL_NOTSAFEPOINT to all functions and lambdas

f8dee84

gbaraldi force-pushed the ojit-llvm-c branch from 62f1dc9 to f8dee84 Compare May 29, 2023 16:49

vchuravy merged commit 957972e into JuliaLang:master May 29, 2023

vtjnash reviewed May 30, 2023

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Expose the Julia JIT with a C API #49858

Expose the Julia JIT with a C API #49858

gbaraldi commented May 17, 2023 •

edited

Loading

pchintalapudi commented May 17, 2023

gbaraldi commented May 18, 2023

maleadt May 22, 2023

gbaraldi May 22, 2023

maleadt May 22, 2023

maleadt commented May 22, 2023

vchuravy May 24, 2023

gbaraldi May 24, 2023

vtjnash May 30, 2023

gbaraldi May 30, 2023

vtjnash May 30, 2023

gbaraldi May 30, 2023

vtjnash May 30, 2023

vtjnash May 30, 2023

vtjnash May 30, 2023

		orc::JITDylib* SearchOrders[3] = {&JD, &GlobalJD, &ExternalJD};
		ArrayRef<orc::JITDylib*> SearchOrder = makeArrayRef(&SearchOrders[0], ExportedSymbolsOnly ? 3 : 1);

Expose the Julia JIT with a C API #49858

Expose the Julia JIT with a C API #49858

Conversation

gbaraldi commented May 17, 2023 • edited Loading

pchintalapudi commented May 17, 2023

gbaraldi commented May 18, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

maleadt commented May 22, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gbaraldi commented May 17, 2023 •

edited

Loading