Module Tolk.Memory
Memory planning.
A schedule allocates a buffer for each intermediate value, though most are needed by a few of its calls only. Planning places the buffers whose lifetimes do not overlap at shared addresses of a few large buffers, the arenas, so the schedule needs less memory.
memory_plan_rewrite ~held_bufs linear is the schedule linear, an Op.t.Linear of calls run in order, with its buffers placed in arenas.
A buffer is planned unless it is among held_bufs (default []), whose contents outlive the schedule, or lives on a disk. A call takes the buffers of its arguments, an argument passing itself if it is a buffer and the buffers it views if it is a view (Op.t.Shrink, Op.t.Bitcast, Op.t.After). A buffer lives from the first call that takes it to the last, a loop around calls (Op.t.End or Op.t.Backedge) counting as one call that takes the buffers of each of its calls and a back edge's flag. A buffer's bytes are free from the call after its last: an arena is one allocation, and the executor runs the work that writes an allocation after every earlier use of it, a copy's included.
Each device has one arena, which its copies' buffers, those of the calls whose body is an Op.t.Store, share with its kernels': the executor runs a device's work in order, so on one device planning copies apart would order nothing more, and two arenas, held at once, would hold both peaks. Across devices the executor orders work by allocation: where a device whose work the host runs reuses bytes a GPU's copy touched, its work, and the run that submits it, wait for the GPU's earlier work. The arena is an Op.t.Buffer of Dtype.t.Int8 as large as its buffers need at once. A buffer's size and phase (Ops.param_arg) together are rounded up to 256 bytes, and its place is found by a Support_memory.Tlsf_allocator, the buffers ending at a call freed before those starting there are placed. Each planned buffer is replaced wherever linear reaches it outside call bodies, inside views included, by the bytes of its arena from its phase past its place (Shape.shrink), viewed as its type (Ops.bitcast), so that it keeps its phase. New arenas take the next slots (Ops.unique_num).
linear is itself if Setting.no_memory_planner is set or no buffer is planned. With Setting.debug at 1 or more, the memory saved is printed on standard output.
taken si is the storage that the entry si of a schedule takes, each once, in the order its calls take it: the buffer, call-local storage or parameter (Op.t.Buffer, Op.t.Alloc, Op.t.Param) that each argument of its calls names or views (Ops.storage_base). A loop around calls takes the storage of each of its calls and its back edge's flag, as memory_plan_rewrite counts a buffer's uses.