From 5e17c90cc3fc2ceccf4956986056f445f45dc8b2 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 28 Jul 2026 05:19:15 -0400 Subject: [PATCH 1/4] Unseed only the chunk in 3-arg seed! The chunk-mode unseed call seed!(xdual, x, i) only needs to clear the N-wide chunk seeded at i: the rest of the array is zeroed up front and every other chunk clears itself. ForwardDiff 0.10 wrote exactly N elements here; the 1.x rewrite made it write from i to the end of the array, i.e. O(n^2/2N) redundant dual writes per chunked gradient/jacobian sweep (~40 GB of memory traffic for gradient! of 100000 elements at chunk 12). Write at most N elements starting at index. The 3-arg seed! form has no other callers in the package. gradient! of sum(abs2, x) at chunk 12: n=1000 502 -> 279 us, n=100000 5.39 -> 2.99 s. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019GcNzbNzaHCqTm4W14eKhN --- src/apiutils.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/apiutils.jl b/src/apiutils.jl index f401a3fc..dc9b8770 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -106,10 +106,13 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, return duals end +# Writes at most N elements starting at `index`: chunk mode only ever needs to +# clear the N-wide chunk it just seeded, so writing through to the end of the +# array would be O(n) redundant work per chunk (O(n^2) per sweep). function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index, seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N} offset = index - 1 - idxs = Iterators.drop(structural_eachindex(duals, x), offset) + idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), offset), N) if isbitstype(V) for idx in idxs duals[idx] = Dual{T,V,N}(x[idx], seed) From b8a9b17f0db226a83962dad981d9b192f60fa375 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 28 Jul 2026 19:46:34 -0400 Subject: [PATCH 2/4] Split unseeding out of seed! into unseed! Review feedback: the zero-Partials default of seed! made the unseeding calls read as seeding. Give them their own name and drop the seed argument, which was never passed for those two methods (the only caller of the explicit-seed form was the allocation test). seed!(duals, x) -> unseed!(duals, x) seed!(duals, x, index) -> unseed!(duals, x, index) The NTuple-of-seeds methods are unchanged. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019GcNzbNzaHCqTm4W14eKhN --- src/apiutils.jl | 41 +++++++++++++++++++++-------------------- src/derivative.jl | 4 ++-- src/gradient.jl | 6 +++--- src/jacobian.jl | 16 ++++++++-------- test/AllocationsTest.jl | 13 +++++++------ 5 files changed, 41 insertions(+), 39 deletions(-) diff --git a/src/apiutils.jl b/src/apiutils.jl index dc9b8770..4617c88e 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -27,7 +27,7 @@ end function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig, y, x) where {F} ydual, xdual = cfg.duals seed!(xdual, x, cfg.seeds) - seed!(ydual, y) + unseed!(ydual, y) f!(ydual, xdual) return ydual end @@ -70,8 +70,10 @@ function structural_eachindex(x::Diagonal, y::AbstractArray) return diagind(x) end -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, - seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N} +# Copies the values of `x` into `duals` with zero partials, i.e. removes any seeds +# `duals` is currently carrying. +function unseed!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} + seed = zero(Partials{N,V}) if isbitstype(V) for idx in structural_eachindex(duals, x) duals[idx] = Dual{T,V,N}(x[idx], seed) @@ -88,16 +90,20 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, return duals end -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, - seeds::NTuple{N,Partials{N,V}}) where {T,V,N} +# Unseeds at most `N` elements starting at `index`: chunk mode only ever needs to clear +# the N-wide chunk it just seeded, so writing through to the end of the array would be +# O(n) redundant work per chunk (O(n^2) per sweep). +function unseed!(duals::AbstractArray{Dual{T,V,N}}, x, index) where {T,V,N} + seed = zero(Partials{N,V}) + idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), N) if isbitstype(V) - for (i, idx) in zip(1:N, structural_eachindex(duals, x)) - duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) + for idx in idxs + duals[idx] = Dual{T,V,N}(x[idx], seed) end else - for (i, idx) in zip(1:N, structural_eachindex(duals, x)) + for idx in idxs if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) + duals[idx] = Dual{T,V,N}(x[idx], seed) else Base._unsetindex!(duals, idx) end @@ -106,21 +112,16 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, return duals end -# Writes at most N elements starting at `index`: chunk mode only ever needs to -# clear the N-wide chunk it just seeded, so writing through to the end of the -# array would be O(n) redundant work per chunk (O(n^2) per sweep). -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index, - seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N} - offset = index - 1 - idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), offset), N) +function seed!(duals::AbstractArray{Dual{T,V,N}}, x, + seeds::NTuple{N,Partials{N,V}}) where {T,V,N} if isbitstype(V) - for idx in idxs - duals[idx] = Dual{T,V,N}(x[idx], seed) + for (i, idx) in zip(1:N, structural_eachindex(duals, x)) + duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) end else - for idx in idxs + for (i, idx) in zip(1:N, structural_eachindex(duals, x)) if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seed) + duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) else Base._unsetindex!(duals, idx) end diff --git a/src/derivative.jl b/src/derivative.jl index b39e2a48..d9fb355a 100644 --- a/src/derivative.jl +++ b/src/derivative.jl @@ -27,7 +27,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba require_one_based_indexing(y) CHK && checktag(T, f!, x) ydual = cfg.duals - seed!(ydual, y) + unseed!(ydual, y) f!(ydual, Dual{T}(x, one(x))) map!(value, y, ydual) return extract_derivative(T, ydual) @@ -65,7 +65,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba result isa DiffResult ? require_one_based_indexing(y) : require_one_based_indexing(result, y) CHK && checktag(T, f!, x) ydual = cfg.duals - seed!(ydual, y) + unseed!(ydual, y) f!(ydual, Dual{T}(x, one(x))) result = extract_value!(T, result, y, ydual) result = extract_derivative!(T, result, ydual) diff --git a/src/gradient.jl b/src/gradient.jl index 0832d354..d76a476e 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -127,14 +127,14 @@ function chunk_mode_gradient_expr(result_definition::Expr) # seed work vectors xdual = cfg.duals seeds = cfg.seeds - seed!(xdual, x) + unseed!(xdual, x) # do first chunk manually to calculate output type seed!(xdual, x, 1, seeds) ydual = f(xdual) $(result_definition) extract_gradient_chunk!(T, result, ydual, 1, N) - seed!(xdual, x, 1) + unseed!(xdual, x, 1) # do middle chunks for c in middlechunks @@ -142,7 +142,7 @@ function chunk_mode_gradient_expr(result_definition::Expr) seed!(xdual, x, i, seeds) ydual = f(xdual) extract_gradient_chunk!(T, result, ydual, i, N) - seed!(xdual, x, i) + unseed!(xdual, x, i) end # do final chunk diff --git a/src/jacobian.jl b/src/jacobian.jl index b8ce58fb..5f3a79fb 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -191,7 +191,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex $(result_definition) out_reshaped = reshape_jacobian(result, ydual, xdual) extract_jacobian_chunk!(T, out_reshaped, ydual, 1, N) - seed!(xdual, x, 1) + unseed!(xdual, x, 1) # do middle chunks for c in middlechunks @@ -199,7 +199,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex seed!(xdual, x, i, seeds) $(compute_ydual) extract_jacobian_chunk!(T, out_reshaped, ydual, i, N) - seed!(xdual, x, i) + unseed!(xdual, x, i) end # do final chunk @@ -216,7 +216,7 @@ end @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(quote xdual = cfg.duals - seed!(xdual, x) + unseed!(xdual, x) end, :(ydual = f(xdual)), :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)), @@ -226,9 +226,9 @@ end @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(quote ydual, xdual = cfg.duals - seed!(xdual, x) + unseed!(xdual, x) end, - :(f!(seed!(ydual, y), xdual)), + :(f!(unseed!(ydual, y), xdual)), :(result = similar(y, length(y), xlen)), :(map!(d -> value(T,d), y, ydual)))) end @@ -236,7 +236,7 @@ end @eval function chunk_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(quote xdual = cfg.duals - seed!(xdual, x) + unseed!(xdual, x) end, :(ydual = f(xdual)), :(), @@ -246,9 +246,9 @@ end @eval function chunk_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(quote ydual, xdual = cfg.duals - seed!(xdual, x) + unseed!(xdual, x) end, - :(f!(seed!(ydual, y), xdual)), + :(f!(unseed!(ydual, y), xdual)), :(), :(extract_value!(T, result, y, ydual)))) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index af8d6e77..d75672b2 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -7,22 +7,23 @@ include(joinpath(dirname(@__FILE__), "utils.jl")) convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,Float64,8},4},2}, 1.3) -@testset "Test seed! allocations" begin +@testset "Test seed!/unseed! allocations" begin x = rand(1000) cfg = ForwardDiff.GradientConfig(nothing, x) duals = cfg.duals seeds = cfg.seeds - seed = cfg.seeds[1] allocs_seed!(args...) = @allocated ForwardDiff.seed!(args...) allocs_seed!(duals, x, seeds) @test iszero(allocs_seed!(duals, x, seeds)) - allocs_seed!(duals, x, seed) - @test iszero(allocs_seed!(duals, x, seed)) allocs_seed!(duals, x, 1, seeds) @test iszero(allocs_seed!(duals, x, 1, seeds)) - allocs_seed!(duals, x, 1, seed) - @test iszero(allocs_seed!(duals, x, 1, seed)) + + allocs_unseed!(args...) = @allocated ForwardDiff.unseed!(args...) + allocs_unseed!(duals, x) + @test iszero(allocs_unseed!(duals, x)) + allocs_unseed!(duals, x, 1) + @test iszero(allocs_unseed!(duals, x, 1)) allocs_convert_test_574() = @allocated convert_test_574() allocs_convert_test_574() From 09dd30a9ea9b3c0d9966ee9b354bb6ae77f01aa0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 3 Aug 2026 21:42:36 +0200 Subject: [PATCH 3/4] Bump version from 1.4.3 to 1.4.4 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 747a19ec..97139e4c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ForwardDiff" uuid = "f6369f11-7733-5829-9624-2563aa707210" -version = "1.4.3" +version = "1.4.4" [deps] CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950" From a337ee6658ed2a26bfa7251f1580da59c0b36625 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 3 Aug 2026 23:23:57 +0200 Subject: [PATCH 4/4] Name the zero-partials seeding by mechanism, and partition the buffer init Keeps the O(n^2/N) fix from the previous commits, but revises the API around it. Rename `unseed!` to `seed_zero_partials!`. Of the 14 call sites the rename touched, 10 are initializing a work buffer that was never seeded -- `cfg.duals` is raw `similar` memory, and `ydual` in the `f!` paths never carried a perturbation -- and only 4 are clearing a chunk. `unseed!` names the minority case and asserts a prior state that usually does not exist; `seed_zero_partials!` names what the function does (write `x`'s values, zero the perturbations), which holds at every site. Give the windowed method an explicit `count = N` rather than an implicit `take(..., N)`, mirroring the `chunksize = N` argument that `seed!(duals, x, index, seeds, chunksize)` already has, so the seeding and zeroing paths take the same window arguments. Factor the shared loop into `_seed_zero_partials!`. The two public methods differed only in the index iterator they pass. Initialize the chunk-mode work buffer as a disjoint partition: seed chunk 1, then zero only the untouched tail (`N + 1`, `xlen - N`), instead of clearing the whole array and overwriting the first N elements. `xlen > N` always holds in chunk mode, since `chunksize(cfg) == structural_length(x)` routes to vector mode, so the two windows are disjoint and together cover `1:xlen` -- the buffer is still fully written before the first `f` call, which is what makes `cfg` reuse safe. This also lets the four `work_array_definition` quotes in jacobian.jl drop to just unpacking `cfg.duals`, removing a line duplicated 4x, and gives `count` an in-package caller. Tests. Add test/SeedTest.jl: the windowed clear is only ever called on a chunk that was just seeded, so clearing too much is harmless and no test written against the public API can tell a bounded implementation from an unbounded one -- which is why this regression survived since #739. The new tests pin the window directly, and were validated by mutation: restoring the unbounded sweep fails 4 of them. Expected structural index sets are written out by hand and `structural_eachindex` is pinned to them once, so a bug in that iterator cannot hide inside the assertions that depend on it; `values_match` compares over `eachindex(x)` for the same reason. Extend the allocation test with the 4-arg form, where `count` is a runtime value and so catches an inference regression at the `_seed_zero_partials!` boundary that the `count`-defaulting forms could hide. Consolidate the second BigFloat block in test/JacobianTest.jl into a loop over the position of the unassigned entry. Existing coverage only ever placed it in the last chunk, which is never cleared, leaving the `Base._unsetindex!` branch of the windowed path unreached; `hole = 5` puts it in a middle chunk. Co-Authored-By: Claude Opus 5 (1M context) --- src/apiutils.jl | 40 +++++++----------- src/derivative.jl | 4 +- src/gradient.jl | 9 ++-- src/jacobian.jl | 32 +++++--------- test/AllocationsTest.jl | 16 ++++--- test/JacobianTest.jl | 33 +++++++++------ test/SeedTest.jl | 93 +++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 5 +++ 8 files changed, 163 insertions(+), 69 deletions(-) create mode 100644 test/SeedTest.jl diff --git a/src/apiutils.jl b/src/apiutils.jl index 4617c88e..0615fdb3 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -27,7 +27,7 @@ end function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig, y, x) where {F} ydual, xdual = cfg.duals seed!(xdual, x, cfg.seeds) - unseed!(ydual, y) + seed_zero_partials!(ydual, y) f!(ydual, xdual) return ydual end @@ -70,32 +70,24 @@ function structural_eachindex(x::Diagonal, y::AbstractArray) return diagind(x) end -# Copies the values of `x` into `duals` with zero partials, i.e. removes any seeds -# `duals` is currently carrying. -function unseed!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} - seed = zero(Partials{N,V}) - if isbitstype(V) - for idx in structural_eachindex(duals, x) - duals[idx] = Dual{T,V,N}(x[idx], seed) - end - else - for idx in structural_eachindex(duals, x) - if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seed) - else - Base._unsetindex!(duals, idx) - end - end - end - return duals +# Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is +# currently carrying and to initialize a freshly allocated work buffer, whose elements must all be +# written before the target function reads them. +seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} = + _seed_zero_partials!(duals, x, structural_eachindex(duals, x)) + +# Zeroes the partials of `count` elements starting at structural position `index`. Chunk mode only +# needs to clear the chunk it just seeded, so writing through to the end of the array would be O(n) +# redundant work per chunk, i.e. O(n^2/N) per sweep. `count` mirrors the `chunksize` argument of +# `seed!(duals, x, index, seeds, chunksize)`. +function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, index, + count = N) where {T,V,N} + idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), count) + return _seed_zero_partials!(duals, x, idxs) end -# Unseeds at most `N` elements starting at `index`: chunk mode only ever needs to clear -# the N-wide chunk it just seeded, so writing through to the end of the array would be -# O(n) redundant work per chunk (O(n^2) per sweep). -function unseed!(duals::AbstractArray{Dual{T,V,N}}, x, index) where {T,V,N} +function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N} seed = zero(Partials{N,V}) - idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), N) if isbitstype(V) for idx in idxs duals[idx] = Dual{T,V,N}(x[idx], seed) diff --git a/src/derivative.jl b/src/derivative.jl index d9fb355a..0c8a6c05 100644 --- a/src/derivative.jl +++ b/src/derivative.jl @@ -27,7 +27,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba require_one_based_indexing(y) CHK && checktag(T, f!, x) ydual = cfg.duals - unseed!(ydual, y) + seed_zero_partials!(ydual, y) f!(ydual, Dual{T}(x, one(x))) map!(value, y, ydual) return extract_derivative(T, ydual) @@ -65,7 +65,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba result isa DiffResult ? require_one_based_indexing(y) : require_one_based_indexing(result, y) CHK && checktag(T, f!, x) ydual = cfg.duals - unseed!(ydual, y) + seed_zero_partials!(ydual, y) f!(ydual, Dual{T}(x, one(x))) result = extract_value!(T, result, y, ydual) result = extract_derivative!(T, result, ydual) diff --git a/src/gradient.jl b/src/gradient.jl index d76a476e..a5ef3dac 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -127,14 +127,15 @@ function chunk_mode_gradient_expr(result_definition::Expr) # seed work vectors xdual = cfg.duals seeds = cfg.seeds - unseed!(xdual, x) - # do first chunk manually to calculate output type + # do first chunk manually to calculate output type. Seeding the first chunk and zeroing the + # remaining elements partitions `xdual`, so every element is initialized exactly once. seed!(xdual, x, 1, seeds) + seed_zero_partials!(xdual, x, N + 1, xlen - N) ydual = f(xdual) $(result_definition) extract_gradient_chunk!(T, result, ydual, 1, N) - unseed!(xdual, x, 1) + seed_zero_partials!(xdual, x, 1) # do middle chunks for c in middlechunks @@ -142,7 +143,7 @@ function chunk_mode_gradient_expr(result_definition::Expr) seed!(xdual, x, i, seeds) ydual = f(xdual) extract_gradient_chunk!(T, result, ydual, i, N) - unseed!(xdual, x, i) + seed_zero_partials!(xdual, x, i) end # do final chunk diff --git a/src/jacobian.jl b/src/jacobian.jl index 5f3a79fb..f14a6a7b 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -184,14 +184,16 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex $(work_array_definition) seeds = cfg.seeds - # do first chunk manually to calculate output type + # do first chunk manually to calculate output type. Seeding the first chunk and zeroing the + # remaining elements partitions `xdual`, so every element is initialized exactly once. seed!(xdual, x, 1, seeds) + seed_zero_partials!(xdual, x, N + 1, xlen - N) $(compute_ydual) ydual isa AbstractArray || throw(JACOBIAN_ERROR) $(result_definition) out_reshaped = reshape_jacobian(result, ydual, xdual) extract_jacobian_chunk!(T, out_reshaped, ydual, 1, N) - unseed!(xdual, x, 1) + seed_zero_partials!(xdual, x, 1) # do middle chunks for c in middlechunks @@ -199,7 +201,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex seed!(xdual, x, i, seeds) $(compute_ydual) extract_jacobian_chunk!(T, out_reshaped, ydual, i, N) - unseed!(xdual, x, i) + seed_zero_partials!(xdual, x, i) end # do final chunk @@ -214,41 +216,29 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex end @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} - $(jacobian_chunk_mode_expr(quote - xdual = cfg.duals - unseed!(xdual, x) - end, + $(jacobian_chunk_mode_expr(:(xdual = cfg.duals), :(ydual = f(xdual)), :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)), :())) end @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} - $(jacobian_chunk_mode_expr(quote - ydual, xdual = cfg.duals - unseed!(xdual, x) - end, - :(f!(unseed!(ydual, y), xdual)), + $(jacobian_chunk_mode_expr(:((ydual, xdual) = cfg.duals), + :(f!(seed_zero_partials!(ydual, y), xdual)), :(result = similar(y, length(y), xlen)), :(map!(d -> value(T,d), y, ydual)))) end @eval function chunk_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} - $(jacobian_chunk_mode_expr(quote - xdual = cfg.duals - unseed!(xdual, x) - end, + $(jacobian_chunk_mode_expr(:(xdual = cfg.duals), :(ydual = f(xdual)), :(), :(extract_value!(T, result, ydual)))) end @eval function chunk_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} - $(jacobian_chunk_mode_expr(quote - ydual, xdual = cfg.duals - unseed!(xdual, x) - end, - :(f!(unseed!(ydual, y), xdual)), + $(jacobian_chunk_mode_expr(:((ydual, xdual) = cfg.duals), + :(f!(seed_zero_partials!(ydual, y), xdual)), :(), :(extract_value!(T, result, y, ydual)))) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index d75672b2..94e7cddd 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -7,7 +7,7 @@ include(joinpath(dirname(@__FILE__), "utils.jl")) convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,Float64,8},4},2}, 1.3) -@testset "Test seed!/unseed! allocations" begin +@testset "Test seed!/seed_zero_partials! allocations" begin x = rand(1000) cfg = ForwardDiff.GradientConfig(nothing, x) duals = cfg.duals @@ -19,11 +19,15 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F allocs_seed!(duals, x, 1, seeds) @test iszero(allocs_seed!(duals, x, 1, seeds)) - allocs_unseed!(args...) = @allocated ForwardDiff.unseed!(args...) - allocs_unseed!(duals, x) - @test iszero(allocs_unseed!(duals, x)) - allocs_unseed!(duals, x, 1) - @test iszero(allocs_unseed!(duals, x, 1)) + # the 4-arg form passes `count` as a runtime value, so it catches an inference regression at the + # `_seed_zero_partials!` boundary that the forms defaulting `count` to `N` could hide + allocs_szp!(args...) = @allocated ForwardDiff.seed_zero_partials!(args...) + allocs_szp!(duals, x) + @test iszero(allocs_szp!(duals, x)) + allocs_szp!(duals, x, 1) + @test iszero(allocs_szp!(duals, x, 1)) + allocs_szp!(duals, x, 1, 4) + @test iszero(allocs_szp!(duals, x, 1, 4)) allocs_convert_test_574() = @allocated convert_test_574() allocs_convert_test_574() diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index 9cc5024c..b6d36180 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -298,18 +298,27 @@ end @test res == I end - # Unassigned (but unused) entry in the input and unassigned entries in the output - resize!(x, 10) - f = (y, x) -> copyto!(y, 1, x, 1, 9) - for chunksize in (1, 2, 10) - y = similar(x, 9) - @test all(i -> !isassigned(y, i), eachindex(y)) - cfg = ForwardDiff.JacobianConfig(f, y, x, ForwardDiff.Chunk{chunksize}()) - res = ForwardDiff.jacobian(f, y, x, cfg) - @test y == x[1:(end-1)] - @test res isa Matrix{BigFloat} - @test res[:, 1:(end-1)] == I - @test all(iszero, res[:, end]) + # Unassigned (but unused) entry in the input and unassigned entries in the output. `hole` is + # varied so the unassigned entry lands in a middle chunk as well as in the last one: only the + # former reaches the `Base._unsetindex!` branch of the windowed seeding path, since the last + # chunk is never cleared. + @testset "unassigned input entry at $hole" for hole in (5, 10) + x = Vector{BigFloat}(undef, 10) + for i in eachindex(x) + i == hole || (x[i] = BigFloat(i)) + end + used = [i for i in eachindex(x) if i != hole] + f = (y, x) -> (for (k, i) in enumerate(used); y[k] = x[i]; end; y) + for chunksize in (1, 2, 10) + y = similar(x, 9) + @test all(i -> !isassigned(y, i), eachindex(y)) + cfg = ForwardDiff.JacobianConfig(f, y, x, ForwardDiff.Chunk{chunksize}()) + res = ForwardDiff.jacobian(f, y, x, cfg) + @test y == x[used] + @test res isa Matrix{BigFloat} + @test res[:, used] == I + @test all(iszero, res[:, hole]) + end end end diff --git a/test/SeedTest.jl b/test/SeedTest.jl new file mode 100644 index 00000000..02b821c3 --- /dev/null +++ b/test/SeedTest.jl @@ -0,0 +1,93 @@ +module SeedTest + +import ForwardDiff +using ForwardDiff: Partials +using LinearAlgebra +using Test + +include("utils.jl") + +# The windowed `seed_zero_partials!` is only ever called to clear a chunk that was just seeded, so +# clearing *too much* is harmless and no test written against the public API can distinguish a +# correctly bounded implementation from an unbounded one. These tests pin the window down directly: +# they seed every structural position with a marker whose partials are all nonzero, clear a window, +# and check exactly which positions lost their marker. +# +# The expected structural index sets are written out by hand rather than obtained from +# `structural_eachindex`, so a bug in that iterator cannot hide inside the assertions depending on +# it; one test ties the two together. Order is significant: `index` and `count` are positions along +# the sequence, not array indices. The sets are heterogeneous by design — `Vector` and `Diagonal` +# enumerate linear indices (the latter via `diagind`), `UpperTriangular` enumerates `CartesianIndex` +# in column-major order. +const SEED_CASES = ( + (rand(10), collect(1:10)), + (UpperTriangular(rand(5, 5)), [CartesianIndex(i, j) for j in 1:5 for i in 1:j]), + (Diagonal(rand(6, 6)), collect(1:7:36)), +) + +# Positions within `sidx` whose partials are zero. +zeroed_positions(duals, sidx) = + [i for (i, idx) in enumerate(sidx) if iszero(ForwardDiff.partials(duals[idx]))] + +# Compares over *every* index of `x`, not just the structural ones, so a bug misplacing values +# outside the structural set is visible. Off-structure reads are safe: the wrapper types return +# `zero(Dual)` without touching the (uninitialized) parent storage. +values_match(duals, x) = all(idx -> ForwardDiff.value(duals[idx]) == x[idx], eachindex(x)) + +function fill_marker!(duals, x, sidx, marker) + D = eltype(duals) + for idx in sidx + duals[idx] = D(x[idx], marker) + end + return duals +end + +@testset "seed_zero_partials!: $(nameof(typeof(x)))" for (x, sidx) in SEED_CASES + cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{3}()) + duals, seeds = cfg.duals, cfg.seeds + N = ForwardDiff.npartials(eltype(duals)) + marker = Partials(ntuple(i -> Float64(i), N)) + nstruct = length(sidx) + + # everything below counts positions along `sidx`, so pin it to the implementation once + @test collect(ForwardDiff.structural_eachindex(duals, x)) == sidx + @test ForwardDiff.structural_length(x) == nstruct + + # `count` defaults to N + fill_marker!(duals, x, sidx, marker) + ForwardDiff.seed_zero_partials!(duals, x, 4) + @test zeroed_positions(duals, sidx) == collect(4:(4 + N - 1)) + @test values_match(duals, x) + + # an explicit `count` narrows the window; a `count` overrunning the end is clamped by + # `Iterators.take` rather than throwing; a zero-width window is a no-op, which is what makes + # `xlen - N` safe as the `count` of chunk mode's tail clear + @testset "index=$index count=$count" for (index, count, expected) in + ((4, 2, 4:5), + (nstruct - 1, N, (nstruct - 1):nstruct), + (1, 0, 1:0)) + fill_marker!(duals, x, sidx, marker) + ForwardDiff.seed_zero_partials!(duals, x, index, count) + @test zeroed_positions(duals, sidx) == collect(expected) + @test values_match(duals, x) + end + + # the 2-arg form clears every structural position + fill_marker!(duals, x, sidx, marker) + ForwardDiff.seed_zero_partials!(duals, x) + @test zeroed_positions(duals, sidx) == collect(1:nstruct) + @test values_match(duals, x) + + # `seed!` and `seed_zero_partials!` must agree on what "the chunk at `index`" is, or chunk mode + # would leave stale seeds behind. `duals` enters each iteration fully cleared. + @testset "round-trips seed! at index=$index" for index in unique((1, 4, nstruct - N + 1)) + ForwardDiff.seed!(duals, x, index, seeds) + @test zeroed_positions(duals, sidx) == + [i for i in 1:nstruct if !(index <= i <= index + N - 1)] + ForwardDiff.seed_zero_partials!(duals, x, index) + @test zeroed_positions(duals, sidx) == collect(1:nstruct) + @test values_match(duals, x) + end +end + +end # module diff --git a/test/runtests.jl b/test/runtests.jl index 2193242d..e39f5e46 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -48,6 +48,11 @@ Random.seed!(SEED) t = @elapsed include("MiscTest.jl") println("##### done (took $t seconds).") end + @testset "Seeding" begin + println("##### Testing seeding...") + t = @elapsed include("SeedTest.jl") + println("##### done (took $t seconds).") + end @testset "Allocations" begin println("##### Testing allocations...") t = @elapsed include("AllocationsTest.jl")