use fuel_ir::dispatch::OpKind;
use fuel_ir::{DType, Layout};
use super::candidate::Candidate;
pub trait AlternativeFilter: Send + Sync {
fn filter(&self, alts: &[Candidate], ctx: &FilterContext) -> Vec<usize>;
fn classification(&self) -> FilterClass;
fn name(&self) -> &'static str;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FilterClass {
Hard,
Soft { min_remaining: usize },
}
#[derive(Clone, Copy, Debug)]
pub struct FilterContext<'a> {
pub op: OpKind,
pub dtypes: &'a [DType],
pub input_layouts: &'a [Layout],
pub tag: Option<&'a str>,
}
impl<'a> FilterContext<'a> {
pub fn new(op: OpKind, dtypes: &'a [DType], input_layouts: &'a [Layout]) -> Self {
Self {
op,
dtypes,
input_layouts,
tag: None,
}
}
pub fn with_tag(mut self, tag: &'a str) -> Self {
self.tag = Some(tag);
self
}
pub fn summary(&self) -> String {
let layouts = self.input_layouts.len();
match self.tag {
Some(t) => format!(
"op={:?} dtypes={:?} input_layouts={} tag={}",
self.op, self.dtypes, layouts, t,
),
None => format!(
"op={:?} dtypes={:?} input_layouts={}",
self.op, self.dtypes, layouts,
),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::Shape;
pub(crate) struct MockFilter {
pub keep: Vec<usize>,
pub class: FilterClass,
pub name: &'static str,
}
impl AlternativeFilter for MockFilter {
fn filter(&self, _alts: &[Candidate], _ctx: &FilterContext) -> Vec<usize> {
self.keep.clone()
}
fn classification(&self) -> FilterClass {
self.class
}
fn name(&self) -> &'static str {
self.name
}
}
#[test]
fn filter_trait_is_dyn_compatible() {
let f: Box<dyn AlternativeFilter> = Box::new(MockFilter {
keep: vec![0, 2],
class: FilterClass::Soft { min_remaining: 1 },
name: "mock",
});
let layouts = [Layout::contiguous(Shape::from(vec![4]))];
let ctx = FilterContext::new(OpKind::AddElementwise, &[DType::F32], &layouts);
assert_eq!(f.filter(&[], &ctx), vec![0, 2]);
assert_eq!(f.name(), "mock");
assert!(matches!(
f.classification(),
FilterClass::Soft { min_remaining: 1 }
));
}
#[test]
fn filter_context_summary_includes_op_and_dtypes() {
let layouts = [Layout::contiguous(Shape::from(vec![4]))];
let ctx = FilterContext::new(
OpKind::MatMul,
&[DType::F32, DType::F32, DType::F32],
&layouts,
);
let s = ctx.summary();
assert!(s.contains("MatMul"));
assert!(s.contains("F32"));
assert!(s.contains("input_layouts=1"));
}
#[test]
fn filter_context_with_tag_appears_in_summary() {
let layouts: [Layout; 0] = [];
let ctx = FilterContext::new(OpKind::AddElementwise, &[DType::F32], &layouts)
.with_tag("NodeId(42)");
assert!(ctx.summary().contains("NodeId(42)"));
}
#[test]
fn filter_class_variants_distinguishable() {
assert_ne!(FilterClass::Hard, FilterClass::Soft { min_remaining: 1 },);
assert_ne!(
FilterClass::Soft { min_remaining: 1 },
FilterClass::Soft { min_remaining: 2 },
);
}
}