osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
1296 lines
46 KiB
Rust
1296 lines
46 KiB
Rust
//! Fusion pattern detection and analysis
|
||
//!
|
||
//! This module analyzes operation streams to detect fusion opportunities.
|
||
//! It supports multiple fusion patterns including elementwise chains,
|
||
//! softmax patterns, layer normalization, and activation chains.
|
||
|
||
use crate::config::FusionConfig;
|
||
use crate::kernel::{DType, KernelSignature, StreamOpKind, StreamOperation, TensorId};
|
||
use crate::stream::OperationStream;
|
||
use std::collections::{HashMap, HashSet};
|
||
|
||
/// Pre-computed analysis context for efficient pattern detection
|
||
/// This avoids rebuilding data structures for each detection pass
|
||
struct AnalysisContext<'a> {
|
||
/// Reference to the operations slice
|
||
ops: &'a [StreamOperation],
|
||
/// Map from output tensor ID to operation index
|
||
output_to_op: HashMap<TensorId, usize>,
|
||
/// Map from operation index to indices of operations that consume its output
|
||
op_consumers: HashMap<usize, Vec<usize>>,
|
||
/// Map from tensor ID to indices of operations that consume it
|
||
tensor_consumers: HashMap<TensorId, Vec<usize>>,
|
||
}
|
||
|
||
impl<'a> AnalysisContext<'a> {
|
||
/// Build analysis context from operations
|
||
fn new(ops: &'a [StreamOperation]) -> Self {
|
||
let capacity = ops.len();
|
||
|
||
// Pre-allocate with estimated capacity
|
||
let mut output_to_op = HashMap::with_capacity(capacity);
|
||
let mut op_consumers: HashMap<usize, Vec<usize>> = HashMap::with_capacity(capacity);
|
||
let mut tensor_consumers: HashMap<TensorId, Vec<usize>> =
|
||
HashMap::with_capacity(capacity * 2);
|
||
|
||
// Single pass to build all maps
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
output_to_op.insert(op.output, idx);
|
||
for &input in &op.inputs {
|
||
tensor_consumers.entry(input).or_default().push(idx);
|
||
}
|
||
}
|
||
|
||
// Build op_consumers from tensor_consumers
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
if let Some(consumers) = tensor_consumers.get(&op.output) {
|
||
op_consumers.insert(idx, consumers.clone());
|
||
}
|
||
}
|
||
|
||
Self {
|
||
ops,
|
||
output_to_op,
|
||
op_consumers,
|
||
tensor_consumers,
|
||
}
|
||
}
|
||
|
||
/// Get the operation at an index
|
||
#[inline]
|
||
fn get_op(&self, idx: usize) -> &StreamOperation {
|
||
&self.ops[idx]
|
||
}
|
||
|
||
/// Get the producer operation index for a tensor
|
||
#[inline]
|
||
fn producer_of(&self, tensor: TensorId) -> Option<usize> {
|
||
self.output_to_op.get(&tensor).copied()
|
||
}
|
||
|
||
/// Get consumer operation indices for an operation
|
||
#[inline]
|
||
fn consumers_of_op(&self, op_idx: usize) -> &[usize] {
|
||
self.op_consumers
|
||
.get(&op_idx)
|
||
.map(|v| v.as_slice())
|
||
.unwrap_or(&[])
|
||
}
|
||
|
||
/// Get consumer operation indices for a tensor
|
||
#[inline]
|
||
fn consumers_of_tensor(&self, tensor: TensorId) -> &[usize] {
|
||
self.tensor_consumers
|
||
.get(&tensor)
|
||
.map(|v| v.as_slice())
|
||
.unwrap_or(&[])
|
||
}
|
||
|
||
/// Check if a tensor is produced by an operation in the stream
|
||
#[inline]
|
||
fn is_internal(&self, tensor: TensorId) -> bool {
|
||
self.output_to_op.contains_key(&tensor)
|
||
}
|
||
}
|
||
|
||
/// Types of fusion patterns that can be detected
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub enum FusionPattern {
|
||
/// Chain of elementwise operations (add, mul, exp, relu, etc.)
|
||
ElementwiseChain,
|
||
/// Softmax pattern: subtract_max -> exp -> sum -> divide
|
||
/// Includes the reduction dimension for proper kernel generation
|
||
SoftmaxPattern { dim: Option<usize> },
|
||
/// Layer normalization pattern: (x - mean) / sqrt(var + eps) * scale + bias
|
||
LayerNormPattern { eps: f32, dim: Option<usize> },
|
||
/// RMS normalization pattern: x / sqrt(mean(x^2) + eps) * scale
|
||
RmsNormPattern { eps: f32, dim: Option<usize> },
|
||
/// Activation chain (e.g., GELU, SiLU with surrounding ops)
|
||
ActivationChain,
|
||
/// Reduction followed by broadcast and element-wise ops
|
||
ReductionBroadcast {
|
||
reduction_op: ReductionType,
|
||
dim: Option<usize>,
|
||
},
|
||
/// Generic fuseable sequence
|
||
Generic,
|
||
}
|
||
|
||
/// Types of reduction operations
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub enum ReductionType {
|
||
Sum,
|
||
Mean,
|
||
Max,
|
||
Min,
|
||
}
|
||
|
||
impl FusionPattern {
|
||
/// Get the expected speedup multiplier for this pattern
|
||
pub fn expected_speedup(&self) -> f64 {
|
||
match self {
|
||
FusionPattern::ElementwiseChain => 2.0,
|
||
FusionPattern::SoftmaxPattern { .. } => 3.0,
|
||
FusionPattern::LayerNormPattern { .. } => 2.5,
|
||
FusionPattern::RmsNormPattern { .. } => 2.5,
|
||
FusionPattern::ActivationChain => 1.8,
|
||
FusionPattern::ReductionBroadcast { .. } => 1.5,
|
||
FusionPattern::Generic => 1.5,
|
||
}
|
||
}
|
||
|
||
/// Get the name of this pattern
|
||
pub fn name(&self) -> &'static str {
|
||
match self {
|
||
FusionPattern::ElementwiseChain => "elementwise_chain",
|
||
FusionPattern::SoftmaxPattern { .. } => "softmax",
|
||
FusionPattern::LayerNormPattern { .. } => "layer_norm",
|
||
FusionPattern::RmsNormPattern { .. } => "rms_norm",
|
||
FusionPattern::ActivationChain => "activation_chain",
|
||
FusionPattern::ReductionBroadcast { .. } => "reduction_broadcast",
|
||
FusionPattern::Generic => "generic",
|
||
}
|
||
}
|
||
|
||
/// Check if this pattern involves reduction operations
|
||
pub fn has_reduction(&self) -> bool {
|
||
matches!(
|
||
self,
|
||
FusionPattern::SoftmaxPattern { .. }
|
||
| FusionPattern::LayerNormPattern { .. }
|
||
| FusionPattern::RmsNormPattern { .. }
|
||
| FusionPattern::ReductionBroadcast { .. }
|
||
)
|
||
}
|
||
|
||
/// Get the reduction dimension if applicable
|
||
pub fn reduction_dim(&self) -> Option<usize> {
|
||
match self {
|
||
FusionPattern::SoftmaxPattern { dim } => *dim,
|
||
FusionPattern::LayerNormPattern { dim, .. } => *dim,
|
||
FusionPattern::RmsNormPattern { dim, .. } => *dim,
|
||
FusionPattern::ReductionBroadcast { dim, .. } => *dim,
|
||
_ => None,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// A detected fusion opportunity
|
||
#[derive(Debug, Clone)]
|
||
pub struct FusionOpportunity {
|
||
/// The type of fusion pattern detected
|
||
pub pattern: FusionPattern,
|
||
/// Indices of operations in the stream that form this opportunity
|
||
pub operation_indices: Vec<usize>,
|
||
/// The operations themselves
|
||
pub operations: Vec<StreamOpKind>,
|
||
/// External input tensor IDs (not produced within the fusion group)
|
||
pub external_inputs: Vec<TensorId>,
|
||
/// Output tensor ID
|
||
pub output: TensorId,
|
||
/// Estimated speedup from fusing these operations
|
||
pub estimated_speedup: f64,
|
||
/// Data type of the operations
|
||
pub dtype: DType,
|
||
}
|
||
|
||
impl FusionOpportunity {
|
||
/// Get the number of operations in this opportunity
|
||
pub fn operation_count(&self) -> usize {
|
||
self.operation_indices.len()
|
||
}
|
||
|
||
/// Check if this opportunity is worthwhile based on minimum ops threshold
|
||
pub fn is_worthwhile(&self, min_ops: usize) -> bool {
|
||
self.operation_indices.len() >= min_ops && self.estimated_speedup > 1.0
|
||
}
|
||
|
||
/// Generate a kernel signature for this opportunity
|
||
pub fn to_kernel_signature(&self) -> KernelSignature {
|
||
KernelSignature::new(
|
||
self.operations.clone(),
|
||
self.external_inputs.len(),
|
||
self.dtype,
|
||
)
|
||
}
|
||
|
||
/// Calculate the benefit score (ops × speedup)
|
||
pub fn benefit_score(&self) -> f64 {
|
||
self.operation_count() as f64 * self.estimated_speedup
|
||
}
|
||
}
|
||
|
||
/// Analyzer for detecting fusion opportunities in operation streams
|
||
#[derive(Debug)]
|
||
pub struct FusionAnalyzer {
|
||
config: FusionConfig,
|
||
}
|
||
|
||
impl FusionAnalyzer {
|
||
/// Create a new fusion analyzer with the given configuration
|
||
pub fn new(config: FusionConfig) -> Self {
|
||
Self { config }
|
||
}
|
||
|
||
/// Create an analyzer with default configuration
|
||
pub fn with_defaults() -> Self {
|
||
Self::new(FusionConfig::default())
|
||
}
|
||
|
||
/// Analyze an operation stream and return fusion opportunities
|
||
pub fn analyze(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
|
||
if !self.config.enabled || stream.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
// Build analysis context once for all detection passes
|
||
let ops = stream.operations();
|
||
let ctx = AnalysisContext::new(&ops);
|
||
|
||
let mut opportunities = Vec::with_capacity(ops.len() / 2);
|
||
|
||
// Detect elementwise chains (primary fusion target)
|
||
opportunities.extend(self.detect_elementwise_chains_opt(&ctx));
|
||
|
||
// Detect advanced patterns if enabled
|
||
if self.config.detect_advanced_patterns {
|
||
opportunities.extend(self.detect_softmax_patterns_opt(&ctx));
|
||
opportunities.extend(self.detect_layer_norm_patterns_opt(&ctx));
|
||
opportunities.extend(self.detect_activation_chains_opt(&ctx));
|
||
}
|
||
|
||
// Sort by benefit score (highest first)
|
||
opportunities.sort_by(|a, b| {
|
||
b.benefit_score()
|
||
.partial_cmp(&a.benefit_score())
|
||
.unwrap_or(std::cmp::Ordering::Equal)
|
||
});
|
||
|
||
// Remove overlapping opportunities (greedy selection)
|
||
self.remove_overlaps(opportunities)
|
||
}
|
||
|
||
/// Original analyze method that doesn't use the optimized context (for backward compatibility)
|
||
#[allow(dead_code)]
|
||
fn analyze_legacy(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
|
||
if !self.config.enabled || stream.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
let mut opportunities = Vec::new();
|
||
|
||
// Detect elementwise chains (primary fusion target)
|
||
opportunities.extend(self.detect_elementwise_chains(stream));
|
||
|
||
// Detect advanced patterns if enabled
|
||
if self.config.detect_advanced_patterns {
|
||
opportunities.extend(self.detect_softmax_patterns(stream));
|
||
opportunities.extend(self.detect_layer_norm_patterns(stream));
|
||
opportunities.extend(self.detect_activation_chains(stream));
|
||
}
|
||
|
||
// Sort by benefit score (highest first)
|
||
opportunities.sort_by(|a, b| {
|
||
b.benefit_score()
|
||
.partial_cmp(&a.benefit_score())
|
||
.unwrap_or(std::cmp::Ordering::Equal)
|
||
});
|
||
|
||
// Remove overlapping opportunities (greedy selection)
|
||
self.remove_overlaps(opportunities)
|
||
}
|
||
|
||
/// Detect chains of elementwise operations
|
||
fn detect_elementwise_chains(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
|
||
let ops = stream.operations();
|
||
if ops.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
let mut opportunities = Vec::new();
|
||
let mut visited: HashSet<usize> = HashSet::new();
|
||
|
||
// Build a map from output tensor to operation index
|
||
let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
output_to_op.insert(op.output, idx);
|
||
}
|
||
|
||
// Build consumer map
|
||
let mut consumers: HashMap<usize, Vec<usize>> = HashMap::new();
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
for &input in &op.inputs {
|
||
if let Some(&producer_idx) = output_to_op.get(&input) {
|
||
consumers.entry(producer_idx).or_default().push(idx);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Find chain start points (fuseable ops with no fuseable producer)
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
if visited.contains(&idx) || !op.op.is_fuseable() {
|
||
continue;
|
||
}
|
||
|
||
// Check if this is a chain start
|
||
let is_start = op.inputs.iter().all(|input| {
|
||
output_to_op
|
||
.get(input)
|
||
.map(|&producer_idx| !ops[producer_idx].op.is_fuseable())
|
||
.unwrap_or(true)
|
||
});
|
||
|
||
if !is_start {
|
||
continue;
|
||
}
|
||
|
||
// Extend the chain greedily
|
||
let mut chain = vec![idx];
|
||
let mut chain_outputs: HashSet<TensorId> = HashSet::new();
|
||
chain_outputs.insert(op.output);
|
||
visited.insert(idx);
|
||
|
||
let mut current_idx = idx;
|
||
|
||
loop {
|
||
// Find the next operation in the chain
|
||
let next_candidates: Vec<_> = consumers
|
||
.get(¤t_idx)
|
||
.map(|c| c.as_slice())
|
||
.unwrap_or(&[])
|
||
.iter()
|
||
.filter(|&&next_idx| {
|
||
!visited.contains(&next_idx)
|
||
&& ops[next_idx].op.is_fuseable()
|
||
&& ops[next_idx].inputs.iter().all(|input| {
|
||
// All inputs must be either external or from our chain
|
||
chain_outputs.contains(input) || !output_to_op.contains_key(input)
|
||
})
|
||
})
|
||
.copied()
|
||
.collect();
|
||
|
||
if next_candidates.is_empty() {
|
||
break;
|
||
}
|
||
|
||
// Take the first valid candidate (could be smarter about this)
|
||
let next_idx = next_candidates[0];
|
||
let next_op = &ops[next_idx];
|
||
|
||
chain.push(next_idx);
|
||
chain_outputs.insert(next_op.output);
|
||
visited.insert(next_idx);
|
||
current_idx = next_idx;
|
||
|
||
// Check max chain length
|
||
if chain.len() >= self.config.max_fusion_ops {
|
||
break;
|
||
}
|
||
}
|
||
|
||
// Create opportunity if chain is long enough
|
||
if chain.len() >= self.config.min_fusion_ops {
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
|
||
let output = ops[*chain.last().unwrap()].output;
|
||
let dtype = ops[chain[0]].dtype;
|
||
|
||
// Calculate speedup with bonus for longer chains
|
||
let base_speedup = FusionPattern::ElementwiseChain.expected_speedup();
|
||
let chain_bonus = 1.0 + (chain.len() as f64 - 2.0) * 0.1;
|
||
let estimated_speedup = base_speedup * chain_bonus;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::ElementwiseChain,
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup,
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Detect softmax patterns
|
||
/// Pattern: (sub max) -> exp -> (sum dim) -> div
|
||
/// Or simpler: exp -> sum -> div
|
||
fn detect_softmax_patterns(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
|
||
let ops = stream.operations();
|
||
let mut opportunities = Vec::new();
|
||
|
||
// Build maps for traversal
|
||
let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
|
||
let mut consumers: HashMap<TensorId, Vec<usize>> = HashMap::new();
|
||
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
output_to_op.insert(op.output, idx);
|
||
for &input in &op.inputs {
|
||
consumers.entry(input).or_default().push(idx);
|
||
}
|
||
}
|
||
|
||
// Look for Exp operations as anchor points
|
||
for (exp_idx, exp_op) in ops.iter().enumerate() {
|
||
if !matches!(exp_op.op, StreamOpKind::Exp) {
|
||
continue;
|
||
}
|
||
|
||
let mut chain = vec![exp_idx];
|
||
let mut dim = None;
|
||
|
||
// Look backward for Sub(x, max) pattern (numerical stability)
|
||
if let Some(&input_id) = exp_op.inputs.first() {
|
||
if let Some(&sub_idx) = output_to_op.get(&input_id) {
|
||
if matches!(ops[sub_idx].op, StreamOpKind::Sub) {
|
||
// Check if second input could be a max reduction
|
||
if ops[sub_idx].inputs.len() >= 2 {
|
||
let second_input = ops[sub_idx].inputs[1];
|
||
if let Some(&max_idx) = output_to_op.get(&second_input) {
|
||
if matches!(ops[max_idx].op, StreamOpKind::Max) {
|
||
chain.insert(0, max_idx);
|
||
chain.insert(1, sub_idx);
|
||
}
|
||
}
|
||
}
|
||
// Even without max, sub before exp could be part of pattern
|
||
if chain.len() == 1 {
|
||
chain.insert(0, sub_idx);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Look forward for Sum operation
|
||
let exp_output = exp_op.output;
|
||
let sum_consumer = consumers.get(&exp_output).and_then(|cons| {
|
||
cons.iter().find(|&&idx| {
|
||
matches!(ops[idx].op, StreamOpKind::Sum | StreamOpKind::SumDim { .. })
|
||
})
|
||
});
|
||
|
||
if let Some(&sum_idx) = sum_consumer {
|
||
chain.push(sum_idx);
|
||
|
||
// Extract dimension from SumDim if present
|
||
if let StreamOpKind::SumDim { dim: d } = &ops[sum_idx].op {
|
||
dim = Some(*d);
|
||
}
|
||
|
||
// Look for Div after Sum
|
||
let sum_output = ops[sum_idx].output;
|
||
let div_consumer = consumers.get(&sum_output).and_then(|cons| {
|
||
cons.iter()
|
||
.find(|&&idx| matches!(ops[idx].op, StreamOpKind::Div))
|
||
});
|
||
|
||
if let Some(&div_idx) = div_consumer {
|
||
chain.push(div_idx);
|
||
}
|
||
}
|
||
|
||
// Only create opportunity if we have at least exp + sum (core softmax)
|
||
if chain.len() >= 2
|
||
&& chain
|
||
.iter()
|
||
.any(|&i| matches!(ops[i].op, StreamOpKind::Sum | StreamOpKind::SumDim { .. }))
|
||
{
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
|
||
let output = ops[*chain.last().unwrap()].output;
|
||
let dtype = exp_op.dtype;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::SoftmaxPattern { dim },
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: FusionPattern::SoftmaxPattern { dim }.expected_speedup(),
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Detect layer normalization patterns
|
||
/// Pattern: mean -> sub -> square -> mean -> (add eps) -> sqrt -> div -> (mul scale) -> (add bias)
|
||
fn detect_layer_norm_patterns(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
|
||
let ops = stream.operations();
|
||
let mut opportunities = Vec::new();
|
||
|
||
// Build maps
|
||
let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
|
||
let mut consumers: HashMap<TensorId, Vec<usize>> = HashMap::new();
|
||
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
output_to_op.insert(op.output, idx);
|
||
for &input in &op.inputs {
|
||
consumers.entry(input).or_default().push(idx);
|
||
}
|
||
}
|
||
|
||
// Look for Mean operations as potential pattern starts
|
||
for (mean_idx, mean_op) in ops.iter().enumerate() {
|
||
if !matches!(
|
||
mean_op.op,
|
||
StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
|
||
) {
|
||
continue;
|
||
}
|
||
|
||
let mut chain = vec![mean_idx];
|
||
let mut dim = None;
|
||
let eps = 1e-5f32; // Default eps
|
||
|
||
// Extract dimension
|
||
if let StreamOpKind::MeanDim { dim: d } = &mean_op.op {
|
||
dim = Some(*d);
|
||
}
|
||
|
||
// Look for Sub(x, mean) consumer
|
||
let mean_output = mean_op.output;
|
||
let sub_consumer = consumers.get(&mean_output).and_then(|cons| {
|
||
cons.iter().find(|&&idx| {
|
||
matches!(ops[idx].op, StreamOpKind::Sub)
|
||
&& ops[idx].inputs.get(1) == Some(&mean_output)
|
||
})
|
||
});
|
||
|
||
if let Some(&sub_idx) = sub_consumer {
|
||
chain.push(sub_idx);
|
||
|
||
// Look for Square(sub_result)
|
||
let sub_output = ops[sub_idx].output;
|
||
let square_consumer = consumers.get(&sub_output).and_then(|cons| {
|
||
cons.iter()
|
||
.find(|&&idx| matches!(ops[idx].op, StreamOpKind::Square))
|
||
});
|
||
|
||
if let Some(&square_idx) = square_consumer {
|
||
chain.push(square_idx);
|
||
|
||
// Look for second Mean (variance computation)
|
||
let square_output = ops[square_idx].output;
|
||
let mean2_consumer = consumers.get(&square_output).and_then(|cons| {
|
||
cons.iter().find(|&&idx| {
|
||
matches!(
|
||
ops[idx].op,
|
||
StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
|
||
)
|
||
})
|
||
});
|
||
|
||
if let Some(&mean2_idx) = mean2_consumer {
|
||
chain.push(mean2_idx);
|
||
|
||
// Continue looking for add_eps -> sqrt -> div pattern
|
||
let var_output = ops[mean2_idx].output;
|
||
|
||
// Look for AddScalar (eps) or Sqrt directly
|
||
if let Some(add_eps_idx) = consumers.get(&var_output).and_then(|cons| {
|
||
cons.iter()
|
||
.find(|&&idx| matches!(ops[idx].op, StreamOpKind::AddScalar(_)))
|
||
}) {
|
||
chain.push(*add_eps_idx);
|
||
}
|
||
|
||
// Continue extending if we find sqrt, div, etc.
|
||
// This is a simplified version - full detection would trace the complete path
|
||
}
|
||
}
|
||
}
|
||
|
||
// Create opportunity if we found enough of the pattern
|
||
if chain.len() >= 3 {
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
|
||
let output = ops[*chain.last().unwrap()].output;
|
||
let dtype = mean_op.dtype;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::LayerNormPattern { eps, dim },
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: FusionPattern::LayerNormPattern { eps, dim }
|
||
.expected_speedup(),
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Detect activation chains (e.g., mul -> sigmoid -> mul for SiLU)
|
||
fn detect_activation_chains(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
|
||
let ops = stream.operations();
|
||
let mut opportunities = Vec::new();
|
||
let mut visited: HashSet<usize> = HashSet::new();
|
||
|
||
// Build output to op map
|
||
let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
output_to_op.insert(op.output, idx);
|
||
}
|
||
|
||
// Look for activation functions
|
||
for (idx, op) in ops.iter().enumerate() {
|
||
if visited.contains(&idx) {
|
||
continue;
|
||
}
|
||
|
||
let is_activation = matches!(
|
||
op.op,
|
||
StreamOpKind::GELU
|
||
| StreamOpKind::SiLU
|
||
| StreamOpKind::Sigmoid
|
||
| StreamOpKind::Tanh
|
||
| StreamOpKind::ReLU
|
||
);
|
||
|
||
if !is_activation || !op.op.is_fuseable() {
|
||
continue;
|
||
}
|
||
|
||
let mut chain = vec![idx];
|
||
visited.insert(idx);
|
||
|
||
// Look backward for preceding elementwise op
|
||
if let Some(&input_id) = op.inputs.first() {
|
||
if let Some(&prev_idx) = output_to_op.get(&input_id) {
|
||
if ops[prev_idx].op.is_fuseable() && !visited.contains(&prev_idx) {
|
||
chain.insert(0, prev_idx);
|
||
visited.insert(prev_idx);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Look forward for following elementwise op
|
||
for (next_idx, next_op) in ops.iter().enumerate() {
|
||
if next_op.inputs.contains(&op.output)
|
||
&& next_op.op.is_fuseable()
|
||
&& !visited.contains(&next_idx)
|
||
{
|
||
chain.push(next_idx);
|
||
visited.insert(next_idx);
|
||
break;
|
||
}
|
||
}
|
||
|
||
if chain.len() >= self.config.min_fusion_ops {
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
|
||
let output = ops[*chain.last().unwrap()].output;
|
||
let dtype = op.dtype;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::ActivationChain,
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: FusionPattern::ActivationChain.expected_speedup(),
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
// ========================================================================
|
||
// Optimized detection methods using AnalysisContext
|
||
// ========================================================================
|
||
|
||
/// Optimized elementwise chain detection using pre-computed context
|
||
fn detect_elementwise_chains_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
|
||
if ctx.ops.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
let mut opportunities = Vec::with_capacity(ctx.ops.len() / 4);
|
||
let mut visited: HashSet<usize> = HashSet::with_capacity(ctx.ops.len());
|
||
|
||
// Find chain start points (fuseable ops with no fuseable producer)
|
||
for (idx, op) in ctx.ops.iter().enumerate() {
|
||
if visited.contains(&idx) || !op.op.is_fuseable() {
|
||
continue;
|
||
}
|
||
|
||
// Check if this is a chain start
|
||
let is_start = op.inputs.iter().all(|input| {
|
||
ctx.producer_of(*input)
|
||
.map(|producer_idx| !ctx.get_op(producer_idx).op.is_fuseable())
|
||
.unwrap_or(true)
|
||
});
|
||
|
||
if !is_start {
|
||
continue;
|
||
}
|
||
|
||
// Extend the chain greedily
|
||
let mut chain = Vec::with_capacity(self.config.max_fusion_ops);
|
||
chain.push(idx);
|
||
let mut chain_outputs: HashSet<TensorId> =
|
||
HashSet::with_capacity(self.config.max_fusion_ops);
|
||
chain_outputs.insert(op.output);
|
||
visited.insert(idx);
|
||
|
||
let mut current_idx = idx;
|
||
|
||
loop {
|
||
// Find the next operation in the chain
|
||
let consumers = ctx.consumers_of_op(current_idx);
|
||
let next_candidate = consumers
|
||
.iter()
|
||
.filter(|&&next_idx| {
|
||
!visited.contains(&next_idx)
|
||
&& ctx.get_op(next_idx).op.is_fuseable()
|
||
&& ctx.get_op(next_idx).inputs.iter().all(|input| {
|
||
chain_outputs.contains(input) || !ctx.is_internal(*input)
|
||
})
|
||
})
|
||
.next();
|
||
|
||
match next_candidate {
|
||
Some(&next_idx) => {
|
||
let next_op = ctx.get_op(next_idx);
|
||
chain.push(next_idx);
|
||
chain_outputs.insert(next_op.output);
|
||
visited.insert(next_idx);
|
||
current_idx = next_idx;
|
||
|
||
if chain.len() >= self.config.max_fusion_ops {
|
||
break;
|
||
}
|
||
}
|
||
None => break,
|
||
}
|
||
}
|
||
|
||
// Create opportunity if chain is long enough
|
||
if chain.len() >= self.config.min_fusion_ops {
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs_opt(&chain, ctx);
|
||
let output = ctx.get_op(*chain.last().unwrap()).output;
|
||
let dtype = ctx.get_op(chain[0]).dtype;
|
||
|
||
let base_speedup = FusionPattern::ElementwiseChain.expected_speedup();
|
||
let chain_bonus = 1.0 + (chain.len() as f64 - 2.0) * 0.1;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::ElementwiseChain,
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: base_speedup * chain_bonus,
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Optimized softmax pattern detection
|
||
fn detect_softmax_patterns_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
|
||
let mut opportunities = Vec::new();
|
||
|
||
for (exp_idx, exp_op) in ctx.ops.iter().enumerate() {
|
||
if !matches!(exp_op.op, StreamOpKind::Exp) {
|
||
continue;
|
||
}
|
||
|
||
let mut chain = vec![exp_idx];
|
||
let mut dim = None;
|
||
|
||
// Look backward for Sub(x, max) pattern
|
||
if let Some(&input_id) = exp_op.inputs.first() {
|
||
if let Some(sub_idx) = ctx.producer_of(input_id) {
|
||
if matches!(ctx.get_op(sub_idx).op, StreamOpKind::Sub) {
|
||
if ctx.get_op(sub_idx).inputs.len() >= 2 {
|
||
let second_input = ctx.get_op(sub_idx).inputs[1];
|
||
if let Some(max_idx) = ctx.producer_of(second_input) {
|
||
if matches!(ctx.get_op(max_idx).op, StreamOpKind::Max) {
|
||
chain.insert(0, max_idx);
|
||
chain.insert(1, sub_idx);
|
||
}
|
||
}
|
||
}
|
||
if chain.len() == 1 {
|
||
chain.insert(0, sub_idx);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Look forward for Sum operation
|
||
let exp_output = exp_op.output;
|
||
let sum_consumer = ctx.consumers_of_tensor(exp_output).iter().find(|&&idx| {
|
||
matches!(
|
||
ctx.get_op(idx).op,
|
||
StreamOpKind::Sum | StreamOpKind::SumDim { .. }
|
||
)
|
||
});
|
||
|
||
if let Some(&sum_idx) = sum_consumer {
|
||
chain.push(sum_idx);
|
||
|
||
if let StreamOpKind::SumDim { dim: d } = &ctx.get_op(sum_idx).op {
|
||
dim = Some(*d);
|
||
}
|
||
|
||
let sum_output = ctx.get_op(sum_idx).output;
|
||
let div_consumer = ctx
|
||
.consumers_of_tensor(sum_output)
|
||
.iter()
|
||
.find(|&&idx| matches!(ctx.get_op(idx).op, StreamOpKind::Div));
|
||
|
||
if let Some(&div_idx) = div_consumer {
|
||
chain.push(div_idx);
|
||
}
|
||
}
|
||
|
||
if chain.len() >= 2
|
||
&& chain.iter().any(|&i| {
|
||
matches!(
|
||
ctx.get_op(i).op,
|
||
StreamOpKind::Sum | StreamOpKind::SumDim { .. }
|
||
)
|
||
})
|
||
{
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs_opt(&chain, ctx);
|
||
let output = ctx.get_op(*chain.last().unwrap()).output;
|
||
let dtype = exp_op.dtype;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::SoftmaxPattern { dim },
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: FusionPattern::SoftmaxPattern { dim }.expected_speedup(),
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Optimized layer norm pattern detection
|
||
fn detect_layer_norm_patterns_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
|
||
let mut opportunities = Vec::new();
|
||
let eps = 1e-5f32;
|
||
|
||
for (mean_idx, mean_op) in ctx.ops.iter().enumerate() {
|
||
if !matches!(
|
||
mean_op.op,
|
||
StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
|
||
) {
|
||
continue;
|
||
}
|
||
|
||
let mut chain = vec![mean_idx];
|
||
let mut dim = None;
|
||
|
||
if let StreamOpKind::MeanDim { dim: d } = &mean_op.op {
|
||
dim = Some(*d);
|
||
}
|
||
|
||
let mean_output = mean_op.output;
|
||
let sub_consumer = ctx.consumers_of_tensor(mean_output).iter().find(|&&idx| {
|
||
matches!(ctx.get_op(idx).op, StreamOpKind::Sub)
|
||
&& ctx.get_op(idx).inputs.get(1) == Some(&mean_output)
|
||
});
|
||
|
||
if let Some(&sub_idx) = sub_consumer {
|
||
chain.push(sub_idx);
|
||
|
||
let sub_output = ctx.get_op(sub_idx).output;
|
||
let square_consumer = ctx
|
||
.consumers_of_tensor(sub_output)
|
||
.iter()
|
||
.find(|&&idx| matches!(ctx.get_op(idx).op, StreamOpKind::Square));
|
||
|
||
if let Some(&square_idx) = square_consumer {
|
||
chain.push(square_idx);
|
||
|
||
let square_output = ctx.get_op(square_idx).output;
|
||
let mean2_consumer =
|
||
ctx.consumers_of_tensor(square_output).iter().find(|&&idx| {
|
||
matches!(
|
||
ctx.get_op(idx).op,
|
||
StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
|
||
)
|
||
});
|
||
|
||
if let Some(&mean2_idx) = mean2_consumer {
|
||
chain.push(mean2_idx);
|
||
|
||
let var_output = ctx.get_op(mean2_idx).output;
|
||
if let Some(&add_eps_idx) = ctx
|
||
.consumers_of_tensor(var_output)
|
||
.iter()
|
||
.find(|&&idx| matches!(ctx.get_op(idx).op, StreamOpKind::AddScalar(_)))
|
||
{
|
||
chain.push(add_eps_idx);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if chain.len() >= 3 {
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs_opt(&chain, ctx);
|
||
let output = ctx.get_op(*chain.last().unwrap()).output;
|
||
let dtype = mean_op.dtype;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::LayerNormPattern { eps, dim },
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: FusionPattern::LayerNormPattern { eps, dim }
|
||
.expected_speedup(),
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Optimized activation chain detection
|
||
fn detect_activation_chains_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
|
||
let mut opportunities = Vec::new();
|
||
let mut visited: HashSet<usize> = HashSet::with_capacity(ctx.ops.len());
|
||
|
||
for (idx, op) in ctx.ops.iter().enumerate() {
|
||
if visited.contains(&idx) {
|
||
continue;
|
||
}
|
||
|
||
let is_activation = matches!(
|
||
op.op,
|
||
StreamOpKind::GELU
|
||
| StreamOpKind::SiLU
|
||
| StreamOpKind::Sigmoid
|
||
| StreamOpKind::Tanh
|
||
| StreamOpKind::ReLU
|
||
);
|
||
|
||
if !is_activation || !op.op.is_fuseable() {
|
||
continue;
|
||
}
|
||
|
||
let mut chain = vec![idx];
|
||
visited.insert(idx);
|
||
|
||
// Look backward
|
||
if let Some(&input_id) = op.inputs.first() {
|
||
if let Some(prev_idx) = ctx.producer_of(input_id) {
|
||
if ctx.get_op(prev_idx).op.is_fuseable() && !visited.contains(&prev_idx) {
|
||
chain.insert(0, prev_idx);
|
||
visited.insert(prev_idx);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Look forward
|
||
for (next_idx, next_op) in ctx.ops.iter().enumerate() {
|
||
if next_op.inputs.contains(&op.output)
|
||
&& next_op.op.is_fuseable()
|
||
&& !visited.contains(&next_idx)
|
||
{
|
||
chain.push(next_idx);
|
||
visited.insert(next_idx);
|
||
break;
|
||
}
|
||
}
|
||
|
||
if chain.len() >= self.config.min_fusion_ops {
|
||
let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
|
||
let external_inputs = self.find_external_inputs_opt(&chain, ctx);
|
||
let output = ctx.get_op(*chain.last().unwrap()).output;
|
||
let dtype = op.dtype;
|
||
|
||
opportunities.push(FusionOpportunity {
|
||
pattern: FusionPattern::ActivationChain,
|
||
operation_indices: chain,
|
||
operations: chain_ops,
|
||
external_inputs,
|
||
output,
|
||
estimated_speedup: FusionPattern::ActivationChain.expected_speedup(),
|
||
dtype,
|
||
});
|
||
}
|
||
}
|
||
|
||
opportunities
|
||
}
|
||
|
||
/// Optimized external input detection using context
|
||
fn find_external_inputs_opt(
|
||
&self,
|
||
chain: &[usize],
|
||
ctx: &AnalysisContext<'_>,
|
||
) -> Vec<TensorId> {
|
||
let chain_set: HashSet<_> = chain.iter().copied().collect();
|
||
let mut external = Vec::with_capacity(chain.len());
|
||
|
||
for &idx in chain {
|
||
for &input in &ctx.get_op(idx).inputs {
|
||
let is_external = ctx
|
||
.producer_of(input)
|
||
.map(|producer_idx| !chain_set.contains(&producer_idx))
|
||
.unwrap_or(true);
|
||
|
||
if is_external && !external.contains(&input) {
|
||
external.push(input);
|
||
}
|
||
}
|
||
}
|
||
|
||
external
|
||
}
|
||
|
||
/// Find external inputs for a chain of operations
|
||
fn find_external_inputs(
|
||
&self,
|
||
chain: &[usize],
|
||
ops: &[StreamOperation],
|
||
output_to_op: &HashMap<TensorId, usize>,
|
||
) -> Vec<TensorId> {
|
||
let chain_set: HashSet<_> = chain.iter().copied().collect();
|
||
let mut external = Vec::new();
|
||
|
||
for &idx in chain {
|
||
for &input in &ops[idx].inputs {
|
||
// External if not produced by an op in the chain
|
||
let is_external = output_to_op
|
||
.get(&input)
|
||
.map(|&producer_idx| !chain_set.contains(&producer_idx))
|
||
.unwrap_or(true);
|
||
|
||
if is_external && !external.contains(&input) {
|
||
external.push(input);
|
||
}
|
||
}
|
||
}
|
||
|
||
external
|
||
}
|
||
|
||
/// Remove overlapping opportunities (greedy: keep highest benefit)
|
||
fn remove_overlaps(&self, mut opportunities: Vec<FusionOpportunity>) -> Vec<FusionOpportunity> {
|
||
let mut result = Vec::new();
|
||
let mut used_ops: HashSet<usize> = HashSet::new();
|
||
|
||
for opp in opportunities {
|
||
// Check if any operation in this opportunity is already used
|
||
let has_overlap = opp
|
||
.operation_indices
|
||
.iter()
|
||
.any(|idx| used_ops.contains(idx));
|
||
|
||
if !has_overlap {
|
||
// Mark operations as used
|
||
for &idx in &opp.operation_indices {
|
||
used_ops.insert(idx);
|
||
}
|
||
result.push(opp);
|
||
}
|
||
}
|
||
|
||
result
|
||
}
|
||
|
||
/// Analyze a single operation for fusion potential
|
||
pub fn analyze_operation(&self, op: &StreamOpKind) -> OperationAnalysis {
|
||
OperationAnalysis {
|
||
is_fuseable: op.is_fuseable(),
|
||
is_sync_point: op.is_sync_point(),
|
||
fusion_benefit: op.fusion_benefit(),
|
||
preferred_patterns: if op.is_fuseable() {
|
||
vec![
|
||
FusionPattern::ElementwiseChain,
|
||
FusionPattern::ActivationChain,
|
||
]
|
||
} else {
|
||
vec![]
|
||
},
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Analysis result for a single operation
|
||
#[derive(Debug)]
|
||
pub struct OperationAnalysis {
|
||
/// Whether the operation can be fused
|
||
pub is_fuseable: bool,
|
||
/// Whether this is a sync point (triggers flush)
|
||
pub is_sync_point: bool,
|
||
/// Benefit score for fusion (0.0-1.0)
|
||
pub fusion_benefit: f64,
|
||
/// Preferred fusion patterns for this operation
|
||
pub preferred_patterns: Vec<FusionPattern>,
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::kernel::TensorId;
|
||
|
||
fn make_op(op: StreamOpKind, inputs: Vec<TensorId>, output: TensorId) -> StreamOperation {
|
||
StreamOperation::new(op, inputs, output, vec![1024], DType::F32)
|
||
}
|
||
|
||
#[test]
|
||
fn test_pattern_speedup() {
|
||
assert!(FusionPattern::ElementwiseChain.expected_speedup() > 1.0);
|
||
assert!(
|
||
FusionPattern::SoftmaxPattern { dim: None }.expected_speedup()
|
||
> FusionPattern::ElementwiseChain.expected_speedup()
|
||
);
|
||
assert!(
|
||
FusionPattern::LayerNormPattern {
|
||
eps: 1e-5,
|
||
dim: None
|
||
}
|
||
.expected_speedup()
|
||
> FusionPattern::ElementwiseChain.expected_speedup()
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_analyzer_creation() {
|
||
let analyzer = FusionAnalyzer::with_defaults();
|
||
assert!(analyzer.config.enabled);
|
||
}
|
||
|
||
#[test]
|
||
fn test_empty_stream() {
|
||
let analyzer = FusionAnalyzer::with_defaults();
|
||
let stream = OperationStream::with_defaults();
|
||
|
||
let opportunities = analyzer.analyze(&stream);
|
||
assert!(opportunities.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_elementwise_chain_detection() {
|
||
let analyzer = FusionAnalyzer::new(FusionConfig::default().with_min_fusion_ops(2));
|
||
let mut stream = OperationStream::with_defaults();
|
||
|
||
let a = TensorId::new();
|
||
let b = TensorId::new();
|
||
let c = TensorId::new();
|
||
let d = TensorId::new();
|
||
let e = TensorId::new();
|
||
|
||
stream.register_input(a);
|
||
stream.register_input(b);
|
||
|
||
// Chain: add(a, b) -> mul(c, a) -> relu(d)
|
||
stream.record(make_op(StreamOpKind::Add, vec![a, b], c));
|
||
stream.record(make_op(StreamOpKind::Mul, vec![c, a], d));
|
||
stream.record(make_op(StreamOpKind::ReLU, vec![d], e));
|
||
|
||
let opportunities = analyzer.analyze(&stream);
|
||
assert!(!opportunities.is_empty());
|
||
|
||
let first = &opportunities[0];
|
||
assert_eq!(first.pattern, FusionPattern::ElementwiseChain);
|
||
assert!(first.operation_count() >= 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_sync_point_breaks_chain() {
|
||
let analyzer = FusionAnalyzer::new(FusionConfig::default().with_min_fusion_ops(2));
|
||
let mut stream = OperationStream::with_defaults();
|
||
|
||
let a = TensorId::new();
|
||
let b = TensorId::new();
|
||
let c = TensorId::new();
|
||
let d = TensorId::new();
|
||
|
||
stream.register_input(a);
|
||
|
||
// add -> matmul (sync point) -> mul
|
||
stream.record(make_op(StreamOpKind::Add, vec![a, a], b));
|
||
stream.record(make_op(StreamOpKind::MatMul, vec![b], c));
|
||
stream.record(make_op(StreamOpKind::Mul, vec![c, c], d));
|
||
|
||
let opportunities = analyzer.analyze(&stream);
|
||
|
||
// MatMul should break the chain
|
||
for opp in &opportunities {
|
||
assert!(
|
||
!opp.operations
|
||
.iter()
|
||
.any(|op| matches!(op, StreamOpKind::MatMul))
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_opportunity_worthwhile() {
|
||
let opp = FusionOpportunity {
|
||
pattern: FusionPattern::ElementwiseChain,
|
||
operation_indices: vec![0, 1],
|
||
operations: vec![StreamOpKind::Add, StreamOpKind::Mul],
|
||
external_inputs: vec![TensorId::new()],
|
||
output: TensorId::new(),
|
||
estimated_speedup: 2.0,
|
||
dtype: DType::F32,
|
||
};
|
||
|
||
assert!(opp.is_worthwhile(2));
|
||
assert!(!opp.is_worthwhile(3));
|
||
}
|
||
|
||
#[test]
|
||
fn test_kernel_signature_generation() {
|
||
let opp = FusionOpportunity {
|
||
pattern: FusionPattern::ElementwiseChain,
|
||
operation_indices: vec![0, 1, 2],
|
||
operations: vec![StreamOpKind::Add, StreamOpKind::Mul, StreamOpKind::ReLU],
|
||
external_inputs: vec![TensorId::new(), TensorId::new()],
|
||
output: TensorId::new(),
|
||
estimated_speedup: 2.5,
|
||
dtype: DType::F32,
|
||
};
|
||
|
||
let sig = opp.to_kernel_signature();
|
||
assert_eq!(sig.ops.len(), 3);
|
||
assert_eq!(sig.num_inputs, 2);
|
||
assert_eq!(sig.dtype, DType::F32);
|
||
}
|
||
|
||
#[test]
|
||
fn test_operation_analysis() {
|
||
let analyzer = FusionAnalyzer::with_defaults();
|
||
|
||
let add_analysis = analyzer.analyze_operation(&StreamOpKind::Add);
|
||
assert!(add_analysis.is_fuseable);
|
||
assert!(!add_analysis.is_sync_point);
|
||
|
||
let matmul_analysis = analyzer.analyze_operation(&StreamOpKind::MatMul);
|
||
assert!(!matmul_analysis.is_fuseable);
|
||
assert!(matmul_analysis.is_sync_point);
|
||
}
|
||
}
|