//! # RTX Fusion - Automatic Kernel Fusion for RustyTorch //! //! This crate provides stream-based automatic kernel fusion for the RustyTorch //! deep learning framework. It intercepts tensor operations, queues them in a //! stream, detects fusion opportunities, and executes optimized fused kernels. //! //! ## Overview //! //! Kernel fusion is a critical optimization for GPU computing. Instead of launching //! separate kernels for each operation (e.g., `add`, `mul`, `relu`), fusion combines //! them into a single kernel that executes in one launch. This provides: //! //! - **Reduced kernel launch overhead**: Each kernel launch has fixed overhead (~5-20μs) //! - **Improved memory bandwidth**: Intermediate results stay in registers/L1 cache //! - **Better GPU utilization**: Fused kernels can achieve higher occupancy //! //! ## Expected Impact //! //! - 10-25% training speedup via reduced kernel launches //! - 40% memory bandwidth savings for elementwise chains //! - 60% reduction in kernel launches for typical transformer blocks //! //! ## Architecture //! //! ```text //! User Code: tensor.add(b).mul(c).relu() //! │ //! ┌──────────▼──────────┐ //! │ Fusion Backend │ ◄── Backend wrapper //! │ (FusionBackend) │ //! └──────────┬──────────┘ //! │ //! ┌───────────────┼───────────────┐ //! │ │ │ //! ▼ ▼ ▼ //! OperationStream FusedKernelCache FusionAnalyzer //! (queue ops) (compiled cache) (detect patterns) //! │ │ │ //! └───────────────┼───────────────┘ //! │ //! ┌──────────▼──────────┐ //! │ CubeCL Codegen │ //! │ (generate kernels) │ //! └──────────┬──────────┘ //! │ //! ┌──────────▼──────────┐ //! │ Inner Backend │ //! │ (CudaBackend etc) │ //! └─────────────────────┘ //! ``` //! //! ## Usage //! //! Wrap any backend with `Fusion` to enable automatic kernel fusion: //! //! ```rust,ignore //! use rtx_fusion::Fusion; //! use rtx_backend_cuda::CudaBackend; //! //! // Create a fused backend //! type FusedBackend = Fusion; //! //! // Use it like any other backend //! let a = FusedBackend::rand([1024, 1024], &device); //! let b = FusedBackend::rand([1024, 1024], &device); //! //! // These operations will be fused into a single kernel //! let c = FusedBackend::add(a.clone(), b); // Queued //! let d = FusedBackend::mul(c, a); // Queued //! let e = FusedBackend::relu(d); // Queued //! //! // Sync triggers execution of the fused kernel //! FusedBackend::sync(&device); //! //! // Check fusion statistics //! let stats = FusedBackend::fusion_stats(); //! println!("Fused executions: {}", stats.fused_executions); //! println!("Kernel launches saved: {}", stats.kernel_launches_saved); //! ``` //! //! ## Fusion Patterns //! //! The analyzer detects several fusion patterns: //! //! - **Elementwise chains**: `add → mul → relu` fused into one kernel //! - **Activation chains**: Operations surrounding GELU/SiLU //! - **Softmax pattern**: `exp → sum → div` (partial fusion) //! - **Layer norm pattern**: Mean/variance computation (partial fusion) //! //! ## Sync Points //! //! Certain operations trigger immediate execution (flush): //! //! - Matrix operations: `matmul`, `bmm`, `conv2d` //! - Reductions: `sum`, `mean`, `max`, `min` //! - Special ops: `flash_attention`, `layer_norm`, `softmax` //! - Data access: `to_data()`, `to_device()` //! - Explicit sync: `sync()` //! //! ## Configuration //! //! Customize fusion behavior via `FusionConfig`: //! //! ```rust,ignore //! use rtx_fusion::FusionConfig; //! //! let config = FusionConfig::default() //! .with_enabled(true) //! .with_min_fusion_ops(2) //! .with_max_fusion_ops(16) //! .with_cache_size(1024); //! ``` #![warn(missing_docs)] mod analyzer; mod backend; mod cache; pub mod codegen; mod config; pub mod cuda_kernels; mod kernel; pub mod runtime_compat; mod stream; mod tensor; // Re-export main types pub use analyzer::{FusionAnalyzer, FusionOpportunity, FusionPattern, OperationAnalysis}; pub use backend::Fusion; pub use cache::{CacheStats, CachedKernel, FusedKernelCache, GlobalKernelCache}; pub use config::{FusionConfig, FusionStats, FusionStatsSnapshot}; pub use cuda_kernels::rms_norm_fused::{rms_norm_cpu, rms_norm_swiglu_cpu, swiglu_cpu}; pub use kernel::{ DType, F32Bits, FusedKernel, KernelSignature, StreamOpKind, StreamOperation, TensorId, }; pub use stream::{DeviceStreams, OperationStream, TensorDependency}; pub use tensor::{FusionTensor, Materializable, MaterializationHandle, TensorState}; #[cfg(feature = "cuda")] pub use cuda_kernels::rms_norm_fused::RmsNormFusedKernel; /// Prelude module for convenient imports pub mod prelude { pub use crate::backend::Fusion; pub use crate::config::{FusionConfig, FusionStatsSnapshot}; pub use crate::kernel::DType; } #[cfg(test)] mod tests; #[cfg(test)] mod lib_tests { use super::*; #[test] fn test_crate_compiles() { // Basic sanity check that the crate compiles let config = FusionConfig::default(); assert!(config.enabled); } #[test] fn test_prelude_imports() { // Verify prelude exports work use crate::prelude::*; let _ = FusionConfig::default(); } }