Initial commit

crates/specialized/rtx-synthesis/src/aot_impl/fusion_patterns.rs

@@ -0,0 +1,685 @@
+//! Fusion Pattern Library
+//!
+//! Provides a comprehensive library of fusion patterns optimized for
+//! different model architectures (Transformer, Mamba, MoE, etc.)
+
+use std::collections::HashMap;
+
+/// Standard fusion patterns for common neural network architectures
+#[derive(Debug, Clone)]
+pub struct FusionPatternLibrary {
+    /// Patterns organized by architecture type
+    patterns: HashMap<Architecture, Vec<FusionPattern>>,
+}
+
+/// Neural network architecture types
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum Architecture {
+    /// Standard Transformer (BERT, GPT-2 style)
+    Transformer,
+    /// LLaMA-style Transformer (RMSNorm, SwiGLU, RoPE)
+    LLaMA,
+    /// Mamba/State Space Model
+    Mamba,
+    /// Mixture of Experts
+    MoE,
+    /// Vision Transformer
+    ViT,
+    /// Convolutional Neural Network
+    CNN,
+    /// Generic/Unknown
+    Generic,
+}
+
+/// A fusion pattern with metadata
+#[derive(Debug, Clone)]
+pub struct FusionPattern {
+    /// Name of the pattern
+    pub name: String,
+    /// Description of what this pattern fuses
+    pub description: String,
+    /// Sequence of operations to match
+    pub operations: Vec<PatternOp>,
+    /// Metal kernel name to use
+    pub metal_kernel: String,
+    /// Estimated performance improvement (1.0 = no improvement)
+    pub speedup_factor: f32,
+    /// Memory bandwidth reduction factor
+    pub memory_reduction: f32,
+    /// Whether this pattern is enabled by default
+    pub enabled: bool,
+}
+
+/// Operations in a pattern (simplified for matching)
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum PatternOp {
+    /// Matrix multiplication
+    MatMul,
+    /// Bias addition
+    BiasAdd,
+    /// ReLU activation
+    ReLU,
+    /// GeLU activation
+    GeLU,
+    /// SiLU/Swish activation
+    SiLU,
+    /// Residual connection
+    Residual,
+    /// Layer normalization
+    LayerNorm,
+    /// RMS normalization
+    RMSNorm,
+    /// Elementwise add
+    Add,
+    /// Elementwise multiply
+    Mul,
+    /// Softmax
+    Softmax,
+    /// Rotary position embedding
+    RoPE,
+    /// SwiGLU gated MLP
+    SwiGLU,
+    /// GeGLU gated MLP
+    GeGLU,
+    /// Causal mask application
+    CausalMask,
+    /// Expert routing (MoE)
+    ExpertRoute,
+    /// Selective scan (Mamba)
+    SelectiveScan,
+    /// Convolution
+    Conv,
+    /// Batch normalization
+    BatchNorm,
+    /// Any activation (wildcard)
+    AnyActivation,
+}
+
+impl FusionPatternLibrary {
+    /// Create a new pattern library with all standard patterns
+    pub fn new() -> Self {
+        let mut patterns = HashMap::new();
+
+        // Transformer patterns
+        patterns.insert(Architecture::Transformer, Self::transformer_patterns());
+
+        // LLaMA patterns
+        patterns.insert(Architecture::LLaMA, Self::llama_patterns());
+
+        // Mamba patterns
+        patterns.insert(Architecture::Mamba, Self::mamba_patterns());
+
+        // MoE patterns
+        patterns.insert(Architecture::MoE, Self::moe_patterns());
+
+        // CNN patterns
+        patterns.insert(Architecture::CNN, Self::cnn_patterns());
+
+        // Generic patterns (always applicable)
+        patterns.insert(Architecture::Generic, Self::generic_patterns());
+
+        Self { patterns }
+    }
+
+    /// Get patterns for a specific architecture
+    pub fn patterns_for(&self, arch: Architecture) -> Vec<&FusionPattern> {
+        let mut result = Vec::new();
+
+        // Always include generic patterns
+        if let Some(generic) = self.patterns.get(&Architecture::Generic) {
+            result.extend(generic.iter());
+        }
+
+        // Add architecture-specific patterns
+        if arch != Architecture::Generic {
+            if let Some(specific) = self.patterns.get(&arch) {
+                result.extend(specific.iter());
+            }
+        }
+
+        result
+    }
+
+    /// Get all patterns
+    pub fn all_patterns(&self) -> impl Iterator<Item = &FusionPattern> {
+        self.patterns.values().flatten()
+    }
+
+    /// Transformer-specific fusion patterns
+    fn transformer_patterns() -> Vec<FusionPattern> {
+        vec![
+            FusionPattern {
+                name: "attention_qkv_projection".to_string(),
+                description: "Fused Q, K, V projections into single GEMM".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::MatMul, PatternOp::MatMul],
+                metal_kernel: "qkv_projection_fused_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.33,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "attention_scores_softmax".to_string(),
+                description: "Fused QK^T / sqrt(d) + causal mask + softmax".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::Softmax],
+                metal_kernel: "attention_scores_causal_f32".to_string(),
+                speedup_factor: 1.5,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "ffn_gelu".to_string(),
+                description: "Fused FFN with GeLU activation".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::GeLU],
+                metal_kernel: "gemm_bias_gelu_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "residual_layernorm".to_string(),
+                description: "Fused residual add + layer normalization".to_string(),
+                operations: vec![PatternOp::Residual, PatternOp::LayerNorm],
+                metal_kernel: "residual_layernorm_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+        ]
+    }
+
+    /// LLaMA-specific fusion patterns
+    fn llama_patterns() -> Vec<FusionPattern> {
+        vec![
+            FusionPattern {
+                name: "llama_rmsnorm".to_string(),
+                description: "Optimized RMS normalization for LLaMA".to_string(),
+                operations: vec![PatternOp::RMSNorm],
+                metal_kernel: "rmsnorm_f32".to_string(),
+                speedup_factor: 1.2,
+                memory_reduction: 0.0,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "llama_residual_rmsnorm".to_string(),
+                description: "Fused residual + RMS normalization".to_string(),
+                operations: vec![PatternOp::Residual, PatternOp::RMSNorm],
+                metal_kernel: "residual_rmsnorm_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "llama_swiglu".to_string(),
+                description: "Fused SwiGLU gated MLP (gate * up projections)".to_string(),
+                operations: vec![PatternOp::SwiGLU],
+                metal_kernel: "swiglu_fused_f32".to_string(),
+                speedup_factor: 1.6,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "llama_rope".to_string(),
+                description: "Fused rotary position embedding".to_string(),
+                operations: vec![PatternOp::RoPE],
+                metal_kernel: "rope_fused_f32".to_string(),
+                speedup_factor: 1.2,
+                memory_reduction: 0.0,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "llama_attention_rope".to_string(),
+                description: "Fused QK projection + RoPE".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::RoPE],
+                metal_kernel: "qk_rope_fused_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.33,
+                enabled: true,
+            },
+        ]
+    }
+
+    /// Mamba-specific fusion patterns
+    fn mamba_patterns() -> Vec<FusionPattern> {
+        vec![
+            FusionPattern {
+                name: "mamba_selective_scan".to_string(),
+                description: "Optimized selective scan kernel".to_string(),
+                operations: vec![PatternOp::SelectiveScan],
+                metal_kernel: "mamba_selective_scan_f32".to_string(),
+                speedup_factor: 2.0,
+                memory_reduction: 0.0,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "mamba_conv_silu".to_string(),
+                description: "Fused causal conv1d + SiLU".to_string(),
+                operations: vec![PatternOp::Conv, PatternOp::SiLU],
+                metal_kernel: "mamba_causal_conv_silu_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "mamba_ssm_update".to_string(),
+                description: "Fused SSM state update".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::Add, PatternOp::Mul],
+                metal_kernel: "mamba_ssm_update_f32".to_string(),
+                speedup_factor: 1.5,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+        ]
+    }
+
+    /// MoE-specific fusion patterns
+    fn moe_patterns() -> Vec<FusionPattern> {
+        vec![
+            FusionPattern {
+                name: "moe_topk_gating".to_string(),
+                description: "Fused top-k expert selection with softmax".to_string(),
+                operations: vec![
+                    PatternOp::MatMul,
+                    PatternOp::Softmax,
+                    PatternOp::ExpertRoute,
+                ],
+                metal_kernel: "moe_topk_gating_f32".to_string(),
+                speedup_factor: 1.5,
+                memory_reduction: 0.33,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "moe_expert_gemm".to_string(),
+                description: "Batched GEMM for expert computation".to_string(),
+                operations: vec![PatternOp::ExpertRoute, PatternOp::MatMul],
+                metal_kernel: "moe_expert_gemm_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.25,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "moe_combine_experts".to_string(),
+                description: "Fused weighted expert combination".to_string(),
+                operations: vec![PatternOp::Mul, PatternOp::Add],
+                metal_kernel: "moe_combine_experts_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+        ]
+    }
+
+    /// CNN-specific fusion patterns
+    fn cnn_patterns() -> Vec<FusionPattern> {
+        vec![
+            FusionPattern {
+                name: "conv_bn_relu".to_string(),
+                description: "Fused convolution + batch norm + ReLU".to_string(),
+                operations: vec![PatternOp::Conv, PatternOp::BatchNorm, PatternOp::ReLU],
+                metal_kernel: "conv_bn_relu_f32".to_string(),
+                speedup_factor: 1.6,
+                memory_reduction: 0.66,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "conv_relu".to_string(),
+                description: "Fused convolution + ReLU".to_string(),
+                operations: vec![PatternOp::Conv, PatternOp::ReLU],
+                metal_kernel: "conv_relu_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+        ]
+    }
+
+    /// Generic fusion patterns (applicable to any architecture)
+    fn generic_patterns() -> Vec<FusionPattern> {
+        vec![
+            FusionPattern {
+                name: "gemm_relu".to_string(),
+                description: "Fused GEMM + ReLU".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::ReLU],
+                metal_kernel: "gemm_relu_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "gemm_gelu".to_string(),
+                description: "Fused GEMM + GeLU".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::GeLU],
+                metal_kernel: "gemm_gelu_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "gemm_silu".to_string(),
+                description: "Fused GEMM + SiLU".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::SiLU],
+                metal_kernel: "gemm_silu_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "gemm_bias_relu".to_string(),
+                description: "Fused GEMM + bias + ReLU".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::ReLU],
+                metal_kernel: "gemm_bias_relu_f32".to_string(),
+                speedup_factor: 1.5,
+                memory_reduction: 0.66,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "gemm_bias_gelu".to_string(),
+                description: "Fused GEMM + bias + GeLU".to_string(),
+                operations: vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::GeLU],
+                metal_kernel: "gemm_bias_gelu_f32".to_string(),
+                speedup_factor: 1.5,
+                memory_reduction: 0.66,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "add_relu".to_string(),
+                description: "Fused elementwise add + ReLU".to_string(),
+                operations: vec![PatternOp::Add, PatternOp::ReLU],
+                metal_kernel: "add_relu_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "add_gelu".to_string(),
+                description: "Fused elementwise add + GeLU".to_string(),
+                operations: vec![PatternOp::Add, PatternOp::GeLU],
+                metal_kernel: "add_gelu_f32".to_string(),
+                speedup_factor: 1.3,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "fma_relu".to_string(),
+                description: "Fused FMA + ReLU (a*b+c then ReLU)".to_string(),
+                operations: vec![PatternOp::Mul, PatternOp::Add, PatternOp::ReLU],
+                metal_kernel: "fma_relu_f32".to_string(),
+                speedup_factor: 1.4,
+                memory_reduction: 0.66,
+                enabled: true,
+            },
+            FusionPattern {
+                name: "geglu".to_string(),
+                description: "Fused GeGLU gated activation".to_string(),
+                operations: vec![PatternOp::GeGLU],
+                metal_kernel: "geglu_fused_f32".to_string(),
+                speedup_factor: 1.5,
+                memory_reduction: 0.50,
+                enabled: true,
+            },
+        ]
+    }
+
+    /// Enable or disable a pattern by name
+    pub fn set_pattern_enabled(&mut self, name: &str, enabled: bool) -> bool {
+        for patterns in self.patterns.values_mut() {
+            for pattern in patterns.iter_mut() {
+                if pattern.name == name {
+                    pattern.enabled = enabled;
+                    return true;
+                }
+            }
+        }
+        false
+    }
+
+    /// Get pattern by name
+    pub fn get_pattern(&self, name: &str) -> Option<&FusionPattern> {
+        for patterns in self.patterns.values() {
+            for pattern in patterns {
+                if pattern.name == name {
+                    return Some(pattern);
+                }
+            }
+        }
+        None
+    }
+
+    /// Estimate total speedup for an architecture
+    pub fn estimate_speedup(&self, arch: Architecture) -> f32 {
+        let patterns = self.patterns_for(arch);
+        if patterns.is_empty() {
+            return 1.0;
+        }
+
+        // Geometric mean of speedup factors (more realistic than arithmetic mean)
+        let product: f32 = patterns
+            .iter()
+            .filter(|p| p.enabled)
+            .map(|p| p.speedup_factor)
+            .product();
+
+        let count = patterns.iter().filter(|p| p.enabled).count();
+        if count == 0 {
+            1.0
+        } else {
+            product.powf(1.0 / count as f32)
+        }
+    }
+}
+
+impl Default for FusionPatternLibrary {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Pattern matching result
+#[derive(Debug, Clone)]
+pub struct PatternMatch {
+    /// The matched pattern
+    pub pattern: FusionPattern,
+    /// Start index in the operation sequence
+    pub start_idx: usize,
+    /// End index in the operation sequence
+    pub end_idx: usize,
+    /// Confidence score (0.0 to 1.0)
+    pub confidence: f32,
+}
+
+/// Pattern matcher for finding fusion opportunities
+#[derive(Debug)]
+pub struct PatternMatcher {
+    library: FusionPatternLibrary,
+    architecture: Architecture,
+}
+
+impl PatternMatcher {
+    /// Create a new pattern matcher
+    pub fn new(architecture: Architecture) -> Self {
+        Self {
+            library: FusionPatternLibrary::new(),
+            architecture,
+        }
+    }
+
+    /// Create with a custom pattern library
+    pub fn with_library(library: FusionPatternLibrary, architecture: Architecture) -> Self {
+        Self {
+            library,
+            architecture,
+        }
+    }
+
+    /// Find all pattern matches in an operation sequence
+    pub fn find_matches(&self, operations: &[PatternOp]) -> Vec<PatternMatch> {
+        let mut matches = Vec::new();
+        let patterns = self.library.patterns_for(self.architecture);
+
+        for pattern in patterns {
+            if !pattern.enabled {
+                continue;
+            }
+
+            // Try to match pattern at each position
+            for start_idx in 0..operations.len() {
+                if let Some((end_idx, confidence)) =
+                    self.try_match_at(operations, &pattern.operations, start_idx)
+                {
+                    matches.push(PatternMatch {
+                        pattern: pattern.clone(),
+                        start_idx,
+                        end_idx,
+                        confidence,
+                    });
+                }
+            }
+        }
+
+        // Sort by confidence and position
+        matches.sort_by(|a, b| {
+            b.confidence
+                .partial_cmp(&a.confidence)
+                .unwrap_or(std::cmp::Ordering::Equal)
+                .then(a.start_idx.cmp(&b.start_idx))
+        });
+
+        matches
+    }
+
+    /// Try to match a pattern at a specific position
+    fn try_match_at(
+        &self,
+        operations: &[PatternOp],
+        pattern: &[PatternOp],
+        start_idx: usize,
+    ) -> Option<(usize, f32)> {
+        if start_idx + pattern.len() > operations.len() {
+            return None;
+        }
+
+        let mut confidence = 1.0;
+
+        for (i, pattern_op) in pattern.iter().enumerate() {
+            let op = &operations[start_idx + i];
+
+            match (pattern_op, op) {
+                // Exact match
+                (a, b) if a == b => {}
+                // Wildcard activation
+                (PatternOp::AnyActivation, PatternOp::ReLU)
+                | (PatternOp::AnyActivation, PatternOp::GeLU)
+                | (PatternOp::AnyActivation, PatternOp::SiLU) => {
+                    confidence *= 0.9; // Slightly lower confidence for wildcards
+                }
+                // No match
+                _ => return None,
+            }
+        }
+
+        Some((start_idx + pattern.len(), confidence))
+    }
+
+    /// Find the best non-overlapping matches
+    pub fn find_best_matches(&self, operations: &[PatternOp]) -> Vec<PatternMatch> {
+        let all_matches = self.find_matches(operations);
+        let mut best = Vec::new();
+        let mut covered = vec![false; operations.len()];
+
+        for m in all_matches {
+            // Check if this match overlaps with existing ones
+            let overlaps = (m.start_idx..m.end_idx).any(|i| covered[i]);
+
+            if !overlaps {
+                // Mark as covered
+                for i in m.start_idx..m.end_idx {
+                    covered[i] = true;
+                }
+                best.push(m);
+            }
+        }
+
+        best
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pattern_library_creation() {
+        let lib = FusionPatternLibrary::new();
+        assert!(lib.patterns_for(Architecture::Transformer).len() > 0);
+        assert!(lib.patterns_for(Architecture::LLaMA).len() > 0);
+        assert!(lib.patterns_for(Architecture::Mamba).len() > 0);
+    }
+
+    #[test]
+    fn test_pattern_matching() {
+        let matcher = PatternMatcher::new(Architecture::Generic);
+
+        let ops = vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::ReLU];
+
+        let matches = matcher.find_matches(&ops);
+        assert!(!matches.is_empty());
+    }
+
+    #[test]
+    fn test_llama_patterns() {
+        let lib = FusionPatternLibrary::new();
+        let llama_patterns = lib.patterns_for(Architecture::LLaMA);
+
+        // Should have SwiGLU pattern
+        assert!(llama_patterns.iter().any(|p| p.name == "llama_swiglu"));
+
+        // Should have RMSNorm pattern
+        assert!(llama_patterns.iter().any(|p| p.name == "llama_rmsnorm"));
+    }
+
+    #[test]
+    fn test_speedup_estimation() {
+        let lib = FusionPatternLibrary::new();
+
+        let speedup = lib.estimate_speedup(Architecture::LLaMA);
+        assert!(speedup > 1.0, "LLaMA should have speedup > 1.0");
+    }
+
+    #[test]
+    fn test_pattern_enable_disable() {
+        let mut lib = FusionPatternLibrary::new();
+
+        assert!(lib.set_pattern_enabled("gemm_relu", false));
+
+        let pattern = lib.get_pattern("gemm_relu").unwrap();
+        assert!(!pattern.enabled);
+    }
+
+    #[test]
+    fn test_best_matches_no_overlap() {
+        let matcher = PatternMatcher::new(Architecture::Generic);
+
+        // Sequence that could have overlapping patterns
+        let ops = vec![
+            PatternOp::MatMul,
+            PatternOp::ReLU,
+            PatternOp::MatMul,
+            PatternOp::GeLU,
+        ];
+
+        let best = matcher.find_best_matches(&ops);
+
+        // Check no overlapping matches
+        for i in 0..best.len() {
+            for j in (i + 1)..best.len() {
+                let range_i = best[i].start_idx..best[i].end_idx;
+                let range_j = best[j].start_idx..best[j].end_idx;
+
+                // Ranges should not overlap
+                assert!(
+                    range_i.end <= range_j.start || range_j.end <= range_i.start,
+                    "Matches should not overlap"
+                );
+            }
+        }
+    }
+}