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

.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

.gitignore

@@ -7,3 +7,4 @@
 /docs/Manifest.toml
 /benchmark_data/
 /Manifest.toml
+/test/trim/Manifest.toml

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

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

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

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
 

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

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)

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}())

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)

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

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"

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

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

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