Extract derivatives at the indices of x, and store those indices in the config

ext/ForwardDiffStaticArraysExt.jl

@@ -5,7 +5,7 @@ using ForwardDiff.LinearAlgebra
 using ForwardDiff.DiffResults
 using ForwardDiff: Dual, partials, npartials, Partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
                    gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
-                   extract_gradient!, extract_jacobian!, extract_value!,
+                   extract_gradient!, extract_jacobian!, extract_value!, structural_indices,
                    vector_mode_gradient, vector_mode_gradient!,
                    vector_mode_jacobian, vector_mode_jacobian!, valtype, value
 using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
@@ -57,7 +57,7 @@ end
 
 @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F}
     T = typeof(Tag(f, eltype(x)))
-    return extract_gradient!(T, result, f(dualize(T, x)))
+    return extract_gradient!(T, result, f(dualize(T, x)), x, structural_indices(x))
 end
 
 # Jacobian
@@ -87,13 +87,13 @@ end
 
 function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
     result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
-    return extract_jacobian!(T, result, ydual, length(x))
+    return extract_jacobian!(T, result, ydual, x, structural_indices(x))
 end
 
 @inline function ForwardDiff.vector_mode_jacobian!(result, f::F, x::StaticArray) where {F}
     T = typeof(Tag(f, eltype(x)))
     ydual = f(dualize(T, x))
-    result = extract_jacobian!(T, result, ydual, length(x))
+    result = extract_jacobian!(T, result, ydual, x, structural_indices(x))
     result = extract_value!(T, result, ydual)
     return result
 end

src/ForwardDiff.jl

@@ -1,7 +1,7 @@
 module ForwardDiff
 
 using DiffRules, DiffResults
-using DiffResults: DiffResult, MutableDiffResult
+using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
 using Preferences
 using Random
 using LinearAlgebra

src/apiutils.jl

@@ -20,14 +20,15 @@ end
 
 function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F}
     xdual = cfg.duals
-    seed!(xdual, x, cfg.seeds)
+    seed!(xdual, x, cfg.indices, cfg.seeds)
     return f(xdual)
 end
 
 function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig, y, x) where {F}
     ydual, xdual = cfg.duals
-    seed!(xdual, x, cfg.seeds)
-    seed_zero_partials!(ydual, y)
+    yindices, xindices = cfg.indices
+    seed!(xdual, x, xindices, cfg.seeds)
+    seed_zero_partials!(ydual, y, yindices)
     f!(ydual, xdual)
     return ydual
 end
@@ -40,104 +41,136 @@ end
     return Expr(:tuple, [:(single_seed(Partials{N,V}, Val{$i}())) for i in 1:N]...)
 end
 
-# Only seed indices that are structurally non-zero
-structural_eachindex(x::AbstractArray) = structural_eachindex(x, x)
-function structural_eachindex(x::AbstractArray, y::AbstractArray)
-    require_one_based_indexing(x, y)
-    eachindex(x, y)
+########################
+# structural positions #
+########################
+
+# The set of linear positions an array stores, so that two arrays can be compared without comparing
+# their positions themselves. `structural_indices` below enumerates the positions of each kind.
+abstract type StructuralKind end
+struct AllEntries <: StructuralKind end
+struct LowerTriangle <: StructuralKind end
+struct UpperTriangle <: StructuralKind end
+struct MainDiagonal <: StructuralKind end
+
+structural_kind(::AbstractArray) = AllEntries()
+structural_kind(::LowerTriangular) = LowerTriangle()
+structural_kind(::UpperTriangular) = UpperTriangle()
+structural_kind(::Diagonal) = MainDiagonal()
+
+# The linear indices of the entries that are seeded, in seeding order. Being linear indices of the
+# array itself, one position indexes the input, the work buffer and the result alike, and is a
+# Jacobian column number as it stands. Configs hold one per work buffer, so that a sweep indexes
+# straight to a chunk.
+function structural_indices(x::AbstractArray)
+    require_one_based_indexing(x)
+    return Base.OneTo(length(x))
 end
