diff --git a/docs/src/man/backends.md b/docs/src/man/backends.md index 6ba5fd8c..9971bd16 100644 --- a/docs/src/man/backends.md +++ b/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 +``` diff --git a/ext/TensorOperationsCUDACoreExt.jl b/ext/TensorOperationsCUDACoreExt.jl index 6584a04a..b1d2593f 100644 --- a/ext/TensorOperationsCUDACoreExt.jl +++ b/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 diff --git a/ext/TensorOperationscuTENSORExt.jl b/ext/TensorOperationscuTENSORExt.jl index 7e88ba35..b67ccbe1 100644 --- a/ext/TensorOperationscuTENSORExt.jl +++ b/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} diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl index bda32ff4..835d8147 100644 --- a/src/implementation/allocator.jl +++ b/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,6 +66,7 @@ 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 @@ -57,6 +74,13 @@ to create temporary tensors, but keeps track of how much additional memory is re 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 diff --git a/test/allocator.jl b/test/allocator.jl index b79dcec5..c1bca423 100644 --- a/test/allocator.jl +++ b/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) diff --git a/test/cutensor.jl b/test/cutensor.jl index dd36f6d8..75d18c13 100644 --- a/test/cutensor.jl +++ b/test/cutensor.jl @@ -13,6 +13,9 @@ if cuTENSOR.functional() using LinearAlgebra: norm using TensorOperations: IndexError using TensorOperations: cuTENSORBackend, CUDAAllocator + using TensorOperations: BufferAllocator, CUDABufferAllocator + using TensorOperations: tensoralloc, tensorfree! + using TensorOperations: allocator_checkpoint!, allocator_reset! @testset "elementary operations" verbose = true begin @testset "tensorcopy" begin @@ -173,6 +176,243 @@ if cuTENSOR.functional() end end + @testset "BufferAllocator" verbose = true begin + DeviceMemory = CUDACore.DeviceMemory + UnifiedMemory = CUDACore.UnifiedMemory + + # is the memory of `A` taken from `buffer`? + function isbufferbacked(A, buffer) + iszero(length(buffer)) && return false + base = UInt(pointer(buffer)) + return base ≤ UInt(pointer(A)) < base + length(buffer) + end + + @testset "Constructor and basic properties" begin + buffer = CUDABufferAllocator(; sizehint = 1024) + @test buffer isa BufferAllocator{CuArray{UInt8, 1, CUDACore.default_memory}} + @test length(buffer) == 1024 + @test isempty(buffer) + @test buffer.offset == 0 + + # explicit memory space + buffer2 = CUDABufferAllocator(; sizehint = 512, memory = UnifiedMemory) + @test buffer2 isa BufferAllocator{CuArray{UInt8, 1, UnifiedMemory}} + @test buffer2 isa + typeof(BufferAllocator{CuArray{UInt8, 1, UnifiedMemory}}(; sizehint = 512)) + @test length(buffer2) == 512 + + # resizing frees the old buffer and allocates a new one + resize!(buffer, 3000) + @test length(buffer) == 4096 + @test isempty(buffer) + buffer.offset = 100 + @test_throws ErrorException resize!(buffer, 8192) + empty!(buffer) + @test length(resize!(buffer, 8192)) == 8192 + end + + @testset "tensoralloc" begin + buffer = CUDABufferAllocator(; sizehint = 4096) + + # temporaries are taken from the buffer, in the memory space of the buffer + C1 = tensoralloc(CuArray{Float32, 2}, (8, 8), Val(true), buffer) + @test C1 isa CuArray{Float32, 2, DeviceMemory} + @test size(C1) == (8, 8) + @test isbufferbacked(C1, buffer) + @test buffer.offset == 8 * 8 * sizeof(Float32) + + # non-temporaries are not + offset = buffer.offset + C2 = tensoralloc(CuArray{Float32, 2}, (8, 8), Val(false), buffer) + @test C2 isa CuArray{Float32, 2} + @test !isbufferbacked(C2, buffer) + @test buffer.offset == offset + + # freeing a buffer-backed tensor does not invalidate the buffer + ptr1 = pointer(C1) + tensorfree!(C1, buffer) + allocator_reset!(buffer, 0) + C3 = tensoralloc(CuArray{Float32, 2}, (8, 8), Val(true), buffer) + @test isbufferbacked(C3, buffer) + @test pointer(C3) == ptr1 + fill!(C3, 1.0f0) + @test all(isone, collect(C3)) + + # unified memory buffers hand out unified memory temporaries + ubuffer = CUDABufferAllocator(; sizehint = 4096, memory = UnifiedMemory) + C4 = tensoralloc(CuArray{Float64, 1}, (16,), Val(true), ubuffer) + @test C4 isa CuArray{Float64, 1, UnifiedMemory} + @test isbufferbacked(C4, ubuffer) + end + + @testset "storage mismatch falls back" begin + # a host buffer cannot back CuArrays + hostbuffer = BufferAllocator(; sizehint = 4096) + C1 = tensoralloc(CuArray{Float32, 2}, (8, 8), Val(true), hostbuffer) + @test C1 isa CuArray{Float32, 2} + @test hostbuffer.offset == 0 + + # a device buffer cannot back Arrays + cubuffer = CUDABufferAllocator(; sizehint = 4096) + C2 = tensoralloc(Array{Float64, 2}, (8, 8), Val(true), cubuffer) + @test C2 isa Array{Float64, 2} + @test cubuffer.offset == 0 + end + + @testset "alignment" begin + # cuTENSOR is sensitive to this: it only selects its fastest kernels for + # 256-byte aligned data, so every temporary has to be padded to that + buffer = CUDABufferAllocator(; sizehint = 8192) + @test iszero(UInt(pointer(buffer)) % 256) + + # a deliberately misaligning allocation of 3 bytes + C1 = tensoralloc(CuArray{UInt8, 1}, (3,), Val(true), buffer) + @test isbufferbacked(C1, buffer) + @test buffer.offset == 3 + for T in (Float32, Float64, ComplexF32, ComplexF64) + C2 = tensoralloc(CuArray{T, 1}, (4,), Val(true), buffer) + @test isbufferbacked(C2, buffer) + @test iszero(UInt(pointer(C2)) % 256) + end + + # host buffers stay at the alignment Julia actually guarantees + @test TensorOperations.buffer_alignment(BufferAllocator()) == 16 + end + + @testset "checkpoint and reset" begin + buffer = CUDABufferAllocator(; sizehint = 4096) + cp0 = allocator_checkpoint!(buffer) + @test cp0 == 0 + + C1 = tensoralloc(CuArray{Float32, 2}, (8, 8), Val(true), buffer) + cp1 = allocator_checkpoint!(buffer) + @test cp1 > cp0 + C2 = tensoralloc(CuArray{Float32, 2}, (8, 8), Val(true), buffer) + @test pointer(C2) != pointer(C1) + + allocator_reset!(buffer, cp1) + @test buffer.offset == cp1 + @test_throws ArgumentError allocator_reset!(buffer, cp1 + 10) + + allocator_reset!(buffer, cp0) + @test isempty(buffer) + end + + @testset "tensor network ($T)" for T in + (Float32, Float64, ComplexF32, ComplexF64) + D1, D2, D3 = 30, 40, 20 + d1, d2 = 2, 3 + + A1 = CuArray(randn(T, D1, d1, D2)) + A2 = CuArray(randn(T, D2, d2, D3)) + ρₗ = CuArray(randn(T, D1, D1)) + ρᵣ = CuArray(randn(T, D3, D3)) + H = CuArray(randn(T, d1, d2, d1, d2)) + + @tensor begin + HRAA1[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + + buffer = CUDABufferAllocator() + @tensor allocator = buffer begin + HRAA2[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test HRAA2 isa CuArray{T, 4} + @test collect(HRAA2) ≈ collect(HRAA1) + + # all temporaries were reclaimed, and the buffer was actually used + @test buffer.offset == 0 + @test buffer.max_offset > 0 + + # The high-water mark only counts the temporaries that actually fit in the + # buffer, so it may still grow while the buffer is warming up, but it has to + # converge to a fixed size after a couple of contractions. + max0 = buffer.max_offset + for _ in 1:5 + @tensor allocator = buffer begin + HRAA3[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test collect(HRAA3) ≈ collect(HRAA1) + end + max1 = buffer.max_offset + @test max1 ≥ max0 + @test length(buffer) ≥ max1 + + for _ in 1:5 + @tensor allocator = buffer begin + HRAA3[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test collect(HRAA3) ≈ collect(HRAA1) + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + + # scalar output + @tensor begin + E1 = ρₗ[a', a] * A1[a, s, b] * A2[b, s', c] * ρᵣ[c, c'] * + H[t, t', s, s'] * conj(A1[a', t, b']) * conj(A2[b', t', c']) + end + @tensor allocator = buffer begin + E2 = ρₗ[a', a] * A1[a, s, b] * A2[b, s', c] * ρᵣ[c, c'] * + H[t, t', s, s'] * conj(A1[a', t, b']) * conj(A2[b', t', c']) + end + @test E1 ≈ E2 + @test buffer.offset == 0 + end + + @testset "host inputs are promoted to CuArray" begin + D1, D2, D3 = 30, 40, 20 + d1, d2 = 2, 3 + T = Float64 + + A1 = randn(T, D1, d1, D2) + A2 = randn(T, D2, d2, D3) + ρₗ = randn(T, D1, D1) + ρᵣ = randn(T, D3, D3) + H = randn(T, d1, d2, d1, d2) + + @tensor begin + HRAA1[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + + for memory in (DeviceMemory, UnifiedMemory) + buffer = CUDABufferAllocator(; memory) + @tensor backend = cuTENSORBackend() allocator = buffer begin + HRAA2[a, s1, s2, c] := ρₗ[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * + ρᵣ[c', c] * H[s1, s2, t1, t2] + end + @test HRAA2 isa CuArray{T, 4} + @test collect(HRAA2) ≈ HRAA1 + @test buffer.offset == 0 + @test buffer.max_offset > 0 + end + end + + @testset "ncon" begin + A = CuArray(randn(Float32, 5, 5)) + B = CuArray(randn(Float32, 5, 5)) + C = CuArray(randn(Float32, 5, 5)) + buffer = CUDABufferAllocator() + + R = ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + @test R isa CuArray{Float32, 2} + @test collect(R) ≈ collect(A) * collect(B) * collect(C) + @test buffer.offset == 0 + + max1 = buffer.max_offset + for _ in 1:5 + ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + end + end + @testset "@cutensor" verbose = true begin @testset "tensorcontract 1" begin A = randn(Float64, (3, 5, 4, 6))