diff --git a/Project.toml b/Project.toml index b0ff2846..99c531b9 100644 --- a/Project.toml +++ b/Project.toml @@ -23,6 +23,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" @@ -33,6 +34,7 @@ TensorOperationsChainRulesCoreExt = "ChainRulesCore" TensorOperationsMooncakeExt = "Mooncake" TensorOperationsCUDACoreExt = "CUDACore" TensorOperationsEnzymeExt = "Enzyme" +TensorOperationsGPUArraysExt = "GPUArrays" TensorOperationscuTENSORExt = "cuTENSOR" TensorOperationsJLArraysExt = "JLArrays" @@ -47,6 +49,7 @@ ChainRulesTestUtils = "1" DynamicPolynomials = "0.5, 0.6" Enzyme = "0.13.183" EnzymeTestUtils = "0.2" +GPUArrays = "11" JLArrays = "0.3" LRUCache = "1" LinearAlgebra = "1.6" diff --git a/ext/TensorOperationsAMDGPUExt.jl b/ext/TensorOperationsAMDGPUExt.jl index a1aaa484..bb64aab1 100644 --- a/ext/TensorOperationsAMDGPUExt.jl +++ b/ext/TensorOperationsAMDGPUExt.jl @@ -59,31 +59,9 @@ function TO.AMDBufferAllocator(; return TO.BufferAllocator{ROCArray{UInt8, 1, buftype}}(; sizehint) end -# AMD buffers can only back `ROCArray`s; the generic implementation already takes care of -# the converse, i.e. that host buffers can never back `ROCArray`s -function TO.buffer_arraytype( - ::Type{<:ROCArray{T, N}}, ::ROCBufferAllocator{B} - ) where {T, N, B} - return ROCArray{T, N, B} -end -TO.buffer_arraytype(::Type{<:Array}, ::ROCBufferAllocator) = nothing - # HIP allocations are 256-byte aligned; matching that keeps rocBLAS kernel selection identical, at ≤255 bytes of padding TO.buffer_alignment(::ROCBufferAllocator) = 256 -# Share the buffer's refcounted `DataRef` at a byte offset, as `reshape` does: that keeps the buffer alive, and -# avoids the `hipPointerGetAttributes` query that `unsafe_wrap` would do per temporary -function TO.unsafe_buffer_wrap( - ::Type{ROCArray{T, N, B}}, buffer::ROCBufferAllocator{B}, start, structure - ) where {T, N, B} - ref = copy(AMDGPU.GPUArrays.storage(buffer.buffer)) - return ROCArray{T, N}(ref, _asdims(structure); offset = Int(start)) -end - -# `structure` is a shape for arrays, but a bare length is accepted for vectors -_asdims(structure::Base.Dims) = structure -_asdims(n::Integer) = (Int(n),) - # mirror the `AMDAllocator` behavior: results and temporaries are `ROCArray`s, even if the inputs are regular host arrays function TO.tensoralloc_add( TC, A::AbstractArray, pA::Index2Tuple, conjA::Bool, diff --git a/ext/TensorOperationsCUDACoreExt.jl b/ext/TensorOperationsCUDACoreExt.jl index b1d2593f..ab3dfae3 100644 --- a/ext/TensorOperationsCUDACoreExt.jl +++ b/ext/TensorOperationsCUDACoreExt.jl @@ -65,24 +65,10 @@ function TO.CUDABufferAllocator(; sizehint::Integer = 0, memory = CUDACore.defau return TO.BufferAllocator{CuArray{UInt8, 1, memory}}(; sizehint) end -# CUDA buffers can only back `CuArray`s; the generic implementation already takes care of -# the converse, i.e. that host buffers can never back `CuArray`s -function TO.buffer_arraytype(::Type{<:CuArray{T, N}}, ::CuBufferAllocator{M}) where {T, N, M} - return CuArray{T, N, M} -end -TO.buffer_arraytype(::Type{<:Array}, ::CuBufferAllocator) = nothing - # CUDA allocations are 256-byte aligned, and cuTENSOR selects noticeably faster kernels for 256-byte aligned data. # The padding this costs is at most 255 bytes per temporary, which is negligible in comparison. TO.buffer_alignment(::CuBufferAllocator) = 256 -function TO.unsafe_buffer_wrap( - ::Type{CuArray{T, N, M}}, buffer::CuBufferAllocator{M}, start, structure - ) where {T, N, M} - ptr = convert(CuPtr{T}, pointer(buffer, start)) - return unsafe_wrap(CuArray{T, N, M}, ptr, structure) -end - # mirror the `CUDAAllocator` behavior: results and temporaries are `CuArray`s, even if the inputs are regular host arrays function TO.tensoralloc_add( TC, A::AbstractArray, pA::Index2Tuple, conjA::Bool, @@ -105,7 +91,7 @@ function TO.tensoralloc_contract( return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype end -# NOTE: this is a no-op for tensors that are backed by the buffer, as `unsafe_wrap` creates a non-owning reference +# NOTE: for tensors backed by the buffer this only releases the reference that `unsafe_buffer_wrap` retained function TO.tensorfree!(C::CuArray, ::CuBufferAllocator) CUDACore.unsafe_free!(C) return nothing diff --git a/ext/TensorOperationsGPUArraysExt.jl b/ext/TensorOperationsGPUArraysExt.jl new file mode 100644 index 00000000..c525bc60 --- /dev/null +++ b/ext/TensorOperationsGPUArraysExt.jl @@ -0,0 +1,36 @@ +module TensorOperationsGPUArraysExt + +using GPUArrays +using TensorOperations +using TensorOperations: TensorOperations as TO + +#------------------------------------------------------------------------------------------- +# BufferAllocator with AbstractGPUArray storage +#------------------------------------------------------------------------------------------- + +const GPUBufferAllocator = TO.BufferAllocator{<:AbstractGPUArray} + +# A GPU buffer backs exactly the arrays its own storage produces, which `buffer_similartype` +# answers for every backend at once. A derived array can only be offset by a whole number of +# elements, so additionally restrict to element types for which the padded offset is +# expressible that way. +function TO.buffer_arraytype( + ::Type{A}, buffer::GPUBufferAllocator + ) where {A <: AbstractArray} + TO.buffer_iselementaddressable(eltype(A), buffer) || return nothing + return TO.buffer_similartype(A, buffer) +end + +# `GPUArrays.derive` is the documented backend hook for producing an array of a different +# type and size backed by the same data, and is what `reshape` and contiguous `view`s go +# through. Sharing the buffer's refcounted storage keeps it alive for as long as the +# temporary is, and going through `derive` rather than a backend's own constructor or +# `unsafe_wrap` keeps this insensitive to how a backend represents the offset internally. +function TO.unsafe_buffer_wrap( + ::Type{A}, buffer::GPUBufferAllocator, start, structure + ) where {A <: AbstractGPUArray} + T = eltype(A) + return GPUArrays.derive(T, buffer.buffer, TO._asdims(structure), Int(start) ÷ sizeof(T)) +end + +end diff --git a/ext/TensorOperationsJLArraysExt.jl b/ext/TensorOperationsJLArraysExt.jl index 89a69512..1006c5bf 100644 --- a/ext/TensorOperationsJLArraysExt.jl +++ b/ext/TensorOperationsJLArraysExt.jl @@ -21,32 +21,6 @@ const JLBuffer = TO.BufferAllocator{JLArray{UInt8, 1}} TO.JLBufferAllocator(; sizehint::Integer = 0) = TO.BufferAllocator{JLArray{UInt8, 1}}(; sizehint) -# `JLArray`s are addressed by element offset, so `T`s whose size does not divide the alignment cannot be buffer-backed -function _iselementaddressable(::Type{T}, buffer::JLBuffer) where {T} - sz = sizeof(T) - alignment = max(Base.datatype_alignment(T), TO.buffer_alignment(buffer)) - return !iszero(sz) && iszero(alignment % sz) -end - -# JLArray buffers can only back `JLArray`s; the generic implementation already takes care of -# the converse, i.e. that host buffers can never back `JLArray`s -function TO.buffer_arraytype(::Type{<:JLArray{T, N}}, buffer::JLBuffer) where {T, N} - return _iselementaddressable(T, buffer) ? JLArray{T, N} : nothing -end -TO.buffer_arraytype(::Type{<:Array}, ::JLBuffer) = nothing - -# Share the buffer's refcounted `DataRef` at an element offset, as `reshape` does, so the buffer outlives the temporary -function TO.unsafe_buffer_wrap( - ::Type{JLArray{T, N}}, buffer::JLBuffer, start, structure - ) where {T, N} - ref = copy(JLArrays.GPUArrays.storage(buffer.buffer)) - return JLArray{T, N}(ref, _asdims(structure); offset = Int(start) ÷ sizeof(T)) -end - -# `structure` is a shape for arrays, but a bare length is accepted for vectors -_asdims(structure::Base.Dims) = structure -_asdims(n::Integer) = (Int(n),) - # mirror the GPU allocator behavior: results and temporaries are `JLArray`s, even if the inputs are regular host arrays function TO.tensoralloc_add( TC, A::AbstractArray, pA::Index2Tuple, conjA::Bool, diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl index 6ba192b1..2f16bb1f 100644 --- a/src/implementation/allocator.jl +++ b/src/implementation/allocator.jl @@ -363,6 +363,30 @@ function _alignup(offset::Integer, alignment::Integer) return (offset + a - one(a)) & ~(a - one(a)) end +# the alignment `tensoralloc` pads a temporary of element type `T` to +_buffer_alignment(::Type{T}, buffer::BufferAllocator) where {T} = + max(Base.datatype_alignment(T), buffer_alignment(buffer)) + +""" + buffer_iselementaddressable(::Type{T}, buffer::BufferAllocator) + +Whether a temporary with element type `T` can be placed at an arbitrary padded offset in +`buffer` while addressing that offset as a whole number of elements. This is what backends +whose arrays carry an element offset rather than a raw pointer need, and it holds exactly +when `sizeof(T)` divides the alignment the offset is padded to. + +Element types for which this fails fall back on a regular allocation, so this is a +restriction on which temporaries a buffer can serve, never on correctness. +""" +function buffer_iselementaddressable(::Type{T}, buffer::BufferAllocator) where {T} + sz = sizeof(T) + return !iszero(sz) && iszero(_buffer_alignment(T, buffer) % sz) +end + +# `structure` is a shape for arrays, but a bare length is accepted for vectors +_asdims(structure::Base.Dims) = structure +_asdims(n::Integer) = (Int(n),) + """ buffer_arraytype(::Type{A}, buffer::BufferAllocator) @@ -371,12 +395,36 @@ Return the concrete array type that is used to serve a temporary allocation of t allocation path is used instead. This only depends on the types involved, such that the choice is resolved at compile time. +The default only backs `Array`s, which is what a host buffer can serve. Storages that are +themselves arrays can use [`TensorOperations.buffer_similartype`](@ref) to answer this +generically; `TensorOperationsGPUArraysExt` does so for every GPU backend at once. + See also [`TensorOperations.unsafe_buffer_wrap`](@ref). """ function buffer_arraytype(::Type{A}, ::BufferAllocator) where {A <: AbstractArray} return A <: Array ? A : nothing end +""" + buffer_similartype(::Type{A}, buffer::BufferAllocator) + +The concrete type `buffer`'s own storage would produce for the element type and rank of `A`, +or `nothing` if that is not a concrete subtype of `A`. This lets a buffer back exactly the +arrays of its own storage kind -- a `CuArray` buffer cannot back an `Array` and vice versa -- +without every storage having to spell that out. + +Resolved from the types alone, via inference on `similar`, so a storage whose `similar` is +not inferrable loses buffer backing rather than becoming incorrect. Note that this is only +meaningful for storages that are arrays of the same kind they hand out: `Memory{UInt8}`, for +one, stays a `Memory` at rank 1 instead of becoming a `Vector`. +""" +function buffer_similartype(::Type{A}, buffer::BufferAllocator) where {A <: AbstractArray} + S = Base.promote_op( + similar, typeof(buffer.buffer), Type{eltype(A)}, Base.Dims{ndims(A)} + ) + return (isconcretetype(S) && S <: A) ? S : nothing +end + """ unsafe_buffer_wrap(::Type{A}, buffer::BufferAllocator, start, structure) -> A @@ -400,8 +448,7 @@ function tensoralloc( T = eltype(AA) nbytes = allocation_size(T, structure) if !iszero(nbytes) # empty temporaries have no meaningful pointer - alignment = max(Base.datatype_alignment(T), buffer_alignment(buffer)) - start = _alignup(buffer.offset, alignment) + start = _alignup(buffer.offset, _buffer_alignment(T, buffer)) offset = start + nbytes sizehint!(buffer, offset)