Fix JLArray buffer-backed temporaries against JLArrays 0.3.2

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"

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,

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

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

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,

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)