Save inference parity implementation and evaluation harness
This commit is contained in:
Georg Bauer
@@ -0,0 +1,807 @@
|
||||
//! Pinned compile.cpp: DFS/parents, available-scalar coalescing and three
|
||||
//! depth-one CSE passes, then original fusion boundaries and Compiled lowering.
|
||||
//! Rebuild descriptors, never mutate captured model data.
|
||||
use super::array::{Array, Dtype};
|
||||
use super::ops::Operation;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::rc::Rc;
|
||||
|
||||
struct Node {
|
||||
array: Array,
|
||||
inputs: Vec<usize>,
|
||||
}
|
||||
|
||||
struct Tape {
|
||||
nodes: HashMap<usize, Node>,
|
||||
order: Vec<usize>,
|
||||
parents: HashMap<usize, Vec<(usize, usize)>>,
|
||||
outputs: HashSet<usize>,
|
||||
}
|
||||
|
||||
impl Tape {
|
||||
/// compile.cpp::split_one: only parents already in the candidate section
|
||||
/// move to the new Broadcast descriptor. Other consumers keep the view.
|
||||
fn split_broadcast(&mut self, id: usize, section: &HashSet<usize>) -> Result<usize, String> {
|
||||
let n = &self.nodes[&id];
|
||||
let inputs = n.inputs.clone();
|
||||
let args = inputs
|
||||
.iter()
|
||||
.map(|id| self.nodes[id].array.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let a = Array::make_shared_operations(
|
||||
n.array.stream().unwrap(),
|
||||
&args,
|
||||
&[n.array.layout().clone()],
|
||||
n.array.operation().unwrap(),
|
||||
)?
|
||||
.remove(0);
|
||||
let split = a.id();
|
||||
self.nodes.insert(split, Node { array: a, inputs });
|
||||
let mut moved = Vec::new();
|
||||
self.parents
|
||||
.get_mut(&id)
|
||||
.unwrap()
|
||||
.retain(|&(parent, index)| {
|
||||
if section.contains(&parent) {
|
||||
self.nodes.get_mut(&parent).unwrap().inputs[index] = split;
|
||||
moved.push((parent, index));
|
||||
false
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
self.parents.insert(split, moved);
|
||||
Ok(split)
|
||||
}
|
||||
|
||||
fn collect_fusable(
|
||||
&mut self,
|
||||
id: usize,
|
||||
depth: usize,
|
||||
stream: super::stream::Stream,
|
||||
shape: &[i32],
|
||||
section: &mut HashSet<usize>,
|
||||
candidates: &mut HashSet<usize>,
|
||||
) -> Result<(), String> {
|
||||
if section.contains(&id) {
|
||||
return Ok(());
|
||||
}
|
||||
let a = &self.nodes[&id].array;
|
||||
let op = a.operation();
|
||||
if depth >= 11
|
||||
|| op
|
||||
.as_ref()
|
||||
.is_none_or(|p| super::fused::operator(p).is_none())
|
||||
|| a.stream() != Some(stream)
|
||||
|| (self.outputs.contains(&id) && a.layout().shape() != shape)
|
||||
{
|
||||
candidates.insert(id);
|
||||
return Ok(());
|
||||
}
|
||||
if depth > 0 && !self.parents[&id].iter().all(|(p, _)| section.contains(p)) {
|
||||
if matches!(&**op.as_ref().unwrap(), Operation::Broadcast) && candidates.len() < 24 {
|
||||
let split = self.split_broadcast(id, section)?;
|
||||
self.collect_fusable(split, depth, stream, shape, section, candidates)?;
|
||||
} else {
|
||||
candidates.insert(id);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
if self.outputs.contains(&id) {
|
||||
candidates.insert(id);
|
||||
} else {
|
||||
candidates.remove(&id);
|
||||
}
|
||||
if candidates.len() >= 24 {
|
||||
return Ok(());
|
||||
}
|
||||
section.insert(id);
|
||||
for input in self.nodes[&id].inputs.clone() {
|
||||
self.collect_fusable(input, depth + 1, stream, shape, section, candidates)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn fused_tape(&self, root: usize, section: &HashSet<usize>) -> (Vec<usize>, Vec<usize>) {
|
||||
let mut inputs = Vec::new();
|
||||
let mut seen_inputs = HashSet::new();
|
||||
let mut tape = Vec::new();
|
||||
let mut seen = HashSet::new();
|
||||
let mut stack = vec![(root, false)];
|
||||
while let Some((id, expanded)) = stack.pop() {
|
||||
if !section.contains(&id) {
|
||||
if seen_inputs.insert(id) {
|
||||
inputs.push(id);
|
||||
}
|
||||
} else if expanded {
|
||||
tape.push(id);
|
||||
} else if seen.insert(id) {
|
||||
stack.push((id, true));
|
||||
stack.extend(self.nodes[&id].inputs.iter().rev().map(|&i| (i, false)));
|
||||
}
|
||||
}
|
||||
(inputs, tape)
|
||||
}
|
||||
|
||||
fn fuse(&mut self, inputs: &[Array], outputs: &[Array]) -> Result<Vec<Array>, String> {
|
||||
let input_ids = inputs.iter().map(Array::id).collect::<HashSet<_>>();
|
||||
let mut output_map = outputs
|
||||
.iter()
|
||||
.map(|a| (a.id(), a.clone()))
|
||||
.collect::<HashMap<_, _>>();
|
||||
let mut keep = Vec::new();
|
||||
let mut global = HashSet::new();
|
||||
for root in self.order.clone().into_iter().rev() {
|
||||
if global.contains(&root) {
|
||||
continue;
|
||||
}
|
||||
let a = &self.nodes[&root].array;
|
||||
let mut section = HashSet::new();
|
||||
if a.operation()
|
||||
.is_some_and(|p| !matches!(&*p, Operation::Broadcast))
|
||||
{
|
||||
let stream = a.stream().unwrap();
|
||||
let shape = a.layout().shape().to_vec();
|
||||
self.collect_fusable(root, 0, stream, &shape, &mut section, &mut HashSet::new())?;
|
||||
}
|
||||
if section.len() <= 1 {
|
||||
keep.push(root);
|
||||
continue;
|
||||
}
|
||||
let (inputs, tape) = self.fused_tape(root, §ion);
|
||||
let mut old_outputs = Vec::new();
|
||||
for &id in &tape[..tape.len() - 1] {
|
||||
if self.outputs.contains(&id) {
|
||||
old_outputs.push(id);
|
||||
self.parents
|
||||
.entry(id)
|
||||
.or_default()
|
||||
.retain(|(p, _)| !section.contains(p));
|
||||
} else {
|
||||
self.parents.remove(&id);
|
||||
}
|
||||
global.insert(id);
|
||||
}
|
||||
old_outputs.push(root);
|
||||
let layouts = old_outputs
|
||||
.iter()
|
||||
.map(|id| self.nodes[id].array.layout().clone())
|
||||
.collect::<Vec<_>>();
|
||||
if layouts
|
||||
.iter()
|
||||
.any(|l| l.shape() != layouts.last().unwrap().shape())
|
||||
{
|
||||
return Err("compilation tried to fuse different output shapes".into());
|
||||
}
|
||||
let args = inputs
|
||||
.iter()
|
||||
.map(|id| self.nodes[id].array.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let constants = args
|
||||
.iter()
|
||||
.map(|a| {
|
||||
a.layout().size() == 1
|
||||
&& a.operation().is_none()
|
||||
&& !input_ids.contains(&a.id())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let slots = inputs
|
||||
.iter()
|
||||
.chain(&tape)
|
||||
.enumerate()
|
||||
.map(|(i, &id)| (id, i))
|
||||
.collect::<HashMap<_, _>>();
|
||||
let statements = tape
|
||||
.iter()
|
||||
.map(|id| {
|
||||
let node = &self.nodes[id];
|
||||
super::fused::Statement {
|
||||
dtype: node.array.layout().dtype(),
|
||||
operator: super::fused::operator(&node.array.operation().unwrap()).unwrap(),
|
||||
inputs: node.inputs.iter().map(|i| slots[i]).collect(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let expr = super::fused::Expression::new(
|
||||
&args,
|
||||
statements,
|
||||
old_outputs.iter().map(|id| slots[id]).collect(),
|
||||
&constants,
|
||||
)?;
|
||||
let new = Array::make_operations(
|
||||
self.nodes[&root].array.stream().unwrap(),
|
||||
&args,
|
||||
&layouts,
|
||||
Operation::Compiled(expr),
|
||||
)?;
|
||||
keep.push(new.last().unwrap().id());
|
||||
for output in &new {
|
||||
self.nodes.insert(
|
||||
output.id(),
|
||||
Node {
|
||||
array: output.clone(),
|
||||
inputs: inputs.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
for (position, input) in inputs.iter().enumerate() {
|
||||
let parents = self.parents.entry(*input).or_default();
|
||||
parents.retain(|(p, _)| !section.contains(p));
|
||||
parents.extend(new.iter().map(|a| (a.id(), position)));
|
||||
}
|
||||
for (&old, new) in old_outputs.iter().zip(new) {
|
||||
self.merge_one(new.id(), old);
|
||||
if let Some(output) = output_map.get_mut(&old) {
|
||||
*output = new;
|
||||
}
|
||||
}
|
||||
}
|
||||
keep.reverse();
|
||||
self.order = keep;
|
||||
Ok(outputs
|
||||
.iter()
|
||||
.map(|a| output_map[&a.id()].clone())
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn capture(outputs: &[Array], original: &[Array]) -> Result<Self, String> {
|
||||
let mut tape = Self {
|
||||
nodes: HashMap::new(),
|
||||
order: Vec::new(),
|
||||
parents: HashMap::new(),
|
||||
outputs: outputs.iter().map(Array::id).collect(),
|
||||
};
|
||||
let forbidden = original.iter().map(Array::id).collect::<HashSet<_>>();
|
||||
let mut seen = HashSet::new();
|
||||
let mut stack = outputs
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|a| (a.clone(), false))
|
||||
.collect::<Vec<_>>();
|
||||
while let Some((a, expanded)) = stack.pop() {
|
||||
let id = a.id();
|
||||
if forbidden.contains(&id) {
|
||||
return Err("compiled function has uncaptured inputs".into());
|
||||
}
|
||||
if seen.contains(&id) {
|
||||
continue;
|
||||
}
|
||||
if expanded {
|
||||
seen.extend(a.outputs().iter().map(Array::id));
|
||||
tape.order.push(id);
|
||||
continue;
|
||||
}
|
||||
let inputs = a.inputs().iter().map(Array::id).collect::<Vec<_>>();
|
||||
// Reference registers the visited output first, then its siblings,
|
||||
// before descending into its inputs (not sorted by array id).
|
||||
let siblings = a.siblings();
|
||||
for (position, &input) in inputs.iter().enumerate() {
|
||||
let parents = tape.parents.entry(input).or_default();
|
||||
parents.push((id, position));
|
||||
parents.extend(siblings.iter().map(|s| (s.id(), position)));
|
||||
}
|
||||
for output in a.outputs() {
|
||||
tape.nodes.insert(
|
||||
output.id(),
|
||||
Node {
|
||||
array: output,
|
||||
inputs: inputs.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
stack.push((a.clone(), true));
|
||||
stack.extend(a.inputs().iter().rev().map(|i| (i.clone(), false)));
|
||||
}
|
||||
Ok(tape)
|
||||
}
|
||||
|
||||
fn merge_one(&mut self, dst: usize, src: usize) {
|
||||
let Some(parents) = self.parents.remove(&src) else {
|
||||
return;
|
||||
};
|
||||
for &(parent, index) in &parents {
|
||||
self.nodes.get_mut(&parent).unwrap().inputs[index] = dst;
|
||||
}
|
||||
let pairs = self.parents.entry(dst).or_default();
|
||||
pairs.extend(parents);
|
||||
pairs.retain(|&(parent, _)| parent != src);
|
||||
}
|
||||
|
||||
fn merge(&mut self, dst: usize, src: usize) {
|
||||
let dst = self.nodes[&dst].array.outputs();
|
||||
let src = self.nodes[&src].array.outputs();
|
||||
for (dst, src) in dst.iter().zip(src.iter()) {
|
||||
self.merge_one(dst.id(), src.id());
|
||||
}
|
||||
}
|
||||
|
||||
fn scalars(&mut self) -> Result<(), String> {
|
||||
let mut scalars = HashMap::<(u64, Dtype), usize>::new();
|
||||
let mut keep = Vec::new();
|
||||
for id in self.order.clone() {
|
||||
let a = &self.nodes[&id].array;
|
||||
if a.is_available()? && a.layout().shape().is_empty() && a.has_data() {
|
||||
let dtype = a.layout().dtype();
|
||||
let mut bits = [0; 8];
|
||||
a.buffer().read(a.offset(), &mut bits[..dtype.itemsize()])?;
|
||||
let dst = *scalars
|
||||
.entry((u64::from_le_bytes(bits), dtype))
|
||||
.or_insert(id);
|
||||
if dst != id {
|
||||
self.merge(dst, id);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
keep.push(id);
|
||||
}
|
||||
self.order = keep;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn equivalent(&self, a: usize, b: usize) -> bool {
|
||||
let a = &self.nodes[&a];
|
||||
let b = &self.nodes[&b];
|
||||
let (Some(pa), Some(pb)) = (a.array.operation(), b.array.operation()) else {
|
||||
return false;
|
||||
};
|
||||
// The original explicitly excludes identical primitive pointers and
|
||||
// does not compare streams at this stage. Fusion checks streams later.
|
||||
!Rc::ptr_eq(&pa, &pb)
|
||||
&& a.inputs == b.inputs
|
||||
&& equivalent_operation(&pa, &pb, &a.array, &b.array)
|
||||
}
|
||||
|
||||
fn merge_parents(&mut self, id: usize, order: &HashMap<usize, usize>) -> bool {
|
||||
let Some(parents) = self.parents.get(&id).cloned() else {
|
||||
return !self.outputs.contains(&id);
|
||||
};
|
||||
let mut mask = vec![false; parents.len()];
|
||||
let groups = if parents.len() > 100 {
|
||||
// Same high-fanout grouping as VecU64Hash, using Rust's native
|
||||
// hash table. Binary/unary variants are distinguished by the
|
||||
// final equivalence check, including their primitive parameters.
|
||||
let mut groups = HashMap::new();
|
||||
for (i, &(parent, _)) in parents.iter().enumerate() {
|
||||
let node = &self.nodes[&parent];
|
||||
let op = node.array.operation().unwrap();
|
||||
groups
|
||||
.entry((node.inputs.clone(), std::mem::discriminant(&*op)))
|
||||
.or_insert_with(Vec::new)
|
||||
.push(i);
|
||||
}
|
||||
groups.into_values().collect::<Vec<_>>()
|
||||
} else {
|
||||
vec![(0..parents.len()).collect()]
|
||||
};
|
||||
for group in groups {
|
||||
for (i, &left) in group.iter().enumerate() {
|
||||
if mask[left] {
|
||||
continue;
|
||||
}
|
||||
for &right in &group[i + 1..] {
|
||||
if mask[right] {
|
||||
continue;
|
||||
}
|
||||
let (mut dst, mut src) = (left, right);
|
||||
if order.get(&parents[src].0).copied().unwrap_or(0)
|
||||
< order.get(&parents[dst].0).copied().unwrap_or(0)
|
||||
{
|
||||
std::mem::swap(&mut dst, &mut src);
|
||||
}
|
||||
let (dst_id, src_id) = (parents[dst].0, parents[src].0);
|
||||
if src_id != dst_id
|
||||
&& !self.outputs.contains(&src_id)
|
||||
&& self.equivalent(src_id, dst_id)
|
||||
{
|
||||
self.merge(dst_id, src_id);
|
||||
mask[src] = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let pairs = self.parents.get_mut(&id).unwrap();
|
||||
let mut i = 0;
|
||||
pairs.retain(|_| {
|
||||
let keep = !mask[i];
|
||||
i += 1;
|
||||
keep
|
||||
});
|
||||
false
|
||||
}
|
||||
|
||||
fn simplify(&mut self) -> Result<(), String> {
|
||||
self.scalars()?;
|
||||
// Copy/StopGradient are absent from this inference graph API: their
|
||||
// no-op contract is represented by returning the input descriptor.
|
||||
let order = self
|
||||
.order
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &id)| (id, i))
|
||||
.collect();
|
||||
for _ in 0..3 {
|
||||
let mut keep = Vec::new();
|
||||
for id in self.order.clone() {
|
||||
let outputs = self.nodes[&id].array.outputs();
|
||||
let mut discard = self.merge_parents(id, &order);
|
||||
for sibling in outputs.iter().filter(|a| a.id() != id) {
|
||||
discard &= self.merge_parents(sibling.id(), &order);
|
||||
}
|
||||
if !discard {
|
||||
keep.push(id);
|
||||
}
|
||||
}
|
||||
self.order = keep;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn rebuild(&self, outputs: &[Array]) -> Result<Vec<Array>, String> {
|
||||
let mut values = HashMap::<usize, Array>::new();
|
||||
for &id in &self.order {
|
||||
let node = &self.nodes[&id];
|
||||
let a = &node.array;
|
||||
let Some(op) = a.operation() else {
|
||||
values.insert(id, a.clone());
|
||||
continue;
|
||||
};
|
||||
let args = node
|
||||
.inputs
|
||||
.iter()
|
||||
.map(|id| {
|
||||
values
|
||||
.get(id)
|
||||
.cloned()
|
||||
.ok_or("simplified input absent from tape".into())
|
||||
})
|
||||
.collect::<Result<Vec<_>, String>>()?;
|
||||
let old = a.outputs();
|
||||
let layouts = old.iter().map(|o| o.layout().clone()).collect::<Vec<_>>();
|
||||
let result = Array::make_shared_operations(
|
||||
a.stream().ok_or("primitive has no stream")?,
|
||||
&args,
|
||||
&layouts,
|
||||
op,
|
||||
)?;
|
||||
for (old, new) in old.iter().zip(result) {
|
||||
values.insert(old.id(), new);
|
||||
}
|
||||
}
|
||||
outputs
|
||||
.iter()
|
||||
.map(|a| {
|
||||
values
|
||||
.get(&a.id())
|
||||
.cloned()
|
||||
.or_else(|| a.operation().is_none().then(|| a.clone()))
|
||||
.ok_or("simplified output absent from tape".into())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn simplify(outputs: &[Array], original_inputs: &[Array]) -> Result<Vec<Array>, String> {
|
||||
let mut tape = Tape::capture(outputs, original_inputs)?;
|
||||
tape.simplify()?;
|
||||
tape.rebuild(outputs)
|
||||
}
|
||||
|
||||
pub(super) fn compile(
|
||||
outputs: &[Array],
|
||||
original_inputs: &[Array],
|
||||
placeholders: &[Array],
|
||||
) -> Result<Vec<Array>, String> {
|
||||
let mut tape = Tape::capture(outputs, original_inputs)?;
|
||||
tape.simplify()?;
|
||||
let outputs = tape.fuse(placeholders, outputs)?;
|
||||
tape.rebuild(&outputs)
|
||||
}
|
||||
|
||||
fn equivalent_operation(a: &Operation, b: &Operation, x: &Array, y: &Array) -> bool {
|
||||
use Operation::*;
|
||||
let same_outputs = || {
|
||||
let x = x.outputs();
|
||||
let y = y.outputs();
|
||||
x.len() == y.len()
|
||||
&& x.iter().zip(&y).all(|(a, b)| {
|
||||
a.layout().shape() == b.layout().shape() && a.layout().dtype() == b.layout().dtype()
|
||||
})
|
||||
};
|
||||
match (a, b) {
|
||||
(Flatten { start: a, end: b }, Flatten { start: c, end: d }) => a == c && b == d,
|
||||
(Unflatten { axis: a, shape: b }, Unflatten { axis: c, shape: d }) => a == c && b == d,
|
||||
(Reshape, Reshape) | (Broadcast, Broadcast) | (AsType, AsType) | (View, View) => {
|
||||
same_outputs()
|
||||
}
|
||||
(Transpose(a), Transpose(b))
|
||||
| (ExpandDims(a), ExpandDims(b))
|
||||
| (Squeeze(a), Squeeze(b))
|
||||
| (Sum(a), Sum(b))
|
||||
| (Min(a), Min(b))
|
||||
| (Max(a), Max(b))
|
||||
| (DynamicSliceUpdate(a), DynamicSliceUpdate(b)) => a == b,
|
||||
(DynamicSlice(a), DynamicSlice(b)) => a == b && same_outputs(),
|
||||
(
|
||||
Slice {
|
||||
starts: a,
|
||||
stops: b,
|
||||
steps: c,
|
||||
},
|
||||
Slice {
|
||||
starts: x,
|
||||
stops: y,
|
||||
steps: z,
|
||||
},
|
||||
)
|
||||
| (
|
||||
SliceUpdate {
|
||||
starts: a,
|
||||
stops: b,
|
||||
steps: c,
|
||||
},
|
||||
SliceUpdate {
|
||||
starts: x,
|
||||
stops: y,
|
||||
steps: z,
|
||||
},
|
||||
) => (a, b, c) == (x, y, z),
|
||||
(Contiguous(a), Contiguous(b))
|
||||
| (Softmax(a), Softmax(b))
|
||||
| (SearchSorted(a), SearchSorted(b))
|
||||
| (CumsumLastF32(a), CumsumLastF32(b)) => a == b,
|
||||
(RmsNorm(a), RmsNorm(b)) => a == b,
|
||||
(
|
||||
SdpaVector {
|
||||
scale: a,
|
||||
causal: b,
|
||||
},
|
||||
SdpaVector {
|
||||
scale: x,
|
||||
causal: y,
|
||||
},
|
||||
) => (a, b) == (x, y),
|
||||
(QuantizedLinearBf16 { bits: a, group: b }, QuantizedLinearBf16 { bits: x, group: y })
|
||||
| (AffineQuantize { bits: a, group: b }, AffineQuantize { bits: x, group: y })
|
||||
| (AffineDequantize { bits: a, group: b }, AffineDequantize { bits: x, group: y })
|
||||
| (
|
||||
GatherQmmBf16 {
|
||||
bits: a, group: b, ..
|
||||
},
|
||||
GatherQmmBf16 {
|
||||
bits: x, group: y, ..
|
||||
},
|
||||
) => (a, b) == (x, y),
|
||||
(
|
||||
Sort {
|
||||
axis: a,
|
||||
indices: b,
|
||||
kth: c,
|
||||
},
|
||||
Sort {
|
||||
axis: x,
|
||||
indices: y,
|
||||
kth: z,
|
||||
},
|
||||
) => (a, b, c) == (x, y, z),
|
||||
(
|
||||
Split {
|
||||
indices: a,
|
||||
axis: b,
|
||||
},
|
||||
Split {
|
||||
indices: x,
|
||||
axis: y,
|
||||
},
|
||||
) => (a, b) == (x, y),
|
||||
(Unary(a), Unary(b)) => std::mem::discriminant(a) == std::mem::discriminant(b),
|
||||
(Binary(a), Binary(b)) => std::mem::discriminant(a) == std::mem::discriminant(b),
|
||||
(Concatenate(a), Concatenate(b))
|
||||
| (ArgMax(a), ArgMax(b))
|
||||
| (GatherAxis(a), GatherAxis(b))
|
||||
| (ScatterAxis(a), ScatterAxis(b))
|
||||
| (GatherTake(a), GatherTake(b)) => a == b,
|
||||
(ArangeI32(a), ArangeI32(b)) => a == b && same_outputs(),
|
||||
(ArangeF32 { start: a, step: b }, ArangeF32 { start: x, step: y }) => {
|
||||
(a, b) == (x, y) && same_outputs()
|
||||
}
|
||||
(ArangeU32, ArangeU32) | (ArangeI64, ArangeI64) => same_outputs(),
|
||||
(RandomBits, RandomBits) => same_outputs(),
|
||||
(Full, Full)
|
||||
| (DenseMatmul, DenseMatmul)
|
||||
| (Select, Select)
|
||||
| (ScatterAddRow, ScatterAddRow)
|
||||
| (AllBoolContiguous, AllBoolContiguous)
|
||||
| (LogSumExp, LogSumExp)
|
||||
| (CummaxLastI64, CummaxLastI64)
|
||||
| (CumsumLastI32, CumsumLastI32)
|
||||
| (DepthwiseConv1d, DepthwiseConv1d)
|
||||
| (PleConv1d, PleConv1d)
|
||||
| (SiluBf16, SiluBf16)
|
||||
| (SwiGluBf16, SwiGluBf16)
|
||||
| (ComputeGBf16, ComputeGBf16) => true,
|
||||
// CustomKernel does NOT override Primitive::is_equivalent in the
|
||||
// reference. Identical shader text is not permission to merge it.
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mtplx_compile_simplify_preserves_reference_parent_rules() {
|
||||
use super::{
|
||||
binary::{Binary, binary},
|
||||
compiled::Replay,
|
||||
ops,
|
||||
stream::{Device, Stream},
|
||||
unary::{Unary, unary},
|
||||
};
|
||||
let stream = Stream::new(0, Device::Gpu);
|
||||
let input = Array::unallocated(&[4], Dtype::F32).unwrap();
|
||||
let args = std::slice::from_ref(&input);
|
||||
let body = |args: &[Array]| {
|
||||
let a = unary(&args[0], Unary::Square, stream)?;
|
||||
let b = unary(&args[0], Unary::Square, stream)?;
|
||||
let c = unary(&a, Unary::Negative, stream)?;
|
||||
let d = unary(&b, Unary::Negative, stream)?;
|
||||
Ok(vec![binary(&c, &d, Binary::Add, stream)?])
|
||||
};
|
||||
let raw = Replay::trace(args, body).unwrap().run(args).unwrap();
|
||||
let merged = Replay::trace_simplified(args, body)
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
assert_ne!(raw[0].inputs()[0].id(), raw[0].inputs()[1].id());
|
||||
assert_eq!(merged[0].inputs()[0].id(), merged[0].inputs()[1].id());
|
||||
let normalizers = Replay::trace_simplified(args, |args| {
|
||||
let a = super::normalization::logsumexp_last(&args[0], true, stream)?;
|
||||
let b = super::normalization::logsumexp_last(&args[0], true, stream)?;
|
||||
Ok(vec![binary(&a, &b, Binary::Add, stream)?])
|
||||
})
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
normalizers[0].inputs()[0].id(),
|
||||
normalizers[0].inputs()[1].id()
|
||||
);
|
||||
// An otherwise equivalent later global output must never disappear.
|
||||
let roots = Replay::trace_simplified(args, |args| {
|
||||
let a = unary(&args[0], Unary::Square, stream)?;
|
||||
let b = unary(&args[0], Unary::Square, stream)?;
|
||||
Ok(vec![a, b])
|
||||
})
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
assert_ne!(roots[0].id(), roots[1].id());
|
||||
// Distinct raw stop parameters remain distinct even if output values and
|
||||
// shapes are equal. compile_simplify compares Slice's stored parameters.
|
||||
let slices = Replay::trace_simplified(args, |args| {
|
||||
let a = ops::slice(&args[0], &[0], &[3], &[2], stream)?;
|
||||
let b = ops::slice(&args[0], &[0], &[4], &[2], stream)?;
|
||||
Ok(vec![binary(&a, &b, Binary::Add, stream)?])
|
||||
})
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
assert_ne!(slices[0].inputs()[0].id(), slices[0].inputs()[1].id());
|
||||
// Equivalent Split primitives merge all siblings in canonical order.
|
||||
let split = Replay::trace_simplified(args, |args| {
|
||||
let a = ops::split(&args[0], &[2], 0, stream)?;
|
||||
let b = ops::split(&args[0], &[2], 0, stream)?;
|
||||
Ok(vec![
|
||||
binary(&a[1], &b[1], Binary::Add, stream)?,
|
||||
binary(&a[0], &b[0], Binary::Multiply, stream)?,
|
||||
])
|
||||
})
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
for root in &split {
|
||||
assert_eq!(root.inputs()[0].id(), root.inputs()[1].id());
|
||||
}
|
||||
// Same immutable primitive pointer must remain separate (reference rule).
|
||||
let shared = Replay::trace_simplified(args, |args| {
|
||||
let a = unary(&args[0], Unary::Square, stream)?;
|
||||
let b = Array::make_shared_operations(
|
||||
stream,
|
||||
args,
|
||||
&[a.layout().clone()],
|
||||
a.operation().unwrap(),
|
||||
)?
|
||||
.remove(0);
|
||||
Ok(vec![binary(&a, &b, Binary::Add, stream)?])
|
||||
})
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
assert_ne!(shared[0].inputs()[0].id(), shared[0].inputs()[1].id());
|
||||
// Exercise the original >100-parent grouping path.
|
||||
let fanout = Replay::trace_simplified(args, |args| {
|
||||
let mut nodes = (0..110)
|
||||
.map(|_| unary(&args[0], Unary::Square, stream))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let mut result = nodes.pop().unwrap();
|
||||
for node in nodes {
|
||||
result = binary(&result, &node, Binary::Add, stream)?;
|
||||
}
|
||||
Ok(vec![result])
|
||||
})
|
||||
.unwrap()
|
||||
.run(args)
|
||||
.unwrap();
|
||||
let mut root = fanout[0].clone();
|
||||
let square = root.inputs()[1].id();
|
||||
for _ in 0..109 {
|
||||
assert_eq!(root.inputs()[1].id(), square);
|
||||
let next = root.inputs()[0].clone();
|
||||
root = next;
|
||||
}
|
||||
assert_eq!(root.id(), square);
|
||||
assert!(Replay::trace_simplified(args, |_| Ok(vec![input.clone()])).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires Metal buffers and external supervisor; scalar ownership only, no GPU dispatch"]
|
||||
fn mtplx_compile_scalar_bits_and_custom_kernel_identity() {
|
||||
use super::{
|
||||
super::gpu::Context,
|
||||
configure_sources, ops,
|
||||
stream::{Device, Stream},
|
||||
};
|
||||
configure_sources().unwrap();
|
||||
let _context = Context::open_qwen(0).unwrap();
|
||||
let stream = Stream::new(0, Device::Gpu);
|
||||
let scalar = |shape: &[i32], dtype: Dtype, bytes: &[u8]| {
|
||||
Array::new(shape, dtype, super::scalar_buffer(bytes).unwrap()).unwrap()
|
||||
};
|
||||
// Available scalar identity is raw bits AND dtype, not float equality:
|
||||
// equal NaN payloads merge, signed zero and distinct dtypes do not.
|
||||
for (left, right, ld, rd, rank, same) in [
|
||||
(
|
||||
0x7fc00001u32,
|
||||
0x7fc00001u32,
|
||||
Dtype::F32,
|
||||
Dtype::F32,
|
||||
0,
|
||||
true,
|
||||
),
|
||||
(0, 0x80000000, Dtype::F32, Dtype::F32, 0, false),
|
||||
(1, 1, Dtype::U32, Dtype::I32, 0, false),
|
||||
(1, 1, Dtype::U32, Dtype::U32, 1, false),
|
||||
] {
|
||||
let shape = if rank == 0 { &[][..] } else { &[1][..] };
|
||||
let a = scalar(shape, ld, &left.to_le_bytes());
|
||||
let b = scalar(shape, rd, &right.to_le_bytes());
|
||||
let roots = [
|
||||
ops::full(&[2], &a, ld, stream).unwrap(),
|
||||
ops::full(&[2], &b, rd, stream).unwrap(),
|
||||
];
|
||||
let mut tape = Tape::capture(&roots, &[]).unwrap();
|
||||
tape.scalars().unwrap();
|
||||
let leaves = tape
|
||||
.order
|
||||
.iter()
|
||||
.filter(|id| tape.nodes[id].array.operation().is_none())
|
||||
.count();
|
||||
assert_eq!(leaves, if same { 1 } else { 2 });
|
||||
}
|
||||
// CustomKernel inherits the reference's default false equivalence.
|
||||
let input = scalar(&[1], Dtype::F32, &1f32.to_le_bytes());
|
||||
let make = || {
|
||||
Array::make_operation(stream, &input, input.layout().clone(), Operation::HyperV3R1).unwrap()
|
||||
};
|
||||
let a = make();
|
||||
let b = make();
|
||||
assert!(!equivalent_operation(
|
||||
&a.operation().unwrap(),
|
||||
&b.operation().unwrap(),
|
||||
&a,
|
||||
&b
|
||||
));
|
||||
// Canonical source descriptors survive simplification unchanged.
|
||||
let x = ops::full(&[2], &input, Dtype::F32, stream).unwrap();
|
||||
let before = x.inputs()[0].id();
|
||||
let optimized = simplify(std::slice::from_ref(&x), &[]).unwrap();
|
||||
assert_eq!(x.inputs()[0].id(), before);
|
||||
assert_ne!(optimized[0].id(), x.id());
|
||||
}
|
||||
Reference in New Issue
Block a user