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,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

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

@@ -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)
 

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