Use FunctionWrappersWrappers fallback for non-isbits types instead of erroring
## Context
PR #1299 reverted AutoSpecialize as the default for ODEFunction because several non-isbits types (BigFloat, SparseConnectivityTracer tracers) couldn't match the pre-compiled FunctionWrapper signatures, causing `NoFunctionWrapperFoundError`.
The core issue: `FunctionWrappersWrapper{FW, false}` (fallback disabled) throws when argument types don't match any pre-compiled wrapper. For non-isbits types this is always the case since wrapper signatures are compiled for `Float64` and `ForwardDiff.Dual{Float64,...}` variants only.
## Proposal
When constructing the `FunctionWrappersWrapper`, detect whether the element type of `u0` is `isbits`. If not, use `FunctionWrappersWrapper{FW, true}` (fallback enabled) instead of `FunctionWrappersWrapper{FW, false}`. The fallback path calls the original unwrapped function via `first(fww.fw).obj[](args...)`, which is equivalent to direct dispatch.
The rationale: non-isbits types already allocate heavily, so the marginal cost of the fallback dispatch path is negligible.
## Benchmark Results
Benchmarked on Julia 1.12 with BenchmarkTools.jl. Compared:
- **Direct call**: `f(du, u, p, t)` — what FullSpecialize gives you
- **FWW fallback (1 sig)**: `FunctionWrappersWrapper` with 1 Float64 signature, `fallback=true`, called with non-matching types
- **FWW fallback (4 sig)**: Same with 4 Float64 signatures (closer to real SciMLBase usage where 4 ForwardDiff variants are pre-compiled)
### Test functions
- **Lorenz** (N=3): Small ODE, ~5ns for Float64
- **Heat equation** (N=100): Larger discretized PDE, ~63ns for Float64
### Results
| Type | Problem | Direct (median) | FWW fb 1sig | FWW fb 4sig | Overhead |
|------|---------|-----------------|-------------|-------------|----------|
| Float64 (isbits, match) | Lorenz N=3 | 5.8 ns, 0 allocs | 19.3 ns (match) | — | 13.5 ns (FWW overhead) |
| Float64 (isbits, match) | Heat N=100 | 63.2 ns, 0 allocs | 74.4 ns (match) | — | 11.2 ns (FWW overhead) |
| **BigFloat** | Lorenz N=3 | 603 ns, 8 allocs | 666 ns, 9 allocs | 654 ns, 9 allocs | **+10% / +8%** |
| **BigFloat** | Heat N=100 | 103.1 μs, 498 allocs | 103.2 μs, 499 allocs | 103.6 μs, 499 allocs | **+0.1% / +0.5%** |
| **GradientTracer** | Lorenz N=3 | 592 ns, 24 allocs | 669 ns, 25 allocs | 649 ns, 25 allocs | **+13% / +10%** |
| **GradientTracer** | Heat N=100 | 23.9 μs, 896 allocs | 24.7 μs, 897 allocs | 24.0 μs, 897 allocs | **+3% / +0.4%** |
| **Dual{F64,GT}** | Lorenz N=3 | 475 ns, 18 allocs | 524 ns, 19 allocs | 535 ns, 19 allocs | **+10% / +13%** |
### Key observations
1. **The fallback adds exactly 1 allocation** (tuple construction for the arg splat) — negligible relative to the hundreds of allocations non-isbits types already produce.
2. **For the smallest possible problem (Lorenz N=3), the overhead is ~50-80ns absolute** (~10-13% relative). This is the worst case — the function body is tiny so the constant overhead of the fallback dispatch is most visible.
3. **For any realistic problem size (N≥100), the overhead drops below 1%** and is within measurement noise. The BigFloat Heat N=100 case shows 103.1μs vs 103.2μs — completely negligible.
4. **The number of pre-compiled signatures (1 vs 4) doesn't meaningfully affect fallback cost** — the linear search through the wrapper tuple is fast since it's compile-time unrolled.
5. **Float64 isbits path is unaffected** — when types match, the existing FunctionWrapper ccall path is used regardless of whether `fallback=true` or `false`.
## Suggested implementation
In `wrapfun_iip` (and `wrapfun_oop`), check `isbitstype(eltype(u0))` and set the fallback flag accordingly:
```julia
function wrapfun_iip(ff, inputs::Tuple{T1, T2, T3, T4}) where {T1, T2, T3, T4}
# ... existing code to build iip_arglists, iip_returnlists ...
# Enable fallback for non-isbits types (BigFloat, tracers, etc.)
# These types already allocate heavily, so fallback overhead is negligible
use_fallback = !isbitstype(eltype(T2))
fwt = map(iip_arglists, iip_returnlists) do A, R
FunctionWrappers.FunctionWrapper{R, A}(Void(ff))
end
return FunctionWrappersWrapper{typeof(fwt), use_fallback}(fwt)
end
```
This would allow AutoSpecialize to be the default for ODEFunction (addressing the goal of #1298) without breaking BigFloat, SparseConnectivityTracer, or any other non-isbits type.
## Benchmark script
<details>
<summary>Full benchmark script (click to expand)</summary>
```julia
using Pkg
Pkg.activate(temp=true)
Pkg.add(["FunctionWrappers", "FunctionWrappersWrappers", "BenchmarkTools", "SparseConnectivityTracer"])
using FunctionWrappers, FunctionWrappersWrappers, BenchmarkTools, SparseConnectivityTracer
using Printf
function lorenz!(du, u, p, t)
σ, ρ, β = p[1], p[2], p[3]
du[1] = σ * (u[2] - u[1])
du[2] = u[1] * (ρ - u[3]) - u[2]
du[3] = u[1] * u[2] - β * u[3]
nothing
end
function heat!(du, u, p, t)
N = length(u)
α = p[1]
dx = 1.0 / (N + 1)
du[1] = α * (u[2] - 2u[1]) / dx^2
for i in 2:(N-1)
du[i] = α * (u[i+1] - 2u[i] + u[i-1]) / dx^2
end
du[N] = α * (-2u[N] + u[N-1]) / dx^2
nothing
end
function make_fww_f64(f; fallback=false)
arglists = (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},)
returnlists = (Nothing,)
fwt = map(arglists, returnlists) do A, R
FunctionWrappers.FunctionWrapper{R, A}(f)
end
FunctionWrappersWrappers.FunctionWrappersWrapper{typeof(fwt), fallback}(fwt)
end
function make_fww_bf(f; fallback=false)
arglists = (Tuple{Vector{BigFloat}, Vector{BigFloat}, Vector{BigFloat}, BigFloat},)
returnlists = (Nothing,)
fwt = map(arglists, returnlists) do A, R
FunctionWrappers.FunctionWrapper{R, A}(f)
end
FunctionWrappersWrappers.FunctionWrappersWrapper{typeof(fwt), fallback}(fwt)
end
function make_fww_f64_multi(f; fallback=false)
arglists = (
Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},
Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},
Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},
Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},
)
returnlists = (Nothing, Nothing, Nothing, Nothing)
fwt = map(arglists, returnlists) do A, R
FunctionWrappers.FunctionWrapper{R, A}(f)
end
FunctionWrappersWrappers.FunctionWrappersWrapper{typeof(fwt), fallback}(fwt)
end
# Float64 baseline
u_f = [1.0, 0.0, 0.0]; du_f = similar(u_f); p_f = [10.0, 28.0, 8/3]; t_f = 0.0
fww_f_match = make_fww_f64(lorenz!)
b_f_direct = @benchmark $lorenz!($du_f, $u_f, $p_f, $t_f)
b_f_match = @benchmark $fww_f_match($du_f, $u_f, $p_f, $t_f)
# BigFloat Lorenz
u_bf = BigFloat.([1.0, 0.0, 0.0]); du_bf = similar(u_bf)
p_bf = BigFloat.([10.0, 28.0, 8/3]); t_bf = BigFloat(0.0)
fww_bf_match = make_fww_bf(lorenz!)
fww_bf_fb = make_fww_f64(lorenz!; fallback=true)
fww_bf_fb_multi = make_fww_f64_multi(lorenz!; fallback=true)
b_bf_direct = @benchmark $lorenz!($du_bf, $u_bf, $p_bf, $t_bf)
b_bf_match = @benchmark $fww_bf_match($du_bf, $u_bf, $p_bf, $t_bf)
b_bf_fb = @benchmark $fww_bf_fb($du_bf, $u_bf, $p_bf, $t_bf)
b_bf_fb_multi = @benchmark $fww_bf_fb_multi($du_bf, $u_bf, $p_bf, $t_bf)
# BigFloat Heat N=100
N_h = 100
u_hbf = BigFloat.(sin.(range(0, π, length=N_h))); du_hbf = similar(u_hbf)
p_hbf = BigFloat.([0.01]); t_hbf = BigFloat(0.0)
fww_hbf_fb = make_fww_f64(heat!; fallback=true)
fww_hbf_fb_multi = make_fww_f64_multi(heat!; fallback=true)
b_hbf_direct = @benchmark $heat!($du_hbf, $u_hbf, $p_hbf, $t_hbf)
b_hbf_fb = @benchmark $fww_hbf_fb($du_hbf, $u_hbf, $p_hbf, $t_hbf)
b_hbf_fb_multi = @benchmark $fww_hbf_fb_multi($du_hbf, $u_hbf, $p_hbf, $t_hbf)
# GradientTracer
GT = SparseConnectivityTracer.GradientTracer{Int, BitSet}
u_gt = [GT(BitSet([1])), GT(BitSet([2])), GT(BitSet([3]))]
du_gt = similar(u_gt); p_gt = [GT(BitSet()), GT(BitSet()), GT(BitSet())]; t_gt = GT(BitSet())
fww_gt_fb = make_fww_f64(lorenz!; fallback=true)
fww_gt_fb_multi = make_fww_f64_multi(lorenz!; fallback=true)
b_gt_direct = @benchmark $lorenz!($du_gt, $u_gt, $p_gt, $t_gt)
b_gt_fb = @benchmark $fww_gt_fb($du_gt, $u_gt, $p_gt, $t_gt)
b_gt_fb_multi = @benchmark $fww_gt_fb_multi($du_gt, $u_gt, $p_gt, $t_gt)
# GradientTracer Heat N=100
u_gt_h = [GT(BitSet([i])) for i in 1:N_h]
du_gt_h = similar(u_gt_h); p_gt_h = [GT(BitSet())]; t_gt_h = GT(BitSet())
fww_gt_h_fb = make_fww_f64(heat!; fallback=true)
fww_gt_h_fb_multi = make_fww_f64_multi(heat!; fallback=true)
b_gt_h_direct = @benchmark $heat!($du_gt_h, $u_gt_h, $p_gt_h, $t_gt_h)
b_gt_h_fb = @benchmark $fww_gt_h_fb($du_gt_h, $u_gt_h, $p_gt_h, $t_gt_h)
b_gt_h_fb_multi = @benchmark $fww_gt_h_fb_multi($du_gt_h, $u_gt_h, $p_gt_h, $t_gt_h)
# Dual tracer
DualGT = SparseConnectivityTracer.Dual{Float64, GT}
u_dgt = [DualGT(1.0, GT(BitSet([1]))), DualGT(0.0, GT(BitSet([2]))), DualGT(0.0, GT(BitSet([3])))]
du_dgt = similar(u_dgt)
p_dgt = [DualGT(10.0, GT(BitSet())), DualGT(28.0, GT(BitSet())), DualGT(8/3, GT(BitSet()))]
t_dgt = DualGT(0.0, GT(BitSet()))
fww_dgt_fb = make_fww_f64(lorenz!; fallback=true)
fww_dgt_fb_multi = make_fww_f64_multi(lorenz!; fallback=true)
b_dgt_direct = @benchmark $lorenz!($du_dgt, $u_dgt, $p_dgt, $t_dgt)
b_dgt_fb = @benchmark $fww_dgt_fb($du_dgt, $u_dgt, $p_dgt, $t_dgt)
b_dgt_fb_multi = @benchmark $fww_dgt_fb_multi($du_dgt, $u_dgt, $p_dgt, $t_dgt)
# Print summary
println("Results:")
for (name, bd, bf1, bf4) in [
("BigFloat Lorenz N=3", b_bf_direct, b_bf_fb, b_bf_fb_multi),
("BigFloat Heat N=100", b_hbf_direct, b_hbf_fb, b_hbf_fb_multi),
("GradTracer Lorenz N=3", b_gt_direct, b_gt_fb, b_gt_fb_multi),
("GradTracer Heat N=100", b_gt_h_direct, b_gt_h_fb, b_gt_h_fb_multi),
("Dual{F64,GT} Lorenz N=3", b_dgt_direct, b_dgt_fb, b_dgt_fb_multi),
]
td = median(bd).time; tf1 = median(bf1).time; tf4 = median(bf4).time
o1 = (tf1-td)/td*100; o4 = (tf4-td)/td*100
println(" $name: direct=$(round(td, sigdigits=4)), fb1=$(round(tf1, sigdigits=4)) ($(round(o1, digits=1))%), fb4=$(round(tf4, sigdigits=4)) ($(round(o4, digits=1))%)")
end
```
</details>
0 条评论