diff --git a/README.md b/README.md index 4c70453..bd19d24 100644 --- a/README.md +++ b/README.md @@ -218,6 +218,9 @@ A `LazyBufferCache` is a `Dict`-like type for the caches which automatically def new cache arrays on demand when they are required. The function `f` maps `size_of_cache = f(size(u))`, which by default creates cache arrays of the same size. +By default the created buffers are not initialized, but a function `initializer!` +can be supplied which is applied to the buffer when it is created, for instance `buf -> fill!(buf, 0.0)`. + Note that `LazyBufferCache` is type-stable and contains no dynamic dispatch. This gives it a ~15ns overhead. The upside of `LazyBufferCache` is that the user does not have to worry about potential issues with chunk sizes and such: `LazyBufferCache` is much easier! diff --git a/docs/src/index.md b/docs/src/index.md index 2e71c21..962d471 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -208,6 +208,9 @@ A `LazyBufferCache` is a `Dict`-like type for the caches, which automatically de new cache arrays on demand when they are required. The function `f` maps `size_of_cache = f(size(u))`, which by default creates cache arrays of the same size. +By default the created buffers are not initialized, but a function `initializer!` +can be supplied which is applied to the buffer when it is created, for instance `buf -> fill!(buf, 0.0)`. + Note that `LazyBufferCache` is type-stable and contains no dynamic dispatch. This gives it a ~15ns overhead. The upside of `LazyBufferCache` is that the user does not have to worry about potential issues with chunk sizes and such: `LazyBufferCache` is much easier! diff --git a/src/PreallocationTools.jl b/src/PreallocationTools.jl index d359f61..ced6354 100644 --- a/src/PreallocationTools.jl +++ b/src/PreallocationTools.jl @@ -197,19 +197,27 @@ end # LazyBufferCache """ - b = LazyBufferCache(f=identity) + b = LazyBufferCache(f = identity; initializer! = identity) A lazily allocated buffer object. Given an array `u`, `b[u]` returns an array of the same type and size `f(size(u))` (defaulting to the same size), which is allocated as needed and then cached within `b` for subsequent usage. +By default the created buffers are not initialized, but a function `initializer!` +can be supplied which is applied to the buffer when it is created, for instance `buf -> fill!(buf, 0.0)`. + Optionally, the size can be explicitly given at calltime using `b[u,s]`, which will return a cache of size `s`. """ -struct LazyBufferCache{F <: Function} +struct LazyBufferCache{F <: Function, I <: Function} bufs::Dict{Any, Any} # a dictionary mapping (type, size) pairs to buffers sizemap::F - LazyBufferCache(f::F = identity) where {F <: Function} = new{F}(Dict(), f) # start with empty dict + initializer!::I + function LazyBufferCache( + f::F = identity; initializer!::I = identity) where { + F <: Function, I <: Function} + new{F, I}(Dict(), f, initializer!) + end # start with empty dict end similar_type(x::AbstractArray, s::Integer) = similar_type(x, (s,)) @@ -222,7 +230,9 @@ end function get_tmp( b::LazyBufferCache, u::T, s = b.sizemap(size(u))) where {T <: AbstractArray} get!(b.bufs, (T, s)) do - similar(u, s) # buffer to allocate if it was not found in b.bufs + buffer = similar(u, s) # buffer to allocate if it was not found in b.bufs + b.initializer!(buffer) + buffer end::similar_type(u, s) # declare type since b.bufs dictionary is untyped end diff --git a/test/lbc.jl b/test/lbc.jl new file mode 100644 index 0000000..8189192 --- /dev/null +++ b/test/lbc.jl @@ -0,0 +1,6 @@ +using PreallocationTools: LazyBufferCache +using Test + +b = LazyBufferCache(Returns(10); initializer! = buf -> fill!(buf, 0)) + +@test b[Float64[]] == zeros(10) diff --git a/test/runtests.jl b/test/runtests.jl index d1315c3..369a492 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,6 +16,7 @@ if GROUP == "All" || GROUP == "Core" @safetestset "DiffCache Resizing" include("core_resizing.jl") @safetestset "DiffCache Nested Duals" include("core_nesteddual.jl") @safetestset "DiffCache Sparsity Support" include("sparsity_support.jl") + @safetestset "LazyBufferCache" include("lbc.jl") @safetestset "GeneralLazyBufferCache" include("general_lbc.jl") end