rustytorch/crates/specialized/rtx-polygraph/src/ir.rs

use crate::error::{PolygraphError, Result};
use indexmap::IndexMap;
use rustc_hash::FxHasher;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::hash::{Hash, Hasher};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NodeId(pub u32);

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum DataType {
    F32,
    F64,
    I32,
    I64,
    C64,
    C128,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Shape {
    dims: Vec<usize>,
}

impl Shape {
    pub fn new(dims: Vec<usize>) -> Self {
        Self { dims }
    }

    pub fn dims(&self) -> &[usize] {
        &self.dims
    }

    pub fn ndim(&self) -> usize {
        self.dims.len()
    }

    pub fn size(&self) -> usize {
        self.dims.iter().product()
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum SparseFormat {
    CSR,
    CSC,
    COO,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum AggregationType {
    Sum,
    Mean,
    Max,
    Min,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum IRNodeType {
    MatMul {
        transpose_a: bool,
        transpose_b: bool,
    },
    SparseMatMul {
        format: SparseFormat,
    },
    GraphConvolution {
        aggregation: AggregationType,
    },
    FFT {
        inverse: bool,
        axes: Vec<usize>,
    },
    ElementwiseMul,
    Add,
    ReLU,
    Conditional {
        condition_id: NodeId,
    },
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct IRNode {
    id: NodeId,
    node_type: IRNodeType,
    inputs: Vec<NodeId>,
    output_dtypes: Vec<DataType>,
    output_shapes: Vec<Shape>,
}

impl IRNode {
    pub fn new(
        id: NodeId,
        node_type: IRNodeType,
        inputs: Vec<NodeId>,
        output_dtypes: Vec<DataType>,
        output_shapes: Vec<Shape>,
    ) -> Self {
        Self {
            id,
            node_type,
            inputs,
            output_dtypes,
            output_shapes,
        }
    }

    pub fn try_new(
        id: NodeId,
        node_type: IRNodeType,
        inputs: Vec<NodeId>,
        output_dtypes: Vec<DataType>,
        output_shapes: Vec<Shape>,
    ) -> Result<Self> {
        // Validate node configuration
        match &node_type {
            IRNodeType::MatMul { .. } | IRNodeType::SparseMatMul { .. } => {
                if inputs.len() < 2 {
                    return Err(PolygraphError::InvalidNodeConfiguration {
                        message: format!(
                            "MatMul operations require at least 2 inputs, got {}",
                            inputs.len()
                        ),
                    });
                }
            }
            IRNodeType::GraphConvolution { .. } => {
                if inputs.len() < 2 {
                    return Err(PolygraphError::InvalidNodeConfiguration {
                        message:
                            "GraphConvolution requires at least 2 inputs (features, adjacency)"
                                .to_string(),
                    });
                }
            }
            IRNodeType::FFT { .. } => {
                if inputs.is_empty() {
                    return Err(PolygraphError::InvalidNodeConfiguration {
                        message: "FFT requires at least 1 input".to_string(),
                    });
                }
            }
            IRNodeType::ElementwiseMul | IRNodeType::Add => {
                if inputs.is_empty() {
                    return Err(PolygraphError::InvalidNodeConfiguration {
                        message: "Elementwise operations require at least 1 input".to_string(),
                    });
                }
            }
            IRNodeType::ReLU => {
                if inputs.len() != 1 {
                    return Err(PolygraphError::InvalidNodeConfiguration {
                        message: format!("ReLU requires exactly 1 input, got {}", inputs.len()),
                    });
                }
            }
            IRNodeType::Conditional { .. } => {
                if inputs.len() < 2 {
                    return Err(PolygraphError::InvalidNodeConfiguration {
                        message: "Conditional requires at least 2 inputs".to_string(),
                    });
                }
            }
        }

        Ok(Self {
            id,
            node_type,
            inputs,
            output_dtypes,
            output_shapes,
        })
    }

    pub fn id(&self) -> NodeId {
        self.id
    }

    pub fn node_type(&self) -> &IRNodeType {
        &self.node_type
    }

    pub fn inputs(&self) -> &[NodeId] {
        &self.inputs
    }

    pub fn outputs(&self) -> &[DataType] {
        &self.output_dtypes
    }

    pub fn output_dtypes(&self) -> &[DataType] {
        &self.output_dtypes
    }

    pub fn output_shapes(&self) -> &[Shape] {
        &self.output_shapes
    }

    pub fn signature(&self) -> String {
        let mut hasher = FxHasher::default();

        // Hash node type, inputs structure, output types and shapes
        // But NOT the specific node IDs
        self.node_type.hash(&mut hasher);
        self.inputs.len().hash(&mut hasher);
        self.output_dtypes.hash(&mut hasher);

        for shape in &self.output_shapes {
            shape.dims.hash(&mut hasher);
        }

        format!("{:016x}", hasher.finish())
    }
}

impl Hash for IRNodeType {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match self {
            Self::MatMul {
                transpose_a,
                transpose_b,
            } => {
                0u8.hash(state);
                transpose_a.hash(state);
                transpose_b.hash(state);
            }
            Self::SparseMatMul { format } => {
                1u8.hash(state);
                format.hash(state);
            }
            Self::GraphConvolution { aggregation } => {
                2u8.hash(state);
                aggregation.hash(state);
            }
            Self::FFT { inverse, axes } => {
                3u8.hash(state);
                inverse.hash(state);
                axes.hash(state);
            }
            Self::ElementwiseMul => 4u8.hash(state),
            Self::Add => 5u8.hash(state),
            Self::ReLU => 6u8.hash(state),
            Self::Conditional { condition_id: _ } => {
                7u8.hash(state);
                // Don't hash the actual condition_id for signature purposes
            }
        }
    }
}

impl Hash for SparseFormat {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match self {
            Self::CSR => 0u8.hash(state),
            Self::CSC => 1u8.hash(state),
            Self::COO => 2u8.hash(state),
        }
    }
}

impl Hash for AggregationType {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match self {
            Self::Sum => 0u8.hash(state),
            Self::Mean => 1u8.hash(state),
            Self::Max => 2u8.hash(state),
            Self::Min => 3u8.hash(state),
        }
    }
}

impl Hash for DataType {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match self {
            Self::F32 => 0u8.hash(state),
            Self::F64 => 1u8.hash(state),
            Self::I32 => 2u8.hash(state),
            Self::I64 => 3u8.hash(state),
            Self::C64 => 4u8.hash(state),
            Self::C128 => 5u8.hash(state),
        }
    }
}

#[derive(Debug, Clone)]
pub struct IRGraph {
    nodes: IndexMap<NodeId, IRNode>,
    edges: HashMap<NodeId, Vec<NodeId>>, // adjacency list
    outputs: Vec<NodeId>,
}

impl IRGraph {
    pub fn new() -> Self {
        Self {
            nodes: IndexMap::new(),
            edges: HashMap::new(),
            outputs: Vec::new(),
        }
    }

    pub fn add_node(&mut self, node: IRNode) {
        let node_id = node.id();

        // Add edges for inputs
        for input_id in node.inputs() {
            self.edges.entry(*input_id).or_default().push(node_id);
        }

        self.nodes.insert(node_id, node);
    }

    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    pub fn contains_node(&self, id: NodeId) -> bool {
        self.nodes.contains_key(&id)
    }

    pub fn get_node(&self, id: NodeId) -> Option<&IRNode> {
        self.nodes.get(&id)
    }

    pub fn mark_output(&mut self, id: NodeId) {
        if !self.outputs.contains(&id) {
            self.outputs.push(id);
        }
    }

    pub fn output_shapes(&self) -> Vec<Vec<Shape>> {
        self.outputs
            .iter()
            .filter_map(|&id| self.get_node(id))
            .map(|node| node.output_shapes().to_vec())
            .collect()
    }

    pub fn output_dtypes(&self) -> Vec<Vec<DataType>> {
        self.outputs
            .iter()
            .filter_map(|&id| self.get_node(id))
            .map(|node| node.output_dtypes().to_vec())
            .collect()
    }

    pub fn estimated_memory_usage(&self) -> usize {
        self.nodes
            .values()
            .map(|node| {
                node.output_shapes()
                    .iter()
                    .zip(node.output_dtypes())
                    .map(|(shape, dtype)| {
                        let element_size = match dtype {
                            DataType::F32 | DataType::I32 => 4,
                            DataType::F64 | DataType::I64 | DataType::C64 => 8,
                            DataType::C128 => 16,
                        };
                        shape.size() * element_size
                    })
                    .sum::<usize>()
            })
            .sum()
    }

    pub fn nodes(&self) -> impl Iterator<Item = &IRNode> {
        self.nodes.values()
    }

    pub fn remove_node(&mut self, id: NodeId) {
        if let Some(node) = self.nodes.shift_remove(&id) {
            // Remove from edges
            for input_id in node.inputs() {
                if let Some(users) = self.edges.get_mut(input_id) {
                    users.retain(|&user_id| user_id != id);
                }
            }
            self.edges.remove(&id);

            // Remove from outputs if present
            self.outputs.retain(|&output_id| output_id != id);
        }
    }

    pub fn users(&self, id: NodeId) -> Option<&[NodeId]> {
        self.edges.get(&id).map(|v| v.as_slice())
    }
}

impl Default for IRGraph {
    fn default() -> Self {
        Self::new()
    }
}