add CUDA buffer support

docs/src/man/backends.md

@@ -137,6 +137,8 @@ Optionally, it can be useful to use the `ManualAllocator`, as the manual memory
 In particular in multi-threaded applications, this can sometimes lead to a significant performance improvement.
 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.
 
 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`.
@@ -162,8 +164,20 @@ with respect to temporaries, as well as input and output tensors.
 
 ```@docs
 TensorOperations.CUDAAllocator
+TensorOperations.CUDABufferAllocator
 ```
 
+Alternatively, the `CUDABufferAllocator` can also be used with the `cuTENSORBackend()`. This is simply a `BufferAllocator` backed by CUDA memory:
+
+```julia
+using TensorOperations, cuTENSOR
+allocator = TensorOperations.CUDABufferAllocator(; sizehint = 2^20)
+@tensor backend = cuTENSORBackend() allocator = allocator A[i,j] := B[i,k] * C[k,j]
+```
+
+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.
+
 ### Custom Allocators
 
 Users can also define their own allocators, to facilitate experimentation with new implementations.
@@ -189,3 +203,12 @@ For allocators that manage reusable buffers or maintain state across multiple co
 
 Here we are guaranteeing that every created checkpoint will be restored, and all temporary allocations that are enclosed within this scope will no longer be accessed.
 Additionally, if multiple checkpoints are created, they will be restored in a first-in-last-out order.
+
+Finally, rather than writing an allocator from scratch, the built-in `BufferAllocator` can be reused with a different kind of storage, by implementing the two methods below.
+The bookkeeping of the buffer itself, i.e. the high-water mark, the automatic resizing and the checkpointing, is then shared with the default implementation.
+
+```@docs
+TensorOperations.buffer_arraytype
+TensorOperations.unsafe_buffer_wrap
+TensorOperations.buffer_alignment
+```

ext/TensorOperationsCUDACoreExt.jl

@@ -54,4 +54,71 @@ function TO.tensorfree!(C::CuArray, ::TO.CUDAAllocator)
     return nothing
 end
 
+#-------------------------------------------------------------------------------------------
+# BufferAllocator with CuArray storage
+#-------------------------------------------------------------------------------------------
+
+const CuBufferAllocator{M} = TO.BufferAllocator{CuArray{UInt8, 1, M}}
+
+# Note: different binding of the same name because type alias cannot be exported from extension
+function TO.CUDABufferAllocator(; sizehint::Integer = 0, memory = CUDACore.default_memory)
+    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,
+        istemp::Val, allocator::CuBufferAllocator
+    )
+    ttype = CuArray{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::CuBufferAllocator
+    )
+    ttype = CuArray{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: this is a no-op for tensors that are backed by the buffer, as `unsafe_wrap` creates a non-owning reference
+function TO.tensorfree!(C::CuArray, ::CuBufferAllocator)
+    CUDACore.unsafe_free!(C)
+    return nothing
+end
+
+function Base.resize!(buffer::CuBufferAllocator{M}, n::Integer) where {M}
+    isempty(buffer) || error("Cannot resize a buffer that still contains elements")
+    n = TO._buffersz(n)
+    if n != length(buffer)
+        CUDACore.unsafe_free!(buffer.buffer) # free before allocating new one to reduce memory pressure
+        buffer.buffer = CuArray{UInt8, 1, M}(undef, n)
+    end
+    return buffer
+end
+
 end

ext/TensorOperationscuTENSORExt.jl

@@ -120,6 +120,9 @@ function TO.tensortrace!(
 end
 
 _custrided(A::AbstractArray, ::DefaultAllocator) = _custrided(A, CUDAAllocator())
+# input tensors are marshalled to the device in the same way as for the default allocator:
+# the buffer is reserved for the temporaries that are created by the tensor operations
+_custrided(A::AbstractArray, ::TO.BufferAllocator) = _custrided(A, CUDAAllocator())
 function _custrided(
         A::AbstractArray, allocator::CUDAAllocator{Mout, Min, Mtemp}
     ) where {Mout, Min, Mtemp}

src/implementation/allocator.jl

@@ -37,6 +37,22 @@ Allocator that uses the AMD memory manager and will thus allocate `ROCArray` ins
 """
 struct AMDAllocator end
 
+"""
+    CUDABufferAllocator(; sizehint = 0, memory = CUDA.default_memory)
+
+Convenience constructor for a [`BufferAllocator`](@ref) that is backed by CUDA memory, and
+which will thus hand out `CuArray` instances that are carved out of a single pre-allocated
+buffer. The `memory` keyword can be any of the CUDA.jl memory types, i.e.
+`CUDA.DeviceMemory`, `CUDA.UnifiedMemory` or `CUDA.HostMemory`, and determines both where the
+buffer itself lives and in which memory space the temporary tensors will be located.
+
+This requires `CUDACore` to be loaded, and is equivalent to spelling out the storage type as
+`BufferAllocator{CuArray{UInt8, 1, memory}}(; sizehint)`.
+
+See also [`TensorOperations.BufferAllocator`](@ref) and [`TensorOperations.CUDAAllocator`](@ref).
+"""
+function CUDABufferAllocator end
+
 """
     ManualAllocator()
 
@@ -50,13 +66,21 @@ struct ManualAllocator end
 
 """
     BufferAllocator(; sizehint = 0)
+    BufferAllocator{Storage}(; sizehint = 0)
 
 Allocator that uses a pre-allocated buffer for storing temporary tensors.
 When the buffer is full, the allocator falls back on Julia's default allocation mechanism
 to create temporary tensors, but keeps track of how much additional memory is required.
 When the buffer is fully reset, the buffer is automatically resized to ensure subsequent
 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`-backed buffers are
