From 99c36e4bafbca103bdbea573c8519497ccf0740d Mon Sep 17 00:00:00 2001 From: Beforerr Date: Fri, 23 Oct 2026 09:01:05 -0600 Subject: [PATCH 1/2] Move SpecialFunctions support to a package extension `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. --- Project.toml | 6 +++-- ext/ForwardDiffSpecialFunctionsExt.jl | 38 +++++++++++++++++++++++++++ src/ForwardDiff.jl | 1 - src/dual.jl | 32 ++++++++++------------ 4 files changed, 56 insertions(+), 21 deletions(-) create mode 100644 ext/ForwardDiffSpecialFunctionsExt.jl diff --git a/Project.toml b/Project.toml index 801792c6..39712763 100644 --- a/Project.toml +++ b/Project.toml @@ -12,12 +12,13 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [weakdeps] +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [extensions] +ForwardDiffSpecialFunctionsExt = "SpecialFunctions" ForwardDiffStaticArraysExt = "StaticArrays" [compat] @@ -42,8 +43,9 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Calculus", "DiffTests", "IrrationalConstants", "JET", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"] +test = ["Calculus", "DiffTests", "IrrationalConstants", "JET", "SparseArrays", "SpecialFunctions", "StaticArrays", "Test", "InteractiveUtils"] diff --git a/ext/ForwardDiffSpecialFunctionsExt.jl b/ext/ForwardDiffSpecialFunctionsExt.jl new file mode 100644 index 00000000..eb6c0df2 --- /dev/null +++ b/ext/ForwardDiffSpecialFunctionsExt.jl @@ -0,0 +1,38 @@ +module ForwardDiffSpecialFunctionsExt + +using ForwardDiff: ForwardDiff, Dual, DiffRules, NaNMath, value, partials, + uses_specialfunctions, unary_dual_definition, binary_dual_definition +import SpecialFunctions + +for (M, f, arity) in DiffRules.diffrules(filter_modules = nothing) + if !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f)) + continue + elseif !uses_specialfunctions(M, f, arity) + continue + end + if arity == 1 + eval(unary_dual_definition(M, f)) + elseif arity == 2 + eval(binary_dual_definition(M, f)) + end +end + +# Functions in SpecialFunctions which return tuples # +# Their derivatives are not defined in DiffRules # +#---------------------------------------------------# + +function SpecialFunctions.logabsgamma(d::Dual{T,<:Real}) where {T} + x = value(d) + y, s = SpecialFunctions.logabsgamma(x) + return (Dual{T}(y, SpecialFunctions.digamma(x) * partials(d)), s) +end + +# Derivatives wrt to first parameter and precision setting are not supported +function SpecialFunctions.gamma_inc(a::Real, d::Dual{T,<:Real}, ind::Integer) where {T} + x = value(d) + p, q = SpecialFunctions.gamma_inc(a, x, ind) + ∂p = exp(-x) * x^(a - 1) / SpecialFunctions.gamma(a) * partials(d) + return (Dual{T}(p, ∂p), Dual{T}(q, -∂p)) +end + +end # module diff --git a/src/ForwardDiff.jl b/src/ForwardDiff.jl index b16b986b..22790735 100644 --- a/src/ForwardDiff.jl +++ b/src/ForwardDiff.jl @@ -8,7 +8,6 @@ using LinearAlgebra using Base: require_one_based_indexing import Printf import NaNMath -import SpecialFunctions import LogExpFunctions import CommonSubexpressions diff --git a/src/dual.jl b/src/dual.jl index 6d13dec3..b0419f1f 100644 --- a/src/dual.jl +++ b/src/dual.jl @@ -473,11 +473,25 @@ Base.float(d::Dual) = convert(float(typeof(d)), d) # General Mathematical Operations # ################################### +mentions(ex::Expr, s::Symbol) = any(a -> mentions(a, s), ex.args) +mentions(ex::Tuple, s::Symbol) = any(a -> mentions(a, s), ex) +mentions(ex, s::Symbol) = ex === s + +# Rules for SpecialFunctions, and rules of other modules whose derivative calls it +# (e.g. `NaNMath.lgamma`), are defined in ForwardDiffSpecialFunctionsExt. +function uses_specialfunctions(M, f, arity) + M === :SpecialFunctions && return true + args = ntuple(i -> Symbol(:x, i), arity) + return mentions(DiffRules.diffrule(M, f, args...), :SpecialFunctions) +end + for (M, f, arity) in DiffRules.diffrules(filter_modules = nothing) if (M, f) in ((:Base, :^), (:NaNMath, :pow), (:Base, :/), (:Base, :+), (:Base, :-), (:Base, :sin), (:Base, :cos)) continue # Skip methods which we define elsewhere. elseif !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f)) continue # Skip rules for methods not defined in the current scope + elseif uses_specialfunctions(M, f, arity) + continue end if arity == 1 eval(unary_dual_definition(M, f)) @@ -864,24 +878,6 @@ function _eigen_hermitian(A::DualMatrixRealComplex{T,<:Real,N}) where {T,N} return Eigen(_to_duals(Val(T), λ, λ_partials), _to_duals(Val(T), Q, Q_partials)) end -# Functions in SpecialFunctions which return tuples # -# Their derivatives are not defined in DiffRules # -#---------------------------------------------------# - -function SpecialFunctions.logabsgamma(d::Dual{T,<:Real}) where {T} - x = value(d) - y, s = SpecialFunctions.logabsgamma(x) - return (Dual{T}(y, SpecialFunctions.digamma(x) * partials(d)), s) -end - -# Derivatives wrt to first parameter and precision setting are not supported -function SpecialFunctions.gamma_inc(a::Real, d::Dual{T,<:Real}, ind::Integer) where {T} - x = value(d) - p, q = SpecialFunctions.gamma_inc(a, x, ind) - ∂p = exp(-x) * x^(a - 1) / SpecialFunctions.gamma(a) * partials(d) - return (Dual{T}(p, ∂p), Dual{T}(q, -∂p)) -end - ################### # Pretty Printing # ################### From 56031c713c67ee2c1f9c608bd5d8063092b18c9b Mon Sep 17 00:00:00 2001 From: Beforerr Date: Fri, 23 Oct 2026 09:04:28 -0600 Subject: [PATCH 2/2] Add a preference to dispatch statically on the default chunk size `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. --- .github/workflows/ci.yml | 11 +++++++++++ .gitignore | 1 + docs/src/user/advanced.md | 5 +++++ src/gradient.jl | 9 +++++++-- src/hessian.jl | 14 +++++++++++--- src/jacobian.jl | 19 +++++++++++++++---- src/prelude.jl | 14 ++++++++++++++ test/trim/Project.toml | 14 ++++++++++++++ test/trim/app.jl | 8 ++++++++ test/trim/build_and_run.jl | 17 +++++++++++++++++ test/trim/functions.jl | 4 ++++ 11 files changed, 107 insertions(+), 9 deletions(-) create mode 100644 test/trim/Project.toml create mode 100644 test/trim/app.jl create mode 100644 test/trim/build_and_run.jl create mode 100644 test/trim/functions.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aab5bd68..33db8cad 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,6 +52,17 @@ jobs: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true + trim: + name: juliac --trim=safe + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + - uses: julia-actions/cache@v3 + - run: julia --color=yes --project=test/trim -e 'using Pkg; Pkg.instantiate()' + - run: julia --color=yes --project=test/trim test/trim/build_and_run.jl docs: name: Documentation runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 69feffb1..bcf2ff11 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ /docs/Manifest.toml /benchmark_data/ /Manifest.toml +/test/trim/Manifest.toml diff --git a/docs/src/user/advanced.md b/docs/src/user/advanced.md index 69468751..01d547a5 100644 --- a/docs/src/user/advanced.md +++ b/docs/src/user/advanced.md @@ -108,6 +108,11 @@ ForwardDiff.gradient(rosenbrock, x)::Vector{Float64} ForwardDiff.gradient!(out, rosenbrock, x) ``` +Alternatively, setting the `static_default_chunk` preference to true makes calls without +a config type-stable, at the cost of compiling the computation for every chunk size the +heuristic can pick on the first call. Statically compiled programs (`juliac --trim`) that +call ForwardDiff without a config, directly or through other packages, require it. + One final question remains: How should you select a chunk size? The answer is essentially "perform your own benchmarks and see what works best for your use case." As stated before, the optimal chunk size is heavily dependent on the target function and length of the input diff --git a/src/gradient.jl b/src/gradient.jl index a5ef3dac..aa63b830 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -13,7 +13,9 @@ This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function gradient(f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} +gradient(f::F, x::AbstractArray) where {F} = with_default_chunk(c -> gradient(f, x, GradientConfig(f, x, c)), x) + +function gradient(f::F, x::AbstractArray, cfg::GradientConfig{T}, ::Val{CHK}=Val{true}()) where {F, T, CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) if chunksize(cfg) == structural_length(x) @@ -32,7 +34,10 @@ Compute `∇f` evaluated at `x` and store the result(s) in `result`, assuming `f This method assumes that `isa(f(x), Real)`. """ -function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK, F} +gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray) where {F} = + with_default_chunk(c -> gradient!(result, f, x, GradientConfig(f, x, c)), x) + +function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::GradientConfig{T}, ::Val{CHK}=Val{true}()) where {T, CHK, F} result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x) CHK && checktag(T, f, x) if chunksize(cfg) == structural_length(x) diff --git a/src/hessian.jl b/src/hessian.jl index 9c755c9a..d36d82b4 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -11,7 +11,9 @@ This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function hessian(f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T,CHK} +hessian(f::F, x::AbstractArray) where {F} = with_default_chunk(c -> hessian(f, x, HessianConfig(f, x, c)), x) + +function hessian(f::F, x::AbstractArray, cfg::HessianConfig{T}, ::Val{CHK}=Val{true}()) where {F, T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}()) @@ -28,7 +30,10 @@ This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} +hessian!(result::AbstractArray, f::F, x::AbstractArray) where {F} = + with_default_chunk(c -> hessian!(result, f, x, HessianConfig(f, x, c)), x) + +function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T}, ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(result, x) CHK && checktag(T, f, x) ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}()) @@ -63,7 +68,10 @@ because `isa(result, DiffResult)`, `cfg` is constructed as `HessianConfig(f, res Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} +hessian!(result::DiffResult, f::F, x::AbstractArray) where {F} = + with_default_chunk(c -> hessian!(result, f, x, HessianConfig(f, result, x, c)), x) + +function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T}, ::Val{CHK}=Val{true}()) where {F,T,CHK} CHK && checktag(T, f, x) ∇f! = InnerGradientForHess(result, cfg, f) jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}()) diff --git a/src/jacobian.jl b/src/jacobian.jl index f14a6a7b..f0b8e185 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -15,7 +15,9 @@ This method assumes that `isa(f(x), AbstractArray)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian(f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} +jacobian(f::F, x::AbstractArray) where {F} = with_default_chunk(c -> jacobian(f, x, JacobianConfig(f, x, c)), x) + +function jacobian(f::F, x::AbstractArray, cfg::JacobianConfig{T}, ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) if chunksize(cfg) == structural_length(x) @@ -33,7 +35,10 @@ stored in `y`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian(f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T, CHK} +jacobian(f!::F, y::AbstractArray, x::AbstractArray) where {F} = + with_default_chunk(c -> jacobian(f!, y, x, JacobianConfig(f!, y, x, c)), x) + +function jacobian(f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T}, ::Val{CHK}=Val{true}()) where {F,T, CHK} require_one_based_indexing(y, x) CHK && checktag(T, f!, x) if chunksize(cfg) == structural_length(x) @@ -54,7 +59,10 @@ This method assumes that `isa(f(x), AbstractArray)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T, CHK} +jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray) where {F} = + with_default_chunk(c -> jacobian!(result, f, x, JacobianConfig(f, x, c)), x) + +function jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::JacobianConfig{T}, ::Val{CHK}=Val{true}()) where {F,T, CHK} result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x) CHK && checktag(T, f, x) if chunksize(cfg) == structural_length(x) @@ -75,7 +83,10 @@ This method assumes that `isa(f(x), AbstractArray)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} +jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray) where {F} = + with_default_chunk(c -> jacobian!(result, f!, y, x, JacobianConfig(f!, y, x, c)), x) + +function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T}, ::Val{CHK}=Val{true}()) where {F,T,CHK} result isa DiffResult ? require_one_based_indexing(y, x) : require_one_based_indexing(result, y, x) CHK && checktag(T, f!, x) if chunksize(cfg) == structural_length(x) diff --git a/src/prelude.jl b/src/prelude.jl index 9e037afa..c8558cbb 100644 --- a/src/prelude.jl +++ b/src/prelude.jl @@ -1,5 +1,6 @@ const NANSAFE_MODE_ENABLED = @load_preference("nansafe_mode", false) const DEFAULT_CHUNK_THRESHOLD = @load_preference("default_chunk_threshold", 12) +const STATIC_DEFAULT_CHUNK = @load_preference("static_default_chunk", false) const AMBIGUOUS_TYPES = (AbstractFloat, Irrational, Integer, Rational, Real, RoundingMode) @@ -23,6 +24,19 @@ function Chunk(x::AbstractArray, threshold::Integer = DEFAULT_CHUNK_THRESHOLD) return Chunk(structural_length(x), threshold) end +if STATIC_DEFAULT_CHUNK + # One branch per chunk size: the Union returned by `Chunk(x)` is too wide to split. + # Exhaustive since `pickchunksize` never exceeds the threshold. + @eval function with_default_chunk(f::F, x::AbstractArray) where {F} + N = pickchunksize(structural_length(x)) + N == 0 && return f(Chunk{0}()) + return Base.@nif $DEFAULT_CHUNK_THRESHOLD d -> (N == d) d -> f(Chunk{d}()) d -> f(Chunk{d}()) + end +else + # Dynamic dispatch compiles only the chunk size that is used. + with_default_chunk(f::F, x::AbstractArray) where {F} = f(Chunk(x)) +end + # Constrained to `N <= threshold`, minimize (in order of priority): # 1. the number of chunks that need to be computed # 2. the number of "left over" perturbations in the final chunk diff --git a/test/trim/Project.toml b/test/trim/Project.toml new file mode 100644 index 00000000..8689b1a0 --- /dev/null +++ b/test/trim/Project.toml @@ -0,0 +1,14 @@ +[deps] +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +ForwardDiff = {path = "../.."} + +[preferences.ForwardDiff] +static_default_chunk = true + +[compat] +JuliaC = "0.3" +julia = "1.12" diff --git a/test/trim/app.jl b/test/trim/app.jl new file mode 100644 index 00000000..6d1a4d78 --- /dev/null +++ b/test/trim/app.jl @@ -0,0 +1,8 @@ +using ForwardDiff +include("functions.jl") + +function (@main)(args::Vector{String})::Cint + n = parse(Int, args[1]) + foreach(a -> foreach(v -> println(Core.stdout, v), a), results(inputs(n))) + return 0 +end diff --git a/test/trim/build_and_run.jl b/test/trim/build_and_run.jl new file mode 100644 index 00000000..360f72df --- /dev/null +++ b/test/trim/build_and_run.jl @@ -0,0 +1,17 @@ +using ForwardDiff, JuliaC, Test +include("functions.jl") + +@testset "juliac --trim=safe" begin + mktempdir() do dir + image = ImageRecipe(output_type = "--output-exe", file = joinpath(@__DIR__, "app.jl"), + trim_mode = "safe", project = @__DIR__) + compile_products(image) + exe = joinpath(dir, "app") + link_products(LinkRecipe(image_recipe = image, outname = exe)) + # below and above the default chunk threshold + for n in (3, 20) + expected = mapreduce(vec, vcat, results(inputs(n))) + @test parse.(Float64, readlines(`$exe $n`)) ≈ expected + end + end +end diff --git a/test/trim/functions.jl b/test/trim/functions.jl new file mode 100644 index 00000000..c6d6bff7 --- /dev/null +++ b/test/trim/functions.jl @@ -0,0 +1,4 @@ +f(x) = sum(sin, x) * x[1]^2 + exp(x[end]) +g(x) = x .* sin(x[1]) .+ reverse(x) +inputs(n) = collect(range(0.1, 1.0; length = n)) +results(x) = (ForwardDiff.gradient(f, x), ForwardDiff.jacobian(g, x), ForwardDiff.hessian(f, x))