From cfff215642e50abdb038ca5ecfed27ab97920ebb Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 10:29:04 -0400 Subject: [PATCH 1/3] Fix plan leak in cuTENSOR `tensortrace!` `plan_trace` builds a `CuTensorPlan` and passes it to the plan-taking `cuTENSOR.reduce!`, which by contract does not free it. The reduction workspace (128 KiB) was therefore only reclaimed by the finalizer, which the GC has no reason to run promptly for device memory: 10^4 traces grew live device memory by 1.22 GiB. After the fix, 9.8 KiB. Also close a matching exception leak in `plan_trace`, where the operation descriptor and plan preference were destroyed after plan creation rather than in a `finally`. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationscuTENSORExt.jl | 22 +++++++++++++++++----- test/cutensor.jl | 26 ++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 5 deletions(-) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index b67ccbe1..925edfd7 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -219,8 +219,17 @@ function TO.tensortrace!( opA = tensorop(A, conjA) # map to reduction operation + # NOTE: the plan-taking `reduce!` does not free the plan, so we have to do it ourselves. + # This mirrors what the planless `cuTENSOR.reduce!` does, and matters here because the + # reduction workspace is sizeable (128 KiB) and would otherwise only be reclaimed by the + # finalizer, which the GC has no reason to run promptly for device memory. plan = plan_trace(A, Ainds, opA, C, Cinds, OP_IDENTITY, OP_ADD) - return reduce!(plan, α, A, β, C) + try + reduce!(plan, α, A, β, C) + finally + CUDACore.unsafe_free!(plan) + end + return C end function cuTENSOR.CuTensorDescriptor( @@ -290,10 +299,13 @@ function plan_trace( plan_pref = Ref{cutensorPlanPreference_t}() cutensorCreatePlanPreference(handle(), plan_pref, algo, jit) - plan = CuTensorPlan(desc[], plan_pref[]; workspacePref = workspace) - cuTENSOR.cutensorDestroyOperationDescriptor(desc[]) - cuTENSOR.cutensorDestroyPlanPreference(plan_pref[]) - return plan + return try + CuTensorPlan(desc[], plan_pref[]; workspacePref = workspace) + finally + # these are no longer needed once the plan exists, and would leak if it does not + cuTENSOR.cutensorDestroyOperationDescriptor(desc[]) + cuTENSOR.cutensorDestroyPlanPreference(plan_pref[]) + end end end diff --git a/test/cutensor.jl b/test/cutensor.jl index 75d18c13..a9d24636 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -576,6 +576,32 @@ if cuTENSOR.functional() end end + @testset "tensortrace! does not leak plans" begin + # `plan_trace` builds a `CuTensorPlan` whose reduction workspace is 128 KiB, and the + # plan-taking `cuTENSOR.reduce!` does not free it. Without an explicit + # `unsafe_free!` the workspace is only reclaimed by the finalizer, which the GC has + # no reason to run promptly for device memory: 10^4 traces leaked over 1 GiB. + 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() + # the leak was exactly 128 KiB per call, i.e. 125 MiB here + @test CUDACore.memory_stats().live - live0 < 2^20 + finally + GC.enable(true) + end + end + @testset "Issues" verbose = true begin @testset "Issue PR #186" begin # https://github.com/Jutho/TensorOperations.jl/pull/186 From 08c9cdec1c37debae562159835aafb3519d10979 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 10:44:18 -0400 Subject: [PATCH 2/3] Trim comments Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationscuTENSORExt.jl | 7 +------ test/cutensor.jl | 8 ++------ 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 925edfd7..4fa9e324 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -219,15 +219,11 @@ function TO.tensortrace!( opA = tensorop(A, conjA) # map to reduction operation - # NOTE: the plan-taking `reduce!` does not free the plan, so we have to do it ourselves. - # This mirrors what the planless `cuTENSOR.reduce!` does, and matters here because the - # reduction workspace is sizeable (128 KiB) and would otherwise only be reclaimed by the - # finalizer, which the GC has no reason to run promptly for device memory. plan = plan_trace(A, Ainds, opA, C, Cinds, OP_IDENTITY, OP_ADD) try reduce!(plan, α, A, β, C) finally - CUDACore.unsafe_free!(plan) + CUDACore.unsafe_free!(plan) # the plan-taking `reduce!` does not free the plan end return C end @@ -302,7 +298,6 @@ function plan_trace( return try CuTensorPlan(desc[], plan_pref[]; workspacePref = workspace) finally - # these are no longer needed once the plan exists, and would leak if it does not cuTENSOR.cutensorDestroyOperationDescriptor(desc[]) cuTENSOR.cutensorDestroyPlanPreference(plan_pref[]) end diff --git a/test/cutensor.jl b/test/cutensor.jl index a9d24636..1708e841 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -577,10 +577,7 @@ if cuTENSOR.functional() end @testset "tensortrace! does not leak plans" begin - # `plan_trace` builds a `CuTensorPlan` whose reduction workspace is 128 KiB, and the - # plan-taking `cuTENSOR.reduce!` does not free it. Without an explicit - # `unsafe_free!` the workspace is only reclaimed by the finalizer, which the GC has - # no reason to run promptly for device memory: 10^4 traces leaked over 1 GiB. + # 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 @@ -595,8 +592,7 @@ if cuTENSOR.functional() @tensor C[a, b] = A[a, c, b, c] end CUDACore.synchronize() - # the leak was exactly 128 KiB per call, i.e. 125 MiB here - @test CUDACore.memory_stats().live - live0 < 2^20 + @test CUDACore.memory_stats().live - live0 < 2^20 # would be 125 MiB if leaking finally GC.enable(true) end From 50fb0a72bf2bbe04260ccc36ad9ddd6dda505bfd Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 10:48:10 -0400 Subject: [PATCH 3/3] Match upstream: drop the try/finally guards cuTENSOR.jl itself uses no `finally` anywhere and accepts leaks on throw, both in `plan_reduction` and in the planless `reduce!`/`contract!`. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationscuTENSORExt.jl | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 4fa9e324..764b6ed9 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -220,11 +220,8 @@ function TO.tensortrace!( # map to reduction operation plan = plan_trace(A, Ainds, opA, C, Cinds, OP_IDENTITY, OP_ADD) - try - reduce!(plan, α, A, β, C) - finally - CUDACore.unsafe_free!(plan) # the plan-taking `reduce!` does not free the plan - end + reduce!(plan, α, A, β, C) + CUDACore.unsafe_free!(plan) # the plan-taking `reduce!` does not free the plan return C end @@ -295,12 +292,10 @@ function plan_trace( plan_pref = Ref{cutensorPlanPreference_t}() cutensorCreatePlanPreference(handle(), plan_pref, algo, jit) - return try - CuTensorPlan(desc[], plan_pref[]; workspacePref = workspace) - finally - cuTENSOR.cutensorDestroyOperationDescriptor(desc[]) - cuTENSOR.cutensorDestroyPlanPreference(plan_pref[]) - end + plan = CuTensorPlan(desc[], plan_pref[]; workspacePref = workspace) + cuTENSOR.cutensorDestroyOperationDescriptor(desc[]) + cuTENSOR.cutensorDestroyPlanPreference(plan_pref[]) + return plan end end