onnx_runtime_session/

lib.rs

//! # `onnx-runtime-session`
//!
//! The user-facing session and inference API for the ORT 2.0 runtime
//! (see `docs/ORT2.md` §20). Design goal: **zero-config by default** — the user
//! never has to know what an execution provider is; the runtime auto-detects
//! hardware and picks a strategy.
//!
//! **Phase 1 skeleton:** the intent-based [`SessionBuilder`] and
//! [`InferenceSession`] surfaces are defined; `build`/`run` bodies are
//! `todo!()` pending the sequential executor (Phase 1 task `ort2-session`).
//!
//! ```ignore
//! let mut session = onnx_runtime_session::load("model.onnx")?;
//! let outputs = session.run(&[("input_ids", &tensor)])?;
//! ```

// SessionError intentionally preserves rich, structured diagnostics in the
// public API; boxing it would be an API and behavior change rather than a lint fix.
#![allow(clippy::result_large_err)]
// Executor calls mirror graph/operator contracts with independently meaningful inputs.
#![allow(clippy::too_many_arguments)]

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use onnx_runtime_ir::{DataType, DeviceType, Shape};

pub use epcontext::{
    CompiledPartition, EpContextPlacement, dump_session_ep_context, load_ep_context_nodes,
};
pub use error::SessionError;
pub use executor::{
    CacheStats, CaptureDecline, CaptureDeclineReport, CapturePathKind, ControlFlowStats,
    DeviceAllocationCounts, DeviceGraphCaptureResult, ExecutionProviderDecline,
    ExecutionProviderFallbackReport, SeamReason, print_exec_phase_profile,
};
pub use onnx_runtime_loader::{
    EpContextDumpConfig, EpContextPartition, Model as EncoderModel, ModelMetadata,
};
pub use tensor::{DeviceBindingTransferStats, DeviceIoBinding, Tensor, cpu_allocator};

mod epcontext;
mod executor;
mod fp16_decode;
pub mod sequence;
mod tensor;

/// A graph output produced by the runtime.
#[derive(Debug)]
pub enum SessionOutput {
    Tensor(Tensor),
    Sequence(sequence::SequenceValue),
}

impl SessionOutput {
    pub fn as_tensor(&self) -> Option<&Tensor> {
        match self {
            Self::Tensor(tensor) => Some(tensor),
            Self::Sequence(_) => None,
        }
    }

    pub fn as_sequence(&self) -> Option<&sequence::SequenceValue> {
        match self {
            Self::Tensor(_) => None,
            Self::Sequence(sequence) => Some(sequence),
        }
    }

    pub fn into_tensor(self) -> Option<Tensor> {
        match self {
            Self::Tensor(tensor) => Some(tensor),
            Self::Sequence(_) => None,
        }
    }

    pub fn into_sequence(self) -> Option<sequence::SequenceValue> {
        match self {
            Self::Tensor(_) => None,
            Self::Sequence(sequence) => Some(sequence),
        }
    }
}

/// Operator-set version associated with an operator dispatch failure.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OpsetVersion {
    /// The model declares this version for the operator's domain.
    Known(u64),
    /// The model has no opset import for the operator's domain.
    Undeclared,
}

impl std::fmt::Display for OpsetVersion {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Known(version) => version.fmt(f),
            Self::Undeclared => f.write_str("<undeclared>"),
        }
    }
}

mod error {
    use super::OpsetVersion;

