Add workspace field to cuTENSOR backend.

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,

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

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