Add an allocator to TensorOperationsAlgorithm

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.17.6"
+version = "0.17.7"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl

@@ -1,18 +1,22 @@
 module TensorAlgebraTensorOperationsExt
 
-using TensorAlgebra: TensorAlgebra as TA
+using TensorAlgebra: TensorAlgebra as TA, TensorOperationsAlgorithm
 using TensorOperations: TensorOperations as TO
 
-"""
-    TensorOperationsAlgorithm(backend::AbstractBackend)
-
-Wrapper type for making a TensorOperations backend work as a TensorAlgebra algorithm.
-"""
-struct TensorOperationsAlgorithm{B <: TO.AbstractBackend} <: TA.ContractAlgorithm
-    backend::B
+# `TensorOperationsAlgorithm` stores `nothing` to mean "TensorOperations' default"; resolve
+# those here, where the defaults can be named.
+function backend(algorithm::TensorOperationsAlgorithm)
+    return @something algorithm.backend TO.DefaultBackend()
+end
+function allocator(algorithm::TensorOperationsAlgorithm)
+    return @something algorithm.allocator TO.DefaultAllocator()
 end
 
-TA.ContractAlgorithm(backend::TO.AbstractBackend) = TensorOperationsAlgorithm(backend)
+# Construct via the `ContractAlgorithm` public constructor seam as well.
+TA.ContractAlgorithm(backend::TO.AbstractBackend) = TensorOperationsAlgorithm(; backend)
+function TA.ContractAlgorithm(backend::TO.AbstractBackend, allocator)
+    return TensorOperationsAlgorithm(; backend, allocator)
+end
 
 # Using TensorOperations backends as TensorAlgebra implementations
 # ----------------------------------------------------------------
@@ -31,7 +35,7 @@ function TA.contract(
     α = true
     return TO.tensorcontract(
         a1, permblocks1, conj1, a2, permblocks2, conj2,
-        permblocks_dest, α, algorithm.backend
+        permblocks_dest, α, backend(algorithm), allocator(algorithm)
     )
 end
 
@@ -47,7 +51,7 @@ function TA.contract(
     α = true
     return TO.tensorcontract(
         a1, permblocks1, conj1, a2, permblocks2, conj2,
-        permblocks_dest, α, algorithm.backend
+        permblocks_dest, α, backend(algorithm), allocator(algorithm)
     )
 end
 
@@ -68,7 +72,7 @@ function TA.contractopadd!(
     a2′ = (op2 === identity || op2 === conj) ? a2 : op2.(a2)
     return TO.tensorcontract!(
         a_dest, a1′, permblocks1, conj1, a2′, permblocks2, conj2,
-        permblocks_dest, α, β, algorithm.backend
+        permblocks_dest, α, β, backend(algorithm), allocator(algorithm)
     )
 end
 

src/TensorAlgebra.jl

@@ -9,7 +9,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public biperm, bipartition, cat_similar, concatenate, concatenate!, contractopadd!, data, datatype, directsum, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
+            "public biperm, bipartition, cat_similar, concatenate, concatenate!, ContractAlgorithm, contractopadd!, data, datatype, directsum, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, TensorOperationsAlgorithm, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
         )
     )
 end

src/contract/contractalgorithm.jl

@@ -11,6 +11,18 @@ end
 Matricize(fusion_style) = Matricize(fusion_style, fusion_style, fusion_style)
 Matricize() = Matricize(ReshapeFusion())
 
+"""
+    TensorOperationsAlgorithm(; backend = nothing, allocator = nothing)
+
+Contract using TensorOperations, with `backend` selecting the contraction kernel and
+`allocator` the allocator for temporary tensors (e.g. `TensorOperations.ManualAllocator()`).
+A `nothing` field uses TensorOperations' default. Only usable with TensorOperations loaded.
+"""
+Base.@kwdef struct TensorOperationsAlgorithm{Backend, Allocator} <: ContractAlgorithm
+    backend::Backend = nothing
+    allocator::Allocator = nothing
+end
+
 function select_contract_algorithm(algorithm, a1, a2)
     return error("Not implemented.")
 end

test/test_exports.jl

@@ -39,9 +39,11 @@ using Test: @test, @testset
             exports,
             [
                 :biperm, :bipartition, :cat_similar,
-                :concatenate, :concatenate!, :contractopadd!, :data, :datatype, :directsum,
+                :concatenate, :concatenate!, :ContractAlgorithm, :contractopadd!, :data,
+                :datatype, :directsum,
                 :flattenlinear, :label_type,
                 :matricizeopperm, :permutedims, :permutedims!, :scalar, :similar_map,
+                :TensorOperationsAlgorithm,
                 :to_range, :tr, :tryflattenlinear, :ungrade, :zero!, :scale!,
                 :permuteddims, :PermutedDims,
             ]

test/test_tensoroperations.jl

@@ -1,5 +1,7 @@
-using TensorAlgebra: Matricize
-using TensorOperations: @tensor, ncon, tensorcontract
+using TensorAlgebra:
+    ContractAlgorithm, Matricize, TensorOperationsAlgorithm, contract, contract!
+using TensorOperations:
+    @tensor, DefaultAllocator, DefaultBackend, ManualAllocator, ncon, tensorcontract
 using Test: @inferred, @test, @testset
 
 @testset "tensorcontract" begin
@@ -123,3 +125,35 @@ end
         @test result1 ≈ result2
     end
 end
+
+@testset "TensorOperationsAlgorithm allocator ($T)" for T in elts
+    a1 = randn(T, 4, 5, 3)
+    a2 = randn(T, 3, 6)
+    labels1 = (:i, :j, :k)
+    labels2 = (:k, :l)
+    ref, ref_labels = contract(a1, labels1, a2, labels2)
+
+    @test TensorOperationsAlgorithm() isa ContractAlgorithm
+
+    @testset "allocator = $(nameof(typeof(alloc)))" for alloc in
+        (
+            DefaultAllocator(),
+            ManualAllocator(),
+        )
+        alg = TensorOperationsAlgorithm(; allocator = alloc)
+        c, labels = contract(a1, labels1, a2, labels2; alg)
+        @test labels == ref_labels
+        @test c ≈ ref
+
+        c_dest = similar(ref)
+        contract!(c_dest, ref_labels, a1, labels1, a2, labels2; alg)
+        @test c_dest ≈ ref
+    end
+
+    # The `ContractAlgorithm(backend, allocator)` constructor seam.
+    seam = ContractAlgorithm(DefaultBackend(), ManualAllocator())
+    @test contract(a1, labels1, a2, labels2; alg = seam)[1] ≈ ref
+
+    # `nothing` fields fall back to the TensorOperations defaults.
+    @test contract(a1, labels1, a2, labels2; alg = TensorOperationsAlgorithm())[1] ≈ ref
+end