Add a TBLIS.jl package extension

[lkdvos]

Adds TensorOperationsTBLISExt, a package extension that provides a TBLISBackend on top of TBLIS.jl. This supersedes the standalone TensorOperationsTBLIS.jl.

Important

This depends on TBLIS.jl v0.3, which is not registered yet (General has 0.1.0 and 0.2.0). CI will not resolve until it is released. Everything below was validated against QuantumKitHub/TBLIS.jl@main (977ce94).

What it does

TBLISBackend() routes tensoradd!, tensortrace! and tensorcontract! through TBLIS, which contracts strided tensors in place rather than reshaping them into matrices, so it avoids the permuted intermediates the BLAS path has to materialize.

It is opt-in: loading TBLIS.jl registers no select_backend method. Arguments TBLIS cannot express — mixed or unsupported element types, non-strided arrays, Diagonal factors — fall back to whatever select_backend would have chosen, so one backend = TBLISBackend() keeps working for a whole @tensor block.

using TensorOperations, TBLIS
TBLIS.set_num_threads(16)
@tensor backend = TensorOperations.TBLISBackend() D[a,b,c,d] := A[a,e,c,f] * B[g,d,e] * conj(C[g,f,b])

Things found while wiring this up

• tblis_tensor_mult ignores the per-tensor conjugation flags, while tblis_tensor_add honours them. Verified directly against tblis 1.3. Worked around by conjugating the output in place when both factors are conjugated, and by materializing a temporary from the allocator when only one is.
• TBLIS.jl's public tblis_tensor constructor is restricted to StridedArray, so it cannot take the StridedViews these kernels work with, and it offers no way to set the conjugation flag. The extension therefore initializes descriptors with tblis_init_tensor_scaled_*, patches the flag through the generated setproperty!, and keeps every referenced buffer rooted with GC.@preserve. Widening that constructor upstream (and letting it take a conjugation flag) would let this drop back to the public API.

Tests

test/tblis.jl, wired into runtests.jl. Compares against StridedNative over all four supported element types for add / trace / contract, every conjugation combination, non-contiguous views, outer products, scalar outputs, argument checking, the fallback paths, @tensor/ncon integration, a GC-pressure stress loop, and threading. All 20 testsets pass.

β == 0 is exercised with NaN-poisoned outputs: TBLIS does ignore C in that case, in both add and mult, so no NaN leaks.

Benchmarks

2× Xeon Gold 6244 (16 physical cores), 16 threads for both OpenBLAS and TBLIS, Julia 1.12.6, tblis_jll 1.3.0. Full tables in benchmarks/README.md; speedup = fastest built-in backend / TBLIS.

Float64:

case	StridedBLAS	TBLIS	speedup
matmul 2000³	33.65 ms	44.56 ms	0.76×
contract aligned (48,48,1024)×(1024,48,48)	20.86 ms	29.15 ms	0.72×
contract permuted (48,1024,48)×(48,1024,48)	56.59 ms	27.76 ms	2.04×
MPS·MPO·env (χ=256, d=4, D=8)	22.59 ms	11.61 ms	1.95×
permuted add 200³	35.50 ms	3.83 ms	9.27×
partial trace (128,32,128,32)	270.4 μs	84.8 μs	3.19×
contract small 8³×8³	13.8 μs	58.1 μs	0.24×

So: permutation-heavy contractions ~2× faster, adds and traces 3-9× faster, plain GEMM ~25% slower, small tensors dominated by descriptor overhead.

ComplexF64 is a different story — contractions run at 0.05–0.15× of BLAS. That is the tblis binaries, not this extension: a bare tblis_tensor_mult on a 1000³ complex matrix product reaches ~25 GFLOP/s against ~500 GFLOP/s for mul! (ComplexF32 is worse still, ~12 vs ~1160 GFLOP/s), and tblis_jll v1.2.0 behaves identically, so it is not a regression from the v1.3 bump. Complex add/trace are unaffected (6.6× and 4.4× wins). This is called out in the TBLISBackend docstring and in the docs.

🤖 Generated with Claude Code