    struct UnsupportedOpRemediation<'a> {
        opset: OpsetVersion,
        domain: &'a str,
    }

    impl std::fmt::Display for UnsupportedOpRemediation<'_> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            if self.opset == OpsetVersion::Undeclared {
                write!(
                    f,
                    "declare an opset_import for domain {:?} in the model, ",
                    self.domain
                )?;
            }
            f.write_str(
                "enable another EP that supports this operator and opset, convert or decompose \
                 the model operator, or file an nxrt issue with the model details",
            )
        }
    }

    fn unsupported_op_remediation(
        opset: OpsetVersion,
        domain: &str,
    ) -> UnsupportedOpRemediation<'_> {
        UnsupportedOpRemediation { opset, domain }
    }

    /// Errors produced by the session layer.
    #[derive(Debug, thiserror::Error)]
    pub enum SessionError {
        #[error("session not initialized")]
        NotInitialized,

        #[error("input not found: {name}")]
        InputNotFound { name: String },

        #[error("unknown session option: {key}")]
        UnknownOption { key: String },

        #[error("invalid value {value:?} for session option {key:?}: expected one of {expected}")]
        InvalidOption {
            key: String,
            value: String,
            expected: String,
        },

        #[error("no model source: set a path or bytes on the builder")]
        NoModelSource,

        #[error("execution provider unavailable: {0}")]
        ExecutionProviderUnavailable(String),

        #[error(
            "CUDA execution required by ONNX_GENAI_REQUIRE_CUDA=1, but CPU fallback is needed: \
             {unsupported_nodes}"
        )]
        HeterogeneousPlacementRequired { unsupported_nodes: String },

        #[error(
            "unsupported operator {domain}::{op_type}: no available execution provider has a \
             kernel; node {node}, opset {opset}; decline reason: {reason}; consulted execution \
             providers (priority order): {execution_providers}. To fix: {remediation}",
            remediation = unsupported_op_remediation(*.opset, .domain)
        )]
        UnsupportedOp {
            op_type: String,
            domain: String,
            node: String,
            opset: OpsetVersion,
            reason: String,
            execution_providers: String,
        },

        #[error("value has a non-static (symbolic) shape and no binding to resolve it: {value}")]
        DynamicShape { value: String },

        #[error(
            "symbol {symbol} bound to conflicting sizes {first} and {second} across bound inputs"
        )]
        SymbolConflict {
            symbol: String,
            first: usize,
            second: usize,
        },

        #[error("input {name}: rank mismatch (graph declares rank {expected}, got {got})")]
        RankMismatch {
            name: String,
            expected: usize,
            got: usize,
        },

        #[error("no inferred shape for value {value} produced by op {op}")]
        UnresolvedShape { value: String, op: String },

        #[error("shape element count overflows usize for value {value} (dims {dims:?})")]
        ShapeOverflow { value: String, dims: Vec<usize> },

        #[error(
            "op {op} produced {got} data-dependent output shape(s) but has {expected} output(s)"
        )]
        OutputShapeCountMismatch {
            op: String,
            expected: usize,
            got: usize,
        },

        #[error(
            "runtime broadcast shape resolution failed for node {node} ({domain}::{op_type}): \
             concrete input shapes {input_shapes:?} are not broadcast-compatible, so no valid \
             elementwise output shape exists. To fix: update the model or runtime inputs so each \
             aligned dimension is equal or one of them is 1"
        )]
        RuntimeBroadcastIncompatible {
            node: String,
            domain: String,
            op_type: String,
            input_shapes: Vec<Vec<usize>>,
        },

        #[error("input {name}: dtype mismatch (expected {expected}, got {got})")]
        DtypeMismatch {
            name: String,
            expected: String,
            got: String,
        },

        #[error("input {name}: shape mismatch (expected {expected:?}, got {got:?})")]
        ShapeMismatch {
            name: String,
            expected: Vec<usize>,
            got: Vec<usize>,
        },

        #[error("internal executor error: {0}")]
        Internal(String),

        #[error("Sequence op {op}: {reason}")]
        SequenceOp { op: String, reason: String },

        #[error("control-flow op {op}: {reason}")]
        ControlFlow { op: String, reason: String },

        #[error(
            "EPContext reference node (main_context=0) has no matching primary \
             (source={source_key:?}, partition_name={partition_name:?})"
        )]
        DanglingEpContext {
            source_key: Option<String>,
            partition_name: Option<String>,
        },

        #[error(transparent)]
        Load(#[from] onnx_runtime_loader::LoaderError),

        #[error(transparent)]
        Ep(#[from] onnx_runtime_ep_api::EpError),

        #[error(transparent)]
        Ir(#[from] onnx_runtime_ir::IrError),

        #[error(transparent)]
        Graph(#[from] onnx_runtime_ir::GraphError),

        #[error(transparent)]
        Optimize(#[from] onnx_runtime_optimizer::OptimizerError),

        #[error(transparent)]
        ShapeInfer(#[from] onnx_runtime_shape_inference::ShapeInferError),
    }

    impl SessionError {
        pub(crate) fn unsupported_op(
            node: &onnx_runtime_ir::Node,
            node_id: onnx_runtime_ir::NodeId,
            opset: u64,
            execution_providers: impl Into<String>,
            reason: impl Into<String>,
        ) -> Self {
            let domain = if node.domain.is_empty() {
                "ai.onnx".to_string()
            } else {
                node.domain.clone()
            };
            let node_display = if node.name.is_empty() {
                format!("<unnamed node #{}>", node_id.0)
            } else {
                format!("{:?}", node.name)
            };
            let opset = if opset == u64::MAX {
                OpsetVersion::Undeclared
            } else {
                OpsetVersion::Known(opset)
            };
            Self::UnsupportedOp {
                op_type: node.op_type.clone(),
                domain,
                node: node_display,
                opset,
                reason: reason.into(),
                execution_providers: execution_providers.into(),
            }
        }
    }

    /// Session `Result` alias.
    pub type Result<T> = std::result::Result<T, SessionError>;
}

use error::Result;

/// Metadata describing a model input or output (§20.2).
#[derive(Clone, Debug)]
pub struct IoMeta {
    pub name: String,
    pub dtype: DataType,
    pub shape: Shape,
}

/// Intent-based device preference (§20.4). The runtime maps this to concrete
/// EPs during `build`.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum DevicePreference {
    /// Pick the best available device automatically.
    #[default]
    Auto,
    /// Prefer CPU execution.
    Cpu,
    /// Prefer a GPU / accelerator, optionally by ordinal.
    Gpu { index: Option<u32> },
    /// Pin to a specific device class + ordinal.
    Explicit { device_type: DeviceType, index: u32 },
}

/// A shape to pre-compile kernels for at session init (§11.3).
#[derive(Clone, Debug)]
pub struct WarmupShape {
    pub input_name: String,
    pub shape: Vec<usize>,
}

/// Decoder-wide numeric precision for the session's decode graph.
///
/// This is a generic, model-agnostic knob selected via
/// [`SessionBuilder::decode_precision`]. The default,
/// [`DecodePrecision::Model`], runs the graph exactly as authored, so default
/// runtime behaviour is byte-identical to a build with no precision knob at all.
///
/// [`DecodePrecision::Fp16`] requests a whole-decoder fp32→fp16 rewrite (see
/// [`fp16_decode`]). It only takes effect on a GPU device and only for an
/// fp32-activation int4/block-32 quantized decoder (fp32-scale `MatMulNBits`);
/// for every other model — including native fp16-activation models — it is a
/// strict no-op, leaving the graph bit-identical.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum DecodePrecision {
    /// Run the decoder at the precision authored in the model graph (default).
    #[default]
    Model,
    /// Cast an fp32-activation int4/block-32 decoder to a fully fp16 graph so a
    /// GPU backend runs it through the fp16-fused decode kernels. No-op unless
    /// the session targets a GPU and the graph is fp32-activation quantized.
    Fp16,
}

/// Graph-optimization level for the session's `optimize` pipeline stage
/// (`docs/ORT2.md` §18). Selected via the generic `"optimization"` session
/// option (see [`SessionBuilder::option`]).
///
/// The default is [`OptimizationLevel::None`]: with optimization off the graph
/// reaches the executor exactly as the loader produced it, so default runtime
/// behavior is byte-identical to a build with no optimizer wired in at all.
///
/// This is a generic, model-agnostic knob — no level ever special-cases a model
/// name or op. Higher levels simply enable more of the device-independent pass
/// pipeline from [`onnx_runtime_optimizer`].
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum OptimizationLevel {
    /// No passes — the `optimize` stage is a no-op (default).
    #[default]
    None,
    /// Structure-preserving passes only: constant folding then dead-node
    /// elimination. No operator fusion, so the op set the executor sees is a
    /// subset of the loaded graph's.
    Basic,
    /// The full device-independent pipeline: constant folding, dead-node
    /// elimination, and operator fusion (which can introduce fused
    /// `com.microsoft` contrib ops such as `LayerNormalization`).
    All,
}

impl OptimizationLevel {
    /// Parse the `"optimization"` option value. Accepts `"none"`, `"basic"`,
    /// and `"all"` (case-insensitive).
    fn parse(value: &str) -> Option<Self> {
        match value.trim().to_ascii_lowercase().as_str() {
            "none" | "off" | "0" => Some(Self::None),
            "basic" => Some(Self::Basic),
            "all" => Some(Self::All),
            _ => None,
        }
    }

    /// The optimizer passes this level enables, in pipeline order. Empty for
    /// [`OptimizationLevel::None`].
    fn passes(self) -> Vec<Box<dyn onnx_runtime_optimizer::OptimizationPass>> {
        use onnx_runtime_optimizer::{ConstantFolding, DeadNodeElimination, OpFusion};
        match self {
            Self::None => Vec::new(),
            Self::Basic => vec![Box::new(ConstantFolding), Box::new(DeadNodeElimination)],
            Self::All => vec![
                Box::new(ConstantFolding),
                Box::new(DeadNodeElimination),
                Box::new(OpFusion::new()),
            ],
        }
    }
}

/// Builder for advanced session configuration (§20.6).
#[derive(Default)]
pub struct SessionBuilder {
    model_path: Option<PathBuf>,
    model_bytes: Option<Vec<u8>>,
    device: DevicePreference,
    execution_provider: Option<std::sync::Arc<dyn onnx_runtime_ep_api::ExecutionProvider>>,
    memory_limit: Option<usize>,
    enable_profiling: bool,
    warmup_shapes: Vec<WarmupShape>,
    decode_precision: DecodePrecision,
    options: HashMap<String, String>,
}

impl SessionBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn model(mut self, path: impl AsRef<Path>) -> Self {
        self.model_path = Some(path.as_ref().to_path_buf());
        self
    }

    pub fn model_bytes(mut self, bytes: &[u8]) -> Self {
        self.model_bytes = Some(bytes.to_vec());
        self
    }

    pub fn device(mut self, pref: DevicePreference) -> Self {
        self.device = pref;
        self
    }

    /// Use an explicitly constructed execution provider instead of device auto-selection.
    pub fn execution_provider(
        mut self,
        execution_provider: std::sync::Arc<dyn onnx_runtime_ep_api::ExecutionProvider>,
    ) -> Self {
        self.execution_provider = Some(execution_provider);
        self
    }

    pub fn memory_limit(mut self, bytes: usize) -> Self {
        self.memory_limit = Some(bytes);
        self
    }

    pub fn profiling(mut self, enable: bool) -> Self {
        self.enable_profiling = enable;
        self
    }

    pub fn warmup(mut self, shapes: Vec<WarmupShape>) -> Self {
        self.warmup_shapes = shapes;
        self
    }

    /// Select the decoder-wide numeric precision (see [`DecodePrecision`]).
    /// Defaults to [`DecodePrecision::Model`] (graph as authored); selecting
    /// [`DecodePrecision::Fp16`] opts into the fp32→fp16 decode rewrite, which
    /// only takes effect on a GPU fp32-activation quantized decoder.
    pub fn decode_precision(mut self, precision: DecodePrecision) -> Self {
        self.decode_precision = precision;
        self
    }

    /// Set a namespaced option. Unknown keys — and unknown values for a known
    /// key — are rejected at [`Self::build`].
    ///
    /// # Recognized options
    ///
    /// | Key                     | Values                       | Default  | Effect |
    /// |-------------------------|------------------------------|----------|--------|
    /// | `"optimization"`        | `"none"`, `"basic"`, `"all"` | `"none"` | Graph optimization level (see [`OptimizationLevel`]). |
    /// | `"ep.context_enable"`   | `"0"`/`"1"`/`"false"`/`"true"` | `"0"`  | Dump a `*_ctx.onnx` EPContext model after compile (§21.4 / §55.4). |
    /// | `"ep.context_file_path"`| any path                     | `<orig>_ctx.onnx` | Output path for the generated context model. |
    /// | `"ep.context_embed_mode"`| `"0"` (external) / `"1"` (embed) | `"1"` | How the compiled blob is stored in each EPContext node. |
    ///
    /// `"optimization"` = `"none"` (the default) leaves the loaded graph
    /// untouched, so behavior is byte-identical to a runtime with no optimizer.
    /// `"basic"` runs constant folding + dead-node elimination; `"all"` adds
    /// operator fusion. When any pass runs, the session re-runs shape inference
    /// on the rewritten graph before compiling so fused/introduced nodes get
    /// inferred shapes.
    pub fn option(mut self, key: &str, value: &str) -> Self {
        self.options.insert(key.to_string(), value.to_string());
        self
    }

    /// Parse every set session option in a single pass, rejecting any unknown
    /// key or unparseable value up front (no silent compat shim — an
    /// unrecognized key is a typo, never a no-op). Returns the resolved
    /// [`OptimizationLevel`] and the EPContext dump config (§21.4 / §55.5)
    /// driven by the `ep.context_*` keys.
    ///
    /// # Recognized keys
    ///
    /// * `"optimization"` → [`OptimizationLevel`] (`none` / `basic` / `all`).
    /// * `"ep.context_enable"` → [`EpContextDumpConfig::enable`]
    ///   (`1`/`0`/`true`/`false`, case-insensitive).
    /// * `"ep.context_file_path"` → [`EpContextDumpConfig::file_path`] (an empty
    ///   value clears it back to the `<orig>_ctx.onnx` default).
    /// * `"ep.context_embed_mode"` → [`EpContextDumpConfig::embed_mode`]
    ///   (`0` external file / `1` embed; any other value is rejected).
    fn parse_options(
        options: &HashMap<String, String>,
    ) -> Result<(OptimizationLevel, EpContextDumpConfig)> {
        let mut level = OptimizationLevel::None;
        let mut ctx = EpContextDumpConfig::default();
        for (key, value) in options {
            match key.as_str() {
                "optimization" => {
                    level = OptimizationLevel::parse(value).ok_or_else(|| {
                        SessionError::InvalidOption {
                            key: key.clone(),
                            value: value.clone(),
                            expected: "none, basic, all".to_string(),
                        }
                    })?;
                }
                "ep.context_enable" => {
                    ctx.enable = parse_bool_option(key, value)?;
                }
                "ep.context_file_path" => {
                    // Empty/unset ⇒ None (fall back to `<orig>_ctx.onnx`).
                    ctx.file_path = if value.trim().is_empty() {
                        None
                    } else {
                        Some(PathBuf::from(value))
                    };
                }
                "ep.context_embed_mode" => {
                    ctx.embed_mode = parse_embed_mode(key, value)?;
                }
                // No compat shim: an unrecognized key is a typo, not a silent
                // no-op.
                _ => return Err(SessionError::UnknownOption { key: key.clone() }),
            }
        }
        Ok((level, ctx))
    }

    /// Build the session: load → detect device → optimize → compile → allocate.
    ///
    /// The `optimize` stage is driven by the `"optimization"` session option and
    /// defaults to [`OptimizationLevel::None`] (a no-op), so the default path is
    /// byte-identical to loading straight into the executor. When optimization
    /// is enabled the pipeline is:
    ///
    /// ```text
    /// load (+ loader shape inference)
    ///   → run optimizer passes (constant-fold / DCE / fusion)
    ///   → re-run shape inference on the rewritten graph
    ///   → compile (kernel per node) → allocate
    /// ```
    ///
    /// The re-inference step is essential: fusion can replace a multi-op
    /// decomposition (e.g. the 9-op LayerNorm) with a single fused node whose
    /// output has no inferred shape yet, and the compile/execute stages require
    /// every value to carry a resolved shape.
    ///
    /// Device selection keeps CPU as the default and selects CUDA only when
    /// explicitly requested in a CUDA-enabled build. "Compile" resolves a
    /// kernel per node into the shape-keyed cache.
    pub fn build(self) -> Result<InferenceSession> {
        let (level, ep_context_config) = Self::parse_options(&self.options)?;

        // Memory limits and profiling remain reserved builder intents.
        let _ = (self.memory_limit, self.enable_profiling);

        let (mut graph, weights, model_dir, model_metadata) =
            match (self.model_path, self.model_bytes) {
                (Some(path), _) => {
                    // The EPContext load path resolves `embed_mode=0` external blob
                    // paths relative to the model file's directory (§55.3), so
                    // retain it (same base dir the loader used for external data).
                    let model_dir = path
                        .parent()
                        .map(Path::to_path_buf)
                        .unwrap_or_else(|| PathBuf::from("."));
                    let bytes = onnx_runtime_loader::read_model_binary(&path)?;
                    let metadata = model_metadata_from_bytes(&bytes)?;
                    let (g, w) =
                        onnx_runtime_loader::load_model_bytes_with_weights(&bytes, &model_dir)?;
                    (g, w, model_dir, metadata)
                }
                (None, Some(bytes)) => {
                    let metadata = model_metadata_from_bytes(&bytes)?;
                    let (g, w) = onnx_runtime_loader::load_model_bytes_with_weights(&bytes, ".")?;
                    (g, w, PathBuf::from("."), metadata)
                }
                (None, None) => return Err(SessionError::NoModelSource),
            };

        // Decoder precision rewrite (opt-in). Applied here — on the freshly
        // loaded graph, before the I/O signature (`IoMeta`) is computed and
        // before EP optimization — so the KV/logits buffers and the executor
        // graph agree on the rewritten dtype. A strict no-op for the default
        // `DecodePrecision::Model`, for non-GPU devices, and for any graph that
        // is not an fp32-activation quantized decoder, so the default path and
        // native fp16 models stay bit-identical.
        fp16_decode::maybe_convert_decode_fp16(
            &mut graph,
            &weights,
            self.decode_precision,
            device_preference_is_gpu(&self.device),
        );

        // Optimize stage. Off by default; only runs when a level is selected.
        optimize_graph(&mut graph, level)?;
        let ep = match self.execution_provider {
            Some(ep) => ep,
            None => select_execution_provider(&self.device)?,
        };

        let mut session = InferenceSession::from_parts(
            graph,
            weights,
            &model_dir,
            ep_context_config,
            model_metadata,
            ep,
        )?;
        if !self.warmup_shapes.is_empty() {
            session.warmup(&self.warmup_shapes)?;
        }
        Ok(session)
    }
}

/// Whether a [`DevicePreference`] targets a GPU / accelerator (non-host) device.
/// Used to gate GPU-only graph rewrites such as the fp16 decode precision mode.
fn device_preference_is_gpu(preference: &DevicePreference) -> bool {
    match preference {
        DevicePreference::Gpu { .. } => true,
        DevicePreference::Explicit { device_type, .. } => !device_type.is_host_accessible(),
        DevicePreference::Auto | DevicePreference::Cpu => false,
    }
}

fn select_execution_provider(
    preference: &DevicePreference,
) -> Result<std::sync::Arc<dyn onnx_runtime_ep_api::ExecutionProvider>> {
    match preference {
        // Keep the zero-config/default behavior CPU-only. CUDA is an explicit
        // opt-in until heterogeneous placement and fallback exist.
        DevicePreference::Auto | DevicePreference::Cpu => executor::auto_detect_cpu_ep(),
        DevicePreference::Explicit {
            device_type: DeviceType::Cpu,
            index: 0,
        } => executor::auto_detect_cpu_ep(),
        DevicePreference::Gpu { index } => cuda_execution_provider(index.unwrap_or(0)),
        DevicePreference::Explicit {
            device_type: DeviceType::Cuda,
            index,
        } => cuda_execution_provider(*index),
        DevicePreference::Explicit { device_type, index } => {
            Err(SessionError::ExecutionProviderUnavailable(format!(
                "{device_type:?}:{index} is not implemented by onnx-runtime-session"
            )))
        }
    }
}

#[cfg(feature = "cuda")]
fn cuda_execution_provider(
    index: u32,
) -> Result<std::sync::Arc<dyn onnx_runtime_ep_api::ExecutionProvider>> {
    let mut ep = onnx_runtime_ep_cuda::CudaExecutionProvider::new(index)?;
    onnx_runtime_ep_api::ExecutionProvider::initialize(&mut ep, &Default::default())?;
    Ok(std::sync::Arc::new(ep))
}

#[cfg(not(feature = "cuda"))]
fn cuda_execution_provider(
    index: u32,
) -> Result<std::sync::Arc<dyn onnx_runtime_ep_api::ExecutionProvider>> {
    Err(SessionError::ExecutionProviderUnavailable(format!(
        "CUDA:{index} requested, but onnx-runtime-session was built without the `cuda` feature"
    )))
}

fn model_metadata_from_bytes(bytes: &[u8]) -> Result<ModelMetadata> {
    let model = onnx_runtime_loader::proto::decode_model(bytes)?;
    Ok(ModelMetadata {
        ir_version: model.ir_version,
        producer_name: model.producer_name,
        producer_version: model.producer_version,
        domain: model.domain,
        model_version: model.model_version,
        doc_string: (!model.doc_string.is_empty()).then_some(model.doc_string),
        graph_name: model.graph.map(|graph| graph.name).unwrap_or_default(),
        metadata_props: model
            .metadata_props
            .into_iter()
            .map(|entry| (entry.key, entry.value))
            .collect(),
    })
}

/// Parse a boolean-ish session-option value (§21.4). Accepts `1`/`0` and
/// `true`/`false` (case-insensitive), mirroring how ORT's C API treats its
/// `int`-typed `ep.context_enable` flag while also allowing the textual form.
/// Any other value is a typo, surfaced as [`SessionError::InvalidOption`].
fn parse_bool_option(key: &str, value: &str) -> Result<bool> {
    match value.trim().to_ascii_lowercase().as_str() {
        "1" | "true" => Ok(true),
        "0" | "false" => Ok(false),
        _ => Err(SessionError::InvalidOption {
            key: key.to_string(),
            value: value.to_string(),
            expected: "0, 1, true, false".to_string(),
        }),
    }
}

/// Parse the `ep.context_embed_mode` option (§21.4): `0` = external sidecar
/// file, `1` = embed the blob inline. Any other value is rejected with
/// [`SessionError::InvalidOption`] (mirroring [`OptimizationLevel::parse`]'s
/// fail-closed rejection rather than silently clamping).
fn parse_embed_mode(key: &str, value: &str) -> Result<u8> {
    match value.trim() {
        "0" => Ok(0),
        "1" => Ok(1),
        _ => Err(SessionError::InvalidOption {
            key: key.to_string(),
            value: value.to_string(),
            expected: "0, 1".to_string(),
        }),
    }
}

/// Run the optimizer passes selected by `level`, then re-run shape inference so
/// any node fusion introduced (whose outputs the loader never saw) gets a fully
/// inferred shape/dtype before compile.
///
/// A no-op when `level` is [`OptimizationLevel::None`] — the graph is returned
/// untouched and no re-inference runs, keeping the default path byte-identical.
fn optimize_graph(graph: &mut onnx_runtime_ir::Graph, level: OptimizationLevel) -> Result<()> {
    let passes = level.passes();
    if passes.is_empty() {
        return Ok(());
    }

    onnx_runtime_optimizer::run_passes(
        graph,
        &passes,
        &onnx_runtime_optimizer::PassContext::new(),
    )?;

    // Fusion emits fused ops in the `com.microsoft` contrib domain; make sure
    // that domain is imported so shape-inference and kernel dispatch pick the
    // contrib-registered rules (they register from opset 1, but recording the
    // import keeps the graph self-consistent and future-proofs versioned rules).
    graph
        .opset_imports
        .entry(onnx_runtime_optimizer::CONTRIB_DOMAIN.to_string())
        .or_insert(1);

    // Re-infer shapes over the rewritten graph: fused nodes' outputs (and any
    // value whose producer changed) must be re-resolved before compile.
    let registry = onnx_runtime_shape_inference::InferenceRegistry::default_registry();
    let opset_imports = graph.opset_imports.clone();
    registry.infer_graph(
        graph,
        &opset_imports,
        onnx_runtime_shape_inference::MergePolicy::Permissive,
    )?;

    Ok(())
}

/// A loaded model ready to run inference (§20.2).
pub struct InferenceSession {
    inputs: Vec<IoMeta>,
    outputs: Vec<IoMeta>,
    model_metadata: ModelMetadata,
    exec: executor::Executor,
    /// EPContext dump config parsed from the `ep.context_*` session options
    /// (§21.4). Drives [`InferenceSession::export_ep_context`]; disabled by
    /// default so an ordinary session never touches the dump path.
    ep_context_config: EpContextDumpConfig,
}

fn io_meta(graph: &onnx_runtime_ir::Graph, values: &[onnx_runtime_ir::ValueId]) -> Vec<IoMeta> {
    values
        .iter()
        .map(|&vid| {
            let v = graph.value(vid);
            IoMeta {
                name: v.name.clone().unwrap_or_default(),
                dtype: v.dtype,
                shape: v.shape.clone(),
            }
        })
        .collect()
}

impl InferenceSession {
    /// Primary entry point: load a model with auto device detection.
    pub fn load(path: impl AsRef<Path>) -> Result<Self> {
        Self::builder().model(path).build()
    }

    /// Load a model from an in-memory buffer.
    pub fn load_bytes(bytes: &[u8]) -> Result<Self> {
        Self::builder().model_bytes(bytes).build()
    }

    /// Build a session directly from an in-memory IR [`Graph`](onnx_runtime_ir::Graph).
    ///
    /// Initializer bytes are read from the graph's inline [`WeightRef`]s, so no
    /// on-disk model or weight store is required. Useful for programmatically
    /// constructed graphs and tests.
    pub fn from_graph(graph: onnx_runtime_ir::Graph) -> Result<Self> {
        // No on-disk model: `embed_mode=0` external EPContext blobs resolve
        // relative to the current directory (consistent with the loader's
        // in-memory `base_dir` default).
        Self::from_parts(
            graph,
            std::sync::Arc::new(onnx_runtime_loader::WeightStore::new()),
            Path::new("."),
            EpContextDumpConfig::default(),
            ModelMetadata::default(),
            executor::auto_detect_cpu_ep()?,
        )
    }

    fn from_parts(
        graph: onnx_runtime_ir::Graph,
        weights: std::sync::Arc<onnx_runtime_loader::WeightStore>,
        model_dir: &Path,
        ep_context_config: EpContextDumpConfig,
        model_metadata: ModelMetadata,
        ep: std::sync::Arc<dyn onnx_runtime_ep_api::ExecutionProvider>,
    ) -> Result<Self> {
        // Establish the canonical-domain invariant for programmatically built
        // graphs (the loader already normalizes at proto-materialization time):
        // the default ONNX domain is `""`, never `"ai.onnx"`. The executor and
        // validators rely on this, comparing `domain.is_empty()` directly.
        let mut graph = graph;
        graph.normalize_domains();

        onnx_runtime_loader::validate_model(&graph)?;

        let inputs = io_meta(&graph, &graph.inputs);
        let outputs = io_meta(&graph, &graph.outputs);
        // EPContext consume path (§55.3): restore any pre-compiled EP contexts
        // before building the executor. Dispatch is a pure `source`-key lookup
        // over the session's selected EP, so a model carrying EPContext nodes
        // for an unloaded compiled EP fails with a clear `NoEpForContext`. The
        // executor then bypasses these nodes (they are pre-compiled, never run
        // as ordinary kernels).
        let eps: [(
            onnx_runtime_ep_api::EpId,
            &dyn onnx_runtime_ep_api::ExecutionProvider,
        ); 1] = [(onnx_runtime_ep_api::EpId(0), ep.as_ref())];
        epcontext::load_ep_context_nodes(&graph, model_dir, &eps)?;

        let exec = executor::Executor::build(graph, weights, ep)?;
        Ok(Self {
            inputs,
            outputs,
            model_metadata,
            exec,
            ep_context_config,
        })
    }

    /// Start a configuration builder.
    pub fn builder() -> SessionBuilder {
        SessionBuilder::new()
    }

    /// Run inference with named inputs, returning the graph outputs in order.
    pub fn run(&mut self, inputs: &[(&str, &Tensor)]) -> Result<Vec<Tensor>> {
        self.exec.run(inputs)
    }

    /// Attach the shared runtime trace context. When enabled, the executor opens
    /// one span per executed op so kernels can attach kernel-variant and
    /// capture-rejection reasons to a live span. Defaults to a disabled no-op
    /// context, so untraced runs pay only a single relaxed atomic load per op.
    pub fn set_trace_context(&mut self, trace: onnx_runtime_tracer::TraceContext) {
        self.exec.set_trace_context(trace);
    }

    /// Run inference and preserve tensor or sequence graph-output types.
    pub fn run_outputs(&mut self, inputs: &[(&str, &Tensor)]) -> Result<Vec<SessionOutput>> {
        self.exec.run_outputs(inputs)
    }

    /// F5 Stage 1 decode-plan memo activity counters `(primed, rebuilt, replayed,
    /// ineligible)` over this session's lifetime. `replayed > 0` after a decode
    /// run proves the memo actually engaged on the real (persistent-KV-binding)
    /// path; the coordinator's on-model A/B reads this to reject a vacuous pass.
    pub fn decode_memo_counts(&self) -> (u64, u64, u64, u64) {
        self.exec.decode_memo_counts()
    }

    /// F5 Stage 2 view-plan activity counters `(views_reused, dispatch_elided)`
    /// over this session's lifetime. Both `> 0` after a decode run prove the
    /// invariant zero-copy view reuse and pure-view dispatch elision actually
    /// fired on the real path (not a vacuous pass); an on-model A/B reads this
    /// alongside [`Self::decode_memo_counts`].
    pub fn decode_view_plan_counts(&self) -> (u64, u64) {
        self.exec.decode_view_plan_counts()
    }

    /// Run with persistent device allocations supplying graph inputs and,
    /// optionally, aliasing graph outputs. Bound outputs are returned as `None`
    /// because their bytes remain resident in the caller-owned allocation.
    pub fn run_with_device_bindings(
        &mut self,
        inputs: &[(&str, &Tensor)],
        bindings: &mut [DeviceIoBinding],
    ) -> Result<Vec<Option<Tensor>>> {
        self.exec.run_with_device_bindings(inputs, bindings)
    }

    /// Allocate a persistent buffer on this session's execution device.
    pub fn allocate_device_binding(
        &self,
        input_name: impl Into<String>,
        output_name: Option<impl Into<String>>,
        dtype: DataType,
        physical_shape: Vec<usize>,
        logical_shape: Vec<usize>,
    ) -> Result<DeviceIoBinding> {
        self.exec.allocate_device_binding(
            input_name.into(),
            output_name.map(Into::into),
            dtype,
            physical_shape,
            logical_shape,
        )
    }

    /// Allocate a persistent buffer for a graph output without also binding it
    /// as an input.
    pub fn allocate_device_output_binding(
        &self,
        output_name: impl Into<String>,
        dtype: DataType,
        physical_shape: Vec<usize>,
        logical_shape: Vec<usize>,
    ) -> Result<DeviceIoBinding> {
        self.exec.allocate_device_output_binding(
            output_name.into(),
            dtype,
            physical_shape,
            logical_shape,
        )
    }

    /// Execute once while recording the kernel launches into a device graph.
    ///
    /// `NotCapturable` means the mandatory all-kernel audit rejected the run
    /// before stream capture began, so callers may safely retry eagerly.
    pub fn try_capture_with_device_bindings(
        &mut self,
        inputs: &[(&str, &Tensor)],
        bindings: &mut [DeviceIoBinding],
    ) -> Result<DeviceGraphCaptureResult> {
        self.exec.try_capture_with_device_bindings(inputs, bindings)
    }

    /// Replay the installed device graph after the caller has refreshed any
    /// persistent scalar inputs. Returns `true` when the graph is still valid for
    /// the next step, or `false` when a control-flow branch flip retired it this
    /// step (the token was produced correctly via eager fallback) and the caller
    /// should re-warm and re-capture.
    pub fn replay_device_graph(&mut self, bindings: &mut [DeviceIoBinding]) -> Result<bool> {
        self.exec.replay_device_graph(bindings)
    }

    /// Invalidate the installed device graph before reset, rewind, shape change,
    /// or binding destruction.
    pub fn reset_device_graph(&mut self) -> Result<bool> {
        self.exec.reset_device_graph()
    }

    /// Number of captured device-graph segments installed by the most recent
    /// [`Self::try_capture_with_device_bindings`] call.
    ///
    /// `1` for a whole-subgraph capture; `>= 2` when the CUDA EP claimed the
    /// subgraph but split it into segments around non-capturable seam nodes.
    pub fn captured_graph_segment_count(&self) -> usize {
        self.exec.captured_segment_count()
    }

    /// Structured, transparent segment boundaries from the most recent capture:
    /// one entry per non-capturable seam node the EP ran eagerly between captured
    /// segments (with its structural seam kind and `CaptureSupport` decline reason).
    /// Empty for a whole-subgraph capture.
    pub fn capture_segmentation(&self) -> &[CaptureDecline] {
        self.exec.capture_segmentation()
    }

    /// Read (without clearing) any latching device capture-safety error recorded
    /// during graph replay, as a raw violation bitmask (zero when none). Callers
    /// poll this at the per-step logits sync to fail before consuming a token
    /// produced from an out-of-range captured replay.
    pub fn check_device_capture_error(&self) -> Result<u32> {
        self.exec.check_device_capture_error()
    }

    pub fn device_allocation_counts(&self) -> Option<DeviceAllocationCounts> {
        self.exec.device_allocation_counts()
    }

    pub fn device_id(&self) -> onnx_runtime_ir::DeviceId {
        self.exec.device_id()
    }

    /// Report why an explicitly requested accelerator session was assigned to
    /// CPU instead. `None` means the requested EP serves the whole graph.
    pub fn execution_provider_fallback_report(&self) -> Option<&ExecutionProviderFallbackReport> {
        self.exec.execution_provider_fallback_report()
    }

    /// Input metadata.
    pub fn inputs(&self) -> &[IoMeta] {
        &self.inputs
    }

    /// Output metadata.
    pub fn outputs(&self) -> &[IoMeta] {
        &self.outputs
    }

    /// Model-level metadata from the source `ModelProto`.
    pub fn model_metadata(&self) -> &ModelMetadata {
        &self.model_metadata
    }

    /// Kernel-cache statistics (§11.1); useful to observe warmup/run reuse.
    pub fn cache_stats(&self) -> CacheStats {
        self.exec.cache_stats()
    }

    /// Control-flow subgraph build/run statistics. A Loop or Scan body with a
    /// stable input-shape signature should build once and run many times.
    pub fn control_flow_stats(&self) -> ControlFlowStats {
        self.exec.control_flow_stats()
    }

    /// Pre-compile kernels for common shapes to avoid first-inference latency
    /// (§11.3). Phase-1 minimal: the compiled plan's shapes already key the
    /// cache, so this repopulates it for the plan; `shapes` are validated to
    /// name real inputs.
    pub fn warmup(&mut self, shapes: &[WarmupShape]) -> Result<()> {
        for ws in shapes {
            if !self.inputs.iter().any(|m| m.name == ws.input_name) {
                return Err(SessionError::InputNotFound {
                    name: ws.input_name.clone(),
                });
            }
        }
        self.exec.warmup()
    }

    /// The EPContext dump configuration parsed from the `ep.context_*` session
    /// options (§21.4). Disabled by default.
    pub fn ep_context_config(&self) -> &EpContextDumpConfig {
        &self.ep_context_config
    }

    /// The session's (post-optimize) compiled graph.
    ///
    /// This is the graph the executor runs and the same one
    /// [`Self::export_ep_context`] serialises — a caller identifying the
    /// [`NodeId`](onnx_runtime_ir::NodeId)s of a compiled partition (the
    /// [`CompiledPartition::covered_nodes`]) must read them from here so they
    /// reference the exact nodes the exporter will splice out. This is the
    /// compiler-integration seam: a real compiling EP inspects this graph to
    /// choose the subgraphs it claims.
    pub fn graph(&self) -> &onnx_runtime_ir::Graph {
        self.exec.graph()
    }

    /// Export a `com.microsoft::EPContext` context-cache model for this session
    /// (§55.4 dump path), driven by the `ep.context_*` session options
    /// ([`Self::ep_context_config`]).
    ///
    /// `orig_path` is the source model path the default output location
    /// (`<orig>_ctx.onnx`) is derived from when `ep.context_file_path` is unset.
    /// `partitions` are the EP-compiled partitions to serialise — each names the
    /// [`ExecutionProvider`](onnx_runtime_ep_api::ExecutionProvider) that
    /// compiled it, so the driver pulls the blob + SDK version via
    /// [`save_context`](onnx_runtime_ep_api::ExecutionProvider::save_context) and
    /// the `source` key via
    /// [`context_source_keys`](onnx_runtime_ep_api::ExecutionProvider::context_source_keys)
    /// (§55.6 — nothing is hardcoded).
    ///
    /// When `ep.context_enable` is `false` (the default) this is a **no-op**: no
    /// EP `save_context` is called and no files are written; it returns the path
    /// it *would* have written to.
    ///
    /// # Compiler-integration seam
    ///
    /// The Phase-1 CPU EP has **no compile step**, so no real EP yet yields
    /// [`CompiledPartition`]s — `partitions` is therefore supplied by the
    /// caller (proven end-to-end with a mock compiling EP in the crate tests).
    /// TODO(compiler): when a real compiling EP lands, collect its partitions
    /// from the compile/placement stage and call this internally at build time
    /// so a session created with `ep.context_enable=1` dumps automatically.
    pub fn export_ep_context(
        &self,
        orig_path: &Path,
        partitions: &[CompiledPartition],
    ) -> Result<PathBuf> {
        let model = EncoderModel::new(self.exec.graph()).with_weights(self.exec.weights().as_ref());
        dump_session_ep_context(&model, orig_path, partitions, &self.ep_context_config)
    }
}

/// Load a model. Auto-detects the best available hardware (§20.2).
///
/// This is the primary entry point — no configuration required.
pub fn load(path: impl AsRef<Path>) -> Result<InferenceSession> {
    InferenceSession::load(path)
}

#[cfg(test)]
mod device_binding_tests {
    use super::*;
    #[cfg(feature = "cuda")]
    use onnx_runtime_ir::Attribute;
    #[cfg(feature = "cuda")]
    use onnx_runtime_ir::static_shape;
    use onnx_runtime_ir::{Graph, Node, NodeId};

    #[test]
    fn persistent_binding_aliases_input_output_and_suppresses_materialization() {
        let mut graph = Graph::new();
        graph.opset_imports.insert("".into(), 13);
        let length = graph.intern_symbol("length");
        let input = graph.create_named_value("input", DataType::Float32, vec![length.into()]);
        graph.add_input(input);
        let output = graph.create_named_value("output", DataType::Float32, vec![length.into()]);
        graph.insert_node(Node::new(
            NodeId(0),
            "Relu",
            vec![Some(input)],
            vec![output],
        ));
        graph.add_output(output);
        let mut session = InferenceSession::from_graph(graph).unwrap();
        let mut binding = session
            .allocate_device_binding("input", Some("output"), DataType::Float32, vec![4], vec![2])
            .unwrap();
        let ptr = binding.device_ptr();
        let bytes = [-2.0f32, 3.0, -4.0, 5.0]
            .into_iter()
            .flat_map(f32::to_le_bytes)
            .collect::<Vec<_>>();
        binding.write_bytes(0, &bytes).unwrap();

        let outputs = session
            .run_with_device_bindings(&[], std::slice::from_mut(&mut binding))
            .unwrap();
        assert_eq!(outputs.len(), 1);
        assert!(outputs[0].is_none());
        assert_eq!(binding.device_ptr(), ptr);
        assert_eq!(binding.logical_shape(), &[2]);
        let values = binding
            .read_bytes()
            .unwrap()
            .chunks_exact(4)
            .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap()))
            .collect::<Vec<_>>();
        assert_eq!(values, vec![0.0, 3.0, -4.0, 5.0]);
        assert_eq!(
            binding.transfer_stats(),
            DeviceBindingTransferStats {
                host_upload_calls: 1,
                host_upload_bytes: 16,
                host_download_calls: 1,
                host_download_bytes: 16,
            }
        );
    }

    #[cfg(feature = "cuda")]
    #[test]
    fn cuda_graph_replay_uses_persistent_io_without_device_allocations() {
        let Ok(mut ep) = onnx_runtime_ep_cuda::CudaExecutionProvider::new(0) else {
            eprintln!("skipping session CUDA graph test: CUDA runtime unavailable");
            return;
        };
        onnx_runtime_ep_api::ExecutionProvider::initialize(&mut ep, &Default::default()).unwrap();

        let mut graph = Graph::new();
        graph.opset_imports.insert("".into(), 13);
        let input = graph.create_named_value("input", DataType::Int64, static_shape([1]));
        graph.add_input(input);
        let output = graph.create_named_value("output", DataType::Float32, static_shape([1]));
        let mut cast = Node::new(NodeId(0), "Cast", vec![Some(input)], vec![output]);
        cast.attributes
            .insert("to".into(), Attribute::Int(DataType::Float32 as i64));
        graph.insert_node(cast);
        graph.add_output(output);

        let mut session = InferenceSession::from_parts(
            graph,
            std::sync::Arc::new(onnx_runtime_loader::WeightStore::new()),
            Path::new("."),
            EpContextDumpConfig::default(),
            ModelMetadata::default(),
            std::sync::Arc::new(ep),
        )
        .unwrap();
        let mut input = session
            .allocate_device_binding("input", None::<String>, DataType::Int64, vec![1], vec![1])
            .unwrap();
        let output = session
            .allocate_device_output_binding("output", DataType::Float32, vec![1], vec![1])
            .unwrap();
        input.write_bytes(0, &7i64.to_le_bytes()).unwrap();
        let mut bindings = vec![input, output];
        session
            .run_with_device_bindings(&[], &mut bindings)
            .unwrap();

        input_write(&mut bindings[0], 11);
        assert!(matches!(
            session
                .try_capture_with_device_bindings(&[], &mut bindings)
                .unwrap(),
            DeviceGraphCaptureResult::Captured(_)
        ));
        assert_eq!(read_bound_f32(&mut bindings[1]), 11.0);

        let before = session.device_allocation_counts().unwrap();
        input_write(&mut bindings[0], 23);
        session.replay_device_graph(&mut bindings).unwrap();
        assert_eq!(read_bound_f32(&mut bindings[1]), 23.0);
        assert_eq!(session.device_allocation_counts().unwrap(), before);
        assert!(session.reset_device_graph().unwrap());

        input_write(&mut bindings[0], 31);
        assert!(matches!(
            session
                .try_capture_with_device_bindings(&[], &mut bindings)
                .unwrap(),
            DeviceGraphCaptureResult::Captured(_)
        ));
        assert_eq!(read_bound_f32(&mut bindings[1]), 31.0);
        input_write(&mut bindings[0], 47);
        session.replay_device_graph(&mut bindings).unwrap();
        assert_eq!(read_bound_f32(&mut bindings[1]), 47.0);
        assert!(session.reset_device_graph().unwrap());
    }

    #[cfg(feature = "cuda")]
    #[test]
    fn segmented_cuda_graph_claims_whole_subgraph_around_eager_seam() {
        let Ok(mut ep) = onnx_runtime_ep_cuda::CudaExecutionProvider::new(0) else {
            eprintln!("skipping segmented session CUDA graph test: CUDA runtime unavailable");
            return;
        };
        onnx_runtime_ep_api::ExecutionProvider::initialize(&mut ep, &Default::default()).unwrap();

        // A decoder-like chain with a deliberately non-capturable node in the
        // middle: input -> Cast(f32) -> Clip(min/max attrs) -> Cast(i64) -> out.
        // Cast is CUDA-graph capture-safe (it skips its trailing sync while the
        // stream is capturing); Clip declines capture, so it forces a segment
        // boundary while remaining CUDA-placed. Over integer inputs and a wide
        // clip the chain round-trips to the identity.
        let n = 4usize;
        let mut graph = Graph::new();
        graph.opset_imports.insert("".into(), 13);
        let input = graph.create_named_value("input", DataType::Int64, static_shape([n]));
        graph.add_input(input);
        let as_float = graph.create_named_value("as_float", DataType::Float32, static_shape([n]));
        let mut cast_in = Node::new(NodeId(0), "Cast", vec![Some(input)], vec![as_float]);
        cast_in
            .attributes
            .insert("to".into(), Attribute::Int(DataType::Float32 as i64));
        graph.insert_node(cast_in);
        let clipped = graph.create_named_value("clipped", DataType::Float32, static_shape([n]));
        let mut clip = Node::new(NodeId(1), "Clip", vec![Some(as_float)], vec![clipped]);
        clip.attributes
            .insert("min".into(), Attribute::Float(-1000.0));
        clip.attributes
            .insert("max".into(), Attribute::Float(1000.0));
        graph.insert_node(clip);
        let output = graph.create_named_value("output", DataType::Int64, static_shape([n]));
        let mut cast_out = Node::new(NodeId(2), "Cast", vec![Some(clipped)], vec![output]);
        cast_out
            .attributes
            .insert("to".into(), Attribute::Int(DataType::Int64 as i64));
        graph.insert_node(cast_out);
        graph.add_output(output);

        let mut session = InferenceSession::from_parts(
            graph,
            std::sync::Arc::new(onnx_runtime_loader::WeightStore::new()),
            Path::new("."),
            EpContextDumpConfig::default(),
            ModelMetadata::default(),
            std::sync::Arc::new(ep),
        )
        .unwrap();

        let input_binding = session
            .allocate_device_binding("input", None::<String>, DataType::Int64, vec![n], vec![n])
            .unwrap();
        let output_binding = session
            .allocate_device_output_binding("output", DataType::Int64, vec![n], vec![n])
            .unwrap();
        let mut bindings = vec![input_binding, output_binding];

        // Warmup / eager reference for input A (also warms the shape-keyed
        // kernels the capture pass requires).
        let values_a = [-2i64, 3, -4, 5];
        write_bound_i64(&mut bindings[0], &values_a);
        session
            .run_with_device_bindings(&[], &mut bindings)
            .unwrap();
        let eager_a = read_bound_i64_vec(&mut bindings[1]);
        assert_eq!(
            eager_a,
            values_a.to_vec(),
            "Cast∘Clip∘Cast round-trips ints"
        );

        // Segmented capture: the whole subgraph is still claimed and run on the
        // CUDA EP even though Clip is not capturable.
        match session
            .try_capture_with_device_bindings(&[], &mut bindings)
            .unwrap()
        {
            DeviceGraphCaptureResult::Captured(outputs) => {
                assert!(
                    outputs.iter().all(Option::is_none),
                    "device-bound outputs must not materialize to host"
                );
            }
            DeviceGraphCaptureResult::NotCapturable(report) => {
                panic!(
                    "expected the CUDA EP to claim the whole subgraph via segmented capture, \
                     got a full decline: {report}"
                );
            }
        }
        // Whole-subgraph claim, split into two captured segments around one seam.
        assert_eq!(
            session.captured_graph_segment_count(),
            2,
            "Clip should split the plan into two captured Cast segments"
        );
        let seams = session.capture_segmentation();
        assert_eq!(seams.len(), 1, "exactly one eager seam node (Clip)");
        assert_eq!(seams[0].op_type, "Clip");
        // Token-exact: the segmented capture pass matches the eager reference.
        assert_eq!(read_bound_i64_vec(&mut bindings[1]), eager_a);

        // Segmented replay for a new input B interleaves the two captured Cast
        // segment graphs with the eager Clip seam, and stays token-exact.
        let values_b = [7i64, -1, 0, -8];
        write_bound_i64(&mut bindings[0], &values_b);
        session.replay_device_graph(&mut bindings).unwrap();
        let replay_b = read_bound_i64_vec(&mut bindings[1]);

        // Independent eager reference for input B.
        assert!(session.reset_device_graph().unwrap());
        write_bound_i64(&mut bindings[0], &values_b);
        session
            .run_with_device_bindings(&[], &mut bindings)
            .unwrap();
        let eager_b = read_bound_i64_vec(&mut bindings[1]);
        assert_eq!(
            replay_b, eager_b,
            "segmented replay must be bit-identical to eager execution"
        );
        assert_eq!(replay_b, values_b.to_vec());
    }

    #[cfg(feature = "cuda")]
    fn write_bound_i64(binding: &mut DeviceIoBinding, values: &[i64]) {
        let bytes = values
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect::<Vec<_>>();
        binding.write_bytes(0, &bytes).unwrap();
    }

    #[cfg(feature = "cuda")]
    fn read_bound_i64_vec(binding: &mut DeviceIoBinding) -> Vec<i64> {
        binding
            .read_bytes()
            .unwrap()
            .chunks_exact(8)
            .map(|bytes| i64::from_le_bytes(bytes.try_into().unwrap()))
            .collect()
    }

    #[cfg(feature = "cuda")]
    fn input_write(binding: &mut DeviceIoBinding, value: i64) {
        binding.write_bytes(0, &value.to_le_bytes()).unwrap();
    }

    #[cfg(feature = "cuda")]
    fn read_bound_f32(binding: &mut DeviceIoBinding) -> f32 {
        let bytes = binding.read_bytes().unwrap();
        f32::from_le_bytes(bytes.try_into().unwrap())
    }
}

