rustytorch/crates/core/rtx-fusion/src/lib.rs

//! # 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<B> Backend  │  ◄── Backend wrapper
//!          │  (FusionBackend)    │
//!          └──────────┬──────────┘
//!                     │
//!     ┌───────────────┼───────────────┐
//!     │               │               │
//!     ▼               ▼               ▼
//! OperationStream  FusedKernelCache  FusionAnalyzer
//! (queue ops)      (compiled cache)  (detect patterns)
//!     │               │               │
//!     └───────────────┼───────────────┘
//!                     │
//!          ┌──────────▼──────────┐
//!          │  CubeCL Codegen     │
//!          │  (generate kernels) │
//!          └──────────┬──────────┘
//!                     │
//!          ┌──────────▼──────────┐
//!          │  Inner Backend <B>  │
//!          │  (CudaBackend etc)  │
//!          └─────────────────────┘
//! ```
//!
//! ## Usage
//!
//! Wrap any backend with `Fusion<B>` to enable automatic kernel fusion:
//!
//! ```rust,ignore
//! use rtx_fusion::Fusion;
//! use rtx_backend_cuda::CudaBackend;
//!
//! // Create a fused backend
//! type FusedBackend = Fusion<CudaBackend>;
//!
//! // 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();
    }
}