From e4b0b217f4d912a88d010dd0133c07005bdf41eb Mon Sep 17 00:00:00 2001 From: Jack Dunham Date: Mon, 19 Oct 2026 12:35:45 -0400 Subject: [PATCH 1/4] Add `workspace` field to cuTENSOR backend. Selects the workspace for the cuTENSOR contraction. Defaults to WORKSPACE_DEFAULT. --- ext/TensorOperationscuTENSORExt.jl | 3 ++- src/backends.jl | 13 +++++++++---- test/cutensor.jl | 14 ++++++++++++++ 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 764b6ed9..74b53387 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -204,7 +204,8 @@ function TO.tensorcontract!( B, Binds, opB, β, C, Cinds, OP_IDENTITY, - OP_IDENTITY + OP_IDENTITY; + workspace = something(backend.workspace, WORKSPACE_DEFAULT) ) end diff --git a/src/backends.jl b/src/backends.jl index d516e618..a44c6bb0 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -96,7 +96,12 @@ struct TBLISBackend <: AbstractBackend end # CuTENSOR backend #----------------- """ - cuTENSORBackend() -Backend for tensor operations that is based on the NVIDIA cuTENSOR library. -""" -struct cuTENSORBackend <: AbstractBackend end + cuTENSORBackend(; workspace = nothing) +Backend for tensor operations that is based on the NVIDIA cuTENSOR library. `workspace` is +cuTENSOR's workspace preference for contractions, e.g. `cuTENSOR.WORKSPACE_MAX`; `nothing` +selects `cuTENSOR.WORKSPACE_DEFAULT`. +""" +struct cuTENSORBackend{W} <: AbstractBackend + workspace::W +end +cuTENSORBackend(; workspace = nothing) = cuTENSORBackend(workspace) diff --git a/test/cutensor.jl b/test/cutensor.jl index 1708e841..ff2913e3 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -77,6 +77,20 @@ if cuTENSOR.functional() @tensor F2[a, b, c, d, e, f] := CuArray(D)[a, b, c] * conj(CuArray(E)[d, e, f]) @test collect(F2) ≈ F1 end + + @testset "tensorcontract workspace = $workspace" for workspace in + (nothing, cuTENSOR.WORKSPACE_MIN, cuTENSOR.WORKSPACE_DEFAULT, cuTENSOR.WORKSPACE_MAX) + backend = cuTENSORBackend(; workspace) + @test backend.workspace === workspace + A = randn(Float32, (3, 20, 5, 3, 4)) + B = randn(Float32, (5, 6, 20, 3)) + @tensor C1[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g] + @tensor backend = backend allocator = CUDAAllocator() begin + C2[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g] + end + @test C2 isa CuArray + @test collect(C2) ≈ C1 + end end @testset "more complicated expressions" verbose = true begin From cf712741b878c13fd7b9c20f5e274cdf76af19b6 Mon Sep 17 00:00:00 2001 From: Jack Dunham Date: Mon, 19 Oct 2026 13:16:52 -0400 Subject: [PATCH 2/4] Use `@kwdef` for cuTENSORBackend; remove type param. --- src/backends.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/backends.jl b/src/backends.jl index a44c6bb0..0d2a121b 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -101,7 +101,6 @@ Backend for tensor operations that is based on the NVIDIA cuTENSOR library. `wor cuTENSOR's workspace preference for contractions, e.g. `cuTENSOR.WORKSPACE_MAX`; `nothing` selects `cuTENSOR.WORKSPACE_DEFAULT`. """ -struct cuTENSORBackend{W} <: AbstractBackend - workspace::W +@kwdef struct cuTENSORBackend <: AbstractBackend + workspace = nothing end -cuTENSORBackend(; workspace = nothing) = cuTENSORBackend(workspace) From de8f744c1b0d2d916792b4ebfd6167f98fd3466d Mon Sep 17 00:00:00 2001 From: Jack Dunham Date: Tue, 20 Oct 2026 17:02:25 -0400 Subject: [PATCH 3/4] Use `Symbol` as field type with lookup via `select_cutensor_workspace`. --- ext/TensorOperationscuTENSORExt.jl | 16 ++++++++++++++-- src/backends.jl | 7 +++---- test/cutensor.jl | 13 +++++++++++-- 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 74b53387..2d7c2ca2 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -12,7 +12,7 @@ using cuTENSOR: handle, stream using cuTENSOR: cutensorWorksizePreference_t, cutensorAlgo_t, cutensorOperationDescriptor_t, cutensorOperator_t, cutensorJitMode_t, cutensorPlanPreference_t, cutensorComputeDescriptorEnum -using cuTENSOR: WORKSPACE_DEFAULT, ALGO_DEFAULT, JIT_MODE_NONE +using cuTENSOR: WORKSPACE_DEFAULT, WORKSPACE_MAX, WORKSPACE_MIN, ALGO_DEFAULT, JIT_MODE_NONE using cuTENSOR: cutensorCreatePlanPreference, cutensorPlan, CuTensorPlan, CuTensorDescriptor, ModeType @@ -205,10 +205,22 @@ function TO.tensorcontract!( β, C, Cinds, OP_IDENTITY, OP_IDENTITY; - workspace = something(backend.workspace, WORKSPACE_DEFAULT) + workspace = select_cutensor_workspace(backend.workspace) ) end +function select_cutensor_workspace(workspace) + if workspace === :min + return WORKSPACE_MIN + elseif workspace === :max + return WORKSPACE_MAX + elseif workspace === :default + return WORKSPACE_DEFAULT + else + throw(ArgumentError("unknown cuTENSOR workspace option $(repr(workspace)); options are `:min`, `:max` or `:default`")) + end +end + function TO.tensortrace!( C::CuStridedView, A::CuStridedView, p::Index2Tuple, q::Index2Tuple, conjA::Bool, diff --git a/src/backends.jl b/src/backends.jl index 0d2a121b..8148e254 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -96,11 +96,10 @@ struct TBLISBackend <: AbstractBackend end # CuTENSOR backend #----------------- """ - cuTENSORBackend(; workspace = nothing) + cuTENSORBackend(; workspace = :default) Backend for tensor operations that is based on the NVIDIA cuTENSOR library. `workspace` is -cuTENSOR's workspace preference for contractions, e.g. `cuTENSOR.WORKSPACE_MAX`; `nothing` -selects `cuTENSOR.WORKSPACE_DEFAULT`. +cuTENSOR's workspace preference for contractions: `:min`, `:default` or `:max`. """ @kwdef struct cuTENSORBackend <: AbstractBackend - workspace = nothing + workspace::Symbol = :default end diff --git a/test/cutensor.jl b/test/cutensor.jl index ff2913e3..246a582a 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -78,8 +78,7 @@ if cuTENSOR.functional() @test collect(F2) ≈ F1 end - @testset "tensorcontract workspace = $workspace" for workspace in - (nothing, cuTENSOR.WORKSPACE_MIN, cuTENSOR.WORKSPACE_DEFAULT, cuTENSOR.WORKSPACE_MAX) + @testset "tensorcontract workspace = $workspace" for workspace in (:min, :default, :max) backend = cuTENSORBackend(; workspace) @test backend.workspace === workspace A = randn(Float32, (3, 20, 5, 3, 4)) @@ -91,6 +90,16 @@ if cuTENSOR.functional() @test C2 isa CuArray @test collect(C2) ≈ C1 end + + @testset "tensorcontract unknown workspace" begin + A = CuArray(randn(Float32, (3, 4))) + B = CuArray(randn(Float32, (4, 5))) + @test_throws ArgumentError begin + @tensor backend = cuTENSORBackend(; workspace = :maximum) begin + C[a, c] := A[a, b] * B[b, c] + end + end + end end @testset "more complicated expressions" verbose = true begin From cbaf56ff7ba86b41f455086eb383f1e7997cb901 Mon Sep 17 00:00:00 2001 From: Jack Dunham Date: Wed, 21 Oct 2026 11:32:32 -0400 Subject: [PATCH 4/4] Validate workspace field at cuTENSORBackend construction. --- ext/TensorOperationscuTENSORExt.jl | 10 +++++----- src/backends.jl | 10 ++++++++++ test/cutensor.jl | 11 +++-------- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 2d7c2ca2..ad52efd5 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/ext/TensorOperationscuTENSORExt.jl @@ -205,19 +205,19 @@ function TO.tensorcontract!( β, C, Cinds, OP_IDENTITY, OP_IDENTITY; - workspace = select_cutensor_workspace(backend.workspace) + workspace = select_cutensor_workspace(backend) ) end -function select_cutensor_workspace(workspace) +function select_cutensor_workspace(backend::cuTENSORBackend) + workspace = backend.workspace if workspace === :min return WORKSPACE_MIN elseif workspace === :max return WORKSPACE_MAX - elseif workspace === :default - return WORKSPACE_DEFAULT else - throw(ArgumentError("unknown cuTENSOR workspace option $(repr(workspace)); options are `:min`, `:max` or `:default`")) + # Symbol validated at `backend` construction, only remaining option is `:default`. + return WORKSPACE_DEFAULT end end diff --git a/src/backends.jl b/src/backends.jl index 8148e254..a49a0a65 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -102,4 +102,14 @@ cuTENSOR's workspace preference for contractions: `:min`, `:default` or `:max`. """ @kwdef struct cuTENSORBackend <: AbstractBackend workspace::Symbol = :default + function cuTENSORBackend(workspace::Symbol) + if workspace ∉ (:min, :max, :default) + throw( + ArgumentError( + "unknown cuTENSOR workspace option $(repr(workspace)); options are `:min`, `:max` or `:default`" + ) + ) + end + return new(workspace) + end end diff --git a/test/cutensor.jl b/test/cutensor.jl index 246a582a..7cff6c64 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -91,14 +91,9 @@ if cuTENSOR.functional() @test collect(C2) ≈ C1 end - @testset "tensorcontract unknown workspace" begin - A = CuArray(randn(Float32, (3, 4))) - B = CuArray(randn(Float32, (4, 5))) - @test_throws ArgumentError begin - @tensor backend = cuTENSORBackend(; workspace = :maximum) begin - C[a, c] := A[a, b] * B[b, c] - end - end + @testset "cuTENSORBackend unknown workspace" begin + @test_throws ArgumentError cuTENSORBackend(; workspace = :maximum) + @test_throws ArgumentError cuTENSORBackend(:maximum) end end