From 548a5bc168dad8619de968a4ecb0c3c78d8cf344 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 7 Aug 2026 15:06:33 -0400 Subject: [PATCH 1/2] 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/2] 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)