diff --git a/docs/src/man/backends.md b/docs/src/man/backends.md index 9971bd16..7ac3cdff 100644 --- a/docs/src/man/backends.md +++ b/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. diff --git a/ext/TensorOperationsAMDGPUExt.jl b/ext/TensorOperationsAMDGPUExt.jl index d77e3669..a1aaa484 100644 --- a/ext/TensorOperationsAMDGPUExt.jl +++ b/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 diff --git a/ext/TensorOperationsJLArraysExt.jl b/ext/TensorOperationsJLArraysExt.jl index 1e8fee10..89a69512 100644 --- a/ext/TensorOperationsJLArraysExt.jl +++ b/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 diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl index 835d8147..6ba192b1 100644 --- a/src/implementation/allocator.jl +++ b/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 diff --git a/test/allocator.jl b/test/allocator.jl index c1bca423..87237558 100644 --- a/test/allocator.jl +++ b/test/allocator.jl @@ -1,9 +1,11 @@ using TensorOperations using TensorOperations: BufferAllocator, DefaultAllocator, ManualAllocator +using TensorOperations: JLBufferAllocator using TensorOperations: tensoralloc, tensorfree!, tensoralloc_add, tensoralloc_contract using TensorOperations: allocator_checkpoint!, allocator_reset! using Test using LinearAlgebra +using JLArrays @testset "BufferAllocator" begin @testset "Constructor and basic properties" begin @@ -148,3 +150,214 @@ using LinearAlgebra @test buffer.max_offset == max1 end end + +# `JLArrays` is the reference GPU array implementation, so a `JLArray`-backed buffer exercises +# the foreign-storage code paths of `BufferAllocator` -- the same ones that `CUDABufferAllocator` +# and `AMDBufferAllocator` rely on -- without requiring any GPU hardware. +@testset "JLArray-backed BufferAllocator" verbose = true begin + # 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 = JLBufferAllocator(; sizehint = 1024) + @test buffer isa BufferAllocator{JLArray{UInt8, 1}} + @test buffer isa typeof(BufferAllocator{JLArray{UInt8, 1}}(; sizehint = 1024)) + @test length(buffer) == 1024 + @test isempty(buffer) + @test buffer.offset == 0 + + # 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 = JLBufferAllocator(; sizehint = 4096) + + # temporaries are taken from the buffer + C1 = tensoralloc(JLArray{Float32, 2}, (8, 8), Val(true), buffer) + @test C1 isa JLArray{Float32, 2} + @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(JLArray{Float32, 2}, (8, 8), Val(false), buffer) + @test C2 isa JLArray{Float32, 2} + @test !isbufferbacked(C2, buffer) + @test buffer.offset == offset + + # freeing a buffer-backed tensor does not invalidate the buffer: `unsafe_buffer_wrap` + # only retains a reference to it, so freeing merely releases that reference again + ptr1 = pointer(C1) + tensorfree!(C1, buffer) + allocator_reset!(buffer, 0) + C3 = tensoralloc(JLArray{Float32, 2}, (8, 8), Val(true), buffer) + @test isbufferbacked(C3, buffer) + @test pointer(C3) == ptr1 + fill!(C3, 1.0f0) + @test all(isone, collect(C3)) + + # a bare length is accepted for vectors + allocator_reset!(buffer, 0) + C4 = tensoralloc(JLArray{Float64, 1}, 16, Val(true), buffer) + @test C4 isa JLArray{Float64, 1} + @test size(C4) == (16,) + @test isbufferbacked(C4, buffer) + end + + @testset "storage mismatch falls back" begin + # a host buffer cannot back JLArrays + hostbuffer = BufferAllocator(; sizehint = 4096) + C1 = tensoralloc(JLArray{Float32, 2}, (8, 8), Val(true), hostbuffer) + @test C1 isa JLArray{Float32, 2} + @test hostbuffer.offset == 0 + + # a JLArray buffer cannot back Arrays + jlbuffer = JLBufferAllocator(; sizehint = 4096) + C2 = tensoralloc(Array{Float64, 2}, (8, 8), Val(true), jlbuffer) + @test C2 isa Array{Float64, 2} + @test jlbuffer.offset == 0 + end + + @testset "alignment" begin + buffer = JLBufferAllocator(; sizehint = 8192) + @test TensorOperations.buffer_alignment(buffer) == 16 + @test iszero(UInt(pointer(buffer)) % 16) + + # a deliberately misaligning allocation of 3 bytes + C1 = tensoralloc(JLArray{UInt8, 1}, (3,), Val(true), buffer) + @test isbufferbacked(C1, buffer) + @test buffer.offset == 3 + for T in (Float32, Float64, ComplexF32, ComplexF64) + C2 = tensoralloc(JLArray{T, 1}, (4,), Val(true), buffer) + @test isbufferbacked(C2, buffer) + @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 + @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 + end + + @testset "checkpoint and reset" begin + buffer = JLBufferAllocator(; sizehint = 4096) + cp0 = allocator_checkpoint!(buffer) + @test cp0 == 0 + + C1 = tensoralloc(JLArray{Float32, 2}, (8, 8), Val(true), buffer) + cp1 = allocator_checkpoint!(buffer) + @test cp1 > cp0 + C2 = tensoralloc(JLArray{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 = JLArray(randn(T, D1, d1, D2)) + A2 = JLArray(randn(T, D2, d2, D3)) + ρₗ = JLArray(randn(T, D1, D1)) + ρᵣ = JLArray(randn(T, D3, D3)) + H = JLArray(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 = JLBufferAllocator() + @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 JLArray{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 "ncon" begin + A = JLArray(randn(Float32, 5, 5)) + B = JLArray(randn(Float32, 5, 5)) + C = JLArray(randn(Float32, 5, 5)) + buffer = JLBufferAllocator() + + R = ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + @test R isa JLArray{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 diff --git a/test/gpu.jl b/test/gpu.jl index 592a1702..8b716906 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -28,6 +28,11 @@ AMDGPU.functional() && push!(ATs, ROCArray) backends = [StridedBLAS(), StridedNative()] +# storage-specific `BufferAllocator` constructor for each of the array types above +bufferallocator(::Type{JLArray}; kwargs...) = TensorOperations.JLBufferAllocator(; kwargs...) +bufferallocator(::Type{CuArray}; kwargs...) = TensorOperations.CUDABufferAllocator(; kwargs...) +bufferallocator(::Type{ROCArray}; kwargs...) = TensorOperations.AMDBufferAllocator(; kwargs...) + @testset "tensoradd! ($AT)" for AT in ATs sz = (3, 5, 4, 6) p = (3, 1, 4, 2) @@ -123,3 +128,85 @@ end end end + +@testset "BufferAllocator ($AT)" for AT in ATs + @testset "tensor network ($T)" for T in (Float32, ComplexF32) + D1, D2, D3 = 30, 40, 20 + d1, d2 = 2, 3 + + A1 = adapt(AT, randn(T, D1, d1, D2)) + A2 = adapt(AT, randn(T, D2, d2, D3)) + ρₗ = adapt(AT, randn(T, D1, D1)) + ρᵣ = adapt(AT, randn(T, D3, D3)) + H = adapt(AT, 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 = bufferallocator(AT) + @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 AT{T, 4} + @test test_result(HRAA2, HRAA1) + + # all temporaries are reclaimed, and the buffer is actually used + @test buffer.offset == 0 + @test buffer.max_offset > 0 + + # the high-water mark only counts temporaries that actually fit, so it may still grow + # while the buffer is warming up, but has to converge to a fixed size afterwards + for _ in 1:3 + @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 test_result(HRAA3, HRAA1) + end + max1 = buffer.max_offset + for _ in 1:3 + @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 + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + @test length(buffer) ≥ 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 + + # chain contraction through `ncon`, which reclaims its temporaries via checkpoints + @testset "ncon" begin + A = adapt(AT, randn(Float32, 5, 5)) + B = adapt(AT, randn(Float32, 5, 5)) + C = adapt(AT, randn(Float32, 5, 5)) + buffer = bufferallocator(AT) + + R = ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + @test R isa AT{Float32, 2} + @test test_result(R, collect(A) * collect(B) * collect(C)) + @test buffer.offset == 0 + + max1 = buffer.max_offset + for _ in 1:3 + ncon([A, B, C], [[-1, 1], [1, 2], [2, -2]]; allocator = buffer) + end + @test buffer.offset == 0 + @test buffer.max_offset == max1 + end +end