|
| 1 | +module ForwardDiffGPUArraysCoreExt |
| 2 | + |
| 3 | +using GPUArraysCore: AbstractGPUArray |
| 4 | +using ForwardDiff: ForwardDiff, Dual, Partials, npartials, partials |
| 5 | + |
| 6 | +struct PartialsFn{T,D<:Dual} |
| 7 | + dual::D |
| 8 | +end |
| 9 | +PartialsFn{T}(dual::Dual) where {T} = PartialsFn{T,typeof(dual)}(dual) |
| 10 | + |
| 11 | +(f::PartialsFn{T})(i) where {T} = partials(T, f.dual, i) |
| 12 | + |
| 13 | +function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x, |
| 14 | + seed::Partials{N,V}) where {T,V,N} |
| 15 | + idxs = collect(ForwardDiff.structural_eachindex(duals, x)) |
| 16 | + duals[idxs] .= Dual{T,V,N}.(view(x, idxs), Ref(seed)) |
| 17 | + return duals |
| 18 | +end |
| 19 | + |
| 20 | +function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x, |
| 21 | + seeds::NTuple{N,Partials{N,V}}) where {T,V,N} |
| 22 | + idxs = collect(Iterators.take(ForwardDiff.structural_eachindex(duals, x), N)) |
| 23 | + duals[idxs] .= Dual{T,V,N}.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs))) |
| 24 | + return duals |
| 25 | +end |
| 26 | + |
| 27 | +function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x, index, |
| 28 | + seed::Partials{N,V}) where {T,V,N} |
| 29 | + offset = index - 1 |
| 30 | + idxs = collect(Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset)) |
| 31 | + duals[idxs] .= Dual{T,V,N}.(view(x, idxs), Ref(seed)) |
| 32 | + return duals |
| 33 | +end |
| 34 | + |
| 35 | +function ForwardDiff.seed!(duals::AbstractGPUArray{Dual{T,V,N}}, x, index, |
| 36 | + seeds::NTuple{N,Partials{N,V}}, chunksize) where {T,V,N} |
| 37 | + offset = index - 1 |
| 38 | + idxs = collect( |
| 39 | + Iterators.take(Iterators.drop(ForwardDiff.structural_eachindex(duals, x), offset), chunksize) |
| 40 | + ) |
| 41 | + duals[idxs] .= Dual{T,V,N}.(view(x, idxs), getindex.(Ref(seeds), 1:length(idxs))) |
| 42 | + return duals |
| 43 | +end |
| 44 | + |
| 45 | +# gradient |
| 46 | +function ForwardDiff.extract_gradient!(::Type{T}, result::AbstractGPUArray, |
| 47 | + dual::Dual) where {T} |
| 48 | + fn = PartialsFn{T}(dual) |
| 49 | + idxs = collect(Iterators.take(ForwardDiff.structural_eachindex(result), npartials(dual))) |
| 50 | + result[idxs] .= fn.(1:length(idxs)) |
| 51 | + return result |
| 52 | +end |
| 53 | + |
| 54 | +function ForwardDiff.extract_gradient_chunk!(::Type{T}, result::AbstractGPUArray, dual, |
| 55 | + index, chunksize) where {T} |
| 56 | + fn = PartialsFn{T}(dual) |
| 57 | + offset = index - 1 |
| 58 | + idxs = collect( |
| 59 | + Iterators.take(Iterators.drop(ForwardDiff.structural_eachindex(result), offset), chunksize) |
| 60 | + ) |
| 61 | + result[idxs] .= fn.(1:length(idxs)) |
| 62 | + return result |
| 63 | +end |
| 64 | + |
| 65 | +end |
0 commit comments