diff --git a/Project.toml b/Project.toml index defbf0a7..4f88083b 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl b/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl index 7cba5cac..2b62f80e 100644 --- a/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl +++ b/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 diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 445d687c..ef3e19e3 100644 --- a/src/TensorAlgebra.jl +++ b/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 diff --git a/src/contract/contractalgorithm.jl b/src/contract/contractalgorithm.jl index 46c929ee..bf7e0d8d 100644 --- a/src/contract/contractalgorithm.jl +++ b/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 diff --git a/test/test_exports.jl b/test/test_exports.jl index 183c3f5f..b45c40cf 100644 --- a/test/test_exports.jl +++ b/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, ] diff --git a/test/test_tensoroperations.jl b/test/test_tensoroperations.jl index 7fb2cc31..68f9c9ee 100644 --- a/test/test_tensoroperations.jl +++ b/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