#[cfg(test)]
mod option_tests {
    use super::*;

    fn opts(pairs: &[(&str, &str)]) -> HashMap<String, String> {
        pairs
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect()
    }

    fn level_of(pairs: &[(&str, &str)]) -> Result<OptimizationLevel> {
        SessionBuilder::parse_options(&opts(pairs)).map(|(level, _)| level)
    }

    fn ctx_of(pairs: &[(&str, &str)]) -> Result<EpContextDumpConfig> {
        SessionBuilder::parse_options(&opts(pairs)).map(|(_, ctx)| ctx)
    }

    #[test]
    fn optimization_defaults_to_none_when_unset() {
        assert_eq!(level_of(&[]).unwrap(), OptimizationLevel::None);
    }

    #[test]
    fn explicit_execution_provider_is_retained_by_builder() {
        let builder =
            SessionBuilder::new().execution_provider(executor::auto_detect_cpu_ep().unwrap());
        assert!(builder.execution_provider.is_some());
    }

    #[test]
    fn optimization_parses_known_values() {
        for (v, want) in [
            ("none", OptimizationLevel::None),
            ("off", OptimizationLevel::None),
            ("BASIC", OptimizationLevel::Basic),
            ("All", OptimizationLevel::All),
        ] {
            assert_eq!(
                level_of(&[("optimization", v)]).unwrap(),
                want,
                "value {v:?}"
            );
        }
    }

