rustytorch/crates/production/rtx-onnx/src/session.rs

//! ONNX Runtime session wrapper
//!
//! Provides a high-level interface for loading and running ONNX models.

use crate::error::{OnnxError, Result};
use crate::execution_provider::ExecutionProviderType;
use crate::tensor_bridge::{ort_to_rtx, rtx_to_ort};
use ort::session::Session;
use ort::session::builder::GraphOptimizationLevel as OrtGraphOptimizationLevel;
use ort::value::{DynValue, TensorElementType};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use tracing::{debug, info};

/// Configuration for creating an ONNX session
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OnnxSessionConfig {
    /// Execution provider to use
    pub execution_provider: ExecutionProviderType,
    /// Number of intra-op threads (0 = auto)
    pub intra_op_threads: Option<usize>,
    /// Number of inter-op threads (0 = auto)
    pub inter_op_threads: Option<usize>,
    /// Graph optimization level
    pub optimization_level: OptimizationLevel,
    /// Enable memory pattern optimization
    pub enable_memory_pattern: bool,
    /// Enable CPU memory arena
    pub enable_cpu_mem_arena: bool,
    /// Enable profiling
    pub enable_profiling: bool,
    /// Output device for tensors
    pub output_device: Device,
}

impl Default for OnnxSessionConfig {
    fn default() -> Self {
        Self {
            execution_provider: ExecutionProviderType::default(),
            intra_op_threads: None,
            inter_op_threads: None,
            optimization_level: OptimizationLevel::All,
            enable_memory_pattern: true,
            enable_cpu_mem_arena: true,
            enable_profiling: false,
            output_device: Device::Cpu,
        }
    }
}

/// Graph optimization level
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum OptimizationLevel {
    /// No optimizations
    Disabled,
    /// Basic optimizations
    Basic,
    /// Extended optimizations
    Extended,
    /// All optimizations
    All,
}

impl From<OptimizationLevel> for OrtGraphOptimizationLevel {
    fn from(level: OptimizationLevel) -> Self {
        match level {
            OptimizationLevel::Disabled => OrtGraphOptimizationLevel::Disable,
            OptimizationLevel::Basic => OrtGraphOptimizationLevel::Level1,
            OptimizationLevel::Extended => OrtGraphOptimizationLevel::Level2,
            OptimizationLevel::All => OrtGraphOptimizationLevel::Level3,
        }
    }
}

/// Model input/output metadata
#[derive(Debug, Clone)]
pub struct IoInfo {
    /// Name of the input/output
    pub name: String,
    /// Shape (None dimensions are dynamic)
    pub shape: Vec<Option<i64>>,
    /// Data type
    pub dtype: TensorElementType,
}

/// ONNX Runtime session wrapper
pub struct OnnxSession {
    /// Underlying ORT session
    session: Session,
    /// Input names
    input_names: Vec<String>,
    /// Output names
    output_names: Vec<String>,
    /// Configuration used to create the session
    config: OnnxSessionConfig,
}

impl OnnxSession {
    /// Load an ONNX model from a file path
    pub fn from_file(path: impl AsRef<Path>, config: OnnxSessionConfig) -> Result<Self> {
        let path = path.as_ref();
        info!("Loading ONNX model from: {}", path.display());

        if !path.exists() {
            return Err(OnnxError::ModelLoad(format!(
                "Model file not found: {}",
                path.display()
            )));
        }

        let session = Self::create_session(path, &config)?;
        Self::from_session(session, config)
    }

    /// Load an ONNX model from bytes
    pub fn from_bytes(model_bytes: &[u8], config: OnnxSessionConfig) -> Result<Self> {
        info!(
            "Loading ONNX model from bytes ({} bytes)",
            model_bytes.len()
        );

        let mut builder =
            Session::builder().map_err(|e| OnnxError::SessionCreation(e.to_string()))?;

        builder = builder
            .with_optimization_level(config.optimization_level.into())
            .map_err(|e| OnnxError::SessionCreation(e.to_string()))?;

        let session = builder
            .commit_from_memory(model_bytes)
            .map_err(|e| OnnxError::SessionCreation(e.to_string()))?;

        Self::from_session(session, config)
    }

    /// Create session from file with configuration
    fn create_session(path: &Path, config: &OnnxSessionConfig) -> Result<Session> {
        let mut builder =
            Session::builder().map_err(|e| OnnxError::SessionCreation(e.to_string()))?;

        // Set optimization level
        builder = builder
            .with_optimization_level(config.optimization_level.into())
            .map_err(|e| OnnxError::SessionCreation(e.to_string()))?;

        // Set thread counts if specified
        if let Some(threads) = config.intra_op_threads {
            builder = builder
                .with_intra_threads(threads)
                .map_err(|e| OnnxError::SessionCreation(e.to_string()))?;
        }

        if let Some(threads) = config.inter_op_threads {
            builder = builder
                .with_inter_threads(threads)
                .map_err(|e| OnnxError::SessionCreation(e.to_string()))?;
        }

        // Configure execution provider
        builder = Self::configure_execution_provider(builder, &config.execution_provider)?;

        builder
            .commit_from_file(path)
            .map_err(|e| OnnxError::SessionCreation(e.to_string()))
    }

