diff --git a/Project.toml b/Project.toml index b0ff2846..55b52b31 100644 --- a/Project.toml +++ b/Project.toml @@ -25,12 +25,14 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" [extensions] TensorOperationsAMDGPUExt = "AMDGPU" TensorOperationsBumperExt = "Bumper" TensorOperationsChainRulesCoreExt = "ChainRulesCore" TensorOperationsMooncakeExt = "Mooncake" +TensorOperationsTBLISExt = "TBLIS" TensorOperationsCUDACoreExt = "CUDACore" TensorOperationsEnzymeExt = "Enzyme" TensorOperationscuTENSORExt = "cuTENSOR" @@ -59,6 +61,7 @@ PtrArrays = "1.2" Random = "1" Strided = "2.6" StridedViews = "0.5" +TBLIS = "0.3" Test = "1" TupleTools = "1.6" VectorInterface = "0.4.1, 0.5, 0.6" @@ -81,8 +84,9 @@ JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" [targets] -test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "ChainRulesCore", "cuRAND", "CUDACore", "cuTENSOR", "Aqua", "Logging", "Bumper", "Mooncake", "Enzyme", "EnzymeTestUtils", "Adapt", "JLArrays", "AMDGPU"] +test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "ChainRulesCore", "cuRAND", "CUDACore", "cuTENSOR", "Aqua", "Logging", "Bumper", "Mooncake", "Enzyme", "EnzymeTestUtils", "Adapt", "JLArrays", "AMDGPU", "TBLIS"] diff --git a/docs/src/man/backends.md b/docs/src/man/backends.md index 9971bd16..af4359c2 100644 --- a/docs/src/man/backends.md +++ b/docs/src/man/backends.md @@ -64,6 +64,7 @@ TensorOperations.BaseCopy TensorOperations.BaseView TensorOperations.StridedNative TensorOperations.StridedBLAS +TensorOperations.TBLISBackend TensorOperations.cuTENSORBackend ``` @@ -74,6 +75,26 @@ On the other hand, the `BaseCopy` and `BaseView` backends are used for arrays th These are designed to be as general as possible, and as a result are not as performant as specific implementations. Nevertheless, they can be useful for debugging purposes or for working with custom tensor types that have limited support for methods outside of `Base`. +The `TBLISBackend` routes the primitive operations through the [TBLIS](https://github.com/devinamatthews/tblis) library. +TBLIS contracts strided tensors in place instead of reshaping them into matrices, and can therefore avoid the intermediate permuted copies that `StridedBLAS` sometimes has to allocate. +It is opt-in, in the sense that loading `TBLIS.jl` does not change the default backend selection, and it is only available through a package extension for [`TBLIS.jl`](https://github.com/QuantumKitHub/TBLIS.jl): + +```julia +using TensorOperations, TBLIS +TBLIS.set_num_threads(8) +@tensor backend = TensorOperations.TBLISBackend() D[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * conj(C[g, f, b]) +``` + +TBLIS requires all tensors in a single operation to share one element type out of `Float32`, `Float64`, `ComplexF32` and `ComplexF64`. +Arguments that do not satisfy this, as well as non-strided arrays, are rejected with an `ArgumentError` instead of being passed on to another backend, so that a contraction which cannot actually reach TBLIS is not silently run somewhere else. + +Note that contracting in place trades throughput for memory rather than being a free win. +With BLAS, TBLIS and `Strided.jl` all given the same number of threads, this backend is roughly on par with `StridedBLAS` for permuted real contractions and slower for other shapes, while allocating no permuted temporaries at all. + +!!! warning + As of `tblis_jll` v1.3, TBLIS has no competitive support for complex element types. + Its complex contraction kernels run an order of magnitude slower than the corresponding BLAS calls, so `StridedBLAS` is the better choice for complex-valued contractions. + Finally, we also provide a `cuTENSORBackend` for use with the `cuTENSOR.jl` library, which is a NVidia GPU-accelerated tensor contraction library. This backend is only available through a package extension for `cuTENSOR`. @@ -89,7 +110,7 @@ Users can also define their own backends, to facilitate experimentation with new This can be done by defining a new type that is a subtype of `AbstractBackend`, and dispatching on this type in the implementation of the primitive tensor operations. In particular, the only required implemented methods are [`tensoradd!`](@ref), [`tensortrace!`](@ref), [`tensorcontract!`](@ref). -For example, [`TensorOperationsTBLIS`](https://github.com/lkdvos/TensorOperationsTBLIS.jl) is a wrapper that provides a backend for tensor contractions using the [TBLIS](https://github.com/devinamatthews/tblis) library. +For example, the `TBLISBackend` above is implemented in exactly this way, as a package extension that only adds methods for these three functions. ## Allocators diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl new file mode 100644 index 00000000..8c1adaf4 --- /dev/null +++ b/ext/TensorOperationsTBLISExt.jl @@ -0,0 +1,221 @@ +module TensorOperationsTBLISExt + +using TensorOperations +using TensorOperations: TensorOperations as TO +using TensorOperations: TBLISBackend, DefaultAllocator, Index2Tuple +using TensorOperations: StridedView, isstrided +using TensorOperations: argcheck_tensoradd, dimcheck_tensoradd, + argcheck_tensortrace, dimcheck_tensortrace, + argcheck_tensorcontract, dimcheck_tensorcontract +using TensorOperations: add_labels, trace_labels, contract_labels +using TensorOperations: tensoralloc_add, tensorfree! + +using TBLIS +using TBLIS: len_type, stride_type, tblis_tensor + +const SV = StridedView + +const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} + +#------------------------------------------------------------------------------------------- +# Wrapping Julia arrays as TBLIS tensors +#------------------------------------------------------------------------------------------- +for (T, init) in ( + (:Float32, :tblis_init_tensor_scaled_s), + (:Float64, :tblis_init_tensor_scaled_d), + (:ComplexF32, :tblis_init_tensor_scaled_c), + (:ComplexF64, :tblis_init_tensor_scaled_z), + ) + @eval function init_tensor!( + p::Ptr{tblis_tensor}, A::StridedView{$T, N}, α::$T, + len::Vector{len_type}, stride::Vector{stride_type} + ) where {N} + return TBLIS.$init(p, α, Cuint(N), pointer(len), pointer(A), pointer(stride)) + end +end + +isconj(A::StridedView{T}, conjA::Bool) where {T} = T <: Complex && (conjA ⊻ (A.op === conj)) +tblis_dims(A::StridedView) = (collect(len_type, size(A)), collect(stride_type, strides(A))) + +""" + tblis_tensor(A::StridedView, α, len, stride, conj) -> Ref{TBLIS.tblis_tensor} + +Descriptor for `α * A`, conjugated when `conj` is set, using `len` and `stride` as the +buffers handed to TBLIS. + +The descriptor only stores raw pointers into `A`, `len` and `stride`, so all three, along +with the returned `Ref`, have to be kept alive by the caller for as long as TBLIS may access +them. Note that `conj` is the *total* conjugation applied to the data of `A`, as computed by +[`isconj`](@ref), not the flag the caller was handed. +""" +function tblis_tensor( + A::StridedView{T, N}, α::T, + len::Vector{len_type}, stride::Vector{stride_type}, conj::Bool + ) where {T <: TBLISFloat, N} + ref = Ref{tblis_tensor}() + GC.@preserve A len stride ref begin + p = Base.unsafe_convert(Ptr{tblis_tensor}, ref) + init_tensor!(p, A, α, len, stride) + conj && setproperty!(p, :conj, Cint(1)) + end + return ref +end + +labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) + +#------------------------------------------------------------------------------------------- +# Argument checking +#------------------------------------------------------------------------------------------- +@noinline function throw_eltype(f, tensors) + return throw( + ArgumentError( + LazyString( + "TBLISBackend requires all tensors of ", f, " to share a single element ", + "type out of Float32, Float64, ComplexF32 and ComplexF64, got ", + join(eltype.(tensors), ", ") + ) + ) + ) +end + +@noinline function throw_strided(f, tensors) + types = join(typeof.(tensors), ", ") + return throw(ArgumentError(lazy"TBLISBackend requires strided arrays for $f, got $types")) +end + +@noinline throw_conj_output(f) = throw( + ArgumentError(lazy"TBLISBackend cannot write into a conjugated view in $f") +) + +function check_arguments(f, C::AbstractArray, As::AbstractArray...) + tensors = (C, As...) + T = eltype(C) + (T <: TBLISFloat && all(A -> eltype(A) === T, As)) || throw_eltype(f, tensors) + all(isstrided, tensors) || throw_strided(f, tensors) + # `tblis_tensor_add` applies the flag of `C` when reading `β * C` but not when writing back + isconj(SV(C), false) && throw_conj_output(f) + return nothing +end + +#------------------------------------------------------------------------------------------- +# Operations +#------------------------------------------------------------------------------------------- +function TO.tensoradd!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::TBLISBackend, allocator = DefaultAllocator() + ) + check_arguments(TO.tensoradd!, C, A) + argcheck_tensoradd(C, A, pA) + dimcheck_tensoradd(C, A, pA) + Base.mightalias(C, A) && + throw(ArgumentError("output tensor must not be aliased with input tensor")) + + T = eltype(C) + einA, einC = add_labels(pA) + Av, Cv = SV(A), SV(C) + lenA, strideA = tblis_dims(Av) + lenC, strideC = tblis_dims(Cv) + GC.@preserve Av Cv lenA strideA lenC strideC begin + tA = tblis_tensor(Av, convert(T, α), lenA, strideA, isconj(Av, conjA)) + tC = tblis_tensor(Cv, convert(T, β), lenC, strideC, false) + TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC)) + end + return C +end + +function TO.tensortrace!( + C::AbstractArray, + A::AbstractArray, p::Index2Tuple, q::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::TBLISBackend, allocator = DefaultAllocator() + ) + check_arguments(TO.tensortrace!, C, A) + argcheck_tensortrace(C, A, p, q) + dimcheck_tensortrace(C, A, p, q) + Base.mightalias(C, A) && + throw(ArgumentError("output tensor must not be aliased with input tensor")) + + T = eltype(C) + einA, einC = trace_labels(p, q) + Av, Cv = SV(A), SV(C) + lenA, strideA = tblis_dims(Av) + lenC, strideC = tblis_dims(Cv) + GC.@preserve Av Cv lenA strideA lenC strideC begin + tA = tblis_tensor(Av, convert(T, α), lenA, strideA, isconj(Av, conjA)) + tC = tblis_tensor(Cv, convert(T, β), lenC, strideC, false) + TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC)) + end + return C +end + +function TO.tensorcontract!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + B::AbstractArray, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, + α::Number, β::Number, + backend::TBLISBackend, allocator = DefaultAllocator() + ) + check_arguments(TO.tensorcontract!, C, A, B) + argcheck_tensorcontract(C, A, pA, B, pB, pAB) + dimcheck_tensorcontract(C, A, pA, B, pB, pAB) + (Base.mightalias(C, A) || Base.mightalias(C, B)) && + throw(ArgumentError("output tensor must not be aliased with input tensor")) + + T = eltype(C) + einA, einB, einC = contract_labels(pA, pB, pAB) + α′ = convert(T, α) + β′ = convert(T, β) + Av, Bv, Cv = SV(A), SV(B), SV(C) + isconjA = isconj(Av, conjA) + isconjB = isconj(Bv, conjB) + + # `tblis_tensor_mult` ignores the conjugation flags, so resolve them into the data first + if isconjA && isconjB + iszero(β′) || conj!(Cv) + tblis_mult!(Cv, Av, Bv, einA, einB, einC, conj(α′), conj(β′)) + conj!(Cv) + elseif isconjA + A′ = materialize_conj(Av, conjA, α′, allocator) + tblis_mult!(Cv, SV(A′), Bv, einA, einB, einC, one(T), β′) + tensorfree!(A′, allocator) + elseif isconjB + B′ = materialize_conj(Bv, conjB, one(T), allocator) + tblis_mult!(Cv, Av, SV(B′), einA, einB, einC, α′, β′) + tensorfree!(B′, allocator) + else + tblis_mult!(Cv, Av, Bv, einA, einB, einC, α′, β′) + end + return C +end + +function tblis_mult!( + C::StridedView{T}, A::StridedView{T}, B::StridedView{T}, + einA, einB, einC, α::T, β::T + ) where {T <: TBLISFloat} + lenA, strideA = tblis_dims(A) + lenB, strideB = tblis_dims(B) + lenC, strideC = tblis_dims(C) + GC.@preserve A B C lenA strideA lenB strideB lenC strideC begin + tA = tblis_tensor(A, α, lenA, strideA, false) + tB = tblis_tensor(B, one(T), lenB, strideB, false) + tC = tblis_tensor(C, β, lenC, strideC, false) + TBLIS.tblis_tensor_mult( + C_NULL, C_NULL, tA, labels(einA), tB, labels(einB), tC, labels(einC) + ) + end + return C +end + +function materialize_conj( + A::StridedView{T, N}, conjA::Bool, α::T, allocator + ) where {T <: TBLISFloat, N} + pA = (ntuple(identity, N), ()) + A′ = tensoralloc_add(T, A, pA, false, Val(true), allocator) + TO.tensoradd!(A′, A, pA, conjA, α, zero(T), TBLISBackend(), allocator) + return A′ +end + +end # module TensorOperationsTBLISExt diff --git a/src/backends.jl b/src/backends.jl index 3b5a29aa..49f622b6 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -71,6 +71,38 @@ struct StridedBLAS <: AbstractBackend end const StridedBackend = Union{StridedNative, StridedBLAS} +# TBLIS backend +#-------------- +""" + TBLISBackend() + +Backend for tensor operations on strided arrays that is based on the +[TBLIS](https://github.com/devinamatthews/tblis) library. +TBLIS performs tensor additions, traces and contractions directly on strided memory, without +the transpositions and temporaries that a BLAS-based approach requires. +This backend is only available through a package extension for +[TBLIS.jl](https://github.com/QuantumKitHub/TBLIS.jl). + +TBLIS requires all tensors in a single operation to share one element type, which moreover +has to be one of `Float32`, `Float64`, `ComplexF32` or `ComplexF64`. +Arguments that do not meet these requirements, including non-strided arrays, are rejected +with an `ArgumentError` rather than passed on to another backend. + +!!! note + Contracting in place trades throughput for memory. + On a 16-core node this backend is roughly on par with + [`StridedBLAS`](@ref TensorOperations.StridedBLAS) for permuted real contractions and + slower for other shapes, but it allocates no permuted temporaries at all. + +!!! warning + As of `tblis_jll` v1.3, TBLIS has no competitive support for complex element types: its + complex contraction kernels run an order of magnitude slower than the corresponding BLAS + calls. + Prefer [`StridedBLAS`](@ref TensorOperations.StridedBLAS) for complex-valued + contractions. +""" +struct TBLISBackend <: AbstractBackend end + # CuTENSOR backend #----------------- """ diff --git a/test/runtests.jl b/test/runtests.jl index 55edbced..3a26e37c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -81,6 +81,10 @@ if !is_buildkite include("butensor.jl") end + @testset "TBLIS extension" verbose = true begin + include("tblis.jl") + end + @testset "Polynomials" begin include("polynomials.jl") end diff --git a/test/tblis.jl b/test/tblis.jl new file mode 100644 index 00000000..7c8610f1 --- /dev/null +++ b/test/tblis.jl @@ -0,0 +1,245 @@ +using TensorOperations +using TensorOperations: TBLISBackend, StridedNative +using TBLIS +using LinearAlgebra +using Random +using Test + +Random.seed!(1234567) + +const eltypes = (Float32, Float64, ComplexF32, ComplexF64) +const tblis = TBLISBackend() +const reference = StridedNative() + +# `Zero()`/`One()` and plain numbers should all work; poison the output with `NaN` whenever +# `β == 0` so that a kernel that computes `0 * C` instead of ignoring `C` is caught. +poison!(C) = fill!(C, convert(eltype(C), NaN)) + +@testset "tensoradd! (eltype = $T)" for T in eltypes + A = randn(T, (3, 5, 4, 6)) + p = ((3, 1), (4, 2)) + for conjA in (false, true), (α, β) in ((1, 0), (randn(T), 0), (randn(T), randn(T))) + C = randn(T, (4, 3, 6, 5)) + Cref = copy(C) + iszero(β) && (poison!(C); poison!(Cref)) + @test tensoradd!(C, A, p, conjA, α, β, tblis) ≈ + tensoradd!(Cref, A, p, conjA, α, β, reference) + end + + # non-contiguous input and output through views + Aview = view(randn(T, (6, 10, 8, 12)), 1:2:6, 1:2:10, 1:2:8, 1:2:12) + C = randn(T, (4, 3, 6, 5)) + Cref = copy(C) + @test tensoradd!(C, Aview, p, true, 2, 1, tblis) ≈ + tensoradd!(Cref, Aview, p, true, 2, 1, reference) + + # an `Adjoint` input already carries a conjugation, which has to combine with `conjA` + Aadj = adjoint(randn(T, (5, 3))) + for conjA in (false, true) + C = randn(T, (5, 3)) + Cref = copy(C) + @test tensoradd!(C, Aadj, ((2, 1), ()), conjA, 2, 1, tblis) ≈ + tensoradd!(Cref, Aadj, ((2, 1), ()), conjA, 2, 1, reference) + end +end + +@testset "tensortrace! (eltype = $T)" for T in eltypes + A = randn(T, (5, 3, 4, 5, 3, 2)) + p = ((6, 3), ()) + q = ((1, 2), (4, 5)) + for conjA in (false, true), (α, β) in ((1, 0), (randn(T), randn(T))) + C = randn(T, (2, 4)) + Cref = copy(C) + iszero(β) && (poison!(C); poison!(Cref)) + @test tensortrace!(C, A, p, q, conjA, α, β, tblis) ≈ + tensortrace!(Cref, A, p, q, conjA, α, β, reference) + end + + # trace all the way down to a scalar + B = randn(T, (4, 4)) + C = fill(convert(T, NaN)) + Cref = fill(convert(T, NaN)) + @test tensortrace!(C, B, ((), ()), ((1,), (2,)), false, 1, 0, tblis)[] ≈ tr(B) +end + +@testset "tensorcontract! (eltype = $T)" for T in eltypes + # open indices of A are its dimensions 3, 1, 4 (sizes 5, 3, 3), contracted are 2 and 5 + # (sizes 20, 4); open indices of B are its dimensions 4 and 2 (sizes 3, 6) + A = randn(T, (3, 20, 5, 3, 4)) + B = randn(T, (4, 6, 20, 3)) + pA = ((3, 1, 4), (2, 5)) + pB = ((3, 1), (4, 2)) + pAB = ((3, 1, 4), (5, 2)) + + for conjA in (false, true), conjB in (false, true), + (α, β) in ((1, 0), (randn(T), 0), (randn(T), randn(T))) + + C = randn(T, (3, 5, 3, 6, 3)) + Cref = copy(C) + iszero(β) && (poison!(C); poison!(Cref)) + @test tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, α, β, tblis) ≈ + tensorcontract!(Cref, A, pA, conjA, B, pB, conjB, pAB, α, β, reference) + end + + # outer product: no contracted indices at all + @testset "outer product" begin + A2 = randn(T, (3, 4)) + B2 = randn(T, (5,)) + pA2, pB2, pAB2 = ((1, 2), ()), ((), (1,)), ((3, 1), (2,)) + C = fill(convert(T, NaN), (5, 3, 4)) + Cref = copy(C) + @test tensorcontract!(C, A2, pA2, true, B2, pB2, false, pAB2, 1, 0, tblis) ≈ + tensorcontract!(Cref, A2, pA2, true, B2, pB2, false, pAB2, 1, 0, reference) + end + + # full contraction: zero-dimensional output + @testset "full contraction" begin + A2 = randn(T, (3, 4)) + B2 = randn(T, (4, 3)) + pA2, pB2, pAB2 = ((), (1, 2)), ((2, 1), ()), ((), ()) + C = fill(convert(T, NaN)) + Cref = fill(convert(T, NaN)) + @test tensorcontract!(C, A2, pA2, true, B2, pB2, true, pAB2, 1, 0, tblis)[] ≈ + tensorcontract!(Cref, A2, pA2, true, B2, pB2, true, pAB2, 1, 0, reference)[] + end + + # non-contiguous factors + @testset "strided views" begin + Av = view(randn(T, (6, 8)), 1:2:6, 1:2:8) + Bv = view(randn(T, (8, 10)), 1:2:8, 1:2:10) + C = fill(convert(T, NaN), (3, 5)) + Cref = copy(C) + pA2, pB2, pAB2 = ((1,), (2,)), ((1,), (2,)), ((1, 2), ()) + @test tensorcontract!(C, Av, pA2, false, Bv, pB2, true, pAB2, 1, 0, tblis) ≈ + tensorcontract!(Cref, Av, pA2, false, Bv, pB2, true, pAB2, 1, 0, reference) + end +end + +@testset "argument checking" begin + A = randn(Float64, (4, 4)) + B = randn(Float64, (4, 4)) + # aliasing must be rejected rather than silently producing garbage + @test_throws ArgumentError tensoradd!(A, A, ((2, 1), ()), false, 1, 0, tblis) + @test_throws ArgumentError tensorcontract!( + A, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) + # shape errors are reported before anything reaches the library + C = randn(Float64, (4, 3)) + @test_throws DimensionMismatch tensoradd!(C, A, ((1, 2), ()), false, 1, 0, tblis) + @test_throws DimensionMismatch tensorcontract!( + C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) +end + +@testset "rejection of unsupported arguments" begin + # element types TBLIS does not know about + for A in ( + randn(Float16, (3, 4)), rand(1:4, (3, 4)), + Rational{Int}.(rand(1:4, (3, 4)), 3), + ) + T = eltype(A) + B = T <: AbstractFloat ? randn(T, (4, 5)) : T.(rand(1:4, (4, 5))) + C = zeros(T, (3, 5)) + @test_throws ArgumentError tensorcontract!( + C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) + @test_throws ArgumentError tensoradd!( + zeros(T, (4, 3)), A, ((2, 1), ()), false, 1, 0, tblis + ) + end + + # mixed element types: TBLIS requires a single type for all tensors + A = randn(Float64, (3, 4)) + B = randn(ComplexF64, (4, 5)) + C = zeros(ComplexF64, (3, 5)) + @test_throws ArgumentError tensorcontract!( + C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) + + # non-strided arrays + D = Diagonal(randn(Float64, 4)) + A = randn(Float64, (3, 4)) + C = zeros(Float64, (3, 4)) + @test_throws ArgumentError tensorcontract!( + C, A, ((1,), (2,)), false, D, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) + + # writing into a conjugated view would silently produce the conjugate of the result + A = randn(ComplexF64, (3, 4)) + Cadj = adjoint(zeros(ComplexF64, (4, 3))) + @test_throws ArgumentError tensoradd!( + Cadj, A, ((1, 2), ()), false, 1, 0, tblis + ) +end + +@testset "@tensor and ncon integration (eltype = $T)" for T in eltypes + A = randn(T, (5, 5, 5, 5)) + B = randn(T, (5, 5, 5)) + C = randn(T, (5, 5, 5)) + + @tensor backend = tblis D[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * conj(C[g, f, b]) + @tensor backend = reference Dref[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * + conj(C[g, f, b]) + @test D ≈ Dref + + network = [[-1, 1, -3, 2], [3, -4, 1], [3, 2, -2]] + conjlist = [false, false, true] + @test ncon([A, B, C], network, conjlist; backend = tblis) ≈ Dref + @test ncon([A, B, C], network; backend = tblis) ≈ + ncon([A, B, C], network; backend = reference) + + # traces and scalar results + @tensor backend = tblis s = A[a, b, a, b] + @tensor backend = reference sref = A[a, b, a, b] + @test s ≈ sref +end + +@testset "garbage collection safety" begin + # A `tblis_tensor` only stores raw pointers into the array and into the buffers holding + # its lengths and strides; make sure nothing goes missing under GC pressure. + T = ComplexF64 + A = randn(T, (12, 8, 6)) + B = randn(T, (8, 6, 10)) + Cref = similar(A, (12, 10)) + tensorcontract!( + Cref, A, ((1,), (2, 3)), true, B, ((1, 2), (3,)), false, ((1, 2), ()), 1, 0, + reference + ) + for i in 1:100 + C = similar(A, (12, 10)) + tensorcontract!( + C, A, ((1,), (2, 3)), true, B, ((1, 2), (3,)), false, ((1, 2), ()), 1, 0, tblis + ) + @test C ≈ Cref + iszero(i % 10) && GC.gc(true) + # keep the allocator busy so that freed buffers get reused promptly + junk = [randn(T, 1024) for _ in 1:8] + @test length(junk) == 8 + end +end + +@testset "threading" begin + nthreads = TBLIS.get_num_threads() + counts = ntuple(identity, min(2, nthreads)) + try + A = randn(Float64, (40, 40, 20)) + B = randn(Float64, (20, 40, 40)) + Cref = similar(A, (40, 40, 40, 40)) + tensorcontract!( + Cref, A, ((1, 2), (3,)), false, B, ((1,), (2, 3)), false, + ((1, 2, 3, 4), ()), 1, 0, reference + ) + for n in counts + TBLIS.set_num_threads(n) + @test TBLIS.get_num_threads() == n + C = similar(A, (40, 40, 40, 40)) + tensorcontract!( + C, A, ((1, 2), (3,)), false, B, ((1,), (2, 3)), false, + ((1, 2, 3, 4), ()), 1, 0, tblis + ) + @test C ≈ Cref + end + finally + TBLIS.set_num_threads(nthreads) + end +end