onnx_runtime_memory/
oracle.rs

1//! The **size oracle**: byte size of an activation value.
2//!
3//! Sizes may be unknown at build time when a shape has symbolic (dynamic)
4//! dimensions. The planner is generic over any `Fn(ValueId) -> Option<usize>`
5//! so the *same* algorithm serves two callers:
6//!
7//! * **build-time** planning from fully-static shapes ([`static_size_oracle`]),
8//!   which returns `None` for any symbolic-shaped value and drives the planner
9//!   to a [`crate::PlanStatus::Deferred`] result; and
10//! * **run-time** planning, where the executor supplies a closure backed by the
11//!   resolved concrete shapes for the current run.
12
13use std::collections::HashMap;
14
15use onnx_runtime_ir::{Dim, Graph, SymbolId, ValueId, as_static_shape};
16
17/// A size oracle closure that sizes values from their fully-static shapes.
18///
19/// Returns `None` for any symbolic-shaped value, which the planner reports as
20/// [`crate::PlanStatus::Deferred`] so the executor can re-plan once shapes
21/// resolve.
22pub fn static_size_oracle(graph: &Graph) -> impl Fn(ValueId) -> Option<usize> + '_ {
23    move |value| static_size(graph, value)
24}
25
26/// Byte size of a value from its *static* shape, or `None` if any dimension is
27/// symbolic (unknown until run time) or the element count overflows `usize`.
28///
29/// Uses [`onnx_runtime_ir::DataType::checked_storage_bytes`] so sub-byte packed
30/// types (`int4`/`uint4`/`float4`) are sized correctly and an overflowing
31/// element count becomes `None` rather than a wrapped under-count.
32pub fn static_size(graph: &Graph, value: ValueId) -> Option<usize> {
33    let val = graph.try_value(value)?;
34    let dims = as_static_shape(&val.shape)?;
35    let mut numel: usize = 1;
36    for d in dims {
37        numel = numel.checked_mul(d)?;
38    }
39    val.dtype.checked_storage_bytes(numel)
40}
41
42/// A size oracle that resolves symbolic dimensions to caller-supplied **upper
43/// bounds**.
44///
45/// [`static_size_oracle`] answers `None` for anything dynamic, which is right
46/// for a plan that must be exact but useless for a *reservation*: an LLM's
47/// activations are dynamic in sequence length, so a reservation computed from
48/// static shapes alone is always zero — and a zero reservation is
49/// indistinguishable from a model that allocates nothing.
50///
51/// A reservation does not need the exact size, it needs the ceiling. Admission
52/// control already knows that ceiling, because it is the largest shape it will
53/// admit. Binding those bounds turns "cannot know" into "cannot exceed".
54///
55/// Symbols with no bound are still `None`, so the planner defers rather than
56/// guessing. Partial knowledge is not a bound.
57pub fn bounded_size_oracle<'a>(
58    graph: &'a Graph,
59    bounds: &'a HashMap<SymbolId, usize>,
60) -> impl Fn(ValueId) -> Option<usize> + 'a {
61    move |value| bounded_size(graph, value, bounds)
62}
63
64/// Byte size of a value with symbolic dimensions resolved through `bounds`.
65///
66/// Returns `None` if any dimension is symbolic and unbound, or if the element
67/// count overflows.
68pub fn bounded_size(
69    graph: &Graph,
70    value: ValueId,
71    bounds: &HashMap<SymbolId, usize>,
72) -> Option<usize> {
73    let val = graph.try_value(value)?;
74    let mut numel: usize = 1;
75    for dim in &val.shape {
76        let extent = match *dim {
77            Dim::Static(extent) => extent,
78            Dim::Symbolic(symbol) => *bounds.get(&symbol)?,
79        };
80        numel = numel.checked_mul(extent)?;
81    }
82    val.dtype.checked_storage_bytes(numel)
83}
onnx_runtime_memory/oracle.rs

onnx_runtime_memory/
oracle.rs