-function structural_eachindex(x::UpperTriangular, y::AbstractArray)
-    require_one_based_indexing(x, y)
-    if size(x) != size(y)
-        throw(DimensionMismatch())
-    end
+function structural_indices(x::Diagonal)
+    require_one_based_indexing(x)
+    return diagind(x)
+end
+function structural_indices(x::UpperTriangular)
+    require_one_based_indexing(x)
     n = size(x, 1)
-    return (CartesianIndex(i, j) for j in 1:n for i in 1:j)
+    return [i + (j - 1) * n for j in 1:n for i in 1:j]
 end
-function structural_eachindex(x::LowerTriangular, y::AbstractArray)
-    require_one_based_indexing(x, y)
-    if size(x) != size(y)
-        throw(DimensionMismatch())
-    end
+function structural_indices(x::LowerTriangular)
+    require_one_based_indexing(x)
     n = size(x, 1)
-    return (CartesianIndex(i, j) for j in 1:n for i in j:n)
+    return [i + (j - 1) * n for j in 1:n for i in j:n]
 end
-function structural_eachindex(x::Diagonal, y::AbstractArray)
-    require_one_based_indexing(x, y)
-    if size(x) != size(y)
-        throw(DimensionMismatch())
-    end
-    return diagind(x)
+
+# The positions of the `count` entries starting at structural position `index`. Allocation-free, and
+# a window overrunning the end is a `BoundsError` rather than a silently truncated chunk.
+structural_chunk(indices, index, count) = view(indices, index:(index + count - 1))
+
+# Does an array of kind `outer` store every position of an array of kind `inner`? Add a method here
+# whenever a kind is added above.
+structural_issubset(inner::StructuralKind, outer::StructuralKind) = inner === outer
+structural_issubset(::StructuralKind, ::AllEntries) = true
+structural_issubset(::MainDiagonal, ::LowerTriangle) = true
+structural_issubset(::MainDiagonal, ::UpperTriangle) = true
+
+# Checks that the structural positions of `x` are positions of `y` as well. Being linear indices,
+# they only constrain how many entries `y` has, not its shape.
+function check_structural_indices(x::AbstractArray, y::AbstractArray)
+    require_one_based_indexing(y)
+    structural_issubset(structural_kind(x), structural_kind(y)) || throw(ArgumentError(LazyString(
+        "an array of type ", nameof(typeof(y)), " does not store every entry of an array of type ",
+        nameof(typeof(x)), ": the two are structured differently")))
+    length(x) == length(y) || throw(DimensionMismatch(
+        lazy"expected an array with $(length(x)) elements, got an array with $(length(y)) elements"))
+    return nothing
 end
 
-# 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)
+# The config's positions were built for its work buffer, so they fit `x` exactly when the two are
+# structurally interchangeable. The kind comparison is a compile-time constant, hence free per call.
+function checkstructure(duals::AbstractArray, x::AbstractArray)
+    structural_kind(duals) === structural_kind(x) || throw(ArgumentError(LazyString(
+        "the config was built for an array of type ", nameof(typeof(duals)),
+        " and cannot be used with an array of type ", nameof(typeof(x)),
+        ": the two are structured differently")))
+    return check_structural_indices(duals, x)
 end
 
