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..dfb17cbd --- /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³ | 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. 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..e14a9299 --- /dev/null +++ b/ext/TensorOperationsTBLISExt.jl @@ -0,0 +1,277 @@ +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, tblis_tensor + +const SV = StridedView + +# TBLIS only knows about these four element types, and requires all tensors taking part in a +# single operation to share it. +const TBLISFloat = Union{Float32, Float64, ComplexF32, ComplexF64} + +#------------------------------------------------------------------------------------------- +# Wrapping Julia arrays as TBLIS tensors +#------------------------------------------------------------------------------------------- +# TBLIS.jl's own `tblis_tensor` constructor is restricted to `StridedArray`, which excludes +# the `StridedView`s that the tensor operations work with, and it has no way to express the +# conjugation flag. We therefore initialize the descriptor through the low-level bindings and +# patch in the flag afterwards. +for (T, init) in ( + (:Float32, :tblis_init_tensor_scaled_s), + (:Float64, :tblis_init_tensor_scaled_d), + (:ComplexF32, :tblis_init_tensor_scaled_c), + (:ComplexF64, :tblis_init_tensor_scaled_z), + ) + @eval function init_tensor!( + p::Ptr{tblis_tensor}, A::StridedView{$T, N}, α::$T, + len::Vector{len_type}, stride::Vector{stride_type} + ) where {N} + return TBLIS.$init(p, α, Cuint(N), pointer(len), pointer(A), pointer(stride)) + end +end + +""" + TBLISTensor(A::StridedView, α, isconj) + +Owning counterpart of `TBLIS.tblis_tensor`, which itself only stores raw pointers into the +array it views and into the buffers holding its lengths and strides. Rooting a `TBLISTensor` +with `GC.@preserve` also roots all of those, as they are reachable from it; pass its `ref` +field wherever the library expects a `Ptr{tblis_tensor}`. +""" +struct TBLISTensor{T, N, A <: StridedView{T, N}} + array::A + len::Vector{len_type} + stride::Vector{stride_type} + ref::Base.RefValue{tblis_tensor} +end + +function TBLISTensor( + A::StridedView{T, N}, α::T, isconj::Bool = hasconj(A) + ) where {T <: TBLISFloat, N} + len = collect(len_type, size(A)) + stride = collect(stride_type, strides(A)) + ref = Ref{tblis_tensor}() + GC.@preserve A len stride ref begin + p = Base.unsafe_convert(Ptr{tblis_tensor}, ref) + init_tensor!(p, A, α, len, stride) + isconj && setproperty!(p, :conj, Cint(1)) + end + return TBLISTensor{T, N, typeof(A)}(A, len, stride, ref) +end + +# A `StridedView` 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) + +# 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.tblis_tensor_add( + C_NULL, C_NULL, tA.ref, labels(einA), tC.ref, labels(einC) + ) + end + return C +end + +function tblis_contract!( + C::StridedView{T}, + A::StridedView{T}, pA::Index2Tuple, + 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.tblis_tensor_mult( + C_NULL, C_NULL, tA.ref, labels(einA), tB.ref, labels(einB), + tC.ref, labels(einC) + ) + end + return C +end + +# Write `α * conj(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