[vchuravy]

Requested by Valentin

Most of the code a differentiated function calls is never differentiated: callees with an EnzymeRules.inactive rule, or whose argument types and return type are all guaranteed const. Enzyme still emits, optimizes, analyzes and code generates it in every thunk, although Julia compiles the same code natively for the primal. In the Enzyme tests of TensorKit.jl (QuantumKitHub/TensorKit.jl#464) that is ~80% of the IR of each thunk: _degeneracystructure, _sectorstructure, LRU cache lookups, fusion-tree iteration, …. Each new sector, element type or tensor shape compiles all of it again.

This PR stops inferring and emitting those callees and calls their native code instead. It's always on for host code on Julia 1.12 and later (Enzyme.Compiler.NATIVE_CALLEES, a constant that is part of the interpreter's cache token). Julia 1.10 and 1.11 keep the old path.

How

• Switch to the native interpreter during inference. typeinf_edge is overridden for EnzymeInterpreter. It resolves a call of a derivative-free method to the CodeInstance Julia's native interpreter infers for it, via return_cached_result. The caller gets the native return type and effects, and neither the method nor its callees are inferred by EnzymeInterpreter. This only applies when the method has no rule other than inactive, isn't a primitive or always-inline function, and Julia would call it through a specialized entry point rather than inline it. The method is inferred natively first and only compiled when it qualifies. Since nothing for that MethodInstance is in the EnzymeInterpreter cache, the optimizer emits :invoke mi.
• Host only. Both steps apply only to interpreters whose method table is GPUCompiler's global one. GPU back-ends stack their own overlay tables.
• Leave them out of the emitted list. On 1.12 and 1.13, GPUCompiler emits exactly the CodeInstances that ci_cache_populate returns (jl_emit_native). For an :invoke of a MethodInstance or CodeInstance that is not in the list, codegen emits only a small specsig stub that boxes the arguments and jl_invokes the callee; nothing of its body or its callees. ci_cache_populate is overridden for EnzymeInterpreter to also leave out derivative-free callees that EnzymeInterpreter did infer, for example ones whose results were already in its cache (native_callee).
• Bind the stubs. bind_native_callees! runs right after emit_llvm (primal module, nested_codegen!, sugar.jl). It replaces each stub by a declare_native! declaration bound to the native specsig entry point, as invoke_codegen! does for custom rules. It runs before the Julia values become symbolic (Keep Julia values symbolic until the module is linked (GPUCompiler 1.8) #3742), because it reads the invoked MethodInstance from the stub's constant global.
• Call sites pass pgcstack (swiftself).
• Declarations are marked nofree, enzyme_inactive and enzymejl_native_inactive. With the last one, materialize_native_invokes! leaves them without a body; nested differentiation doesn't differentiate them either.
• GC-rooted struct returns get an sret type with the tracked pointers stripped (the roots array keeps the objects alive). Enzyme may move the result buffer to the heap to keep it for the reverse pass, and LateLowerGC requires a stack buffer for an sret with tracked pointers.

Gotchas found on the way

• The stubs pass the MethodInstance to jl_invoke, not the CodeInstance. Each populate therefore publishes the table of what it left out, and the bind that follows its emission consumes it. That avoids matching a stub against a callee left out by another compilation.
• Not every callee has a specialized native entry. Julia's codegen gives a function a specialized signature only when that's worthwhile (uses_specsig with prefer_specsig off): an unboxed argument or return, a union return, or all-singleton arguments. With only boxed arguments and a boxed return it uses the jl_fptr_args ABI, although jl_read_codeinst_invoke still reports a specialized entry. Calling that through the derived specsig crashed (#_sortperm#29(…, ::Vector{Int})). native_uses_specsig mirrors the rule, and varargs methods are left out as well. Custom rules never hit this, because they always take an unboxed annotation argument.

Measurements

After rebasing onto main (3de895b): Julia 1.13.1, with 1.10 for reference

TensorKit.jl#464, single test files run standalone with CI=true (the CI configuration), main (9e86a1c) vs this PR, all runs concurrently on one machine with the profiler on. All tests pass in every configuration.

test	1.10, main	1.13, main	1.13, this PR
enzyme/qr	487 s	499 s	262 s (−47%)
enzyme/mul (25260 tests)	1297 s	1356 s	720 s (−47%)

Where the time goes in enzyme/mul (sampled CPU time):

phase	1.10, main	1.13, main	1.13, this PR
Enzyme LLVM optimization before AD	211 s	266 s	57 s
Enzyme LLVM optimization after AD	109 s	93 s	20 s
machine code for the thunk (ORC)	65 s	79 s	18 s
GPUCompiler check_ir	41 s	92 s	24 s
other Enzyme/GPUCompiler Julia-side IR passes	112 s	129 s	51 s
EnzymeInterpreter inference	147 s	29 s	5 s
Enzyme AD itself (C++)	34 s	50 s	37 s
native Julia JIT (test harness, FiniteDifferences, TensorKit primal per type)	185 s	255 s	256 s

The PR removes nearly all the work on code that is never differentiated; the AD step itself costs the same. After it, ordinary native compilation of the test harness and primal code is the largest cost (~40% of mul), which this PR doesn't touch. Julia 1.10 spends about as long as 1.13 without the PR, so it would benefit similarly if the path worked there.

The measurements below are from earlier revisions, when an ENZYME_NATIVE_CALLEES switch still existed.

Latest (all three commits, Julia 1.12.7), three configurations at the same time

	total	first call (warmup)	other 215 calls
off (ENZYME_NATIVE_CALLEES=0)	1556 s	75 s	1481 s
populate step only (864c32b)	829 s	108 s	721 s
+ native inference (default)	802 s (−48%)	73 s	729 s

All 216 calls pass in each configuration. Switching to the native interpreter mainly removes the extra warmup cost of the populate-only version. svd_trunc on Z2 takes 13.2 / 9.1 / 8.0 s.

Earlier (first commit, opt-in)

TensorKit.jl#464 enzyme-factorizations CI matrix, 216 test_reverse calls (each compiles one new thunk), Julia 1.12.7, flag off vs on, run at the same time on one machine:

	total
off (main)	1736 s
ENZYME_NATIVE_CALLEES=1	798 s (−54%)

• All 216 calls pass, and none is slower.
• By function: qr_compact −66%, project_(anti)hermitian! −67%, lq_compact −60%, eig_trunc −47%, proj_eigh_trunc −47%, svd_trunc −40%.
• By sector: all non-trivial sectors are −58 to −59%, Trivial −11%.
• For svd_trunc on the Z2 sector the emitted module is 16k instructions instead of 72k (3.3k after post-optimization), and the call takes 11.3 s → 6.7 s.

These numbers were measured before rebasing onto current main; after the rebase, svd_trunc Z2 takes 6.5 s.

Julia 1.13.1

72 calls (Trivial and Z2 sectors, including the ~60 s one-time warmup), flag off vs on at the same time: 340 s → 208 s (−39%), all pass. The only change needed for 1.13 is codeinfo_for_const taking the world range and edges of the CodeInstance; GPUCompiler's ci_cache_populate and the stubs codegen emits are the same as on 1.12.

Tests

New test test/native_callees.jl: a @noinline inactive callee returning a struct with a Vector field (rooted sret), and the boxed-ABI sortperm case, in forward and reverse mode.

On the final commit, with the feature on (now the default), these groups pass locally on 1.12: native_callees rules core advanced basic sugar errors callconv mixed typeunstable. That's 15201 pass and 49 broken (the same as with it off), 0 errors.

I didn't run the GPU, threads or integration suites with the flag on. The path on Julia 1.10 and 1.11 is unchanged.

Open questions / follow-ups

• Globals: a function whose argument types are all inactive can still write a global that active code later reads as active data. That derivative would be lost; this is the main correctness risk of the change. In TensorKit the only such writes are the structure LRU caches, which hold integer data.
• Upstream shape: a hook in GPUCompiler that filters which CodeInstances get emitted would be cleaner than overriding ci_cache_populate.

🤖 Generated with Claude Code

[github-actions]

Benchmark Results

	main	148fca8...	main / 148fca8...
basics/make_zero/namedtuple	0.0518 ± 0.002 μs	0.055 ± 0.0034 μs	0.941 ± 0.068
basics/make_zero/struct	0.262 ± 0.0059 μs	0.267 ± 0.0074 μs	0.982 ± 0.035
basics/overhead	4.34 ± 0.01 ns	4.34 ± 0.01 ns	1 ± 0.0033
basics/remake_zero!/namedtuple	0.223 ± 0.0051 μs	0.223 ± 0.0069 μs	0.997 ± 0.039
basics/remake_zero!/struct	0.227 ± 0.009 μs	0.229 ± 0.012 μs	0.99 ± 0.064
fold_broadcast/multidim_sum_bcast/1D	0.496 ± 0.0054 μs	0.501 ± 0.0074 μs	0.99 ± 0.018
fold_broadcast/multidim_sum_bcast/2D	0.365 ± 0.0042 μs	0.37 ± 0.0052 μs	0.988 ± 0.018
inline_abi/call	15.8 ± 0.05 ns	17.6 ± 0.05 ns	0.894 ± 0.0038
inline_abi/call_alloc	0.0434 ± 0.001 μs	0.0453 ± 0.0012 μs	0.958 ± 0.034
inline_abi/loop	4.01 ± 0.0043 μs	4.01 ± 0.0043 μs	1 ± 0.0015
time_to_load	1.54 ± 0.014 s	1.67 ± 0.04 s	0.922 ± 0.024

Benchmark Plots

A plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/39576052072/artifacts/12341809715.

[vchuravy]

Two red checks on 864c32b that aren't from this PR:

• Integration Tests - Flux - Julia 1.11: the GRU/LSTM EnzymeInternalError (broadcast_unalias) that already fails on main (also red on Do not fold loads of non-const globals on Julia 1.10 #3709 and Only write printed type names into the IR when asked to #3711). This PR only changes code on Julia 1.12/1.13, so nothing in it runs on 1.11. No fix exists for it yet.
• Integration Tests - Comrade - Julia 1.12: the runner pod was rejected for lack of memory before any test ran ("Node didn't have enough resource: memory"). I'll re-run it once this workflow run finishes.

---

Generated by Claude Code

[vchuravy]

Julia 1.13 - windows-latest fails in test/optimize.jl:306, "Loop-invariant load hoisted past a read-only-or-throw call": bytes(2000) - bytes(10) is 63998, limit 47760. This comes from this PR. Linux 1.12 and 1.13 should fail the same way once they run.

The cause is that lookup_or_throw(::Vector{Int}, ::Int) is derivative-free, so bind_native_callees! now replaces it with a body-less native declaration. Enzyme used to infer enzyme_ReadOnlyOrThrow from its body (it writes only on the bounds-error path). The declaration has no body, so it only gets nofree. LICM then can't hoist the Memory load past the call, and the tape caches a pointer on every iteration. I haven't tested the following fix yet: in bind_native_callees!, take the attribute from the native CodeInstance's effects, roughly

effects = Core.Compiler.decode_effects(ci.ipo_purity_bits)
if Core.Compiler.is_effect_free(effects)
    push!(fattrs, StringAttribute("enzyme_ReadOnlyOrThrow"))
end

This keeps the information the body used to give, for every native callee and not just this test.

---

Generated by Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 93.98148% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.93%. Comparing base (3de895b) to head (504c17c).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
src/compiler/native_callees.jl	93.89%	21 Missing ⚠️
src/errors.jl	0.00%	2 Missing ⚠️
src/rules/customrules.jl	93.54%	2 Missing ⚠️
src/compiler/interpreter.jl	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #3768       +/-   ##
===========================================
+ Coverage   49.20%   79.93%   +30.72%     
===========================================
  Files          58       72       +14     
  Lines       23504    25335     +1831     
===========================================
+ Hits        11566    20252     +8686     
+ Misses      11938     5083     -6855     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vchuravy]

Stack on #3771

[vchuravy]

We should also enable this for 1.10 and 1.11.

[vchuravy]

Rebase the PR on #3771 and enable this for 1.10/1.11 as well.

[vchuravy]

#3771 had been rewritten (cea484c, 45f6fc0, bbd149d), so I brought it into this branch with a merge commit (65bbfbb). This session isn't allowed to rewrite history, so a real rebase and force-push wasn't possible here. If you want linear history, a local git rebase --onto origin/claude/project-thread-my4sax of the 9 non-merge commits on this branch will do it.

1.10 and 1.11 were already enabled in 2bfda4d. Two adjustments came with the merge:

• Native callees are now bound after record_julia_values!. When GPUCompiler 2.x leaves a stub's slot as a declaration, stub_callee reads it from the context.
• GPUCompiler 2.x has no ci_cache_populate; it uses CompilerCaching.get_codeinfos. So on 1.12 and 1.13 the native path is now limited to GPUCompiler 1.x, which is what CI resolves (v1.23.0). Hooking the 2.x path is a follow-up.

---

Generated by Claude Code

[vchuravy]

@claude Hook the 2.x path now

[vchuravy]

The GPUCompiler 2.x path is hooked in ebe6733. CompilerCaching.get_codeinfos is now overridden for EnzymeInterpreter, the same way ci_cache_populate is on 1.x. It walks the :invoke edges, skips derivative-free callees and records them in the EnzymeContext, and otherwise falls back to the default. Since #3771 the default CI matrix resolves GPUCompiler 2.x, and the gpucompiler: '1' jobs keep covering 1.x.

---

Generated by Claude Code

[vchuravy]

CI status on 7bc78da: every Linux and Windows job for 1.10 through 1.13 passes, including the GPUCompiler 2.x default jobs and the GPUCompiler=1 jobs. All integration suites pass except Flux on 1.11. The macOS jobs are still queued.

• Flux 1.11: only the LSTM, GRU and Chain(LSTM, LSTM) testsets error. This is the same failure as on main and is already noted above, so it isn't caused by this PR.
• codecov/project (-0.77%): patch coverage is 94% and codecov/patch passes. The project-level drop comes from lines outside the diff. I expect two causes. First, derivative-free callees now bind to native code, so the compiler paths that used to process their bodies no longer run in the test suite. Second, the macOS coverage uploads haven't arrived yet. I'm not adding tests just to raise that number.

---

Generated by Claude Code

[vchuravy]

@claude rebase

[github-actions]

Your PR no longer requires formatting changes. Thank you for your contribution!

[vchuravy]

CI on 08d10d6: everything that has finished is green except two checks.

• codecov/project (80.47%, -0.09% vs main): most of the 17 uncovered patch lines in native_callees.jl are the 1.10/1.11 path, which stays off until the Enzyme C++ fix (no sret on functions that never return in FixupJuliaCallingConvention) ships in a new Enzyme_jll. Turning 1.10/1.11 on will cover them. With both fixes applied locally, TensorKit's full enzyme-factorizations matrix passes on 1.10 (-55% time) and 1.11 (-64%).
• Integration Tests - Flux - Julia 1.11: the GRU/LSTM failure that also fails on main, as noted earlier.

Integration Tests - Lux - Julia 1.11 was cancelled after hanging in "Load Julia packages from cache" before any test ran; I re-ran it.

---

Generated by Claude Code

[vchuravy]

Every Julia 1.10 unit job is red on 3656e7f. The failing test is the new nc_branch one in test/native_callees.jl, which fails LLVM verification. Julia 1.11 and later pass.

The error is "Call parameter type does not match function signature!". In preprocess_julia_nc_shared_*, a call to ijl_bounds_error_unboxed_int gets a [5 x i64] addrspace(11)* where the callee takes an i8 addrspace(11)*.

The cause is in Enzyme. When an alloca is moved to the heap, RecursivelyReplaceAddressSpace replaces a typed-pointer addrspacecast that also changes the pointee type, and drops that change. Marking the jl_invoke call read-only-or-throw does not stop the clone in CI, including with the libEnzyme built from Enzyme main.

The fix is EnzymeAD/Enzyme#3491. Its lit test reproduces this exact error on LLVM 15 without the fix and passes with it. This PR stays red on 1.10 until an Enzyme_jll release carries that fix, unless the test is gated in the meantime.

---

Generated by Claude Code

[vchuravy]

On 9d97ea0 the Documentation build itself passes: cross-references, doctests and rendering all succeed. Only the preview deploy fails. deploydocs gets git@github.com: Permission denied (publickey) when it fetches with DOCUMENTER_KEY, which means the repository's deploy key or secret is rejected. That is a CI credential problem, not something in this PR's code.

Downgrade 1.10 is still the typed-pointer failure from my earlier comment, which needs EnzymeAD/Enzyme#3491 in a released Enzyme_jll.

---

Generated by Claude Code