    /// Configure the execution provider for the session
    fn configure_execution_provider(
        builder: ort::session::builder::SessionBuilder,
        provider: &ExecutionProviderType,
    ) -> Result<ort::session::builder::SessionBuilder> {
        match provider {
            ExecutionProviderType::CPU(_opts) => {
                debug!("Using CPU execution provider");
                // CPU is always available as fallback
                Ok(builder)
            }
            #[cfg(feature = "cuda")]
            ExecutionProviderType::CUDA(opts) => {
                use ort::execution_providers::CUDAExecutionProvider;
                debug!(
                    "Configuring CUDA execution provider (device {})",
                    opts.device_id
                );
                builder
                    .with_execution_providers([CUDAExecutionProvider::default()
                        .with_device_id(opts.device_id)
                        .build()])
                    .map_err(|e| OnnxError::ExecutionProvider(e.to_string()))
            }
            #[cfg(feature = "coreml")]
            ExecutionProviderType::CoreML(_opts) => {
                use ort::execution_providers::CoreMLExecutionProvider;
                debug!("Configuring CoreML execution provider");
                builder
                    .with_execution_providers([CoreMLExecutionProvider::default().build()])
                    .map_err(|e| OnnxError::ExecutionProvider(e.to_string()))
            }
            #[cfg(feature = "tensorrt")]
            ExecutionProviderType::TensorRT(opts) => {
                use ort::execution_providers::TensorRTExecutionProvider;
                debug!("Configuring TensorRT execution provider");
                builder
                    .with_execution_providers([TensorRTExecutionProvider::default()
                        .with_device_id(opts.device_id)
                        .build()])
                    .map_err(|e| OnnxError::ExecutionProvider(e.to_string()))
            }
            #[cfg(feature = "directml")]
            ExecutionProviderType::DirectML(opts) => {
                use ort::execution_providers::DirectMLExecutionProvider;
                debug!("Configuring DirectML execution provider");
                builder
                    .with_execution_providers([DirectMLExecutionProvider::default()
                        .with_device_id(opts.device_id)
                        .build()])
                    .map_err(|e| OnnxError::ExecutionProvider(e.to_string()))
            }
        }
    }

    /// Create OnnxSession from an existing ORT session
    fn from_session(session: Session, config: OnnxSessionConfig) -> Result<Self> {
        // Extract input names
        let input_names: Vec<String> = session
            .inputs()
            .iter()
            .map(|i| i.name().to_string())
            .collect();

        // Extract output names
        let output_names: Vec<String> = session
            .outputs()
            .iter()
            .map(|o| o.name().to_string())
            .collect();

        info!(
            "Loaded ONNX model with {} inputs, {} outputs",
            input_names.len(),
            output_names.len()
        );
        debug!("Inputs: {:?}", input_names);
        debug!("Outputs: {:?}", output_names);

        Ok(Self {
            session,
            input_names,
            output_names,
            config,
        })
    }

    /// Run inference with rtx tensors
    pub fn run(&mut self, inputs: HashMap<String, &Tensor>) -> Result<HashMap<String, Tensor>> {
        // Validate inputs
        for name in &self.input_names {
            if !inputs.contains_key(name) {
                return Err(OnnxError::MissingInput(name.clone()));
            }
        }

        // Convert inputs to ORT values
        let mut ort_inputs: Vec<(String, DynValue)> = Vec::new();
        for (name, tensor) in &inputs {
            let value = rtx_to_ort(tensor)?;
            ort_inputs.push((name.clone(), value));
        }

        // Run inference
        let ort_outputs = self
            .session
            .run(ort_inputs)
            .map_err(|e| OnnxError::Inference(e.to_string()))?;

        // Convert outputs to rtx tensors using into_iter to take ownership
        let mut result = HashMap::new();
        for (name, value) in ort_outputs {
            let tensor = ort_to_rtx(value, &self.config.output_device)?;
            result.insert(name.to_string(), tensor);
        }

        Ok(result)
    }

    /// Get input names
    pub fn input_names(&self) -> &[String] {
        &self.input_names
    }

    /// Get output names
    pub fn output_names(&self) -> &[String] {
        &self.output_names
    }

    /// Get the configuration used to create this session
    pub fn config(&self) -> &OnnxSessionConfig {
        &self.config
    }

    /// Get the execution provider name
    pub fn execution_provider_name(&self) -> &'static str {
        self.config.execution_provider.name()
    }
}

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

    #[test]
    fn test_default_config() {
        let config = OnnxSessionConfig::default();
        assert!(matches!(
            config.execution_provider,
            ExecutionProviderType::CPU(_)
        ));
        assert!(matches!(config.optimization_level, OptimizationLevel::All));
    }

    #[test]
    fn test_optimization_level_conversion() {
        assert!(matches!(
            OrtGraphOptimizationLevel::from(OptimizationLevel::Disabled),
            OrtGraphOptimizationLevel::Disable
        ));
        assert!(matches!(
            OrtGraphOptimizationLevel::from(OptimizationLevel::All),
            OrtGraphOptimizationLevel::Level3
        ));
    }
}