Skip to content

QuantumKitHub/TensorOperationsTBLIS.jl

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

14 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

TensorOperationsTBLIS.jl

Julia wrapper for tblis with TensorOperations.jl.

CI CI (Julia nightly)

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

tblisbackend = tblisBackend()
α = 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:

using TensorOperationsTBLIS
tblis_set_num_threads(4)
@show tblis_get_num_threads()

Notes

  • 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.