    #[test]
    fn unknown_option_key_is_rejected() {
        let err = level_of(&[("optimisation", "all")]).unwrap_err();
        assert!(matches!(err, SessionError::UnknownOption { key } if key == "optimisation"));
    }

    #[test]
    fn invalid_optimization_value_is_rejected() {
        let err = level_of(&[("optimization", "aggressive")]).unwrap_err();
        assert!(matches!(
            err,
            SessionError::InvalidOption { key, value, .. } if key == "optimization" && value == "aggressive"
        ));
    }

    #[test]
    fn none_level_selects_no_passes() {
        assert!(OptimizationLevel::None.passes().is_empty());
        assert_eq!(OptimizationLevel::Basic.passes().len(), 2);
        assert_eq!(OptimizationLevel::All.passes().len(), 3);
    }

    // ── EPContext dump options (§21.4 / §55.5) ────────────────────────────────

    #[test]
    fn ep_context_defaults_to_disabled() {
        let ctx = ctx_of(&[]).unwrap();
        assert_eq!(ctx, EpContextDumpConfig::default());
        assert!(!ctx.enable);
        assert_eq!(ctx.file_path, None);
        assert_eq!(ctx.embed_mode, 1);
    }

    #[test]
    fn ep_context_enable_parses_bool_forms() {
        for (v, want) in [
            ("1", true),
            ("0", false),
            ("true", true),
            ("TRUE", true),
            ("false", false),
            ("False", false),
        ] {
            let ctx = ctx_of(&[("ep.context_enable", v)]).unwrap();
            assert_eq!(ctx.enable, want, "value {v:?}");
        }
    }