+supported through [`TensorOperations.CUDABufferAllocator`](@ref).
+
 !!! warning
     This allocator is **not** thread-safe, and it is the user's responsibility to avoid running
     the same allocator on concurrent jobs. For concurrent usage, it is recommended to either
@@ -80,8 +104,24 @@ end
 const DefaultStorageType = @static isdefined(Core, :Memory) ? Memory{UInt8} : Vector{UInt8}
 BufferAllocator(; kwargs...) = BufferAllocator{DefaultStorageType}(; kwargs...)
 
-# allocate buffers in sizes that are powers of 2
-_buffersz(x::Integer) = iszero(x) ? x : Base.nextpow(2, x)
+# `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.
+@static if isdefined(Sys, :PAGESIZE)
+    _pagesize() = Int(Sys.PAGESIZE)
+else
+    _pagesize() = 4096
+end
+
+# Allocate buffers in sizes that are a multiple of the page size.
+# Below a single page, powers of two are used instead, as rounding every small buffer up to a full page would be wasteful.
+# The result is always an `Int`, as that is what the storage constructors expect.
+function _buffersz(x::Integer)
+    iszero(x) && return 0
+    pagesize = _pagesize()
+    x ≤ pagesize && return Int(Base.nextpow(2, x))
+    return Int(cld(x, pagesize) * pagesize)
+end
 
 # ------------------------------------------------------------------------------------------
 # Generic implementation
@@ -270,26 +310,87 @@ end
 allocation_size(::Type{T}, structure::Base.Dims) where {T} = prod(structure) * sizeof(T)
 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.
+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.
+
+Defaults to `16`, which is the alignment that Julia guarantees for its allocations,
+and which covers the natural alignment of all standard element types.
+
+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
+function _alignup(offset::Integer, alignment::Integer)
+    a = oftype(offset, alignment)
+    return (offset + a - one(a)) & ~(a - one(a))
+end
+
+"""
+    buffer_arraytype(::Type{A}, buffer::BufferAllocator)
+
+Return the concrete array type that is used to serve a temporary allocation of type `A` from
+`buffer`, or `nothing` if `buffer` cannot back arrays of type `A`, in which case the regular
+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).
+"""
+function buffer_arraytype(::Type{A}, ::BufferAllocator) where {A <: AbstractArray}
+    return A <: Array ? A : nothing
+end
+
+"""
+    unsafe_buffer_wrap(::Type{A}, buffer::BufferAllocator, start, structure) -> A
+
+Wrap the memory of `buffer`, starting at byte offset `start`, into an array of type `A` with
+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.
+"""
+function unsafe_buffer_wrap(
+        ::Type{A}, buffer::BufferAllocator, start, structure
+    ) where {A <: Array}
+    ptr = convert(Ptr{eltype(A)}, pointer(buffer, start))
+    return Base.unsafe_wrap(Array, ptr, structure)
+end
+
 function tensoralloc(
         ::Type{A}, structure, ::Val{istemp}, buffer::BufferAllocator
     ) where {A <: AbstractArray, istemp}
