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/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/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/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 # ################### 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))