From 548a5bc168dad8619de968a4ecb0c3c78d8cf344 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 15:06:33 -0400 Subject: [PATCH 1/5] Fix JLArray buffer-backed temporaries against JLArrays 0.3.2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JLArrays 0.3.2 changed `JLArray`'s `offset` field from a count of elements to a count of bytes, and `pointer` from `pointer(x.data) + x.offset*elsize(x)` to `pointer(x.data) + x.offset`. `unsafe_buffer_wrap` built the temporary with the constructor directly, so under 0.3.2 every buffer-backed temporary landed at byte `start ÷ sizeof(T)` instead of `start`: misaligned, and overlapping both each other and earlier temporaries, which silently corrupted results. Go through `GPUArrays.derive` instead, the documented backend hook for producing an array of a different type and size backed by the same data. Its offset is expressed in elements on both versions, so this is insensitive to how the backend stores it. `ROCArray` already used a byte offset directly and was unaffected. Compat gets a `0.3.1` lower bound, matching the versions this was verified on. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 2 +- ext/TensorOperationsJLArraysExt.jl | 13 +++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index b0ff2846..3316ab3d 100644 --- a/Project.toml +++ b/Project.toml @@ -47,7 +47,7 @@ ChainRulesTestUtils = "1" DynamicPolynomials = "0.5, 0.6" Enzyme = "0.13.183" EnzymeTestUtils = "0.2" -JLArrays = "0.3" +JLArrays = "0.3.1" LRUCache = "1" LinearAlgebra = "1.6" Logging = "1.6" diff --git a/ext/TensorOperationsJLArraysExt.jl b/ext/TensorOperationsJLArraysExt.jl index 89a69512..d20a0c54 100644 --- a/ext/TensorOperationsJLArraysExt.jl +++ b/ext/TensorOperationsJLArraysExt.jl @@ -21,7 +21,8 @@ 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 +# Derived arrays are offset by a whole number of elements, 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)) @@ -35,12 +36,16 @@ function TO.buffer_arraytype(::Type{<:JLArray{T, N}}, buffer::JLBuffer) where {T 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 +# `GPUArrays.derive` is the documented backend hook for producing an array of a different type +# and size backed by the same data, as `reshape` and contiguous `view`s do, so the buffer +# outlives the temporary. Going through it rather than the `JLArray` constructor keeps this +# insensitive to whether the offset is stored per element or in bytes. 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)) + return JLArrays.GPUArrays.derive( + T, buffer.buffer, _asdims(structure), Int(start) ÷ sizeof(T) + ) end # `structure` is a shape for arrays, but a bare length is accepted for vectors From 22b85b0466b5ea299793dd66ca44a31280ae3122 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 15:31:14 -0400 Subject: [PATCH 2/5] Share the buffer-backing implementation across GPU backends `unsafe_buffer_wrap` was written three times, once per backend, and each reached for a different private detail: `JLArray`'s constructor with an element offset, `ROCArray`'s with a byte offset, and `unsafe_wrap` on a raw `CuPtr`. The first of those is what broke against JLArrays 0.3.2. Route all three through `GPUArrays.derive` in a new `TensorOperationsGPUArraysExt` instead. It is the documented backend hook for producing an array of a different type and size backed by the same data, its offset is in elements on every backend and version, and sharing the storage's refcounted handle keeps the buffer alive for as long as a temporary derived from it -- which the `CuArray` path, wrapping a bare pointer, did not do. `buffer_arraytype` likewise collapses into one method, via a new `buffer_similartype` that asks the buffer's own storage what it produces for the requested element type and rank. This resolves at compile time and reproduces each backend's answer exactly, memory space and buffer type included. It stays out of the core default because a storage is not always the kind of array it hands out: `similar(::Memory{UInt8}, T, Dims{1})` is a `Memory`, not a `Vector`. Since `derive` offsets by whole elements, the element-addressability restriction that only `JLArray` applied now covers all three, as `buffer_iselementaddressable`. Element types it rejects fall back on a regular allocation, so this costs buffer backing for exotic types, never correctness. Verified against JLArrays 0.3.1 and 0.3.2, so the compat bound stays at "0.3". Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 5 ++- ext/TensorOperationsAMDGPUExt.jl | 22 ------------- ext/TensorOperationsCUDACoreExt.jl | 16 +-------- ext/TensorOperationsGPUArraysExt.jl | 36 ++++++++++++++++++++ ext/TensorOperationsJLArraysExt.jl | 31 ------------------ src/implementation/allocator.jl | 51 +++++++++++++++++++++++++++-- 6 files changed, 90 insertions(+), 71 deletions(-) create mode 100644 ext/TensorOperationsGPUArraysExt.jl diff --git a/Project.toml b/Project.toml index 3316ab3d..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,7 +49,8 @@ ChainRulesTestUtils = "1" DynamicPolynomials = "0.5, 0.6" Enzyme = "0.13.183" EnzymeTestUtils = "0.2" -JLArrays = "0.3.1" +GPUArrays = "11" +JLArrays = "0.3" LRUCache = "1" LinearAlgebra = "1.6" Logging = "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 d20a0c54..1006c5bf 100644 --- a/ext/TensorOperationsJLArraysExt.jl +++ b/ext/TensorOperationsJLArraysExt.jl @@ -21,37 +21,6 @@ const JLBuffer = TO.BufferAllocator{JLArray{UInt8, 1}} TO.JLBufferAllocator(; sizehint::Integer = 0) = TO.BufferAllocator{JLArray{UInt8, 1}}(; sizehint) -# Derived arrays are offset by a whole number of elements, 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 - -# `GPUArrays.derive` is the documented backend hook for producing an array of a different type -# and size backed by the same data, as `reshape` and contiguous `view`s do, so the buffer -# outlives the temporary. Going through it rather than the `JLArray` constructor keeps this -# insensitive to whether the offset is stored per element or in bytes. -function TO.unsafe_buffer_wrap( - ::Type{JLArray{T, N}}, buffer::JLBuffer, start, structure - ) where {T, N} - return JLArrays.GPUArrays.derive( - T, buffer.buffer, _asdims(structure), 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) From 6353a79f5300a908cf215510a6ecead8242cd4cf Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 10 Aug 2026 09:29:04 -0400 Subject: [PATCH 3/5] Make the buffer's element-offset restriction unnecessary `GPUArrays.derive` counts its offset in elements, which previously restricted buffer-backed temporaries to element types whose size divides the alignment the offset is padded to: `buffer_iselementaddressable` rejected the others, and the GPUArrays extension had to override `buffer_arraytype` to apply that check. Pad the offset to a multiple of `sizeof(T)` instead. This is the only padding that can express a byte offset as an element offset on every backend -- some carry the offset in bytes, others in elements, and no way of calling `derive` can represent a non-element-aligned offset in the latter -- and it costs no additional padding for element types whose size divides the alignment, which is all the standard ones. Oversized and oddly sized element types are now served from the buffer as well, rather than falling back on a regular allocation. With the check gone, what is left of the extension's `buffer_arraytype` is the `similar`-based rule, which is the right default for any storage: promote it to the default implementation, and keep an `Array` answer for `Memory` storage, the one storage whose `similar` stays a `Memory` at rank 1 while the temporaries it hands out are `Array`s at every rank. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationsGPUArraysExt.jl | 13 +---- src/implementation/allocator.jl | 82 ++++++++++++++--------------- test/allocator.jl | 43 +++++++++++---- 3 files changed, 75 insertions(+), 63 deletions(-) diff --git a/ext/TensorOperationsGPUArraysExt.jl b/ext/TensorOperationsGPUArraysExt.jl index c525bc60..f676b34d 100644 --- a/ext/TensorOperationsGPUArraysExt.jl +++ b/ext/TensorOperationsGPUArraysExt.jl @@ -10,22 +10,13 @@ using TensorOperations: TensorOperations as TO 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. +# The offset `derive` takes is counted in elements, which the padding that `tensoralloc` +# applies guarantees the byte offset to be a whole number of. function TO.unsafe_buffer_wrap( ::Type{A}, buffer::GPUBufferAllocator, start, structure ) where {A <: AbstractGPUArray} diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl index 2f16bb1f..85632077 100644 --- a/src/implementation/allocator.jl +++ b/src/implementation/allocator.jl @@ -345,10 +345,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,30 +358,24 @@ 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 -_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) +# The alignment `tensoralloc` pads a temporary of element type `T` to: at least the natural +# alignment of `T` and the alignment the buffer asks for, and always a multiple of `sizeof(T)` +# on top of that. The latter keeps the byte offset expressible as a whole number of elements, +# which is how the arrays of some backends carry their offset. For all standard element types +# `sizeof(T)` already divides the alignment, so this costs no additional padding; it only +# kicks in for oversized or oddly sized element types, which would otherwise have to fall +# back on a regular allocation. +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 @@ -395,36 +390,33 @@ 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. +The default answer is the type that `buffer`'s own storage would produce for the element type +and rank of `A`, and `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, and covers all GPU backends +at once. It is resolved from the types alone, via inference on `similar`, so a storage whose +`similar` is not inferrable loses buffer backing rather than becoming incorrect. 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} +function buffer_arraytype(::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 +@static if isdefined(Core, :Memory) + # `Memory` is the one storage for which `similar` does not answer this: it stays a `Memory` + # at rank 1, while the temporaries a `Memory`-backed buffer hands out are `Array`s at every + # rank + function buffer_arraytype( + ::Type{A}, ::BufferAllocator{<:Memory} + ) where {A <: AbstractArray} + return A <: Array ? A : nothing + end +end + """ unsafe_buffer_wrap(::Type{A}, buffer::BufferAllocator, start, structure) -> A @@ -432,6 +424,10 @@ 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 an implementation for arrays +that carry an element offset rather than a pointer can use `start ÷ sizeof(eltype(A))` without +losing bytes to the division. """ function unsafe_buffer_wrap( ::Type{A}, buffer::BufferAllocator, start, structure diff --git a/test/allocator.jl b/test/allocator.jl index 87237558..ffff9419 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, whatever its own storage is -- note that + # `similar` of the default `Memory` storage stays a `Memory` at rank 1 + 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,26 @@ 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 + # `GPUArrays.derive` takes an element offset, so the padding of an element type whose + # size does not divide the alignment is a multiple of that size instead: such types are + # still served from the buffer, at an offset that survives the conversion to elements + # rather than silently 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 From 37f19fdd67b336e62c53baa96d8aae977ab95e21 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 10 Aug 2026 10:28:36 -0400 Subject: [PATCH 4/5] Trim the comments and docstrings down Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorOperationsGPUArraysExt.jl | 10 +++------- src/implementation/allocator.jl | 30 +++++++++++------------------ test/allocator.jl | 11 +++++------ 3 files changed, 19 insertions(+), 32 deletions(-) diff --git a/ext/TensorOperationsGPUArraysExt.jl b/ext/TensorOperationsGPUArraysExt.jl index f676b34d..3bca5852 100644 --- a/ext/TensorOperationsGPUArraysExt.jl +++ b/ext/TensorOperationsGPUArraysExt.jl @@ -10,13 +10,9 @@ using TensorOperations: TensorOperations as TO const GPUBufferAllocator = TO.BufferAllocator{<:AbstractGPUArray} -# `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. -# The offset `derive` takes is counted in elements, which the padding that `tensoralloc` -# applies guarantees the byte offset to be a whole number of. +# `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} diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl index 85632077..5336826d 100644 --- a/src/implementation/allocator.jl +++ b/src/implementation/allocator.jl @@ -366,13 +366,9 @@ function _alignup(offset::Integer, alignment::Integer) return cld(offset, a) * a end -# The alignment `tensoralloc` pads a temporary of element type `T` to: at least the natural -# alignment of `T` and the alignment the buffer asks for, and always a multiple of `sizeof(T)` -# on top of that. The latter keeps the byte offset expressible as a whole number of elements, -# which is how the arrays of some backends carry their offset. For all standard element types -# `sizeof(T)` already divides the alignment, so this costs no additional padding; it only -# kicks in for oversized or oddly sized element types, which would otherwise have to fall -# back on a regular allocation. +# 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) @@ -390,12 +386,11 @@ 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 answer is the type that `buffer`'s own storage would produce for the element type -and rank of `A`, and `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, and covers all GPU backends -at once. It is resolved from the types alone, via inference on `similar`, so a storage whose -`similar` is not inferrable loses buffer backing rather than becoming incorrect. +The default answer is the type `buffer`'s own storage would produce for the element type and +rank of `A`, as inferred from `similar`, and `nothing` if that is not a concrete subtype of +`A`: a buffer backs exactly the arrays of its own storage kind, which covers all GPU backends +at once. A storage whose `similar` is not inferrable loses buffer backing rather than +becoming incorrect. See also [`TensorOperations.unsafe_buffer_wrap`](@ref). """ @@ -407,9 +402,7 @@ function buffer_arraytype(::Type{A}, buffer::BufferAllocator) where {A <: Abstra end @static if isdefined(Core, :Memory) - # `Memory` is the one storage for which `similar` does not answer this: it stays a `Memory` - # at rank 1, while the temporaries a `Memory`-backed buffer hands out are `Array`s at every - # rank + # `similar` of a `Memory` stays a `Memory` at rank 1, while the temporaries are `Array`s function buffer_arraytype( ::Type{A}, ::BufferAllocator{<:Memory} ) where {A <: AbstractArray} @@ -425,9 +418,8 @@ 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 an implementation for arrays -that carry an element offset rather than a pointer can use `start ÷ sizeof(eltype(A))` without -losing bytes to the division. +`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 diff --git a/test/allocator.jl b/test/allocator.jl index ffff9419..1ba32670 100644 --- a/test/allocator.jl +++ b/test/allocator.jl @@ -80,8 +80,8 @@ using JLArrays end @testset "buffer_arraytype" begin - # a host buffer serves `Array`s at every rank, whatever its own storage is -- note that - # `similar` of the default `Memory` storage stays a `Memory` at rank 1 + # 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}, @@ -261,10 +261,9 @@ end @test iszero(UInt(pointer(C2)) % 16) end - # `GPUArrays.derive` takes an element offset, so the padding of an element type whose - # size does not divide the alignment is a multiple of that size instead: such types are - # still served from the buffer, at an offset that survives the conversion to elements - # rather than silently truncating onto the previous temporary + # 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) === From 626ac14372404acf412a18d6777e2e87a64b2d57 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 10 Aug 2026 11:12:16 -0400 Subject: [PATCH 5/5] Infer `buffer_arraytype` from `unsafe_buffer_wrap` itself Asking `similar` what the storage would produce needed a `Memory` exception, since a `Memory` stays a `Memory` at rank 1 while the temporaries it hands out are `Array`s. Infer the wrap instead: the answer is then the type that is actually handed out, and no method applying is itself the `nothing` answer, so a storage only has to implement `unsafe_buffer_wrap` to be served from the buffer. Dispatch that host method on host storages, so that a device buffer asked for an `Array` has no method rather than a pointer conversion that would throw. Co-Authored-By: Claude Opus 5 (1M context) --- src/implementation/allocator.jl | 32 ++++++++++++-------------------- 1 file changed, 12 insertions(+), 20 deletions(-) diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl index 5336826d..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. @@ -386,30 +389,19 @@ 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 answer is the type `buffer`'s own storage would produce for the element type and -rank of `A`, as inferred from `similar`, and `nothing` if that is not a concrete subtype of -`A`: a buffer backs exactly the arrays of its own storage kind, which covers all GPU backends -at once. A storage whose `similar` is not inferrable loses buffer backing rather than -becoming incorrect. - -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}, buffer::BufferAllocator) where {A <: AbstractArray} S = Base.promote_op( - similar, typeof(buffer.buffer), Type{eltype(A)}, Base.Dims{ndims(A)} + unsafe_buffer_wrap, Type{A}, typeof(buffer), Int, Base.Dims{ndims(A)} ) return (isconcretetype(S) && S <: A) ? S : nothing end -@static if isdefined(Core, :Memory) - # `similar` of a `Memory` stays a `Memory` at rank 1, while the temporaries are `Array`s - function buffer_arraytype( - ::Type{A}, ::BufferAllocator{<:Memory} - ) where {A <: AbstractArray} - return A <: Array ? A : nothing - end -end - """ unsafe_buffer_wrap(::Type{A}, buffer::BufferAllocator, start, structure) -> A @@ -422,7 +414,7 @@ that the requested range actually fits within the buffer. 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)