77 lines
2.5 KiB
Rust
77 lines
2.5 KiB
Rust
//! # RTX-Polygraph: Unified IR and Super-Fusion
|
|
//!
|
|
//! RTX-Polygraph provides a unified intermediate representation (IR) for cross-domain
|
|
//! tensor operations and intelligent fusion across different computation domains including:
|
|
//! - Dense tensor operations (BLAS-like)
|
|
//! - Sparse tensor operations
|
|
//! - Graph neural network operations
|
|
//! - FFT and signal processing operations
|
|
//! - Control flow operations
|
|
//!
|
|
//! The crate enables aggressive fusion opportunities that can achieve ≥25% step-time
|
|
//! reduction through intelligent cross-domain kernel fusion and caching.
|
|
//!
|
|
//! ## Core Features
|
|
//!
|
|
//! - **Unified IR**: Single representation for heterogeneous operations
|
|
//! - **Cross-Domain Fusion**: Fuse operations across different domains (e.g., GNN + Dense)
|
|
//! - **Intelligent Caching**: Fast kernel cache with signature-based keys
|
|
//! - **Optimization Passes**: Dead code elimination, memory optimization, fusion passes
|
|
//! - **Type Safety**: Comprehensive error handling and validation
|
|
//!
|
|
//! ## Example Usage
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_polygraph::{
|
|
//! ir::{IRNode, IRNodeType, NodeId, DataType, Shape},
|
|
//! fusion::FusionAnalyzer,
|
|
//! cache::KernelCache,
|
|
//! };
|
|
//!
|
|
//! // Create a computation graph with fusible operations
|
|
//! let mut analyzer = FusionAnalyzer::new();
|
|
//!
|
|
//! let matmul1 = IRNode::new(
|
|
//! NodeId(1),
|
|
//! IRNodeType::MatMul { transpose_a: false, transpose_b: false },
|
|
//! vec![NodeId(0)],
|
|
//! vec![DataType::F32],
|
|
//! vec![Shape::new(vec![64, 128])],
|
|
//! );
|
|
//!
|
|
//! let matmul2 = IRNode::new(
|
|
//! NodeId(2),
|
|
//! IRNodeType::MatMul { transpose_a: false, transpose_b: false },
|
|
//! vec![NodeId(1)],
|
|
//! vec![DataType::F32],
|
|
//! vec![Shape::new(vec![64, 256])],
|
|
//! );
|
|
//!
|
|
//! analyzer.add_node(matmul1).unwrap();
|
|
//! analyzer.add_node(matmul2).unwrap();
|
|
//!
|
|
//! // Find fusion opportunities
|
|
//! let opportunities = analyzer.find_fusion_opportunities();
|
|
//! assert!(!opportunities.is_empty());
|
|
//! ```
|
|
|
|
pub mod cache;
|
|
pub mod error;
|
|
pub mod fusion;
|
|
pub mod ir;
|
|
pub mod ops;
|
|
pub mod optimizer;
|
|
|
|
pub use error::{PolygraphError, Result};
|
|
|
|
// Re-export GraphError alias for backwards compatibility
|
|
pub use error::PolygraphError as GraphError;
|
|
|
|
// Re-export commonly used types
|
|
pub use cache::{CacheKey, CachedKernel, KernelCache};
|
|
pub use fusion::{FusionAnalyzer, FusionOpportunity, FusionType};
|
|
pub use ir::{DataType, IRGraph, IRNode, IRNodeType, NodeId, Shape};
|
|
pub use optimizer::{
|
|
DeadCodeEliminationPass, FusionPass, MemoryOptimizationPass, OptimizationPass,
|
|
};
|