Buffers

Buffers module

This module contains functions to handle buffers.

The Buffer object is used to store data of type T with an offset, while the ThreadsBuffer object is used to store data of type T with an offset for each thread.

The buffers are used to store data in a contiguous memory block and to avoid memory allocation in loops. The buffers can be used with alloc! to allocate tensors of given dimensions, drop! to drop tensors from the buffer, and reset! to reset the buffer to the initial state.

Alternativelly, the buffers can be reshaped with reshape_buf! to use the same memory block for different tensors or to allocate tensors with a specific offset.

The size of the buffer can be extended if necessary, and the buffer can be set to be extendable (default) or not at construction with Buffer or later with set_extendable!.

In any case, the ::ThreadsBuffer buffers should be released after use with Buffers.release! or reset!.

If some functions complain about tensors being aliases or if the tensors will be used in C, the neuralyze function can be used to wipe the memory about the origin of the tensor. Do not use this function if the size of the tensor might be changed in between, i.e., neuralyze the tensor only after all necessary allocations are done.

Buffer{T}

Buffer object to store data of type T with an offset.

The buffer allocates an extra element at the beginning which is used to check if the buffer can be extended and to ensure that pointers to the allocated arrays will never point to the same memory as the buffer.

If the buffer is used with reshape_buf!, the offset is set to zero.

MAllocBuffer{T}

Buffer object to store data of type T with an offset. The buffer memory is manually allocated and deallocated and is not extendable.

The buffer allocates an extra element at the beginning which is used to check if the buffer can be extended and to ensure that pointers to the allocated arrays will never point to the same memory as the buffer.

If the buffer is used with reshape_buf!, the offset is set to zero.

Example

julia> @buffer buf(100) begin
         A = alloc!(buf, 10, 10)
         B = alloc!(buf, 10, 10)
         C = alloc!(buf, 10, 10)
         A .= 1.0
         B .= 2.0
         @tensor C[i,j] = A[i,l] * B[l,j]
       end

ThreadsBuffer{T}

Buffer object to store data of type T for each thread.

By default, the buffer is created for nthreads() threads, i.e., each thread has its own buffer Buffer.

Create the buffer with ThreadsBuffer{T}(len, nbuf=Threads.nthreads()) and use it with alloc!, drop!, reset!, etc.

Example

julia> buf = Buffer(10000)
julia> C = alloc!(buf, 10, 10, 20) # 10x10x20 destination tensor on a single thread
julia> @threadsbuffer tbuf(1000) begin # 1000 elements buffer for nthreads() threads each
julia>  Threads.@threads for k = 1:20
          A = alloc!(tbuf, 10, 10) # 10x10 tensor
          B = alloc!(tbuf, 10, 10) # 10x10 tensor
          rand!(A)
          rand!(B)
          @tensor C[:,:,k][i,j] = A[i,l] * B[l,j]
          reset!(tbuf)
        end
       end

ThreadsMAllocBuffer{T}

MAllocBuffer object to store data of type T for each thread.

By default, the buffer is created for nthreads() threads, i.e., each thread has its own buffer MAllocBuffer.

Create the buffer with ThreadsMAllocBuffer{T}(len, nbuf=Threads.nthreads()) and use it with alloc!, drop!, reset!, etc.

Example

@buffer buf(10000) begin
  C = alloc!(buf, 10, 10, 20) # 10x10x20 destination tensor on a single thread
  @threadsbuffer tbuf(1000) begin # 1000 elements buffer for nthreads() threads each
    @sync for k = 1:20
      Threads.@spawn begin
        A = alloc!(tbuf, 10, 10) # 10x10 tensor
        B = alloc!(tbuf, 10, 10) # 10x10 tensor
        rand!(A)
        rand!(B)
        @tensor C[:,:,k][i,j] = A[i,l] * B[l,j]
        reset!(tbuf)
      end 
    end
  end # free threadsbuffer tbuf
end # free buffer buf

alloc!(buf, dims...; extend=true)

Allocate tensor of given dimensions in buffer buf.

The tensor is allocated in the buffer starting at the current offset. The offset is increased by the length of the tensor. If extend=true, the buffer is extended if necessary. For ThreadsBuffer, the tensor is allocated in the buffer of the current thread.

Return the allocated tensor.

julia> buf = Buffer(100000)
julia> A = alloc!(buf, 10, 10, 20) # 10x10x20 tensor
julia> B = alloc!(buf, 10, 10, 10) # 10x10x10 tensor starting after A
julia> C = alloc!(buf, 10, 20) # 10x20 tensor starting after B
julia> rand!(B)
julia> rand!(C)
julia> An = neuralyze(A) # tensor without origin
julia> @tensor An[i,j,k] = B[i,j,l] * C[l,k]

drop!(buf, tensor...)

Drop tensor(s) from buffer buf.

Only last tensors can be dropped. For ThreadsBuffer, drop tensors from the buffer of the current thread.

free!(buf)

Free the memory of buffer buf.

isextendable(buf)

Check if buffer buf is extendable.

nbuffers(buf::AbstractThreadsBuffer)

Return the number of buffers in buf::ThreadsBuffer.

neuralyze(tensor::AbstractArray)

Wipe the memory about origin of tensor.

tensor is a (contiguous!) array that is a (possibly reshaped) view of a larger array. Return the same tensor pointing to the same memory, but without the information about the origin. To be used together with alloc! or reshape_buf! to trick Base.mightalias.

Example

