Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,714 @@
|
||||
//! Metal Kernel Fusion DSL
|
||||
//!
|
||||
//! Provides a declarative DSL for composing fused Metal operations.
|
||||
//! This module bridges the fusion analysis from rtx-polygraph with
|
||||
//! the actual Metal kernel execution from rtx-tensor.
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use rtx_synthesis::aot_impl::metal_fusion::{FusedPipeline, FusionBuilder};
|
||||
//!
|
||||
//! let pipeline = FusionBuilder::new()
|
||||
//! .gemm(a, b)
|
||||
//! .bias(bias_vec)
|
||||
//! .activation(Activation::ReLU)
|
||||
//! .build(device)?;
|
||||
//!
|
||||
//! let output = pipeline.execute()?;
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Activation functions for fused operations
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum Activation {
|
||||
/// No activation function
|
||||
None,
|
||||
/// ReLU: max(0, x)
|
||||
ReLU,
|
||||
/// GeLU: Gaussian Error Linear Unit
|
||||
GeLU,
|
||||
/// SiLU/Swish: x * sigmoid(x)
|
||||
SiLU,
|
||||
/// Tanh activation
|
||||
Tanh,
|
||||
/// Sigmoid activation
|
||||
Sigmoid,
|
||||
}
|
||||
|
||||
impl Activation {
|
||||
/// Get the Metal kernel suffix for this activation
|
||||
pub fn metal_suffix(&self) -> &'static str {
|
||||
match self {
|
||||
Activation::None => "",
|
||||
Activation::ReLU => "_relu",
|
||||
Activation::GeLU => "_gelu",
|
||||
Activation::SiLU => "_silu",
|
||||
Activation::Tanh => "_tanh",
|
||||
Activation::Sigmoid => "_sigmoid",
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this activation requires a fused kernel
|
||||
pub fn requires_fusion(&self) -> bool {
|
||||
!matches!(self, Activation::None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalization types for fused operations
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum Normalization {
|
||||
/// Layer normalization
|
||||
LayerNorm { epsilon: f32 },
|
||||
/// RMS normalization (used in LLaMA, Gemma, etc.)
|
||||
RMSNorm { epsilon: f32 },
|
||||
/// Batch normalization
|
||||
BatchNorm { epsilon: f32, momentum: f32 },
|
||||
}
|
||||
|
||||
/// Operation types that can be fused
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum FusedOp {
|
||||
/// Matrix multiplication: C = A @ B
|
||||
Gemm {
|
||||
transpose_a: bool,
|
||||
transpose_b: bool,
|
||||
},
|
||||
/// Bias addition: C = C + bias
|
||||
BiasAdd,
|
||||
/// Activation function
|
||||
Activation(Activation),
|
||||
/// Residual connection: output = input + residual
|
||||
Residual,
|
||||
/// Normalization layer
|
||||
Normalization(Normalization),
|
||||
/// Elementwise addition: C = A + B
|
||||
Add,
|
||||
/// Elementwise multiplication: C = A * B
|
||||
Mul,
|
||||
/// Fused multiply-add: C = A * B + C
|
||||
FMA,
|
||||
/// Gated activation (SwiGLU, GeGLU)
|
||||
GatedActivation { gate_activation: Activation },
|
||||
/// Rotary position embedding
|
||||
RoPE { head_dim: usize },
|
||||
/// Causal attention mask + softmax
|
||||
CausalAttention { num_heads: usize, head_dim: usize },
|
||||
}
|
||||
|
||||
impl FusedOp {
|
||||
/// Check if this operation can fuse with the previous one
|
||||
pub fn can_fuse_with(&self, prev: &FusedOp) -> bool {
|
||||
match (prev, self) {
|
||||
// GEMM can fuse with bias and activation
|
||||
(FusedOp::Gemm { .. }, FusedOp::BiasAdd) => true,
|
||||
(FusedOp::Gemm { .. }, FusedOp::Activation(_)) => true,
|
||||
(FusedOp::BiasAdd, FusedOp::Activation(_)) => true,
|
||||
|
||||
// Residual can fuse with normalization
|
||||
(FusedOp::Residual, FusedOp::Normalization(_)) => true,
|
||||
|
||||
// Add/Mul can fuse with activation
|
||||
(FusedOp::Add, FusedOp::Activation(_)) => true,
|
||||
(FusedOp::Mul, FusedOp::Activation(_)) => true,
|
||||
(FusedOp::FMA, FusedOp::Activation(_)) => true,
|
||||
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the memory bandwidth savings from fusion (0.0 to 1.0)
|
||||
pub fn fusion_savings(&self, prev: &FusedOp) -> f32 {
|
||||
if !self.can_fuse_with(prev) {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
match (prev, self) {
|
||||
// GEMM + bias + activation saves writing intermediate
|
||||
(FusedOp::Gemm { .. }, FusedOp::BiasAdd) => 0.33,
|
||||
(FusedOp::Gemm { .. }, FusedOp::Activation(_)) => 0.50,
|
||||
(FusedOp::BiasAdd, FusedOp::Activation(_)) => 0.33,
|
||||
|
||||
// Residual + norm avoids intermediate storage
|
||||
(FusedOp::Residual, FusedOp::Normalization(_)) => 0.50,
|
||||
|
||||
// Elementwise chains
|
||||
(FusedOp::Add, FusedOp::Activation(_)) => 0.50,
|
||||
(FusedOp::Mul, FusedOp::Activation(_)) => 0.50,
|
||||
(FusedOp::FMA, FusedOp::Activation(_)) => 0.33,
|
||||
|
||||
_ => 0.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Represents a sequence of operations to be fused
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FusionChain {
|
||||
/// Operations in the chain
|
||||
pub ops: Vec<FusedOp>,
|
||||
/// Input tensor shapes
|
||||
pub input_shapes: Vec<Vec<usize>>,
|
||||
/// Output tensor shape
|
||||
pub output_shape: Vec<usize>,
|
||||
}
|
||||
|
||||
impl FusionChain {
|
||||
/// Create a new empty fusion chain
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
ops: Vec::new(),
|
||||
input_shapes: Vec::new(),
|
||||
output_shape: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add an operation to the chain
|
||||
pub fn add_op(&mut self, op: FusedOp) {
|
||||
self.ops.push(op);
|
||||
}
|
||||
|
||||
/// Check if the chain can be executed as a single fused kernel
|
||||
pub fn is_fusable(&self) -> bool {
|
||||
if self.ops.len() < 2 {
|
||||
return false;
|
||||
}
|
||||
|
||||
for i in 1..self.ops.len() {
|
||||
if !self.ops[i].can_fuse_with(&self.ops[i - 1]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Get the Metal kernel name for this fusion chain
|
||||
pub fn metal_kernel_name(&self) -> Option<String> {
|
||||
if !self.is_fusable() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Build kernel name based on operations
|
||||
let mut name_parts = Vec::new();
|
||||
|
||||
for op in &self.ops {
|
||||
match op {
|
||||
FusedOp::Gemm { .. } => name_parts.push("gemm"),
|
||||
FusedOp::BiasAdd => name_parts.push("bias"),
|
||||
FusedOp::Activation(act) => {
|
||||
name_parts.push(match act {
|
||||
Activation::ReLU => "relu",
|
||||
Activation::GeLU => "gelu",
|
||||
Activation::SiLU => "silu",
|
||||
Activation::Tanh => "tanh",
|
||||
Activation::Sigmoid => "sigmoid",
|
||||
Activation::None => continue,
|
||||
});
|
||||
}
|
||||
FusedOp::Residual => name_parts.push("residual"),
|
||||
FusedOp::Normalization(norm) => {
|
||||
name_parts.push(match norm {
|
||||
Normalization::LayerNorm { .. } => "layernorm",
|
||||
Normalization::RMSNorm { .. } => "rmsnorm",
|
||||
Normalization::BatchNorm { .. } => "batchnorm",
|
||||
});
|
||||
}
|
||||
FusedOp::Add => name_parts.push("add"),
|
||||
FusedOp::Mul => name_parts.push("mul"),
|
||||
FusedOp::FMA => name_parts.push("fma"),
|
||||
FusedOp::GatedActivation { gate_activation } => {
|
||||
name_parts.push(match gate_activation {
|
||||
Activation::SiLU => "swiglu",
|
||||
Activation::GeLU => "geglu",
|
||||
_ => "gated",
|
||||
});
|
||||
}
|
||||
FusedOp::RoPE { .. } => name_parts.push("rope"),
|
||||
FusedOp::CausalAttention { .. } => name_parts.push("causal_attn"),
|
||||
}
|
||||
}
|
||||
|
||||
Some(format!("{}_fused_f32", name_parts.join("_")))
|
||||
}
|
||||
|
||||
/// Calculate total memory bandwidth savings from fusion
|
||||
pub fn total_savings(&self) -> f32 {
|
||||
if self.ops.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut savings = 0.0;
|
||||
for i in 1..self.ops.len() {
|
||||
savings += self.ops[i].fusion_savings(&self.ops[i - 1]);
|
||||
}
|
||||
|
||||
savings
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FusionChain {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Builder for constructing fused operation pipelines
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FusionBuilder {
|
||||
chain: FusionChain,
|
||||
}
|
||||
|
||||
impl FusionBuilder {
|
||||
/// Create a new fusion builder
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
chain: FusionChain::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a GEMM operation
|
||||
pub fn gemm(mut self, transpose_a: bool, transpose_b: bool) -> Self {
|
||||
self.chain.add_op(FusedOp::Gemm {
|
||||
transpose_a,
|
||||
transpose_b,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a bias addition
|
||||
pub fn bias(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::BiasAdd);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add an activation function
|
||||
pub fn activation(mut self, act: Activation) -> Self {
|
||||
self.chain.add_op(FusedOp::Activation(act));
|
||||
self
|
||||
}
|
||||
|
||||
/// Add ReLU activation (convenience method)
|
||||
pub fn relu(self) -> Self {
|
||||
self.activation(Activation::ReLU)
|
||||
}
|
||||
|
||||
/// Add GeLU activation (convenience method)
|
||||
pub fn gelu(self) -> Self {
|
||||
self.activation(Activation::GeLU)
|
||||
}
|
||||
|
||||
/// Add SiLU/Swish activation (convenience method)
|
||||
pub fn silu(self) -> Self {
|
||||
self.activation(Activation::SiLU)
|
||||
}
|
||||
|
||||
/// Add a residual connection
|
||||
pub fn residual(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::Residual);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add layer normalization
|
||||
pub fn layer_norm(mut self, epsilon: f32) -> Self {
|
||||
self.chain
|
||||
.add_op(FusedOp::Normalization(Normalization::LayerNorm { epsilon }));
|
||||
self
|
||||
}
|
||||
|
||||
/// Add RMS normalization
|
||||
pub fn rms_norm(mut self, epsilon: f32) -> Self {
|
||||
self.chain
|
||||
.add_op(FusedOp::Normalization(Normalization::RMSNorm { epsilon }));
|
||||
self
|
||||
}
|
||||
|
||||
/// Add elementwise addition
|
||||
pub fn add(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::Add);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add elementwise multiplication
|
||||
pub fn mul(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::Mul);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add fused multiply-add
|
||||
pub fn fma(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::FMA);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add SwiGLU gated activation
|
||||
pub fn swiglu(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::GatedActivation {
|
||||
gate_activation: Activation::SiLU,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Add GeGLU gated activation
|
||||
pub fn geglu(mut self) -> Self {
|
||||
self.chain.add_op(FusedOp::GatedActivation {
|
||||
gate_activation: Activation::GeLU,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Add rotary position embedding
|
||||
pub fn rope(mut self, head_dim: usize) -> Self {
|
||||
self.chain.add_op(FusedOp::RoPE { head_dim });
|
||||
self
|
||||
}
|
||||
|
||||
/// Add causal attention computation
|
||||
pub fn causal_attention(mut self, num_heads: usize, head_dim: usize) -> Self {
|
||||
self.chain.add_op(FusedOp::CausalAttention {
|
||||
num_heads,
|
||||
head_dim,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the fusion chain
|
||||
pub fn build(self) -> FusionChain {
|
||||
self.chain
|
||||
}
|
||||
|
||||
/// Check if the current chain is fusable
|
||||
pub fn is_fusable(&self) -> bool {
|
||||
self.chain.is_fusable()
|
||||
}
|
||||
|
||||
/// Get the Metal kernel name for the current chain
|
||||
pub fn kernel_name(&self) -> Option<String> {
|
||||
self.chain.metal_kernel_name()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FusionBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Fusion pattern matcher for automatic kernel selection
|
||||
#[derive(Debug)]
|
||||
pub struct FusionPatternMatcher {
|
||||
/// Known fusion patterns and their kernel names
|
||||
patterns: HashMap<Vec<FusedOp>, String>,
|
||||
}
|
||||
|
||||
impl FusionPatternMatcher {
|
||||
/// Create a new pattern matcher with standard patterns
|
||||
pub fn new() -> Self {
|
||||
let mut patterns = HashMap::new();
|
||||
|
||||
// GEMM + Activation patterns
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
},
|
||||
FusedOp::Activation(Activation::ReLU),
|
||||
],
|
||||
"gemm_relu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
},
|
||||
FusedOp::Activation(Activation::GeLU),
|
||||
],
|
||||
"gemm_gelu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
},
|
||||
FusedOp::Activation(Activation::SiLU),
|
||||
],
|
||||
"gemm_silu_f32".to_string(),
|
||||
);
|
||||
|
||||
// GEMM + Bias + Activation patterns
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
},
|
||||
FusedOp::BiasAdd,
|
||||
FusedOp::Activation(Activation::ReLU),
|
||||
],
|
||||
"gemm_bias_relu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
},
|
||||
FusedOp::BiasAdd,
|
||||
FusedOp::Activation(Activation::GeLU),
|
||||
],
|
||||
"gemm_bias_gelu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
},
|
||||
FusedOp::BiasAdd,
|
||||
FusedOp::Activation(Activation::SiLU),
|
||||
],
|
||||
"gemm_bias_silu_f32".to_string(),
|
||||
);
|
||||
|
||||
// Residual + Normalization patterns
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Residual,
|
||||
FusedOp::Normalization(Normalization::RMSNorm { epsilon: 1e-6 }),
|
||||
],
|
||||
"residual_rmsnorm_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![
|
||||
FusedOp::Residual,
|
||||
FusedOp::Normalization(Normalization::LayerNorm { epsilon: 1e-6 }),
|
||||
],
|
||||
"residual_layernorm_f32".to_string(),
|
||||
);
|
||||
|
||||
// Elementwise patterns
|
||||
patterns.insert(
|
||||
vec![FusedOp::Add, FusedOp::Activation(Activation::ReLU)],
|
||||
"add_relu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![FusedOp::Add, FusedOp::Activation(Activation::GeLU)],
|
||||
"add_gelu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![FusedOp::Add, FusedOp::Activation(Activation::SiLU)],
|
||||
"add_silu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![FusedOp::FMA, FusedOp::Activation(Activation::ReLU)],
|
||||
"fma_relu_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![FusedOp::FMA, FusedOp::Activation(Activation::GeLU)],
|
||||
"fma_gelu_f32".to_string(),
|
||||
);
|
||||
|
||||
// Gated activation patterns
|
||||
patterns.insert(
|
||||
vec![FusedOp::GatedActivation {
|
||||
gate_activation: Activation::SiLU,
|
||||
}],
|
||||
"swiglu_fused_f32".to_string(),
|
||||
);
|
||||
|
||||
patterns.insert(
|
||||
vec![FusedOp::GatedActivation {
|
||||
gate_activation: Activation::GeLU,
|
||||
}],
|
||||
"geglu_fused_f32".to_string(),
|
||||
);
|
||||
|
||||
Self { patterns }
|
||||
}
|
||||
|
||||
/// Match a fusion chain to a known pattern
|
||||
pub fn match_pattern(&self, chain: &FusionChain) -> Option<&String> {
|
||||
// Try to match with epsilon variations for normalization
|
||||
for (pattern, kernel_name) in &self.patterns {
|
||||
if Self::patterns_match(&chain.ops, pattern) {
|
||||
return Some(kernel_name);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Check if two operation sequences match (ignoring epsilon values)
|
||||
fn patterns_match(ops: &[FusedOp], pattern: &[FusedOp]) -> bool {
|
||||
if ops.len() != pattern.len() {
|
||||
return false;
|
||||
}
|
||||
|
||||
ops.iter().zip(pattern.iter()).all(|(op, pat)| {
|
||||
match (op, pat) {
|
||||
// Normalization matches if type is same (ignore epsilon)
|
||||
(
|
||||
FusedOp::Normalization(Normalization::LayerNorm { .. }),
|
||||
FusedOp::Normalization(Normalization::LayerNorm { .. }),
|
||||
) => true,
|
||||
(
|
||||
FusedOp::Normalization(Normalization::RMSNorm { .. }),
|
||||
FusedOp::Normalization(Normalization::RMSNorm { .. }),
|
||||
) => true,
|
||||
(
|
||||
FusedOp::Normalization(Normalization::BatchNorm { .. }),
|
||||
FusedOp::Normalization(Normalization::BatchNorm { .. }),
|
||||
) => true,
|
||||
// Exact match for other operations
|
||||
_ => op == pat,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Get all available fusion patterns
|
||||
pub fn available_patterns(&self) -> impl Iterator<Item = (&Vec<FusedOp>, &String)> {
|
||||
self.patterns.iter()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FusionPatternMatcher {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics for fusion analysis
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct FusionStats {
|
||||
/// Number of operations analyzed
|
||||
pub total_ops: usize,
|
||||
/// Number of fusion opportunities found
|
||||
pub fusion_opportunities: usize,
|
||||
/// Estimated memory bandwidth savings
|
||||
pub bandwidth_savings_percent: f32,
|
||||
/// Number of kernel launches saved
|
||||
pub kernel_launches_saved: usize,
|
||||
}
|
||||
|
||||
impl FusionStats {
|
||||
/// Create new stats from a fusion chain
|
||||
pub fn from_chain(chain: &FusionChain) -> Self {
|
||||
let total_ops = chain.ops.len();
|
||||
let fusion_opportunities = if chain.is_fusable() { 1 } else { 0 };
|
||||
let bandwidth_savings_percent = chain.total_savings() * 100.0;
|
||||
let kernel_launches_saved = if chain.is_fusable() {
|
||||
chain.ops.len().saturating_sub(1)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
Self {
|
||||
total_ops,
|
||||
fusion_opportunities,
|
||||
bandwidth_savings_percent,
|
||||
kernel_launches_saved,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_activation_suffix() {
|
||||
assert_eq!(Activation::ReLU.metal_suffix(), "_relu");
|
||||
assert_eq!(Activation::GeLU.metal_suffix(), "_gelu");
|
||||
assert_eq!(Activation::SiLU.metal_suffix(), "_silu");
|
||||
assert_eq!(Activation::None.metal_suffix(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fusion_builder() {
|
||||
let chain = FusionBuilder::new()
|
||||
.gemm(false, false)
|
||||
.bias()
|
||||
.relu()
|
||||
.build();
|
||||
|
||||
assert_eq!(chain.ops.len(), 3);
|
||||
assert!(chain.is_fusable());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gemm_activation_fusion() {
|
||||
let chain = FusionBuilder::new().gemm(false, false).gelu().build();
|
||||
|
||||
assert!(chain.is_fusable());
|
||||
assert!(chain.total_savings() > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_residual_norm_fusion() {
|
||||
let chain = FusionBuilder::new().residual().rms_norm(1e-6).build();
|
||||
|
||||
assert!(chain.is_fusable());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pattern_matcher() {
|
||||
let matcher = FusionPatternMatcher::new();
|
||||
|
||||
let chain = FusionBuilder::new().gemm(false, false).relu().build();
|
||||
|
||||
let kernel = matcher.match_pattern(&chain);
|
||||
assert!(kernel.is_some());
|
||||
assert_eq!(kernel.unwrap(), "gemm_relu_f32");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_swiglu_pattern() {
|
||||
let matcher = FusionPatternMatcher::new();
|
||||
|
||||
let chain = FusionBuilder::new().swiglu().build();
|
||||
|
||||
let kernel = matcher.match_pattern(&chain);
|
||||
assert!(kernel.is_some());
|
||||
assert_eq!(kernel.unwrap(), "swiglu_fused_f32");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fusion_stats() {
|
||||
let chain = FusionBuilder::new()
|
||||
.gemm(false, false)
|
||||
.bias()
|
||||
.relu()
|
||||
.build();
|
||||
|
||||
let stats = FusionStats::from_chain(&chain);
|
||||
assert_eq!(stats.total_ops, 3);
|
||||
assert_eq!(stats.fusion_opportunities, 1);
|
||||
assert_eq!(stats.kernel_launches_saved, 2);
|
||||
assert!(stats.bandwidth_savings_percent > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unfusable_chain() {
|
||||
let mut chain = FusionChain::new();
|
||||
chain.add_op(FusedOp::RoPE { head_dim: 64 });
|
||||
chain.add_op(FusedOp::Gemm {
|
||||
transpose_a: false,
|
||||
transpose_b: false,
|
||||
});
|
||||
|
||||
// RoPE followed by GEMM is not fusable
|
||||
assert!(!chain.is_fusable());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user