Fix: ncon should add allocator_checkpoint! and allocator_reset! to avoid leaking buffer space

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 <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd1994@gmail.com>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 
 [deps]

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
 

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