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..3bca5852 --- /dev/null +++ b/ext/TensorOperationsGPUArraysExt.jl @@ -0,0 +1,23 @@ +module TensorOperationsGPUArraysExt + +using GPUArrays +using TensorOperations +using TensorOperations: TensorOperations as TO + +#------------------------------------------------------------------------------------------- +# BufferAllocator with AbstractGPUArray storage +#------------------------------------------------------------------------------------------- + +const GPUBufferAllocator = TO.BufferAllocator{<:AbstractGPUArray} + +# `GPUArrays.derive` is the backend hook that `reshape` and contiguous `view`s go through: it +# shares the buffer's refcounted storage, which keeps it alive for as long as the temporary, +# and is 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..ed96f7e3 100644 --- a/src/implementation/allocator.jl +++ b/src/implementation/allocator.jl @@ -108,8 +108,8 @@ contractions will now fit in the buffer. The optional type parameter `Storage` determines the container that backs the buffer, and must have single-byte elements. It defaults to `Memory{UInt8}` (or `Vector{UInt8}` on Julia versions without `Memory`), which hands out regular `Array` temporaries. Other storage types -can be supported by implementing [`TensorOperations.buffer_arraytype`](@ref) and -[`TensorOperations.unsafe_buffer_wrap`](@ref); in particular, `CuArray`-, `ROCArray`- and +can be supported by implementing [`TensorOperations.unsafe_buffer_wrap`](@ref); +in particular, `CuArray`-, `ROCArray`- and `JLArray`-backed buffers are supported through [`TensorOperations.CUDABufferAllocator`](@ref), [`TensorOperations.AMDBufferAllocator`](@ref) and [`TensorOperations.JLBufferAllocator`](@ref). @@ -136,6 +136,9 @@ end const DefaultStorageType = @static isdefined(Core, :Memory) ? Memory{UInt8} : Vector{UInt8} BufferAllocator(; kwargs...) = BufferAllocator{DefaultStorageType}(; kwargs...) +# storages whose `pointer` is a host pointer, so that an `Array` can be wrapped around it +const HostStorageType = @static isdefined(Core, :Memory) ? Union{Memory, Array} : Array + # `Sys.PAGESIZE` only exists on sufficiently recent Julia versions; fall back on the standard # page size otherwise, as this only serves as a granularity for rounding buffer sizes. # Note the conversion to `Int`: the underlying `Clong` is 32 bits wide on Windows. @@ -345,10 +348,11 @@ allocation_size(::Type{T}, structure::Int) where {T} = structure * sizeof(T) """ buffer_alignment(buffer::BufferAllocator) -The alignment, in bytes, to which the temporaries handed out by `buffer` are padded. +The minimum alignment, in bytes, to which the temporaries handed out by `buffer` are padded. This has to be a power of two, and currently there is no point in making it larger than the alignment of the buffer's own base pointer, as the padding would then not -actually buy any alignment. +actually buy any alignment. Element types whose size is not a divisor of it are padded +further, to a multiple of that size as well. Defaults to `16`, which is the alignment that Julia guarantees for its allocations, and which covers the natural alignment of all standard element types. @@ -357,12 +361,26 @@ See also [`TensorOperations.buffer_arraytype`](@ref). """ buffer_alignment(::BufferAllocator) = 16 -# round `offset` up to the next multiple of `alignment`, which has to be a power of 2 +# round `offset` up to the next multiple of `alignment` function _alignup(offset::Integer, alignment::Integer) a = oftype(offset, alignment) - return (offset + a - one(a)) & ~(a - one(a)) + # the bit trick only holds for powers of two, which `_buffer_alignment` is free not to be + ispow2(a) && return (offset + a - one(a)) & ~(a - one(a)) + return cld(offset, a) * a +end + +# The alignment `tensoralloc` pads a temporary of element type `T` to. The multiple of +# `sizeof(T)` keeps the offset expressible in elements, which is how some backends carry it, +# and costs no additional padding for element types whose size divides the alignment. +function _buffer_alignment(::Type{T}, buffer::BufferAllocator) where {T} + alignment = max(Base.datatype_alignment(T), buffer_alignment(buffer)) + return iszero(sizeof(T)) ? alignment : lcm(sizeof(T), alignment) 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,10 +389,17 @@ 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. -See also [`TensorOperations.unsafe_buffer_wrap`](@ref). +The default answer is inferred from [`TensorOperations.unsafe_buffer_wrap`](@ref): the type +that wrapping `buffer`'s memory actually produces, or `nothing` if that is not a concrete +subtype of `A`, which includes the case where no method applies. A storage therefore only has +to implement `unsafe_buffer_wrap` for its arrays to be served from the buffer, and one whose +wrap is not inferrable loses buffer backing rather than becoming incorrect. """ -function buffer_arraytype(::Type{A}, ::BufferAllocator) where {A <: AbstractArray} - return A <: Array ? A : nothing +function buffer_arraytype(::Type{A}, buffer::BufferAllocator) where {A <: AbstractArray} + S = Base.promote_op( + unsafe_buffer_wrap, Type{A}, typeof(buffer), Int, Base.Dims{ndims(A)} + ) + return (isconcretetype(S) && S <: A) ? S : nothing end """ @@ -384,9 +409,12 @@ Wrap the memory of `buffer`, starting at byte offset `start`, into an array of t shape `structure`. Here, `A` is the type returned by [`TensorOperations.buffer_arraytype`](@ref), and it is the caller's responsibility to ensure that the requested range actually fits within the buffer. + +`start` is guaranteed to be a multiple of `sizeof(eltype(A))`, so arrays that carry an element +offset rather than a pointer can use `start ÷ sizeof(eltype(A))` without losing bytes. """ function unsafe_buffer_wrap( - ::Type{A}, buffer::BufferAllocator, start, structure + ::Type{A}, buffer::BufferAllocator{<:HostStorageType}, start, structure ) where {A <: Array} ptr = convert(Ptr{eltype(A)}, pointer(buffer, start)) return Base.unsafe_wrap(Array, ptr, structure) @@ -400,8 +428,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) diff --git a/test/allocator.jl b/test/allocator.jl index 87237558..1ba32670 100644 --- a/test/allocator.jl +++ b/test/allocator.jl @@ -79,6 +79,22 @@ using JLArrays @test length(BufferAllocator(; sizehint = UInt(100))) == 128 end + @testset "buffer_arraytype" begin + # a host buffer serves `Array`s at every rank, including the rank where `similar` of + # the default `Memory` storage stays a `Memory` + for buffer in (BufferAllocator(), BufferAllocator{Vector{UInt8}}(; sizehint = 1024)) + for A in ( + Vector{Float64}, Matrix{ComplexF64}, Array{Float32, 3}, + Vector{NTuple{4, Float64}}, + ) + @test TensorOperations.buffer_arraytype(A, buffer) === A + end + @static if isdefined(Core, :Memory) + @test TensorOperations.buffer_arraytype(Memory{Float64}, buffer) === nothing + end + end + end + @testset "Checkpoint and reset" begin buffer = BufferAllocator(sizehint = 128) L = length(buffer) @@ -245,17 +261,25 @@ end @test iszero(UInt(pointer(C2)) % 16) end - # `JLArray`s address their data by an element offset, which the 16-byte padding can only - # express for element types that are at most that large: bigger ones fall back on a - # regular allocation rather than silently landing on a truncated offset + # element types whose size does not divide the alignment are padded to a multiple of + # that size instead, so that the offset survives the conversion to elements that + # `GPUArrays.derive` takes rather than truncating onto the previous temporary @test TensorOperations.buffer_arraytype(JLArray{ComplexF64, 1}, buffer) === JLArray{ComplexF64, 1} - @test TensorOperations.buffer_arraytype(JLArray{NTuple{4, Float64}, 1}, buffer) === nothing - offset = buffer.offset - C3 = tensoralloc(JLArray{NTuple{4, Float64}, 1}, (4,), Val(true), buffer) - @test C3 isa JLArray{NTuple{4, Float64}, 1} - @test !isbufferbacked(C3, buffer) - @test buffer.offset == offset + @test TensorOperations.buffer_arraytype(JLArray{NTuple{4, Float64}, 1}, buffer) === + JLArray{NTuple{4, Float64}, 1} + for T in (NTuple{4, Float64}, NTuple{3, Float32}) + empty!(buffer) + C3 = tensoralloc(JLArray{UInt8, 1}, (3,), Val(true), buffer) + C4 = tensoralloc(JLArray{T, 1}, (4,), Val(true), buffer) + @test isbufferbacked(C4, buffer) + start = UInt(pointer(C4)) - UInt(pointer(buffer)) + @test iszero(start % sizeof(T)) + @test iszero(start % TensorOperations.buffer_alignment(buffer)) + # no overlap with the 3 bytes that `C3` occupies + @test start ≥ 3 + @test buffer.offset == start + 4 * sizeof(T) + end end @testset "checkpoint and reset" begin