From 85dc10dcb1db16a2be5b2cceb7b426716f711f80 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 10:44:34 -0400 Subject: [PATCH 1/5] 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 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..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 6075b9479591fac6d92783e6d9c056a8a48d685b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 11:03:17 -0400 Subject: [PATCH 2/5] 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 257efc518c83a9aba72a044ce82d98887fa13e3b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 12:01:49 -0400 Subject: [PATCH 3/5] 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 3885eb9c9a4ba3ebaa917e64788f2bca9953e076 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 12:13:39 -0400 Subject: [PATCH 4/5] 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 581c2b6add15d9717b6e3d0d6158bebd062ffb6a Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 12:36:03 -0400 Subject: [PATCH 5/5] 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(β′))