|
| 1 | +module TensorOperationsTBLISExt |
| 2 | + |
| 3 | +using TensorOperations |
| 4 | +using TensorOperations: TensorOperations as TO |
| 5 | +using TensorOperations: TBLISBackend, DefaultAllocator, Index2Tuple |
| 6 | +using TensorOperations: StridedView, isstrided |
| 7 | +using TensorOperations: argcheck_tensoradd, dimcheck_tensoradd, |
| 8 | + argcheck_tensortrace, dimcheck_tensortrace, |
| 9 | + argcheck_tensorcontract, dimcheck_tensorcontract |
| 10 | +using TensorOperations: add_labels, trace_labels, contract_labels |
| 11 | +using TensorOperations: tensoralloc_add, tensorfree! |
| 12 | + |
| 13 | +using TBLIS |
| 14 | +using TBLIS: len_type, stride_type, tblis_tensor |
| 15 | + |
| 16 | +using LinearAlgebra.BLAS: BlasFloat |
| 17 | +const SV = StridedView |
| 18 | + |
| 19 | +const TBLISFloat = BlasFloat |
| 20 | + |
| 21 | +# `tblis_jll` does ship a Windows build, but it aborts the process with |
| 22 | +# "posix_memalign: Invalid argument" from `tblis::MemoryPool::acquire` as soon as a |
| 23 | +# contraction reaches the GEMM kernels, so nothing is handed to the library there. |
| 24 | +const PLATFORM_SUPPORTED = !Sys.iswindows() |
| 25 | + |
| 26 | +#------------------------------------------------------------------------------------------- |
| 27 | +# Wrapping Julia arrays as TBLIS tensors |
| 28 | +#------------------------------------------------------------------------------------------- |
| 29 | +for (T, init) in ( |
| 30 | + (:Float32, :tblis_init_tensor_scaled_s), |
| 31 | + (:Float64, :tblis_init_tensor_scaled_d), |
| 32 | + (:ComplexF32, :tblis_init_tensor_scaled_c), |
| 33 | + (:ComplexF64, :tblis_init_tensor_scaled_z), |
| 34 | + ) |
| 35 | + @eval function init_tensor!( |
| 36 | + p::Ptr{tblis_tensor}, A::StridedView{$T, N}, α::$T, |
| 37 | + len::Vector{len_type}, stride::Vector{stride_type} |
| 38 | + ) where {N} |
| 39 | + return TBLIS.$init(p, α, Cuint(N), pointer(len), pointer(A), pointer(stride)) |
| 40 | + end |
| 41 | +end |
| 42 | + |
| 43 | +isconj(A::StridedView{T}, conjA::Bool) where {T} = T <: Complex && (conjA ⊻ (A.op === conj)) |
| 44 | +tblis_dims(A::StridedView) = (collect(len_type, size(A)), collect(stride_type, strides(A))) |
| 45 | + |
| 46 | +""" |
| 47 | + tblis_tensor_ref(A::StridedView, α, len, stride, conj) -> Ref{TBLIS.tblis_tensor} |
| 48 | +
|
| 49 | +Descriptor for `α * A`, conjugated when `conj` is set, using `len` and `stride` as the |
| 50 | +buffers handed to TBLIS. |
| 51 | +
|
| 52 | +The descriptor only stores raw pointers into `A`, `len` and `stride`, so all three, along |
| 53 | +with the returned `Ref`, have to be kept alive by the caller for as long as TBLIS may access |
| 54 | +them. Note that `conj` is the *total* conjugation applied to the data of `A`, as computed by |
| 55 | +[`isconj`](@ref), not the flag the caller was handed. |
| 56 | +""" |
| 57 | +function tblis_tensor_ref( |
| 58 | + A::StridedView{T, N}, α::T, |
| 59 | + len::Vector{len_type}, stride::Vector{stride_type}, conj::Bool |
| 60 | + ) where {T <: TBLISFloat, N} |
| 61 | + ref = Ref{tblis_tensor}() |
| 62 | + GC.@preserve A len stride ref begin |
| 63 | + p = Base.unsafe_convert(Ptr{tblis_tensor}, ref) |
| 64 | + init_tensor!(p, A, α, len, stride) |
| 65 | + conj && setproperty!(p, :conj, Cint(1)) |
| 66 | + end |
| 67 | + return ref |
| 68 | +end |
| 69 | + |
| 70 | +labels(ein::Tuple{Vararg{Char}}) = String(UInt8[c for c in ein]) |
| 71 | + |
| 72 | +#------------------------------------------------------------------------------------------- |
| 73 | +# Argument checking |
| 74 | +#------------------------------------------------------------------------------------------- |
| 75 | +@noinline function throw_unsupported_platform(f) |
| 76 | + return throw( |
| 77 | + ArgumentError( |
| 78 | + LazyString( |
| 79 | + "TBLISBackend is not supported on ", Base.BUILD_TRIPLET, ": the tblis_jll ", |
| 80 | + "binaries for this platform abort the process from inside the library. ", |
| 81 | + "Use another backend, such as StridedBLAS(), for ", f |
| 82 | + ) |
| 83 | + ) |
| 84 | + ) |
| 85 | +end |
| 86 | + |
| 87 | +@noinline function throw_eltype(f, tensors) |
| 88 | + return throw( |
| 89 | + ArgumentError( |
| 90 | + LazyString( |
| 91 | + "TBLISBackend requires all tensors of ", f, " to share a single element ", |
| 92 | + "type out of Float32, Float64, ComplexF32 and ComplexF64, got ", |
| 93 | + join(eltype.(tensors), ", ") |
| 94 | + ) |
| 95 | + ) |
| 96 | + ) |
| 97 | +end |
| 98 | + |
| 99 | +@noinline function throw_strided(f, tensors) |
| 100 | + types = join(typeof.(tensors), ", ") |
| 101 | + return throw(ArgumentError(lazy"TBLISBackend requires strided arrays for $f, got $types")) |
| 102 | +end |
| 103 | + |
| 104 | +@noinline throw_conj_output(f) = throw( |
| 105 | + ArgumentError(lazy"TBLISBackend cannot write into a conjugated view in $f") |
| 106 | +) |
| 107 | + |
| 108 | +function check_arguments(f, C::AbstractArray, As::AbstractArray...) |
| 109 | + PLATFORM_SUPPORTED || throw_unsupported_platform(f) |
| 110 | + tensors = (C, As...) |
| 111 | + T = eltype(C) |
| 112 | + (T <: TBLISFloat && all(A -> eltype(A) === T, As)) || throw_eltype(f, tensors) |
| 113 | + all(isstrided, tensors) || throw_strided(f, tensors) |
| 114 | + # `tblis_tensor_add` applies the flag of `C` when reading `β * C` but not when writing back |
| 115 | + isconj(SV(C), false) && throw_conj_output(f) |
| 116 | + return nothing |
| 117 | +end |
| 118 | + |
| 119 | +#------------------------------------------------------------------------------------------- |
| 120 | +# Operations |
| 121 | +#------------------------------------------------------------------------------------------- |
| 122 | +function TO.tensoradd!( |
| 123 | + C::AbstractArray, |
| 124 | + A::AbstractArray, pA::Index2Tuple, conjA::Bool, |
| 125 | + α::Number, β::Number, |
| 126 | + backend::TBLISBackend, allocator = DefaultAllocator() |
| 127 | + ) |
| 128 | + check_arguments(TO.tensoradd!, C, A) |
| 129 | + argcheck_tensoradd(C, A, pA) |
| 130 | + dimcheck_tensoradd(C, A, pA) |
| 131 | + Base.mightalias(C, A) && |
| 132 | + throw(ArgumentError("output tensor must not be aliased with input tensor")) |
| 133 | + |
| 134 | + T = eltype(C) |
| 135 | + einA, einC = add_labels(pA) |
| 136 | + Av, Cv = SV(A), SV(C) |
| 137 | + lenA, strideA = tblis_dims(Av) |
| 138 | + lenC, strideC = tblis_dims(Cv) |
| 139 | + GC.@preserve Av Cv lenA strideA lenC strideC begin |
| 140 | + tA = tblis_tensor_ref(Av, convert(T, α), lenA, strideA, isconj(Av, conjA)) |
| 141 | + tC = tblis_tensor_ref(Cv, convert(T, β), lenC, strideC, false) |
| 142 | + TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC)) |
| 143 | + end |
| 144 | + return C |
| 145 | +end |
| 146 | + |
| 147 | +function TO.tensortrace!( |
| 148 | + C::AbstractArray, |
| 149 | + A::AbstractArray, p::Index2Tuple, q::Index2Tuple, conjA::Bool, |
| 150 | + α::Number, β::Number, |
| 151 | + backend::TBLISBackend, allocator = DefaultAllocator() |
| 152 | + ) |
| 153 | + check_arguments(TO.tensortrace!, C, A) |
| 154 | + argcheck_tensortrace(C, A, p, q) |
| 155 | + dimcheck_tensortrace(C, A, p, q) |
| 156 | + Base.mightalias(C, A) && |
| 157 | + throw(ArgumentError("output tensor must not be aliased with input tensor")) |
| 158 | + |
| 159 | + T = eltype(C) |
| 160 | + einA, einC = trace_labels(p, q) |
| 161 | + Av, Cv = SV(A), SV(C) |
| 162 | + lenA, strideA = tblis_dims(Av) |
| 163 | + lenC, strideC = tblis_dims(Cv) |
| 164 | + GC.@preserve Av Cv lenA strideA lenC strideC begin |
| 165 | + tA = tblis_tensor_ref(Av, convert(T, α), lenA, strideA, isconj(Av, conjA)) |
| 166 | + tC = tblis_tensor_ref(Cv, convert(T, β), lenC, strideC, false) |
| 167 | + TBLIS.tblis_tensor_add(C_NULL, C_NULL, tA, labels(einA), tC, labels(einC)) |
| 168 | + end |
| 169 | + return C |
| 170 | +end |
| 171 | + |
| 172 | +function TO.tensorcontract!( |
| 173 | + C::AbstractArray, |
| 174 | + A::AbstractArray, pA::Index2Tuple, conjA::Bool, |
| 175 | + B::AbstractArray, pB::Index2Tuple, conjB::Bool, |
| 176 | + pAB::Index2Tuple, |
| 177 | + α::Number, β::Number, |
| 178 | + backend::TBLISBackend, allocator = DefaultAllocator() |
| 179 | + ) |
| 180 | + check_arguments(TO.tensorcontract!, C, A, B) |
| 181 | + argcheck_tensorcontract(C, A, pA, B, pB, pAB) |
| 182 | + dimcheck_tensorcontract(C, A, pA, B, pB, pAB) |
| 183 | + (Base.mightalias(C, A) || Base.mightalias(C, B)) && |
| 184 | + throw(ArgumentError("output tensor must not be aliased with input tensor")) |
| 185 | + |
| 186 | + T = eltype(C) |
| 187 | + einA, einB, einC = contract_labels(pA, pB, pAB) |
| 188 | + α′ = convert(T, α) |
| 189 | + β′ = convert(T, β) |
| 190 | + Av, Bv, Cv = SV(A), SV(B), SV(C) |
| 191 | + isconjA = isconj(Av, conjA) |
| 192 | + isconjB = isconj(Bv, conjB) |
| 193 | + |
| 194 | + # `tblis_tensor_mult` ignores the conjugation flags, so resolve them into the data first |
| 195 | + if isconjA && isconjB |
| 196 | + iszero(β′) || conj!(Cv) |
| 197 | + tblis_mult!(Cv, Av, Bv, einA, einB, einC, conj(α′), conj(β′)) |
| 198 | + conj!(Cv) |
| 199 | + elseif isconjA |
| 200 | + A′ = materialize_conj(Av, conjA, α′, allocator) |
| 201 | + tblis_mult!(Cv, SV(A′), Bv, einA, einB, einC, one(T), β′) |
| 202 | + tensorfree!(A′, allocator) |
| 203 | + elseif isconjB |
| 204 | + B′ = materialize_conj(Bv, conjB, one(T), allocator) |
| 205 | + tblis_mult!(Cv, Av, SV(B′), einA, einB, einC, α′, β′) |
| 206 | + tensorfree!(B′, allocator) |
| 207 | + else |
| 208 | + tblis_mult!(Cv, Av, Bv, einA, einB, einC, α′, β′) |
| 209 | + end |
| 210 | + return C |
| 211 | +end |
| 212 | + |
| 213 | +function tblis_mult!( |
| 214 | + C::StridedView{T}, A::StridedView{T}, B::StridedView{T}, |
| 215 | + einA, einB, einC, α::T, β::T |
| 216 | + ) where {T <: TBLISFloat} |
| 217 | + lenA, strideA = tblis_dims(A) |
| 218 | + lenB, strideB = tblis_dims(B) |
| 219 | + lenC, strideC = tblis_dims(C) |
| 220 | + GC.@preserve A B C lenA strideA lenB strideB lenC strideC begin |
| 221 | + tA = tblis_tensor_ref(A, α, lenA, strideA, false) |
| 222 | + tB = tblis_tensor_ref(B, one(T), lenB, strideB, false) |
| 223 | + tC = tblis_tensor_ref(C, β, lenC, strideC, false) |
| 224 | + TBLIS.tblis_tensor_mult( |
| 225 | + C_NULL, C_NULL, tA, labels(einA), tB, labels(einB), tC, labels(einC) |
| 226 | + ) |
| 227 | + end |
| 228 | + return C |
| 229 | +end |
| 230 | + |
| 231 | +function materialize_conj( |
| 232 | + A::StridedView{T, N}, conjA::Bool, α::T, allocator |
| 233 | + ) where {T <: TBLISFloat, N} |
| 234 | + pA = (ntuple(identity, N), ()) |
| 235 | + A′ = tensoralloc_add(T, A, pA, false, Val(true), allocator) |
| 236 | + TO.tensoradd!(A′, A, pA, conjA, α, zero(T), TBLISBackend(), allocator) |
| 237 | + return A′ |
| 238 | +end |
| 239 | + |
| 240 | +end # module TensorOperationsTBLISExt |
0 commit comments