Skip to content

Support juliac --trim=safe: SpecialFunctions extension, opt-in static default chunk - #864

Open
Beforerr wants to merge 2 commits into
JuliaDiff:masterfrom
Beforerr:trim-safe
Open

Beforerr wants to merge 2 commits into
JuliaDiff:masterfrom
Beforerr:trim-safe

Conversation

@Beforerr

@Beforerr Beforerr commented Oct 23, 2026 •

Copy link
Copy Markdown

This PR makes ForwardDiff usable under juliac --trim=safe and removes a library load at startup for programs that never use special functions.

Problem

  1. using ForwardDiff always loads SpecialFunctions and OpenSpecFun_jll. A trimmed executable keeps OpenSpecFun_jll.__init__, so every launch opens libopenspecfun, even when no special function is ever called. In a minimal trimmed program, that roughly doubles startup time.
  2. gradient(f, x), jacobian(f, x), jacobian(f!, y, x) and hessian(f, x) infer as Any for x::Vector{Float64}, so --trim=safe rejects them, including when another package makes the call. Only calls with an explicit Chunk{N} config compile.

Root cause

For (2): Chunk(x) returns a Union of 13 Chunk{N} types. That is too wide to union-split, so building the default config is a dynamic dispatch.

Fix

  1. SpecialFunctions becomes a weak dependency. ForwardDiffSpecialFunctionsExt defines the Dual methods for SpecialFunctions, plus logabsgamma and gamma_inc, and rules of other modules whose derivative calls SpecialFunctions (today only NaNMath.lgamma). User-visible change: NaNMath.lgamma on a Dual now needs SpecialFunctions loaded.
  2. Opt-in static dispatch on the default chunk size. With the new static_default_chunk preference, methods without a config branch on the chunk size before building it, one branch per size pickchunksize can return, so they infer concretely. It is off by default because the first call then compiles every chunk size (first hessian(f, rand(9)): 0.5 s → 3.5 s); with it off, behavior is unchanged. Documented in advanced.md.

Verification

  • New CI job builds and runs a --trim=safe executable from test/trim with the preference set, comparing results for n = 3 and 20 with the JIT. Passes locally on Julia 1.14.0-alpha1; without the preference the build fails with unresolved calls to hessian/HessianConfig.
  • With the preference off, first-call latency matches master.

`import SpecialFunctions` made every `using ForwardDiff` load SpecialFunctions
and OpenSpecFun_jll, whose `__init__` opens libopenspecfun. In a
`juliac --trim=safe` executable that initializer is kept, so every launch pays
for a library the program may never call.

The Dual methods for SpecialFunctions are now defined in
ForwardDiffSpecialFunctionsExt, together with the rules of other modules whose
derivative calls SpecialFunctions (`NaNMath.lgamma`, whose derivative is
`digamma`). With SpecialFunctions loaded, the set of methods on `Dual` is
unchanged.
@Beforerr Beforerr changed the title Support juliac --trim=safe: SpecialFunctions extension and static default chunk dispatch Support juliac --trim=safe: static default chunk dispatch Oct 23, 2026
`Chunk(x)` returns a Union of 13 `Chunk{N}` types, too wide to union-split, so
building the default config is a dynamic dispatch. `gradient(f, x)`,
`jacobian(f, x)`, `jacobian(f!, y, x)` and `hessian(f, x)` infer as `Any` and
fail `juliac --trim=safe`.

With the new `static_default_chunk` preference, calls without a config branch on
the chunk size before building it, with one branch per size `pickchunksize` can
return. Each branch is statically dispatched, so these calls infer concretely and
pass `--trim=safe`, including when made by other packages. The first call then
compiles every chunk size, e.g. a first `hessian(f, rand(9))` takes 3.5 s instead
of 0.5 s, so the preference is off by default and the default path is unchanged.

A new CI job builds and runs a `--trim=safe` executable from `test/trim`, which
sets the preference.
@Beforerr Beforerr changed the title Support juliac --trim=safe: static default chunk dispatch Support juliac --trim=safe: SpecialFunctions extension, opt-in static default chunk Oct 23, 2026
@codecov

codecov Bot commented Oct 23, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.10345% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.22%. Comparing base (a3c0f4f) to head (56031c7).

Files with missing lines Patch % Lines
src/prelude.jl 20.00% 4 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master     #864      +/-   ##
==========================================
- Coverage   91.23%   91.22%   -0.01%     
==========================================
  Files          11       12       +1     
  Lines        1072     1117      +45     
==========================================
+ Hits          978     1019      +41     
- Misses         94       98       +4     

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

@devmotion

Copy link
Copy Markdown
Member

There's no point in having a SpecialFunctions extension - it's a hard dependency of DiffRules and hence the extension is always loaded. One would first need something like JuliaDiff/DiffRules.jl#120 that removes such hard dependencies in DiffRules.

@Beforerr

Beforerr commented Oct 23, 2026 •

Copy link
Copy Markdown
Author

Hi @devmotion . I checked that DiffRules declares SpecialFunctions in [deps] but never imports it, so the extension stays unloaded unless the user loads SpecialFunctions.

With this PR

using ForwardDiff
haskey(Base.loaded_modules, SpecialFunctions) = false
get_extension(ForwardDiff, :ForwardDiffSpecialFunctionsExt) = nothing

In a minimal trimmed program, removing SpecialFunctions hard dep halved the startup time.

This branch has not been deployed

No deployments
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.

2 participants