    #[test]
    fn ep_context_enable_rejects_garbage() {
        let err = ctx_of(&[("ep.context_enable", "yes")]).unwrap_err();
        assert!(matches!(
            err,
            SessionError::InvalidOption { key, value, .. }
                if key == "ep.context_enable" && value == "yes"
        ));
    }

    #[test]
    fn ep_context_file_path_parses_and_empty_clears() {
        let ctx = ctx_of(&[("ep.context_file_path", "/out/net_ctx.onnx")]).unwrap();
        assert_eq!(ctx.file_path, Some(PathBuf::from("/out/net_ctx.onnx")));

        // Empty value falls back to the `<orig>_ctx.onnx` default (None).
        let ctx = ctx_of(&[("ep.context_file_path", "")]).unwrap();
        assert_eq!(ctx.file_path, None);
    }

    #[test]
    fn ep_context_embed_mode_parses_and_rejects() {
        assert_eq!(
            ctx_of(&[("ep.context_embed_mode", "0")])
                .unwrap()
                .embed_mode,
            0
        );
        assert_eq!(
            ctx_of(&[("ep.context_embed_mode", "1")])
                .unwrap()
                .embed_mode,
            1
        );

        let err = ctx_of(&[("ep.context_embed_mode", "2")]).unwrap_err();
        assert!(matches!(
            err,
            SessionError::InvalidOption { key, value, expected }
                if key == "ep.context_embed_mode" && value == "2" && expected == "0, 1"
        ));
    }

    #[test]
    fn ep_context_options_combine_with_optimization() {
        let (level, ctx) = SessionBuilder::parse_options(&opts(&[
            ("optimization", "all"),
            ("ep.context_enable", "1"),
            ("ep.context_file_path", "/tmp/out_ctx.onnx"),
            ("ep.context_embed_mode", "0"),
        ]))
        .unwrap();
        assert_eq!(level, OptimizationLevel::All);
        assert!(ctx.enable);
        assert_eq!(ctx.file_path, Some(PathBuf::from("/tmp/out_ctx.onnx")));
        assert_eq!(ctx.embed_mode, 0);
    }
}