Skip to main content

onnx_runtime_session/
executor.rs

1//! The sequential CPU executor (Track D, `docs/ORT2.md` §20, §11.3).
2//!
3//! Turns a loaded [`Graph`] plus its live [`WeightStore`] into a runnable plan:
4//! resolve every value's concrete shape from the actual bound inputs at
5//! `run`, size a device buffer per value from those *resolved* shapes, resolve
6//! a kernel per node through the execution provider (keyed by the resolved
7//! input shapes), then walk the topological order binding
8//! [`TensorView`]/[`TensorMut`] windows over those buffers and invoking each
9//! kernel. It is generic over any [`Graph`] and any [`ExecutionProvider`]; the
10//! Phase-1 build wires in the CPU EP only, but nothing here is op- or
11//! model-specific.
12//!
13//! ## Symbolic → concrete shape resolution (§3.2, §11)
14//!
15//! Real models carry *symbolic* input dims (e.g. `batch`, `max_seq_len`): the
16//! loader produces a [`Shape`] whose dims are a mix of [`Dim::Static`] and
17//! [`Dim::Symbolic`]. This executor is model-agnostic about them — a symbol is
18//! whatever [`SymbolId`] the graph interned. At `run` it reads the actual shape
19//! of each bound input, **binds** the graph's symbols to concrete sizes from
20//! those inputs (conflicts across inputs are an error), and **substitutes**
21//! those bindings into every value's loader shape to obtain a fully-concrete
22//! shape. Buffers are sized from the resolved shapes and become run-scoped when
23//! shapes are dynamic (reused when the resolved shape is unchanged, re-allocated
24//! when it changes). A fully-static graph is simply the special case where
25//! there are no symbols: resolution is a no-op and every buffer/kernel is
26//! materialized once at build.
27//!
28//! The session does **not** infer op output shapes — that is the loader's job
29//! (the loader runs `onnx-runtime-shape-inference` at load time). If a value's
30//! loader shape still contains an unbound symbol after substitution, the
31//! session resolves genuinely data-dependent extents just-in-time during
32//! execution (see [`dynamic_output_shapes`]); anything it still cannot size is
33//! reported as [`SessionError::UnresolvedShape`] naming the value and its
34//! producing op, rather than guessing.
35//!
36//! ## Holden's precondition (ep-api safety review #1) — enforced here
37//!
38//! A [`TensorView`] carries no backing length, so it cannot self-check storage
39//! bounds. This executor owns every buffer, so it is the layer that *can*: for
40//! **every** input and output view of **every** node it calls
41//! [`strided::view_in_bounds`] (or, for sub-byte dtypes, the `storage_bytes`
42//! equivalent in [`view_bounds`]) against the **run-scoped resolved** buffer and
43//! refuses to dispatch on failure. That check is the sole thing that makes
44//! ep-cpu's unchecked pointer derefs sound.
45
46use std::collections::{HashMap, HashSet};
47use std::sync::Arc;
48
49use onnx_runtime_ep_api::{
50    DeviceBuffer, DevicePtr, DevicePtrMut, ExecutionProvider, KernelMatch, TensorMut, TensorView,
51};
52use onnx_runtime_ep_cpu::strided::view_in_bounds;
53use onnx_runtime_ep_cpu::CpuExecutionProvider;
54use onnx_runtime_ir::{
55    as_static_shape, compute_contiguous_strides, DataType, Dim, Graph, Node, NodeId, Shape,
56    SymbolId, TensorLayout, ValueId,
57};
58use onnx_runtime_loader::WeightStore;
59use onnx_runtime_shape_inference::{InferenceRegistry, MergePolicy};
60
61use crate::error::{Result, SessionError};
62use crate::sequence::{
63    concat_axis, split_axis, stack_new_axis, SeqTensor, SequenceValue,
64};
65use crate::tensor::{host_bytes, write_host, Tensor};
66
67/// A per-node compiled entry: the structural facts the run loop needs without
68/// re-deriving them from the graph. Shapes are **not** baked here — they are
69/// resolved per run from the bound inputs (see module docs).
70#[derive(Debug)]
71pub(crate) struct NodePlan {
72    pub node_id: NodeId,
73    /// Positional input value ids in ONNX signature order. An omitted optional
74    /// input (ONNX empty-string input name → `None` slot) is preserved as
75    /// `None` so a later present input is never misread as the omitted one
76    /// (e.g. `Slice(data, starts, ends, "", steps)`). Trailing `None`s are
77    /// trimmed — a truly absent trailing optional simply lowers the arity.
78    pub inputs: Vec<Option<ValueId>>,
79    /// Output value ids, in positional order.
80    pub outputs: Vec<ValueId>,
81    /// Element types of the inputs, positional (matches `inputs`). A `None`
82    /// input slot carries a placeholder dtype that kernels never read.
83    pub input_dtypes: Vec<DataType>,
84    /// Element types of the outputs.
85    pub output_dtypes: Vec<DataType>,
86}
87
88/// Cache key for a compiled kernel (§11.1). Keyed by the concrete node and its
89/// **resolved** (concrete) input shapes: attributes are fixed per node, so this
90/// is correct, and the shape component makes it *shape-keyed* — a re-run with
91/// the same resolved shapes hits, a different shape (e.g. a new batch/seq)
92/// misses and re-compiles. This preserves Chew's guarantee: a kernel is never
93/// reused for a shape it was not compiled for.
94#[derive(Clone, PartialEq, Eq, Hash, Debug)]
95struct KernelKey {
96    node: u32,
97    shapes: Vec<Vec<usize>>,
98}
99
100/// Observable kernel-cache statistics (§11.1) — enough to prove reuse in tests.
101#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
102pub struct CacheStats {
103    /// Distinct compiled entries currently held.
104    pub entries: usize,
105    /// Lookups served from an existing entry.
106    pub hits: u64,
107    /// Lookups that compiled a new kernel.
108    pub misses: u64,
109}
110
111/// Observable control-flow executor statistics. These counters make subgraph
112/// reuse deterministic to test without relying on timing.
113#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
114pub struct ControlFlowStats {
115    /// Child executors built, including shape-signature rebuilds.
116    pub subgraph_builds: u64,
117    /// Child subgraph invocations served by those executors.
118    pub subgraph_runs: u64,
119}
120
121/// Shape-keyed kernel cache (§11.1). Owns the compiled kernels for the session.
122#[derive(Default)]
123pub(crate) struct KernelCache {
124    entries: HashMap<KernelKey, Box<dyn onnx_runtime_ep_api::Kernel>>,
125    hits: u64,
126    misses: u64,
127}
128
129impl KernelCache {
130    fn stats(&self) -> CacheStats {
131        CacheStats {
132            entries: self.entries.len(),
133            hits: self.hits,
134            misses: self.misses,
135        }
136    }
137
138    /// Return the cached kernel for `(node, resolved_input_shapes)`, verifying
139    /// EP support and compiling+inserting it on a miss. The EP support check
140    /// lives on the miss path so a re-planned shape is re-validated exactly
141    /// once per distinct shape.
142    fn get_or_create(
143        &mut self,
144        node_id: NodeId,
145        node: &Node,
146        input_shapes: &[Vec<usize>],
147        opset: u64,
148        ep: &CpuExecutionProvider,
149    ) -> Result<&dyn onnx_runtime_ep_api::Kernel> {
150        let key = KernelKey {
151            node: node_id.0,
152            shapes: input_shapes.to_vec(),
153        };
154        if self.entries.contains_key(&key) {
155            self.hits += 1;
156        } else {
157            // Verify the EP claims this op at these concrete shapes/layouts
158            // before compiling — same gate the static path used at build.
159            let shape_dims: Vec<Shape> = input_shapes
160                .iter()
161                .map(|s| s.iter().map(|&d| Dim::Static(d)).collect())
162                .collect();
163            let layouts = vec![TensorLayout::contiguous(); input_shapes.len()];
164            if !matches!(
165                ep.supports_op(node, &shape_dims, &layouts),
166                KernelMatch::Supported { .. }
167            ) {
168                return Err(SessionError::unsupported_op(node, node_id, opset, ep.name()));
169            }
170            let kernel = ep.get_kernel(node, input_shapes, opset)?;
171            self.entries.insert(key.clone(), kernel);
172            self.misses += 1;
173        }
174        Ok(self.entries.get(&key).expect("just inserted").as_ref())
175    }
176}
177
178/// The compiled, runnable graph: buffers + plan + kernel cache. Owned by the
179/// public [`InferenceSession`](crate::InferenceSession).
180pub(crate) struct Executor {
181    graph: Graph,
182    /// Kept alive so external-weight memory maps outlive buffer population —
183    /// **and**, since the weight-streaming change, so borrowed initializer
184    /// buffers that alias this store's mmap bytes stay valid for the executor's
185    /// whole lifetime. `weights` MUST outlive every live use of `buffers`: a
186    /// borrowed `DeviceBuffer` in `buffers` points into `weights`' mmap/inline
187    /// storage. Teardown is safe because `Executor::drop` **drains and
188    /// deallocates `buffers` first** (a borrowed deallocate is a no-op free), so
189    /// no buffer still aliases `weights` when the `Arc<WeightStore>` field is
190    /// dropped afterwards — no use-after-free regardless of field drop order.
191    weights: Arc<WeightStore>,
192    ep: Arc<CpuExecutionProvider>,
193    /// One device buffer per backed value. Static values are allocated once at
194    /// build; dynamic (symbol-shaped) values are allocated per run and cached
195    /// here so a run whose resolved shape is unchanged reuses the allocation.
196    buffers: HashMap<ValueId, DeviceBuffer>,
197    /// The concrete shape each live buffer in [`Self::buffers`] is currently
198    /// sized for — the reuse key for run-scoped buffers.
199    buffer_shapes: HashMap<ValueId, Vec<usize>>,
200    /// Loader-produced (possibly symbolic) shape of every value.
201    value_shapes: HashMap<ValueId, Shape>,
202    /// Element type of every value.
203    value_dtypes: HashMap<ValueId, DataType>,
204    /// Topologically ordered execution plan (structure only; shapes per run).
205    plan: Vec<NodePlan>,
206    /// name → value id for the graph inputs the caller must supply.
207    input_index: HashMap<String, ValueId>,
208    /// Value ids the caller must supply at `run` (graph inputs minus initializers).
209    required_inputs: Vec<ValueId>,
210    /// Whether any value in the graph carries a symbolic dim. A fully-static
211    /// graph is materialized eagerly at build; a symbolic graph defers buffer
212    /// allocation and kernel compilation to the first `run` that fixes shapes.
213    has_symbols: bool,
214    cache: KernelCache,
215    /// name → value id for every named value in this graph (inputs, outputs,
216    /// initializers and interior SSA values). Used to resolve outer-scope
217    /// captures referenced by name from a nested control-flow subgraph body.
218    name_index: HashMap<String, ValueId>,
219    /// Compiled child executors for this graph's control-flow subgraph bodies,
220    /// keyed by `(control-flow node, subgraph attr key)`. Built lazily on first
221    /// execution (once concrete input shapes are known) and **reused across
222    /// Loop/Scan iterations** — the whole point of the efficiency directive: a
223    /// body's topo-sort, buffer sizing and kernel compilation happen once, then
224    /// every iteration is just a re-bind + dispatch. Rebuilt only if a later
225    /// invocation's external input shapes differ from the ones it was compiled
226    /// for (a shape-varying loop body — rare).
227    subgraph_execs: HashMap<(NodeId, String), CompiledSubgraph>,
228    control_flow_stats: ControlFlowStats,
229    /// Run-scoped zero-copy **view** metadata (§5.4). A value id present here is
230    /// a strided view aliasing another value's buffer (a layout/movement-op
231    /// output such as `Slice`) rather than an owner in [`Self::buffers`]. Built
232    /// during the run loop and cleared at the start of every run.
233    views: HashMap<ValueId, ValueView>,
234    /// Run-scoped set of buffer-owning value ids that have ≥1 live view alias.
235    /// A pinned buffer must not be reused or deallocated for the remainder of
236    /// the run (conservative liveness: a source buffer outlives every view that
237    /// aliases it, guaranteeing no use-after-free). Cleared each run.
238    pinned: HashSet<ValueId>,
239    /// Value ids whose runtime value is a **sequence of tensors** rather than a
240    /// single tensor (produced by `SequenceEmpty`/`SequenceConstruct`/
241    /// `SequenceInsert`/`SequenceErase`/`SplitToSequence`). Computed once at
242    /// build; these values own no [`DeviceBuffer`] and are skipped by buffer
243    /// sizing — their storage lives in [`Self::sequences`] at run time.
244    sequence_values: HashSet<ValueId>,
245    /// Run-scoped storage for sequence values: `value id → SequenceValue`. A
246    /// [`SequenceValue`] holds its elements as `Arc`-shared immutable tensors,
247    /// so a sequence op that inserts/erases/etc. shares element `Arc`s with the
248    /// source rather than deep-copying bytes (see [`crate::sequence`] for the
249    /// no-copy + no-race invariants). Cleared each run.
250    sequences: HashMap<ValueId, SequenceValue>,
251    /// Run-scoped **zero-copy** backing for a *tensor* value whose bytes are a
252    /// shared sequence element (the output of `SequenceAt`): the tensor aliases
253    /// the element's `Arc` instead of owning a `DeviceBuffer`, so no bytes are
254    /// copied out of the sequence. A downstream kernel reads it through a
255    /// [`TensorView`] over the `Arc`'s bytes; it is materialized to owned bytes
256    /// only at the graph-output/control-flow boundary. Cleared each run.
257    seq_elem_values: HashMap<ValueId, Arc<SeqTensor>>,
258}
259
260/// Run-scoped metadata for a zero-copy view value: it owns no buffer but
261/// borrows `source`'s buffer with the given (real, possibly non-contiguous or
262/// negative-strided) geometry. `strides`/`byte_offset` are expressed relative
263/// to `source`'s allocation base, so a view-of-a-view is flattened to a single
264/// hop whose `source` is always a real buffer owner (never itself a view).
265#[derive(Clone, Debug)]
266struct ValueView {
267    source: ValueId,
268    shape: Vec<usize>,
269    strides: Vec<i64>,
270    byte_offset: usize,
271}
272
273/// Per-input geometry the run loop resolves once per node: the raw base pointer
274/// of the backing (root) buffer plus the real view (shape, element strides —
275/// possibly non-contiguous or negative — and byte offset) to read it through.
276/// A plain owned value yields contiguous strides at offset 0; a view value
277/// yields its recorded strides/offset over its source buffer. `present` is false
278/// for an omitted optional input (an absent placeholder).
279struct InInfo {
280    present: bool,
281    dtype: DataType,
282    shape: Vec<usize>,
283    strides: Vec<i64>,
284    byte_offset: usize,
285    base_ptr: *const std::ffi::c_void,
286    /// Length in bytes of the backing (root) allocation, for the bounds gate.
287    root_len: usize,
288}
289
290/// A cached child executor for one control-flow subgraph body, plus the
291/// external-input shape signature it was compiled for (so a shape change forces
292/// a rebuild rather than a silent shape mismatch).
293struct CompiledSubgraph {
294    exec: Executor,
295    /// Ordered names of every external input: formal parameters first, then
296    /// captured outer-scope values.
297    input_names: Vec<String>,
298    /// Concrete shapes the child was last compiled for, in `input_names` order.
299    built_shapes: Vec<Vec<usize>>,
300}
301
302/// Invocation-invariant binding metadata for one selected subgraph. Loop/Scan
303/// prepare this once outside the iteration loop, including one-time capture
304/// materialization, then only rebind the changing formal tensors each step.
305struct PreparedSubgraph {
306    key: (NodeId, String),
307    formal_names: Vec<String>,
308    capture_names: Vec<String>,
309    /// Direct captures plus transitive captures needed only by nested bodies.
310    captures: HashMap<String, Tensor>,
311}
312
313/// The `[shape, strides, byte_offset]` storage-bounds gate (Holden's
314/// precondition). Uses [`view_in_bounds`] for fixed-width dtypes and a
315/// `storage_bytes` check for sub-byte packed dtypes (which have no integral
316/// per-element byte size).
317fn view_bounds(
318    shape: &[usize],
319    strides: &[i64],
320    byte_offset: usize,
321    dtype: DataType,
322    buffer_len: usize,
323) -> Result<()> {
324    let esize = dtype.byte_size();
325    if esize == 0 {
326        // Sub-byte (int4/uint4) or variable-width: size via `storage_bytes`.
327        let numel: usize = shape.iter().product();
328        let need = byte_offset + dtype.storage_bytes(numel);
329        if need > buffer_len {
330            return Err(SessionError::from(
331                onnx_runtime_ep_api::EpError::InvalidTensorView {
332                    reason: format!(
333                        "sub-byte view needs {need} bytes but backing allocation is {buffer_len}"
334                    ),
335                },
336            ));
337        }
338        return Ok(());
339    }
340    view_in_bounds(shape, strides, byte_offset, esize, buffer_len)?;
341    Ok(())
342}
343
344/// Gather a strided view over `src` into a fresh contiguous row-major byte
345/// buffer. `strides` are in **elements** (may be negative); `byte_offset` is the
346/// byte position of the element origin within `src`. `esize` is the element
347/// size in bytes (fixed-width types only — callers exclude sub-byte dtypes).
348/// This is the materialization copy that turns a zero-copy view back into a
349/// contiguous tensor for a strided-unaware consumer or the output boundary.
350fn gather_view(
351    src: &[u8],
352    shape: &[usize],
353    strides: &[i64],
354    byte_offset: usize,
355    esize: usize,
356) -> Vec<u8> {
357    let n: usize = shape.iter().product();
358    let mut out = vec![0u8; n * esize];
359    if n == 0 {
360        return out;
361    }
362    let rank = shape.len();
363    let mut idx = vec![0usize; rank];
364    let mut w = 0usize;
365    loop {
366        let mut off = byte_offset as i64;
367        for d in 0..rank {
368            off += strides[d] * idx[d] as i64 * esize as i64;
369        }
370        let s = off as usize;
371        out[w..w + esize].copy_from_slice(&src[s..s + esize]);
372        w += esize;
373        // Advance the row-major index; stop when it wraps to all-zero.
374        let mut carried = true;
375        for axis in (0..rank).rev() {
376            idx[axis] += 1;
377            if idx[axis] < shape[axis] {
378                carried = false;
379                break;
380            }
381            idx[axis] = 0;
382        }
383        if carried {
384            break;
385        }
386    }
387    out
388}
389
390/// Element count of a shape with overflow checking. A malicious or corrupt
391/// shape whose dims multiply past `usize::MAX` would silently wrap under a plain
392/// `iter().product()`, under-sizing the backing buffer. Returns
393/// [`SessionError::ShapeOverflow`] instead so the caller allocates nothing.
394fn checked_numel(dims: &[usize], value: impl FnOnce() -> String) -> Result<usize> {
395    let mut acc = 1usize;
396    for &d in dims {
397        acc = match acc.checked_mul(d) {
398            Some(n) => n,
399            None => {
400                return Err(SessionError::ShapeOverflow {
401                    value: value(),
402                    dims: dims.to_vec(),
403                })
404            }
405        };
406    }
407    Ok(acc)
408}
409
410/// Byte size of `numel` elements of `dtype` with overflow checking. Even when
411/// the element *count* fits in `usize` (guarded by [`checked_numel`]), the
412/// element-count → bytes multiply can still wrap for a fixed-width dtype and
413/// under-size the backing buffer. Returns [`SessionError::ShapeOverflow`] so the
414/// caller allocates nothing rather than a wrapped, undersized buffer.
415fn checked_storage_bytes(
416    dtype: DataType,
417    numel: usize,
418    value: impl FnOnce() -> String,
419    dims: &[usize],
420) -> Result<usize> {
421    dtype
422        .checked_storage_bytes(numel)
423        .ok_or_else(|| SessionError::ShapeOverflow {
424            value: value(),
425            dims: dims.to_vec(),
426        })
427}
428
429/// The effective operator-set version governing `node` — the graph's imported
430/// opset for the node's domain. The default ONNX domain is spelled both `""`
431/// and `"ai.onnx"`; both map to the same import. Loader and programmatic session
432/// entry points validate this invariant before executor construction.
433fn effective_opset(graph: &Graph, node: &Node) -> u64 {
434    let domain = node.domain.as_str();
435    graph
436        .opset_imports
437        .get(domain)
438        .or_else(|| {
439            if domain.is_empty() {
440                graph.opset_imports.get("ai.onnx")
441            } else if domain == "ai.onnx" {
442                graph.opset_imports.get("")
443            } else {
444                None
445            }
446        })
447        .copied()
448        .unwrap_or_else(|| {
449            unreachable!(
450                "internal invariant violated: node #{} ({}::{}) has no opset import",
451                node.id.0,
452                if node.domain.is_empty() {
453                    "ai.onnx"
454                } else {
455                    &node.domain
456                },
457                node.op_type
458            )
459        })
460}
461
462/// Substitute concrete symbol bindings into a (possibly symbolic) shape.
463/// Returns `None` if any dim is a symbol with no binding.
464fn substitute(shape: &Shape, bindings: &HashMap<SymbolId, usize>) -> Option<Vec<usize>> {
465    shape
466        .iter()
467        .map(|d| match d {
468            Dim::Static(n) => Some(*n),
469            Dim::Symbolic(s) => bindings.get(s).copied(),
470        })
471        .collect()
472}
473
474/// Decode a host buffer's integer elements as `i64` for `dtype`, or `None` if
475/// the dtype is not an integer the shape math understands. Used to read the
476/// *values* of shape-defining inputs (e.g. `Slice` starts/ends) at run time.
477fn buffer_as_i64(buffer: &DeviceBuffer, dtype: DataType) -> Option<Vec<i64>> {
478    bytes_as_i64(crate::tensor::host_bytes(buffer), dtype)
479}
480
481/// Decode raw little-endian integer bytes as `i64` for `dtype`, or `None` if the
482/// dtype is not an integer the shape math understands. Shared by the owned-buffer
483/// and materialized-view integer-input readers.
484fn bytes_as_i64(bytes: &[u8], dtype: DataType) -> Option<Vec<i64>> {
485    match dtype {
486        DataType::Int64 => Some(
487            bytes
488                .chunks_exact(8)
489                .map(|c| i64::from_le_bytes(c.try_into().unwrap()))
490                .collect(),
491        ),
492        DataType::Int32 => Some(
493            bytes
494                .chunks_exact(4)
495                .map(|c| i32::from_le_bytes(c.try_into().unwrap()) as i64)
496                .collect(),
497        ),
498        _ => None,
499    }
500}
501
502/// Compute the concrete output shapes of a *data-dependent* shape op from its
503/// already-resolved input shapes and the runtime *values* of its integer
504/// inputs. This is the executor's fallback for the rare value whose shape the
505/// loader's static (symbolic) inference could not pin down — e.g. a `Slice`
506/// whose `ends` is produced by a runtime `Shape → Min → Cast` chain, so its
507/// extent is only known once those upstream nodes have executed.
508///
509/// Model-agnostic: it dispatches on the op type alone. Returns `None` for ops
510/// this executor cannot resolve dynamically, which surfaces as
511/// [`SessionError::UnresolvedShape`] exactly as before.
512fn dynamic_output_shapes(
513    node: &Node,
514    input_shapes: &[Vec<usize>],
515    input_values: &[Option<Vec<i64>>],
516) -> Option<Vec<Vec<usize>>> {
517    match node.op_type.as_str() {
518        // Opset-10+ `Slice`: data, starts, ends, [axes], [steps] as inputs. The
519        // per-axis element count mirrors the `Slice` kernel's clamp semantics
520        // exactly (ONNX reference), so the buffer we size here matches what the
521        // kernel writes.
522        "Slice" => {
523            let data_shape = input_shapes.first()?;
524            let starts = input_values.get(1)?.as_ref()?;
525            let ends = input_values.get(2)?.as_ref()?;
526            let (axes, steps) = onnx_runtime_ep_cpu::slice_axes_steps(
527                starts.len(),
528                input_values.get(3).and_then(|v| v.as_deref()),
529                input_values.get(4).and_then(|v| v.as_deref()),
530            );
531            // Reuse the exact kernel geometry helper so the buffer we size here
532            // always matches what the Slice kernel writes. Any error (length
533            // mismatch, out-of-range axis, zero step) means "cannot resolve".
534            let plan =
535                onnx_runtime_ep_cpu::slice_plan(data_shape, starts, ends, &axes, &steps).ok()?;
536            let count: Vec<usize> = plan.iter().map(|p| p.count).collect();
537            Some(vec![count])
538        }
539        _ => None,
540    }
541}
542
543impl Executor {
544    /// Compile a graph + weights into a runnable executor on the CPU EP.
545    pub(crate) fn build(
546        graph: Graph,
547        weights: Arc<WeightStore>,
548        ep: Arc<CpuExecutionProvider>,
549    ) -> Result<Self> {
550        // Topological order up front: also validates the graph is a DAG.
551        let order = graph.topological_order()?;
552
553        let mut value_shapes: HashMap<ValueId, Shape> = HashMap::new();
554        let mut value_dtypes: HashMap<ValueId, DataType> = HashMap::new();
555        let mut buffers: HashMap<ValueId, DeviceBuffer> = HashMap::new();
556        let mut buffer_shapes: HashMap<ValueId, Vec<usize>> = HashMap::new();
557
558        // 1) Initializers: always concrete. Record dims and back each with a
559        //    device buffer. Where the mmap'd weight bytes are already suitably
560        //    aligned we **borrow** them zero-copy (no RAM allocation, no copy)
561        //    so a model whose weights exceed RAM still runs — the OS pages the
562        //    mmap in/out on demand. Unaligned or empty slices fall back to the
563        //    original allocate + copy path (correctness first).
564        let init_align = TensorLayout::contiguous().alignment;
565        for (&vid, weight) in &graph.initializers {
566            let dtype = weight.dtype();
567            let dims = weight.dims().to_vec();
568            let bytes = weights.bytes(weight).ok_or_else(|| {
569                SessionError::Internal(format!("weight bytes unavailable for value#{}", vid.0))
570            })?;
571            // Only borrow when the value has NO producer. The borrowed
572            // `DeviceBuffer` aliases read-only mmap/inline storage, so it must
573            // never be written. A legitimate initializer always has
574            // `producer == None`; a malformed graph can reuse an initializer's
575            // `ValueId` as a node output (see loader `validate_no_initializer_producer`),
576            // giving it a producer — a kernel would then write through
577            // `as_mut_ptr()` into read-only mmap (SIGSEGV / aliasing UB). In
578            // that case fall back to the owned writable copy below.
579            let producer_less = graph.value(vid).producer.is_none();
580            let buf = if producer_less
581                && !bytes.is_empty()
582                && (bytes.as_ptr() as usize).is_multiple_of(init_align)
583            {
584                // Zero-copy: alias the aligned mmap bytes. `weights` (an owned
585                // Arc field on this executor) outlives `buffers`, so the pointer
586                // stays valid for every run; the borrowed buffer is read-only
587                // (initializers are read-only SSA inputs, never a mutable output)
588                // and its Drop/deallocate never frees the mmap.
589                // SAFETY: `bytes` borrows `weights`' live mmap/inline storage,
590                // is `bytes.len()` long and `init_align`-aligned (checked above),
591                // is treated as read-only, and `weights` outlives every use.
592                unsafe {
593                    DeviceBuffer::from_borrowed_parts(
594                        bytes.as_ptr() as *mut std::ffi::c_void,
595                        ep.device_id(),
596                        bytes.len(),
597                        init_align,
598                    )
599                }
600            } else {
601                let mut owned = ep.allocate(bytes.len().max(1), init_align)?;
602                write_host(&mut owned, bytes)?;
603                owned
604            };
605            value_dtypes.insert(vid, dtype);
606            value_shapes.insert(vid, dims.iter().map(|&d| Dim::Static(d)).collect());
607            buffer_shapes.insert(vid, dims);
608            buffers.insert(vid, buf);
609        }
610
611        // 2) Record the loader shape + dtype of every remaining value (graph
612        //    inputs and node outputs). No allocation yet — shapes may be
613        //    symbolic and are only sized once resolved.
614        for &vid in &graph.inputs {
615            value_shapes
616                .entry(vid)
617                .or_insert_with(|| graph.value(vid).shape.clone());
618            value_dtypes.entry(vid).or_insert(graph.value(vid).dtype);
619        }
620        for &nid in &order {
621            for &out in &graph.node(nid).outputs {
622                value_shapes
623                    .entry(out)
624                    .or_insert_with(|| graph.value(out).shape.clone());
625                value_dtypes.entry(out).or_insert(graph.value(out).dtype);
626            }
627        }
628
629        let has_symbols = value_shapes.values().any(|s| as_static_shape(s).is_none());
630
631        // Sequence-typed values own no tensor buffer: a Sequence op stores its
632        // list in `sequences` at run time. Mark every value produced by a
633        // sequence-producing op so buffer sizing skips it (and so a Sequence
634        // graph output is diagnosed cleanly rather than read as tensor bytes).
635        let mut sequence_values: HashSet<ValueId> = HashSet::new();
636        for &nid in &order {
637            let node = graph.node(nid);
638            if produces_sequence_output(&node.op_type, &node.domain) {
639                for &out in &node.outputs {
640                    sequence_values.insert(out);
641                }
642            }
643        }
644
645        // 3) Build the structural per-node plan.
646        let mut plan = Vec::with_capacity(order.len());
647        for &nid in &order {
648            let node = graph.node(nid);
649            // EPContext nodes are pre-compiled: they bypass placement and were
650            // already restored through their owning EP by the session's
651            // consume path (§55.3). They must never be resolved as ordinary
652            // kernels — the CPU EP has no `EPContext` kernel — so skip them
653            // here.
654            if onnx_runtime_loader::is_ep_context_op(&node.op_type, &node.domain) {
655                continue;
656            }
657            // Preserve positional input arity: keep interior `None` (omitted
658            // optional) slots so a later present input is not misread as the
659            // omitted one, but trim trailing `None`s (a trailing omitted
660            // optional just lowers the arity, matching ONNX semantics).
661            let mut slots: Vec<Option<ValueId>> = node.inputs.clone();
662            while matches!(slots.last(), Some(None)) {
663                slots.pop();
664            }
665            let inputs = slots;
666            let outputs: Vec<ValueId> = node.outputs.clone();
667            let input_dtypes: Vec<DataType> = inputs
668                .iter()
669                .map(|v| v.map(|vid| value_dtypes[&vid]).unwrap_or(DataType::Float32))
670                .collect();
671            let output_dtypes: Vec<DataType> = outputs.iter().map(|v| value_dtypes[v]).collect();
672            plan.push(NodePlan {
673                node_id: nid,
674                inputs,
675                outputs,
676                input_dtypes,
677                output_dtypes,
678            });
679        }
680
681        // 4) name → value id and the set of caller-required inputs.
682        let mut input_index = HashMap::new();
683        let mut required_inputs = Vec::new();
684        for &vid in &graph.inputs {
685            if graph.initializers.contains_key(&vid) {
686                continue; // pre-filled; not a caller input
687            }
688            required_inputs.push(vid);
689            if let Some(name) = &graph.value(vid).name {
690                input_index.insert(name.clone(), vid);
691            }
692        }
693
694        // Full name → value id map (every named value in the graph), used to
695        // resolve a nested subgraph's outer-scope captures by name.
696        let mut name_index = HashMap::new();
697        for (vid, value) in graph.values.iter() {
698            if let Some(name) = &value.name {
699                name_index.insert(name.clone(), vid);
700            }
701        }
702
703        let mut exec = Self {
704            graph,
705            weights,
706            ep,
707            buffers,
708            buffer_shapes,
709            value_shapes,
710            value_dtypes,
711            plan,
712            input_index,
713            required_inputs,
714            has_symbols,
715            cache: KernelCache::default(),
716            name_index,
717            subgraph_execs: HashMap::new(),
718            control_flow_stats: ControlFlowStats::default(),
719            views: HashMap::new(),
720            pinned: HashSet::new(),
721            sequence_values,
722            sequences: HashMap::new(),
723            seq_elem_values: HashMap::new(),
724        };
725
726        // 5) Fully-static graphs are materialized eagerly (buffers + the whole
727        //    "compiled plan" of kernels), so the first `run` sees only cache
728        //    hits. Symbolic graphs cannot be sized until a `run` fixes their
729        //    shapes, so their buffers/kernels are created on first use.
730        if !exec.has_symbols {
731            let empty = HashMap::new();
732            let resolved = exec.resolve_all(&empty)?;
733            exec.size_buffers(&resolved)?;
734            exec.compile_all(&resolved)?;
735        }
736        Ok(exec)
737    }
738
739    /// Allocate `vid`'s buffer for `dims`, or reuse the existing allocation when
740    /// it is already sized for `dims` (the run-scoped reuse path).
741    fn ensure_buffer(&mut self, vid: ValueId, dtype: DataType, dims: &[usize]) -> Result<()> {
742        if self.buffer_shapes.get(&vid).map(|s| s.as_slice()) == Some(dims) {
743            return Ok(()); // identical shape → reuse allocation
744        }
745        if let Some(old) = self.buffers.remove(&vid) {
746            self.ep.deallocate(old)?;
747        }
748        let numel = checked_numel(dims, || format!("value#{}", vid.0))?;
749        let size = checked_storage_bytes(dtype, numel, || format!("value#{}", vid.0), dims)?;
750        let buf = self
751            .ep
752            .allocate(size.max(1), TensorLayout::contiguous().alignment)?;
753        self.buffers.insert(vid, buf);
754        self.buffer_shapes.insert(vid, dims.to_vec());
755        Ok(())
756    }
757
758    /// Resolve every value's concrete shape by substituting `bindings` into its
759    /// loader shape. A value whose shape stays symbolic (unbound) cannot be
760    /// sized: report it as an uninferred shape, naming its producing op.
761    fn resolve_all(
762        &self,
763        bindings: &HashMap<SymbolId, usize>,
764    ) -> Result<HashMap<ValueId, Vec<usize>>> {
765        let mut resolved = HashMap::with_capacity(self.value_shapes.len());
766        for (&vid, shape) in &self.value_shapes {
767            // Sequence-typed values have no meaningful tensor shape and are
768            // never buffer-sized; skip them so a static graph does not trip the
769            // unresolved-shape check on a sequence value.
770            if self.sequence_values.contains(&vid) {
771                continue;
772            }
773            match substitute(shape, bindings) {
774                Some(dims) => {
775                    resolved.insert(vid, dims);
776                }
777                None => {
778                    let value = self.graph.value(vid);
779                    let name = value
780                        .name
781                        .clone()
782                        .unwrap_or_else(|| format!("value#{}", vid.0));
783                    let op = value
784                        .producer
785                        .map(|nid| self.graph.node(nid).op_type.clone())
786                        .unwrap_or_else(|| "<graph input>".to_string());
787                    return Err(SessionError::UnresolvedShape { value: name, op });
788                }
789            }
790        }
791        Ok(resolved)
792    }
793
794    /// Like [`Self::resolve_all`] but never errors: values whose shape stays
795    /// symbolic (a data-dependent extent the loader could not pin down) are
796    /// simply omitted, to be resolved just-in-time during execution once their
797    /// producing node's inputs are concrete.
798    fn resolve_soft(&self, bindings: &HashMap<SymbolId, usize>) -> HashMap<ValueId, Vec<usize>> {
799        let mut resolved = HashMap::with_capacity(self.value_shapes.len());
800        for (&vid, shape) in &self.value_shapes {
801            if let Some(dims) = substitute(shape, bindings) {
802                resolved.insert(vid, dims);
803            }
804        }
805        resolved
806    }
807
808    /// Size (allocate or reuse) a backing buffer for every value from its
809    /// resolved concrete shape. Initializers already hold their weights and are
810    /// left untouched. Values whose shape is not (yet) in `resolved` — the
811    /// data-dependent ones filled in during execution — are skipped here and
812    /// sized just-in-time in the run loop.
813    fn size_buffers(&mut self, resolved: &HashMap<ValueId, Vec<usize>>) -> Result<()> {
814        let vids: Vec<ValueId> = self.value_shapes.keys().copied().collect();
815        for vid in vids {
816            if self.graph.initializers.contains_key(&vid) {
817                continue;
818            }
819            // Sequence-typed values own no tensor buffer (their list lives in
820            // `sequences` at run time), so never size one for them.
821            if self.sequence_values.contains(&vid) {
822                continue;
823            }
824            let dtype = self.value_dtypes[&vid];
825            let Some(dims) = resolved.get(&vid).cloned() else {
826                continue;
827            };
828            self.ensure_buffer(vid, dtype, &dims)?;
829        }
830        Ok(())
831    }
832
833    /// Resolved input shapes of a plan node, in positional order. An omitted
834    /// optional input (`None` slot) has no shape; it takes an empty shape,
835    /// which the run loop only ever pairs with an absent placeholder view.
836    fn node_input_shapes(
837        plan: &NodePlan,
838        resolved: &HashMap<ValueId, Vec<usize>>,
839    ) -> Vec<Vec<usize>> {
840        plan.inputs
841            .iter()
842            .map(|v| v.map(|vid| resolved[&vid].clone()).unwrap_or_default())
843            .collect()
844    }
845
846    /// Populate the kernel cache for the compiled plan against `resolved` shapes.
847    fn compile_all(&mut self, resolved: &HashMap<ValueId, Vec<usize>>) -> Result<()> {
848        for i in 0..self.plan.len() {
849            let node_id = self.plan[i].node_id;
850            let node = self.graph.node(node_id);
851            // Control-flow ops (If/Loop/Scan) are not leaf kernels — they execute
852            // nested subgraphs through the executor's own path, so they have no
853            // entry in the EP kernel registry and must not be compiled here.
854            if is_control_flow_op(&node.op_type, &node.domain) {
855                continue;
856            }
857            // Sequence ops are executor-handled (they operate on sequence-of-
858            // tensor values, not tensor views) — they have no EP kernel and must
859            // not be compiled here, exactly like control-flow ops.
860            if is_sequence_op(&node.op_type, &node.domain) {
861                continue;
862            }
863            let input_shapes = Self::node_input_shapes(&self.plan[i], resolved);
864            let node = self.graph.node(node_id);
865            let opset = effective_opset(&self.graph, node);
866            self.cache
867                .get_or_create(node_id, node, &input_shapes, opset, &self.ep)?;
868        }
869        Ok(())
870    }
871
872    pub(crate) fn cache_stats(&self) -> CacheStats {
873        self.cache.stats()
874    }
875
876    pub(crate) fn control_flow_stats(&self) -> ControlFlowStats {
877        self.control_flow_stats
878    }
879
880    /// The compiled graph, retained for the §55.4 EPContext dump path: the
881    /// exporter needs the (post-optimize) graph to serialise a `*_ctx.onnx`
882    /// context-cache model with compiled partitions spliced out.
883    pub(crate) fn graph(&self) -> &Graph {
884        &self.graph
885    }
886
887    /// Live weight bytes backing the graph, needed alongside [`Self::graph`] so
888    /// the EPContext dump can encode initializers into the context model.
889    pub(crate) fn weights(&self) -> &Arc<WeightStore> {
890        &self.weights
891    }
892
893    /// Warmup: re-touch the shape-keyed cache for the compiled plan so the first
894    /// real `run` sees only cache hits (§11.3). Only meaningful for fully-static
895    /// graphs, whose plan shapes are known at build; symbolic graphs cannot be
896    /// pre-compiled without a concrete shape and warm up on their first `run`.
897    pub(crate) fn warmup(&mut self) -> Result<()> {
898        if self.has_symbols {
899            return Ok(());
900        }
901        let empty = HashMap::new();
902        let resolved = self.resolve_all(&empty)?;
903        self.compile_all(&resolved)
904    }
905
906    /// Bind the graph's symbols to concrete sizes from the actual bound-input
907    /// shapes, validating rank and static dims and detecting symbol conflicts.
908    fn bind_symbols(
909        &self,
910        inputs: &[(&str, &Tensor)],
911    ) -> Result<HashMap<SymbolId, usize>> {
912        let mut bindings: HashMap<SymbolId, usize> = HashMap::new();
913        for (name, tensor) in inputs {
914            let vid = *self
915                .input_index
916                .get(*name)
917                .ok_or_else(|| SessionError::InputNotFound {
918                    name: (*name).to_string(),
919                })?;
920            let want_dtype = self.value_dtypes[&vid];
921            if tensor.dtype != want_dtype {
922                return Err(SessionError::DtypeMismatch {
923                    name: (*name).to_string(),
924                    expected: format!("{want_dtype:?}"),
925                    got: format!("{:?}", tensor.dtype),
926                });
927            }
928            let decl = &self.value_shapes[&vid];
929            if decl.len() != tensor.shape.len() {
930                return Err(SessionError::RankMismatch {
931                    name: (*name).to_string(),
932                    expected: decl.len(),
933                    got: tensor.shape.len(),
934                });
935            }
936            for (dim, &actual) in decl.iter().zip(&tensor.shape) {
937                match dim {
938                    Dim::Static(n) => {
939                        if *n != actual {
940                            return Err(SessionError::ShapeMismatch {
941                                name: (*name).to_string(),
942                                expected: as_static_shape(decl).unwrap_or_default(),
943                                got: tensor.shape.clone(),
944                            });
945                        }
946                    }
947                    Dim::Symbolic(s) => {
948                        if let Some(&prev) = bindings.get(s) {
949                            if prev != actual {
950                                let sym = self
951                                    .symbol_name(*s)
952                                    .unwrap_or_else(|| format!("symbol#{}", s.0));
953                                return Err(SessionError::SymbolConflict {
954                                    symbol: sym,
955                                    first: prev,
956                                    second: actual,
957                                });
958                            }
959                        } else {
960                            bindings.insert(*s, actual);
961                        }
962                    }
963                }
964            }
965        }
966        Ok(bindings)
967    }
968
969    /// Human-readable name of a symbol, if the graph recorded one.
970    fn symbol_name(&self, s: SymbolId) -> Option<String> {
971        self.graph
972            .symbol_constraints
973            .get(&s)
974            .and_then(|c| c.name.clone())
975    }
976
977    /// Sequential topological executor.
978    pub(crate) fn run(&mut self, inputs: &[(&str, &Tensor)]) -> Result<Vec<Tensor>> {
979        self.run_scoped(inputs, &HashMap::new())
980    }
981
982    /// Execute the graph with `inputs` bound by name, plus an `outer_scope` of
983    /// enclosing named values a nested control-flow subgraph body may capture.
984    /// The top-level session `run` passes an empty scope; a control-flow body's
985    /// child executor is invoked with its enclosing graph's live values so a
986    /// deeply-nested body can still reach an outer capture (ONNX lexical scope).
987    fn run_scoped(
988        &mut self,
989        inputs: &[(&str, &Tensor)],
990        outer_scope: &HashMap<String, Tensor>,
991    ) -> Result<Vec<Tensor>> {
992        // Zero-copy view metadata is run-scoped: a value that aliased another's
993        // buffer last run must not leak into this one (buffers may be resized).
994        self.views.clear();
995        self.pinned.clear();
996        // Sequence values and their zero-copy element-backed tensors are equally
997        // run-scoped (element Arcs from a prior run must not leak in).
998        self.sequences.clear();
999        self.seq_elem_values.clear();
1000
1001        // --- Resolve shapes from the actual bound inputs --------------------
1002        let bindings = self.bind_symbols(inputs)?;
1003
1004        // Every required input must be supplied.
1005        let provided: Vec<ValueId> = inputs
1006            .iter()
1007            .filter_map(|(name, _)| self.input_index.get(*name).copied())
1008            .collect();
1009        for &vid in &self.required_inputs {
1010            if !provided.contains(&vid) {
1011                let name = self
1012                    .graph
1013                    .value(vid)
1014                    .name
1015                    .clone()
1016                    .unwrap_or_else(|| format!("value#{}", vid.0));
1017                return Err(SessionError::InputNotFound { name });
1018            }
1019        }
1020
1021        // Substitute the bindings into every value → concrete shapes, then size
1022        // the run-scoped buffers from them (reused when unchanged). Values with a
1023        // data-dependent shape stay unresolved here and are filled in during the
1024        // execution loop, once their producing node's inputs are concrete.
1025        let mut resolved = self.resolve_soft(&bindings);
1026        self.size_buffers(&resolved)?;
1027
1028        // --- Bind input bytes into their (now correctly sized) buffers ------
1029        for (name, tensor) in inputs {
1030            let vid = self.input_index[*name];
1031            let buf = self
1032                .buffers
1033                .get_mut(&vid)
1034                .expect("input value has a buffer");
1035            write_host(buf, tensor.as_bytes())?;
1036        }
1037
1038        // --- Execute nodes ---------------------------------------------------
1039        // Iterate by index so a control-flow node can take `&mut self` (it must
1040        // build/reuse child executors) while an ordinary kernel node uses the
1041        // disjoint-field borrow split inside `exec_kernel_node`.
1042        for pi in 0..self.plan.len() {
1043            let node_id = self.plan[pi].node_id;
1044            let node = self.graph.node(node_id);
1045            if is_control_flow_op(&node.op_type, &node.domain) {
1046                self.exec_control_flow(pi, &mut resolved, outer_scope)?;
1047            } else if is_sequence_op(&node.op_type, &node.domain) {
1048                self.exec_sequence_node(pi, &mut resolved)?;
1049            } else {
1050                self.exec_kernel_node(pi, &mut resolved)?;
1051            }
1052        }
1053
1054        // --- Collect graph outputs into owned tensors -----------------------
1055        // A view output (a layout op whose result aliases an input buffer) is
1056        // materialized to contiguous owned bytes here — external consumers and
1057        // the Python/DLPack boundary expect contiguous tensors.
1058        let mut results = Vec::with_capacity(self.graph.outputs.len());
1059        for &vid in &self.graph.outputs {
1060            // A Sequence value cannot be returned through the tensor-typed
1061            // `run` boundary. Diagnose it clearly instead of misreading it as
1062            // tensor bytes; consumers extract tensors via SequenceAt /
1063            // ConcatFromSequence before the graph output.
1064            if self.sequence_values.contains(&vid) {
1065                let name = self
1066                    .graph
1067                    .try_value(vid)
1068                    .and_then(|v| v.name.clone())
1069                    .unwrap_or_else(|| format!("value#{}", vid.0));
1070                return Err(SessionError::SequenceOp {
1071                    op: "<graph output>".to_string(),
1072                    reason: format!(
1073                        "graph output {name} is a Sequence value, which cannot be \
1074                         returned through the tensor `run` API. To fix: end the graph \
1075                         with ConcatFromSequence or SequenceAt to produce tensor \
1076                         output(s)"
1077                    ),
1078                });
1079            }
1080            let dtype = self.value_dtypes[&vid];
1081            let shape = resolved[&vid].clone();
1082            let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
1083            results.push(Tensor::from_raw_in(self.ep.clone(), dtype, shape, &bytes)?);
1084        }
1085        Ok(results)
1086    }
1087
1088    /// Execute one ordinary (leaf-kernel) plan node: resolve any data-dependent
1089    /// output shapes, size buffers, build the input/output views (with Holden's
1090    /// bounds gate), resolve the shape-keyed kernel, and dispatch it.
1091    fn exec_kernel_node(
1092        &mut self,
1093        pi: usize,
1094        resolved: &mut HashMap<ValueId, Vec<usize>>,
1095    ) -> Result<()> {
1096        // Small owned copies of the plan facts so the buffer/view/cache fields
1097        // can be mutated below without fighting a borrow of `self.plan`.
1098        let node_id = self.plan[pi].node_id;
1099        let inputs = self.plan[pi].inputs.clone();
1100        let outputs = self.plan[pi].outputs.clone();
1101        let input_dtypes = self.plan[pi].input_dtypes.clone();
1102        let output_dtypes = self.plan[pi].output_dtypes.clone();
1103
1104        let input_shapes: Vec<Vec<usize>> = inputs
1105            .iter()
1106            .map(|v| v.map(|vid| resolved[&vid].clone()).unwrap_or_default())
1107            .collect();
1108
1109        // Data-dependent shapes: if any output's shape is still unresolved,
1110        // compute it now from the concrete input shapes + the runtime *values*
1111        // of this node's integer inputs. Buffers are NOT sized here — a view
1112        // output needs none, and the compute path sizes them just below.
1113        if outputs.iter().any(|v| !resolved.contains_key(v)) {
1114            let input_values: Vec<Option<Vec<i64>>> = inputs
1115                .iter()
1116                .enumerate()
1117                .map(|(i, v)| {
1118                    v.and_then(|vid| self.input_i64(vid, &input_shapes[i], input_dtypes[i]))
1119                })
1120                .collect();
1121            let node = self.graph.node(node_id);
1122            let out_shapes = dynamic_output_shapes(node, &input_shapes, &input_values)
1123                .ok_or_else(|| {
1124                    let vid = outputs
1125                        .iter()
1126                        .find(|v| !resolved.contains_key(v))
1127                        .copied()
1128                        .unwrap_or(outputs[0]);
1129                    let value = self.graph.value(vid);
1130                    SessionError::UnresolvedShape {
1131                        value: value
1132                            .name
1133                            .clone()
1134                            .unwrap_or_else(|| format!("value#{}", vid.0)),
1135                        op: node.op_type.clone(),
1136                    }
1137                })?;
1138            if out_shapes.len() != outputs.len() {
1139                return Err(SessionError::OutputShapeCountMismatch {
1140                    op: self.graph.node(node_id).op_type.clone(),
1141                    expected: outputs.len(),
1142                    got: out_shapes.len(),
1143                });
1144            }
1145            for (oi, &ovid) in outputs.iter().enumerate() {
1146                resolved.insert(ovid, out_shapes[oi].clone());
1147            }
1148        }
1149
1150        let output_shapes: Vec<Vec<usize>> =
1151            outputs.iter().map(|v| resolved[v].clone()).collect();
1152
1153        // Resolve each input's real geometry (root buffer + strides/offset) and
1154        // bounds-check it. View inputs read through their recorded strides.
1155        let mut in_infos: Vec<InInfo> = Vec::with_capacity(inputs.len());
1156        for (i, slot) in inputs.iter().enumerate() {
1157            let Some(vid) = *slot else {
1158                in_infos.push(InInfo {
1159                    present: false,
1160                    dtype: input_dtypes[i],
1161                    shape: Vec::new(),
1162                    strides: Vec::new(),
1163                    byte_offset: 0,
1164                    base_ptr: std::ptr::null(),
1165                    root_len: 0,
1166                });
1167                continue;
1168            };
1169            // A tensor input backed by a shared sequence element (SequenceAt
1170            // output) owns no DeviceBuffer: read it zero-copy through a
1171            // contiguous view over the element's immutable `Arc` bytes. The Arc
1172            // is held live in `self.seq_elem_values` for the whole run, so the
1173            // pointer stays valid across this kernel dispatch.
1174            if let Some(elem) = self.seq_elem_values.get(&vid) {
1175                let shape = input_shapes[i].clone();
1176                let strides = compute_contiguous_strides(&shape);
1177                let root_len = elem.data.len();
1178                let base_ptr = elem.as_ptr() as *const std::ffi::c_void;
1179                view_bounds(&shape, &strides, 0, input_dtypes[i], root_len)?;
1180                in_infos.push(InInfo {
1181                    present: true,
1182                    dtype: input_dtypes[i],
1183                    shape,
1184                    strides,
1185                    byte_offset: 0,
1186                    base_ptr,
1187                    root_len,
1188                });
1189                continue;
1190            }
1191            let root = self.root_of(vid);
1192            let buf = self.buffers.get(&root).ok_or_else(|| {
1193                SessionError::Internal(format!("missing buffer for input value#{}", vid.0))
1194            })?;
1195            let root_len = buf.len();
1196            let base_ptr = buf.as_ptr();
1197            let (shape, strides, byte_offset) = match self.views.get(&vid) {
1198                Some(view) => (view.shape.clone(), view.strides.clone(), view.byte_offset),
1199                None => {
1200                    let shape = input_shapes[i].clone();
1201                    let strides = compute_contiguous_strides(&shape);
1202                    (shape, strides, 0)
1203                }
1204            };
1205            view_bounds(&shape, &strides, byte_offset, input_dtypes[i], root_len)?;
1206            in_infos.push(InInfo {
1207                present: true,
1208                dtype: input_dtypes[i],
1209                shape,
1210                strides,
1211                byte_offset,
1212                base_ptr,
1213                root_len,
1214            });
1215        }
1216
1217        let ep = self.ep.clone();
1218
1219        // Bind the mutated fields as disjoint locals so `self` is never borrowed
1220        // whole while the kernel (from `cache`) and the buffers/views are held.
1221        let graph = &self.graph;
1222        let cache = &mut self.cache;
1223        let buffers = &mut self.buffers;
1224        let buffer_shapes = &mut self.buffer_shapes;
1225        let views_meta = &mut self.views;
1226        let pinned = &mut self.pinned;
1227
1228        // Build the (possibly strided) input views once; they feed both the
1229        // view-output probe and, on the compute path, the kernel itself.
1230        let mut views: Vec<TensorView> = Vec::with_capacity(in_infos.len());
1231        for info in &in_infos {
1232            if !info.present {
1233                views.push(TensorView::absent(info.dtype));
1234                continue;
1235            }
1236            views.push(
1237                TensorView::new(
1238                    DevicePtr(info.base_ptr),
1239                    info.dtype,
1240                    &info.shape,
1241                    &info.strides,
1242                    onnx_runtime_ir::DeviceId::cpu(),
1243                )
1244                .with_byte_offset(info.byte_offset),
1245            );
1246        }
1247
1248        let node = graph.node(node_id);
1249        let opset = effective_opset(graph, node);
1250        let kernel = cache.get_or_create(node_id, node, &input_shapes, opset, &ep)?;
1251
1252        // --- Zero-copy view fast path ---------------------------------------
1253        // Ask the kernel whether its outputs are strided views over its inputs
1254        // (a layout/movement op such as Slice). If so, record view metadata
1255        // aliasing the source buffer and skip compute + allocation entirely.
1256        if let Some(specs) = kernel.view_outputs(&views, outputs.len()) {
1257            drop(views);
1258            if specs.len() != outputs.len() {
1259                return Err(SessionError::Internal(format!(
1260                    "op '{}' returned {} view outputs for {} outputs",
1261                    node.op_type,
1262                    specs.len(),
1263                    outputs.len()
1264                )));
1265            }
1266            for (oi, spec) in specs.into_iter().enumerate() {
1267                let ovid = outputs[oi];
1268                let Some(in_vid) = inputs.get(spec.input_index).copied().flatten() else {
1269                    return Err(SessionError::Internal(format!(
1270                        "op '{}' view output {} references invalid input index {}",
1271                        node.op_type, oi, spec.input_index
1272                    )));
1273                };
1274                let root = match views_meta.get(&in_vid) {
1275                    Some(v) => v.source,
1276                    None => in_vid,
1277                };
1278                let root_len = buffers.get(&root).map(|b| b.len()).ok_or_else(|| {
1279                    SessionError::Internal(format!(
1280                        "view source value#{} has no buffer",
1281                        root.0
1282                    ))
1283                })?;
1284                // Bounds-gate the composed view against the source allocation.
1285                view_bounds(
1286                    &spec.shape,
1287                    &spec.strides,
1288                    spec.byte_offset,
1289                    output_dtypes[oi],
1290                    root_len,
1291                )?;
1292                // The output becomes a view: drop any buffer it used to own so a
1293                // later run re-sizes cleanly, then record the alias and pin the
1294                // source (conservative liveness — a source with any live view is
1295                // never reused/freed for the rest of the run; no use-after-free).
1296                // A freshly-produced output can never already be pinned (its
1297                // viewers run strictly after it under SSA topo order).
1298                debug_assert!(
1299                    !pinned.contains(&ovid),
1300                    "value#{} is pinned as a live view source yet is being reproduced",
1301                    ovid.0
1302                );
1303                if let Some(old) = buffers.remove(&ovid) {
1304                    ep.deallocate(old)?;
1305                }
1306                buffer_shapes.remove(&ovid);
1307                views_meta.insert(
1308                    ovid,
1309                    ValueView {
1310                        source: root,
1311                        shape: spec.shape.clone(),
1312                        strides: spec.strides,
1313                        byte_offset: spec.byte_offset,
1314                    },
1315                );
1316                pinned.insert(root);
1317                resolved.insert(ovid, spec.shape);
1318            }
1319            return Ok(());
1320        }
1321
1322        // --- Compute path ----------------------------------------------------
1323        // Size (allocate or reuse) each output's contiguous buffer, JIT-sizing
1324        // data-dependent ones. A value that was a view on a prior run has no
1325        // buffer here and is freshly allocated.
1326        for (oi, &ovid) in outputs.iter().enumerate() {
1327            let dims = &output_shapes[oi];
1328            let numel = checked_numel(dims, || format!("value#{}", ovid.0))?;
1329            let need = checked_storage_bytes(
1330                output_dtypes[oi],
1331                numel,
1332                || format!("value#{}", ovid.0),
1333                dims,
1334            )?
1335            .max(1);
1336            let fits = buffers.get(&ovid).map(|b| b.len() == need).unwrap_or(false);
1337            if !fits {
1338                // Never free a buffer that has a live view alias (would dangle
1339                // the viewer). Unreachable under SSA topo order, but enforced.
1340                debug_assert!(
1341                    !pinned.contains(&ovid),
1342                    "value#{} is pinned as a live view source yet is being resized",
1343                    ovid.0
1344                );
1345                if let Some(old) = buffers.remove(&ovid) {
1346                    ep.deallocate(old)?;
1347                }
1348                let buf = ep.allocate(need, TensorLayout::contiguous().alignment)?;
1349                buffers.insert(ovid, buf);
1350            }
1351        }
1352
1353        // Auto-materialization gate: a strided (view) input feeding a kernel
1354        // that does not accept strided input on that slot is gathered into a
1355        // private contiguous temp so contiguous-assuming kernels stay correct.
1356        // Temps must outlive the views that borrow them.
1357        let mut mat: Vec<Option<(Vec<u8>, Vec<i64>)>> = Vec::with_capacity(in_infos.len());
1358        for (i, info) in in_infos.iter().enumerate() {
1359            if !info.present {
1360                mat.push(None);
1361                continue;
1362            }
1363            let contiguous = onnx_runtime_ir::is_contiguous(&info.shape, &info.strides);
1364            if contiguous || kernel.supports_strided_input(i) {
1365                mat.push(None);
1366                continue;
1367            }
1368            let esize = info.dtype.byte_size();
1369            if esize == 0 {
1370                return Err(SessionError::from(
1371                    onnx_runtime_ep_api::EpError::InvalidTensorView {
1372                        reason: format!(
1373                            "cannot materialize sub-byte strided input {i} of op '{}'",
1374                            node.op_type
1375                        ),
1376                    },
1377                ));
1378            }
1379            let src = unsafe {
1380                std::slice::from_raw_parts(info.base_ptr as *const u8, info.root_len)
1381            };
1382            let gathered = gather_view(src, &info.shape, &info.strides, info.byte_offset, esize);
1383            let strides = compute_contiguous_strides(&info.shape);
1384            mat.push(Some((gathered, strides)));
1385        }
1386
1387        // Rebuild input views, swapping any materialized slot to its contiguous
1388        // temp (offset 0, contiguous strides over the fresh buffer).
1389        drop(views);
1390        let mut views: Vec<TensorView> = Vec::with_capacity(in_infos.len());
1391        for (i, info) in in_infos.iter().enumerate() {
1392            if !info.present {
1393                views.push(TensorView::absent(info.dtype));
1394                continue;
1395            }
1396            match &mat[i] {
1397                Some((buf, strides)) => views.push(TensorView::new(
1398                    DevicePtr(buf.as_ptr() as *const std::ffi::c_void),
1399                    info.dtype,
1400                    &info.shape,
1401                    strides,
1402                    onnx_runtime_ir::DeviceId::cpu(),
1403                )),
1404                None => views.push(
1405                    TensorView::new(
1406                        DevicePtr(info.base_ptr),
1407                        info.dtype,
1408                        &info.shape,
1409                        &info.strides,
1410                        onnx_runtime_ir::DeviceId::cpu(),
1411                    )
1412                    .with_byte_offset(info.byte_offset),
1413                ),
1414            }
1415        }
1416
1417        // Take output buffers out so they can be borrowed `&mut` disjointly from
1418        // the input reads (SSA guarantees outputs are disjoint from inputs).
1419        let out_strides: Vec<Vec<i64>> = output_shapes
1420            .iter()
1421            .map(|s| compute_contiguous_strides(s))
1422            .collect();
1423        let mut out_bufs: Vec<(ValueId, DeviceBuffer)> = Vec::with_capacity(outputs.len());
1424        for &vid in &outputs {
1425            let buf = buffers.remove(&vid).ok_or_else(|| {
1426                SessionError::Internal(format!("missing buffer for output value#{}", vid.0))
1427            })?;
1428            out_bufs.push((vid, buf));
1429        }
1430        let mut outs: Vec<TensorMut> = Vec::with_capacity(out_bufs.len());
1431        for (i, (_, buf)) in out_bufs.iter_mut().enumerate() {
1432            view_bounds(
1433                &output_shapes[i],
1434                &out_strides[i],
1435                0,
1436                output_dtypes[i],
1437                buf.len(),
1438            )?;
1439            let ptr = buf.as_mut_ptr();
1440            outs.push(TensorMut::new(
1441                DevicePtrMut(ptr),
1442                output_dtypes[i],
1443                &output_shapes[i],
1444                &out_strides[i],
1445                onnx_runtime_ir::DeviceId::cpu(),
1446            ));
1447        }
1448
1449        kernel.execute(&views, &mut outs)?;
1450
1451        drop(views);
1452        drop(outs);
1453        for (vid, buf) in out_bufs {
1454            buffers.insert(vid, buf);
1455        }
1456        Ok(())
1457    }
1458
1459    /// Read the integer *values* of input `vid` as `i64`, materializing a view
1460    /// first if needed. Used to resolve data-dependent output shapes (e.g. a
1461    /// `Slice` whose `ends` is produced at runtime). Returns `None` if the value
1462    /// has no readable buffer/view or its dtype is not an integer.
1463    fn input_i64(&self, vid: ValueId, shape: &[usize], dtype: DataType) -> Option<Vec<i64>> {
1464        if self.views.contains_key(&vid) || self.seq_elem_values.contains_key(&vid) {
1465            let bytes = self.contiguous_bytes(vid, shape, dtype).ok()?;
1466            bytes_as_i64(&bytes, dtype)
1467        } else {
1468            self.buffers.get(&vid).and_then(|b| buffer_as_i64(b, dtype))
1469        }
1470    }
1471}
1472
1473// === Sequence-of-tensors ops: SequenceEmpty / SequenceConstruct /
1474// SequenceInsert / SequenceErase / SequenceAt / SequenceLength /
1475// SplitToSequence / ConcatFromSequence ===
1476//
1477// These are handled at the executor level (like control-flow ops) rather than as
1478// leaf kernels, because they operate on a *sequence-of-tensors* runtime value
1479// that a `Kernel` — which sees only individual tensor views — cannot represent.
1480//
1481// ## No-copy design
1482//
1483// A sequence stores its elements as `Arc`-shared **immutable** [`SeqTensor`]s
1484// (see [`crate::sequence`]). Insert/Erase/Construct build a NEW list that SHARES
1485// the surviving element `Arc`s — only handles (a refcount bump), never element
1486// bytes, are cloned. `SequenceAt` yields the shared element `Arc` and backs its
1487// output tensor value with that same allocation (`seq_elem_values`), so a
1488// downstream kernel reads it through a zero-copy [`TensorView`] and no bytes are
1489// copied out of the sequence until the graph-output boundary. The only copies
1490// are unavoidable boundary crossings: a *tensor → sequence* entry (a produced
1491// `DeviceBuffer`, reused across runs, cannot be aliased so its bytes are moved
1492// into the element `Arc` exactly once) and the single-alloc `Split`/`Concat`
1493// data movement.
1494//
1495// ## No-race design
1496//
1497// Elements are immutable after construction and only ever shared read-only
1498// through `Arc`; there is no interior mutability, so concurrent readers cannot
1499// race (the only cross-thread state is `Arc`'s atomic refcount).
1500impl Executor {
1501    /// Execute one Sequence-op plan node.
1502    fn exec_sequence_node(
1503        &mut self,
1504        pi: usize,
1505        resolved: &mut HashMap<ValueId, Vec<usize>>,
1506    ) -> Result<()> {
1507        let node_id = self.plan[pi].node_id;
1508        let inputs = self.plan[pi].inputs.clone();
1509        let outputs = self.plan[pi].outputs.clone();
1510        let op = self.graph.node(node_id).op_type.clone();
1511
1512        match op.as_str() {
1513            "SequenceEmpty" => {
1514                let dtype_attr = self.graph.node(node_id).attr("dtype").and_then(|a| a.as_int());
1515                let dtype = match dtype_attr {
1516                    None => DataType::Float32, // ONNX default element type.
1517                    Some(raw) => DataType::from_onnx(raw as i32).ok_or_else(|| {
1518                        SessionError::SequenceOp {
1519                            op: op.clone(),
1520                            reason: format!(
1521                                "attribute 'dtype' = {raw} is not a known ONNX \
1522                                 TensorProto.DataType. To fix: use a valid element \
1523                                 dtype id (e.g. 1=float32, 7=int64)"
1524                            ),
1525                        }
1526                    })?,
1527                };
1528                self.sequences
1529                    .insert(outputs[0], SequenceValue::empty(dtype));
1530                Ok(())
1531            }
1532            "SequenceConstruct" => {
1533                let mut items = Vec::with_capacity(inputs.len());
1534                for slot in &inputs {
1535                    let vid = slot.ok_or_else(|| self.seq_missing_input(&op))?;
1536                    items.push(self.read_seq_element(vid, resolved)?);
1537                }
1538                let seq = SequenceValue::construct(items).map_err(seq_err)?;
1539                self.sequences.insert(outputs[0], seq);
1540                Ok(())
1541            }
1542            "SequenceInsert" => {
1543                let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1544                let tvid = inputs.get(1).copied().flatten().ok_or_else(|| {
1545                    self.seq_missing_input(&op)
1546                })?;
1547                let tensor = self.read_seq_element(tvid, resolved)?;
1548                let position = match inputs.get(2).copied().flatten() {
1549                    Some(pvid) => Some(self.read_scalar_i64(pvid, resolved, &op)?),
1550                    None => None,
1551                };
1552                let out = seq.insert(tensor, position).map_err(seq_err)?;
1553                self.sequences.insert(outputs[0], out);
1554                Ok(())
1555            }
1556            "SequenceErase" => {
1557                let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1558                let position = match inputs.get(1).copied().flatten() {
1559                    Some(pvid) => Some(self.read_scalar_i64(pvid, resolved, &op)?),
1560                    None => None,
1561                };
1562                let out = seq.erase(position).map_err(seq_err)?;
1563                self.sequences.insert(outputs[0], out);
1564                Ok(())
1565            }
1566            "SequenceAt" => {
1567                let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1568                let pvid = inputs.get(1).copied().flatten().ok_or_else(|| {
1569                    SessionError::SequenceOp {
1570                        op: op.clone(),
1571                        reason: "requires a 'position' input. To fix: supply the \
1572                                 index tensor of the element to read"
1573                            .to_string(),
1574                    }
1575                })?;
1576                let pos = self.read_scalar_i64(pvid, resolved, &op)?;
1577                let elem = seq.at(pos).map_err(seq_err)?;
1578                self.store_seq_element_output(outputs[0], elem, resolved)
1579            }
1580            "SequenceLength" => {
1581                let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1582                let len = seq.len() as i64;
1583                self.store_raw_tensor_output(
1584                    outputs[0],
1585                    DataType::Int64,
1586                    Vec::new(),
1587                    &len.to_le_bytes(),
1588                    resolved,
1589                )
1590            }
1591            "SplitToSequence" => self.exec_split_to_sequence(&op, &inputs, &outputs, resolved),
1592            "ConcatFromSequence" => {
1593                self.exec_concat_from_sequence(node_id, &op, &inputs, &outputs, resolved)
1594            }
1595            other => Err(SessionError::SequenceOp {
1596                op: other.to_string(),
1597                reason: "unrecognized Sequence op (executor routing bug)".to_string(),
1598            }),
1599        }
1600    }
1601
1602    /// `SplitToSequence`: split a tensor into a sequence along `axis`.
1603    fn exec_split_to_sequence(
1604        &mut self,
1605        op: &str,
1606        inputs: &[Option<ValueId>],
1607        outputs: &[ValueId],
1608        resolved: &mut HashMap<ValueId, Vec<usize>>,
1609    ) -> Result<()> {
1610        let node = self.graph.node(self.plan_node_of(outputs[0]));
1611        let axis_attr = node.attr("axis").and_then(|a| a.as_int()).unwrap_or(0);
1612        let keepdims = node.attr("keepdims").and_then(|a| a.as_int()).unwrap_or(1) != 0;
1613
1614        let ivid = inputs.first().copied().flatten().ok_or_else(|| self.seq_missing_input(op))?;
1615        let dtype = self.value_dtypes[&ivid];
1616        let esize = dtype.byte_size();
1617        if esize == 0 {
1618            return Err(SessionError::SequenceOp {
1619                op: op.to_string(),
1620                reason: format!(
1621                    "sub-byte dtype {dtype:?} is not supported for SplitToSequence. \
1622                     To fix: Cast to a byte-addressable dtype before splitting"
1623                ),
1624            });
1625        }
1626        let shape = resolved
1627            .get(&ivid)
1628            .cloned()
1629            .ok_or_else(|| self.seq_unresolved(op, ivid))?;
1630        let rank = shape.len();
1631        if rank == 0 {
1632            return Err(SessionError::SequenceOp {
1633                op: op.to_string(),
1634                reason: "cannot split a scalar (rank-0) tensor. To fix: split a \
1635                         tensor with at least one dimension"
1636                    .to_string(),
1637            });
1638        }
1639        let axis = normalize_axis(axis_attr, rank).ok_or_else(|| SessionError::SequenceOp {
1640            op: op.to_string(),
1641            reason: format!(
1642                "attribute 'axis' = {axis_attr} is out of range for a rank-{rank} \
1643                 input (valid range is [{}, {}])",
1644                -(rank as i64),
1645                rank as i64 - 1
1646            ),
1647        })?;
1648        let axis_dim = shape[axis];
1649        let bytes = self.contiguous_bytes(ivid, &shape, dtype)?;
1650
1651        // Determine per-chunk sizes and whether to squeeze the split axis.
1652        let mut squeeze = false;
1653        let sizes: Vec<usize> = match inputs.get(1).copied().flatten() {
1654            None => {
1655                // No 'split': one element per index along axis; keepdims=0 drops
1656                // the (size-1) axis from each element.
1657                squeeze = !keepdims;
1658                vec![1; axis_dim]
1659            }
1660            Some(svid) => {
1661                let sshape = resolved.get(&svid).cloned().unwrap_or_default();
1662                let svals = self.read_i64_vec(svid, &sshape, op)?;
1663                let is_scalar = sshape.is_empty();
1664                if is_scalar {
1665                    let chunk = *svals.first().ok_or_else(|| SessionError::SequenceOp {
1666                        op: op.to_string(),
1667                        reason: "'split' scalar is empty".to_string(),
1668                    })?;
1669                    if chunk <= 0 {
1670                        return Err(SessionError::SequenceOp {
1671                            op: op.to_string(),
1672                            reason: format!(
1673                                "'split' chunk size {chunk} must be positive"
1674                            ),
1675                        });
1676                    }
1677                    let chunk = chunk as usize;
1678                    let mut v = Vec::new();
1679                    let mut rem = axis_dim;
1680                    while rem > 0 {
1681                        let k = rem.min(chunk);
1682                        v.push(k);
1683                        rem -= k;
1684                    }
1685                    v
1686                } else {
1687                    let v: Vec<usize> = svals.iter().map(|&x| x.max(0) as usize).collect();
1688                    let sum: usize = v.iter().sum();
1689                    if sum != axis_dim {
1690                        return Err(SessionError::SequenceOp {
1691                            op: op.to_string(),
1692                            reason: format!(
1693                                "'split' sizes {v:?} sum to {sum} but axis {axis} has \
1694                                 extent {axis_dim}. To fix: make the split sizes sum to \
1695                                 the axis length"
1696                            ),
1697                        });
1698                    }
1699                    v
1700                }
1701            }
1702        };
1703
1704        let parts = split_axis(&bytes, &shape, axis, &sizes, esize);
1705        let items: Vec<std::sync::Arc<SeqTensor>> = parts
1706            .into_iter()
1707            .map(|(mut sh, data)| {
1708                if squeeze {
1709                    sh.remove(axis);
1710                }
1711                SeqTensor::shared(dtype, sh, data)
1712            })
1713            .collect();
1714        self.sequences.insert(
1715            outputs[0],
1716            SequenceValue {
1717                elem_dtype: dtype,
1718                items,
1719            },
1720        );
1721        Ok(())
1722    }
1723
1724    /// `ConcatFromSequence`: concatenate (or stack, when `new_axis=1`) a
1725    /// sequence's tensors into one freshly-allocated output.
1726    fn exec_concat_from_sequence(
1727        &mut self,
1728        node_id: NodeId,
1729        op: &str,
1730        inputs: &[Option<ValueId>],
1731        outputs: &[ValueId],
1732        resolved: &mut HashMap<ValueId, Vec<usize>>,
1733    ) -> Result<()> {
1734        let node = self.graph.node(node_id);
1735        let axis_attr = node.attr("axis").and_then(|a| a.as_int()).ok_or_else(|| {
1736            SessionError::SequenceOp {
1737                op: op.to_string(),
1738                reason: "requires the mandatory 'axis' attribute. To fix: set 'axis'"
1739                    .to_string(),
1740            }
1741        })?;
1742        let new_axis = node.attr("new_axis").and_then(|a| a.as_int()).unwrap_or(0) != 0;
1743
1744        let seq = self.get_sequence(inputs.first().copied().flatten(), op)?;
1745        if seq.is_empty() {
1746            return Err(SessionError::SequenceOp {
1747                op: op.to_string(),
1748                reason: "cannot concatenate an empty sequence (output shape is \
1749                         undefined). To fix: guard with SequenceLength"
1750                    .to_string(),
1751            });
1752        }
1753        let dtype = seq.elem_dtype;
1754        let esize = dtype.byte_size();
1755        if esize == 0 {
1756            return Err(SessionError::SequenceOp {
1757                op: op.to_string(),
1758                reason: format!(
1759                    "sub-byte dtype {dtype:?} is not supported for ConcatFromSequence"
1760                ),
1761            });
1762        }
1763        let elem_shapes: Vec<Vec<usize>> = seq.items.iter().map(|t| t.shape.clone()).collect();
1764        let elem_datas: Vec<&[u8]> = seq.items.iter().map(|t| t.data.as_slice()).collect();
1765        let rank = elem_shapes[0].len();
1766
1767        let (oshape, out) = if new_axis {
1768            // Stack: every element must share one shape; new axis in [-rank-1, rank].
1769            for (i, s) in elem_shapes.iter().enumerate() {
1770                if s != &elem_shapes[0] {
1771                    return Err(SessionError::SequenceOp {
1772                        op: op.to_string(),
1773                        reason: format!(
1774                            "ConcatFromSequence(new_axis=1) requires identical element \
1775                             shapes, but element {i} has shape {s:?} vs {:?}",
1776                            elem_shapes[0]
1777                        ),
1778                    });
1779                }
1780            }
1781            let axis = normalize_axis(axis_attr, rank + 1).ok_or_else(|| {
1782                SessionError::SequenceOp {
1783                    op: op.to_string(),
1784                    reason: format!(
1785                        "'axis' = {axis_attr} is out of range for new_axis=1 stacking \
1786                         of rank-{rank} elements (valid range is [{}, {}])",
1787                        -(rank as i64) - 1,
1788                        rank as i64
1789                    ),
1790                }
1791            })?;
1792            stack_new_axis(&elem_datas, &elem_shapes[0], axis, esize)
1793        } else {
1794            let axis = normalize_axis(axis_attr, rank).ok_or_else(|| SessionError::SequenceOp {
1795                op: op.to_string(),
1796                reason: format!(
1797                    "'axis' = {axis_attr} is out of range for rank-{rank} elements \
1798                     (valid range is [{}, {}])",
1799                    -(rank as i64),
1800                    rank as i64 - 1
1801                ),
1802            })?;
1803            // Elements must agree on every dimension except the concat axis.
1804            for (i, s) in elem_shapes.iter().enumerate() {
1805                let mismatch = s.len() != rank
1806                    || s.iter().enumerate().any(|(d, &v)| d != axis && v != elem_shapes[0][d]);
1807                if mismatch {
1808                    return Err(SessionError::SequenceOp {
1809                        op: op.to_string(),
1810                        reason: format!(
1811                            "ConcatFromSequence requires elements to match on all axes \
1812                             except {axis}, but element {i} has shape {s:?} vs {:?}",
1813                            elem_shapes[0]
1814                        ),
1815                    });
1816                }
1817            }
1818            concat_axis(&elem_datas, &elem_shapes, axis, esize)
1819        };
1820        drop(seq);
1821        self.store_raw_tensor_output(outputs[0], dtype, oshape, &out, resolved)
1822    }
1823
1824    /// Build (or share) an `Arc<SeqTensor>` for a tensor value entering a
1825    /// sequence. If the value is already a shared sequence element (a
1826    /// `SequenceAt` result), its `Arc` is **shared** with no copy; otherwise its
1827    /// contiguous bytes are moved into a fresh element once (the tensor→sequence
1828    /// entry boundary).
1829    fn read_seq_element(
1830        &self,
1831        vid: ValueId,
1832        resolved: &HashMap<ValueId, Vec<usize>>,
1833    ) -> Result<std::sync::Arc<SeqTensor>> {
1834        if let Some(elem) = self.seq_elem_values.get(&vid) {
1835            return Ok(std::sync::Arc::clone(elem)); // zero-copy share
1836        }
1837        let dtype = self.value_dtypes[&vid];
1838        let shape = resolved
1839            .get(&vid)
1840            .cloned()
1841            .ok_or_else(|| self.seq_unresolved("Sequence", vid))?;
1842        let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
1843        Ok(SeqTensor::shared(dtype, shape, bytes))
1844    }
1845
1846    /// Fetch (clone) the sequence value bound to `vid` (cheap — `Arc` handle
1847    /// clones), or an actionable error if the input is missing / not a sequence.
1848    fn get_sequence(&self, vid: Option<ValueId>, op: &str) -> Result<SequenceValue> {
1849        let vid = vid.ok_or_else(|| self.seq_missing_input(op))?;
1850        self.sequences.get(&vid).cloned().ok_or_else(|| SessionError::SequenceOp {
1851            op: op.to_string(),
1852            reason: format!(
1853                "input value#{} is not a live sequence. To fix: ensure it is produced \
1854                 by a Sequence-producing op (SequenceEmpty/Construct/Insert/Erase/\
1855                 SplitToSequence)",
1856                vid.0
1857            ),
1858        })
1859    }
1860
1861    /// Read a scalar `i64`/`i32` position input.
1862    fn read_scalar_i64(
1863        &self,
1864        vid: ValueId,
1865        resolved: &HashMap<ValueId, Vec<usize>>,
1866        op: &str,
1867    ) -> Result<i64> {
1868        let shape = resolved.get(&vid).cloned().unwrap_or_default();
1869        let dtype = self.value_dtypes[&vid];
1870        let vals = self.input_i64(vid, &shape, dtype).ok_or_else(|| SessionError::SequenceOp {
1871            op: op.to_string(),
1872            reason: format!(
1873                "position input has dtype {dtype:?}, expected an integer (int32/int64). \
1874                 To fix: provide an int64 scalar index"
1875            ),
1876        })?;
1877        vals.first().copied().ok_or_else(|| SessionError::SequenceOp {
1878            op: op.to_string(),
1879            reason: "position input is empty; expected a single scalar index".to_string(),
1880        })
1881    }
1882
1883    /// Read an `i64` vector from an integer tensor input (SplitToSequence's
1884    /// `split`).
1885    fn read_i64_vec(&self, vid: ValueId, shape: &[usize], op: &str) -> Result<Vec<i64>> {
1886        let dtype = self.value_dtypes[&vid];
1887        self.input_i64(vid, shape, dtype).ok_or_else(|| SessionError::SequenceOp {
1888            op: op.to_string(),
1889            reason: format!(
1890                "'split' input has dtype {dtype:?}, expected int32/int64. To fix: \
1891                 provide integer split sizes"
1892            ),
1893        })
1894    }
1895
1896    /// Back a tensor *output* value with a shared sequence element (SequenceAt):
1897    /// the value owns no buffer — a downstream consumer reads the element's
1898    /// `Arc` bytes zero-copy. Any stale buffer/view for the value is released.
1899    fn store_seq_element_output(
1900        &mut self,
1901        vid: ValueId,
1902        elem: std::sync::Arc<SeqTensor>,
1903        resolved: &mut HashMap<ValueId, Vec<usize>>,
1904    ) -> Result<()> {
1905        if let Some(old) = self.buffers.remove(&vid) {
1906            self.ep.deallocate(old)?;
1907        }
1908        self.buffer_shapes.remove(&vid);
1909        self.views.remove(&vid);
1910        resolved.insert(vid, elem.shape.clone());
1911        self.value_dtypes.insert(vid, elem.dtype);
1912        self.seq_elem_values.insert(vid, elem);
1913        Ok(())
1914    }
1915
1916    /// Store freshly-computed contiguous bytes into a tensor output value
1917    /// (SequenceLength / ConcatFromSequence): (re)allocate its buffer, copy the
1918    /// bytes once, and record its dtype/shape.
1919    fn store_raw_tensor_output(
1920        &mut self,
1921        vid: ValueId,
1922        dtype: DataType,
1923        dims: Vec<usize>,
1924        bytes: &[u8],
1925        resolved: &mut HashMap<ValueId, Vec<usize>>,
1926    ) -> Result<()> {
1927        // A value that was seq-element-backed on a prior run must not shadow the
1928        // fresh buffer we write here.
1929        self.seq_elem_values.remove(&vid);
1930        self.views.remove(&vid);
1931        let need = bytes.len().max(1);
1932        let fits = self.buffers.get(&vid).map(|b| b.len() == need).unwrap_or(false);
1933        if !fits {
1934            if let Some(old) = self.buffers.remove(&vid) {
1935                self.ep.deallocate(old)?;
1936            }
1937            let buf = self.ep.allocate(need, TensorLayout::contiguous().alignment)?;
1938            self.buffers.insert(vid, buf);
1939        }
1940        let buf = self.buffers.get_mut(&vid).expect("just ensured");
1941        write_host(buf, bytes)?;
1942        self.value_dtypes.insert(vid, dtype);
1943        self.buffer_shapes.insert(vid, dims.clone());
1944        resolved.insert(vid, dims);
1945        Ok(())
1946    }
1947
1948    /// The producing node id of value `vid` (Sequence ops always have a
1949    /// producer).
1950    fn plan_node_of(&self, vid: ValueId) -> NodeId {
1951        self.graph
1952            .value(vid)
1953            .producer
1954            .expect("sequence op output has a producer")
1955    }
1956
1957    fn seq_missing_input(&self, op: &str) -> SessionError {
1958        SessionError::SequenceOp {
1959            op: op.to_string(),
1960            reason: "a required input is missing (omitted None slot). To fix: connect \
1961                     all required inputs of this Sequence op"
1962                .to_string(),
1963        }
1964    }
1965
1966    fn seq_unresolved(&self, op: &str, vid: ValueId) -> SessionError {
1967        let name = self
1968            .graph
1969            .try_value(vid)
1970            .and_then(|v| v.name.clone())
1971            .unwrap_or_else(|| format!("value#{}", vid.0));
1972        SessionError::SequenceOp {
1973            op: op.to_string(),
1974            reason: format!(
1975                "input {name} has no resolved shape yet. To fix: ensure its producer \
1976                 runs before this Sequence op"
1977            ),
1978        }
1979    }
1980}
1981
1982/// Map a [`crate::sequence::SeqOpError`] into an actionable `SessionError`.
1983fn seq_err(e: crate::sequence::SeqOpError) -> SessionError {
1984    SessionError::SequenceOp {
1985        op: e.op.to_string(),
1986        reason: e.reason,
1987    }
1988}
1989
1990/// Normalize a possibly-negative ONNX `axis` against `rank`, returning the
1991/// non-negative axis or `None` when out of `[-rank, rank-1]`.
1992fn normalize_axis(axis: i64, rank: usize) -> Option<usize> {
1993    let r = rank as i64;
1994    let a = if axis < 0 { axis + r } else { axis };
1995    if a < 0 || a >= r {
1996        None
1997    } else {
1998        Some(a as usize)
1999    }
2000}
2001
2002
2003/// Whether `(op_type, domain)` is one of the standard subgraph-bearing
2004/// control-flow ops the executor handles recursively (default `ai.onnx`
2005/// domain). Kept in lock-step with the loader's `validate_no_control_flow`
2006/// allow-list.
2007fn is_control_flow_op(op_type: &str, domain: &str) -> bool {
2008    (domain.is_empty() || domain == "ai.onnx") && matches!(op_type, "If" | "Loop" | "Scan")
2009}
2010
2011/// Whether `(op_type, domain)` is an ONNX **Sequence** op the executor handles
2012/// directly (default `ai.onnx` domain). Like control-flow ops these are handled
2013/// at the executor level rather than as leaf [`Kernel`](onnx_runtime_ep_api::Kernel)s
2014/// because a `Kernel` sees only tensor views, never a *sequence-of-tensors*
2015/// runtime value. Kept as a small self-contained routing predicate (mirroring
2016/// [`is_control_flow_op`]) so it never collides with the EP kernel registry.
2017fn is_sequence_op(op_type: &str, domain: &str) -> bool {
2018    (domain.is_empty() || domain == "ai.onnx")
2019        && matches!(
2020            op_type,
2021            "SequenceEmpty"
2022                | "SequenceConstruct"
2023                | "SequenceInsert"
2024                | "SequenceErase"
2025                | "SequenceAt"
2026                | "SequenceLength"
2027                | "SplitToSequence"
2028                | "ConcatFromSequence"
2029        )
2030}
2031
2032/// Whether a Sequence op yields a *sequence* value (vs. a tensor). Used at build
2033/// to mark sequence-typed values so they are excluded from tensor buffer sizing.
2034fn produces_sequence_output(op_type: &str, domain: &str) -> bool {
2035    (domain.is_empty() || domain == "ai.onnx")
2036        && matches!(
2037            op_type,
2038            "SequenceEmpty"
2039                | "SequenceConstruct"
2040                | "SequenceInsert"
2041                | "SequenceErase"
2042                | "SplitToSequence"
2043        )
2044}
2045
2046/// Read a single scalar `i64`/`i32` element from a length-1 tensor (Loop's `M`).
2047fn tensor_scalar_i64(t: &Tensor) -> Option<i64> {
2048    match t.dtype {
2049        DataType::Int64 => t
2050            .as_bytes()
2051            .get(..8)
2052            .map(|c| i64::from_le_bytes(c.try_into().unwrap())),
2053        DataType::Int32 => t
2054            .as_bytes()
2055            .get(..4)
2056            .map(|c| i32::from_le_bytes(c.try_into().unwrap()) as i64),
2057        _ => None,
2058    }
2059}
2060
2061/// Read a single scalar bool from a length-1 `BOOL` tensor (a `BOOL` is one
2062/// byte; any nonzero is true, per ONNX).
2063fn tensor_scalar_bool(t: &Tensor) -> Option<bool> {
2064    if t.dtype != DataType::Bool {
2065        return None;
2066    }
2067    t.as_bytes().first().map(|&b| b != 0)
2068}
2069
2070/// Build a length-1 `i64` scalar tensor (Loop's `iter_num` body input).
2071fn scalar_i64_tensor(v: i64) -> Result<Tensor> {
2072    Tensor::from_raw(DataType::Int64, vec![], &v.to_le_bytes())
2073}
2074
2075/// Build a scalar `BOOL` tensor (Loop's `cond` body input).
2076fn scalar_bool_tensor(v: bool) -> Result<Tensor> {
2077    Tensor::from_raw(DataType::Bool, vec![], &[u8::from(v)])
2078}
2079
2080fn missing_capture_error(attr_key: &str, name: &str) -> SessionError {
2081    SessionError::Internal(format!(
2082        "control-flow body '{attr_key}' captures free variable '{name}', but it is not \
2083         available in the enclosing scope. RULES #1: a subgraph may only reference outer \
2084         values that are graph inputs, initializers, or produced by an upstream node in an \
2085         enclosing graph; '{name}' matches none of these"
2086    ))
2087}
2088
2089/// Names a graph or any nested body needs from its enclosing lexical scope.
2090/// A nested requirement stops propagating when this graph defines that name,
2091/// because the nested body will bind the local value at execution time.
2092fn required_outer_names(graph: &Graph) -> HashSet<String> {
2093    let formal_set: HashSet<ValueId> = graph.inputs.iter().copied().collect();
2094    let local_names: HashSet<&str> = graph
2095        .values
2096        .iter()
2097        .filter_map(|(_, value)| value.name.as_deref())
2098        .collect();
2099    let mut required = HashSet::new();
2100    for (vid, value) in graph.values.iter() {
2101        if value.producer.is_none()
2102            && !formal_set.contains(&vid)
2103            && !graph.initializers.contains_key(&vid)
2104            && let Some(name) = &value.name
2105        {
2106            required.insert(name.clone());
2107        }
2108    }
2109    for nested in graph.subgraphs.values() {
2110        for name in required_outer_names(nested) {
2111            if !local_names.contains(name.as_str()) {
2112                required.insert(name);
2113            }
2114        }
2115    }
2116    required
2117}
2118
2119// === Control-flow (subgraph-executing) ops: If / Loop / Scan ===
2120//
2121// These are handled at the executor level rather than as leaf kernels because
2122// they must recursively execute a nested ONNX [`Graph`] with the enclosing
2123// scope bound — something a `Kernel` (which sees only tensor views, never the
2124// session/graph context) cannot do. Each body is compiled to a child
2125// [`Executor`] once and **reused across iterations** (see [`CompiledSubgraph`]).
2126impl Executor {
2127    /// Materialize value `vid`'s current bytes into an owned host [`Tensor`],
2128    /// using its resolved concrete shape and recorded dtype.
2129    fn value_tensor(
2130        &self,
2131        vid: ValueId,
2132        resolved: &HashMap<ValueId, Vec<usize>>,
2133    ) -> Result<Tensor> {
2134        let dtype = self.value_dtypes[&vid];
2135        let shape = resolved.get(&vid).cloned().ok_or_else(|| {
2136            let name = self
2137                .graph
2138                .try_value(vid)
2139                .and_then(|v| v.name.clone())
2140                .unwrap_or_else(|| format!("value#{}", vid.0));
2141            SessionError::UnresolvedShape {
2142                value: name,
2143                op: "<control-flow input>".to_string(),
2144            }
2145        })?;
2146        // A view value owns no buffer; materialize its strided bytes contiguous.
2147        let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
2148        Tensor::from_raw_in(self.ep.clone(), dtype, shape, &bytes)
2149    }
2150
2151    /// The buffer-owning (root) value backing `vid`: `vid` itself if it owns a
2152    /// buffer, or the `source` recorded in its view metadata (always a root,
2153    /// since views are flattened at creation).
2154    fn root_of(&self, vid: ValueId) -> ValueId {
2155        match self.views.get(&vid) {
2156            Some(v) => v.source,
2157            None => vid,
2158        }
2159    }
2160
2161    /// Contiguous row-major bytes of `vid` for `shape`/`dtype`, materializing a
2162    /// view (strided gather over its source buffer) or truncating an owned
2163    /// buffer to its logical size. This is the single materialization seam used
2164    /// by the graph-output boundary and control-flow scope capture.
2165    fn contiguous_bytes(
2166        &self,
2167        vid: ValueId,
2168        shape: &[usize],
2169        dtype: DataType,
2170    ) -> Result<Vec<u8>> {
2171        let numel: usize = shape.iter().product();
2172        let n = dtype.storage_bytes(numel);
2173        // A tensor value backed by a shared sequence element (SequenceAt output)
2174        // owns no buffer; its bytes are the element's contiguous bytes. This is
2175        // the one materialization point where they are copied out (the boundary
2176        // back into owned tensors); the compute path reads them zero-copy.
2177        if let Some(elem) = self.seq_elem_values.get(&vid) {
2178            return Ok(elem.data[..n.min(elem.data.len())].to_vec());
2179        }
2180        if let Some(view) = self.views.get(&vid) {
2181            let buf = self.buffers.get(&view.source).ok_or_else(|| {
2182                SessionError::Internal(format!(
2183                    "view value#{} aliases missing source buffer value#{}",
2184                    vid.0, view.source.0
2185                ))
2186            })?;
2187            let esize = dtype.byte_size();
2188            if esize == 0 {
2189                // Sub-byte views are never created (Slice falls back to copy),
2190                // so reaching here is an internal invariant violation.
2191                return Err(SessionError::Internal(format!(
2192                    "cannot materialize sub-byte view value#{}",
2193                    vid.0
2194                )));
2195            }
2196            Ok(gather_view(
2197                host_bytes(buf),
2198                &view.shape,
2199                &view.strides,
2200                view.byte_offset,
2201                esize,
2202            ))
2203        } else {
2204            let buf = self.buffers.get(&vid).ok_or_else(|| {
2205                SessionError::Internal(format!("value#{} not produced", vid.0))
2206            })?;
2207            Ok(host_bytes(buf)[..n].to_vec())
2208        }
2209    }
2210
2211    /// Store a control-flow op's produced output `tensor` into this graph's
2212    /// output value `vid`: (re)size the backing buffer, copy the bytes, and
2213    /// record the runtime dtype/shape so the caller (and the final output
2214    /// collection) reads them back correctly. Control-flow output shapes are
2215    /// data-dependent (the loader never inferred inside the body), so they are
2216    /// resolved here, exactly as the JIT data-dependent path does for kernels.
2217    fn store_output_tensor(
2218        &mut self,
2219        vid: ValueId,
2220        tensor: &Tensor,
2221        resolved: &mut HashMap<ValueId, Vec<usize>>,
2222    ) -> Result<()> {
2223        self.store_output_bytes(
2224            vid,
2225            tensor.dtype,
2226            tensor.shape.clone(),
2227            tensor.as_bytes(),
2228            resolved,
2229        )
2230    }
2231
2232    fn store_output_bytes(
2233        &mut self,
2234        vid: ValueId,
2235        dtype: DataType,
2236        dims: Vec<usize>,
2237        bytes: &[u8],
2238        resolved: &mut HashMap<ValueId, Vec<usize>>,
2239    ) -> Result<()> {
2240        let numel = checked_numel(&dims, || format!("value#{}", vid.0))?;
2241        let need = checked_storage_bytes(dtype, numel, || format!("value#{}", vid.0), &dims)?
2242            .max(1);
2243        let fits = self.buffers.get(&vid).map(|b| b.len() == need).unwrap_or(false);
2244        if !fits {
2245            if let Some(old) = self.buffers.remove(&vid) {
2246                self.ep.deallocate(old)?;
2247            }
2248            let buf = self
2249                .ep
2250                .allocate(need, TensorLayout::contiguous().alignment)?;
2251            self.buffers.insert(vid, buf);
2252        }
2253        let buf = self.buffers.get_mut(&vid).expect("just ensured");
2254        write_host(buf, bytes)?;
2255        self.value_dtypes.insert(vid, dtype);
2256        self.buffer_shapes.insert(vid, dims.clone());
2257        resolved.insert(vid, dims);
2258        Ok(())
2259    }
2260
2261    /// Prepare one selected control-flow subgraph and materialize only the free
2262    /// variables that body actually captures. This avoids copying every named
2263    /// value in the enclosing graph and, for Loop/Scan, keeps captures stable
2264    /// across all iterations.
2265    fn prepare_subgraph(
2266        &self,
2267        node_id: NodeId,
2268        attr_key: &str,
2269        resolved: &HashMap<ValueId, Vec<usize>>,
2270        outer_scope: &HashMap<String, Tensor>,
2271    ) -> Result<PreparedSubgraph> {
2272        let key = (node_id, attr_key.to_string());
2273        let body = self.graph.subgraphs.get(&key).ok_or_else(|| {
2274            SessionError::Internal(format!(
2275                "control-flow node #{} references missing subgraph '{attr_key}'",
2276                node_id.0
2277            ))
2278        })?;
2279
2280        let formal_names: Vec<String> = body
2281            .inputs
2282            .iter()
2283            .map(|&vid| {
2284                body.value(vid)
2285                    .name
2286                    .clone()
2287                    .unwrap_or_else(|| format!("value#{}", vid.0))
2288            })
2289            .collect();
2290        let formal_set: HashSet<ValueId> = body.inputs.iter().copied().collect();
2291        let mut capture_names = Vec::new();
2292        for (vid, value) in body.values.iter() {
2293            if value.producer.is_none()
2294                && !formal_set.contains(&vid)
2295                && !body.initializers.contains_key(&vid)
2296                && let Some(name) = &value.name
2297            {
2298                capture_names.push(name.clone());
2299            }
2300        }
2301        capture_names.sort();
2302
2303        let mut scope_names = required_outer_names(body);
2304        scope_names.extend(capture_names.iter().cloned());
2305        let mut captures = HashMap::with_capacity(scope_names.len());
2306        for name in scope_names {
2307            let tensor = if let Some(&vid) = self.name_index.get(&name) {
2308                let materialized = self.buffers.contains_key(&vid)
2309                    || self.views.contains_key(&vid)
2310                    || self.seq_elem_values.contains_key(&vid);
2311                if resolved.contains_key(&vid) && materialized {
2312                    self.value_tensor(vid, resolved)?
2313                } else {
2314                    outer_scope.get(&name).cloned().ok_or_else(|| {
2315                        missing_capture_error(attr_key, &name)
2316                    })?
2317                }
2318            } else {
2319                outer_scope.get(&name).cloned().ok_or_else(|| {
2320                    missing_capture_error(attr_key, &name)
2321                })?
2322            };
2323            captures.insert(name, tensor);
2324        }
2325
2326        Ok(PreparedSubgraph {
2327            key,
2328            formal_names,
2329            capture_names,
2330            captures,
2331        })
2332    }
2333
2334    /// Compile a control-flow body subgraph to a child [`Executor`], turning
2335    /// captured outer-scope names into extra graph inputs (so they are supplied
2336    /// and written every run), seeding the concrete external-input shapes, and
2337    /// running shape inference so the body's interior buffers can be sized.
2338    fn build_subgraph_exec(
2339        &self,
2340        prepared: &PreparedSubgraph,
2341        externals: &[&Tensor],
2342    ) -> Result<CompiledSubgraph> {
2343        let key = &prepared.key;
2344        let body = self.graph.subgraphs.get(key).ok_or_else(|| {
2345            SessionError::Internal(format!(
2346                "control-flow node #{} has no registered subgraph '{}'",
2347                key.0 .0, key.1
2348            ))
2349        })?;
2350        let mut g = body.clone();
2351
2352        // name → value id within the body.
2353        let mut body_names: HashMap<String, ValueId> = HashMap::new();
2354        for (vid, value) in g.values.iter() {
2355            if let Some(n) = &value.name {
2356                body_names.insert(n.clone(), vid);
2357            }
2358        }
2359
2360        // Captures become graph inputs so they are required + written each run.
2361        for cname in &prepared.capture_names {
2362            let vid = *body_names.get(cname).ok_or_else(|| {
2363                SessionError::Internal(format!(
2364                    "control-flow body '{}' lost capture value '{cname}'",
2365                    key.1
2366                ))
2367            })?;
2368            if !g.inputs.contains(&vid) {
2369                g.add_input(vid);
2370            }
2371        }
2372
2373        // Seed each external input's concrete static shape + runtime dtype so
2374        // shape inference can flow through the body.
2375        let all_names = prepared
2376            .formal_names
2377            .iter()
2378            .chain(prepared.capture_names.iter());
2379        for (name, tensor) in all_names.zip(externals.iter()) {
2380            let vid = *body_names.get(name).ok_or_else(|| {
2381                SessionError::Internal(format!(
2382                    "control-flow body '{}' missing formal/captured input '{name}'",
2383                    key.1
2384                ))
2385            })?;
2386            let v = g.value_mut(vid);
2387            v.dtype = tensor.dtype;
2388            v.shape = tensor.shape.iter().map(|&d| Dim::Static(d)).collect();
2389        }
2390
2391        // Run shape inference over the seeded body (best-effort: interior shapes
2392        // that stay data-dependent are resolved just-in-time at run, exactly as
2393        // for the top-level graph).
2394        let registry = InferenceRegistry::default_registry();
2395        let opset_imports = self.graph.opset_imports.clone();
2396        registry.infer_graph(&mut g, &opset_imports, MergePolicy::Permissive)?;
2397
2398        let exec = Executor::build(g, self.weights.clone(), self.ep.clone())?;
2399        Ok(CompiledSubgraph {
2400            exec,
2401            input_names: prepared
2402                .formal_names
2403                .iter()
2404                .chain(prepared.capture_names.iter())
2405                .cloned()
2406                .collect(),
2407            built_shapes: externals.iter().map(|t| t.shape.clone()).collect(),
2408        })
2409    }
2410
2411    /// Run a prepared control-flow body with changing formal inputs. Captures and
2412    /// signature metadata are reused; only a concrete shape change rebuilds the
2413    /// child executor.
2414    fn run_subgraph(
2415        &mut self,
2416        prepared: &PreparedSubgraph,
2417        formal_inputs: &[&Tensor],
2418    ) -> Result<Vec<Tensor>> {
2419        if prepared.formal_names.len() != formal_inputs.len() {
2420            return Err(SessionError::Internal(format!(
2421                "control-flow body '{}' expects {} formal input(s) but {} were supplied",
2422                prepared.key.1,
2423                prepared.formal_names.len(),
2424                formal_inputs.len()
2425            )));
2426        }
2427
2428        let mut externals: Vec<&Tensor> =
2429            Vec::with_capacity(formal_inputs.len() + prepared.capture_names.len());
2430        externals.extend_from_slice(formal_inputs);
2431        for name in &prepared.capture_names {
2432            externals.push(prepared.captures.get(name).expect("prepared capture must be present"));
2433        }
2434        let rebuild = match self.subgraph_execs.get(&prepared.key) {
2435            Some(cs) => {
2436                cs.built_shapes.len() != externals.len()
2437                    || cs
2438                        .built_shapes
2439                        .iter()
2440                        .zip(externals.iter())
2441                        .any(|(built, tensor)| built != &tensor.shape)
2442            }
2443            None => true,
2444        };
2445        if rebuild {
2446            let child = self.build_subgraph_exec(prepared, &externals)?;
2447            self.subgraph_execs.insert(prepared.key.clone(), child);
2448            self.control_flow_stats.subgraph_builds += 1;
2449        }
2450
2451        self.control_flow_stats.subgraph_runs += 1;
2452        let cs = self.subgraph_execs.get_mut(&prepared.key).expect("child present");
2453        let inputs: Vec<(&str, &Tensor)> = cs
2454            .input_names
2455            .iter()
2456            .map(String::as_str)
2457            .zip(externals)
2458            .collect();
2459        cs.exec.run_scoped(&inputs, &prepared.captures)
2460    }
2461
2462    /// Dispatch a control-flow plan node to its op-specific handler.
2463    fn exec_control_flow(
2464        &mut self,
2465        pi: usize,
2466        resolved: &mut HashMap<ValueId, Vec<usize>>,
2467        outer_scope: &HashMap<String, Tensor>,
2468    ) -> Result<()> {
2469        let node = self.graph.node(self.plan[pi].node_id).clone();
2470        match node.op_type.as_str() {
2471            "If" => self.exec_if(&node, resolved, outer_scope),
2472            "Loop" => self.exec_loop(&node, resolved, outer_scope),
2473            "Scan" => self.exec_scan(&node, resolved, outer_scope),
2474            other => Err(SessionError::Internal(format!(
2475                "exec_control_flow reached non-control-flow op {other:?}"
2476            ))),
2477        }
2478    }
2479
2480    /// ONNX `If`: read the scalar `cond`, execute exactly one branch subgraph
2481    /// (0 formal inputs), and route the branch's outputs to `If`'s outputs.
2482    fn exec_if(
2483        &mut self,
2484        node: &Node,
2485        resolved: &mut HashMap<ValueId, Vec<usize>>,
2486        outer_scope: &HashMap<String, Tensor>,
2487    ) -> Result<()> {
2488        let cond_vid = node.inputs.first().and_then(|s| *s).ok_or_else(|| {
2489            SessionError::Internal("If node is missing its required 'cond' input".to_string())
2490        })?;
2491        let cond_t = self.value_tensor(cond_vid, resolved)?;
2492        let cond = tensor_scalar_bool(&cond_t).ok_or_else(|| SessionError::Internal(format!(
2493            "If: 'cond' must be a BOOL scalar, got dtype {:?} shape {:?}",
2494            cond_t.dtype, cond_t.shape
2495        )))?;
2496
2497        let attr_key = if cond { "then_branch" } else { "else_branch" };
2498        let prepared = self.prepare_subgraph(node.id, attr_key, resolved, outer_scope)?;
2499        let outs = self.run_subgraph(&prepared, &[])?;
2500
2501        if outs.len() != node.outputs.len() {
2502            return Err(SessionError::OutputShapeCountMismatch {
2503                op: format!("If/{attr_key}"),
2504                expected: node.outputs.len(),
2505                got: outs.len(),
2506            });
2507        }
2508        for (vid, t) in node.outputs.iter().zip(outs.iter()) {
2509            self.store_output_tensor(*vid, t, resolved)?;
2510        }
2511        Ok(())
2512    }
2513
2514    /// ONNX `Loop`: inputs `[M?, cond?, v_initial...]`, body signature
2515    /// `(iter_num, cond_in, carried...) -> (cond_out, carried..., scan_out...)`.
2516    /// Iterates while `cond` is true and `iter < M`, threading loop-carried
2517    /// values across iterations and stacking each scan output along a new
2518    /// leading iteration axis.
2519    fn exec_loop(
2520        &mut self,
2521        node: &Node,
2522        resolved: &mut HashMap<ValueId, Vec<usize>>,
2523        outer_scope: &HashMap<String, Tensor>,
2524    ) -> Result<()> {
2525        // Inputs: [M, cond, v_initial...]. M and cond may be omitted (None slot)
2526        // or an empty-name optional; absence means "unbounded" / "true".
2527        let m: Option<i64> = match node.inputs.first().and_then(|s| *s) {
2528            Some(vid) => {
2529                let t = self.value_tensor(vid, resolved)?;
2530                let m = tensor_scalar_i64(&t).ok_or_else(|| SessionError::Internal(format!(
2531                    "Loop: trip-count 'M' must be an INT64/INT32 scalar, got dtype {:?}",
2532                    t.dtype
2533                )))?;
2534                Some(m)
2535            }
2536            None => None,
2537        };
2538        let mut cond: Option<bool> = match node.inputs.get(1).and_then(|s| *s) {
2539            Some(vid) => {
2540                let t = self.value_tensor(vid, resolved)?;
2541                Some(tensor_scalar_bool(&t).ok_or_else(|| SessionError::Internal(format!(
2542                    "Loop: 'cond' must be a BOOL scalar, got dtype {:?}",
2543                    t.dtype
2544                )))?)
2545            }
2546            None => None,
2547        };
2548
2549        // Initial loop-carried dependencies (inputs after M and cond).
2550        let mut carried: Vec<Tensor> = Vec::new();
2551        for slot in node.inputs.iter().skip(2) {
2552            let vid = slot.ok_or_else(|| SessionError::Internal(
2553                "Loop: an interior loop-carried input is omitted (empty), which ONNX does not \
2554                 allow — every v_initial must be provided".to_string(),
2555            ))?;
2556            carried.push(self.value_tensor(vid, resolved)?);
2557        }
2558        let num_carried = carried.len();
2559        // Loop outputs = carried finals ++ scan outputs. Scan-output count is
2560        // whatever remains after the carried finals.
2561        let num_outputs = node.outputs.len();
2562        if num_outputs < num_carried {
2563            return Err(SessionError::Internal(format!(
2564                "Loop: node declares {num_outputs} output(s) but has {num_carried} loop-carried \
2565                 dependency(ies); outputs must be carried-finals followed by scan-outputs"
2566            )));
2567        }
2568        let num_scan = num_outputs - num_carried;
2569        let expected_iterations = m.and_then(|n| usize::try_from(n).ok());
2570        let mut scan_acc: Vec<TensorStackAccumulator> = (0..num_scan)
2571            .map(|_| TensorStackAccumulator::new(expected_iterations))
2572            .collect();
2573        let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope)?;
2574        let mut iter_tensor = scalar_i64_tensor(0)?;
2575        let mut cond_tensor = scalar_bool_tensor(cond.unwrap_or(true))?;
2576
2577        let mut iter: i64 = 0;
2578        loop {
2579            if let Some(m) = m
2580                && iter >= m
2581            {
2582                break;
2583            }
2584            if cond == Some(false) {
2585                break;
2586            }
2587
2588            iter_tensor.overwrite_bytes(&iter.to_le_bytes())?;
2589            cond_tensor.overwrite_bytes(&[u8::from(cond.unwrap_or(true))])?;
2590            let mut formal: Vec<&Tensor> = Vec::with_capacity(2 + num_carried);
2591            formal.push(&iter_tensor);
2592            formal.push(&cond_tensor);
2593            formal.extend(carried.iter());
2594
2595            let outs = self.run_subgraph(&prepared, &formal)?;
2596            drop(formal);
2597            // Body outputs: cond_out, carried..., scan_out...
2598            let expected = 1 + num_carried + num_scan;
2599            if outs.len() != expected {
2600                return Err(SessionError::OutputShapeCountMismatch {
2601                    op: "Loop/body".to_string(),
2602                    expected,
2603                    got: outs.len(),
2604                });
2605            }
2606            let mut it = outs.into_iter();
2607            let cond_out = it.next().expect("cond_out present");
2608            cond = Some(tensor_scalar_bool(&cond_out).ok_or_else(|| SessionError::Internal(
2609                format!(
2610                    "Loop: body's first output 'cond_out' must be a BOOL scalar, got dtype {:?}",
2611                    cond_out.dtype
2612                ),
2613            ))?);
2614            carried.clear();
2615            carried.extend((&mut it).take(num_carried));
2616            for acc in scan_acc.iter_mut() {
2617                acc.push(it.next().expect("scan output present"))?;
2618            }
2619
2620            iter += 1;
2621        }
2622
2623        // Emit outputs: carried finals, then stacked scan outputs.
2624        for (i, t) in carried.iter().enumerate() {
2625            self.store_output_tensor(node.outputs[i], t, resolved)?;
2626        }
2627        for (s, acc) in scan_acc.into_iter().enumerate() {
2628            let (dtype, shape, bytes) = acc.finish();
2629            self.store_output_bytes(
2630                node.outputs[num_carried + s],
2631                dtype,
2632                shape,
2633                &bytes,
2634                resolved,
2635            )?;
2636        }
2637        Ok(())
2638    }
2639
2640    /// ONNX `Scan`: inputs `[initial_state..., scan_input...]`, body signature
2641    /// `(state..., scan_slice...) -> (state..., scan_out_slice...)`. Iterates
2642    /// over the scan axis, threading state and stacking scan outputs along their
2643    /// axis. Phase-1 scope: scan axis 0 for every scan input/output, forward
2644    /// direction (the common exporter output); anything else is rejected
2645    /// clearly rather than silently mis-scanned.
2646    fn exec_scan(
2647        &mut self,
2648        node: &Node,
2649        resolved: &mut HashMap<ValueId, Vec<usize>>,
2650        outer_scope: &HashMap<String, Tensor>,
2651    ) -> Result<()> {
2652        let num_scan_inputs = node
2653            .attr("num_scan_inputs")
2654            .and_then(|a| a.as_int())
2655            .ok_or_else(|| SessionError::Internal(
2656                "Scan: required attribute 'num_scan_inputs' is missing or not an INT".to_string(),
2657            ))? as usize;
2658
2659        // Reject the axis/direction knobs this Phase-1 implementation does not
2660        // yet honor, rather than silently ignoring them (RULES #1/#5).
2661        for attr in ["scan_input_axes", "scan_output_axes"] {
2662            if let Some(a) = node.attr(attr)
2663                && let Some(axes) = a.as_ints()
2664                && axes.iter().any(|&ax| ax != 0)
2665            {
2666                return Err(SessionError::Internal(format!(
2667                    "Scan: attribute '{attr}' = {axes:?} requests a non-zero scan axis, \
2668                     which this runtime does not yet support. Expected axis 0 for every \
2669                     scan input/output; re-export with axis 0 or wait for full Scan-axis \
2670                     support"
2671                )));
2672            }
2673        }
2674        for attr in ["scan_input_directions", "scan_output_directions"] {
2675            if let Some(a) = node.attr(attr)
2676                && let Some(dirs) = a.as_ints()
2677                && dirs.iter().any(|&d| d != 0)
2678            {
2679                return Err(SessionError::Internal(format!(
2680                    "Scan: attribute '{attr}' = {dirs:?} requests reverse iteration, which \
2681                     this runtime does not yet support (forward only). Re-export forward or \
2682                     wait for reverse-Scan support"
2683                )));
2684            }
2685        }
2686
2687        let total_inputs = node.inputs.len();
2688        if total_inputs < num_scan_inputs {
2689            return Err(SessionError::Internal(format!(
2690                "Scan: node has {total_inputs} input(s) but num_scan_inputs={num_scan_inputs}"
2691            )));
2692        }
2693        let num_state = total_inputs - num_scan_inputs;
2694
2695        // Initial state (threaded across iterations) + scan inputs (sliced).
2696        let mut state: Vec<Tensor> = Vec::with_capacity(num_state);
2697        for slot in node.inputs.iter().take(num_state) {
2698            let vid = slot.ok_or_else(|| SessionError::Internal(
2699                "Scan: an initial-state input is omitted (empty), which ONNX does not allow"
2700                    .to_string(),
2701            ))?;
2702            state.push(self.value_tensor(vid, resolved)?);
2703        }
2704        let mut scan_inputs: Vec<Tensor> = Vec::with_capacity(num_scan_inputs);
2705        for slot in node.inputs.iter().skip(num_state) {
2706            let vid = slot.ok_or_else(|| SessionError::Internal(
2707                "Scan: a scan input is omitted (empty), which ONNX does not allow".to_string(),
2708            ))?;
2709            scan_inputs.push(self.value_tensor(vid, resolved)?);
2710        }
2711
2712        // Sequence length = extent of scan axis 0; all scan inputs must agree.
2713        let seq_len = scan_inputs
2714            .first()
2715            .and_then(|t| t.shape.first().copied())
2716            .ok_or_else(|| SessionError::Internal(
2717                "Scan: requires at least one scan input with rank >= 1".to_string(),
2718            ))?;
2719        for (i, t) in scan_inputs.iter().enumerate() {
2720            let this = t.shape.first().copied().unwrap_or(0);
2721            if this != seq_len {
2722                return Err(SessionError::Internal(format!(
2723                    "Scan: scan input #{i} has scan-axis length {this} but the first scan input has \
2724                     {seq_len}; all scan inputs must share the same scan-axis length"
2725                )));
2726            }
2727        }
2728
2729        let num_outputs = node.outputs.len();
2730        if num_outputs < num_state {
2731            return Err(SessionError::Internal(format!(
2732                "Scan: declares {num_outputs} output(s) but has {num_state} state variable(s); \
2733                 outputs must be final-state followed by scan-outputs"
2734            )));
2735        }
2736        let num_scan_out = num_outputs - num_state;
2737        let mut scan_acc: Vec<TensorStackAccumulator> = (0..num_scan_out)
2738            .map(|_| TensorStackAccumulator::new(Some(seq_len)))
2739            .collect();
2740        let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope)?;
2741        let mut scan_slices = Vec::with_capacity(num_scan_inputs);
2742        if seq_len != 0 {
2743            for t in &scan_inputs {
2744                let (shape, bytes) = leading_slice(t, 0)?;
2745                scan_slices.push(Tensor::from_raw_in(
2746                    self.ep.clone(),
2747                    t.dtype,
2748                    shape,
2749                    bytes,
2750                )?);
2751            }
2752        }
2753        for step in 0..seq_len {
2754            if step != 0 {
2755                for (source, slice) in scan_inputs.iter().zip(scan_slices.iter_mut()) {
2756                    let (_, bytes) = leading_slice(source, step)?;
2757                    slice.overwrite_bytes(bytes)?;
2758                }
2759            }
2760            let mut formal: Vec<&Tensor> = Vec::with_capacity(num_state + num_scan_inputs);
2761            formal.extend(state.iter());
2762            formal.extend(scan_slices.iter());
2763
2764            let outs = self.run_subgraph(&prepared, &formal)?;
2765            drop(formal);
2766            let expected = num_state + num_scan_out;
2767            if outs.len() != expected {
2768                return Err(SessionError::OutputShapeCountMismatch {
2769                    op: "Scan/body".to_string(),
2770                    expected,
2771                    got: outs.len(),
2772                });
2773            }
2774            let mut it = outs.into_iter();
2775            state.clear();
2776            state.extend((&mut it).take(num_state));
2777            for acc in scan_acc.iter_mut() {
2778                acc.push(it.next().expect("scan output present"))?;
2779            }
2780        }
2781
2782        for (i, t) in state.iter().enumerate() {
2783            self.store_output_tensor(node.outputs[i], t, resolved)?;
2784        }
2785        for (s, acc) in scan_acc.into_iter().enumerate() {
2786            let (dtype, shape, bytes) = acc.finish();
2787            self.store_output_bytes(node.outputs[num_state + s], dtype, shape, &bytes, resolved)?;
2788        }
2789        Ok(())
2790    }
2791}
2792
2793/// Borrow the `index`-th contiguous slice along a tensor's leading axis while
2794/// dropping that axis from the returned shape.
2795fn leading_slice(t: &Tensor, index: usize) -> Result<(Vec<usize>, &[u8])> {
2796    if t.shape.is_empty() {
2797        return Err(SessionError::Internal(
2798            "Scan: cannot slice a scalar scan input along axis 0".to_string(),
2799        ));
2800    }
2801    let outer = t.shape[0];
2802    if index >= outer {
2803        return Err(SessionError::Internal(format!(
2804            "Scan: slice index {index} out of range for scan-axis length {outer}"
2805        )));
2806    }
2807    let inner_shape = t.shape[1..].to_vec();
2808    let inner_numel: usize = inner_shape.iter().product();
2809    let esize = t.dtype.byte_size();
2810    // Sub-byte dtypes have no clean per-element byte stride for slicing; reject.
2811    if esize == 0 {
2812        return Err(SessionError::Internal(format!(
2813            "Scan: sub-byte dtype {:?} scan inputs are not supported",
2814            t.dtype
2815        )));
2816    }
2817    let slice_bytes = inner_numel * esize;
2818    let start = index * slice_bytes;
2819    let bytes = &t.as_bytes()[start..start + slice_bytes];
2820    Ok((inner_shape, bytes))
2821}
2822
2823/// Single-allocation accumulator for Loop/Scan scan outputs. Each iteration's
2824/// temporary output is copied directly into its final stacked byte position and
2825/// then dropped, avoiding a retained tensor allocation per step and a second
2826/// full stacking pass.
2827struct TensorStackAccumulator {
2828    expected_len: Option<usize>,
2829    dtype: Option<DataType>,
2830    elem_shape: Vec<usize>,
2831    len: usize,
2832    bytes: Vec<u8>,
2833}
2834
2835impl TensorStackAccumulator {
2836    fn new(expected_len: Option<usize>) -> Self {
2837        Self {
2838            expected_len,
2839            dtype: None,
2840            elem_shape: Vec::new(),
2841            len: 0,
2842            bytes: Vec::new(),
2843        }
2844    }
2845
2846    fn push(&mut self, tensor: Tensor) -> Result<()> {
2847        if let Some(dtype) = self.dtype {
2848            if tensor.shape != self.elem_shape || tensor.dtype != dtype {
2849                return Err(SessionError::Internal(format!(
2850                    "Loop/Scan: scan output slice {} has shape {:?} dtype {:?} but the first slice \
2851                     is shape {:?} dtype {:?}; every iteration's scan output must match",
2852                    self.len, tensor.shape, tensor.dtype, self.elem_shape, dtype
2853                )));
2854            }
2855        } else {
2856            if tensor.dtype.byte_size() == 0 {
2857                return Err(SessionError::Internal(format!(
2858                    "Loop/Scan: sub-byte dtype {:?} scan outputs are not supported",
2859                    tensor.dtype
2860                )));
2861            }
2862            self.dtype = Some(tensor.dtype);
2863            self.elem_shape = tensor.shape.clone();
2864            if let Some(expected) = self.expected_len {
2865                self.bytes.reserve(expected.saturating_mul(tensor.as_bytes().len()));
2866            }
2867        }
2868        self.bytes.extend_from_slice(tensor.as_bytes());
2869        self.len += 1;
2870        Ok(())
2871    }
2872
2873    fn finish(self) -> (DataType, Vec<usize>, Vec<u8>) {
2874        if self.len == 0 {
2875            return (DataType::Float32, vec![0], Vec::new());
2876        }
2877        let dtype = self.dtype.expect("non-empty accumulator has dtype");
2878        let mut shape = Vec::with_capacity(1 + self.elem_shape.len());
2879        shape.push(self.len);
2880        shape.extend(self.elem_shape);
2881        (dtype, shape, self.bytes)
2882    }
2883}
2884
2885impl Drop for Executor {
2886    fn drop(&mut self) {
2887        // Free every buffer via the owning EP (DeviceBuffer has no Drop).
2888        for (_, buf) in self.buffers.drain() {
2889            let _ = self.ep.deallocate(buf);
2890        }
2891    }
2892}
2893
2894/// Instantiate and initialize the Phase-1 CPU execution provider (§20.7,
2895/// CPU-only auto-detection). A GPU/accelerator EP would be prepended here in a
2896/// later phase; for Phase 1 the CPU EP is the sole, always-available backend.
2897pub(crate) fn auto_detect_cpu_ep() -> Result<Arc<CpuExecutionProvider>> {
2898    let mut ep = CpuExecutionProvider::new();
2899    ep.initialize(&Default::default())?;
2900    Ok(Arc::new(ep))
2901}
2902
2903#[cfg(test)]
2904mod tests {
2905    use super::*;
2906
2907    /// Holden's precondition: the dispatch-boundary gate must reject a view that
2908    /// addresses bytes past its backing allocation, rather than letting a kernel
2909    /// dereference out of bounds (UB).
2910    #[test]
2911    fn view_bounds_rejects_out_of_bounds_view() {
2912        // A [2, 3] f32 view needs 24 bytes; give it a 16-byte backing length.
2913        let shape = [2usize, 3];
2914        let strides = compute_contiguous_strides(&shape);
2915        let err = view_bounds(&shape, &strides, 0, DataType::Float32, 16);
2916        assert!(err.is_err(), "gate must reject an oversized view");
2917
2918        // Exactly-fitting length is accepted.
2919        assert!(view_bounds(&shape, &strides, 0, DataType::Float32, 24).is_ok());
2920    }
2921
2922    /// A negative byte offset region (via a byte_offset that pushes the origin
2923    /// past the buffer) is also rejected.
2924    #[test]
2925    fn view_bounds_rejects_offset_overrun() {
2926        let shape = [4usize];
2927        let strides = compute_contiguous_strides(&shape);
2928        // 4 f32 = 16 bytes; origin at byte 8 leaves only 8 bytes → overrun.
2929        assert!(view_bounds(&shape, &strides, 8, DataType::Float32, 16).is_err());
2930        assert!(view_bounds(&shape, &strides, 0, DataType::Float32, 16).is_ok());
2931    }
2932
2933    /// Symbol substitution: static dims pass through, bound symbols resolve, an
2934    /// unbound symbol yields `None` (the uninferred-shape signal).
2935    #[test]
2936    fn substitute_resolves_bound_symbols_only() {
2937        let mut bindings = HashMap::new();
2938        bindings.insert(SymbolId(0), 7usize);
2939        let shape = vec![Dim::Symbolic(SymbolId(0)), Dim::Static(4)];
2940        assert_eq!(substitute(&shape, &bindings), Some(vec![7, 4]));
2941
2942        let unbound = vec![Dim::Symbolic(SymbolId(1)), Dim::Static(4)];
2943        assert_eq!(substitute(&unbound, &bindings), None);
2944    }
2945
2946    /// H-D1: element-count multiplication must be overflow-checked so a huge or
2947    /// malicious shape reports `ShapeOverflow` instead of wrapping `usize` and
2948    /// under-sizing the buffer.
2949    #[test]
2950    fn checked_numel_detects_overflow() {
2951        // Well-formed shapes multiply normally.
2952        assert_eq!(checked_numel(&[2, 3, 4], || "v".into()).unwrap(), 24);
2953        assert_eq!(checked_numel(&[], || "v".into()).unwrap(), 1);
2954
2955        // A product past usize::MAX overflows.
2956        let huge = [usize::MAX, 2];
2957        let err = checked_numel(&huge, || "value#9".into());
2958        assert!(matches!(
2959            err,
2960            Err(SessionError::ShapeOverflow { .. })
2961        ));
2962    }
2963
2964    /// H-D1 (byte layer): even when the element *count* fits in `usize`, the
2965    /// count → bytes multiply can wrap for a fixed-width dtype. The allocation
2966    /// path must report `ShapeOverflow` rather than under-allocating.
2967    #[test]
2968    fn checked_storage_bytes_detects_byte_overflow() {
2969        // `usize::MAX / 4` elements fit in usize (pass checked_numel) but
2970        // `* 8` bytes for Float64 wraps — this is the exploited under-alloc.
2971        let numel = usize::MAX / 4;
2972        let err = checked_storage_bytes(DataType::Float64, numel, || "value#9".into(), &[numel]);
2973        assert!(matches!(err, Err(SessionError::ShapeOverflow { .. })));
2974
2975        // A well-formed size passes through unchanged.
2976        assert_eq!(
2977            checked_storage_bytes(DataType::Float32, 4, || "v".into(), &[4]).unwrap(),
2978            16
2979        );
2980    }
2981
2982    /// The data-dependent shape sizer must return exactly one shape per output
2983    /// so the run loop's `out_shapes[oi]` indexing can never misindex. Slice is
2984    /// single-output, so it returns a 1-element Vec; the run loop additionally
2985    /// guards the count (see `OutputShapeCountMismatch`).
2986    #[test]
2987    fn dynamic_output_shapes_slice_is_single_output() {
2988        let node = Node::new(NodeId(0), "Slice", vec![], vec![]);
2989        let input_shapes = vec![vec![4usize, 2]];
2990        let input_values = vec![
2991            None,           // data (unused by sizer)
2992            Some(vec![1]),  // starts
2993            Some(vec![3]),  // ends
2994            Some(vec![0]),  // axes
2995            Some(vec![1]),  // steps
2996        ];
2997        let out = dynamic_output_shapes(&node, &input_shapes, &input_values).unwrap();
2998        assert_eq!(out.len(), 1, "Slice must resolve exactly one output shape");
2999        assert_eq!(out[0], vec![2, 2]);
3000
3001        // An op the sizer cannot resolve returns None (surfaces as UnresolvedShape).
3002        let other = Node::new(NodeId(1), "Conv", vec![], vec![]);
3003        assert!(dynamic_output_shapes(&other, &input_shapes, &input_values).is_none());
3004    }
3005
3006    /// The effective opset is read from the graph's import for the op's domain,
3007    /// with the default and `ai.onnx` spellings treated as one.
3008    #[test]
3009    fn effective_opset_reads_graph_import() {
3010        let mut graph = Graph::default();
3011        graph.opset_imports.insert(String::new(), 12);
3012        let node = Node::new(NodeId(0), "Softmax", vec![], vec![]);
3013        assert_eq!(effective_opset(&graph, &node), 12);
3014
3015        graph.opset_imports.insert(String::new(), 0);
3016        assert_eq!(effective_opset(&graph, &node), 0);
3017
3018    }
3019
3020    #[test]
3021    #[should_panic(expected = "internal invariant violated")]
3022    fn effective_opset_requires_validated_import() {
3023        effective_opset(
3024            &Graph::default(),
3025            &Node::new(NodeId(0), "Softmax", vec![], vec![]),
3026        );
3027    }
3028
3029    // --- weight-streaming: zero-copy borrowed initializer buffers -----------
3030
3031    use onnx_runtime_ir::{static_shape, WeightRef};
3032    use std::path::PathBuf;
3033
3034    /// A writable scratch dir under the workspace `target/` (never `/tmp`).
3035    fn weightstream_tmp_dir() -> PathBuf {
3036        let dir = PathBuf::from(concat!(
3037            env!("CARGO_MANIFEST_DIR"),
3038            "/../../target/weightstream_test"
3039        ));
3040        std::fs::create_dir_all(&dir).expect("create weight-streaming test dir");
3041        dir
3042    }
3043
3044    fn f32_le(data: &[f32]) -> Vec<u8> {
3045        data.iter().flat_map(|v| v.to_le_bytes()).collect()
3046    }
3047
3048    /// (b) An aligned external-data initializer is backed **zero-copy** by a
3049    /// borrowed buffer whose data pointer EQUALS the WeightStore's mmap slice —
3050    /// no allocation, no copy. A model larger than RAM relies on this.
3051    #[test]
3052    fn aligned_external_initializer_is_borrowed_zero_copy() {
3053        let align = TensorLayout::contiguous().alignment;
3054        let path = weightstream_tmp_dir().join("aligned_init.bin");
3055        let w_data = [1.0f32, 2.0, 3.0, 4.0];
3056        std::fs::write(&path, f32_le(&w_data)).unwrap();
3057
3058        let mut store = WeightStore::new();
3059        store.map_external(&path).unwrap();
3060
3061        let mut g = Graph::new();
3062        g.opset_imports.insert(String::new(), 17);
3063        let w = g.create_named_value("W", DataType::Float32, static_shape([4]));
3064        g.set_initializer(
3065            w,
3066            WeightRef::External {
3067                path: path.clone(),
3068                offset: 0, // mmap base is page-aligned -> 0 is `align`-aligned
3069                length: 16,
3070                dtype: DataType::Float32,
3071                dims: vec![4],
3072            },
3073        );
3074        let y = g.create_value(DataType::Float32, static_shape([4]));
3075        g.insert_node(Node::new(NodeId(0), "Relu", vec![Some(w)], vec![y]));
3076        g.add_output(y);
3077
3078        let ep = auto_detect_cpu_ep().unwrap();
3079        let exec = Executor::build(g, Arc::new(store), ep).unwrap();
3080
3081        let weight = &exec.graph.initializers[&w];
3082        let src = exec.weights().bytes(weight).unwrap();
3083        assert!(
3084            (src.as_ptr() as usize).is_multiple_of(align),
3085            "mmap window must be aligned for this test to exercise the zero-copy path"
3086        );
3087        let buf = &exec.buffers[&w];
3088        assert!(buf.is_borrowed(), "aligned initializer must be borrowed, not copied");
3089        assert_eq!(
3090            buf.as_ptr() as *const u8,
3091            src.as_ptr(),
3092            "zero-copy: the buffer must alias the mmap bytes (no copy)"
3093        );
3094
3095        let _ = std::fs::remove_file(&path);
3096    }
3097
3098    /// (c) An unaligned external-data initializer falls back to an owned copy
3099    /// (buffer ptr != slice ptr) and is still numerically correct end-to-end.
3100    #[test]
3101    fn unaligned_external_initializer_falls_back_to_owned_copy() {
3102        let align = TensorLayout::contiguous().alignment;
3103        let path = weightstream_tmp_dir().join("unaligned_init.bin");
3104        // Prefix the weight window with 8 bytes so it starts at offset 8, which
3105        // is not a multiple of `align` (64) -> forces the copy fallback.
3106        let offset = 8usize;
3107        let w_data = [5.0f32, 6.0, 7.0, 8.0];
3108        let mut file = vec![0u8; offset];
3109        file.extend_from_slice(&f32_le(&w_data));
3110        std::fs::write(&path, &file).unwrap();
3111
3112        let mut store = WeightStore::new();
3113        store.map_external(&path).unwrap();
3114
3115        let mut g = Graph::new();
3116        g.opset_imports.insert(String::new(), 17);
3117        let w = g.create_named_value("W", DataType::Float32, static_shape([4]));
3118        g.set_initializer(
3119            w,
3120            WeightRef::External {
3121                path: path.clone(),
3122                offset,
3123                length: 16,
3124                dtype: DataType::Float32,
3125                dims: vec![4],
3126            },
3127        );
3128        let x = g.create_named_value("X", DataType::Float32, static_shape([4]));
3129        g.add_input(x);
3130        let y = g.create_value(DataType::Float32, static_shape([4]));
3131        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(x), Some(w)], vec![y]));
3132        g.add_output(y);
3133
3134        let ep = auto_detect_cpu_ep().unwrap();
3135        let mut exec = Executor::build(g, Arc::new(store), ep).unwrap();
3136
3137        let weight = &exec.graph.initializers[&w];
3138        let src = exec.weights().bytes(weight).unwrap();
3139        assert!(
3140            !(src.as_ptr() as usize).is_multiple_of(align),
3141            "window must be unaligned for this test to exercise the fallback"
3142        );
3143        let buf = &exec.buffers[&w];
3144        assert!(
3145            !buf.is_borrowed(),
3146            "unaligned initializer must fall back to an owned copy"
3147        );
3148        assert_ne!(
3149            buf.as_ptr() as *const u8,
3150            src.as_ptr(),
3151            "fallback: the buffer must be a fresh copy, not an alias"
3152        );
3153
3154        // The copy is numerically correct: Y = X + W.
3155        let x_tensor = Tensor::from_f32(&[4], &[10.0, 20.0, 30.0, 40.0]).unwrap();
3156        let out = exec.run(&[("X", &x_tensor)]).unwrap();
3157        assert_eq!(out.len(), 1);
3158        let got = out[0].to_vec_f32();
3159        let want = [15.0f32, 26.0, 37.0, 48.0];
3160        assert_eq!(got.len(), want.len());
3161        for (g, w) in got.iter().zip(want.iter()) {
3162            assert!((g - w).abs() < 1e-5, "got {g}, want {w}");
3163        }
3164
3165        let _ = std::fs::remove_file(&path);
3166    }
3167
3168    /// (d) Soundness guard: even when an initializer's mmap bytes are aligned
3169    /// (so the zero-copy path would otherwise fire), the executor must NOT
3170    /// borrow them if the value also has a producer — i.e. a malformed graph
3171    /// reused the initializer's `ValueId` as a node output. Borrowing yields a
3172    /// read-only buffer; a kernel writing that output would write through the
3173    /// mmap (SIGSEGV / aliasing UB). The build must fall back to an owned,
3174    /// writable copy instead.
3175    #[test]
3176    fn producer_backed_initializer_is_not_borrowed() {
3177        let align = TensorLayout::contiguous().alignment;
3178        let path = weightstream_tmp_dir().join("producer_backed_init.bin");
3179        let w_data = [1.0f32, 2.0, 3.0, 4.0];
3180        std::fs::write(&path, f32_le(&w_data)).unwrap();
3181
3182        let mut store = WeightStore::new();
3183        store.map_external(&path).unwrap();
3184
3185        let mut g = Graph::new();
3186        g.opset_imports.insert(String::new(), 17);
3187        let x = g.create_named_value("X", DataType::Float32, static_shape([4]));
3188        g.add_input(x);
3189        let w = g.create_named_value("W", DataType::Float32, static_shape([4]));
3190        g.set_initializer(
3191            w,
3192            WeightRef::External {
3193                path: path.clone(),
3194                offset: 0, // aligned: without the producer guard this would borrow
3195                length: 16,
3196                dtype: DataType::Float32,
3197                dims: vec![4],
3198            },
3199        );
3200        // Reuse the initializer's ValueId as a node output -> gives `w` a
3201        // producer, exactly the malformed shape the loader also rejects.
3202        g.insert_node(Node::new(NodeId(0), "Identity", vec![Some(x)], vec![w]));
3203        let y = g.create_value(DataType::Float32, static_shape([4]));
3204        g.insert_node(Node::new(NodeId(1), "Add", vec![Some(x), Some(w)], vec![y]));
3205        g.add_output(y);
3206
3207        assert!(
3208            g.value(w).producer.is_some(),
3209            "test setup: initializer value must have a producer",
3210        );
3211
3212        let ep = auto_detect_cpu_ep().unwrap();
3213        let exec = Executor::build(g, Arc::new(store), ep).unwrap();
3214
3215        let weight = &exec.graph.initializers[&w];
3216        let src = exec.weights().bytes(weight).unwrap();
3217        assert!(
3218            (src.as_ptr() as usize).is_multiple_of(align),
3219            "mmap window must be aligned so only the producer guard prevents borrowing",
3220        );
3221        let buf = &exec.buffers[&w];
3222        assert!(
3223            !buf.is_borrowed(),
3224            "producer-backed initializer must fall back to an owned writable copy",
3225        );
3226        assert_ne!(
3227            buf.as_ptr() as *const u8,
3228            src.as_ptr(),
3229            "producer-backed initializer must not alias read-only mmap bytes",
3230        );
3231
3232        let _ = std::fs::remove_file(&path);
3233    }
3234}