Examples of improper usage

In the following, we will use our functions:

julia> using AllocArrays
julia> function some_allocating_function(a)
           b = similar(a)
           b .= a
           c = similar(a)
           c .= a
           return (; b, c)
       endsome_allocating_function (generic function with 1 method)
julia> function basic_reduction(a)
           (; b, c) = some_allocating_function(a)
           return sum(b .+ c)
       endbasic_reduction (generic function with 1 method)

Wrong: allowing escapes

Allocations inside @no_escape must not escape!

julia> function bad_function_1(a)
           b = BumperAllocator(2^25) # 32 MiB
           output = []
           with_allocator(b) do
               result = some_allocating_function(a)
               push!(output, result.b) # wrong! `b` is escaping `@no_escape`!
               reset!(b)
           end
           return sum(output...)
       endbad_function_1 (generic function with 1 method)
julia> bad_function_1(CheckedAllocArray([1]))ERROR: InvalidMemoryException: Array accessed after its memory has been deallocated.

Here is a corrected version:

julia> function good_function_1(a)
           b = BumperAllocator(2^25) # 32 MiB
       
           # note, we are not inside `with_allocator`, so we are not making buffer-backed memory
           output = similar(a)
       
           with_allocator(b) do
               result = some_allocating_function(a)
               output .= result.b # OK! we are copying buffer-backed memory into our heap-allocated memory
               reset!(b)
           end
           return sum(output)
       endgood_function_1 (generic function with 1 method)
julia> good_function_1(AllocArray([1]))1

Wrong: resetting a buffer in active use

julia> function bad_function_2(a)
           b = BumperAllocator(2^25) # 32 MiB
           output = Channel(Inf)
           with_allocator(b) do
               @sync for _ = 1:10
                   Threads.@spawn begin
                       scalar = basic_reduction(a)
                       put!(output, scalar)
                       reset!(b) # wrong! we cannot reset here as `b` is being used on other tasks
                   end
               end
           end
           close(output)
           return sum(collect(output))
       endbad_function_2 (generic function with 1 method)
julia> bad_function_2(CheckedAllocArray([1]))20

Here is a corrected version:

julia> function good_function_2(a)
           b = BumperAllocator(2^25) # 32 MiB
           output = Channel(Inf)
           with_allocator(b) do
               @sync for _ = 1:10
                   Threads.@spawn begin
                       scalar = basic_reduction(a)
                       put!(output, scalar)
                   end
               end
               reset!(b) # OK! resetting once we no longer need the allocations
           end
           close(output)
           return sum(collect(output))
       endgood_function_2 (generic function with 1 method)
julia> good_function_2(AllocArray([1]))20

Or, if we need to reset multiple times as we process the data, we could do a serial version:

julia> function good_function_2b(a)
           b = BumperAllocator(2^25) # 32 MiB
           output = Channel(Inf)
           with_allocator(b) do
               for _ = 1:10
                   scalar = basic_reduction(a)
                   put!(output, scalar)
                   reset!(b) # OK to reset here! buffer-backed memory is not being used
               end
           end
           close(output)
           return sum(collect(output))
       endgood_function_2b (generic function with 1 method)
julia> good_function_2b(AllocArray([1]))20

We could also do something in-between, by launching tasks in batches of n, then resetting the buffer between them. Or we could use multiple buffers.

Wrong: neglecting to reset the buffer

As shown above, we must be careful about when we reset the buffer. However, if we never reset it (analogous to never garbage collecting), we run into another problem which is we will run out memory!

function bad_function_3(a, N)
    b = BumperAllocator(2^25) # 32 MiB
    output = Channel(Inf)
    with_allocator(b) do
        for _ = 1:N # bad! we are going to allocate `N` times without resetting!
            # if `N` is very large, we will run out of memory.
            scalar = basic_reduction(a)
            put!(output, scalar)
        end
    end
    close(output)
    return sum(collect(output))
end

This can be fixed by resetting appropriately, as in e.g. good_function_2b above.