-    if istemp
-        T = eltype(A)
-        offset = buffer.offset + allocation_size(T, structure)
-        sizehint!(buffer, offset)
-
-        # grow buffer if empty
-        isempty(buffer) && resize!(buffer, buffer.max_offset)
-
-        # Use pointer if there is enough space
-        if offset < length(buffer)
-            ptr = convert(Ptr{T}, pointer(buffer, buffer.offset))
-            buffer.offset = offset
-            return Base.unsafe_wrap(Array, ptr, structure)
+    AA = buffer_arraytype(A, buffer)
+    if istemp && AA !== nothing
+        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)
+            offset = start + nbytes
+            sizehint!(buffer, offset)
+
+            # grow buffer if empty: this should never shrink the buffer, as that would
+            # discard the size that was requested through `sizehint` or `resize!`
+            if isempty(buffer) && buffer.max_offset > length(buffer)
+                resize!(buffer, buffer.max_offset)
+            end
+
+            # Use pointer if there is enough space
+            if offset <= length(buffer)
+                buffer.offset = offset
+                return unsafe_buffer_wrap(AA, buffer, start, structure)
+            end
         end
+
+        # Allocate in the same memory space as the buffer if it does not fit
+        return AA(undef, structure)
     end
 
-    # Allocate default if not
+    # Allocate default if the buffer cannot back this type of array
     return A(undef, structure)
 end
 

test/allocator.jl

@@ -27,7 +27,7 @@ using LinearAlgebra
         @test buffer3 isa BufferAllocator{Vector{UInt8}}
         @test length(buffer3) >= 512
 
-        # sizehint! grows to next power-of-two elements (UInt8)
+        # buffers smaller than a page are rounded up to a power of two (UInt8 elements)
         resize!(buffer, 3000)
         @test length(buffer) == 4096
         @test isempty(buffer)
@@ -50,6 +50,33 @@ using LinearAlgebra
         @test buffer.max_offset == 1024
     end
 
+    @testset "Buffer sizes" begin
+        P = TensorOperations._pagesize()
+
+        # below a page, sizes are rounded up to a power of two
+        @test length(BufferAllocator(; sizehint = 0)) == 0
+        @test length(BufferAllocator(; sizehint = 100)) == 128
+        @test length(BufferAllocator(; sizehint = P)) == P
+
+        # above a page, sizes are rounded up to a multiple of a page, and not all the way
+        # up to the next power of two
+        @test length(BufferAllocator(; sizehint = P + 1)) == 2P
+        @test length(BufferAllocator(; sizehint = 10P + 1)) == 11P
+        buffer = BufferAllocator(; sizehint = 4P + 123)
+        @test length(buffer) == 5P
+        resize!(buffer, 2^20 + 1)
+        @test length(buffer) == cld(2^20 + 1, P) * P
+        @test length(buffer) < 2^21
+
+        # sizes are normalized to `Int`, whatever integer type they are computed from:
+        # the storage constructors do not accept e.g. the `Int32` that `Clong` is on Windows
+        @test TensorOperations._pagesize() isa Int
+        @test TensorOperations._buffersz(Int32(100)) === 128
+        @test TensorOperations._buffersz(UInt(10P + 1)) === 11P
+        @test length(BufferAllocator(; sizehint = Int32(100))) == 128
+        @test length(BufferAllocator(; sizehint = UInt(100))) == 128
+    end
+
     @testset "Checkpoint and reset" begin
         buffer = BufferAllocator(sizehint = 128)
         L = length(buffer)