diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 764b6ed9..ad52efd5 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 @@ -204,10 +204,23 @@ function TO.tensorcontract!( B, Binds, opB, β, C, Cinds, OP_IDENTITY, - OP_IDENTITY + OP_IDENTITY; + workspace = select_cutensor_workspace(backend) ) end +function select_cutensor_workspace(backend::cuTENSORBackend) + workspace = backend.workspace + if workspace === :min + return WORKSPACE_MIN + elseif workspace === :max + return WORKSPACE_MAX + else + # Symbol validated at `backend` construction, only remaining option is `:default`. + return WORKSPACE_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 d516e618..a49a0a65 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -96,7 +96,20 @@ 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 = :default) +Backend for tensor operations that is based on the NVIDIA cuTENSOR library. `workspace` is +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 1708e841..7cff6c64 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -77,6 +77,24 @@ 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 (:min, :default, :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 + + @testset "cuTENSORBackend unknown workspace" begin + @test_throws ArgumentError cuTENSORBackend(; workspace = :maximum) + @test_throws ArgumentError cuTENSORBackend(:maximum) + end end @testset "more complicated expressions" verbose = true begin