diff --git a/Project.toml b/Project.toml index 01c6e0f..3b39f9e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorOperations" uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" -version = "5.6.1" +version = "5.6.2" authors = ["Lukas Devos ", "Maarten Van Damme ", "Jutho Haegeman "] [deps] diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl index 1216bd5..468ec69 100644 --- a/src/implementation/ncon.jl +++ b/src/implementation/ncon.jl @@ -42,13 +42,16 @@ function ncon( tensors, network = resolve_traces(tensors, network) tree = order === nothing ? ncontree(network) : indexordertree(network, order) + allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator() + cp = allocator_checkpoint!(allocator) + A, IA, conjA = contracttree(tensors, network, conjlist, tree[1]; kwargs...) B, IB, conjB = contracttree(tensors, network, conjlist, tree[2]; kwargs...) IC = tuple(output′...) C = tensorcontract(IC, A, IA, conjA, B, IB, conjB; kwargs...) - allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator() tree[1] isa Int || tensorfree!(A, allocator) tree[2] isa Int || tensorfree!(B, allocator) + allocator_reset!(allocator, cp) return length(IC) == 0 ? tensorscalar(C) : C end diff --git a/test/allocator.jl b/test/allocator.jl index 20c9e9d..b79dcec 100644 --- a/test/allocator.jl +++ b/test/allocator.jl @@ -99,4 +99,25 @@ using LinearAlgebra tensoralloc(Array{UInt8, 2}, (L + 1, 1), Val(true), buffer) @test length(buffer) > L end + + @testset "ncon does not leak buffer space" begin + buffer = BufferAllocator() + A = randn(5, 5) + B = randn(5, 5) + C = randn(5, 5) + # chain contraction A*B*C -> at least one intermediate tensor allocated as temp + R = ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + @test R ≈ A * B * C + # offset must return to 0: intermediates were reclaimed via allocator_reset! + @test buffer.offset == 0 + @test isempty(buffer) + + # repeated calls must not grow the high-water mark beyond a single call's needs + max1 = buffer.max_offset + for _ in 1:5 + ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + end end