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

benchmarks/Project.toml

@@ -0,0 +1,11 @@
+[deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+
+[compat]
+BenchmarkTools = "1"
+TBLIS = "0.3"

benchmarks/README.md

@@ -0,0 +1,61 @@
+# Benchmarks
+
+## `tblis_backend.jl`
+
+Compares `TensorOperations.TBLISBackend()` against the built-in `StridedBLAS()` and
+`StridedNative()` backends. BLAS and TBLIS are handed the same number of threads.
+
+```console
+$ julia --project=benchmarks -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()'
+$ julia --project=benchmarks benchmarks/tblis_backend.jl 16
+```
+
+The optional argument is the thread count, defaulting to half of `Sys.CPU_THREADS`.
+
+### Reference numbers
+
+2× Intel Xeon Gold 6244 (16 physical cores total), 16 threads for both OpenBLAS and TBLIS,
+Julia 1.12.6, `tblis_jll` v1.3.0.
+
+`eltype = Float64`
+
+| case                                       | StridedBLAS | StridedNative |    TBLIS | speedup |
+| :----------------------------------------- | ----------: | ------------: | -------: | ------: |
+| matmul 2000³                                |    33.65 ms |        4.44 s | 44.56 ms |   0.76× |
+| contract aligned (48,48,1024)×(1024,48,48)  |    20.86 ms |        3.61 s | 29.15 ms |   0.72× |
+| contract permuted (48,1024,48)×(48,1024,48) |    56.59 ms |        3.49 s | 27.76 ms |   2.04× |
+| MPS·MPO·env (χ=256, d=4, D=8)               |    22.59 ms |     685.24 ms | 11.61 ms |   1.95× |
+| permuted add 200³                           |    35.50 ms |      35.47 ms |  3.83 ms |   9.27× |
+| partial trace (128,32,128,32)               |    270.4 μs |      270.8 μs |  84.8 μs |   3.19× |
+| contract small 8³×8³                        |     13.8 μs |       36.6 μs |  58.1 μs |   0.24× |
+
+`eltype = ComplexF64`
+
+| case                                          | StridedBLAS | StridedNative |     TBLIS | speedup |
+| :-------------------------------------------- | ----------: | ------------: | --------: | ------: |
+| matmul 2000³                                   |   125.49 ms |        9.28 s |    2.55 s |   0.05× |
+| contract aligned (48,48,1024)×(1024,48,48)     |   100.07 ms |        6.40 s |    1.74 s |   0.06× |
+| contract permuted (48,1024,48)×(48,1024,48)    |   168.72 ms |        6.15 s |    1.76 s |   0.10× |
+| MPS·MPO·env (χ=256, d=4, D=8)                  |    56.65 ms |        1.24 s | 381.66 ms |   0.15× |
+| permuted add 200³                              |    81.03 ms |      81.47 ms |  12.37 ms |   6.55× |
+| partial trace (128,32,128,32)                  |    451.5 μs |      452.7 μs |  102.4 μs |   4.41× |
+| contract small 8³×8³                           |     23.9 μs |       61.5 μs |   68.0 μs |   0.35× |
+| contract conj(A) (48,1024,48)×(48,1024,48)     |   166.13 ms |        6.71 s |    1.75 s |   0.10× |
+| contract conj(A)conj(B) (48,1024,48)×(48,1024,48) | 181.13 ms |        6.84 s |    1.76 s |   0.10× |
+
+`speedup` is the fastest built-in backend divided by TBLIS, so values above 1 mean TBLIS wins.
+
+### Reading the numbers
+
+* Contractions that need a **permutation** before they can be handed to GEMM — the common
+  case in tensor-network codes — are roughly **2× faster** with TBLIS, because it contracts
+  the strided memory directly instead of materializing permuted copies.
+* **Additions and traces** are where TBLIS is furthest ahead (up to 9×), since `Strided.jl`
+  loses to a blocked, threaded copy for these memory-bound shapes.
+* Plain **GEMM-shaped** contractions stay ~25-40% faster with OpenBLAS, and **small tensors**
+  are dominated by the per-call overhead of building the TBLIS descriptors.
+* **Complex contractions are 10-20× slower than BLAS.** This is a property of the `tblis`
+  binaries rather than of this extension: a bare `tblis_tensor_mult` on a 1000³ complex
+  matrix product reaches ~25 GFLOP/s against ~500 GFLOP/s for `mul!`, and the same holds for
+  `tblis_jll` v1.2.0, so it is not a regression from the v1.3 bump. Complex `add`/`trace` are
+  unaffected. Prefer `StridedBLAS` for complex-valued contractions.

benchmarks/tblis_backend.jl

@@ -0,0 +1,243 @@
+# Compare the `TBLISBackend` against the built-in strided backends on a handful of
+# representative workloads.
+#
+#     julia --project=benchmarks benchmarks/tblis_backend.jl [nthreads]
+#
+# BLAS and TBLIS are given the same number of threads, so that the comparison is about the
+# algorithms rather than about the thread pools.
+
+using TensorOperations
+using TensorOperations: TBLISBackend, StridedBLAS, StridedNative
+using TBLIS
+using BenchmarkTools
+using LinearAlgebra
+using Printf
+using Random
+
+const NTHREADS = length(ARGS) ≥ 1 ? parse(Int, ARGS[1]) : max(1, Sys.CPU_THREADS ÷ 2)
+
+BLAS.set_num_threads(NTHREADS)
+TBLIS.set_num_threads(NTHREADS)
+
+const BACKENDS = (
+    "StridedBLAS" => StridedBLAS(),
+    "StridedNative" => StridedNative(),
+    "TBLIS" => TBLISBackend(),
+)
+
+Random.seed!(1234)
+
+# Each case prepares its tensors once and exposes a kernel that writes into a preallocated
+# output, so that the timings are not dominated by allocating the result.
+struct Case{F, T <: Tuple}
+    name::String
+    kernel::F
+    tensors::T
+end
+
+run!(c::Case, backend) = c.kernel(backend, c.tensors...)
+
+function benchmark(cases; seconds = 2.0)
+    results = Dict{String, Dict{String, Float64}}()
+    for c in cases
+        print("  ", rpad(c.name, 42))
+        results[c.name] = Dict{String, Float64}()
+        # sanity check: every backend has to agree with StridedBLAS on the answer
+        ref = copy(run!(c, StridedBLAS()))
+        for (label, backend) in BACKENDS
+            out = run!(c, backend)
+            isapprox(out, ref; rtol = 1.0e-6) ||
+                @warn "$(c.name): $label disagrees with StridedBLAS"
+            results[c.name][label] = @belapsed(
+                run!($c, $backend), samples = 10_000, seconds = seconds
+            )
+            print(".")
+        end
+        println()
+    end
+    return results
+end
+
+function prettytime(t)
+    t < 1.0e-6 && return @sprintf("%.1f ns", t * 1.0e9)
+    t < 1.0e-3 && return @sprintf("%.1f μs", t * 1.0e6)
+    t < 1.0 && return @sprintf("%.2f ms", t * 1.0e3)
+    return @sprintf("%.2f s", t)
+end
+
+function report(cases, results)
+    @printf(
+        "\n| %-40s | %12s | %13s | %12s | %8s |\n",
+        "case", "StridedBLAS", "StridedNative", "TBLIS", "speedup"
+    )
+    @printf("|%s|%s|%s|%s|%s|\n", "-"^42, "-"^14, "-"^15, "-"^14, "-"^10)
+    for c in cases
+        r = results[c.name]
+        best_builtin = min(r["StridedBLAS"], r["StridedNative"])
+        @printf(
+            "| %-40s | %12s | %13s | %12s | %7.2fx |\n", c.name,
+            prettytime(r["StridedBLAS"]), prettytime(r["StridedNative"]),
+            prettytime(r["TBLIS"]), best_builtin / r["TBLIS"]
+        )
+    end
+    println("\nspeedup = fastest built-in backend / TBLIS  (> 1 means TBLIS wins)")
+    return nothing
+end
+
+#-------------------------------------------------------------------------------------------
+# Cases
+#-------------------------------------------------------------------------------------------
+function cases(T)
+    out = Case[]
+
+    # plain matrix multiplication: no permutations at all, pure BLAS territory
+    let n = 2000
+        A, B, C = randn(T, n, n), randn(T, n, n), zeros(T, n, n)
+        push!(
+            out, Case(
+                "matmul $(n)³",
+                (bk, C, A, B) -> (@tensor backend = bk C[i, j] = A[i, k] * B[k, j]),
+                (C, A, B)
+            )
+        )
+    end
+
+    # contraction whose contracted index already sits at the matrix boundary: BLAS only has
+    # to reshape, not permute
+    let m = 48, k = 1024
+        A, B = randn(T, m, m, k), randn(T, k, m, m)
+        C = zeros(T, m, m, m, m)
+        push!(
+            out, Case(
+                "contract aligned ($m,$m,$k)×($k,$m,$m)",
+                (bk, C, A, B) -> (
+                    @tensor backend = bk C[a, b, c, d] = A[a, b, x] *
+                        B[x, c, d]
+                ),
+                (C, A, B)
+            )
+        )
+    end
+
+    # same flop count, but now both factors need a non-trivial permutation before they can
+    # be fed to GEMM, and the result needs one afterwards -- this is what TBLIS avoids
+    let m = 48, k = 1024
+        A, B = randn(T, m, k, m), randn(T, m, k, m)
+        C = zeros(T, m, m, m, m)
+        push!(
+            out, Case(
+                "contract permuted ($m,$k,$m)×($m,$k,$m)",
+                (bk, C, A, B) -> (
+                    @tensor backend = bk C[c, a, d, b] = A[a, x, b] *
+                        B[c, x, d]
+                ),
+                (C, A, B)
+            )
+        )
+    end
+
+    # a DMRG-style effective Hamiltonian application, evaluated pairwise by the macro
+    let χ = 256, d = 4, D = 8
+        E = randn(T, χ, D, χ)       # left environment
+        A = randn(T, χ, d, χ)       # MPS tensor
+        O = randn(T, D, d, D, d)    # MPO tensor
+        F = randn(T, χ, D, χ)       # right environment
+        C = zeros(T, χ, d, χ)
+        push!(
+            out, Case(
+                "MPS·MPO·env (χ=$χ, d=$d, D=$D)",
+                (bk, C, E, A, O, F) -> (
+                    @tensor backend = bk C[a, s, c] = E[a, x, y] *
+                        A[y, t, z] * O[x, s, w, t] * F[c, w, z]
+                ),
+                (C, E, A, O, F)
+            )
+        )
+    end
+
+    # additions and traces both go through `tblis_tensor_add`
+    let n = 200
+        A, C = randn(T, n, n, n), zeros(T, n, n, n)
+        push!(
+            out, Case(
+                "permuted add $(n)³",
+                (bk, C, A) -> (@tensor backend = bk C[c, a, b] = A[a, b, c]),
+                (C, A)
+            )
+        )
+    end
+    let n = 128, k = 32
+        A, C = randn(T, n, k, n, k), zeros(T, n, n)
+        push!(
+            out, Case(
+                "partial trace ($n,$k,$n,$k)",
+                (bk, C, A) -> (@tensor backend = bk C[a, b] = A[a, x, b, x]),
+                (C, A)
+            )
+        )
+    end
+
+    # small tensors, where the per-call overhead dominates
+    let n = 8
+        A, B = randn(T, n, n, n), randn(T, n, n, n)
+        C = zeros(T, n, n, n, n)
+        push!(
+            out, Case(
+                "contract small $(n)³×$(n)³",
+                (bk, C, A, B) -> (
+                    @tensor backend = bk C[c, a, d, b] = A[a, x, b] *
+                        B[c, x, d]
+                ),
+                (C, A, B)
+            )
+        )
+    end
+
+    if T <: Complex
+        # BLAS folds conjugation into GEMM; TBLIS ignores the conjugation flags in its
+        # `mult` kernel, so the extension has to work around it
+        let m = 48, k = 1024
+            A, B = randn(T, m, k, m), randn(T, m, k, m)
+            C = zeros(T, m, m, m, m)
+            push!(
+                out, Case(
+                    "contract conj(A) ($m,$k,$m)×($m,$k,$m)",
+                    (bk, C, A, B) -> (
+                        @tensor backend = bk C[c, a, d, b] =
+                            conj(A[a, x, b]) * B[c, x, d]
+                    ),
+                    (C, A, B)
+                )
+            )
+            push!(
+                out, Case(
+                    "contract conj(A)conj(B) ($m,$k,$m)×($m,$k,$m)",
+                    (bk, C, A, B) -> (
+                        @tensor backend = bk C[c, a, d, b] =
+                            conj(A[a, x, b]) * conj(B[c, x, d])
+                    ),
+                    (C, A, B)
+                )
+            )
+        end
+    end
+
+    return out
+end
+
+#-------------------------------------------------------------------------------------------
+# Main
+#-------------------------------------------------------------------------------------------
+println(
+    "threads: BLAS = $(BLAS.get_num_threads()), TBLIS = $(TBLIS.get_num_threads()), ",
+    "julia = $(Threads.nthreads())"
+)
+println("BLAS: ", BLAS.get_config())
+
+for T in (Float64, ComplexF64)
+    println("\n", "="^100)
+    println("eltype = $T")
+    println("="^100)
+    cs = cases(T)
+    report(cs, benchmark(cs))
+end

docs/src/man/backends.md

@@ -64,6 +64,7 @@ TensorOperations.BaseCopy
 TensorOperations.BaseView
 TensorOperations.StridedNative
 TensorOperations.StridedBLAS
+TensorOperations.TBLISBackend
 TensorOperations.cuTENSORBackend
 ```
 
@@ -74,6 +75,17 @@ 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, which 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 -- loading `TBLIS.jl` does not change the default backend selection -- and 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])
+```
+
+Because TBLIS requires all tensors in a single operation to share one element type out of `Float32`, `Float64`, `ComplexF32` and `ComplexF64`, operations that do not satisfy this -- as well as non-strided arrays -- are delegated to the backend that `select_backend` would have chosen, so that the backend can be set once for a whole `@tensor` block.
+
 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 +101,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