Unseed only the chunk, and split unseed! out of seed!

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"

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)
+    seed_zero_partials!(ydual, y)
     f!(ydual, xdual)
     return ydual
 end
@@ -70,14 +70,30 @@ 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. 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
+
+function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N}
+    seed = zero(Partials{N,V})
     if isbitstype(V)
-        for idx in structural_eachindex(duals, x)
+        for idx in idxs
             duals[idx] = Dual{T,V,N}(x[idx], seed)
         end
     else
-        for idx in structural_eachindex(duals, x)
+        for idx in idxs
             if isassigned(x, idx)
                 duals[idx] = Dual{T,V,N}(x[idx], seed)
             else
@@ -106,26 +122,6 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
     return duals
 end
 
-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)
-    if isbitstype(V)
-        for idx in idxs
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
-        end
-    else
-        for idx in idxs
-            if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seed)
-            else
-                Base._unsetindex!(duals, idx)
-            end
-        end
-    end
-    return duals
-end
-
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
     offset = index - 1

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)
+    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
-    seed!(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)

src/gradient.jl

@@ -127,22 +127,23 @@ function chunk_mode_gradient_expr(result_definition::Expr)
         # seed work vectors
         xdual = cfg.duals
         seeds = cfg.seeds
-        seed!(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)
-        seed!(xdual, x, 1)
+        seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
         for c in middlechunks
             i = ((c - 1) * N + 1)
             seed!(xdual, x, i, seeds)
             ydual = f(xdual)
             extract_gradient_chunk!(T, result, ydual, i, N)
-            seed!(xdual, x, i)
+            seed_zero_partials!(xdual, x, i)
         end
 
         # do final chunk

src/jacobian.jl

@@ -184,22 +184,24 @@ 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)
-        seed!(xdual, x, 1)
+        seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
         for c in middlechunks
             i = ((c - 1) * N + 1)
             seed!(xdual, x, i, seeds)
             $(compute_ydual)
             extract_jacobian_chunk!(T, out_reshaped, ydual, i, N)
-            seed!(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
-                                   seed!(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
-                                   seed!(xdual, x)
-                               end,
-                               :(f!(seed!(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
-                                   seed!(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
-                                   seed!(xdual, x)
-                               end,
-                               :(f!(seed!(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

test/AllocationsTest.jl

@@ -7,22 +7,27 @@ 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!/seed_zero_partials! 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))
+
+    # 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()

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
 

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

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