Add AMDGPU and JLArrays buffer support

docs/src/man/backends.md

@@ -138,7 +138,8 @@ In particular in multi-threaded applications, this can sometimes lead to a signi
 On the other hand, for repeated (but thread-safe!) `@tensor` calls, the `BufferAllocator` is a lightweight slab allocator that pre-allocates a buffer for temporaries, falling back to Julia's default if needed.
 Upon repeated use it will automatically resize the buffer to accommodate the requested temporaries, avoiding repeated reallocation.
 The container that backs the buffer is a type parameter, such that the same slab strategy can also be used for other kinds of memory.
-In particular, a `CuArray`-backed buffer, which is constructed through `CUDABufferAllocator`, hands out `CuArray` temporaries that are carved out of a single device allocation, thus avoiding repeated round-trips to the CUDA memory pool.
+In particular, a `CuArray`- or `ROCArray`-backed buffer, which is constructed through `CUDABufferAllocator` or `AMDBufferAllocator`, hands out device-array temporaries that are carved out of a single device allocation, thus avoiding repeated round-trips to the GPU memory pool.
+A `JLArray`-backed buffer is available as well through `JLBufferAllocator`, which is mostly useful for testing since `JLArrays` requires no GPU hardware.
 
 Finally, users can also opt to use the `Bumper.jl` system, which pre-allocates a slab of memory that can be re-used afterwards.
 This is available through a package extension for `Bumper`.
@@ -178,6 +179,26 @@ allocator = TensorOperations.CUDABufferAllocator(; sizehint = 2^20)
 Here, all temporaries are taken from the pre-allocated device buffer, which is grown automatically until it is large enough to hold all temporaries of a single `@tensor` block.
 Just like `CUDAAllocator`, this will produce `CuArray` outputs even when the inputs are regular host arrays.
 
+The AMDGPU counterparts are `AMDAllocator`, which allocates every temporary through the AMD memory manager, and `AMDBufferAllocator`, which serves them from a single pre-allocated `ROCArray`:
+
+```@docs
+TensorOperations.AMDAllocator
+TensorOperations.AMDBufferAllocator
+```
+
+```julia
+using TensorOperations, AMDGPU
+allocator = TensorOperations.AMDBufferAllocator(; sizehint = 2^20)
+@tensor allocator = allocator A[i,j] := B[i,k] * C[k,j]
+```
+
+Finally, `JLBufferAllocator` provides the same thing on top of `JLArrays`, the reference GPU array implementation.
+Since it requires no GPU hardware, it is mostly useful to exercise the device-storage code paths in tests:
+
+```@docs
+TensorOperations.JLBufferAllocator
+```
+
 ### Custom Allocators
 
 Users can also define their own allocators, to facilitate experimentation with new implementations.

ext/TensorOperationsAMDGPUExt.jl

@@ -46,4 +46,80 @@ function TO.tensorfree!(C::ROCArray, ::TO.AMDAllocator)
     return nothing
 end
 
+#-------------------------------------------------------------------------------------------
+# BufferAllocator with ROCArray storage
+#-------------------------------------------------------------------------------------------
+
+const ROCBufferAllocator{B} = TO.BufferAllocator{ROCArray{UInt8, 1, B}}
+
+# Note: separate binding from the alias above, as a type alias cannot be added to the parent module from an extension
+function TO.AMDBufferAllocator(;
+        sizehint::Integer = 0, buftype = AMDGPU.Mem.HIPBuffer
+    )
+    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,
+        istemp::Val, allocator::ROCBufferAllocator
+    )
+    ttype = ROCArray{TC, TO.numind(pA)}
+    structure = TO.tensoradd_structure(A, pA, conjA)
+    return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype
+end
+
+function TO.tensoralloc_contract(
+        TC,
+        A::AbstractArray, pA::Index2Tuple, conjA::Bool,
+        B::AbstractArray, pB::Index2Tuple, conjB::Bool,
+        pAB::Index2Tuple,
+        istemp::Val, allocator::ROCBufferAllocator
+    )
+    ttype = ROCArray{TC, TO.numind(pAB)}
+    structure = TO.tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)
+    return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype
+end
+
+# NOTE: for tensors backed by the buffer this only releases the reference that `unsafe_buffer_wrap` retained
+function TO.tensorfree!(C::ROCArray, ::ROCBufferAllocator)
+    AMDGPU.unsafe_free!(C)
+    return nothing
+end
+
+function Base.resize!(buffer::ROCBufferAllocator{B}, n::Integer) where {B}
+    isempty(buffer) || error("Cannot resize a buffer that still contains elements")
+    n = TO._buffersz(n)
+    if n != length(buffer)
+        AMDGPU.unsafe_free!(buffer.buffer) # free before allocating new one to reduce memory pressure
+        buffer.buffer = ROCArray{UInt8, 1, B}(undef, n)
+    end
+    return buffer
+end
+
 end

