Tensor Operations

MTensorOperations

Utilities for operations on tensors. This module provides wrappers around TensorOperations.jl and StridedViews.jl macros and functions.

The functions and macros are reexported by TensorTools module.

@buftensor buf tensor_expr

Use a Buffer from Buffers.jl to handle allocation of temporary tensors in tensor operations. The buffer offset is saved before the expression and restored afterwards, reclaiming all temporary allocations.

This is analogous to @butensor from TensorOperations (which uses Bumper.jl), but uses ElemCo's Buffers package instead.

The buffer can be pre-allocated to avoid runtime growth:

buf = Buffer(nvir^3 * 4)  # pre-allocate for estimated temp usage
@buftensor buf begin
  C[a,b,c] = A[a,d] * B[d,b,c]
  C[a,b,c] += D[a,d] * E[d,b,c]
end

If the buffer argument is omitted, the variable buf4tensor is used by default:

buf4tensor = Buffer(nvir^3 * 4)
@buftensor C[a,b,c] = A[a,d] * B[d,b,c]

Additional @tensor keyword arguments can be passed:

@buftensor buf opt=true C[a,b,c] := A[a,d] * B[d,e] * E[e,b,c]

mtensor(ex)

Macro for tensor operations with manual allocator.

@mview(ex)

StridedView based version of @view.

mview(arr, args...)

StridedView based version of view.

The data array is enforced to be a vector, such that the view is always a StridedView{..., Vector{...},...}.

mview(arr::StridedView, args...)

StridedView based version of view, for StridedView input.

Simply calls StridedViews.sview.

_buf_mightalias(A::StridedView, B::StridedView)

Check if two StridedViews might alias, using pointer-range overlap instead of Base.mightalias. This avoids false positives when both views reference non-overlapping regions of the same buffer.

_is_tensor_expr(ex)

Check if an expression looks like a tensor expression (indexed assignment or begin block) rather than a buffer variable/expression.

_memory_span(sv::StridedView{T})

Compute the byte range (lo, hi) of memory accessed by a StridedView, accounting for non-unit strides.

replace_ref_begin_end!(ex::Expr)

Replace begin and end in reference expression ex with firstindex and lastindex.

This is needed for macros that generate code with begin and end in references, because those macros are expanded before the actual indices are known.

TensorOperations.blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, buf::AbstractBuffer)

Override blas_contract! for buffer allocator. Saves and restores the buffer offset around the generic implementation to reclaim per-contraction internal temporaries (permutation buffers).

TensorOperations.stridedtensorcontract!(C, A, pA, B, pB, pAB, α, β,
                                         backend::StridedBLAS, allocator::AbstractBuffer)

Override for buffer allocator: uses pointer-range overlap check instead of Base.mightalias to avoid false aliasing when arrays share a buffer parent.

TensorOperations.tensoralloc(::Type{A}, structure, ::Val{istemp}, buf::AbstractBuffer)

Allocate tensor temporaries from a Buffer instead of using Julia's default allocator. Only temporary tensors (istemp=true) with non-zero dimensions are allocated from the buffer; output tensors that escape the @tensor block use default allocation.

TensorOperations.tensorfree!(C, ::AbstractBuffer)

No-op for buffer allocator. Temporary memory is reclaimed by offset restore in blas_contract! and @buftensor.