diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index b67ccbe1..764b6ed9 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -220,7 +220,9 @@ function TO.tensortrace!( # map to reduction operation plan = plan_trace(A, Ainds, opA, C, Cinds, OP_IDENTITY, OP_ADD) - return reduce!(plan, α, A, β, C) + reduce!(plan, α, A, β, C) + CUDACore.unsafe_free!(plan) # the plan-taking `reduce!` does not free the plan + return C end function cuTENSOR.CuTensorDescriptor( diff --git a/test/cutensor.jl b/test/cutensor.jl index 75d18c13..1708e841 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -576,6 +576,28 @@ if cuTENSOR.functional() end end + @testset "tensortrace! does not leak plans" begin + # each leaked plan holds on to a 128 KiB reduction workspace + A = CuArray(randn(Float32, 64, 64, 64, 64)) + C = CUDACore.zeros(Float32, 64, 64) + @tensor C[a, b] = A[a, c, b, c] # warm up + + GC.gc(true) + CUDACore.reclaim() + CUDACore.synchronize() + GC.enable(false) + try + live0 = CUDACore.memory_stats().live + for _ in 1:1000 + @tensor C[a, b] = A[a, c, b, c] + end + CUDACore.synchronize() + @test CUDACore.memory_stats().live - live0 < 2^20 # would be 125 MiB if leaking + finally + GC.enable(true) + end + end + @testset "Issues" verbose = true begin @testset "Issue PR #186" begin # https://github.com/Jutho/TensorOperations.jl/pull/186