Module `Tolk.Allocations`

Buffer allocation for tensor graphs.

Decides which tensor computations need explicit buffer allocations and transforms a lazy tensor-level Tolk_uop.Ops.t.Sink into a Tolk_uop.Ops.t.Call with allocated buffers.

The transformation runs in three phases:

Tag. Identify nodes that need realization (Tolk_uop.Ops.t.Contiguous, Tolk_uop.Ops.t.After+Tolk_uop.Ops.t.Store, and non-trivial bases of the sink's children).
Allocate. Replace tagged nodes with explicit Tolk_uop.Ops.t.Buffer + Tolk_uop.Ops.t.Store + Tolk_uop.Ops.t.After sequences. When movement ops on a buffer collapse to a contiguous range, a Tolk_uop.Ops.t.Slice is used instead.
Finalize. Strip internal bookkeeping, collect the resulting stores, replace input buffers with Tolk_uop.Ops.t.Param nodes for cache-key normalisation, and wrap everything in a Tolk_uop.Ops.t.Call.

The returned buffer_map tracks which original tensor nodes map to which allocated buffers, keyed by Tolk_uop.Uop.tag.

val transform_to_call : 
  Tolk_uop.Uop.t ->
  Tolk_uop.Uop.t * (int, Tolk_uop.Uop.t) Stdlib.Hashtbl.t

transform_to_call big_sink is (call, buffer_map).

big_sink must be a Tolk_uop.Ops.t.Sink node representing the lazy tensor graph to be realized.

call is a Tolk_uop.Ops.t.Call whose callee is a parameterised sink (input buffers replaced by Tolk_uop.Ops.t.Param nodes) and whose arguments are the original buffer and bind nodes.

buffer_map maps original tensor nodes to their allocated buffers, keyed by Tolk_uop.Uop.tag. Downstream scheduling uses this to resolve tensor references to concrete buffers.