Julia wrapper for TBLIS with TensorOperations.jl.
Currently provides implementations of tensorcontract!
, tensoradd!
and tensortrace!
for array types compatible with Strided.jl, i.e. StridedView{<:BlasFloat}
.
These can be accessed through the backend system of TensorOperations, i.e.
using TensorOperations
using TensorOperationsTBLIS
tblisbackend = TBLIS()
α = 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 = tblisbackend 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:
TensorOperationsTBLIS.set_num_threads(4)
@show TensorOperationsTBLIS.get_num_threads()
- This implementation of TBLIS for TensorOperations.jl is only supported from v5 of TensorOperations.jl onwards. For v4, an earlier version of this package exists. For older versions, you could look for BliContractor.jl or TBLIS.jl.