|
1 |
| -module PreallocationTools |
2 |
| - |
3 |
| -using ForwardDiff, ArrayInterfaceCore, LabelledArrays, Adapt |
4 |
| - |
5 |
| -struct DiffCache{T<:AbstractArray, S<:AbstractArray} |
6 |
| - du::T |
7 |
| - dual_du::S |
8 |
| -end |
9 |
| - |
10 |
| -function DiffCache(u::AbstractArray{T}, siz, chunk_sizes) where {T} |
11 |
| - x = adapt(ArrayInterfaceCore.parameterless_type(u), zeros(T, prod(chunk_sizes .+ 1)*prod(siz))) |
12 |
| - DiffCache(u, x) |
13 |
| -end |
14 |
| - |
15 |
| -""" |
16 |
| -
|
17 |
| -`dualcache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); levels::Int = 1)` |
18 |
| -`dualcache(u::AbstractArray; N::AbstractArray{<:Int})` |
19 |
| -
|
20 |
| -Builds a `DualCache` object that stores both a version of the cache for `u` |
21 |
| -and for the `Dual` version of `u`, allowing use of pre-cached vectors with |
22 |
| -forward-mode automatic differentiation. Supports nested AD via keyword `levels` |
23 |
| -or specifying an array of chunk_sizes. |
24 |
| -
|
25 |
| -""" |
26 |
| -dualcache(u::AbstractArray, N::Int=ForwardDiff.pickchunksize(length(u)); levels::Int = 1) = DiffCache(u, size(u), N*ones(Int, levels)) |
27 |
| -dualcache(u::AbstractArray, N::AbstractArray{<:Int}) = DiffCache(u, size(u), N) |
28 |
| -dualcache(u::AbstractArray, ::Type{Val{N}}; levels::Int = 1) where N = dualcache(u,N;levels) |
29 |
| -dualcache(u::AbstractArray, ::Val{N}; levels::Int = 1) where N = dualcache(u,N;levels) |
30 |
| - |
31 |
| -""" |
32 |
| -
|
33 |
| -`get_tmp(dc::DiffCache, u)` |
34 |
| -
|
35 |
| -Returns the `Dual` or normal cache array stored in `dc` based on the type of `u`. |
36 |
| -
|
37 |
| -""" |
38 |
| -function get_tmp(dc::DiffCache, u::T) where T<:ForwardDiff.Dual |
39 |
| - nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*length(dc.du) |
40 |
| - if nelem > length(dc.dual_du) |
41 |
| - enlargedualcache!(dc, nelem) |
42 |
| - end |
43 |
| - ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) |
44 |
| -end |
45 |
| - |
46 |
| -function get_tmp(dc::DiffCache, u::AbstractArray{T}) where T<:ForwardDiff.Dual |
47 |
| - nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*length(dc.du) |
48 |
| - if nelem > length(dc.dual_du) |
49 |
| - enlargedualcache!(dc, nelem) |
50 |
| - end |
51 |
| - ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) |
52 |
| -end |
53 |
| - |
54 |
| -function get_tmp(dc::DiffCache, u::LabelledArrays.LArray{T,N,D,Syms}) where {T,N,D,Syms} |
55 |
| - nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*length(dc.du) |
56 |
| - if nelem > length(dc.dual_du) |
57 |
| - enlargedualcache!(dc, nelem) |
58 |
| - end |
59 |
| - _x = ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) |
60 |
| - LabelledArrays.LArray{T,N,D,Syms}(_x) |
61 |
| -end |
62 |
| - |
63 |
| -get_tmp(dc::DiffCache, u::Number) = dc.du |
64 |
| -get_tmp(dc::DiffCache, u::AbstractArray) = dc.du |
65 |
| - |
66 |
| -function enlargedualcache!(dc, nelem) #warning comes only once per dualcache. |
67 |
| - chunksize = div(nelem, length(dc.du)) - 1 |
68 |
| - @warn "The supplied dualcache was too small and was enlarged. This incurrs allocations |
69 |
| - on the first call to get_tmp. If few calls to get_tmp occur and optimal performance is essential, |
70 |
| - consider changing 'N'/chunk size of this dualcache to $chunksize." |
71 |
| - resize!(dc.dual_du, nelem) |
72 |
| -end |
73 |
| - |
74 |
| - |
75 |
| -""" |
76 |
| - b = LazyBufferCache(f=identity) |
77 |
| -
|
78 |
| -A lazily allocated buffer object. Given an array `u`, `b[u]` returns an array of the |
79 |
| -same type and size `f(size(u))` (defaulting to the same size), which is allocated as |
80 |
| -needed and then cached within `b` for subsequent usage. |
81 |
| -
|
82 |
| -""" |
83 |
| -struct LazyBufferCache{F<:Function} |
84 |
| - bufs::Dict # a dictionary mapping types to buffers |
85 |
| - sizemap::F |
86 |
| - LazyBufferCache(f::F=identity) where {F<:Function} = new{F}(Dict(), f) # start with empty dict |
87 |
| -end |
88 |
| - |
89 |
| -# override the [] method |
90 |
| -function Base.getindex(b::LazyBufferCache, u::T) where {T<:AbstractArray} |
91 |
| - s = b.sizemap(size(u)) # required buffer size |
92 |
| - buf = get!(b.bufs, (T, s)) do |
93 |
| - similar(u, s) # buffer to allocate if it was not found in b.bufs |
94 |
| - end::T # declare type since b.bufs dictionary is untyped |
95 |
| - return buf |
96 |
| -end |
97 |
| - |
98 |
| -export dualcache, get_tmp, LazyBufferCache |
99 |
| - |
100 |
| -end |
| 1 | +module PreallocationTools |
| 2 | + |
| 3 | +using ForwardDiff, ArrayInterfaceCore, LabelledArrays, Adapt |
| 4 | + |
| 5 | +struct DiffCache{T <: AbstractArray, S <: AbstractArray} |
| 6 | + du::T |
| 7 | + dual_du::S |
| 8 | +end |
| 9 | + |
| 10 | +function DiffCache(u::AbstractArray{T}, siz, chunk_sizes) where {T} |
| 11 | + x = adapt(ArrayInterfaceCore.parameterless_type(u), |
| 12 | + zeros(T, prod(chunk_sizes .+ 1) * prod(siz))) |
| 13 | + DiffCache(u, x) |
| 14 | +end |
| 15 | + |
| 16 | +""" |
| 17 | +
|
| 18 | +`dualcache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); levels::Int = 1)` |
| 19 | +`dualcache(u::AbstractArray; N::AbstractArray{<:Int})` |
| 20 | +
|
| 21 | +Builds a `DualCache` object that stores both a version of the cache for `u` |
| 22 | +and for the `Dual` version of `u`, allowing use of pre-cached vectors with |
| 23 | +forward-mode automatic differentiation. Supports nested AD via keyword `levels` |
| 24 | +or specifying an array of chunk_sizes. |
| 25 | +
|
| 26 | +""" |
| 27 | +function dualcache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); |
| 28 | + levels::Int = 1) |
| 29 | + DiffCache(u, size(u), N * ones(Int, levels)) |
| 30 | +end |
| 31 | +dualcache(u::AbstractArray, N::AbstractArray{<:Int}) = DiffCache(u, size(u), N) |
| 32 | +function dualcache(u::AbstractArray, ::Type{Val{N}}; levels::Int = 1) where {N} |
| 33 | + dualcache(u, N; levels) |
| 34 | +end |
| 35 | +dualcache(u::AbstractArray, ::Val{N}; levels::Int = 1) where {N} = dualcache(u, N; levels) |
| 36 | + |
| 37 | +""" |
| 38 | +
|
| 39 | +`get_tmp(dc::DiffCache, u)` |
| 40 | +
|
| 41 | +Returns the `Dual` or normal cache array stored in `dc` based on the type of `u`. |
| 42 | +
|
| 43 | +""" |
| 44 | +function get_tmp(dc::DiffCache, u::T) where {T <: ForwardDiff.Dual} |
| 45 | + nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du) |
| 46 | + if nelem > length(dc.dual_du) |
| 47 | + enlargedualcache!(dc, nelem) |
| 48 | + end |
| 49 | + ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) |
| 50 | +end |
| 51 | + |
| 52 | +function get_tmp(dc::DiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual} |
| 53 | + nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du) |
| 54 | + if nelem > length(dc.dual_du) |
| 55 | + enlargedualcache!(dc, nelem) |
| 56 | + end |
| 57 | + ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) |
| 58 | +end |
| 59 | + |
| 60 | +function get_tmp(dc::DiffCache, |
| 61 | + u::LabelledArrays.LArray{T, N, D, Syms}) where {T, N, D, Syms} |
| 62 | + nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du) |
| 63 | + if nelem > length(dc.dual_du) |
| 64 | + enlargedualcache!(dc, nelem) |
| 65 | + end |
| 66 | + _x = ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) |
| 67 | + LabelledArrays.LArray{T, N, D, Syms}(_x) |
| 68 | +end |
| 69 | + |
| 70 | +get_tmp(dc::DiffCache, u::Number) = dc.du |
| 71 | +get_tmp(dc::DiffCache, u::AbstractArray) = dc.du |
| 72 | + |
| 73 | +function enlargedualcache!(dc, nelem) #warning comes only once per dualcache. |
| 74 | + chunksize = div(nelem, length(dc.du)) - 1 |
| 75 | + @warn "The supplied dualcache was too small and was enlarged. This incurrs allocations |
| 76 | + on the first call to get_tmp. If few calls to get_tmp occur and optimal performance is essential, |
| 77 | + consider changing 'N'/chunk size of this dualcache to $chunksize." |
| 78 | + resize!(dc.dual_du, nelem) |
| 79 | +end |
| 80 | + |
| 81 | +""" |
| 82 | + b = LazyBufferCache(f=identity) |
| 83 | +
|
| 84 | +A lazily allocated buffer object. Given an array `u`, `b[u]` returns an array of the |
| 85 | +same type and size `f(size(u))` (defaulting to the same size), which is allocated as |
| 86 | +needed and then cached within `b` for subsequent usage. |
| 87 | +
|
| 88 | +""" |
| 89 | +struct LazyBufferCache{F <: Function} |
| 90 | + bufs::Dict # a dictionary mapping types to buffers |
| 91 | + sizemap::F |
| 92 | + LazyBufferCache(f::F = identity) where {F <: Function} = new{F}(Dict(), f) # start with empty dict |
| 93 | +end |
| 94 | + |
| 95 | +# override the [] method |
| 96 | +function Base.getindex(b::LazyBufferCache, u::T) where {T <: AbstractArray} |
| 97 | + s = b.sizemap(size(u)) # required buffer size |
| 98 | + buf = get!(b.bufs, (T, s)) do |
| 99 | + similar(u, s) # buffer to allocate if it was not found in b.bufs |
| 100 | + end::T # declare type since b.bufs dictionary is untyped |
| 101 | + return buf |
| 102 | +end |
| 103 | + |
| 104 | +export dualcache, get_tmp, LazyBufferCache |
| 105 | + |
| 106 | +end |
0 commit comments