class

`SushiAI::Graph::TensorPool`

Owns device memory buffers backing graph values across steps.

Declared in: include/SushiAI/graph/lowering.hpp

Maintains persistent tensor storages and arena slices across replays.

Public member functions

explicit TensorPool(SushiBLAS::Engine &engine) noexcept

Constructs an empty pool drawing memory from engine.

Parameters

engine: The engine whose context owns the allocator; must outlive the pool.

void adopt(MemoryPlan plan)

Allocates arenas from the specified memory plan.

Must be called before any planned value is materialised.

Parameters

plan: Arena placement plan to adopt.

Exceptions

Error: If a plan was already adopted or allocation fails.

bool has_plan() const noexcept

Whether a plan has been adopted.

const ArenaSet & arenas() const

The adopted plan and its arenas.

Returns

The arena set.

Exceptions

Error: If no plan has been adopted.

SushiBLAS::Tensor create(const Shape &shape, DType dtype) const

Allocates an unmanaged standalone tensor.

Parameters

shape: Extents of the tensor.
dtype: Element data type.

Returns

Newly allocated tensor instance.

void bind(ValueId id, SushiBLAS::Tensor tensor)

Binds an externally owned tensor to the given value id.

Parameters

id: Value identifier to bind.
tensor: External tensor to associate.

Exceptions

Error: If the value is already bound.

bool is_bound(ValueId id) const noexcept

Whether id has a tensor behind it yet.

SushiBLAS::Tensor & tensor(ValueId id)

The tensor behind id.

Parameters

id: The value to look up.

Returns

A reference stable for the pool's lifetime.

Exceptions

Error: If id is not bound.

const SushiBLAS::Tensor & tensor(ValueId id) const

Const overload of tensor().

SushiBLAS::Tensor & materialise(ValueId id, const Value &value)

Returns id's tensor, allocating it on first request.

Parameters

id: The value to materialise.
value: The value's description, read for shape and dtype.

Returns

A reference stable for the pool's lifetime.

void materialise_leaves(const Graph &graph)

Allocates uninitialised storage for unbound graph leaves.

Parameters

graph: Graph whose unbound input leaves are allocated.

void alias(ValueId target, ValueId source)

Aliases target value storage to an existing source tensor.

Parameters

target: Value identifier to alias.
source: Already-bound source value identifier.

Exceptions

Error: If source is unbound or target is already bound elsewhere.

SushiBLAS::Tensor & ones(const Shape &shape, DType dtype)

Retrieves or creates a shared constant all-ones tensor.

Parameters

shape: Extents of the tensor.
dtype: Element data type.

Returns

Stable reference to the allocated all-ones tensor.

SushiBLAS::Tensor & scratch(NodeId node, int slot, const Shape &shape, DType dtype)

Retrieves or creates a per-node scratch buffer.

Keyed by node and slot, reused across replays.

Returns

Stable reference to scratch tensor.

void request_constant(SushiBLAS::Tensor &tensor, double value)

Queues tensor to be filled with value once.

Parameters

tensor: The tensor to fill; must outlive the flush.
value: The constant.

void flush_constants()

Runs every queued constant fill and waits for it.

Records the fills into the engine's (by then empty) task graph, executes, and waits. Call it after the step has been compiled - never during recording, because it drains the builder.

std::size_t tensor_count() const noexcept

How many tensors the pool has handed out.

std::size_t allocated_bytes() const noexcept

Bytes actually allocated, excluding bound externals.

Every arena in full, plus the elements of every tensor that kept its own Storage. This is the number that shrinks when planning works; compare it against unpacked_bytes().

std::size_t unpacked_bytes() const noexcept

What the same tensors would cost one allocation each.

The counterfactual the saving is measured against: the sum of every tensor's own elements, packed or not. Equal to allocated_bytes() when no plan has been adopted.

SushiBLAS::Engine & engine() const noexcept

The engine this pool draws memory from.