From 9eb01b2da20bc3839db020fa7a39f3c4e9fe99f9 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 10:44:34 -0400 Subject: [PATCH 1/7] Add a TBLIS.jl package extension Adds `TBLISBackend`, which routes `tensoradd!`, `tensortrace!` and `tensorcontract!` through the TBLIS library via TBLIS.jl. TBLIS contracts strided tensors in place, so it avoids the permuted intermediates that the BLAS-based backend has to materialize. The backend is opt-in: loading TBLIS.jl does not register a `select_backend` method. Arguments TBLIS cannot express, that is mixed or unsupported element types, non-strided arrays and conjugated outputs, throw an `ArgumentError` rather than being handed to another backend, so that a contraction which never reaches TBLIS cannot pass for one that did. Conjugation is carried as an ordinary runtime flag on the `TBLISTensor` descriptor rather than in the type of a `StridedView`, which keeps the entry points type stable without branching over `conj` variants. A conjugation already present on the input, as for an `Adjoint`, combines with the requested one. Two library quirks are worked around: * `tblis_tensor_mult` ignores the per-tensor conjugation flags, unlike `tblis_tensor_add`. When both factors are conjugated this is resolved by conjugating the output in place, and otherwise by materializing the conjugated factor into a temporary from the allocator. * TBLIS.jl only exposes its tensor constructor for `StridedArray` and offers no way to set the conjugation flag, so the descriptor is initialized through the low-level bindings and every referenced buffer is rooted with `GC.@preserve`. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 6 +- docs/src/man/backends.md | 23 ++- ext/TensorOperationsTBLISExt.jl | 288 ++++++++++++++++++++++++++++++++ src/backends.jl | 32 ++++ test/runtests.jl | 4 + test/tblis.jl | 244 +++++++++++++++++++++++++++ 6 files changed, 595 insertions(+), 2 deletions(-) create mode 100644 ext/TensorOperationsTBLISExt.jl create mode 100644 test/tblis.jl diff --git a/Project.toml b/Project.toml index 99c531b9..64255196 100644 --- a/Project.toml +++ b/Project.toml @@ -26,12 +26,14 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" 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" TensorOperationsGPUArraysExt = "GPUArrays" @@ -62,6 +64,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" @@ -84,8 +87,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 7ac3cdff..6fe1ecb8 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..5fe17536 --- /dev/null +++ b/ext/TensorOperationsTBLISExt.jl @@ -0,0 +1,288 @@ +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 + +# TBLIS only knows about these four element types, and requires all tensors taking part in a +# single operation to share it. +const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} + +#------------------------------------------------------------------------------------------- +# Wrapping Julia arrays as TBLIS tensors +#------------------------------------------------------------------------------------------- +# TBLIS.jl's own `tblis_tensor` constructor is restricted to `StridedArray`, which excludes +# the `StridedView`s that the tensor operations work with, and it has no way to express the +# conjugation flag. We therefore initialize the descriptor through the low-level bindings and +# patch in the flag afterwards. +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 + +""" + TBLISTensor(A::StridedView, α, isconj) + +Owning counterpart of `TBLIS.tblis_tensor`, which itself only stores raw pointers into the +array it views and into the buffers holding its lengths and strides. +Rooting a `TBLISTensor` with `GC.@preserve` also roots all of those, as they are reachable +from it; pass its `ref` field wherever the library expects a `Ptr{tblis_tensor}`. + +Conjugation is carried in the descriptor rather than in the type of the viewed array, so +`isconj` is an ordinary runtime flag and callers never have to branch on it to stay type +stable. +""" +struct TBLISTensor{T, N, A <: StridedView{T, N}} + array::A + len::Vector{len_type} + stride::Vector{stride_type} + ref::Base.RefValue{tblis_tensor} +end + +function TBLISTensor( + A::StridedView{T, N}, α::T, isconj::Bool + ) where {T <: TBLISFloat, N} + len = collect(len_type, size(A)) + stride = collect(stride_type, strides(A)) + 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) + isconj && setproperty!(p, :conj, Cint(1)) + end + return TBLISTensor{T, N, typeof(A)}(A, len, stride, ref) +end + +# A `StridedView` may already carry a conjugation in its `op` field, for instance when it +# views an `Adjoint`. What TBLIS needs is the total conjugation applied to the raw data, so +# that flag and the one requested by the caller combine. For real element types `conj` is the +# identity and flagging it would only confuse the library. +hasconj(A::StridedView{T}) where {T} = T <: Complex && (A.op === conj || A.op === adjoint) +resolve_conj(A::StridedView{T}, conjA::Bool) where {T} = + T <: Complex && (conjA ⊻ hasconj(A)) + +# TBLIS takes the index labels as one `char` per dimension; `add_labels` and friends hand us +# `Char` tuples that are guaranteed to be ASCII. +labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) + +#------------------------------------------------------------------------------------------- +# Argument checking +#------------------------------------------------------------------------------------------- +# Rather than quietly handing unsupported arguments to another backend, which would make a +# `backend = TBLISBackend()` that never reaches TBLIS look like it did, reject them. +@noinline function throw_eltype(f, tensors) + types = join(eltype.(tensors), ", ") + return throw( + ArgumentError( + "TBLISBackend requires all tensors of $f to share a single element type out of \ + Float32, Float64, ComplexF32 and ComplexF64, got $types" + ) + ) +end + +@noinline function throw_strided(f, tensors) + types = join(typeof.(tensors), ", ") + return throw( + ArgumentError("TBLISBackend requires strided arrays for $f, got $types") + ) +end + +@noinline throw_conj_output(f) = throw( + ArgumentError("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) + hasconj(SV(C)) && throw_conj_output(f) + return nothing +end + +#------------------------------------------------------------------------------------------- +# Entry points +#------------------------------------------------------------------------------------------- +function TO.tensoradd!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::TBLISBackend, allocator = DefaultAllocator() + ) + check_arguments(TO.tensoradd!, C, A) + tblis_add!(SV(C), SV(A), conjA, pA, α, β) + 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) + tblis_trace!(SV(C), SV(A), conjA, p, q, α, β) + 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) + tblis_contract!(SV(C), SV(A), pA, conjA, SV(B), pB, conjB, pAB, α, β, allocator) + return C +end + +#------------------------------------------------------------------------------------------- +# StridedView implementation +#------------------------------------------------------------------------------------------- +# `tblis_tensor_add` computes `C[einC] = β * C[einC] + α * op(A)[einA]` and does honour the +# per-tensor conjugation flag. Labels repeated within `einA` but absent from `einC` are +# traced over, so both `tensoradd!` and `tensortrace!` map onto it directly. +function tblis_add!( + C::StridedView{T}, A::StridedView{T}, conjA::Bool, pA::Index2Tuple, + α::Number, β::Number + ) where {T <: TBLISFloat} + 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")) + + einA, einC = add_labels(pA) + return unsafe_add!( + C, A, resolve_conj(A, conjA), einA, einC, convert(T, α), convert(T, β) + ) +end + +function tblis_trace!( + C::StridedView{T}, A::StridedView{T}, conjA::Bool, + p::Index2Tuple, q::Index2Tuple, α::Number, β::Number + ) where {T <: TBLISFloat} + 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")) + + einA, einC = trace_labels(p, q) + return unsafe_add!( + C, A, resolve_conj(A, conjA), einA, einC, convert(T, α), convert(T, β) + ) +end + +# `isconjA` is the total conjugation to apply to the raw data of `A`, as resolved by +# `resolve_conj`. +function unsafe_add!( + C::StridedView{T}, A::StridedView{T}, isconjA::Bool, einA, einC, α::T, β::T + ) where {T <: TBLISFloat} + tA = TBLISTensor(A, α, isconjA) + tC = TBLISTensor(C, β, false) + GC.@preserve tA tC begin + TBLIS.tblis_tensor_add( + C_NULL, C_NULL, tA.ref, labels(einA), tC.ref, labels(einC) + ) + end + return C +end + +function tblis_contract!( + C::StridedView{T}, + A::StridedView{T}, pA::Index2Tuple, conjA::Bool, + B::StridedView{T}, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, + α::Number, β::Number, allocator + ) where {T <: TBLISFloat} + 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")) + + einA, einB, einC = contract_labels(pA, pB, pAB) + α′ = convert(T, α) + β′ = convert(T, β) + isconjA = resolve_conj(A, conjA) + isconjB = resolve_conj(B, conjB) + + # `tblis_tensor_mult` silently ignores the conjugation flags of its arguments (verified + # against tblis 1.3), so conjugation has to be resolved before calling into it. + if isconjA && isconjB + # conj(A) * conj(B) == conj(A * B), so conjugating the output in place lets both + # factors through unconjugated, at the cost of two passes over C. That is much + # cheaper than materializing a conjugated copy of both A and B. + iszero(β′) || conj!(C) + unsafe_mult!(C, A, B, einA, einB, einC, conj(α′), conj(β′)) + conj!(C) + elseif isconjA + A′ = materialize_conj(A, α′, allocator) + try + unsafe_mult!(C, SV(A′), B, einA, einB, einC, one(T), β′) + finally + tensorfree!(A′, allocator) + end + elseif isconjB + B′ = materialize_conj(B, one(T), allocator) + try + unsafe_mult!(C, A, SV(B′), einA, einB, einC, α′, β′) + finally + tensorfree!(B′, allocator) + end + else + unsafe_mult!(C, A, B, einA, einB, einC, α′, β′) + end + return C +end + +function unsafe_mult!( + C::StridedView{T}, A::StridedView{T}, B::StridedView{T}, + einA, einB, einC, α::T, β::T + ) where {T <: TBLISFloat} + # TBLIS scales the product by the scalars of both factors, so α rides along on A alone. + # The conjugation flags are all `false` because `tblis_tensor_mult` ignores them and the + # caller has resolved the conjugations already. + tA = TBLISTensor(A, α, false) + tB = TBLISTensor(B, one(T), false) + tC = TBLISTensor(C, β, false) + GC.@preserve tA tB tC begin + TBLIS.tblis_tensor_mult( + C_NULL, C_NULL, tA.ref, labels(einA), tB.ref, labels(einB), + tC.ref, labels(einC) + ) + end + return C +end + +# Write `α * conj(raw A)` into a fresh temporary that keeps the index order of `A`, so that +# the labels computed for `A` remain valid for it and it can be fed to `mult` unconjugated. +function materialize_conj(A::StridedView{T, N}, α::T, allocator) where {T <: TBLISFloat, N} + pA = (ntuple(identity, N), ()) + A′ = tensoralloc_add(T, A, pA, false, Val(true), allocator) + einA, einC = add_labels(pA) + unsafe_add!(SV(A′), A, true, einA, einC, α, zero(T)) + 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..8b91a69d --- /dev/null +++ b/test/tblis.jl @@ -0,0 +1,244 @@ +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() + 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 (1, 2) + 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 From e151725e30fd5fce4d4a747c340046b689900bfd Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 11:03:17 -0400 Subject: [PATCH 2/7] Address review: drop the descriptor struct and flatten the call paths * Replace the `TBLISTensor` struct with a `tblis_tensor` function returning the `Ref{tblis_tensor}` directly. The caller owns the length and stride buffers and keeps them rooted, so nothing has to be wrapped to stay alive. * `isconj` specializes on `StridedView` to fold any conjugation the view already carries into the one the caller requested. * Inline `tblis_add!` and `unsafe_add!` into `tensoradd!` and `tensortrace!`. `materialize_conj` now goes through `tensoradd!` instead of a private copy of the same call sequence. * Drop the `try`/`finally` around the conjugated temporaries and free them on the success path only, matching how the cuTENSOR extension handles this. * Guard the threading test so it only asks for as many threads as the machine reports, for single-core CI runners. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationsTBLISExt.jl | 230 ++++++++++++++------------------ test/tblis.jl | 5 +- 2 files changed, 103 insertions(+), 132 deletions(-) diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 5fe17536..52d3a18f 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -23,9 +23,9 @@ const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} # Wrapping Julia arrays as TBLIS tensors #------------------------------------------------------------------------------------------- # TBLIS.jl's own `tblis_tensor` constructor is restricted to `StridedArray`, which excludes -# the `StridedView`s that the tensor operations work with, and it has no way to express the -# conjugation flag. We therefore initialize the descriptor through the low-level bindings and -# patch in the flag afterwards. +# the `StridedView`s used below, and it has no way to express the conjugation flag. We +# therefore initialize the descriptor through the low-level bindings and patch in the flag +# afterwards. for (T, init) in ( (:Float32, :tblis_init_tensor_scaled_s), (:Float64, :tblis_init_tensor_scaled_d), @@ -40,47 +40,42 @@ for (T, init) in ( end end -""" - TBLISTensor(A::StridedView, α, isconj) +# Total conjugation to apply to the raw data of `A`: a `StridedView` may already carry one in +# its `op` field, as it does for an `Adjoint`, and that combines with the one the caller asks +# for. Conjugation therefore never has to appear in the type of the view, which keeps the +# callers free of union splits over `conj` variants. For real element types `conj` is the +# identity and flagging it would only confuse the library. +isconj(A::StridedView{T}, conjA::Bool) where {T} = T <: Complex && (conjA ⊻ (A.op === conj)) -Owning counterpart of `TBLIS.tblis_tensor`, which itself only stores raw pointers into the -array it views and into the buffers holding its lengths and strides. -Rooting a `TBLISTensor` with `GC.@preserve` also roots all of those, as they are reachable -from it; pass its `ref` field wherever the library expects a `Ptr{tblis_tensor}`. +# Lengths and strides in the layout TBLIS expects. The descriptor stores raw pointers into +# these, so the caller has to keep them alive for as long as TBLIS may look at them. +tblis_dims(A::StridedView) = + (collect(len_type, size(A)), collect(stride_type, strides(A))) -Conjugation is carried in the descriptor rather than in the type of the viewed array, so -`isconj` is an ordinary runtime flag and callers never have to branch on it to stay type -stable. """ -struct TBLISTensor{T, N, A <: StridedView{T, N}} - array::A - len::Vector{len_type} - stride::Vector{stride_type} - ref::Base.RefValue{tblis_tensor} -end + 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. -function TBLISTensor( - A::StridedView{T, N}, α::T, isconj::Bool +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} - len = collect(len_type, size(A)) - stride = collect(stride_type, strides(A)) 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) - isconj && setproperty!(p, :conj, Cint(1)) + conj && setproperty!(p, :conj, Cint(1)) end - return TBLISTensor{T, N, typeof(A)}(A, len, stride, ref) + return ref end -# A `StridedView` may already carry a conjugation in its `op` field, for instance when it -# views an `Adjoint`. What TBLIS needs is the total conjugation applied to the raw data, so -# that flag and the one requested by the caller combine. For real element types `conj` is the -# identity and flagging it would only confuse the library. -hasconj(A::StridedView{T}) where {T} = T <: Complex && (A.op === conj || A.op === adjoint) -resolve_conj(A::StridedView{T}, conjA::Bool) where {T} = - T <: Complex && (conjA ⊻ hasconj(A)) - # TBLIS takes the index labels as one `char` per dimension; `add_labels` and friends hand us # `Char` tuples that are guaranteed to be ASCII. labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) @@ -116,13 +111,21 @@ function check_arguments(f, C::AbstractArray, As::AbstractArray...) T = eltype(C) (T <: TBLISFloat && all(A -> eltype(A) === T, As)) || throw_eltype(f, tensors) all(isstrided, tensors) || throw_strided(f, tensors) - hasconj(SV(C)) && throw_conj_output(f) + isconj(SV(C), false) && throw_conj_output(f) return nothing end #------------------------------------------------------------------------------------------- -# Entry points +# Operations #------------------------------------------------------------------------------------------- +# The `StridedView` wrapper is what makes the whole range of strided inputs usable here: an +# `Adjoint` of a complex array has neither `strides` nor `pointer`, and a `ReshapedArray` need +# not have `strides` either, whereas `StridedView` normalizes all of them and exposes any +# conjugation as its `op` field. It is also exactly what the `isstrided` check above admits. + +# `tblis_tensor_add` computes `C[einC] = β * C[einC] + α * op(A)[einA]` and does honour the +# per-tensor conjugation flag. Labels repeated within `einA` but absent from `einC` are traced +# over, so both `tensoradd!` and `tensortrace!` map onto it directly. function TO.tensoradd!( C::AbstractArray, A::AbstractArray, pA::Index2Tuple, conjA::Bool, @@ -130,7 +133,21 @@ function TO.tensoradd!( backend::TBLISBackend, allocator = DefaultAllocator() ) check_arguments(TO.tensoradd!, C, A) - tblis_add!(SV(C), SV(A), conjA, pA, α, β) + 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 @@ -141,7 +158,21 @@ function TO.tensortrace!( backend::TBLISBackend, allocator = DefaultAllocator() ) check_arguments(TO.tensortrace!, C, A) - tblis_trace!(SV(C), SV(A), conjA, p, q, α, β) + 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 @@ -154,78 +185,18 @@ function TO.tensorcontract!( backend::TBLISBackend, allocator = DefaultAllocator() ) check_arguments(TO.tensorcontract!, C, A, B) - tblis_contract!(SV(C), SV(A), pA, conjA, SV(B), pB, conjB, pAB, α, β, allocator) - return C -end - -#------------------------------------------------------------------------------------------- -# StridedView implementation -#------------------------------------------------------------------------------------------- -# `tblis_tensor_add` computes `C[einC] = β * C[einC] + α * op(A)[einA]` and does honour the -# per-tensor conjugation flag. Labels repeated within `einA` but absent from `einC` are -# traced over, so both `tensoradd!` and `tensortrace!` map onto it directly. -function tblis_add!( - C::StridedView{T}, A::StridedView{T}, conjA::Bool, pA::Index2Tuple, - α::Number, β::Number - ) where {T <: TBLISFloat} - 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")) - - einA, einC = add_labels(pA) - return unsafe_add!( - C, A, resolve_conj(A, conjA), einA, einC, convert(T, α), convert(T, β) - ) -end - -function tblis_trace!( - C::StridedView{T}, A::StridedView{T}, conjA::Bool, - p::Index2Tuple, q::Index2Tuple, α::Number, β::Number - ) where {T <: TBLISFloat} - 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")) - - einA, einC = trace_labels(p, q) - return unsafe_add!( - C, A, resolve_conj(A, conjA), einA, einC, convert(T, α), convert(T, β) - ) -end - -# `isconjA` is the total conjugation to apply to the raw data of `A`, as resolved by -# `resolve_conj`. -function unsafe_add!( - C::StridedView{T}, A::StridedView{T}, isconjA::Bool, einA, einC, α::T, β::T - ) where {T <: TBLISFloat} - tA = TBLISTensor(A, α, isconjA) - tC = TBLISTensor(C, β, false) - GC.@preserve tA tC begin - TBLIS.tblis_tensor_add( - C_NULL, C_NULL, tA.ref, labels(einA), tC.ref, labels(einC) - ) - end - return C -end - -function tblis_contract!( - C::StridedView{T}, - A::StridedView{T}, pA::Index2Tuple, conjA::Bool, - B::StridedView{T}, pB::Index2Tuple, conjB::Bool, - pAB::Index2Tuple, - α::Number, β::Number, allocator - ) where {T <: TBLISFloat} 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, β) - isconjA = resolve_conj(A, conjA) - isconjB = resolve_conj(B, conjB) + Av, Bv, Cv = SV(A), SV(B), SV(C) + isconjA = isconj(Av, conjA) + isconjB = isconj(Bv, conjB) # `tblis_tensor_mult` silently ignores the conjugation flags of its arguments (verified # against tblis 1.3), so conjugation has to be resolved before calling into it. @@ -233,55 +204,52 @@ function tblis_contract!( # conj(A) * conj(B) == conj(A * B), so conjugating the output in place lets both # factors through unconjugated, at the cost of two passes over C. That is much # cheaper than materializing a conjugated copy of both A and B. - iszero(β′) || conj!(C) - unsafe_mult!(C, A, B, einA, einB, einC, conj(α′), conj(β′)) - conj!(C) + iszero(β′) || conj!(Cv) + tblis_mult!(Cv, Av, Bv, einA, einB, einC, conj(α′), conj(β′)) + conj!(Cv) elseif isconjA - A′ = materialize_conj(A, α′, allocator) - try - unsafe_mult!(C, SV(A′), B, einA, einB, einC, one(T), β′) - finally - tensorfree!(A′, allocator) - end + 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(B, one(T), allocator) - try - unsafe_mult!(C, A, SV(B′), einA, einB, einC, α′, β′) - finally - tensorfree!(B′, allocator) - end + B′ = materialize_conj(Bv, conjB, one(T), allocator) + tblis_mult!(Cv, Av, SV(B′), einA, einB, einC, α′, β′) + tensorfree!(B′, allocator) else - unsafe_mult!(C, A, B, einA, einB, einC, α′, β′) + tblis_mult!(Cv, Av, Bv, einA, einB, einC, α′, β′) end return C end -function unsafe_mult!( +# Shared by the four conjugation branches above, which have already folded any conjugation +# into the data, so the descriptors are built unconjugated. TBLIS scales the product by the +# scalars of both factors, so α rides along on A alone. +function tblis_mult!( C::StridedView{T}, A::StridedView{T}, B::StridedView{T}, einA, einB, einC, α::T, β::T ) where {T <: TBLISFloat} - # TBLIS scales the product by the scalars of both factors, so α rides along on A alone. - # The conjugation flags are all `false` because `tblis_tensor_mult` ignores them and the - # caller has resolved the conjugations already. - tA = TBLISTensor(A, α, false) - tB = TBLISTensor(B, one(T), false) - tC = TBLISTensor(C, β, false) - GC.@preserve tA tB tC begin + 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.ref, labels(einA), tB.ref, labels(einB), - tC.ref, labels(einC) + C_NULL, C_NULL, tA, labels(einA), tB, labels(einB), tC, labels(einC) ) end return C end -# Write `α * conj(raw A)` into a fresh temporary that keeps the index order of `A`, so that -# the labels computed for `A` remain valid for it and it can be fed to `mult` unconjugated. -function materialize_conj(A::StridedView{T, N}, α::T, allocator) where {T <: TBLISFloat, N} +# Write `α * conj(A)` into a fresh temporary that keeps the index order of `A`, so that the +# labels computed for `A` remain valid for it and it can be fed to `mult` unconjugated. +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) - einA, einC = add_labels(pA) - unsafe_add!(SV(A′), A, true, einA, einC, α, zero(T)) + TO.tensoradd!(A′, A, pA, conjA, α, zero(T), TBLISBackend(), allocator) return A′ end diff --git a/test/tblis.jl b/test/tblis.jl index 8b91a69d..3bb87adb 100644 --- a/test/tblis.jl +++ b/test/tblis.jl @@ -220,6 +220,9 @@ end @testset "threading" begin nthreads = TBLIS.get_num_threads() + # a runner with a single core cannot honour a request for more than one thread, so only + # ask for counts it can actually provide + counts = ntuple(identity, min(2, nthreads)) try A = randn(Float64, (40, 40, 20)) B = randn(Float64, (20, 40, 40)) @@ -228,7 +231,7 @@ end Cref, A, ((1, 2), (3,)), false, B, ((1,), (2, 3)), false, ((1, 2, 3, 4), ()), 1, 0, reference ) - for n in (1, 2) + for n in counts TBLIS.set_num_threads(n) @test TBLIS.get_num_threads() == n C = similar(A, (40, 40, 40, 40)) From e227b3f658f80f9b45007c38927c3aae36cfc0c9 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 12:01:49 -0400 Subject: [PATCH 3/7] Address review: document the conjugated-output case, trim comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Note at each `SV` conversion that it is there to support inputs such as `Adjoint`, which have no `strides` or `pointer` of their own. * Explain why the `C` descriptor is always built unconjugated: TBLIS applies the flag of `C` when reading the `β * C` term but not when writing the result back, so a conjugated output would conjugate half the operation. `check_arguments` rejects such a `C` for that reason. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationsTBLISExt.jl | 23 ++++++++--------------- test/tblis.jl | 2 -- 2 files changed, 8 insertions(+), 17 deletions(-) diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 52d3a18f..1c59b8b5 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -40,17 +40,8 @@ for (T, init) in ( end end -# Total conjugation to apply to the raw data of `A`: a `StridedView` may already carry one in -# its `op` field, as it does for an `Adjoint`, and that combines with the one the caller asks -# for. Conjugation therefore never has to appear in the type of the view, which keeps the -# callers free of union splits over `conj` variants. For real element types `conj` is the -# identity and flagging it would only confuse the library. isconj(A::StridedView{T}, conjA::Bool) where {T} = T <: Complex && (conjA ⊻ (A.op === conj)) - -# Lengths and strides in the layout TBLIS expects. The descriptor stores raw pointers into -# these, so the caller has to keep them alive for as long as TBLIS may look at them. -tblis_dims(A::StridedView) = - (collect(len_type, size(A)), collect(stride_type, strides(A))) +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} @@ -102,6 +93,10 @@ end ) end +# A conjugated output cannot be expressed: TBLIS applies the flag of `C` when reading the +# `β * C` term but not when writing the result back, so it would conjugate half the +# operation. Hence `check_arguments` rejects such a `C` and every `C` descriptor below is +# built unconjugated. @noinline throw_conj_output(f) = throw( ArgumentError("TBLISBackend cannot write into a conjugated view in $f") ) @@ -118,11 +113,6 @@ end #------------------------------------------------------------------------------------------- # Operations #------------------------------------------------------------------------------------------- -# The `StridedView` wrapper is what makes the whole range of strided inputs usable here: an -# `Adjoint` of a complex array has neither `strides` nor `pointer`, and a `ReshapedArray` need -# not have `strides` either, whereas `StridedView` normalizes all of them and exposes any -# conjugation as its `op` field. It is also exactly what the `isstrided` check above admits. - # `tblis_tensor_add` computes `C[einC] = β * C[einC] + α * op(A)[einA]` and does honour the # per-tensor conjugation flag. Labels repeated within `einA` but absent from `einC` are traced # over, so both `tensoradd!` and `tensortrace!` map onto it directly. @@ -140,6 +130,7 @@ function TO.tensoradd!( T = eltype(C) einA, einC = add_labels(pA) + # `SV` supports inputs such as `Adjoint`, which have no `strides` or `pointer` of their own Av, Cv = SV(A), SV(C) lenA, strideA = tblis_dims(Av) lenC, strideC = tblis_dims(Cv) @@ -165,6 +156,7 @@ function TO.tensortrace!( T = eltype(C) einA, einC = trace_labels(p, q) + # `SV` supports inputs such as `Adjoint`, which have no `strides` or `pointer` of their own Av, Cv = SV(A), SV(C) lenA, strideA = tblis_dims(Av) lenC, strideC = tblis_dims(Cv) @@ -194,6 +186,7 @@ function TO.tensorcontract!( einA, einB, einC = contract_labels(pA, pB, pAB) α′ = convert(T, α) β′ = convert(T, β) + # `SV` supports inputs such as `Adjoint`, which have no `strides` or `pointer` of their own Av, Bv, Cv = SV(A), SV(B), SV(C) isconjA = isconj(Av, conjA) isconjB = isconj(Bv, conjB) diff --git a/test/tblis.jl b/test/tblis.jl index 3bb87adb..7c8610f1 100644 --- a/test/tblis.jl +++ b/test/tblis.jl @@ -220,8 +220,6 @@ end @testset "threading" begin nthreads = TBLIS.get_num_threads() - # a runner with a single core cannot honour a request for more than one thread, so only - # ask for counts it can actually provide counts = ntuple(identity, min(2, nthreads)) try A = randn(Float64, (40, 40, 20)) From 420f5ab8d4277bab532855647f592cb50a394dde Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 12:13:39 -0400 Subject: [PATCH 4/7] Drop the explanatory comments from the extension Keeps the section banners and the `tblis_tensor` docstring. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationsTBLISExt.jl | 30 ------------------------------ 1 file changed, 30 deletions(-) diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 1c59b8b5..1fe490ac 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -15,17 +15,11 @@ using TBLIS: len_type, stride_type, tblis_tensor const SV = StridedView -# TBLIS only knows about these four element types, and requires all tensors taking part in a -# single operation to share it. const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} #------------------------------------------------------------------------------------------- # Wrapping Julia arrays as TBLIS tensors #------------------------------------------------------------------------------------------- -# TBLIS.jl's own `tblis_tensor` constructor is restricted to `StridedArray`, which excludes -# the `StridedView`s used below, and it has no way to express the conjugation flag. We -# therefore initialize the descriptor through the low-level bindings and patch in the flag -# afterwards. for (T, init) in ( (:Float32, :tblis_init_tensor_scaled_s), (:Float64, :tblis_init_tensor_scaled_d), @@ -67,15 +61,11 @@ function tblis_tensor( return ref end -# TBLIS takes the index labels as one `char` per dimension; `add_labels` and friends hand us -# `Char` tuples that are guaranteed to be ASCII. labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) #------------------------------------------------------------------------------------------- # Argument checking #------------------------------------------------------------------------------------------- -# Rather than quietly handing unsupported arguments to another backend, which would make a -# `backend = TBLISBackend()` that never reaches TBLIS look like it did, reject them. @noinline function throw_eltype(f, tensors) types = join(eltype.(tensors), ", ") return throw( @@ -93,10 +83,6 @@ end ) end -# A conjugated output cannot be expressed: TBLIS applies the flag of `C` when reading the -# `β * C` term but not when writing the result back, so it would conjugate half the -# operation. Hence `check_arguments` rejects such a `C` and every `C` descriptor below is -# built unconjugated. @noinline throw_conj_output(f) = throw( ArgumentError("TBLISBackend cannot write into a conjugated view in $f") ) @@ -113,9 +99,6 @@ end #------------------------------------------------------------------------------------------- # Operations #------------------------------------------------------------------------------------------- -# `tblis_tensor_add` computes `C[einC] = β * C[einC] + α * op(A)[einA]` and does honour the -# per-tensor conjugation flag. Labels repeated within `einA` but absent from `einC` are traced -# over, so both `tensoradd!` and `tensortrace!` map onto it directly. function TO.tensoradd!( C::AbstractArray, A::AbstractArray, pA::Index2Tuple, conjA::Bool, @@ -130,7 +113,6 @@ function TO.tensoradd!( T = eltype(C) einA, einC = add_labels(pA) - # `SV` supports inputs such as `Adjoint`, which have no `strides` or `pointer` of their own Av, Cv = SV(A), SV(C) lenA, strideA = tblis_dims(Av) lenC, strideC = tblis_dims(Cv) @@ -156,7 +138,6 @@ function TO.tensortrace!( T = eltype(C) einA, einC = trace_labels(p, q) - # `SV` supports inputs such as `Adjoint`, which have no `strides` or `pointer` of their own Av, Cv = SV(A), SV(C) lenA, strideA = tblis_dims(Av) lenC, strideC = tblis_dims(Cv) @@ -186,17 +167,11 @@ function TO.tensorcontract!( einA, einB, einC = contract_labels(pA, pB, pAB) α′ = convert(T, α) β′ = convert(T, β) - # `SV` supports inputs such as `Adjoint`, which have no `strides` or `pointer` of their own Av, Bv, Cv = SV(A), SV(B), SV(C) isconjA = isconj(Av, conjA) isconjB = isconj(Bv, conjB) - # `tblis_tensor_mult` silently ignores the conjugation flags of its arguments (verified - # against tblis 1.3), so conjugation has to be resolved before calling into it. if isconjA && isconjB - # conj(A) * conj(B) == conj(A * B), so conjugating the output in place lets both - # factors through unconjugated, at the cost of two passes over C. That is much - # cheaper than materializing a conjugated copy of both A and B. iszero(β′) || conj!(Cv) tblis_mult!(Cv, Av, Bv, einA, einB, einC, conj(α′), conj(β′)) conj!(Cv) @@ -214,9 +189,6 @@ function TO.tensorcontract!( return C end -# Shared by the four conjugation branches above, which have already folded any conjugation -# into the data, so the descriptors are built unconjugated. TBLIS scales the product by the -# scalars of both factors, so α rides along on A alone. function tblis_mult!( C::StridedView{T}, A::StridedView{T}, B::StridedView{T}, einA, einB, einC, α::T, β::T @@ -235,8 +207,6 @@ function tblis_mult!( return C end -# Write `α * conj(A)` into a fresh temporary that keeps the index order of `A`, so that the -# labels computed for `A` remain valid for it and it can be fed to `mult` unconjugated. function materialize_conj( A::StridedView{T, N}, conjA::Bool, α::T, allocator ) where {T <: TBLISFloat, N} From 43583a753d4c75267159a2458840bb88748394e4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 12:36:03 -0400 Subject: [PATCH 5/7] Note the two TBLIS conjugation restrictions and make the errors lazy One-line comments for the two library restrictions that are not visible from the code, and `LazyString` messages for the three argument-check errors. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationsTBLISExt.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 1fe490ac..8c1adaf4 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -67,24 +67,24 @@ labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) # Argument checking #------------------------------------------------------------------------------------------- @noinline function throw_eltype(f, tensors) - types = join(eltype.(tensors), ", ") return throw( ArgumentError( - "TBLISBackend requires all tensors of $f to share a single element type out of \ - Float32, Float64, ComplexF32 and ComplexF64, got $types" + 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("TBLISBackend requires strided arrays for $f, got $types") - ) + return throw(ArgumentError(lazy"TBLISBackend requires strided arrays for $f, got $types")) end @noinline throw_conj_output(f) = throw( - ArgumentError("TBLISBackend cannot write into a conjugated view in $f") + ArgumentError(lazy"TBLISBackend cannot write into a conjugated view in $f") ) function check_arguments(f, C::AbstractArray, As::AbstractArray...) @@ -92,6 +92,7 @@ function check_arguments(f, C::AbstractArray, As::AbstractArray...) 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 @@ -171,6 +172,7 @@ function TO.tensorcontract!( 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(β′)) From f9dc9be04fae11fd1d121477f194ea79ddb8044d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 12 Aug 2026 09:48:15 -0400 Subject: [PATCH 6/7] Reject the TBLIS backend on Windows and fix the 1.12 constructor warning The Windows build of tblis_jll aborts the process with "posix_memalign: Invalid argument" out of tblis::MemoryPool::acquire as soon as a contraction reaches the GEMM kernels. The abort does not tear the process down cleanly either, so CI hung until the 6h runner timeout instead of failing. Reject the backend up front on Windows with an ArgumentError that points at StridedBLAS, skip the TBLIS test file there, and document the restriction. Rename the tblis_tensor helper to tblis_tensor_ref: tblis_tensor is a struct in TBLIS, so the method was extending its constructor through an implicit `using` binding, which Julia 1.12 warns about on load. Also cap the CI jobs at 90 minutes so that a native library that aborts or deadlocks fails the job instead of wedging a runner for six hours. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci.yml | 4 ++++ docs/src/man/backends.md | 3 +++ ext/TensorOperationsTBLISExt.jl | 36 ++++++++++++++++++++++++--------- src/backends.jl | 5 +++++ test/runtests.jl | 17 ++++++++++++++-- 5 files changed, 54 insertions(+), 11 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b2c4e62e..9c080b81 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,6 +17,9 @@ jobs: test: name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} runs-on: ${{ matrix.os }} + # a native library that aborts or deadlocks otherwise wedges the runner for the full + # 6h default; the slowest passing job takes about 45 minutes + timeout-minutes: 90 strategy: fail-fast: false matrix: @@ -49,6 +52,7 @@ jobs: needs: test name: Julia nightly - ${{ matrix.os }} - ${{ matrix.arch }} runs-on: ${{ matrix.os }} + timeout-minutes: 90 strategy: fail-fast: false matrix: diff --git a/docs/src/man/backends.md b/docs/src/man/backends.md index 6fe1ecb8..cd3ea8fc 100644 --- a/docs/src/man/backends.md +++ b/docs/src/man/backends.md @@ -88,6 +88,9 @@ TBLIS.set_num_threads(8) 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. +The backend is not supported on Windows. +`tblis_jll` does ship a Windows build, but it aborts the process from inside the library as soon as a contraction reaches its GEMM kernels, so calls there are rejected with an `ArgumentError` instead. + 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. diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 8c1adaf4..e2cdf13f 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -17,6 +17,11 @@ const SV = StridedView const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} +# `tblis_jll` does ship a Windows build, but it aborts the process with +# "posix_memalign: Invalid argument" from `tblis::MemoryPool::acquire` as soon as a +# contraction reaches the GEMM kernels, so nothing is handed to the library there. +const PLATFORM_SUPPORTED = !Sys.iswindows() + #------------------------------------------------------------------------------------------- # Wrapping Julia arrays as TBLIS tensors #------------------------------------------------------------------------------------------- @@ -38,7 +43,7 @@ isconj(A::StridedView{T}, conjA::Bool) where {T} = T <: Complex && (conjA ⊻ (A 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} + tblis_tensor_ref(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. @@ -48,7 +53,7 @@ with the returned `Ref`, have to be kept alive by the caller for as long as TBLI 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( +function tblis_tensor_ref( A::StridedView{T, N}, α::T, len::Vector{len_type}, stride::Vector{stride_type}, conj::Bool ) where {T <: TBLISFloat, N} @@ -66,6 +71,18 @@ labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) #------------------------------------------------------------------------------------------- # Argument checking #------------------------------------------------------------------------------------------- +@noinline function throw_unsupported_platform(f) + return throw( + ArgumentError( + LazyString( + "TBLISBackend is not supported on ", Base.BUILD_TRIPLET, ": the tblis_jll ", + "binaries for this platform abort the process from inside the library. ", + "Use another backend, such as StridedBLAS(), for ", f + ) + ) + ) +end + @noinline function throw_eltype(f, tensors) return throw( ArgumentError( @@ -88,6 +105,7 @@ end ) function check_arguments(f, C::AbstractArray, As::AbstractArray...) + PLATFORM_SUPPORTED || throw_unsupported_platform(f) tensors = (C, As...) T = eltype(C) (T <: TBLISFloat && all(A -> eltype(A) === T, As)) || throw_eltype(f, tensors) @@ -118,8 +136,8 @@ function TO.tensoradd!( 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) + tA = tblis_tensor_ref(Av, convert(T, α), lenA, strideA, isconj(Av, conjA)) + tC = tblis_tensor_ref(Cv, convert(T, β), lenC, strideC, false) TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC)) end return C @@ -143,8 +161,8 @@ function TO.tensortrace!( 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) + tA = tblis_tensor_ref(Av, convert(T, α), lenA, strideA, isconj(Av, conjA)) + tC = tblis_tensor_ref(Cv, convert(T, β), lenC, strideC, false) TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC)) end return C @@ -199,9 +217,9 @@ function tblis_mult!( 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) + tA = tblis_tensor_ref(A, α, lenA, strideA, false) + tB = tblis_tensor_ref(B, one(T), lenB, strideB, false) + tC = tblis_tensor_ref(C, β, lenC, strideC, false) TBLIS.tblis_tensor_mult( C_NULL, C_NULL, tA, labels(einA), tB, labels(einB), tC, labels(einC) ) diff --git a/src/backends.jl b/src/backends.jl index 49f622b6..9acc836b 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -88,6 +88,11 @@ 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. +!!! warning + This backend is not supported on Windows: the `tblis_jll` binaries for that platform + abort the process from inside the library, so calls are rejected with an + `ArgumentError` there as well. + !!! note Contracting in place trades throughput for memory. On a 16-core node this backend is roughly on par with diff --git a/test/runtests.jl b/test/runtests.jl index 3a26e37c..10f2728f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -81,8 +81,21 @@ if !is_buildkite include("butensor.jl") end - @testset "TBLIS extension" verbose = true begin - include("tblis.jl") + if Sys.iswindows() + # the Windows build of `tblis_jll` aborts the process from inside the library, so + # the extension refuses to call into it and there is nothing else to test + using TBLIS + @testset "TBLIS extension (unsupported platform)" begin + A = randn(Float64, (3, 4)) + @test_throws ArgumentError tensoradd!( + zeros(Float64, (4, 3)), A, ((2, 1), ()), false, 1, 0, + TensorOperations.TBLISBackend() + ) + end + else + @testset "TBLIS extension" verbose = true begin + include("tblis.jl") + end end @testset "Polynomials" begin From 9e7a7ac4fd69508fa8e3e8fdb9fc6b13484fd0b6 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 13 Aug 2026 03:29:16 -0400 Subject: [PATCH 7/7] code review --- .github/workflows/ci.yml | 2 -- ext/TensorOperationsTBLISExt.jl | 3 ++- src/backends.jl | 33 +++++++++------------------------ 3 files changed, 11 insertions(+), 27 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9c080b81..356ce20a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,8 +17,6 @@ jobs: test: name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} runs-on: ${{ matrix.os }} - # a native library that aborts or deadlocks otherwise wedges the runner for the full - # 6h default; the slowest passing job takes about 45 minutes timeout-minutes: 90 strategy: fail-fast: false diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index e2cdf13f..8e44c45c 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -13,9 +13,10 @@ using TensorOperations: tensoralloc_add, tensorfree! using TBLIS using TBLIS: len_type, stride_type, tblis_tensor +using LinearAlgebra.BLAS: BlasFloat const SV = StridedView -const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} +const TBLISFloat = BlasFloat # `tblis_jll` does ship a Windows build, but it aborts the process with # "posix_memalign: Invalid argument" from `tblis::MemoryPool::acquire` as soon as a diff --git a/src/backends.jl b/src/backends.jl index 9acc836b..d516e618 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -76,35 +76,20 @@ const StridedBackend = Union{StridedNative, StridedBLAS} """ 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. +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). -!!! warning - This backend is not supported on Windows: the `tblis_jll` binaries for that platform - abort the process from inside the library, so calls are rejected with an - `ArgumentError` there as well. +TBLIS requires all tensors in a single operation to be of the same element type, and only supports `BLASFloat`s. -!!! 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 + This backend is currently not supported on Windows: the `tblis_jll` binaries for that platform + abort the process from inside the library, so calls are rejected with an `ArgumentError`. !!! 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. + 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