Drop functions nothing calls after AD before optimizing

[vchuravy]

Stacked on #3736 → #3732 → #3734 (Enzyme_jll 0.0.300).

Originally this PR addressed #3570 by turning functions with a custom rule into declarations during AD. EnzymeAD/Enzyme#3352 and #3353 (in Enzyme_jll 0.0.300) fix that in Enzyme itself: nothing infers readnone or writeonly on the parameters of a function with a custom rule any more. The regression tests from this PR are now in #3736, which passes without any change here. What is left is the compile-time part of the original idea, as suggested in #3610 (comment).

Change

Enzyme replaces the calls it differentiates to a function with a custom rule with the rule, so after enzyme! the function and everything it calls are often unused. They kept external linkage until the end of compile_unhooked_impl, so nothing could delete them before then. Right after AD, this makes every function except the entry points private and runs GlobalDCEPass. Calls Enzyme kept, e.g. from a constant callee, still reach the body.

The function bodies are not dropped before AD. Enzyme keeps calling the primal for calls where everything is Const: rules imported from ChainRules assert a non-Const return, and the lock rules from #3513 run only for active locks. So a body is still needed whenever such a call survives.

Effect

A function with a custom rule whose body calls 400 distinct @noinline functions (the shape from #3455), Julia 1.13:

path	defined functions after AD	reaching post_optimize! before → after
CPU (autodiff, forward and reverse)	408 / 409	1 → 1
CUDA (autodiff_deferred in a kernel, 100 callees)	104	104 → 1

• CPU: GPUCompiler's cleanup after our hook already removed these before _thunk runs post_optimize!. The change only makes that happen earlier.
• Non-native targets: post_optimize! runs inside compile_unhooked_impl, before that cleanup, and optimized all of them. That post_optimize! went from 0.029 s to 0.001 s; with trivial bodies the total compile time (~20 s) does not change measurably.
• Not addressed: the bulk of the cost described in Compile-time overhead of custom rules: dead recursive codegen of rule-bearing functions + duplicated nested codegen of rules #3455 is inference and codegen of these functions before AD, which this does not change.

Testing

• Full test suite, Julia 1.12.7: passes.
• Full test suite, Julia 1.13.1: one error, advanced.jl "Copy Broadcast arg" (IllegalTypeAnalysisException). It fails the same way without this change in the same environment, and does not fail in Update Enzyme_jll to 0.0.300 #3734's CI.
• test/cuda.jl, Julia 1.13.1 (Quadro RTX 4000): 61/61 pass.

🤖 Generated with Claude Code

[github-actions]

Benchmark Results

	main	042b6bd...	main / 042b6bd...
basics/make_zero/namedtuple	0.0516 ± 0.0023 μs	0.0521 ± 0.0021 μs	0.991 ± 0.06
basics/make_zero/struct	0.274 ± 0.0067 μs	0.276 ± 0.0069 μs	0.994 ± 0.035
basics/overhead	4.34 ± 0.01 ns	4.34 ± 0.01 ns	1 ± 0.0033
basics/remake_zero!/namedtuple	0.222 ± 0.0051 μs	0.223 ± 0.0044 μs	0.995 ± 0.03
basics/remake_zero!/struct	0.223 ± 0.0079 μs	0.223 ± 0.007 μs	0.998 ± 0.047
fold_broadcast/multidim_sum_bcast/1D	0.502 ± 0.0056 μs	0.5 ± 0.006 μs	1 ± 0.016
fold_broadcast/multidim_sum_bcast/2D	0.37 ± 0.0044 μs	0.372 ± 0.0039 μs	0.994 ± 0.016
inline_abi/call	16.4 ± 0.05 ns	17.3 ± 0.05 ns	0.947 ± 0.004
inline_abi/call_alloc	0.0437 ± 0.00093 μs	0.0438 ± 0.0013 μs	1 ± 0.036
inline_abi/loop	4.01 ± 0.0041 μs	4.01 ± 0.0041 μs	1 ± 0.0015
time_to_load	1.49 ± 0.022 s	1.51 ± 0.0086 s	0.986 ± 0.015

Benchmark Plots

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

[kshyatt]

I hit a version of this from a slightly different angle, and it seems one more change would help here:

Trigger: On Julia 1.13.1, a struct with inline roots is passed as a pointer to its data half plus its roots. If a function with a custom rule only uses the pointer fields in its primal body, EnzymeDetectReadonlyOrThrow marks the data-half argument readnone. DSE then removes the caller's store of the bits fields into that buffer, before AD. The rule then reads uninitialized memory for every non-pointer field. The derivative values themselves can still come out right (here they only use the array), but anything the rule reads from the bits fields is garbage. Minimal version:

using Enzyme, LinearAlgebra
import Enzyme.EnzymeRules

struct TM; data::Vector{Float64}; n::Int; end
f1(A::TM) = norm(A.data)
function EnzymeRules.forward(config::EnzymeRules.FwdConfigWidth{1}, ::Const{typeof(f1)}, ::Type{RT}, A::Annotation{TM}) where {RT}
    @show A.val.n A.dval.n          # 7 expected
    n = f1(A.val)
    return Duplicated(n, dot(A.val.data, A.dval.data) / n)
end
g1(A) = f1(A)
autodiff(Forward, g1, Duplicated(TM(randn(3), 7), TM(ones(3), 7)))

On 1.13.1 this prints garbage for both ns. It's fine on 1.12.7, where the store survives even though the parameter is also inferred readnone. This PR doesn't quite fix it because middle_optimize! gained an API.EnzymeDetectReadonlyOrThrow(mod) call.

Routing that call through the wrapper fixes it:

     function middle_optimize!(second_stage = false)
         ...
-        API.EnzymeDetectReadonlyOrThrow(mod)
+        detect_readonly_or_throw!(mod)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 76.53%. Comparing base (d7a2f8f) to head (042b6bd).
⚠️ Report is 2 commits behind head on vc/custom-rule-regression-tests.

Additional details and impacted files

@@                         Coverage Diff                         @@
##           vc/custom-rule-regression-tests    #3610      +/-   ##
===================================================================
- Coverage                            80.03%   76.53%   -3.51%     
===================================================================
  Files                                   71       71              
  Lines                                24172    24180       +8     
===================================================================
- Hits                                 19347    18505     -842     
- Misses                                4825     5675     +850     

☔ 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]

The better fix is not not codegen these functions to begin with / or if that is not possible to drop the bodies of the function since we are going to replace them with custom rules later.

HT Billy

[wsmoses]

I think it may still be worthwhile (and possibly necessary for sret bits) here to run the attribute detection first [and keep those attributes, minus changing readnone -> readonly, and (non sret/returnroots) writeonly -> nothing