Add a TBLIS.jl package extension

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"]

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
 

ext/TensorOperationsTBLISExt.jl

@@ -0,0 +1,221 @@
+module TensorOperationsTBLISExt
+
+using TensorOperations
+using TensorOperations: TensorOperations as TO
+using TensorOperations: TBLISBackend, DefaultAllocator, Index2Tuple
+using TensorOperations: StridedView, isstrided
+using TensorOperations: argcheck_tensoradd, dimcheck_tensoradd,
+    argcheck_tensortrace, dimcheck_tensortrace,
+    argcheck_tensorcontract, dimcheck_tensorcontract
+using TensorOperations: add_labels, trace_labels, contract_labels
+using TensorOperations: tensoralloc_add, tensorfree!
+
+using TBLIS
+using TBLIS: len_type, stride_type, tblis_tensor
+
+const SV = StridedView
+
+const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64}
+
+#-------------------------------------------------------------------------------------------
+# Wrapping Julia arrays as TBLIS tensors
+#-------------------------------------------------------------------------------------------
+for (T, init) in (
+        (:Float32, :tblis_init_tensor_scaled_s),
+        (:Float64, :tblis_init_tensor_scaled_d),
+        (:ComplexF32, :tblis_init_tensor_scaled_c),
+        (:ComplexF64, :tblis_init_tensor_scaled_z),
+    )
+    @eval function init_tensor!(
+            p::Ptr{tblis_tensor}, A::StridedView{$T, N}, α::$T,
+            len::Vector{len_type}, stride::Vector{stride_type}
+        ) where {N}
+        return TBLIS.$init(p, α, Cuint(N), pointer(len), pointer(A), pointer(stride))
+    end
+end
+
+isconj(A::StridedView{T}, conjA::Bool) where {T} = T <: Complex && (conjA ⊻ (A.op === conj))
+tblis_dims(A::StridedView) = (collect(len_type, size(A)), collect(stride_type, strides(A)))
+
+"""
+    tblis_tensor(A::StridedView, α, len, stride, conj) -> Ref{TBLIS.tblis_tensor}
+
+Descriptor for `α * A`, conjugated when `conj` is set, using `len` and `stride` as the
+buffers handed to TBLIS.
+
+The descriptor only stores raw pointers into `A`, `len` and `stride`, so all three, along
+with the returned `Ref`, have to be kept alive by the caller for as long as TBLIS may access
+them. Note that `conj` is the *total* conjugation applied to the data of `A`, as computed by
+[`isconj`](@ref), not the flag the caller was handed.
+"""
+function tblis_tensor(
+        A::StridedView{T, N}, α::T,
+        len::Vector{len_type}, stride::Vector{stride_type}, conj::Bool
+    ) where {T <: TBLISFloat, N}
+    ref = Ref{tblis_tensor}()
+    GC.@preserve A len stride ref begin
+        p = Base.unsafe_convert(Ptr{tblis_tensor}, ref)
+        init_tensor!(p, A, α, len, stride)
+        conj && setproperty!(p, :conj, Cint(1))
+    end
+    return ref
+end
+
+labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein])
+
+#-------------------------------------------------------------------------------------------
+# Argument checking
+#-------------------------------------------------------------------------------------------
+@noinline function throw_eltype(f, tensors)
+    return throw(
+        ArgumentError(
+            LazyString(
+                "TBLISBackend requires all tensors of ", f, " to share a single element ",
+                "type out of Float32, Float64, ComplexF32 and ComplexF64, got ",
+                join(eltype.(tensors), ", ")
+            )
+        )
+    )
+end
+
+@noinline function throw_strided(f, tensors)
+    types = join(typeof.(tensors), ", ")
+    return throw(ArgumentError(lazy"TBLISBackend requires strided arrays for $f, got $types"))
+end
+
+@noinline throw_conj_output(f) = throw(
+    ArgumentError(lazy"TBLISBackend cannot write into a conjugated view in $f")
+)
+
+function check_arguments(f, C::AbstractArray, As::AbstractArray...)
+    tensors = (C, As...)
+    T = eltype(C)
+    (T <: TBLISFloat && all(A -> eltype(A) === T, As)) || throw_eltype(f, tensors)
+    all(isstrided, tensors) || throw_strided(f, tensors)
+    # `tblis_tensor_add` applies the flag of `C` when reading `β * C` but not when writing back
+    isconj(SV(C), false) && throw_conj_output(f)
+    return nothing
+end
+
+#-------------------------------------------------------------------------------------------
+# Operations
+#-------------------------------------------------------------------------------------------
+function TO.tensoradd!(
+        C::AbstractArray,
+        A::AbstractArray, pA::Index2Tuple, conjA::Bool,
+        α::Number, β::Number,
+        backend::TBLISBackend, allocator = DefaultAllocator()
+    )
+    check_arguments(TO.tensoradd!, C, A)
+    argcheck_tensoradd(C, A, pA)
+    dimcheck_tensoradd(C, A, pA)
+    Base.mightalias(C, A) &&
+        throw(ArgumentError("output tensor must not be aliased with input tensor"))
+
+    T = eltype(C)
+    einA, einC = add_labels(pA)
+    Av, Cv = SV(A), SV(C)
+    lenA, strideA = tblis_dims(Av)
+    lenC, strideC = tblis_dims(Cv)
+    GC.@preserve Av Cv lenA strideA lenC strideC begin
+        tA = tblis_tensor(Av, convert(T, α), lenA, strideA, isconj(Av, conjA))
+        tC = tblis_tensor(Cv, convert(T, β), lenC, strideC, false)
+        TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC))
+    end
+    return C
+end
+
+function TO.tensortrace!(
+        C::AbstractArray,
+        A::AbstractArray, p::Index2Tuple, q::Index2Tuple, conjA::Bool,
+        α::Number, β::Number,
+        backend::TBLISBackend, allocator = DefaultAllocator()
+    )
+    check_arguments(TO.tensortrace!, C, A)
+    argcheck_tensortrace(C, A, p, q)
+    dimcheck_tensortrace(C, A, p, q)
+    Base.mightalias(C, A) &&
+        throw(ArgumentError("output tensor must not be aliased with input tensor"))
+
+    T = eltype(C)
+    einA, einC = trace_labels(p, q)
+    Av, Cv = SV(A), SV(C)
+    lenA, strideA = tblis_dims(Av)
+    lenC, strideC = tblis_dims(Cv)
+    GC.@preserve Av Cv lenA strideA lenC strideC begin
+        tA = tblis_tensor(Av, convert(T, α), lenA, strideA, isconj(Av, conjA))
+        tC = tblis_tensor(Cv, convert(T, β), lenC, strideC, false)
+        TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC))
+    end
+    return C
+end
+
+function TO.tensorcontract!(
+        C::AbstractArray,
+        A::AbstractArray, pA::Index2Tuple, conjA::Bool,
+        B::AbstractArray, pB::Index2Tuple, conjB::Bool,
+        pAB::Index2Tuple,
+        α::Number, β::Number,
+        backend::TBLISBackend, allocator = DefaultAllocator()
+    )
+    check_arguments(TO.tensorcontract!, C, A, B)
+    argcheck_tensorcontract(C, A, pA, B, pB, pAB)
+    dimcheck_tensorcontract(C, A, pA, B, pB, pAB)
+    (Base.mightalias(C, A) || Base.mightalias(C, B)) &&
+        throw(ArgumentError("output tensor must not be aliased with input tensor"))
+
+    T = eltype(C)
+    einA, einB, einC = contract_labels(pA, pB, pAB)
+    α′ = convert(T, α)
+    β′ = convert(T, β)
+    Av, Bv, Cv = SV(A), SV(B), SV(C)
+    isconjA = isconj(Av, conjA)
+    isconjB = isconj(Bv, conjB)
+
+    # `tblis_tensor_mult` ignores the conjugation flags, so resolve them into the data first
+    if isconjA && isconjB
+        iszero(β′) || conj!(Cv)
+        tblis_mult!(Cv, Av, Bv, einA, einB, einC, conj(α′), conj(β′))
+        conj!(Cv)
+    elseif isconjA
+        A′ = materialize_conj(Av, conjA, α′, allocator)
+        tblis_mult!(Cv, SV(A′), Bv, einA, einB, einC, one(T), β′)
+        tensorfree!(A′, allocator)
+    elseif isconjB
+        B′ = materialize_conj(Bv, conjB, one(T), allocator)
+        tblis_mult!(Cv, Av, SV(B′), einA, einB, einC, α′, β′)
+        tensorfree!(B′, allocator)
+    else
+        tblis_mult!(Cv, Av, Bv, einA, einB, einC, α′, β′)
+    end
+    return C
+end
+
+function tblis_mult!(
+        C::StridedView{T}, A::StridedView{T}, B::StridedView{T},
+        einA, einB, einC, α::T, β::T
+    ) where {T <: TBLISFloat}
+    lenA, strideA = tblis_dims(A)
+    lenB, strideB = tblis_dims(B)
+    lenC, strideC = tblis_dims(C)
+    GC.@preserve A B C lenA strideA lenB strideB lenC strideC begin
+        tA = tblis_tensor(A, α, lenA, strideA, false)
+        tB = tblis_tensor(B, one(T), lenB, strideB, false)
+        tC = tblis_tensor(C, β, lenC, strideC, false)
+        TBLIS.tblis_tensor_mult(
+            C_NULL, C_NULL, tA, labels(einA), tB, labels(einB), tC, labels(einC)
+        )
+    end
+    return C
+end
+
+function materialize_conj(
+        A::StridedView{T, N}, conjA::Bool, α::T, allocator
+    ) where {T <: TBLISFloat, N}
+    pA = (ntuple(identity, N), ())
+    A′ = tensoralloc_add(T, A, pA, false, Val(true), allocator)
+    TO.tensoradd!(A′, A, pA, conjA, α, zero(T), TBLISBackend(), allocator)
+    return A′
+end
+
+end # module TensorOperationsTBLISExt

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
 #-----------------
 """

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