Reduce the cost of NaN-safe mode by checking the scalar factor once

[KristofferC]

Related to #857.

In NaN-safe mode every partial is currently masked with iszero(partial) ? zero(y) : y. The mask only changes the result when the scalar factor is non-finite (0 * Inf, 0 * NaN) for multiplication, or zero/NaN (0 / 0, 0 / NaN) for division. This PR checks the factor once per tuple operation (*, /, _mul_partials) and only takes the masked path in that case; otherwise plain *// is used over the tuple. _isfinite_all recurses into nested Duals (value and partials) so Hessians also hit the fast path. Division by a nested Dual always uses the masked path. With NaN-safe mode disabled nothing changes (the checks are constant false).

Benchmarks

Minimum time, Apple Silicon, Julia 1.13-DEV:

Benchmark	NaN-safe off	NaN-safe on (master)	NaN-safe on (this PR)
Rosenbrock gradient, n=1000	298 μs	487 μs (+63%)	359 μs (+20%)
Rosenbrock Hessian, n=20	8.8 μs	23.8 μs (+170%)	9.9 μs (+12%)
Ackley gradient, n=1000	788 μs	847 μs (+7%)	808 μs (+3%)
(A*x).^2 Jacobian, n=30	5.3 μs	5.3 μs	5.3 μs

The remaining overhead is one fsub + fcmp + well-predicted branch per operation.

Notes

• Dual{3}^Dual{3} is no longer SLP-vectorized in NaN-safe mode on aarch64 (three scalar fmuls), which adds two failures to SIMDTest.jl with NaN-safe mode enabled locally (that file already has failures on master on this machine).

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.15%. Comparing base (a3c0f4f) to head (762ff77).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #859      +/-   ##
==========================================
- Coverage   91.23%   91.15%   -0.08%     
==========================================
  Files          11       11              
  Lines        1072     1085      +13     
==========================================
+ Hits          978      989      +11     
- Misses         94       96       +2     

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

I looked into this PR (and a few alternatives) in some detail. The speedups are real for several benchmarks, but I'm hesitant to merge it in its current form: Conceptually, I think it checks the wrong thing, and its performance depends on what LLVM does with the branch, which leads to large regressions in some cases.

Correctness

