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" diff --git a/src/apiutils.jl b/src/apiutils.jl index f401a3fc..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) - 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 diff --git a/src/derivative.jl b/src/derivative.jl index b39e2a48..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 - 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) diff --git a/src/gradient.jl b/src/gradient.jl index 0832d354..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 - 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 @@ -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) - seed!(xdual, x, i) + seed_zero_partials!(xdual, x, i) end # do final chunk diff --git a/src/jacobian.jl b/src/jacobian.jl index b8ce58fb..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) - seed!(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) - 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 diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index af8d6e77..94e7cddd 100644 --- a/test/AllocationsTest.jl +++ b/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() 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")