From 9d349c9879d170e75d02d0375b87cd04a97066ca Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 1 Aug 2026 22:38:12 -0400 Subject: [PATCH 1/2] 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 any `select_backend` method. Arguments TBLIS cannot express -- mixed or unsupported element types, non-strided arrays, `Diagonal` factors -- are delegated back to the regular backend selection, so a single `backend = TBLISBackend()` keeps working for a whole `@tensor` block. Two library quirks are worked around in the extension: * `tblis_tensor_mult` ignores the per-tensor conjugation flags (unlike `tblis_tensor_add`, which honours them). When both factors are conjugated this is resolved by conjugating the output in place; when only one is, that factor is materialized into a temporary from the allocator. * TBLIS.jl only exposes its tensor constructor for `StridedArray`, and offers no way to set the conjugation flag or to own the length/stride buffers, so the extension builds the `tblis_tensor` struct itself and keeps every referenced buffer rooted with `GC.@preserve`. Includes a benchmark script and measurements under `benchmarks/`. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 6 +- benchmarks/Project.toml | 11 ++ benchmarks/README.md | 61 ++++++++ benchmarks/tblis_backend.jl | 243 +++++++++++++++++++++++++++++ docs/src/man/backends.md | 14 +- ext/TensorOperationsTBLISExt.jl | 269 ++++++++++++++++++++++++++++++++ src/backends.jl | 24 +++ test/runtests.jl | 4 + test/tblis.jl | 235 ++++++++++++++++++++++++++++ 9 files changed, 865 insertions(+), 2 deletions(-) create mode 100644 benchmarks/Project.toml create mode 100644 benchmarks/README.md create mode 100644 benchmarks/tblis_backend.jl create mode 100644 ext/TensorOperationsTBLISExt.jl create mode 100644 test/tblis.jl diff --git a/Project.toml b/Project.toml index caa37aae..bddf0b92 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/benchmarks/Project.toml b/benchmarks/Project.toml new file mode 100644 index 00000000..aca538d4 --- /dev/null +++ b/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" diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 00000000..65b70f27 --- /dev/null +++ b/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³ | 34.25 ms | 4.64 s | 49.22 ms | 0.70× | +| contract aligned (48,48,1024)×(1024,48,48) | 23.56 ms | 3.76 s | 33.13 ms | 0.71× | +| contract permuted (48,1024,48)×(48,1024,48) | 61.15 ms | 3.57 s | 31.42 ms | 1.95× | +| MPS·MPO·env (χ=256, d=4, D=8) | 25.88 ms | 698.69 ms | 13.87 ms | 1.87× | +| permuted add 200³ | 38.83 ms | 38.44 ms | 4.15 ms | 9.26× | +| partial trace (128,32,128,32) | 280.4 μs | 266.4 μs | 99.7 μs | 2.67× | +| contract small 8³×8³ | 17.1 μs | 40.5 μs | 56.5 μs | 0.30× | + +`eltype = ComplexF64` + +| case | StridedBLAS | StridedNative | TBLIS | speedup | +| :-------------------------------------------- | ----------: | ------------: | --------: | ------: | +| matmul 2000³ | 133.10 ms | 9.57 s | 2.84 s | 0.05× | +| contract aligned (48,48,1024)×(1024,48,48) | 86.58 ms | 6.69 s | 1.78 s | 0.05× | +| contract permuted (48,1024,48)×(48,1024,48) | 195.12 ms | 6.28 s | 1.79 s | 0.11× | +| MPS·MPO·env (χ=256, d=4, D=8) | 58.09 ms | 1.27 s | 384.48 ms | 0.15× | +| permuted add 200³ | 86.81 ms | 87.04 ms | 16.50 ms | 5.26× | +| partial trace (128,32,128,32) | 527.0 μs | 463.7 μs | 120.9 μs | 3.84× | +| contract small 8³×8³ | 26.8 μs | 62.6 μs | 73.3 μs | 0.37× | +| contract conj(A) (48,1024,48)×(48,1024,48) | 174.06 ms | 6.60 s | 1.75 s | 0.10× | +| contract conj(A)conj(B) (48,1024,48)×(48,1024,48) | 190.03 ms | 7.26 s | 1.79 s | 0.11× | + +`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 ~30% 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. diff --git a/benchmarks/tblis_backend.jl b/benchmarks/tblis_backend.jl new file mode 100644 index 00000000..18efe6b7 --- /dev/null +++ b/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 diff --git a/docs/src/man/backends.md b/docs/src/man/backends.md index 6ba5fd8c..65a2cf2f 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,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 diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl new file mode 100644 index 00000000..50b9cb1b --- /dev/null +++ b/ext/TensorOperationsTBLISExt.jl @@ -0,0 +1,269 @@ +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!, select_backend + +using TBLIS +using TBLIS: len_type, stride_type + +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_type(::Type{Float32}) = TBLIS.TYPE_SINGLE +tblis_type(::Type{Float64}) = TBLIS.TYPE_DOUBLE +tblis_type(::Type{ComplexF32}) = TBLIS.TYPE_SCOMPLEX +tblis_type(::Type{ComplexF64}) = TBLIS.TYPE_DCOMPLEX + +""" + 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. + +Use [`unsafe_tblis_tensor`](@ref) to obtain the plain C struct that the library expects, and +keep the `TBLISTensor` alive for at least as long as TBLIS may access it. +""" +struct TBLISTensor{T, N, A <: StridedView{T, N}} + array::A + len::Vector{len_type} + stride::Vector{stride_type} + scalar::T + isconj::Bool +end + +function TBLISTensor(A::StridedView{T, N}, α::T, isconj::Bool = hasconj(A)) where {T, N} + len = collect(len_type, size(A)) + stride = collect(stride_type, strides(A)) + return TBLISTensor{T, N, typeof(A)}(A, len, stride, α, isconj) +end + +# A `StridedView` carries its conjugation in an `op` field rather than in the data, which is +# exactly the representation TBLIS wants. 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) + +function unsafe_tblis_tensor(t::TBLISTensor{T, N}) where {T, N} + return TBLIS.tblis_tensor( + tblis_type(T), Cint(t.isconj), TBLIS.tblis_scalar(t.scalar), + Ptr{Cvoid}(pointer(t.array)), Cuint(N), + pointer(t.len), pointer(t.stride) + ) +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]) + +#------------------------------------------------------------------------------------------- +# Entry points: force the strided implementation on AbstractArray with the TBLIS backend +#------------------------------------------------------------------------------------------- +# Anything TBLIS cannot express -- mixed or unsupported element types, non-strided arrays, +# `Diagonal` factors -- is handed back to the regular backend selection mechanism rather than +# erroring, so that a single `backend = TBLISBackend()` keeps working for a whole `@tensor` +# block in which only some of the operations are TBLIS-compatible. +tblis_compatible(::AbstractArray...) = false +function tblis_compatible(A::AbstractArray{T}, As::AbstractArray{T}...) where {T} + return T <: TBLISFloat && isstrided(A) && all(isstrided, As) +end + +function TO.tensoradd!( + C::AbstractArray, + A::AbstractArray, pA::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::TBLISBackend, allocator = DefaultAllocator() + ) + if !tblis_compatible(C, A) + fallback = select_backend(TO.tensoradd!, C, A) + return TO.tensoradd!(C, A, pA, conjA, α, β, fallback, allocator) + end + + # resolve conj flags and absorb into StridedView constructor to avoid type instabilities + if conjA + tblis_add!(SV(C), conj(SV(A)), pA, α, β) + else + tblis_add!(SV(C), SV(A), pA, α, β) + end + return C +end + +function TO.tensortrace!( + C::AbstractArray, + A::AbstractArray, p::Index2Tuple, q::Index2Tuple, conjA::Bool, + α::Number, β::Number, + backend::TBLISBackend, allocator = DefaultAllocator() + ) + if !tblis_compatible(C, A) + fallback = select_backend(TO.tensortrace!, C, A) + return TO.tensortrace!(C, A, p, q, conjA, α, β, fallback, allocator) + end + + if conjA + tblis_trace!(SV(C), conj(SV(A)), p, q, α, β) + else + tblis_trace!(SV(C), SV(A), p, q, α, β) + 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() + ) + if !tblis_compatible(C, A, B) + fallback = select_backend(TO.tensorcontract!, C, A, B) + return TO.tensorcontract!( + C, A, pA, conjA, B, pB, conjB, pAB, α, β, fallback, allocator + ) + end + + if conjA && conjB + tblis_contract!(SV(C), conj(SV(A)), pA, conj(SV(B)), pB, pAB, α, β, allocator) + elseif conjA + tblis_contract!(SV(C), conj(SV(A)), pA, SV(B), pB, pAB, α, β, allocator) + elseif conjB + tblis_contract!(SV(C), SV(A), pA, conj(SV(B)), pB, pAB, α, β, allocator) + else + tblis_contract!(SV(C), SV(A), pA, SV(B), pB, pAB, α, β, allocator) + end + 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}, 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, einA, einC, convert(T, α), convert(T, β)) +end + +function tblis_trace!( + C::StridedView{T}, A::StridedView{T}, 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, einA, einC, convert(T, α), convert(T, β)) +end + +function unsafe_add!( + C::StridedView{T}, A::StridedView{T}, einA, einC, α::T, β::T + ) where {T <: TBLISFloat} + tA = TBLISTensor(A, α) + tC = TBLISTensor(C, β) + GC.@preserve tA tC begin + tblis_tensor_add( + unsafe_tblis_tensor(tA), labels(einA), + unsafe_tblis_tensor(tC), labels(einC) + ) + end + return C +end + +function tblis_contract!( + C::StridedView{T}, + A::StridedView{T}, pA::Index2Tuple, + B::StridedView{T}, pB::Index2Tuple, + 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, β) + + # `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 hasconj(A) && hasconj(B) + # 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 hasconj(A) + A′ = materialize_conj(A, α′, allocator) + try + unsafe_mult!(C, SV(A′), B, einA, einB, einC, one(T), β′) + finally + tensorfree!(A′, allocator) + end + elseif hasconj(B) + 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 passed as `false` because `tblis_tensor_mult` ignores them + # and the callers have 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_tensor_mult( + unsafe_tblis_tensor(tA), labels(einA), + unsafe_tblis_tensor(tB), labels(einB), + unsafe_tblis_tensor(tC), labels(einC) + ) + end + 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. +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, einA, einC, α, zero(T)) + return A′ +end + +end # module TensorOperationsTBLISExt diff --git a/src/backends.jl b/src/backends.jl index 3b5a29aa..b89e8acf 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -71,6 +71,30 @@ 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, which 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`. Operations that do not +meet these requirements are silently delegated to the backend that +[`select_backend`](@ref TensorOperations.select_backend) would have picked. + +!!! note + The contraction kernels shipped by `tblis_jll` are an order of magnitude slower than BLAS + for complex element types, so this backend pays off mostly for real contractions that + require a permutation, and for additions and traces of any element type. See + `benchmarks/README.md` for measurements. +""" +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..0ec2f3df --- /dev/null +++ b/test/tblis.jl @@ -0,0 +1,235 @@ +using TensorOperations +using TensorOperations: TBLISBackend, StridedNative, BaseView +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) +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 "fallback for 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)) + Cref = zeros(T, (3, 5)) + @test tensorcontract!( + C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) == tensorcontract!( + Cref, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, + TensorOperations.DefaultBackend() + ) + 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)) + Cref = zeros(ComplexF64, (3, 5)) + @test tensorcontract!( + C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) ≈ tensorcontract!( + Cref, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1, 0, reference + ) + + # non-strided arrays + D = Diagonal(randn(Float64, 4)) + A = randn(Float64, (3, 4)) + C = zeros(Float64, (3, 4)) + Cref = zeros(Float64, (3, 4)) + @test tensorcontract!( + C, A, ((1,), (2,)), false, D, ((1,), (2,)), false, ((1, 2), ()), 1, 0, tblis + ) ≈ tensorcontract!( + Cref, A, ((1,), (2,)), false, D, ((1,), (2,)), false, ((1, 2), ()), 1, 0, reference + ) +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 357daa75da97611ba31b60d56e1c43806c1d7ddb Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 1 Aug 2026 22:53:03 -0400 Subject: [PATCH 2/2] Build the TBLIS descriptors through the low-level bindings TBLIS.jl v0.3 regenerated its C bindings, so `tblis_tensor` is now an opaque 64-byte blob with pointer-based accessors rather than a struct with individual fields, and its high-level constructor only accepts `StridedArray`. Initialize the descriptor with `tblis_init_tensor_scaled_*` instead -- which works for any `StridedView` -- patch in the conjugation flag through the generated `setproperty!`, and hand the resulting `Ref` to the low-level `tblis_tensor_add` and `tblis_tensor_mult`. Benchmark numbers re-measured against the same release candidate. Co-Authored-By: Claude Opus 5 (1M context) --- benchmarks/README.md | 34 +++++++++--------- ext/TensorOperationsTBLISExt.jl | 64 ++++++++++++++++++--------------- 2 files changed, 53 insertions(+), 45 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 65b70f27..dfb17cbd 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -21,27 +21,27 @@ Julia 1.12.6, `tblis_jll` v1.3.0. | case | StridedBLAS | StridedNative | TBLIS | speedup | | :----------------------------------------- | ----------: | ------------: | -------: | ------: | -| matmul 2000³ | 34.25 ms | 4.64 s | 49.22 ms | 0.70× | -| contract aligned (48,48,1024)×(1024,48,48) | 23.56 ms | 3.76 s | 33.13 ms | 0.71× | -| contract permuted (48,1024,48)×(48,1024,48) | 61.15 ms | 3.57 s | 31.42 ms | 1.95× | -| MPS·MPO·env (χ=256, d=4, D=8) | 25.88 ms | 698.69 ms | 13.87 ms | 1.87× | -| permuted add 200³ | 38.83 ms | 38.44 ms | 4.15 ms | 9.26× | -| partial trace (128,32,128,32) | 280.4 μs | 266.4 μs | 99.7 μs | 2.67× | -| contract small 8³×8³ | 17.1 μs | 40.5 μs | 56.5 μs | 0.30× | +| 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³ | 133.10 ms | 9.57 s | 2.84 s | 0.05× | -| contract aligned (48,48,1024)×(1024,48,48) | 86.58 ms | 6.69 s | 1.78 s | 0.05× | -| contract permuted (48,1024,48)×(48,1024,48) | 195.12 ms | 6.28 s | 1.79 s | 0.11× | -| MPS·MPO·env (χ=256, d=4, D=8) | 58.09 ms | 1.27 s | 384.48 ms | 0.15× | -| permuted add 200³ | 86.81 ms | 87.04 ms | 16.50 ms | 5.26× | -| partial trace (128,32,128,32) | 527.0 μs | 463.7 μs | 120.9 μs | 3.84× | -| contract small 8³×8³ | 26.8 μs | 62.6 μs | 73.3 μs | 0.37× | -| contract conj(A) (48,1024,48)×(48,1024,48) | 174.06 ms | 6.60 s | 1.75 s | 0.10× | -| contract conj(A)conj(B) (48,1024,48)×(48,1024,48) | 190.03 ms | 7.26 s | 1.79 s | 0.11× | +| 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. @@ -52,7 +52,7 @@ Julia 1.12.6, `tblis_jll` v1.3.0. 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 ~30% faster with OpenBLAS, and **small tensors** +* 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 diff --git a/ext/TensorOperationsTBLISExt.jl b/ext/TensorOperationsTBLISExt.jl index 50b9cb1b..e14a9299 100644 --- a/ext/TensorOperationsTBLISExt.jl +++ b/ext/TensorOperationsTBLISExt.jl @@ -11,7 +11,7 @@ using TensorOperations: add_labels, trace_labels, contract_labels using TensorOperations: tensoralloc_add, tensorfree!, select_backend using TBLIS -using TBLIS: len_type, stride_type +using TBLIS: len_type, stride_type, tblis_tensor const SV = StridedView @@ -22,33 +22,51 @@ const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} #------------------------------------------------------------------------------------------- # Wrapping Julia arrays as TBLIS tensors #------------------------------------------------------------------------------------------- -tblis_type(::Type{Float32}) = TBLIS.TYPE_SINGLE -tblis_type(::Type{Float64}) = TBLIS.TYPE_DOUBLE -tblis_type(::Type{ComplexF32}) = TBLIS.TYPE_SCOMPLEX -tblis_type(::Type{ComplexF64}) = TBLIS.TYPE_DCOMPLEX +# 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. - -Use [`unsafe_tblis_tensor`](@ref) to obtain the plain C struct that the library expects, and -keep the `TBLISTensor` alive for at least as long as TBLIS may access it. +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}`. """ struct TBLISTensor{T, N, A <: StridedView{T, N}} array::A len::Vector{len_type} stride::Vector{stride_type} - scalar::T - isconj::Bool + ref::Base.RefValue{tblis_tensor} end -function TBLISTensor(A::StridedView{T, N}, α::T, isconj::Bool = hasconj(A)) where {T, N} +function TBLISTensor( + A::StridedView{T, N}, α::T, isconj::Bool = hasconj(A) + ) where {T <: TBLISFloat, N} len = collect(len_type, size(A)) stride = collect(stride_type, strides(A)) - return TBLISTensor{T, N, typeof(A)}(A, len, stride, α, isconj) + 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` carries its conjugation in an `op` field rather than in the data, which is @@ -56,14 +74,6 @@ end # flagging it would only confuse the library. hasconj(A::StridedView{T}) where {T} = T <: Complex && (A.op === conj || A.op === adjoint) -function unsafe_tblis_tensor(t::TBLISTensor{T, N}) where {T, N} - return TBLIS.tblis_tensor( - tblis_type(T), Cint(t.isconj), TBLIS.tblis_scalar(t.scalar), - Ptr{Cvoid}(pointer(t.array)), Cuint(N), - pointer(t.len), pointer(t.stride) - ) -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]) @@ -183,9 +193,8 @@ function unsafe_add!( tA = TBLISTensor(A, α) tC = TBLISTensor(C, β) GC.@preserve tA tC begin - tblis_tensor_add( - unsafe_tblis_tensor(tA), labels(einA), - unsafe_tblis_tensor(tC), labels(einC) + TBLIS.tblis_tensor_add( + C_NULL, C_NULL, tA.ref, labels(einA), tC.ref, labels(einC) ) end return C @@ -247,10 +256,9 @@ function unsafe_mult!( tB = TBLISTensor(B, one(T), false) tC = TBLISTensor(C, β, false) GC.@preserve tA tB tC begin - tblis_tensor_mult( - unsafe_tblis_tensor(tA), labels(einA), - unsafe_tblis_tensor(tB), labels(einB), - unsafe_tblis_tensor(tC), labels(einC) + TBLIS.tblis_tensor_mult( + C_NULL, C_NULL, tA.ref, labels(einA), tB.ref, labels(einB), + tC.ref, labels(einC) ) end return C