diff --git a/Project.toml b/Project.toml index 0b1d989..c347d41 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "PreallocationTools" uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46" -version = "1.2.2" +version = "1.3.0" authors = ["Chris Rackauckas "] [deps] diff --git a/docs/src/index.md b/docs/src/index.md index 962d471..eaebbdc 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -55,6 +55,20 @@ the `levels` keyword (`N` for each level computed using based on the size of the state vector) or by specifying `N` as an array of integers of chunk sizes, which enables full control of chunk sizes on all differentiation levels. +For resizeable multidimensional workspaces, keep the backing `DiffCache` +one-dimensional and reshape it when a multidimensional cache is needed: + +```julia +storage = DiffCache(zeros(10), 2) +cache = reshape(storage, 2, 5) + +resize!(storage, 12) +cache = reshape(storage, 2, 6) +``` + +The reshaped cache shares the normal cache storage with `storage`, while the +dual cache storage remains vector-backed for `get_tmp`. + ### DiffCache Example 1: Direct Usage ```julia diff --git a/src/PreallocationTools.jl b/src/PreallocationTools.jl index 087a998..61a73fc 100644 --- a/src/PreallocationTools.jl +++ b/src/PreallocationTools.jl @@ -164,6 +164,28 @@ function get_tmp(dc::DiffCache, ::Type{T}) where {T <: Number} end end +""" + reshape(dc::DiffCache, dims...) + reshape(dc::DiffCache, dims) + +Return a `DiffCache` whose normal cache has shape `dims`. + +This is useful for vector-backed caches that need to be resized. Resize the +backing `DiffCache` with `resize!`, then call `reshape` again with the updated +dimensions. The returned cache shares the normal cache storage with `dc`; the +raw dual cache storage remains vector-backed so `get_tmp` can reinterpret it for +the requested automatic differentiation element type. +""" +function Base.reshape(dc::DiffCache, dims::Tuple{Vararg{Integer}}) + shape = map(Int, dims) + return DiffCache(_resizeable_reshape(dc.du, shape), dc.dual_du, dc.any_du, dc.warn_on_resize) +end + +Base.reshape(dc::DiffCache, dims::Integer...) = reshape(dc, dims) + +_resizeable_reshape(a::AbstractVector, shape) = reshape(view(a, :), shape) +_resizeable_reshape(a::AbstractArray, shape) = reshape(a, shape) + get_tmp(dc, u) = dc """ @@ -178,9 +200,18 @@ function _restructure(normal_cache::Array, duals) end function _restructure(normal_cache::AbstractArray, duals) + if _has_vector_view_parent(normal_cache) + return reshape(duals, size(normal_cache)...) + end return ArrayInterface.restructure(normal_cache, duals) end +function _has_vector_view_parent(a) + a isa SubArray{<:Any, 1, <:AbstractVector} && return true + p = applicable(parent, a) ? parent(a) : nothing + return p isa SubArray{<:Any, 1, <:AbstractVector} +end + """ enlargediffcache!(dc::DiffCache, nelem::Integer) @@ -305,15 +336,13 @@ Base.getindex(b::GeneralLazyBufferCache, u::T) where {T} = get_tmp(b, u) function Base.resize!(dc::DiffCache, n::Integer) # Only resize if the array is a vector if dc.du isa AbstractVector + dual_length = length(dc.du) == 0 ? n : cld(length(dc.dual_du), length(dc.du)) * n resize!(dc.du, n) else throw(ArgumentError("resize! is only supported for DiffCache with vector arrays, got $(typeof(dc.du))")) end - # dual_du is often pre-allocated for ForwardDiff dual numbers, - # and may need special handling based on chunk size - # Only resize if it's a vector if dc.dual_du isa AbstractVector - resize!(dc.dual_du, n) + resize!(dc.dual_du, dual_length) end # Always resize the any_du cache resize!(dc.any_du, n) diff --git a/test/core_resizing.jl b/test/core_resizing.jl index 3b34920..6f832ce 100644 --- a/test/core_resizing.jl +++ b/test/core_resizing.jl @@ -1,4 +1,5 @@ -using Test, PreallocationTools, ForwardDiff, LinearAlgebra +using Test, PreallocationTools, ForwardDiff, LinearAlgebra, LabelledArrays, + RecursiveArrayTools #test for downsizing cache randmat = rand(5, 3) @@ -56,7 +57,7 @@ analyticalsolution = [3.0 0; 0 0] # Test resize! functionality for DiffCache @testset "resize! for DiffCache" begin u = rand(10) - dc = DiffCache(u) + dc = DiffCache(u, 2) # Initial size @test length(dc.du) == 10 @@ -65,15 +66,127 @@ analyticalsolution = [3.0 0; 0 0] # Resize to larger resize!(dc, 20) @test length(dc.du) == 20 + @test length(dc.dual_du) == 20 * 3 # Resize to smaller resize!(dc, 5) @test length(dc.du) == 5 + @test length(dc.dual_du) == 5 * 3 # Test that it returns the cache itself @test resize!(dc, 8) === dc end +@testset "reshape standard DiffCache arrays" begin + dc = DiffCache(collect(1.0:10.0), 2) + cache = reshape(dc, 2, 5) + + normal_tmp = get_tmp(cache, 1.0) + @test size(normal_tmp) == (2, 5) + @test vec(normal_tmp) == collect(1.0:10.0) + + normal_tmp[2, 3] = 42.0 + @test dc.du[6] == 42.0 + + dual = ForwardDiff.Dual{Nothing}(1.0, 1.0, 0.0) + dual_tmp = get_tmp(cache, dual) + @test size(dual_tmp) == (2, 5) + @test eltype(dual_tmp) <: ForwardDiff.Dual + dual_tmp[1] = ForwardDiff.Dual{Nothing}(2.0, 3.0, 4.0) + @test dc.dual_du[1:3] == [2.0, 3.0, 4.0] + + dual_type_tmp = get_tmp(cache, typeof(dual)) + @test size(dual_type_tmp) == (2, 5) + @test eltype(dual_type_tmp) == typeof(dual) + dual_type_tmp[2] = ForwardDiff.Dual{Nothing}(5.0, 6.0, 7.0) + @test dc.dual_du[4:6] == [5.0, 6.0, 7.0] + + dual_array_tmp = get_tmp(cache, fill(dual, 2, 5)) + @test size(dual_array_tmp) == (2, 5) + @test eltype(dual_array_tmp) == typeof(dual) + dual_array_tmp[3] = ForwardDiff.Dual{Nothing}(8.0, 9.0, 10.0) + @test dc.dual_du[7:9] == [8.0, 9.0, 10.0] + + resize!(dc, 12) + @test length(dc.dual_du) == 36 + + resized_cache = reshape(dc, (2, 6)) + @test size(get_tmp(resized_cache, 1.0)) == (2, 6) + @test size(get_tmp(resized_cache, dual)) == (2, 6) + + matrix_dc = DiffCache(zeros(2, 5), 2) + matrix_cache = reshape(matrix_dc, 5, 2) + matrix_dual_tmp = get_tmp(matrix_cache, dual) + @test size(matrix_dual_tmp) == (5, 2) + matrix_dual_tmp[1] = ForwardDiff.Dual{Nothing}(5.0, 6.0, 7.0) + @test matrix_dc.dual_du[1:3] == [5.0, 6.0, 7.0] +end + +@testset "reshape vector-backed non-standard DiffCache arrays" begin + dual = ForwardDiff.Dual{Nothing}(1.0, 1.0, 0.0) + + parent = collect(1.0:10.0) + view_dc = DiffCache(view(parent, 1:10), 2) + view_cache = reshape(view_dc, 2, 5) + view_normal_tmp = get_tmp(view_cache, 1.0) + @test size(view_normal_tmp) == (2, 5) + view_normal_tmp[2, 4] = 21.0 + @test view_dc.du[8] == 21.0 + + view_dual_tmp = get_tmp(view_cache, dual) + @test size(view_dual_tmp) == (2, 5) + @test eltype(view_dual_tmp) == typeof(dual) + view_dual_tmp[4] = ForwardDiff.Dual{Nothing}(11.0, 12.0, 13.0) + @test view_dc.dual_du[10:12] == [11.0, 12.0, 13.0] + + lvector_dc = DiffCache(LVector(a = 1.0, b = 2.0, c = 3.0, d = 4.0), 2) + lvector_cache = reshape(lvector_dc, 2, 2) + lvector_normal_tmp = get_tmp(lvector_cache, 1.0) + @test size(lvector_normal_tmp) == (2, 2) + @test vec(lvector_normal_tmp) == [1.0, 2.0, 3.0, 4.0] + lvector_normal_tmp[2, 2] = 31.0 + @test lvector_dc.du[4] == 31.0 + + lvector_dual_tmp = get_tmp(lvector_cache, dual) + @test size(lvector_dual_tmp) == (2, 2) + @test eltype(lvector_dual_tmp) == typeof(dual) + lvector_dual_tmp[2] = ForwardDiff.Dual{Nothing}(14.0, 15.0, 16.0) + @test lvector_dc.dual_du[4:6] == [14.0, 15.0, 16.0] + + partition_dc = DiffCache(ArrayPartition([1.0, 2.0], [3.0, 4.0, 5.0]), 2) + partition_cache = reshape(partition_dc, 5) + partition_normal_tmp = get_tmp(partition_cache, 1.0) + @test size(partition_normal_tmp) == (5,) + @test collect(partition_normal_tmp) == [1.0, 2.0, 3.0, 4.0, 5.0] + partition_normal_tmp[4] = 41.0 + @test partition_dc.du[4] == 41.0 + + partition_dual_tmp = get_tmp(partition_cache, dual) + @test size(partition_dual_tmp) == (5,) + @test eltype(partition_dual_tmp) == typeof(dual) + partition_dual_tmp[3] = ForwardDiff.Dual{Nothing}(17.0, 18.0, 19.0) + @test partition_dc.dual_du[7:9] == [17.0, 18.0, 19.0] +end + +@testset "reshape non-vector non-standard DiffCache arrays" begin + dual = ForwardDiff.Dual{Nothing}(1.0, 1.0, 0.0) + + larray_dc = DiffCache(LArray((2, 2); a = 1.0, b = 2.0, c = 3.0, d = 4.0), 2) + larray_cache = reshape(larray_dc, 4) + larray_normal_tmp = get_tmp(larray_cache, 1.0) + @test size(larray_normal_tmp) == (4,) + @test collect(larray_normal_tmp) == [1.0, 2.0, 3.0, 4.0] + larray_normal_tmp[3] = 51.0 + @test larray_dc.du[3] == 51.0 + + larray_dual_tmp = get_tmp(larray_cache, dual) + @test size(larray_dual_tmp) == (4,) + @test larray_dual_tmp isa LArray + @test eltype(larray_dual_tmp) == typeof(dual) + larray_dual_tmp[1] = ForwardDiff.Dual{Nothing}(20.0, 21.0, 22.0) + @test larray_dual_tmp[1] == ForwardDiff.Dual{Nothing}(20.0, 21.0, 22.0) +end + # Test warn_on_resize option @testset "warn_on_resize option" begin # Default: warn_on_resize = true diff --git a/test/qa/Project.toml b/test/qa/Project.toml index 1dcbb93..b734928 100644 --- a/test/qa/Project.toml +++ b/test/qa/Project.toml @@ -1,4 +1,5 @@ [deps] +AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" @@ -9,6 +10,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" PreallocationTools = {path = "../.."} [compat] +AllocCheck = "0.2" Aqua = "0.8.11" SafeTestsets = "0.1, 1" SciMLTesting = "1.6" diff --git a/test/qa/allocation_tests.jl b/test/qa/allocation_tests.jl new file mode 100644 index 0000000..b21eac2 --- /dev/null +++ b/test/qa/allocation_tests.jl @@ -0,0 +1,31 @@ +using AllocCheck, PreallocationTools, Test + +_qa_reshape_varargs(cache, rows, cols) = reshape(cache, rows, cols) +_qa_reshape_tuple(cache, dims) = reshape(cache, dims) +_qa_get_tmp_number(cache, x) = get_tmp(cache, x) +_qa_get_tmp_type(cache, ::Type{T}) where {T} = get_tmp(cache, T) +_qa_first_tmp_value(cache, x) = get_tmp(cache, x)[1] + +function _qa_test_no_allocs(f, argtypes) + allocs = check_allocs(f, argtypes) + @test isempty(allocs) + return allocs +end + +@testset "reshaped DiffCache" begin + storage = DiffCache(collect(1.0:10.0), 2) + reshaped = reshape(storage, 2, 5) + + parent = collect(1.0:10.0) + view_storage = DiffCache(view(parent, 1:10), 2) + view_reshaped = reshape(view_storage, 2, 5) + + _qa_test_no_allocs(_qa_reshape_varargs, Tuple{typeof(storage), Int, Int}) + _qa_test_no_allocs(_qa_reshape_tuple, Tuple{typeof(storage), Tuple{Int, Int}}) + _qa_test_no_allocs(_qa_reshape_varargs, Tuple{typeof(view_storage), Int, Int}) + + _qa_test_no_allocs(_qa_get_tmp_number, Tuple{typeof(reshaped), Float64}) + _qa_test_no_allocs(_qa_get_tmp_type, Tuple{typeof(reshaped), Type{Float64}}) + _qa_test_no_allocs(_qa_first_tmp_value, Tuple{typeof(reshaped), Float64}) + _qa_test_no_allocs(_qa_get_tmp_number, Tuple{typeof(view_reshaped), Float64}) +end diff --git a/test/qa/qa.jl b/test/qa/qa.jl index 60ec209..bd513ba 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -45,3 +45,7 @@ run_qa( ), ), ) + +@testset "AllocCheck" begin + include("allocation_tests.jl") +end