use smallvec::SmallVec;
use crate::{Element, Tensor};
use super::batch_norm;
use super::pattern::Pattern;
use super::reduce_window;
use super::view::View;
use super::window;
#[derive(Debug, Clone)]
pub(crate) struct Candidate {
pub(crate) pattern: Pattern,
pub(crate) root: usize,
pub(crate) interiors: SmallVec<[usize; 8]>,
pub(crate) named: SmallVec<[usize; 4]>,
}
#[derive(Debug, Clone)]
pub(crate) struct Candidates {
length: usize,
all: Vec<Candidate>,
}
impl Candidates {
pub(crate) fn discover<E: Element>(view: &View<Tensor<E>>) -> Self {
let mut all = Vec::new();
discover_one(view, &mut all, window::match_at);
discover_one(view, &mut all, reduce_window::match_at);
discover_one(view, &mut all, batch_norm::match_training);
discover_one(view, &mut all, batch_norm::match_inference);
Self {
length: view.len(),
all,
}
}
pub(crate) fn length(&self) -> usize {
self.length
}
pub(crate) fn iter(&self) -> impl Iterator<Item = &Candidate> {
self.all.iter()
}
}
fn discover_one<E: Element>(
view: &View<Tensor<E>>,
all: &mut Vec<Candidate>,
matcher: fn(usize, &View<Tensor<E>>) -> Option<Candidate>,
) {
for index in 0..view.len() {
if !view.wanted(index) {
continue;
}
let Some(candidate) = matcher(index, view) else {
continue;
};
if !view.closed(index, &candidate.interiors, &candidate.named) {
continue;
}
all.push(candidate);
}
}