julia> buf = Buffer(100000)
julia> A = alloc(buf, 10, 10, 20) # 10x10x20 tensor
julia> B = alloc(buf, 10, 10, 10) # 10x10x10 tensor starting after A
julia> C = alloc(buf, 10, 20) # 10x20 tensor starting after B
julia> rand!(B)
julia> rand!(C)
julia> An = neuralyze(A) # tensor without origin but pointing to the same memory
julia> @tensor An[i,j,k] = B[i,j,l] * C[l,k]

pseudo_alloc!(buf, dims...)

Pseudo allocation function to calculate length for buffer.

The function is used in combination with @print_buffer_usage.

pseudo_drop!(buf, lens...)

Pseudo drop function to calculate length for buffer.

The function is used in combination with @print_buffer_usage.

pseudo_reset!(buf)

Pseudo reset function to calculate length for buffer.

The function is used in combination with @print_buffer_usage.

repair!(buf::AbstractThreadsBuffer)

Repair ThreadsBuffer buf by releasing all buffers and resetting the pool.

This function should be used after the threaded loop if the buffers were not released properly.

reset!(buf)

Reset buffer buf to the initial state. For ThreadsBuffer, reset the buffer of the current thread and release it.

reshape_buf!(buf, dims...; offset=0, extend=true)

Reshape (part of) a buffer to given dimensions (without copying), using offset.

For ThreadsBuffer, reshape the buffer of the current thread. Call reset!(::ThreadsBuffer) or release! after use.

It can be used, e.g., for itermediates in tensor contractions.

Example

julia> buf = Buffer(100000)
julia> A = reshape_buf!(buf, 10, 10, 20) # 10x10x20 tensor
julia> B = reshape_buf!(buf, 10, 10, 10, offset=2000) # 10x10x10 tensor starting at 2001
julia> B .= rand(10,10,10)
julia> C = rand(10,20)
julia> @tensor A[i,j,k] = B[i,j,l] * C[l,k]

set_extendable!(buf, extend=true)

Set buffer buf to be extendable or not.

used(buf)

Return the number of elements used in buffer buf.

If the buffer is used with reshape_buf!, -1 is returned.

For ThreadsBuffer, return the number of elements used in the buffer of the current thread.

with_buffer(f::Function, buf::ThreadsBuffer)

Execute function f with buffer buf.

The buffer is released after the function is executed.

Example

julia> buf = Buffer(10000)
julia> C = alloc!(buf, 10, 10, 20) # 10x10x20 destination tensor on a single thread
julia> tbuf = ThreadsBuffer(1000)
julia> Threads.@threads for k = 1:20
          with_buffer(tbuf) do bu
            A = alloc!(bu, 10, 10) # 10x10 tensor
            B = alloc!(bu, 10, 10) # 10x10 tensor
            rand!(A)
            rand!(B)
            @tensor C[:,:,k][i,j] = A[i,l] * B[l,j]
          end
        end

@buffer(specs, specs2, ex)

Create two MAllocBuffer objects with the given specifications specs and specs2 and execute the expression ex with the buffers.

@buffer(specs, ex)

Create a MAllocBuffer object with the given specifications specs and execute the expression ex with the buffer.

The specs have the following format:

• buf(100) creates a buffer buf of type Float64 with length 100.
• buf(Int, 100) creates a buffer buf of type Int with length 100.

The buffer is freed after the expression is executed.

Example

@buffer buf(Float64, 300) begin
  A = alloc!(buf, 10, 10)
  B = alloc!(buf, 10, 10)
  C = alloc!(buf, 10, 10)
  rand!(A)
  rand!(B)
  @tensor C[i,j] = A[i,l] * B[l,j]
end

@print_buffer_usage(buf, ex)

Print buffer buf usage in expression ex.

The macro generates a body of a function that calculates the length of buffer buf in expression ex. It is possible to use the macro with multiple buffers, e.g., @print_buffer_usage buf1 buf2 begin ... end.

All function calls with buf as an argument are replaced with pseudo_<function> calls. pseudo_alloc!,pseudo_drop!, and pseudo_reset! functions are pre-defined, custom pseudo_functions can be defined if necessary.

Example

```julia buf = Buffer(100000) @printbufferusage buf begin if true A = alloc!(buf, 10, 10, 20) else A = alloc!(buf, 10, 10, 30) end B = alloc!(buf, 10, 10, 10) if true C = alloc!(buf, 10, 20) else C = alloc!(buf, 10, 30) end rand!(B) rand!(C) An = neuralyze(A) @tensor An[i,j,k] = B[i,j,l] * C[l,k] drop!(buf, B, C) reset!(buf) end

@threadsbuffer(specs, specs2, ex)

Create two ThreadsMAllocBuffer objects with the given specifications and execute the expression ex.

@threadsbuffer(specs, ex)

Create a ThreadsMAllocBuffer object with the given specifications and execute the expression ex.

The specifications specs can be:

• buf(100) creates a buffer buf of type Float64 with length 100.
• buf(Int, 100) creates a buffer buf of type Int with length 100.
• buf(Int, 100, 4) creates a buffer buf of type Int with length 100 for 4 threads.

_buffer_usage(ex, bufs)

Allocations and deallocations together with corresponding ifs in expression ex.

current_buffer(buf::AbstractThreadsBuffer)

Return the buffer of the current thread.

If the buffer is not available, wait until it is released.

current_buffer_index(buf::AbstractThreadsBuffer)

Return the index of the buffer of the current thread.

If the buffer is not available, wait until it is released.

iscontiguous_tensor(tensor::AbstractArray)

Check if tensor is contiguous.

Return true if tensor is a Vector or a SubArray that is contiguous.

release!(buf::AbstractThreadsBuffer)

Release buffer of the current thread.