-function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N}
-    seed = zero(Partials{N,V})
+checkstructure(cfg::AbstractConfig, x) = checkstructure(cfg.duals, x)
+
+# The `f!(y, x)` configs hold a buffer for the output as well, and it is seeded too.
+function checkstructure(cfg::AbstractConfig, y, x)
+    ydual, xdual = cfg.duals
+    checkstructure(ydual, y)
+    return checkstructure(xdual, x)
+end
+
+###########
+# seeding #
+###########
+
+# Mirrors an unassigned entry of `x` into the work buffer. `Base._unsetindex!` is implemented for
+# `Array` alone, a structured wrapper being a view onto a parent with no slot of its own to unset.
+_unsetindex!(duals::Array, idx) = Base._unsetindex!(duals, idx)
+_unsetindex!(duals::AbstractArray, idx) = throw(ArgumentError(LazyString(
+    "cannot differentiate at an input with an unassigned entry at index ", idx,
+    ": that would leave an entry of the ", nameof(typeof(duals)),
+    " work buffer unassigned, which is only possible for an Array")))
+
+# The two seeding operations below differ only in the `Dual` they build for the `i`th position of
+# their window, so they share the walk. The `isbitstype` branch keeps the common case a plain store:
+# `isassigned` is a `try`/`catch` for most array types, and a bits element type is never unassigned.
+@inline function _seed!(make_dual::F, duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {F,T,V,N}
     if isbitstype(V)
-        for idx in idxs
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
+        for (i, idx) in enumerate(idxs)
+            duals[idx] = make_dual(x[idx], i)
         end
     else
-        for idx in idxs
+        for (i, idx) in enumerate(idxs)
             if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seed)
+                duals[idx] = make_dual(x[idx], i)
             else
-                Base._unsetindex!(duals, idx)
+                _unsetindex!(duals, idx)
             end
         end
     end
     return duals
 end
 
-function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
-               seeds::NTuple{N,Partials{N,V}}) where {T,V,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])
-        end
-    else
-        for (i, idx) in zip(1:N, structural_eachindex(duals, x))
-            if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-            else
-                Base._unsetindex!(duals, idx)
-            end
-        end
+function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N}
+    seed = zero(Partials{N,V})
+    return _seed!(duals, x, idxs) do value, _
+        Dual{T,V,N}(value, seed)
     end
-    return duals
 end
 
-function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
+seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, indices, index, count = N) where {T,V,N} =
+    seed_zero_partials!(duals, x, structural_chunk(indices, index, count))
+
+seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices,
+      seeds::NTuple{N,Partials{N,V}}) where {T,V,N} = seed!(duals, x, indices, 1, seeds)
+
+function seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
-    offset = index - 1
-    idxs = Iterators.drop(structural_eachindex(duals, x), offset)
-    if isbitstype(V)
-        for (i, idx) in zip(1:chunksize, idxs)
-            duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-        end
-    else
-        for (i, idx) in zip(1:chunksize, idxs)
-            if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-            else
-                Base._unsetindex!(duals, idx)
-            end
-        end
+    return _seed!(duals, x, structural_chunk(indices, index, chunksize)) do value, i
+        Dual{T,V,N}(value, seeds[i])
     end
-    return duals
 end

src/config.jl

@@ -59,8 +59,11 @@ Base.eltype(cfg::AbstractConfig) = eltype(typeof(cfg))
 # DerivativeConfig #
 ####################
 
-struct DerivativeConfig{T,D} <: AbstractConfig{1}
+# `indices` holds the structural positions of `duals`, so that the sweeps do not derive them per
+# chunk. They are linear indices throughout.
+struct DerivativeConfig{T,D,I} <: AbstractConfig{1}
     duals::D
+    indices::I
 end
 
 """
@@ -84,19 +87,21 @@ function DerivativeConfig(f::F,
                           x::X,
                           tag::T = Tag(f, X)) where {F,X<:Real,Y<:Real,T}
     duals = similar(y, Dual{T,Y,1})
-    return DerivativeConfig{T,typeof(duals)}(duals)
+    indices = structural_indices(duals)
+    return DerivativeConfig{T,typeof(duals),typeof(indices)}(duals, indices)
 end
 
 checktag(::DerivativeConfig{T},f,x) where {T} = checktag(T,f,x)
-Base.eltype(::Type{DerivativeConfig{T,D}}) where {T,D} = eltype(D)
+Base.eltype(::Type{DerivativeConfig{T,D,I}}) where {T,D,I} = eltype(D)
 
 ##################
 # GradientConfig #
 ##################
 
-struct GradientConfig{T,V,N,D} <: AbstractConfig{N}
+struct GradientConfig{T,V,N,D,I} <: AbstractConfig{N}
     seeds::NTuple{N,Partials{N,V}}
     duals::D
+    indices::I
 end
 
 """
