Skip to content

The JIT driver: calls compiled into Julia's JIT, cached by CodeInstance - #20

Draft
vchuravy wants to merge 2 commits into
vc/jc-18-emit-intofrom
vc/jc-19-jit
Draft

vchuravy wants to merge 2 commits into
vc/jc-18-emit-intofrom
vc/jc-19-jit

Conversation

@vchuravy

Copy link
Copy Markdown
Contributor

Stacked on #19. This completes plan P4's drivers. The runtime hook remains, on its own stack.

jit = JuliaCodegen.JIT(; config=CodegenConfig(), pipeline=JuliaCppPipeline(), relocatable=true)
jf  = JuliaCodegen.jit_compile(jit, Tuple{typeof(f), Int})   # ::JITFunction (ci, fptr, specptr, fallbacks)
jf(f, 3)                                                      # through the jlcall entry
JuliaCodegen.jit_call(jit, f, 3)

It uses no unmerged Julia API; everything goes through LLVM.jl's JuliaOJIT.

  • IR: JITInterpreter is native inference with a cache owner of its own (:JuliaCodegen). Its CodeInstances keep uncompressed sources (may_discard_trees/may_compress are false). collect_codeinfos(interp, mi) gathers the closed set that the call :invokes. Functions with a constant result have no kept source, so they are re-inferred with typeinf_code. Varargs callees are left to the runtime. emit_into! (emit_into!: compile a closed set of CodeInstances into the caller's module #19) compiles the set.
  • Cache: keyed by the CodeInstance, looked up in the interpreter's cache first, because inferring again returns a fresh CodeInstance. Redefining a method therefore recompiles what depends on it, which the plan's "invalidation" item required.
  • Linking:
    • every set goes into a module of its own, and only the root's entries are external, under process-unique names (jlcodegen_N, shared by all JITs because they share Julia's JIT);
    • the module is checked to be relocatable before linking;
    • literal objects are then bound to their addresses and rooted in the JIT;
    • jl_get_pgcstack_resolved is bound to the runtime's getter (moved from the test JIT into the driver);
    • the pipeline runs, and the module goes into a JITDylib of its own where supported.
  • Concurrency: one lock per JIT.

Not included: publishing compiled code into CodeInstances (ci.invoke/specptr), so that native callers use it. That belongs to the hook stack.

Tests (test/jitdriver.jl)

  • a differential corpus through jit_call;
  • cache hits for the same CodeInstance;
  • invalidation: redefining a callee gives a new CodeInstance and new code;
  • fallbacks: a stdcall callee is recorded with :fallback and throws with :error;
  • varargs callees through the runtime;
  • the specsig entry called as a C function (gcstack_arg=false);
  • literals surviving a full GC;
  • concurrent compilation from 8 threaded tasks (one cache entry).

The JIT tests also pass under JULIACODEGEN_GC_STRESS=1 and with -t 4.

Also: the memory test's allocation count is now the least of three measurements. It failed once in a full nightly run and passed on reruns.

Tested locally with julia +1.12, +1.13 and +nightly: 6955, 6961 and 6960 tests pass.

🤖 Generated with Claude Code

`JIT(; config, pipeline, relocatable)`, `jit_compile(jit, tt) -> JITFunction`
and `jit_call(jit, f, args...)` compile calls with JuliaCodegen into Julia's
JIT through LLVM.jl's `JuliaOJIT`, without unmerged Julia APIs:

- IR: `JITInterpreter`, native inference whose CodeInstances have an owner
  of their own and keep their uncompressed sources. `collect_codeinfos`
  gathers the closed set a call needs, re-inferring the code of functions
  with a constant result, and `emit_into!` compiles it.
- Cache: keyed by the CodeInstance, looked up in the interpreter's cache
  first, so redefining a method recompiles what depends on it.
- Linking: each set gets a module of its own; only the root's entries are
  external, under names unique in the process. The literal objects are bound
  to their addresses and kept alive. `jl_get_pgcstack_resolved` is bound
  (moved here from the test JIT). The pipeline runs.
- One lock per JIT.

`JITFunction` has the jlcall entry (`fptr`, which calling it uses), the
specsig entry (`specptr`; a C function with `gcstack_arg=false`) and the
fallbacks.

Also: the memory test's allocation count is the least of three
measurements; it was flaky in the full nightly suite.

Assisted-by: Claude Code (Opus 5.5)
Assisted-by: Claude Code (Opus 5.5)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant