Skip to content

Finer grained precompile native code cache (part 1)#58592

Open
xal-0 wants to merge 10 commits intoJuliaLang:masterfrom
xal-0:aotcompile-cache
Open

Finer grained precompile native code cache (part 1)#58592
xal-0 wants to merge 10 commits intoJuliaLang:masterfrom
xal-0:aotcompile-cache

Conversation

@xal-0
Copy link
Copy Markdown
Member

@xal-0 xal-0 commented May 30, 2025

Overview

This pull request adds a new mode for precompiling sysimages and pkgimages that caches the results of compiling each LLVM module, reducing the time spent emitting native code for CodeInstances that generate identical LLVM IR. For now, it works only when using the ahead-of-time compiler, but the approach is also valid for JITed code.

Usage

Set JULIA_NATIVE_CACHE=<dir> to look for and store compiled objects in <dir>. When JULIA_IMAGE_TIMINGS=1 is also set, the cache hit rate will be printed, like:

[...]
cache hits: 260/261 (99%)
added 2520 B to cache
[...]

Internals

Normally, jl_emit_native emits every CodeInstance into a separate LLVM module before combining everything into a single module. When the fine-grained cache is enabled, the modules are serialized to bitcode separately. The cache key for each module is computed from the hash of the serialized bitcode and the LLVM version, and the compiled object file is the value.

The partitionModule pass is not run, since multi-threaded compilation is done by having JULIA_IMAGE_THREADS worker threads take from the queue of serialized modules.

When the fine-grained cache is used, we generate a single jl_image_shard_t table for the entire image. The gvar_offsets are resolved by the linker.

Currently, the cache uses LLVM's FileCache, a thread safe key-value store that uses one file per key and write-and-rename to implement atomic updates. It's a convenient choice for development because the contents can be easily objdumped, but the long term plan is to switch to a more appropriate database, be it LLVMCAS when it is merged, or sqlite.

Current limitations

  • The multiversioning pass requires a combined module. If it is requested, the native code cache will be disabled.
  • Only .o outputs are cached. The fine-grained cache cannot be used if --output-bc, --output-unopt-bc or --output-asm are specified.
  • It is unclear how much of a (compile time) performance impact using seperate modules will have. When we understand how to maximize cache hits I'd like to use some heuristic to emit code into shared modules to mitigate this.
  • Cache entries do not currently expire.
  • The cache directory hits file system bottlenecks very fast on Windows. There is much wasted space.
  • The APIs intended for external use also require the use of a single module (jl_get_llvm_module_impl, jl_get_llvm_function_impl, jl_emit_native_impl with llvmmod set).
  • The cache hit rate is predictably quite low because of how we generate names during codegen. We'll need to change how this works and delay the uniquing to linking to improve this.

Plan

  • Support compiling split modules, emitting a single shard table (this PR)
  • Generate names that are predictable and only unique within one module; rename while linking.
  • Find a good heuristic to generate fewer modules
  • Switch to a better KV store
  • Cache eviction
  • Support multiversioning

@xal-0 xal-0 added compiler:codegen Generation of LLVM IR and native code compiler:precompilation Precompilation of modules feature Indicates new feature / enhancement requests labels May 30, 2025
@vchuravy
Copy link
Copy Markdown
Member

vchuravy commented Jun 1, 2025

This is fantastic! I have long been wanting to explore more fine-grained compilation caching using approaches like llvm-cas. From a cursory look this seems similar to how GPUCompilers on-disk cache is working (and maybe there is room for unifying both approaches)

My big picture questions would be:

  1. How is data handled?
  2. How is provenance handled? E.g. this relies on the uniqueness of CodeInstances? IIRC that uniqueness is only guaranteed during precompilation
  3. Cache collision/invalidation?

x-ref: https://github.com/JuliaGPU/GPUCompiler.jl/blob/c3ba85b62daeb572b2faa800013094b21f94a4f7/src/execution.jl#L167

Another concern is the many small files since they often perform very badly, so something along the line of llvm-cas or sqllite as a blob storage might be good.

@xal-0
Copy link
Copy Markdown
Member Author

xal-0 commented Jun 2, 2025

It's a little different from the GPUCompiler approach, since it makes no attempt
to predict if a CodeInstance will ultimately emit the same code. Instead, we
emit LLVM as normal and compute the ModuleHash, which we do anyway to move LLVM
modules between threads.

The idea with this PR is to do the minimum thing that will give a correct
result, then work to get more cache hits (starting with
globalUniqueGeneratedNames), so we can evaluate if this type of cache is the
right approach at all.

llvm-cas possibly being merged soon is the reason I have so far avoided pulling
in an embeddable key-value store, though sqlite might be the better choice if we
find other uses for having it around.

So far, there is no invalidation, pending a switch to a real KV store.

@IanButterworth
Copy link
Copy Markdown
Member

[WIP] Precompile native code cache
This is the first iteration of a cache for native code. ...

This confusingly sounds like something we already have. It might be good to make it clearer in the title and top comment how this differentiates from the current situation.

@gbaraldi gbaraldi changed the title [WIP] Precompile native code cache [WIP] Finer grained precompile native code cache Jun 3, 2025
@StefanKarpinski
Copy link
Copy Markdown
Member

2. How is provenance handled? E.g. this relies on the uniqueness of CodeInstances? IIRC that uniqueness is only guaranteed during precompilation

It sounds like that isn't an issue for this approach, but it could be possible to have a cache keyed by CodeInstances that is only used during precompilation, where uniqueness is guaranteed. Yes, precompilation is already generating a cached code, but we can have multiple layers of caching, and this could be used to speed up precompilation.

@xal-0 xal-0 changed the title [WIP] Finer grained precompile native code cache Finer grained precompile native code cache (part 1) Jun 9, 2025
@xal-0
Copy link
Copy Markdown
Member Author

xal-0 commented Jun 9, 2025

Reverted two commits that were an attempt at getting better cache hit rates with the simplest possible change, in favour of doing a link-time rename step.

@xal-0 xal-0 marked this pull request as ready for review June 9, 2025 21:53
src/codegen.cpp Outdated
Comment on lines +1579 to +1590
static StringMap<uint64_t> fresh_name_map;
static std::mutex fresh_name_lock;

static void freshen_name(std::string &name)
{
uint64_t n;
{
std::lock_guard guard{fresh_name_lock};
n = fresh_name_map[name]++;
}
raw_string_ostream(name) << "_" << n;
}
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am worried that this map could grow big. What issue are you trying to solve here? Just that the order of generation causes different names?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this was an attempt at getting reasonable hit rates with the minimum amount of work; it turned out not to be worth it. The new plan is to abandon globally unique names at this stage and resolve things after code generation, but before linking.

@xal-0 xal-0 mentioned this pull request Nov 4, 2025
xal-0 added a commit that referenced this pull request Feb 24, 2026
# Overview

This PR overhauls the way linking works in Julia, both in the JIT and
AOT. The point is to enable us to generate LLVM IR that depends only on
the source IR, eliminating both nondeterminism and statefulness. This
serves two purposes. First, if the IR is predictable, we can cache
compile objects using the bitcode hash as a key, like how the ThinLTO
cache works. #58592 was an early experiment along these lines. Second,
we can reuse work that was done in a previous session, like pkgimages,
but for the JIT.

We accomplish this by generating names that are unique only within the
current LLVM module, removing most uses of the
`globalUniqueGeneratedNames` counter. The replacement for
`jl_codegen_params_t`, `jl_codegen_output_t`, represents a Julia
"translation unit", and tracks the information we'll need to link the
compiled module into the running session. When linking, we manipulate
the JITLink [LinkGraph](https://llvm.org/docs/JITLink.html#linkgraph)
(after compilation) instead of renaming functions in the LLVM IR
(before).

## Example

```
julia> @noinline foo(x) = x + 2.0
       baz(x) = foo(foo(x))

       code_llvm(baz, (Int64,); dump_module=true, optimize=false)
```

Nightly:
```llvm
[...]
@"+Core.Float64#774" = private unnamed_addr constant ptr @"+Core.Float64#774.jit"
@"+Core.Float64#774.jit" = private alias ptr, inttoptr (i64 4797624416 to ptr)

; Function Signature: baz(Int64)
;  @ REPL[1]:2 within `baz`
define double @julia_baz_772(i64 signext %"x::Int64") #0 {
top:
  %pgcstack = call ptr @julia.get_pgcstack()
  %0 = call double @j_foo_775(i64 signext %"x::Int64")
  %1 = call double @j_foo_776(double %0)
  ret double %1
}

; Function Attrs: noinline optnone
define nonnull ptr @jfptr_baz_773(ptr %"function::Core.Function", ptr noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") #1 {
top:
  %pgcstack = call ptr @julia.get_pgcstack()
  %0 = getelementptr inbounds i8, ptr %"args::Any[]", i32 0
  %1 = load ptr, ptr %0, align 8
  %.unbox = load i64, ptr %1, align 8
  %2 = call double @julia_baz_772(i64 signext %.unbox)
  %"+Core.Float64#774" = load ptr, ptr @"+Core.Float64#774", align 8
  %Float64 = ptrtoint ptr %"+Core.Float64#774" to i64
  %3 = inttoptr i64 %Float64 to ptr
  %current_task = getelementptr inbounds i8, ptr %pgcstack, i32 -152
  %"box::Float64" = call noalias nonnull align 8 dereferenceable(8) ptr @julia.gc_alloc_obj(ptr %current_task, i64 8, ptr %3) #5
  store double %2, ptr %"box::Float64", align 8
  ret ptr %"box::Float64"
}
[...]
```

Diff after this PR. Notice how each symbol gets the lowest possible
integer suffix that will make it unique to the module, and how the two
specializations for `foo` get different names:
```diff
@@ -4,18 +4,18 @@
 target triple = "arm64-apple-darwin24.6.0"
 
-@"+Core.Float64#774" = external global ptr
+@"+Core.Float64#_0" = external global ptr
 
 ; Function Signature: baz(Int64)
 ;  @ REPL[1]:2 within `baz`
-define double @julia_baz_772(i64 signext %"x::Int64") #0 {
+define double @julia_baz_0(i64 signext %"x::Int64") #0 {
 top:
   %pgcstack = call ptr @julia.get_pgcstack()
-  %0 = call double @j_foo_775(i64 signext %"x::Int64")
-  %1 = call double @j_foo_776(double %0)
+  %0 = call double @j_foo_0(i64 signext %"x::Int64")
+  %1 = call double @j_foo_1(double %0)
   ret double %1
 }
 
 ; Function Attrs: noinline optnone
-define nonnull ptr @jfptr_baz_773(ptr %"function::Core.Function", ptr noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") #1 {
+define nonnull ptr @jfptr_baz_0(ptr %"function::Core.Function", ptr noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") #1 {
 top:
   %pgcstack = call ptr @julia.get_pgcstack()
@@ -23,7 +23,7 @@
   %1 = load ptr, ptr %0, align 8
   %.unbox = load i64, ptr %1, align 8
-  %2 = call double @julia_baz_772(i64 signext %.unbox)
-  %"+Core.Float64#774" = load ptr, ptr @"+Core.Float64#774", align 8
-  %Float64 = ptrtoint ptr %"+Core.Float64#774" to i64
+  %2 = call double @julia_baz_0(i64 signext %.unbox)
+  %"+Core.Float64#_0" = load ptr, ptr @"+Core.Float64#_0", align 8
+  %Float64 = ptrtoint ptr %"+Core.Float64#_0" to i64
   %3 = inttoptr i64 %Float64 to ptr
   %current_task = getelementptr inbounds i8, ptr %pgcstack, i32 -152
@@ -39,8 +39,8 @@
 
 ; Function Signature: foo(Int64)
-declare double @j_foo_775(i64 signext) #3
+declare double @j_foo_0(i64 signext) #3
 
 ; Function Signature: foo(Float64)
-declare double @j_foo_776(double) #4
+declare double @j_foo_1(double) #4
 
 attributes #0 = { "frame-pointer"="all" "julia.fsig"="baz(Int64)" "probe-stack"="inline-asm" }
```

## List of changes
- Many sources of statefulness and nondeterminism in the emitted LLVM IR
have been eliminated, namely:
  - Function symbols defined for CodeInstances
  - Global symbols referring to data on the Julia heap
- Undefined function symbols referring to invoked external CodeInstances

- `jl_codeinst_params_t` has become `jl_codegen_output_t`. It now
represents one Julia "translation unit". More than one CodeInstance can
be emitted to the same `jl_codegen_output_t`, if desired, though in the
JIT every CI gets its own right now. One motivation behind this is to
allow us to emit code on multiple threads and avoid the bitcode
serialize/deserialize step we currently do, if that proves worthwhile.
  
When we are done emitting to a `jl_codegen_output_t`, we call
`.finish()`, which discards the intermediate state and returns only the
LLVM module and the info needed for linking (`jl_linker_info_t`).

- The new `JLMaterializationUnit` wraps emitting Julia LLVM modules and
the associated `jl_linker_info_t`. It informs ORC that we can
materialize symbols for the CIs defined by that output, and picks
globally unique names for them. When it is materialized, it resolves all
the call targets and generates trampolines for CodeInstances that are
invoked but have the wrong calling convention, or are not yet compiled.

- We now postpone linking decisions to after codegen whenever possible.
For example, `emit_invoke` no longer tries to find a compiled version of
the CodeInstance, and it no longer generates trampolines to adapt
calling conventions. `jl_analyze_workqueue`'s job has been absorbed into
`JuliaOJIT::linkOutput`.

- Some `image_codegen` differences have been removed:
- Codegen no longer cares if a compiled CodeInstance came from an image.
During ahead-of-time linking, we generate thunk functions that load the
address from the fvars table.

- In `jl_emit_native_impl`, emit every CodeInstance into one
`jl_codegen_output_t`. We now defer the creation of the `llvm::Linker`
for llvmcalls, which has construction cost that grows with the size of
the destination module, until the very end.

- RTDyld is removed completely, since we cannot control linking like we
can with JITLink. Since #60105, platforms that previous used the
optimized memory manager now use the new one.

### General refactoring
- Adapt the `jl_callingconv_t` enum from `staticdata.c` into
`jl_invoke_api_t` and use it in more places. There is one enumerator for
each special `jl_callptr_t` function that can go in a CodeInstance's
`invoke` field, as well as one that indicates an invoke wrapper should
be there. There is a convenience function for reading an invoke pointer
and getting the API type, and vice versa.
- Avoid using magic string values, and try to directly pass pointers to
LLVM `Function *` or ORC string pool entries when possible.

## Future work
- `DLSymOptimizer` should be mostly removed, in favour of emitting raw
ccalls and redirecting them to the appropriate target during linking.

- We should support ahead-of-time linking multiple
`jl_codegen_output_t`s together, in order to parallelize LLVM IR
emission when compiling a system image.
      
- We still pass strings to `emit_call_specfun_other`, even though the
prototype for the function is now created by
`jl_codegen_output_t::get_call_target`. We should hold on to the calling
convention info so it doesn't have to be recomputed.
mlechu pushed a commit to mlechu/julia that referenced this pull request Feb 24, 2026
# Overview

This PR overhauls the way linking works in Julia, both in the JIT and
AOT. The point is to enable us to generate LLVM IR that depends only on
the source IR, eliminating both nondeterminism and statefulness. This
serves two purposes. First, if the IR is predictable, we can cache
compile objects using the bitcode hash as a key, like how the ThinLTO
cache works. JuliaLang#58592 was an early experiment along these lines. Second,
we can reuse work that was done in a previous session, like pkgimages,
but for the JIT.

We accomplish this by generating names that are unique only within the
current LLVM module, removing most uses of the
`globalUniqueGeneratedNames` counter. The replacement for
`jl_codegen_params_t`, `jl_codegen_output_t`, represents a Julia
"translation unit", and tracks the information we'll need to link the
compiled module into the running session. When linking, we manipulate
the JITLink [LinkGraph](https://llvm.org/docs/JITLink.html#linkgraph)
(after compilation) instead of renaming functions in the LLVM IR
(before).

## Example

```
julia> @noinline foo(x) = x + 2.0
       baz(x) = foo(foo(x))

       code_llvm(baz, (Int64,); dump_module=true, optimize=false)
```

Nightly:
```llvm
[...]
@"+Core.Float64#774" = private unnamed_addr constant ptr @"+Core.Float64#774.jit"
@"+Core.Float64#774.jit" = private alias ptr, inttoptr (i64 4797624416 to ptr)

; Function Signature: baz(Int64)
;  @ REPL[1]:2 within `baz`
define double @julia_baz_772(i64 signext %"x::Int64") #0 {
top:
  %pgcstack = call ptr @julia.get_pgcstack()
  %0 = call double @j_foo_775(i64 signext %"x::Int64")
  %1 = call double @j_foo_776(double %0)
  ret double %1
}

; Function Attrs: noinline optnone
define nonnull ptr @jfptr_baz_773(ptr %"function::Core.Function", ptr noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") JuliaLang#1 {
top:
  %pgcstack = call ptr @julia.get_pgcstack()
  %0 = getelementptr inbounds i8, ptr %"args::Any[]", i32 0
  %1 = load ptr, ptr %0, align 8
  %.unbox = load i64, ptr %1, align 8
  %2 = call double @julia_baz_772(i64 signext %.unbox)
  %"+Core.Float64#774" = load ptr, ptr @"+Core.Float64#774", align 8
  %Float64 = ptrtoint ptr %"+Core.Float64#774" to i64
  %3 = inttoptr i64 %Float64 to ptr
  %current_task = getelementptr inbounds i8, ptr %pgcstack, i32 -152
  %"box::Float64" = call noalias nonnull align 8 dereferenceable(8) ptr @julia.gc_alloc_obj(ptr %current_task, i64 8, ptr %3) JuliaLang#5
  store double %2, ptr %"box::Float64", align 8
  ret ptr %"box::Float64"
}
[...]
```

Diff after this PR. Notice how each symbol gets the lowest possible
integer suffix that will make it unique to the module, and how the two
specializations for `foo` get different names:
```diff
@@ -4,18 +4,18 @@
 target triple = "arm64-apple-darwin24.6.0"
 
-@"+Core.Float64#774" = external global ptr
+@"+Core.Float64#_0" = external global ptr
 
 ; Function Signature: baz(Int64)
 ;  @ REPL[1]:2 within `baz`
-define double @julia_baz_772(i64 signext %"x::Int64") #0 {
+define double @julia_baz_0(i64 signext %"x::Int64") #0 {
 top:
   %pgcstack = call ptr @julia.get_pgcstack()
-  %0 = call double @j_foo_775(i64 signext %"x::Int64")
-  %1 = call double @j_foo_776(double %0)
+  %0 = call double @j_foo_0(i64 signext %"x::Int64")
+  %1 = call double @j_foo_1(double %0)
   ret double %1
 }
 
 ; Function Attrs: noinline optnone
-define nonnull ptr @jfptr_baz_773(ptr %"function::Core.Function", ptr noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") JuliaLang#1 {
+define nonnull ptr @jfptr_baz_0(ptr %"function::Core.Function", ptr noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") JuliaLang#1 {
 top:
   %pgcstack = call ptr @julia.get_pgcstack()
@@ -23,7 +23,7 @@
   %1 = load ptr, ptr %0, align 8
   %.unbox = load i64, ptr %1, align 8
-  %2 = call double @julia_baz_772(i64 signext %.unbox)
-  %"+Core.Float64#774" = load ptr, ptr @"+Core.Float64#774", align 8
-  %Float64 = ptrtoint ptr %"+Core.Float64#774" to i64
+  %2 = call double @julia_baz_0(i64 signext %.unbox)
+  %"+Core.Float64#_0" = load ptr, ptr @"+Core.Float64#_0", align 8
+  %Float64 = ptrtoint ptr %"+Core.Float64#_0" to i64
   %3 = inttoptr i64 %Float64 to ptr
   %current_task = getelementptr inbounds i8, ptr %pgcstack, i32 -152
@@ -39,8 +39,8 @@
 
 ; Function Signature: foo(Int64)
-declare double @j_foo_775(i64 signext) JuliaLang#3
+declare double @j_foo_0(i64 signext) JuliaLang#3
 
 ; Function Signature: foo(Float64)
-declare double @j_foo_776(double) JuliaLang#4
+declare double @j_foo_1(double) JuliaLang#4
 
 attributes #0 = { "frame-pointer"="all" "julia.fsig"="baz(Int64)" "probe-stack"="inline-asm" }
```

## List of changes
- Many sources of statefulness and nondeterminism in the emitted LLVM IR
have been eliminated, namely:
  - Function symbols defined for CodeInstances
  - Global symbols referring to data on the Julia heap
- Undefined function symbols referring to invoked external CodeInstances

- `jl_codeinst_params_t` has become `jl_codegen_output_t`. It now
represents one Julia "translation unit". More than one CodeInstance can
be emitted to the same `jl_codegen_output_t`, if desired, though in the
JIT every CI gets its own right now. One motivation behind this is to
allow us to emit code on multiple threads and avoid the bitcode
serialize/deserialize step we currently do, if that proves worthwhile.
  
When we are done emitting to a `jl_codegen_output_t`, we call
`.finish()`, which discards the intermediate state and returns only the
LLVM module and the info needed for linking (`jl_linker_info_t`).

- The new `JLMaterializationUnit` wraps emitting Julia LLVM modules and
the associated `jl_linker_info_t`. It informs ORC that we can
materialize symbols for the CIs defined by that output, and picks
globally unique names for them. When it is materialized, it resolves all
the call targets and generates trampolines for CodeInstances that are
invoked but have the wrong calling convention, or are not yet compiled.

- We now postpone linking decisions to after codegen whenever possible.
For example, `emit_invoke` no longer tries to find a compiled version of
the CodeInstance, and it no longer generates trampolines to adapt
calling conventions. `jl_analyze_workqueue`'s job has been absorbed into
`JuliaOJIT::linkOutput`.

- Some `image_codegen` differences have been removed:
- Codegen no longer cares if a compiled CodeInstance came from an image.
During ahead-of-time linking, we generate thunk functions that load the
address from the fvars table.

- In `jl_emit_native_impl`, emit every CodeInstance into one
`jl_codegen_output_t`. We now defer the creation of the `llvm::Linker`
for llvmcalls, which has construction cost that grows with the size of
the destination module, until the very end.

- RTDyld is removed completely, since we cannot control linking like we
can with JITLink. Since JuliaLang#60105, platforms that previous used the
optimized memory manager now use the new one.

### General refactoring
- Adapt the `jl_callingconv_t` enum from `staticdata.c` into
`jl_invoke_api_t` and use it in more places. There is one enumerator for
each special `jl_callptr_t` function that can go in a CodeInstance's
`invoke` field, as well as one that indicates an invoke wrapper should
be there. There is a convenience function for reading an invoke pointer
and getting the API type, and vice versa.
- Avoid using magic string values, and try to directly pass pointers to
LLVM `Function *` or ORC string pool entries when possible.

## Future work
- `DLSymOptimizer` should be mostly removed, in favour of emitting raw
ccalls and redirecting them to the appropriate target during linking.

- We should support ahead-of-time linking multiple
`jl_codegen_output_t`s together, in order to parallelize LLVM IR
emission when compiling a system image.
      
- We still pass strings to `emit_call_specfun_other`, even though the
prototype for the function is now created by
`jl_codegen_output_t::get_call_target`. We should hold on to the calling
convention info so it doesn't have to be recomputed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

compiler:codegen Generation of LLVM IR and native code compiler:precompilation Precompilation of modules feature Indicates new feature / enhancement requests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants