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

//! Intermediate representation for GPU kernels

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};

/// Unique identifier for IR nodes
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NodeId(u64);

static NEXT_NODE_ID: AtomicU64 = AtomicU64::new(1);

impl NodeId {
    /// Create a new unique node ID
    pub fn new(id: u64) -> Self {
        Self(id)
    }

    /// Generate a new unique node ID
    pub fn generate() -> Self {
        Self(NEXT_NODE_ID.fetch_add(1, Ordering::Relaxed))
    }

    /// Get the raw ID value
    pub fn id(&self) -> u64 {
        self.0
    }
}

/// Operation types in the IR
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum OperationType {
    // Basic operations
    Add,
    Sub,
    Mul,
    Div,

    // Memory operations
    Load,
    Store,

    // Control flow
    Branch,
    Loop,

    // Function operations
    Call,
    Return,

    // GPU-specific operations
    ThreadIdx,
    BlockIdx,
    BlockDim,
    GridDim,

    // Optimization-related operations
    Constant,
    Nop,

    // Extended operations for testing
    FusedConvBnRelu,
    Quantize,
    Dequantize,
    InstanceNorm1D,

    // Additional operations for comprehensive testing
    Conv2D,
    ReLU,
    BatchNorm,
    Identity,
}

/// IR Graph node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IRNode {
    pub id: NodeId,
    pub operation: OperationType,
    pub inputs: Vec<NodeId>,
    pub outputs: Vec<NodeId>,
    pub metadata: HashMap<String, String>,
}

impl IRNode {
    /// Create a new IR node
    pub fn new(operation: OperationType) -> Self {
        Self {
            id: NodeId::generate(),
            operation,
            inputs: Vec::new(),
            outputs: Vec::new(),
            metadata: HashMap::new(),
        }
    }

    /// Create an IR node with specific ID
    pub fn with_id(id: NodeId, operation: OperationType) -> Self {
        Self {
            id,
            operation,
            inputs: Vec::new(),
            outputs: Vec::new(),
            metadata: HashMap::new(),
        }
    }
}

/// IR Graph structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IRGraph {
    pub name: String,
    pub nodes: HashMap<NodeId, IRNode>,
    pub edges: Vec<(NodeId, NodeId)>,
    pub inputs: Vec<NodeId>,
    pub outputs: Vec<NodeId>,
}

impl IRGraph {
    /// Create a new empty IR graph
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            nodes: HashMap::new(),
            edges: Vec::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
        }
    }

    /// Add a node to the graph
    pub fn add_node(&mut self, node: IRNode) -> NodeId {
        let id = node.id;
        self.nodes.insert(id, node);
        id
    }

    /// Add an edge between two nodes
    pub fn add_edge(&mut self, from: NodeId, to: NodeId) {
        self.edges.push((from, to));
    }

    /// Get the number of nodes in the graph
    pub fn num_nodes(&self) -> usize {
        self.nodes.len()
    }

    /// Get the number of edges in the graph
    pub fn num_edges(&self) -> usize {
        self.edges.len()
    }

    /// Remove a node from the graph
    pub fn remove_node(&mut self, node_id: NodeId) {
        self.nodes.remove(&node_id);
        self.edges
            .retain(|(from, to)| *from != node_id && *to != node_id);
    }

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

    /// Get mutable node by ID
    pub fn get_node_mut(&mut self, id: NodeId) -> Option<&mut IRNode> {
        self.nodes.get_mut(&id)
    }

    /// Validate the graph structure
    pub fn validate(&self) -> Result<(), String> {
        // Check that all edges reference valid nodes
        for (from, to) in &self.edges {
            if !self.nodes.contains_key(from) {
                return Err(format!(
                    "Edge references non-existent source node: {from:?}"
                ));
            }
            if !self.nodes.contains_key(to) {
                return Err(format!("Edge references non-existent target node: {to:?}"));
            }
        }

        // Check that input/output nodes exist
        for input_id in &self.inputs {
            if !self.nodes.contains_key(input_id) {
                return Err(format!("Input references non-existent node: {input_id:?}"));
            }
        }

        for output_id in &self.outputs {
            if !self.nodes.contains_key(output_id) {
                return Err(format!(
                    "Output references non-existent node: {output_id:?}"
                ));
            }
        }

        Ok(())
    }

    /// Clone the graph with a new name
    pub fn clone_with_name(&self, new_name: &str) -> Self {
        let mut cloned = self.clone();
        cloned.name = new_name.to_string();
        cloned
    }

    /// Add an operation to the graph with attributes
    pub fn add_operation(
        &mut self,
        name: &str,
        operation: OperationType,
        inputs: &[NodeId],
        attributes: HashMap<String, AttributeValue>,
    ) -> Result<NodeId, String> {
        let node_id = NodeId::generate();
        let mut node = IRNode::with_id(node_id, operation);
        node.inputs = inputs.to_vec();

        // Store attributes in metadata
        for (key, value) in attributes {
            node.metadata.insert(key, value.to_string());
        }

        // Store name
        node.metadata.insert("name".to_string(), name.to_string());

        self.add_node(node);
        Ok(node_id)
    }

    /// Add an input to the graph
    pub fn add_input(&mut self, name: &str) -> Result<NodeId, String> {
        let input_id = self.add_operation(name, OperationType::Identity, &[], HashMap::new())?;
        self.inputs.push(input_id);
        Ok(input_id)
    }

    /// Add an output to the graph
    pub fn add_output(&mut self, name: &str, node_id: NodeId) -> Result<(), String> {
        if !self.nodes.contains_key(&node_id) {
            return Err(format!("Node {node_id:?} does not exist in graph"));
        }

        if let Some(node) = self.get_node_mut(node_id) {
            node.metadata
                .insert("output_name".to_string(), name.to_string());
        }

        self.outputs.push(node_id);
        Ok(())
    }
}

