Merge pull request #21 from longemen3000/tag

add check argument

Author: ChrisRackauckas

README.md

@@ -17,6 +17,7 @@ It's pretty much the same as ForwardDiff.jl except it is threaded. The API is th
 PolyesterForwardDiff.threaded_gradient!(f, dx, x, ForwardDiff.Chunk(8));
 PolyesterForwardDiff.threaded_jacobian!(g, dx, x, ForwardDiff.Chunk(8));
 PolyesterForwardDiff.threaded_jacobian!(g!, y, dx, x, ForwardDiff.Chunk(8));
+PolyesterForwardDiff.threaded_gradient!(f, dx, x, ForwardDiff.Chunk(8),Val{true}()); #To enable tag checking
 ```
 
 ## Citing

src/PolyesterForwardDiff.jl

@@ -14,11 +14,16 @@ function cld_fast(n::T, d::T) where {T}
     x += n != d*x
 end
 
+tag(check::Val{false},f,x) = nothing
+tag(check::Val{true},f,x::AbstractArray{V}) where V = ForwardDiff.Tag(f, V)
+
 store_val!(r::Base.RefValue{T}, x::T) where {T} = (r[] = x)
 store_val!(r::Ptr{T}, x::T) where {T} = Base.unsafe_store!(r, x)
 
-function evaluate_chunks!(f::F, (r,Δx,x), start, stop, ::ForwardDiff.Chunk{C}) where {F,C}
-    cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{C}(), nothing)
+function evaluate_chunks!(f::F, (r,Δx,x), start, stop, ::ForwardDiff.Chunk{C}, check::Val{B}) where {F,C,B}
+    Tag = tag(check, f, x)
+    TagType = typeof(Tag)
+    cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{C}(), Tag)
     N = length(x)
     last_stop = cld_fast(N, C)
     is_last = last_stop == stop
@@ -31,34 +36,35 @@ function evaluate_chunks!(f::F, (r,Δx,x), start, stop, ::ForwardDiff.Chunk{C})
         i = (c-1) * C + 1
         ForwardDiff.seed!(xdual, x, i, seeds)
         ydual = f(xdual)
-        ForwardDiff.extract_gradient_chunk!(Nothing, Δx, ydual, i, C)
+        ForwardDiff.extract_gradient_chunk!(TagType, Δx, ydual, i, C)
         ForwardDiff.seed!(xdual, x, i)
     end
     if is_last
         lastchunksize = C + N - last_stop*C
         lastchunkindex = N - lastchunksize + 1
         ForwardDiff.seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
         _ydual = f(xdual)
-        ForwardDiff.extract_gradient_chunk!(Nothing, Δx, _ydual, lastchunkindex, lastchunksize)
+        ForwardDiff.extract_gradient_chunk!(TagType, Δx, _ydual, lastchunkindex, lastchunksize)
         store_val!(r, ForwardDiff.value(_ydual))
     end
 end
 
-function threaded_gradient!(f::F, Δx::AbstractVector, x::AbstractVector, ::ForwardDiff.Chunk{C}) where {F,C}
+function threaded_gradient!(f::F, Δx::AbstractVector, x::AbstractVector, ::ForwardDiff.Chunk{C}, check = Val{false}()) where {F,C}
     N = length(x)
     d = cld_fast(N, C)
     r = Ref{eltype(Δx)}()
     batch((d,min(d,Threads.nthreads())), r, Δx, x) do rΔxx,start,stop
-        evaluate_chunks!(f, rΔxx, start, stop, ForwardDiff.Chunk{C}())
+        evaluate_chunks!(f, rΔxx, start, stop, ForwardDiff.Chunk{C}(), check)
     end
     r[]
 end
 
 #### in-place jac, out-of-place f ####
 
-function evaluate_jacobian_chunks!(f::F, (Δx,x), start, stop, ::ForwardDiff.Chunk{C}) where {F,C}
-    cfg = ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk{C}(), nothing)
-
+function evaluate_jacobian_chunks!(f::F, (Δx,x), start, stop, ::ForwardDiff.Chunk{C}, check::Val{B}) where {F,C,B}
+    Tag = tag(check, f, x)
+    TagType = typeof(Tag)
+    cfg = ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk{C}(), Tag)
     # figure out loop bounds
     N = length(x)
     last_stop = cld_fast(N, C)
@@ -81,7 +87,7 @@ function evaluate_jacobian_chunks!(f::F, (Δx,x), start, stop, ::ForwardDiff.Chu
 
         # extract part of the Jacobian
         Δx_reshaped = ForwardDiff.reshape_jacobian(Δx, ydual, xdual)
-        ForwardDiff.extract_jacobian_chunk!(Nothing, Δx_reshaped, ydual, i, C)
+        ForwardDiff.extract_jacobian_chunk!(TagType, Δx_reshaped, ydual, i, C)
         ForwardDiff.seed!(xdual, x, i)
     end
 
@@ -98,23 +104,25 @@ function evaluate_jacobian_chunks!(f::F, (Δx,x), start, stop, ::ForwardDiff.Chu
 
         # extract part of the Jacobian
         _Δx_reshaped = ForwardDiff.reshape_jacobian(Δx, _ydual, xdual)
-        ForwardDiff.extract_jacobian_chunk!(Nothing, _Δx_reshaped, _ydual, lastchunkindex, lastchunksize)
+        ForwardDiff.extract_jacobian_chunk!(TagType, _Δx_reshaped, _ydual, lastchunkindex, lastchunksize)
     end
 end
 
-function threaded_jacobian!(f::F, Δx::AbstractArray, x::AbstractArray, ::ForwardDiff.Chunk{C}) where {F,C}
+function threaded_jacobian!(f::F, Δx::AbstractArray, x::AbstractArray, ::ForwardDiff.Chunk{C}, check = Val{false}()) where {F,C}
     N = length(x)
     d = cld_fast(N, C)
     batch((d,min(d,Threads.nthreads())), Δx, x) do Δxx,start,stop
-        evaluate_jacobian_chunks!(f, Δxx, start, stop, ForwardDiff.Chunk{C}())
+        evaluate_jacobian_chunks!(f, Δxx, start, stop, ForwardDiff.Chunk{C}(), check)
     end
     return Δx
 end
 
 # # #### in-place jac, in-place f ####
 
-function evaluate_f_and_jacobian_chunks!(f!::F, (y,Δx,x), start, stop, ::ForwardDiff.Chunk{C}) where {F,C}
-    cfg = ForwardDiff.JacobianConfig(f!, y, x, ForwardDiff.Chunk{C}(), nothing)
+function evaluate_f_and_jacobian_chunks!(f!::F, (y,Δx,x), start, stop, ::ForwardDiff.Chunk{C}, check::Val{B}) where {F,C,B}
+    Tag = tag(check, f!, x)
+    TagType = typeof(Tag)
+    cfg = ForwardDiff.JacobianConfig(f!, y, x, ForwardDiff.Chunk{C}(), Tag)
 
     # figure out loop bounds
     N = length(x)
@@ -138,7 +146,7 @@ function evaluate_f_and_jacobian_chunks!(f!::F, (y,Δx,x), start, stop, ::Forwar
         f!(ForwardDiff.seed!(ydual, y), xdual)
 
         # extract part of the Jacobian
-        ForwardDiff.extract_jacobian_chunk!(Nothing, Δx_reshaped, ydual, i, C)
+        ForwardDiff.extract_jacobian_chunk!(TagType, Δx_reshaped, ydual, i, C)
         ForwardDiff.seed!(xdual, x, i)
     end
 
@@ -154,16 +162,16 @@ function evaluate_f_and_jacobian_chunks!(f!::F, (y,Δx,x), start, stop, ::Forwar
         f!(ForwardDiff.seed!(ydual, y), xdual)
 
         # extract part of the Jacobian
-        ForwardDiff.extract_jacobian_chunk!(Nothing, Δx_reshaped, ydual, lastchunkindex, lastchunksize)
+        ForwardDiff.extract_jacobian_chunk!(TagType, Δx_reshaped, ydual, lastchunkindex, lastchunksize)
         map!(ForwardDiff.value, y, ydual)
     end
 end
 
-function threaded_jacobian!(f!::F, y::AbstractArray, Δx::AbstractArray, x::AbstractArray, ::ForwardDiff.Chunk{C}) where {F,C}
+function threaded_jacobian!(f!::F, y::AbstractArray, Δx::AbstractArray, x::AbstractArray, ::ForwardDiff.Chunk{C},check = Val{false}()) where {F,C}
     N = length(x)
     d = cld_fast(N, C)
     batch((d,min(d,Threads.nthreads())), y, Δx, x) do yΔxx,start,stop
-        evaluate_f_and_jacobian_chunks!(f!, yΔxx, start, stop, ForwardDiff.Chunk{C}())
+        evaluate_f_and_jacobian_chunks!(f!, yΔxx, start, stop, ForwardDiff.Chunk{C}(), check)
     end
     Δx
 end

test/runtests.jl

@@ -27,8 +27,13 @@ g(x) = A*x
 PolyesterForwardDiff.threaded_jacobian!(g, dx, x, ForwardDiff.Chunk(8));
 ForwardDiff.jacobian!(dxref, g, x, ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk(8), nothing));
 @test dx ≈ dxref
+PolyesterForwardDiff.threaded_jacobian!(g, dx, x, ForwardDiff.Chunk(8),Val{true}());
+@test dx ≈ dxref
 
 PolyesterForwardDiff.threaded_jacobian!(g!, y, dx, x, ForwardDiff.Chunk(8));
 ForwardDiff.jacobian!(dxref, g!, yref, x, ForwardDiff.JacobianConfig(g!, yref, x, ForwardDiff.Chunk(8), nothing));
 @test dx ≈ dxref
 @test y ≈ yref
+PolyesterForwardDiff.threaded_jacobian!(g!, y, dx, x, ForwardDiff.Chunk(8),Val{true}());
+@test dx ≈ dxref
+@test y ≈ yref