TensorOperationsTBLIS.jl

Julia wrapper for tblis with TensorOperations.jl.

Currently provides implementations of tensorcontract!, tensoradd! and tensortrace! for StridedArray{<:BlasFloat}. These can be accessed through the backend system of TensorOperations, i.e.

using TensorOperations
using TensorOperationsTBLIS

α = randn()
A = randn(5, 5, 5, 5, 5, 5)
B = randn(5, 5, 5)
C = randn(5, 5, 5)
D = zeros(5, 5, 5)

@tensor backend = tblis begin
    D2[a, b, c] = A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b]
    E2[a, b, c] := A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b]
end

Additionally, the number of threads used by tblis can be set by:

using TensorOperationsTBLIS
tblis_set_num_threads(4)
@show tblis_get_num_threads()

Notes

• This implementation of tblis for TensorOperations.jl is only supported from v4.0.0 of TensorOperations.jl onwards. For older versions, you could look for BliContractor.jl or TBLIS.jl.