686 lines
23 KiB
Rust
686 lines
23 KiB
Rust
//! 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"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|