rustytorch/crates/production/rtx-onnx-codegen/src/ir/graph.rs

//! ONNX graph representation.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::error::{Error, Result};
use crate::ir::{Node, Tensor, ValueInfo};
use crate::onnx_proto::GraphProto;

/// An ONNX computation graph.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OnnxGraph {
    /// Graph name.
    pub name: String,

    /// ONNX opset version.
    pub opset_version: i64,

    /// Graph inputs.
    pub inputs: Vec<ValueInfo>,

    /// Graph outputs.
    pub outputs: Vec<ValueInfo>,

    /// Computation nodes (operations).
    pub nodes: Vec<Node>,

    /// Initializers (weights/constants).
    pub initializers: HashMap<String, Tensor>,

    /// Value info for intermediate tensors.
    pub value_info: HashMap<String, ValueInfo>,
}

impl OnnxGraph {
    /// Create from ONNX GraphProto.
    pub fn from_proto(proto: GraphProto, opset_version: i64) -> Result<Self> {
        let name = proto.name.clone();

        // Parse inputs
        let inputs: Vec<ValueInfo> = proto
            .input
            .iter()
            .filter(|i| !proto.initializer.iter().any(|init| init.name == i.name))
            .map(ValueInfo::from_proto)
            .collect::<Result<_>>()?;

        // Parse outputs
        let outputs: Vec<ValueInfo> = proto
            .output
            .iter()
            .map(ValueInfo::from_proto)
            .collect::<Result<_>>()?;

        // Parse nodes
        let nodes: Vec<Node> = proto
            .node
            .iter()
            .map(Node::from_proto)
            .collect::<Result<_>>()?;

        // Parse initializers (weights)
        let mut initializers = HashMap::new();
        for init in &proto.initializer {
            let tensor = Tensor::from_proto(init)?;
            initializers.insert(tensor.name.clone(), tensor);
        }

        // Parse value info
        let mut value_info = HashMap::new();
        for vi in &proto.value_info {
            let info = ValueInfo::from_proto(vi)?;
            value_info.insert(info.name.clone(), info);
        }

        log::debug!(
            "Graph '{}': {} inputs, {} outputs, {} nodes, {} initializers",
            name,
            inputs.len(),
            outputs.len(),
            nodes.len(),
            initializers.len()
        );

        Ok(Self {
            name,
            opset_version,
            inputs,
            outputs,
            nodes,
            initializers,
            value_info,
        })
    }

    /// Get input names.
    pub fn input_names(&self) -> Vec<&str> {
        self.inputs.iter().map(|i| i.name.as_str()).collect()
    }

    /// Get output names.
    pub fn output_names(&self) -> Vec<&str> {
        self.outputs.iter().map(|o| o.name.as_str()).collect()
    }

    /// Check if a tensor is an initializer (weight).
    pub fn is_initializer(&self, name: &str) -> bool {
        self.initializers.contains_key(name)
    }

    /// Get an initializer by name.
    pub fn get_initializer(&self, name: &str) -> Option<&Tensor> {
        self.initializers.get(name)
    }

    /// Validate the graph structure.
    pub fn validate(&self) -> Result<()> {
        // Check all inputs are defined
        let mut defined: std::collections::HashSet<&str> =
            self.inputs.iter().map(|i| i.name.as_str()).collect();

        // Add initializers
        for name in self.initializers.keys() {
            defined.insert(name);
        }

        // Check each node's inputs and track outputs
        for node in &self.nodes {
            for input in &node.inputs {
                if !input.is_empty() && !defined.contains(input.as_str()) {
                    return Err(Error::InvalidModel(format!(
                        "Node '{}' ({:?}) references undefined input '{}'",
                        node.name, node.kind, input
                    )));
                }
            }

            for output in &node.outputs {
                defined.insert(output);
            }
        }

        // Check all outputs are defined
        for output in &self.outputs {
            if !defined.contains(output.name.as_str()) {
                return Err(Error::InvalidModel(format!(
                    "Output '{}' is not defined",
                    output.name
                )));
            }
        }

        Ok(())
    }

    /// Get statistics about the graph.
    pub fn stats(&self) -> GraphStats {
        let mut op_counts = HashMap::new();
        let mut unsupported = Vec::new();

        for node in &self.nodes {
            *op_counts.entry(format!("{:?}", node.kind)).or_insert(0) += 1;

            if !node.kind.is_supported() {
                unsupported.push(format!("{:?}", node.kind));
            }
        }

        let total_params: usize = self
            .initializers
            .values()
            .map(super::tensor::Tensor::numel)
            .sum();

        GraphStats {
            num_inputs: self.inputs.len(),
            num_outputs: self.outputs.len(),
            num_nodes: self.nodes.len(),
            num_initializers: self.initializers.len(),
            total_parameters: total_params,
            op_counts,
            unsupported_ops: unsupported,
        }
    }
}

/// Statistics about an ONNX graph.
#[derive(Debug, Clone)]
pub struct GraphStats {
    /// Number of inputs.
    pub num_inputs: usize,
    /// Number of outputs.
    pub num_outputs: usize,
    /// Number of nodes.
    pub num_nodes: usize,
    /// Number of initializers.
    pub num_initializers: usize,
    /// Total number of parameters.
    pub total_parameters: usize,
    /// Count of each operator type.
    pub op_counts: HashMap<String, usize>,
    /// List of unsupported operators.
    pub unsupported_ops: Vec<String>,
}

impl std::fmt::Display for GraphStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Graph Statistics:")?;
        writeln!(f, "  Inputs: {}", self.num_inputs)?;
        writeln!(f, "  Outputs: {}", self.num_outputs)?;
        writeln!(f, "  Nodes: {}", self.num_nodes)?;
        writeln!(f, "  Initializers: {}", self.num_initializers)?;
        writeln!(f, "  Total Parameters: {}", self.total_parameters)?;

        if !self.unsupported_ops.is_empty() {
            writeln!(f, "  Unsupported ops: {:?}", self.unsupported_ops)?;
        }

        Ok(())
    }
}