@@ -120,19 +125,23 @@ function GradientConfig(f::F,
                         ::T = Tag(f, V)) where {F,V,N,T}
     seeds = construct_seeds(Partials{N,V})
     duals = similar(x, Dual{T,V,N})
-    return GradientConfig{T,V,N,typeof(duals)}(seeds, duals)
+    indices = structural_indices(duals)
+    return GradientConfig{T,V,N,typeof(duals),typeof(indices)}(seeds, duals, indices)
 end
 
 checktag(::GradientConfig{T},f,x) where {T} = checktag(T,f,x)
-Base.eltype(::Type{GradientConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N}
+Base.eltype(::Type{GradientConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,V,N}
 
 ##################
 # JacobianConfig #
 ##################
 
-struct JacobianConfig{T,V,N,D} <: AbstractConfig{N}
+# `indices` mirrors `duals`: the structural positions of the one work buffer of an `f(x)`, or a
+# `(y, x)` pair of position vectors for the two buffers of an `f!(y, x)`.
+struct JacobianConfig{T,V,N,D,I} <: AbstractConfig{N}
     seeds::NTuple{N,Partials{N,V}}
     duals::D
+    indices::I
 end
 
 """
@@ -157,7 +166,8 @@ function JacobianConfig(f::F,
                         ::T = Tag(f, V)) where {F,V,N,T}
     seeds = construct_seeds(Partials{N,V})
     duals = similar(x, Dual{T,V,N})
-    return JacobianConfig{T,V,N,typeof(duals)}(seeds, duals)
+    indices = structural_indices(duals)
+    return JacobianConfig{T,V,N,typeof(duals),typeof(indices)}(seeds, duals, indices)
 end
 
 """
@@ -185,19 +195,20 @@ function JacobianConfig(f::F,
     yduals = similar(y, Dual{T,Y,N})
     xduals = similar(x, Dual{T,X,N})
     duals = (yduals, xduals)
-    return JacobianConfig{T,X,N,typeof(duals)}(seeds, duals)
+    indices = (structural_indices(yduals), structural_indices(xduals))
+    return JacobianConfig{T,X,N,typeof(duals),typeof(indices)}(seeds, duals, indices)
 end
 
 checktag(::JacobianConfig{T},f,x) where {T} = checktag(T,f,x)
-Base.eltype(::Type{JacobianConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N}
+Base.eltype(::Type{JacobianConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,V,N}
 
 #################
 # HessianConfig #
 #################
 
-struct HessianConfig{T,V,N,DG,DJ} <: AbstractConfig{N}
-    jacobian_config::JacobianConfig{T,V,N,DJ}
-    gradient_config::GradientConfig{T,Dual{T,V,N},N,DG}
+struct HessianConfig{T,V,N,DG,DJ,IG,IJ} <: AbstractConfig{N}
+    jacobian_config::JacobianConfig{T,V,N,DJ,IJ}
+    gradient_config::GradientConfig{T,Dual{T,V,N},N,DG,IG}
 end
 
 """
@@ -253,5 +264,5 @@ function HessianConfig(f::F,
 end
 
 checktag(::HessianConfig{T},f,x) where {T} = checktag(T,f,x)
-Base.eltype(::Type{HessianConfig{T,V,N,DG,DJ}}) where {T,V,N,DG,DJ} =
+Base.eltype(::Type{HessianConfig{T,V,N,DG,DJ,IG,IJ}}) where {T,V,N,DG,DJ,IG,IJ} =
     Dual{T,Dual{T,V,N},N}