ext/TensorOperationsJLArraysExt.jl

@@ -2,8 +2,87 @@ module TensorOperationsJLArraysExt
 
 using JLArrays
 using TensorOperations
+using TensorOperations: TensorOperations as TO
 
-TensorOperations.tensoradd_type(TC, A::JLArray, pA::Index2Tuple, conjA::Bool) =
-    JLArray{TC, sum(length.(pA))}
+#-------------------------------------------------------------------------------------------
+# Allocator
+#-------------------------------------------------------------------------------------------
+
+TO.tensoradd_type(TC, A::JLArray, pA::Index2Tuple, conjA::Bool) =
+    JLArray{TC, TO.numind(pA)}
+
+#-------------------------------------------------------------------------------------------
+# BufferAllocator with JLArray storage
+#-------------------------------------------------------------------------------------------
+
+const JLBuffer = TO.BufferAllocator{JLArray{UInt8, 1}}
+
+# Note: separate binding from the alias above, as a type alias cannot be added to the parent module from an extension
+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,
+        istemp::Val, allocator::JLBuffer
+    )
+    ttype = JLArray{TC, TO.numind(pA)}
+    structure = TO.tensoradd_structure(A, pA, conjA)
+    return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype
+end
+
+function TO.tensoralloc_contract(
+        TC,
+        A::AbstractArray, pA::Index2Tuple, conjA::Bool,
+        B::AbstractArray, pB::Index2Tuple, conjB::Bool,
+        pAB::Index2Tuple,
+        istemp::Val, allocator::JLBuffer
+    )
+    ttype = JLArray{TC, TO.numind(pAB)}
+    structure = TO.tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)
+    return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype
+end
+
+# NOTE: for tensors backed by the buffer this only releases the reference that `unsafe_buffer_wrap` retained
+function TO.tensorfree!(C::JLArray, ::JLBuffer)
+    JLArrays.unsafe_free!(C)
+    return nothing
+end
+
+function Base.resize!(buffer::JLBuffer, n::Integer)
+    isempty(buffer) || error("Cannot resize a buffer that still contains elements")
+    n = TO._buffersz(n)
+    if n != length(buffer)
+        JLArrays.unsafe_free!(buffer.buffer) # free before allocating new one to reduce memory pressure
+        buffer.buffer = JLArray{UInt8, 1}(undef, n)
+    end
+    return buffer
+end
 
 end

src/implementation/allocator.jl

@@ -53,6 +53,37 @@ See also [`TensorOperations.BufferAllocator`](@ref) and [`TensorOperations.CUDAA
 """
 function CUDABufferAllocator end
 
+"""
+    AMDBufferAllocator(; sizehint = 0, buftype = AMDGPU.Mem.HIPBuffer)
+
+Convenience constructor for a [`BufferAllocator`](@ref) that is backed by AMD memory, and which
+will thus hand out `ROCArray` instances that are carved out of a single pre-allocated buffer.
+The `buftype` keyword can be any of the AMDGPU.jl buffer types, i.e.
+`AMDGPU.Mem.HIPBuffer` or `AMDGPU.Mem.HostBuffer`, and determines both where the buffer itself
+lives and in which memory space the temporary tensors will be located.
+
+This requires `AMDGPU` to be loaded, and is equivalent to spelling out the storage type as
+`BufferAllocator{ROCArray{UInt8, 1, buftype}}(; sizehint)`.
+
+See also [`TensorOperations.BufferAllocator`](@ref) and [`TensorOperations.AMDAllocator`](@ref).
+"""
+function AMDBufferAllocator end
+
+"""
+    JLBufferAllocator(; sizehint = 0)
+
+Convenience constructor for a [`BufferAllocator`](@ref) that is backed by a `JLArray`, and which
+will thus hand out `JLArray` instances that are carved out of a single pre-allocated buffer.
+As `JLArrays` is the reference GPU array implementation, this is mostly useful for testing the
+foreign-storage code paths of `BufferAllocator` without requiring actual GPU hardware.
+
+This requires `JLArrays` to be loaded, and is equivalent to spelling out the storage type as
+`BufferAllocator{JLArray{UInt8, 1}}(; sizehint)`.
+
+See also [`TensorOperations.BufferAllocator`](@ref).
+"""
+function JLBufferAllocator end
+
 """
     ManualAllocator()
 
@@ -78,8 +109,9 @@ The optional type parameter `Storage` determines the container that backs the bu
 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`-backed buffers are
-supported through [`TensorOperations.CUDABufferAllocator`](@ref).
+[`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).
 
 !!! warning
     This allocator is **not** thread-safe, and it is the user's responsibility to avoid running