For IEEEFloats, the PR seems to give the same results as master up to the sign of zero: A zero partial multiplied by a finite negative factor yields -0.0 instead of 0.0 (e.g., ForwardDiff.hessian(v -> -3 * v[1] * v[2], [1.0, 2.0]) has -0.0 on the diagonal, whereas it's 0.0 on master).

However, I think checking the factor is conceptually wrong. The math underlying NaN-safe mode is about the partials: A dual number with a zero partial is merely unperturbed in that direction, i.e., a regular value, and hence its contribution has to be exactly zero, regardless of the factor. Master implements exactly this, for every partial. The PR replaces it with a shortcut ("0 * x can only be non-zero if x is non-finite") that holds for IEEE floats but ties correctness to isfinite (and _isfinite_all) being the right predicate for every Real type that ends up as a factor. The signed zeros are just a symptom of the rule being a different one.

Additionally, it adds a new requirement for number types: On master, NaN-safe mode only uses iszero of the partials, whereas with this PR every factor has to support isfinite. For custom Real types that don't implement it, Base's fallback isfinite(::Real) throws (MethodError: no method matching decompose(...)). Thus with this PR, e.g., ForwardDiff.derivative(x -> x * x, x0) fails for such types in NaN-safe mode whereas it works on master. If NaN-safe mode becomes the default (#858), this would be breaking for such types.

I'd recommend to stick with checking for zero partials (see below). But if the factor check is pursued further, I'd at least restrict it to IEEEFloats, for which "0 * x is zero if and only if x is finite" is guaranteed, and to Duals built from them, and use master's per-partial mask for all other types. With dispatch this is resolved at compile time:

_mul_needs_mask(x::Base.IEEEFloat) = !isfinite(x)
_mul_needs_mask(x::Dual) = _mul_needs_mask(value(x)) | any(_mul_needs_mask, partials(x).values)
_mul_needs_mask(::Real) = true  # always use the per-partial mask

For division, _div_needs_mask(::Real) = true for Dual divisors should be kept. Extending the factor check recursively is not straightforward since the predicate would have to mirror the arithmetic of Dual division rather than the math: E.g., Dual(0.0, 0.0) / Dual(1e-200, 1.0) == Dual(0.0, NaN) even though the divisor is finite and nonzero at all levels, since vy * vy underflows in _div_partials. More generally, this shows that the factor check couples the masking to implementation details of the factor's type - the per-partial check doesn't.

Performance

The benchmarks in the PR description are minimum times, and I saw quite some variation between runs. So I tried to get more robust numbers (Apple M-series, Julia 1.13.1; every configuration in a fresh process, 10 interleaved rounds, medians with bootstrap 95% CIs; all ratios relative to NaN-safe mode disabled on master).

Whole programs (the benchmarks from #858):

benchmark	NaN-safe on, master	NaN-safe on, this PR
gradient, Rosenbrock, n = 1000	1.47×	1.17×
Hessian, Rosenbrock, n = 10	1.59×	0.93×
gradient, sum(exp, x), n = 100	1.03×	0.88×
Jacobian, cumprod, n = 10	1.07×	0.99×
derivative, sin(x) * exp(x) / x	1.00×	1.12×
gradient, sum(xi -> xi^x[1], x), n = 100	1.04×	1.33×

For the Rosenbrock benchmarks the PR removes most of the overhead. Some ratios are even below 1×, which suggests that part of the differences is caused by code generation and layout effects rather than the cost of the mask itself. For the same reason I'd take the powsum regression with a grain of salt: When I investigated it, a small unrelated change to ^ on master (only adding an iszero(vy) branch around the computation of powval) caused a similar slowdown (1.41×) of the same gradient with chunk size 8.

Single operations in a loop over 1024 Duals with N partials:

N	op	NaN-safe on, master	NaN-safe on, this PR
3	x * y	1.34×	1.85×
3	x / y	1.29×	1.53×
3	a * x (a::Float64)	1.26×	1.51×
3	x^2	1.27×	1.50×
4	x * y	1.51×	0.89×
4	x / y	1.40×	0.85×
4	a * x	1.47×	1.34×
4	x^2	1.14×	1.04×
8	x * y	1.29×	1.02×
8	x / y	1.38×	1.02×
8	a * x	0.99×	3.86×
8	x^2	1.00×	3.79×

The 3.8× regressions for N = 8 are caused by an inlining failure: With both the masked and the unmasked tuple operations in the method body, Real * Dual{8} exceeds the inlining budget, and the loop calls it and passes the Dual through the stack:

using ForwardDiff, Chairmarks
using ForwardDiff: Dual

f(a, x) = a * x
function g!(out, as, xs)
    @inbounds for i in eachindex(out, as, xs)
        out[i] = f(as[i], xs[i])
    end
    return out
end

xs = [Dual(rand(), ntuple(_ -> rand(), 8)...) for _ in 1:1024]
as = rand(1024)
out = similar(xs)

@code_typed g!(out, as, xs)  # this PR: `invoke Main.f(...)` in the loop body
@b g!($out, $as, $xs)        # NaN-safe mode: master 1.18 μs, this PR 5.18 μs

The other differences are caused by what LLVM does with the branch, which seems hard to control from Julia:

• For N = 3, the loop vectorizer vectorizes across elements and if-converts the branch. Thus both paths are computed, and the finiteness check is added on top of the mask.
• The branch is also the reason why Dual{3}^Dual{3} is no longer SLP-vectorized (the note in the PR description).

Alternatives

I tried a few variants to check whether the approach can be made robust. None of them was faster than master across the board:

• Computing the products once and zeroing partials only if the factor is non-finite: LLVM (SimplifyCFG/the loop vectorizer) frequently speculates the cheap fix-up into selects, which then is strictly more expensive than master's mask (up to 2.7×).
• Moving the masked path to a @noinline function to keep the branch: With tuple arguments, the tuples go through the stack also on the fast path (2-6×). With scalar arguments, it's fast for N = 3 but for larger N the additional call sites push the operations over the inlining budget and they are not inlined anymore (up to 9×).
• Branch hints with llvm.expect in llvmcall: LowerExpectIntrinsicPass lowers them before the llvmcall function is inlined, so they never reach the branch.
• Short-circuiting if all partials are zero (i.e., skipping all computations on the partials in that case, in addition to the per-partial mask): Correct, type-stable, and bit-identical to master in NaN-safe mode, but slower than master in the whole-program benchmarks (e.g., 1.69× vs 1.47× for the Rosenbrock gradient) - even though in that benchmark 98.8% of the operations on partials only involve zero partials.

Recommendation

Hence I'd recommend to stick with checking for zero partials, as on master:

• Conceptually, it implements exactly the underlying math (a zero partial means unperturbed, so its contribution is zero) for every number type, without relying on properties of the factor's type or implementation details of its arithmetic.
• Its lowering is predictable: The per-partial mask is branch-free and SIMD-friendly (fcmeq + bic on aarch64). There's no branch that LLVM could if-convert, speculate, or keep, and no additional code that could push operations over the inlining budget. Its cost is moderate and consistent: In none of the benchmarks above it's much slower than NaN-safe mode disabled (at most 1.59×), whereas the gains of this PR come with regressions of up to 3.8× that depend on N, the call site, and LLVM's heuristics.

All measurements are from a single machine (aarch64), I haven't checked x86.