/// Attribute values for graph operations
#[derive(Debug, Clone)]
pub enum AttributeValue {
    Int(i64),
    Float(f64),
    String(String),
    Bool(bool),
    IntArray(Vec<i64>),
    FloatArray(Vec<f64>),
}

impl ToString for AttributeValue {
    fn to_string(&self) -> String {
        match self {
            Self::Int(v) => v.to_string(),
            Self::Float(v) => v.to_string(),
            Self::String(v) => v.clone(),
            Self::Bool(v) => v.to_string(),
            Self::IntArray(v) => format!("{v:?}"),
            Self::FloatArray(v) => format!("{v:?}"),
        }
    }
}

/// IR node types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum IrNode {
    /// Function definition
    Function {
        name: String,
        params: Vec<IrType>,
        body: Box<Self>,
    },
    /// Binary operation
    BinaryOp {
        op: BinaryOperator,
        left: Box<Self>,
        right: Box<Self>,
    },
    /// Memory load
    Load { address: Box<Self>, ty: IrType },
    /// Memory store
    Store {
        address: Box<Self>,
        value: Box<Self>,
    },
    /// Constant value
    Constant(IrValue),
    /// Thread index
    ThreadIdx(Dimension),
    /// Block index
    BlockIdx(Dimension),
    /// Block dimension
    BlockDim(Dimension),
    /// Grid dimension
    GridDim(Dimension),
}

/// Binary operators
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum BinaryOperator {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    And,
    Or,
    Xor,
    Shl,
    Shr,
}

/// IR types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum IrType {
    I8,
    I16,
    I32,
    I64,
    U8,
    U16,
    U32,
    U64,
    F16,
    F32,
    F64,
    Ptr(Box<Self>),
}

/// IR values
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum IrValue {
    I32(i32),
    I64(i64),
    F32(f32),
    F64(f64),
}

/// Thread/block dimensions
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum Dimension {
    X,
    Y,
    Z,
}

/// IR builder for constructing kernel IR
pub struct IrBuilder {
    nodes: Vec<IrNode>,
}

impl IrBuilder {
    /// Create new IR builder
    pub fn new() -> Self {
        Self { nodes: Vec::new() }
    }

    /// Add a node to the IR
    pub fn add_node(&mut self, node: IrNode) {
        self.nodes.push(node);
    }

    /// Build the IR
    pub fn build(self) -> Vec<IrNode> {
        self.nodes
    }
}

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