Skip to content

Commit

Permalink
add check argument
Browse files Browse the repository at this point in the history
  • Loading branch information
@longemen3000
longemen3000 authored Apr 12, 2024
1 parent fbad79f commit 959559f
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 19 deletions.
1 change: 1 addition & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
46 changes: 27 additions & 19 deletions src/PolyesterForwardDiff.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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

Expand All @@ -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)
Expand All @@ -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

Expand All @@ -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
Expand Down
3 changes: 3 additions & 0 deletions test/runtests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,11 @@ 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));
PolyesterForwardDiff.threaded_jacobian!(g!, y, dx, x, ForwardDiff.Chunk(8),Val{true}());
@test dx ≈ dxref
@test y ≈ yref

0 comments on commit 959559f

Please sign in to comment.