Add a TBLIS.jl package extension

.github/workflows/ci.yml

@@ -17,6 +17,7 @@ jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}
+    timeout-minutes: 90
     strategy:
       fail-fast: false
       matrix:
@@ -49,6 +50,7 @@ jobs:
     needs: test
     name: Julia nightly - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}
+    timeout-minutes: 90
     strategy:
       fail-fast: false
       matrix:

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

docs/src/man/backends.md

@@ -64,6 +64,7 @@ TensorOperations.BaseCopy
 TensorOperations.BaseView
 TensorOperations.StridedNative
 TensorOperations.StridedBLAS
+TensorOperations.TBLISBackend
 TensorOperations.cuTENSORBackend
 ```
 
@@ -74,6 +75,29 @@ 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.
+
+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.
+
+!!! 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 +113,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,240 @@
+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
+
+using LinearAlgebra.BLAS: BlasFloat
+const SV = StridedView
+
+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
+# contraction reaches the GEMM kernels, so nothing is handed to the library there.
+const PLATFORM_SUPPORTED = !Sys.iswindows()
+
+#-------------------------------------------------------------------------------------------
+# 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_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.
+
+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_ref(
+        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_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(
+            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...)
+    PLATFORM_SUPPORTED || throw_unsupported_platform(f)
+    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_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
+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_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
+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_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)
+        )
+    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,28 @@ 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 be of the same element type, and only supports `BLASFloat`s.
+
+!!! 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.
+"""
+struct TBLISBackend <: AbstractBackend end
+
 # CuTENSOR backend
 #-----------------
 """

test/runtests.jl

@@ -81,6 +81,23 @@ if !is_buildkite
         include("butensor.jl")
     end
 
+    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
         include("polynomials.jl")
     end