use std::collections::HashMap;
use std::sync::RwLock;
use fuel_kernel_seam_types::{OpAttrs, OpTag, PatternNode, matmul_roles};
use crate::registry::{FusedOpId, FusedOpParams};
use crate::{Graph, Node, NodeId, Op};
#[derive(Clone, Debug)]
pub struct RuntimeFusedOpEntry {
pub id: FusedOpId,
pub name: String,
pub region: PatternNode,
}
#[derive(Clone, Debug, PartialEq)]
pub enum RuntimeFusedError {
NonContiguousBinds(Vec<u8>),
UnRepresentable(OpTag),
NonConcreteRegion,
IdSpaceExhausted,
InvalidRelAttrs { tag: OpTag, error: RelAttrError },
}
static RUNTIME_FUSED_OPS: RwLock<Vec<RuntimeFusedOpEntry>> = RwLock::new(Vec::new());
fn hash_index() -> &'static RwLock<HashMap<u64, FusedOpId>> {
static IDX: std::sync::OnceLock<RwLock<HashMap<u64, FusedOpId>>> = std::sync::OnceLock::new();
IDX.get_or_init(|| RwLock::new(HashMap::new()))
}
fn push_placeholder_leaves(graph: &mut Graph, arity: usize) -> Vec<NodeId> {
(0..arity)
.map(|_| {
graph.push(Node {
op: Op::Const,
inputs: vec![],
shape: fuel_ir::Shape::from_dims(&[1]),
dtype: fuel_ir::DType::F32,
})
})
.collect()
}
fn region_base_map_hash(region: &PatternNode) -> Option<u64> {
let n_binds = region.bind_indices().len();
let scalars = vec![0.0; count_scalar_slots(region)];
let graph: crate::SharedGraph = std::sync::Arc::new(RwLock::new(Graph::new()));
let sink = {
let mut g = graph.write().ok()?;
let inputs = push_placeholder_leaves(&mut g, n_binds);
try_emit_region(&mut g, region, &inputs, &scalars).ok()?
};
let roots = crate::opt::lower_to_base_map(&graph, &[sink]);
let root = *roots.first()?;
let g = graph.read().ok()?;
Some(crate::opt::base_map_hash(&g, root))
}
pub fn register_runtime_fused(
name: impl Into<String>,
region: PatternNode,
) -> Result<FusedOpId, RuntimeFusedError> {
let name = name.into();
validate_recipe(®ion)?;
let hash = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
region_base_map_hash(®ion)
}))
.unwrap_or(None);
if hash.is_none() {
eprintln!(
"register_runtime_fused: base-map hash unavailable for {name:?}; \
registering without dedup (allocate-fresh fallback)"
);
}
let mut idx = hash_index().write().unwrap();
if let Some(h) = hash
&& let Some(&existing) = idx.get(&h)
{
return Ok(existing);
}
let mut w = RUNTIME_FUSED_OPS.write().unwrap();
let raw = FusedOpId::RUNTIME_FUSED_BASE as usize + w.len();
if raw > u16::MAX as usize {
return Err(RuntimeFusedError::IdSpaceExhausted);
}
let id = FusedOpId(raw as u16);
w.push(RuntimeFusedOpEntry { id, name, region });
drop(w);
if let Some(h) = hash {
idx.insert(h, id);
}
Ok(id)
}
pub fn runtime_region(id: FusedOpId) -> Option<PatternNode> {
if !id.is_runtime() {
return None;
}
let idx = (id.0 - FusedOpId::RUNTIME_FUSED_BASE) as usize;
RUNTIME_FUSED_OPS
.read()
.unwrap()
.get(idx)
.map(|e| e.region.clone())
}
pub fn runtime_name(id: FusedOpId) -> Option<String> {
if !id.is_runtime() {
return None;
}
let idx = (id.0 - FusedOpId::RUNTIME_FUSED_BASE) as usize;
RUNTIME_FUSED_OPS
.read()
.unwrap()
.get(idx)
.map(|e| e.name.clone())
}
pub fn runtime_entries() -> Vec<RuntimeFusedOpEntry> {
RUNTIME_FUSED_OPS.read().unwrap().clone()
}
#[doc(hidden)]
pub fn clear_runtime_fused_for_tests() {
RUNTIME_FUSED_OPS.write().unwrap().clear();
hash_index().write().unwrap().clear();
}
fn tag_to_op(tag: OpTag, attrs: &OpAttrs) -> Option<Op> {
use OpTag as T;
use fuel_ir::DType;
use std::str::FromStr;
Some(match tag {
T::Add => Op::Add,
T::Sub => Op::Sub,
T::Mul => Op::Mul,
T::Div => Op::Div,
T::Maximum => Op::Maximum,
T::Minimum => Op::Minimum,
T::Pow => Op::Pow,
T::Rem => Op::Rem,
T::Neg => Op::Neg,
T::Abs => Op::Abs,
T::Sqr => Op::Sqr,
T::Sqrt => Op::Sqrt,
T::Rsqrt => Op::Rsqrt,
T::Recip => Op::Recip,
T::Exp => Op::Exp,
T::Log => Op::Log,
T::Sin => Op::Sin,
T::Cos => Op::Cos,
T::Tanh => Op::Tanh,
T::Sigmoid => Op::Sigmoid,
T::Silu => Op::Silu,
T::GeluTanh => Op::GeluTanh,
T::Gelu => Op::Gelu,
T::Relu => Op::Relu,
T::Erf => Op::Erf,
T::Step => Op::Step,
T::Floor => Op::Floor,
T::Ceil => Op::Ceil,
T::Round => Op::Round,
T::Sign => Op::Sign,
T::AddScalar => Op::AddScalar(*attrs.scalars.first()?),
T::MulScalar => Op::MulScalar(*attrs.scalars.first()?),
T::PowI => Op::PowI(*attrs.scalars.first()? as i32),
T::Equal => Op::Equal,
T::Ne => Op::Ne,
T::Lt => Op::Lt,
T::Le => Op::Le,
T::Gt => Op::Gt,
T::Ge => Op::Ge,
T::Where => Op::Where,
T::Cast => Op::Cast(DType::from_str(attrs.cast_dtype.as_deref()?).ok()?),
T::MatMul => {
if attrs.lhs_roles.is_empty() && attrs.rhs_roles.is_empty() {
Op::MatMul
} else {
let (canon_lhs, canon_rhs) =
matmul_roles(attrs.lhs_roles.len(), attrs.rhs_roles.len());
if attrs.lhs_roles.len() == attrs.rhs_roles.len()
&& attrs.lhs_roles.len() >= 2
&& attrs.lhs_roles == canon_lhs
&& attrs.rhs_roles == canon_rhs
{
Op::MatMul
} else {
return None;
}
}
}
T::LogSoftmaxLastDim => Op::LogSoftmaxLastDim,
T::Transpose => Op::Transpose,
T::Permute => Op::Permute(attrs.perm.iter().map(|&x| x as usize).collect()),
T::Reshape => Op::Reshape(shape_from_attr(attrs)?),
T::BroadcastTo => Op::BroadcastTo(shape_from_attr(attrs)?),
T::ReduceSumTo => Op::ReduceSumTo(shape_from_attr(attrs)?),
T::ReduceMaxTo => Op::ReduceMaxTo(shape_from_attr(attrs)?),
T::Unsqueeze => Op::Unsqueeze {
dim: *attrs.dims.first()? as usize,
},
T::Squeeze => Op::Squeeze {
dim: *attrs.dims.first()? as usize,
},
T::Slice => Op::Slice {
dim: attrs.axis? as usize,
start: attrs.slice_start? as usize,
len: attrs.slice_len? as usize,
},
T::Concat => Op::Concat {
dim: attrs.axis? as usize,
},
T::Flip => Op::Flip {
dim: attrs.axis? as usize,
},
T::Roll => Op::Roll {
dim: attrs.axis? as usize,
shift: attrs.roll_shift?,
},
T::Pad => Op::Pad {
padding: attrs
.pad_amounts
.iter()
.map(|&(b, e)| (b as usize, e as usize))
.collect(),
mode: match attrs.pad_mode? {
0 => crate::PadMode::Constant,
1 => crate::PadMode::Reflect,
2 => crate::PadMode::Replicate,
_ => return None,
},
value: attrs.pad_value.unwrap_or(0.0),
},
T::Triu => Op::Triu {
diagonal: attrs.axis?,
},
T::Tril => Op::Tril {
diagonal: attrs.axis?,
},
T::SumDim => Op::SumDim(attrs.axis? as usize),
T::MaxDim => Op::MaxDim(attrs.axis? as usize),
T::MeanDim => Op::MeanDim(attrs.axis? as usize),
T::SumAll => Op::SumAll,
T::MaxAll => Op::MaxAll,
T::MinAll => Op::MinAll,
T::MeanAll => Op::MeanAll,
T::CumSum => Op::CumSum {
dim: attrs.axis? as usize,
},
T::Iota => Op::Iota {
len: *attrs.target_shape.first()? as usize,
},
T::IndexSelect => Op::IndexSelect {
dim: attrs.axis? as usize,
},
T::Gather => Op::Gather {
dim: attrs.axis? as usize,
},
T::IndexAdd => Op::IndexAdd {
dim: attrs.axis? as usize,
},
T::ScatterAdd => Op::ScatterAdd {
dim: attrs.axis? as usize,
},
T::Scan => Op::Scan {
n_carries: attrs.scan_n_carries? as usize,
n_xs: attrs.scan_n_xs? as usize,
bound: attrs.scan_bound? as usize,
emit: match attrs.scan_emit? {
0 => crate::ScanEmit::All,
1 => crate::ScanEmit::Final,
_ => return None,
},
early_exit: if attrs.scan_early_exit.unwrap_or(false) {
Some(crate::ScanPredicate)
} else {
None
},
},
T::ScanPlaceholder => Op::ScanPlaceholder {
role: match attrs.scan_role? {
fuel_kernel_seam_types::SCAN_ROLE_CARRY => crate::ScanRole::Carry,
fuel_kernel_seam_types::SCAN_ROLE_ELEM => crate::ScanRole::Elem,
_ => return None,
},
index: attrs.scan_index? as usize,
},
T::View => Op::View {
slot: attrs.view_slot?,
},
T::Fused => {
let fid =
crate::registry::default_registry().id_for_name(attrs.fused_op.as_deref()?)?;
Op::Fused(fid, fused_params_for(fid)?)
}
T::MaskedFill => {
let v = *attrs.scalars.first()?;
let dtype = match attrs.cast_dtype.as_deref() {
Some(name) => DType::from_str(name).ok()?,
None => DType::F32,
};
Op::MaskedFill {
value: masked_fill_scalar(v, dtype)?,
}
}
_ => return None,
})
}
fn fused_params_for(fid: FusedOpId) -> Option<FusedOpParams> {
use crate::registry::FusedOps;
Some(match fid {
FusedOps::SOFTMAX_LAST_DIM => FusedOpParams::SoftmaxLastDim,
FusedOps::SOFTMAX_LAST_DIM_BACKWARD => FusedOpParams::SoftmaxLastDimBackward,
_ => return None,
})
}
fn masked_fill_scalar(value: f64, dtype: fuel_ir::DType) -> Option<fuel_ir::Scalar> {
fuel_ir::Scalar::from_f64(value, dtype).ok()
}
fn shape_from_attr(attrs: &OpAttrs) -> Option<fuel_ir::Shape> {
if attrs.target_shape.is_empty() {
if attrs.rank0_target {
return Some(fuel_ir::Shape::from_dims(&[]));
}
return None;
}
let dims: Vec<usize> = attrs.target_shape.iter().map(|&d| d as usize).collect();
Some(fuel_ir::Shape::from_dims(&dims))
}
fn scalar_slot_arity(tag: OpTag) -> usize {
matches!(tag, OpTag::AddScalar | OpTag::MulScalar) as usize
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RelAttrError {
RelAbsConflict { field: &'static str },
Expr(fuel_kernel_seam_types::shape_expr::ShapeExprError),
SymbolicGap { field: &'static str },
Negative { field: &'static str, value: i64 },
AxisLastUnsupported { tag: OpTag },
NoChildOperand,
FrameNotExpressible {
field: &'static str,
frame: Vec<i64>,
missing_ctor: u8,
},
}
fn has_rel_attrs(attrs: &OpAttrs) -> bool {
attrs.target_shape_rel.is_some()
|| attrs.slice_start_rel.is_some()
|| attrs.slice_len_rel.is_some()
|| attrs.axis_last
|| attrs.scalar_rel.is_some()
}
fn rel_abs_conflict_field(tag: OpTag, attrs: &OpAttrs) -> Option<&'static str> {
if attrs.target_shape_rel.is_some() && !attrs.target_shape.is_empty() {
return Some("target_shape");
}
if attrs.slice_start_rel.is_some() && attrs.slice_start.is_some() {
return Some("slice_start");
}
if attrs.slice_len_rel.is_some() && attrs.slice_len.is_some() {
return Some("slice_len");
}
if attrs.scalar_rel.is_some() && !attrs.scalars.is_empty() {
return Some("scalars");
}
if attrs.axis_last {
match tag {
OpTag::Unsqueeze | OpTag::Squeeze => {
if !attrs.dims.is_empty() {
return Some("dims");
}
}
_ => {
if attrs.axis.is_some() {
return Some("axis");
}
}
}
}
None
}
fn bind_broadcast_frame(bind_shapes: &[Vec<i64>]) -> Option<Vec<i64>> {
use fuel_kernel_seam_types::shape_expr::SYMBOLIC;
let rank = bind_shapes.iter().map(Vec::len).max()?;
let mut frame = vec![1i64; rank];
for s in bind_shapes {
if s.iter().any(|&e| e == SYMBOLIC || e < 0) {
return None;
}
let pad = rank - s.len(); for (i, &e) in s.iter().enumerate() {
let f = &mut frame[pad + i];
if *f == e || e == 1 {
continue;
}
if *f == 1 {
*f = e;
continue;
}
return None; }
}
Some(frame)
}
fn same_as_frame_guard(
tag: OpTag,
se: &fuel_kernel_seam_types::shape_expr::ShapeExpr,
bind_shapes: &[Vec<i64>],
) -> Result<(), RelAttrError> {
use fuel_kernel_seam_types::shape_expr::{ShapeExpr, TAG_DIMS};
match se {
ShapeExpr::SameAs { .. } => {}
ShapeExpr::WithDim { .. } | ShapeExpr::Dims(_) => return Ok(()),
}
if tag != OpTag::BroadcastTo {
return Ok(());
}
let Some(frame) = bind_broadcast_frame(bind_shapes) else {
return Ok(()); };
if bind_shapes.contains(&frame) {
return Ok(()); }
Err(RelAttrError::FrameNotExpressible {
field: "target_shape",
frame,
missing_ctor: TAG_DIMS,
})
}
pub fn resolve_rel_attrs(
tag: OpTag,
attrs: &OpAttrs,
bind_shapes: &[Vec<i64>],
child_shapes: &[Vec<i64>],
) -> Result<OpAttrs, RelAttrError> {
use fuel_kernel_seam_types::shape_expr::{self, Dim, DimValue, LAST, ShapeValue, resolve_axis};
if let Some(field) = rel_abs_conflict_field(tag, attrs) {
return Err(RelAttrError::RelAbsConflict { field });
}
let mut out = attrs.clone();
if let Some(se) = &attrs.target_shape_rel {
match shape_expr::eval_shape(se, bind_shapes, &[]).map_err(RelAttrError::Expr)? {
ShapeValue::Concrete(s) => {
if let Some(&bad) = s.iter().find(|&&e| e < 0) {
return Err(RelAttrError::Negative {
field: "target_shape",
value: bad,
});
}
same_as_frame_guard(tag, se, bind_shapes)?;
out.target_shape = s;
}
ShapeValue::Gap => {
return Err(RelAttrError::SymbolicGap {
field: "target_shape",
});
}
}
out.target_shape_rel = None;
}
let eval_dim_field = |d: &Dim, field: &'static str| -> Result<u64, RelAttrError> {
match shape_expr::eval_dim(d, bind_shapes, &[]).map_err(RelAttrError::Expr)? {
DimValue::Concrete(v) if v < 0 => Err(RelAttrError::Negative { field, value: v }),
DimValue::Concrete(v) => Ok(v as u64),
DimValue::Gap => Err(RelAttrError::SymbolicGap { field }),
}
};
if let Some(d) = &attrs.slice_start_rel {
out.slice_start = Some(eval_dim_field(d, "slice_start")?);
out.slice_start_rel = None;
}
if let Some(d) = &attrs.slice_len_rel {
out.slice_len = Some(eval_dim_field(d, "slice_len")?);
out.slice_len_rel = None;
}
if let Some(d) = &attrs.scalar_rel {
match shape_expr::eval_dim(d, bind_shapes, &[]).map_err(RelAttrError::Expr)? {
DimValue::Concrete(v) => out.scalars = vec![v as f64],
DimValue::Gap => return Err(RelAttrError::SymbolicGap { field: "scalars" }),
}
out.scalar_rel = None;
}
if attrs.axis_last {
let rank = child_shapes
.first()
.ok_or(RelAttrError::NoChildOperand)?
.len();
use OpTag as T;
match tag {
T::SumDim
| T::MaxDim
| T::MeanDim
| T::CumSum
| T::Concat
| T::Flip
| T::Slice
| T::Roll
| T::IndexSelect
| T::Gather
| T::IndexAdd
| T::ScatterAdd => {
let a = resolve_axis(LAST, rank).map_err(RelAttrError::Expr)?;
out.axis = Some(a as i64);
}
T::Unsqueeze => {
out.dims = vec![rank as u8];
}
T::Squeeze => {
let a = resolve_axis(LAST, rank).map_err(RelAttrError::Expr)?;
out.dims = vec![a as u8];
}
other => return Err(RelAttrError::AxisLastUnsupported { tag: other }),
}
out.axis_last = false;
}
Ok(out)
}
pub fn count_scalar_slots(node: &PatternNode) -> usize {
match node {
PatternNode::Op {
op,
operands,
attrs,
} => {
let own = if attrs.scalars.is_empty() && attrs.scalar_rel.is_none() {
scalar_slot_arity(*op)
} else {
0
};
own + operands.iter().map(count_scalar_slots).sum::<usize>()
}
_ => 0,
}
}
pub(crate) fn validate_recipe(region: &PatternNode) -> Result<(), RuntimeFusedError> {
let binds = region.bind_indices();
let n = binds.len() as u8;
if binds != (0..n).collect::<Vec<_>>() {
return Err(RuntimeFusedError::NonContiguousBinds(binds));
}
validate_representable(region)
}
fn validate_representable(region: &PatternNode) -> Result<(), RuntimeFusedError> {
let n_binds = region.bind_indices().len();
validate_node(region, n_binds)
}
fn validate_node(node: &PatternNode, n_binds: usize) -> Result<(), RuntimeFusedError> {
match node {
PatternNode::Op {
op,
operands,
attrs,
} => {
let probed;
let attrs = if has_rel_attrs(attrs) {
probed = rel_probe(*op, attrs, n_binds)
.map_err(|error| RuntimeFusedError::InvalidRelAttrs { tag: *op, error })?;
&probed
} else {
attrs
};
let representable = if attrs.scalars.is_empty() && scalar_slot_arity(*op) > 0 {
let mut filled = attrs.clone();
filled.scalars = vec![0.0; scalar_slot_arity(*op)];
tag_to_op(*op, &filled).is_some()
} else {
tag_to_op(*op, attrs).is_some()
};
if !representable {
return Err(RuntimeFusedError::UnRepresentable(*op));
}
for o in operands {
validate_node(o, n_binds)?;
}
Ok(())
}
PatternNode::Bind { .. } => Ok(()),
PatternNode::Any | PatternNode::SeeThrough { .. } => {
Err(RuntimeFusedError::NonConcreteRegion)
}
}
}
fn rel_probe(tag: OpTag, attrs: &OpAttrs, n_binds: usize) -> Result<OpAttrs, RelAttrError> {
use fuel_kernel_seam_types::shape_expr::ShapeExprError as E;
let probe: Vec<i64> = vec![2, 4];
let bind_shapes = vec![probe.clone(); n_binds];
match resolve_rel_attrs(tag, attrs, &bind_shapes, std::slice::from_ref(&probe)) {
Ok(resolved) => Ok(resolved),
Err(
e @ (RelAttrError::RelAbsConflict { .. }
| RelAttrError::AxisLastUnsupported { .. }
| RelAttrError::Expr(E::OperandOutOfRange { .. } | E::ParamOutOfRange { .. })),
) => Err(e),
Err(_) => Ok(dummy_fill_rel(tag, attrs)),
}
}
fn dummy_fill_rel(tag: OpTag, attrs: &OpAttrs) -> OpAttrs {
let mut out = attrs.clone();
if out.target_shape_rel.take().is_some() && out.target_shape.is_empty() {
out.target_shape = vec![1];
}
if out.slice_start_rel.take().is_some() && out.slice_start.is_none() {
out.slice_start = Some(0);
}
if out.slice_len_rel.take().is_some() && out.slice_len.is_none() {
out.slice_len = Some(1);
}
if out.scalar_rel.take().is_some() && out.scalars.is_empty() {
out.scalars = vec![1.0];
}
if out.axis_last {
out.axis_last = false;
match tag {
OpTag::Unsqueeze | OpTag::Squeeze => {
if out.dims.is_empty() {
out.dims = vec![0];
}
}
_ => {
if out.axis.is_none() {
out.axis = Some(0);
}
}
}
}
out
}
pub fn decompose_region(graph: &mut Graph, node_id: NodeId) -> NodeId {
let (fid, node_scalars) = match &graph.node(node_id).op {
Op::Fused(id, FusedOpParams::Runtime { scalars }) => (*id, scalars.clone()),
Op::Fused(id, _) => (*id, Vec::new()),
_ => return node_id,
};
let region = match runtime_region(fid) {
Some(r) => r,
None => return node_id,
};
if node_scalars.len() != count_scalar_slots(®ion) {
return node_id;
}
let inputs = graph.node(node_id).inputs.clone();
let bind_shapes = bind_operand_shapes(graph, &inputs);
let mut cursor = node_scalars.as_slice();
emit(
graph,
®ion,
&inputs,
&bind_shapes,
&mut cursor,
&mut Vec::new(),
)
.unwrap_or(node_id)
}
pub fn emit_region(
graph: &mut Graph,
region: &PatternNode,
inputs: &[NodeId],
scalars: &[f64],
) -> NodeId {
try_emit_region(graph, region, inputs, scalars).expect(
"rel-attr resolution failed — emit_region callers pass concrete-attr or \
shape-compatible pre-validated regions; fallible callers use try_emit_region",
)
}
pub fn try_emit_region(
graph: &mut Graph,
region: &PatternNode,
inputs: &[NodeId],
scalars: &[f64],
) -> Result<NodeId, RelAttrError> {
let bind_shapes = bind_operand_shapes(graph, inputs);
let mut cursor = scalars;
emit(
graph,
region,
inputs,
&bind_shapes,
&mut cursor,
&mut Vec::new(),
)
}
fn shape_expr_operand(shape: &fuel_ir::Shape) -> Vec<i64> {
use fuel_kernel_seam_types::shape_expr::SYMBOLIC;
(0..shape.rank())
.map(|a| {
if shape.extent(a).is_dynamic() {
SYMBOLIC
} else {
shape.dims()[a] as i64
}
})
.collect()
}
fn bind_operand_shapes(graph: &Graph, inputs: &[NodeId]) -> Vec<Vec<i64>> {
inputs
.iter()
.map(|&id| shape_expr_operand(&graph.node(id).shape))
.collect()
}
fn emit<'r>(
graph: &mut Graph,
node: &'r PatternNode,
inputs: &[NodeId],
bind_shapes: &[Vec<i64>],
scalars: &mut &[f64],
memo: &mut Vec<(&'r PatternNode, NodeId)>,
) -> Result<NodeId, RelAttrError> {
match node {
PatternNode::Bind { index } => Ok(inputs[*index as usize]),
PatternNode::Op {
op,
operands,
attrs,
} => {
let sharable = count_scalar_slots(node) == 0;
if sharable && let Some(&(_, id)) = memo.iter().find(|(p, _)| *p == node) {
return Ok(id);
}
let arity = scalar_slot_arity(*op);
let filled;
let attrs = if attrs.scalars.is_empty() && arity > 0 && attrs.scalar_rel.is_none() {
let (take, rest) = scalars.split_at(arity);
*scalars = rest;
filled = OpAttrs {
scalars: take.to_vec(),
..attrs.clone()
};
&filled
} else {
attrs
};
let mut child_ids = Vec::with_capacity(operands.len());
for o in operands {
child_ids.push(emit(graph, o, inputs, bind_shapes, scalars, memo)?);
}
let mut child_shapes: Vec<fuel_ir::Shape> = child_ids
.iter()
.map(|&c| graph.node(c).shape.clone())
.collect();
let child_dtypes: Vec<fuel_ir::DType> =
child_ids.iter().map(|&c| graph.node(c).dtype).collect();
let resolved;
let attrs = if has_rel_attrs(attrs) {
let child_ops: Vec<Vec<i64>> =
child_shapes.iter().map(shape_expr_operand).collect();
resolved = resolve_rel_attrs(*op, attrs, bind_shapes, &child_ops)?;
&resolved
} else {
attrs
};
let mut prim =
tag_to_op(*op, attrs).expect("region validated re-emittable at registration");
if let Op::MaskedFill { value } = &prim
&& let Some(&dt) = child_dtypes.first()
&& value.dtype() != dt
&& let Some(fixed) = masked_fill_scalar(value.to_f64(), dt)
{
prim = Op::MaskedFill { value: fixed };
}
if let Op::BroadcastTo(target) = &prim
&& let Some(cs) = child_shapes.first()
&& target.rank() > cs.rank()
{
let mut padded: Vec<usize> = vec![1; target.rank() - cs.rank()];
padded.extend_from_slice(cs.dims());
let pad_shape = fuel_ir::Shape::from_dims(&padded);
let pad = graph.push(Node {
op: Op::Reshape(pad_shape.clone()),
inputs: vec![child_ids[0]],
shape: pad_shape.clone(),
dtype: child_dtypes[0],
});
child_ids[0] = pad;
child_shapes[0] = pad_shape;
}
fn fallback_sd(
cs: &[fuel_ir::Shape],
cd: &[fuel_ir::DType],
) -> (fuel_ir::Shape, fuel_ir::DType) {
(
cs.first()
.cloned()
.unwrap_or_else(|| fuel_ir::Shape::from_dims(&[])),
cd.first().copied().unwrap_or(fuel_ir::DType::F32),
)
}
let mut scan_bundle: Option<Vec<fuel_ir::storage::OutputViewSpec>> = None;
let (s, d) = match &prim {
Op::Scan {
n_carries: _,
n_xs,
bound,
early_exit,
..
} => {
let n_trailing = if early_exit.is_some() { 3 } else { 2 };
let len = child_ids.len();
if len >= 1 + *n_xs + n_trailing {
let carry_shape = child_shapes[0].clone();
let carry_dtype = child_dtypes[0];
let body_y_i = len - n_trailing + 1;
let y_shape = &child_shapes[body_y_i];
let y_dtype = child_dtypes[body_y_i];
let mut ys_dims = Vec::with_capacity(1 + y_shape.rank());
ys_dims.push(*bound);
ys_dims.extend_from_slice(y_shape.dims());
let ys_shape = fuel_ir::Shape::from_dims(&ys_dims);
scan_bundle = Some(vec![
fuel_ir::storage::OutputViewSpec::contiguous(y_dtype, ys_shape.clone()),
fuel_ir::storage::OutputViewSpec::contiguous(carry_dtype, carry_shape),
]);
(ys_shape, y_dtype)
} else {
fallback_sd(&child_shapes, &child_dtypes)
}
}
Op::View { slot } => child_ids
.first()
.and_then(|&producer| {
graph
.output_views(producer)
.and_then(|v| v.get(*slot as usize))
.map(|spec| (spec.shape.clone(), spec.dtype))
})
.unwrap_or_else(|| fallback_sd(&child_shapes, &child_dtypes)),
Op::Fused(fid, params) => match crate::registry::default_registry().entry(*fid) {
Some(e) => (
(e.shape_rule)(&child_shapes, params),
(e.dtype_rule)(&child_dtypes, params),
),
None => fallback_sd(&child_shapes, &child_dtypes),
},
Op::ScanPlaceholder { .. } => (
shape_from_attr(attrs)
.unwrap_or_else(|| fallback_sd(&child_shapes, &child_dtypes).0),
inputs
.first()
.map(|&b| graph.node(b).dtype)
.unwrap_or_else(|| fallback_sd(&child_shapes, &child_dtypes).1),
),
_ => crate::shape::primitive_shape(&prim, &child_shapes, &child_dtypes)
.unwrap_or_else(|_| fallback_sd(&child_shapes, &child_dtypes)),
};
let out = graph.push(Node {
op: prim,
inputs: child_ids,
shape: s,
dtype: d,
});
if let Some(specs) = scan_bundle
&& let Ok((_bytes, views)) = fuel_ir::storage::compose_bundle(&specs)
{
let _ = graph.set_output_views(out, std::sync::Arc::from(views.into_boxed_slice()));
}
if sharable {
memo.push((node, out));
}
Ok(out)
}
PatternNode::Any | PatternNode::SeeThrough { .. } => {
unreachable!("region validated concrete (Op/Bind) at registration")
}
}
}
pub fn runtime_lowering_decompose(
graph: &mut Graph,
node_id: NodeId,
_params: &FusedOpParams,
) -> NodeId {
decompose_region(graph, node_id)
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::{DType, Shape};
fn relu_add_region() -> PatternNode {
PatternNode::Op {
op: OpTag::Relu,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::Add,
attrs: OpAttrs::default(),
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 1 },
],
}],
}
}
fn relu_mul_region() -> PatternNode {
PatternNode::Op {
op: OpTag::Relu,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::Mul,
attrs: OpAttrs::default(),
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 1 },
],
}],
}
}
#[test]
fn register_allocates_a_runtime_id_and_keeps_the_region() {
let id = register_runtime_fused("test::relu_mul", relu_mul_region()).unwrap();
assert!(id.is_runtime(), "allocated id is in the runtime range");
assert_eq!(runtime_region(id), Some(relu_mul_region()));
assert_eq!(runtime_name(id).as_deref(), Some("test::relu_mul"));
}
#[test]
fn register_runtime_fused_dedups_structurally_identical_regions() {
let id1 = register_runtime_fused("dedup::a", relu_add_region()).unwrap();
let id2 = register_runtime_fused("dedup::b", relu_add_region()).unwrap(); assert_eq!(
id1, id2,
"an identical region must resolve to the same FusedOpId, not a duplicate"
);
}
#[test]
fn register_rejects_non_contiguous_binds() {
let region = PatternNode::Op {
op: OpTag::Add,
attrs: OpAttrs::default(),
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 2 },
],
};
assert_eq!(
register_runtime_fused("bad", region),
Err(RuntimeFusedError::NonContiguousBinds(vec![0, 2]))
);
}
#[test]
fn register_rejects_unrepresentable_region() {
let region = PatternNode::Op {
op: OpTag::Clamp,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Bind { index: 0 }],
};
assert_eq!(
register_runtime_fused("bad", region),
Err(RuntimeFusedError::UnRepresentable(OpTag::Clamp))
);
}
#[test]
fn tag_to_op_reconstructs_shape_changing_ops() {
use fuel_ir::Shape;
let attrs = OpAttrs {
axis: Some(1),
slice_start: Some(2),
slice_len: Some(3),
..OpAttrs::default()
};
assert!(matches!(
super::tag_to_op(OpTag::Slice, &attrs),
Some(Op::Slice {
dim: 1,
start: 2,
len: 3
})
));
let attrs = OpAttrs {
axis: Some(0),
..OpAttrs::default()
};
assert!(matches!(
super::tag_to_op(OpTag::Concat, &attrs),
Some(Op::Concat { dim: 0 })
));
let attrs = OpAttrs {
target_shape: vec![6],
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::Reshape, &attrs),
Some(Op::Reshape(Shape::from_dims(&[6])))
);
let attrs = OpAttrs {
target_shape: vec![2, 3],
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::BroadcastTo, &attrs),
Some(Op::BroadcastTo(Shape::from_dims(&[2, 3])))
);
let attrs = OpAttrs {
target_shape: vec![2, 1],
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::ReduceMaxTo, &attrs),
Some(Op::ReduceMaxTo(Shape::from_dims(&[2, 1])))
);
}
#[test]
fn shape_target_ops_represent_rank0_via_the_marker() {
use fuel_ir::Shape;
let empty = Shape::from_dims(&[]);
let rank0 = OpAttrs {
rank0_target: true,
..OpAttrs::default()
};
assert_eq!(super::shape_from_attr(&rank0), Some(empty.clone()));
assert_eq!(
super::tag_to_op(OpTag::ReduceSumTo, &rank0),
Some(Op::ReduceSumTo(empty.clone()))
);
assert_eq!(
super::tag_to_op(OpTag::ReduceMaxTo, &rank0),
Some(Op::ReduceMaxTo(empty.clone()))
);
assert_eq!(
super::tag_to_op(OpTag::Reshape, &rank0),
Some(Op::Reshape(empty.clone()))
);
assert_eq!(
super::tag_to_op(OpTag::BroadcastTo, &rank0),
Some(Op::BroadcastTo(empty))
);
assert_eq!(super::shape_from_attr(&OpAttrs::default()), None);
assert_eq!(
super::tag_to_op(OpTag::ReduceSumTo, &OpAttrs::default()),
None
);
}
#[test]
fn tag_to_op_reconstructs_reductions_dtype_and_matmul() {
use fuel_ir::DType;
assert!(matches!(
super::tag_to_op(
OpTag::MeanDim,
&OpAttrs {
axis: Some(1),
..OpAttrs::default()
}
),
Some(Op::MeanDim(1))
));
assert!(matches!(
super::tag_to_op(OpTag::MatMul, &OpAttrs::default()),
Some(Op::MatMul)
));
let attrs = OpAttrs {
cast_dtype: Some("f16".into()),
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::Cast, &attrs),
Some(Op::Cast(DType::F16))
);
assert!(matches!(
super::tag_to_op(OpTag::Lt, &OpAttrs::default()),
Some(Op::Lt)
));
}
#[test]
fn tag_to_op_reconstructs_masked_fill() {
use fuel_ir::Scalar;
let attrs = OpAttrs {
scalars: vec![-1.0],
cast_dtype: Some("f16".into()),
..OpAttrs::default()
};
match super::tag_to_op(OpTag::MaskedFill, &attrs) {
Some(Op::MaskedFill { value }) => {
assert_eq!(value.dtype(), DType::F16, "dtype rides cast_dtype");
assert_eq!(
value,
Scalar::from_f64(-1.0, DType::F16).unwrap(),
"value rides scalars[0]"
);
}
other => panic!("expected MaskedFill, got {other:?}"),
}
assert_eq!(
super::tag_to_op(OpTag::MaskedFill, &OpAttrs::default()),
None
);
let dummy = OpAttrs {
scalars: vec![1.0],
cast_dtype: Some("f4".into()),
..OpAttrs::default()
};
assert_eq!(super::tag_to_op(OpTag::MaskedFill, &dummy), None);
let poly = OpAttrs {
scalars: vec![1.0],
..OpAttrs::default()
};
assert!(matches!(
super::tag_to_op(OpTag::MaskedFill, &poly),
Some(Op::MaskedFill { .. })
));
assert_eq!(super::scalar_slot_arity(OpTag::MaskedFill), 0);
}
#[test]
fn tag_to_op_reconstructs_powi() {
assert!(matches!(
super::tag_to_op(
OpTag::PowI,
&OpAttrs {
scalars: vec![3.0],
..OpAttrs::default()
}
),
Some(Op::PowI(3))
));
assert!(matches!(
super::tag_to_op(
OpTag::PowI,
&OpAttrs {
scalars: vec![-1.0],
..OpAttrs::default()
}
),
Some(Op::PowI(-1))
));
assert_eq!(super::tag_to_op(OpTag::PowI, &OpAttrs::default()), None);
assert_eq!(super::scalar_slot_arity(OpTag::PowI), 0);
}
#[test]
fn masked_fill_region_registers_and_emits_with_operand_dtype() {
let region = PatternNode::Op {
op: OpTag::MaskedFill,
attrs: OpAttrs {
scalars: vec![1.0],
..OpAttrs::default()
},
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 1 },
],
};
register_runtime_fused("mf::poly", region.clone())
.expect("MaskedFill region registers (was UnRepresentable before A2)");
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F16,
});
let mask = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::U8,
});
let root = emit_region(&mut g, ®ion, &[x, mask], &[]);
match &g.node(root).op {
Op::MaskedFill { value } => {
assert_eq!(
value.dtype(),
DType::F16,
"fill Scalar re-resolved to operand[0] dtype"
);
assert_eq!(
value.to_f64(),
1.0,
"fill VALUE preserved through the re-resolution"
);
}
other => panic!("expected MaskedFill, got {other:?}"),
}
assert_eq!(g.node(root).dtype, DType::F16);
}
#[test]
fn powi_region_registers_and_emits() {
let region = PatternNode::Op {
op: OpTag::PowI,
attrs: OpAttrs {
scalars: vec![3.0],
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
};
register_runtime_fused("powi::cube", region.clone())
.expect("PowI region registers (was UnRepresentable before A3)");
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let root = emit_region(&mut g, ®ion, &[x], &[]);
assert!(
matches!(g.node(root).op, Op::PowI(3)),
"emit reconstructs Op::PowI(3) from scalars[0]",
);
assert_eq!(
g.node(root).shape,
Shape::from_dims(&[2, 3]),
"PowI is shape-preserving"
);
assert_eq!(g.node(root).dtype, DType::F32);
}
#[test]
fn tag_to_op_matmul_resolves_canonical_roles() {
let attrs = OpAttrs {
lhs_roles: vec![1, 3],
rhs_roles: vec![3, 2],
..OpAttrs::default()
};
assert!(matches!(
super::tag_to_op(OpTag::MatMul, &attrs),
Some(Op::MatMul)
));
let attrs4 = OpAttrs {
lhs_roles: vec![0, 0, 1, 3],
rhs_roles: vec![0, 0, 3, 2],
..OpAttrs::default()
};
assert!(matches!(
super::tag_to_op(OpTag::MatMul, &attrs4),
Some(Op::MatMul)
));
}
#[test]
fn tag_to_op_matmul_empty_roles_implicit_accept() {
assert!(matches!(
super::tag_to_op(OpTag::MatMul, &OpAttrs::default()),
Some(Op::MatMul)
));
}
#[test]
fn tag_to_op_matmul_rejects_noncanonical_roles() {
let transposed = OpAttrs {
lhs_roles: vec![3, 1],
rhs_roles: vec![3, 2],
..OpAttrs::default()
};
assert_eq!(super::tag_to_op(OpTag::MatMul, &transposed), None);
let multi_k = OpAttrs {
lhs_roles: vec![3, 3],
rhs_roles: vec![3, 2],
..OpAttrs::default()
};
assert_eq!(super::tag_to_op(OpTag::MatMul, &multi_k), None);
let freen_before_k = OpAttrs {
lhs_roles: vec![1, 3],
rhs_roles: vec![2, 3],
..OpAttrs::default()
};
assert_eq!(super::tag_to_op(OpTag::MatMul, &freen_before_k), None);
}
#[test]
fn tag_to_op_reconstructs_max_dim() {
assert!(matches!(
super::tag_to_op(
OpTag::MaxDim,
&OpAttrs {
axis: Some(1),
..OpAttrs::default()
}
),
Some(Op::MaxDim(1))
));
assert_eq!(super::tag_to_op(OpTag::MaxDim, &OpAttrs::default()), None);
assert_eq!(super::scalar_slot_arity(OpTag::MaxDim), 0);
}
#[test]
fn max_dim_axis_last_resolves_to_rank_minus_one() {
let attrs = OpAttrs {
axis_last: true,
..OpAttrs::default()
};
let resolved =
super::resolve_rel_attrs(OpTag::MaxDim, &attrs, &[vec![2, 3, 4]], &[vec![2, 3, 4]])
.expect("axis_last must resolve on MaxDim, not AxisLastUnsupported");
assert_eq!(resolved.axis, Some(2), "rank-3 operand → LAST = axis 2");
assert!(
!resolved.axis_last,
"rel carrier must be cleared post-resolve"
);
}
#[test]
fn tag_to_op_still_rejects_basis_gap_and_scan() {
assert_eq!(
super::tag_to_op(OpTag::Iota, &OpAttrs::default()),
None,
"Iota needs a len (target_shape) — empty attrs is a miss"
);
}
#[test]
fn tag_to_op_reconstructs_nested_fused() {
use crate::registry::{FusedOpParams, FusedOps};
let sm = OpAttrs {
fused_op: Some("SoftmaxLastDim".into()),
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::Fused, &sm),
Some(Op::Fused(
FusedOps::SOFTMAX_LAST_DIM,
FusedOpParams::SoftmaxLastDim
)),
"SoftmaxLastDim selector reconstructs the param-less nested fused op",
);
let smb = OpAttrs {
fused_op: Some("SoftmaxLastDimBackward".into()),
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::Fused, &smb),
Some(Op::Fused(
FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
FusedOpParams::SoftmaxLastDimBackward,
)),
"SoftmaxLastDimBackward selector reconstructs its param-less nested fused op",
);
assert_eq!(
super::tag_to_op(OpTag::Fused, &OpAttrs::default()),
None,
"an unset fused_op selector is an honest miss",
);
let bogus = OpAttrs {
fused_op: Some("NotARealFusedOp".into()),
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::Fused, &bogus),
None,
"an unknown name is an honest miss"
);
let paged = OpAttrs {
fused_op: Some("PagedAttn".into()),
..OpAttrs::default()
};
assert_eq!(
super::tag_to_op(OpTag::Fused, &paged),
None,
"a registered but param-carrying fused op is an honest miss (not in the fid->params map)",
);
}
#[test]
fn emit_reconstructs_nested_fused_shape_and_dtype() {
use crate::registry::{FusedOpParams, FusedOps};
use fuel_ir::{DType, Shape};
let region = PatternNode::Op {
op: OpTag::Fused,
attrs: OpAttrs {
fused_op: Some("SoftmaxLastDim".into()),
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
};
assert!(
super::validate_recipe(®ion).is_ok(),
"the nested-fused recipe validates as re-emittable (C-T2 tag_to_op arm)",
);
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3, 4]),
dtype: DType::F32,
});
let root = emit_region(&mut g, ®ion, &[x], &[]);
assert_eq!(
g.node(root).op,
Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
"the nested fused node is reconstructed as-is (2a)",
);
assert_eq!(
g.node(root).shape,
Shape::from_dims(&[2, 3, 4]),
"shape from softmax shape_rule (passthrough), NOT primitive_shape (which misses Fused)",
);
assert_eq!(
g.node(root).dtype,
DType::F32,
"dtype from softmax dtype_rule (passthrough)"
);
assert_eq!(
g.node(root).inputs,
vec![x],
"the fused node's single operand is the bound input"
);
}
#[test]
fn validate_representable_now_accepts_a_slice_region() {
let region = PatternNode::Op {
op: OpTag::Concat,
attrs: OpAttrs {
axis: Some(0),
..OpAttrs::default()
},
operands: vec![
PatternNode::Op {
op: OpTag::Neg,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::Slice,
attrs: OpAttrs {
axis: Some(0),
slice_start: Some(0),
slice_len: Some(1),
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
}],
},
PatternNode::Bind { index: 0 },
],
};
assert!(
super::validate_representable(®ion).is_ok(),
"slice/concat region must now validate"
);
}
#[test]
fn emit_gets_shape_right_for_a_reduction_region() {
use fuel_ir::{DType, Shape};
let region = PatternNode::Op {
op: OpTag::ReduceSumTo,
attrs: OpAttrs {
target_shape: vec![2, 1],
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
};
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 5]),
dtype: DType::F32,
});
let root = emit_region(&mut g, ®ion, &[x], &[]);
assert!(matches!(g.node(root).op, Op::ReduceSumTo(_)));
assert_eq!(
g.node(root).shape,
Shape::from_dims(&[2, 1]),
"emit must use the reduced shape, not operand[0]"
);
assert_eq!(g.node(root).dtype, DType::F32);
}
#[test]
fn emit_gets_dtype_right_for_a_cast_region() {
use fuel_ir::{DType, Shape};
let region = PatternNode::Op {
op: OpTag::Cast,
attrs: OpAttrs {
cast_dtype: Some("f16".into()),
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
};
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 3]),
dtype: DType::F32,
});
let root = emit_region(&mut g, ®ion, &[x], &[]);
assert!(matches!(g.node(root).op, Op::Cast(DType::F16)));
assert_eq!(
g.node(root).dtype,
DType::F16,
"emit must take Cast's target dtype, not operand[0]'s"
);
assert_eq!(g.node(root).shape, Shape::from_dims(&[3, 3]));
}
#[test]
fn emit_zero_operand_representable_region_is_panic_free() {
let region = PatternNode::Op {
op: OpTag::Add,
attrs: OpAttrs::default(),
operands: vec![],
};
let mut g = Graph::new();
let root = emit_region(&mut g, ®ion, &[], &[]);
assert!(
matches!(g.node(root).op, Op::Add),
"emit returns a node, not a panic"
);
}
#[test]
fn decompose_region_re_emits_relu_add() {
let id = register_runtime_fused("test::relu_add::decompose", relu_add_region()).unwrap();
let mut g = Graph::new();
let s = Shape::from_dims(&[4]);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![] }),
inputs: vec![a, b],
shape: s.clone(),
dtype: DType::F32,
});
let root = decompose_region(&mut g, fused);
assert!(matches!(g.node(root).op, Op::Relu));
let add_id = g.node(root).inputs[0];
assert!(matches!(g.node(add_id).op, Op::Add));
assert_eq!(g.node(add_id).inputs, vec![a, b]);
assert_eq!(g.node(root).shape, s);
assert_eq!(g.node(add_id).shape, s);
}
fn tanh_mul_scalar_slot_region() -> PatternNode {
PatternNode::Op {
op: OpTag::Tanh,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::MulScalar,
attrs: OpAttrs::default(), operands: vec![PatternNode::Bind { index: 0 }],
}],
}
}
#[test]
fn slot_template_registers_and_counts() {
let id =
register_runtime_fused("test::tanh_mul_scalar::slot", tanh_mul_scalar_slot_region())
.expect("a slot template is registrable");
let region = runtime_region(id).unwrap();
assert_eq!(count_scalar_slots(®ion), 1, "one open slot");
}
#[test]
fn decompose_fills_slots_from_the_node_scalars() {
let id =
register_runtime_fused("test::tanh_mul_scalar::fill", tanh_mul_scalar_slot_region())
.unwrap();
let mut g = Graph::new();
let s = Shape::from_dims(&[4]);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![2.5] }),
inputs: vec![a],
shape: s.clone(),
dtype: DType::F32,
});
let root = decompose_region(&mut g, fused);
assert!(matches!(g.node(root).op, Op::Tanh));
let ms = g.node(root).inputs[0];
assert!(
matches!(g.node(ms).op, Op::MulScalar(v) if v == 2.5),
"slot filled with the node's live scalar, got {:?}",
g.node(ms).op,
);
assert_eq!(g.node(ms).inputs, vec![a]);
}
#[test]
fn decompose_slot_count_mismatch_is_a_fixpoint_not_a_crash() {
let id = register_runtime_fused(
"test::tanh_mul_scalar::mismatch",
tanh_mul_scalar_slot_region(),
)
.unwrap();
let mut g = Graph::new();
let s = Shape::from_dims(&[4]);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![] }),
inputs: vec![a],
shape: s.clone(),
dtype: DType::F32,
});
assert_eq!(
decompose_region(&mut g, fused),
fused,
"mismatch ⇒ fixpoint"
);
}
fn mul_scalar_reduced_count_region() -> PatternNode {
use fuel_kernel_seam_types::shape_expr::{Dim, LAST};
PatternNode::Op {
op: OpTag::MulScalar,
attrs: OpAttrs {
scalar_rel: Some(Dim::Extent {
operand: 0,
axis: LAST,
}),
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
}
}
#[test]
fn scalar_rel_is_shape_derived_not_a_cursor_slot() {
let region = mul_scalar_reduced_count_region();
assert_eq!(
count_scalar_slots(®ion),
0,
"a scalar_rel node is shape-derived, never a params-cursor slot",
);
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 5]),
dtype: DType::F32,
});
let root = try_emit_region(&mut g, ®ion, &[x], &[])
.expect("scalar_rel resolves against the input shape");
assert!(
matches!(g.node(root).op, Op::MulScalar(v) if v == 5.0),
"scalar_rel = Extent(0, LAST) over x[3,5] ⇒ MulScalar(5.0), got {:?}",
g.node(root).op,
);
assert_eq!(g.node(root).inputs, vec![x]);
let mut g3 = Graph::new();
let x3 = g3.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 4, 7]),
dtype: DType::F32,
});
let root3 = try_emit_region(&mut g3, ®ion, &[x3], &[]).unwrap();
assert!(
matches!(g3.node(root3).op, Op::MulScalar(v) if v == 7.0),
"same datum, x[2,4,7] ⇒ MulScalar(7.0), got {:?}",
g3.node(root3).op,
);
}
#[test]
fn scalar_rel_and_baked_scalar_together_is_a_typed_conflict() {
use fuel_kernel_seam_types::shape_expr::{Dim, LAST};
let region = PatternNode::Op {
op: OpTag::MulScalar,
attrs: OpAttrs {
scalars: vec![2.0],
scalar_rel: Some(Dim::Extent {
operand: 0,
axis: LAST,
}),
..OpAttrs::default()
},
operands: vec![PatternNode::Bind { index: 0 }],
};
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 5]),
dtype: DType::F32,
});
assert_eq!(
try_emit_region(&mut g, ®ion, &[x], &[]),
Err(RelAttrError::RelAbsConflict { field: "scalars" }),
"scalar_rel XOR scalars — both set is a typed decline",
);
}
fn op_node(op: OpTag, attrs: OpAttrs, operands: Vec<PatternNode>) -> PatternNode {
PatternNode::Op {
op,
attrs,
operands,
}
}
fn bind(i: u8) -> PatternNode {
PatternNode::Bind { index: i }
}
fn assert_structural_eq(g: &Graph, a: NodeId, b: NodeId) {
if a == b {
return; }
let na = g.node(a);
let nb = g.node(b);
assert_eq!(na.op, nb.op, "op mismatch: {:?} vs {:?}", na.op, nb.op);
assert_eq!(
na.shape, nb.shape,
"shape mismatch at {:?} vs {:?}",
na.op, nb.op
);
assert_eq!(na.dtype, nb.dtype, "dtype mismatch at {:?}", na.op);
assert_eq!(
na.inputs.len(),
nb.inputs.len(),
"arity mismatch at {:?}",
na.op
);
for (&ia, &ib) in na.inputs.iter().zip(nb.inputs.iter()) {
assert_structural_eq(g, ia, ib);
}
}
fn frozen_legacy_softmax_decompose(
graph: &mut Graph,
id: NodeId,
_params: &FusedOpParams,
) -> NodeId {
let (x_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.shape.clone(), n.dtype)
};
let dims = x_shape.dims().to_vec();
let rank = dims.len();
let last = rank - 1;
let mut keepdim_dims = dims.clone();
keepdim_dims[last] = 1;
let keepdim_shape = Shape::from_dims(&keepdim_dims);
let m_id = graph.push(Node {
op: Op::ReduceMaxTo(keepdim_shape.clone()),
inputs: vec![x_id],
shape: keepdim_shape.clone(),
dtype,
});
let mb_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![m_id],
shape: x_shape.clone(),
dtype,
});
let s_id = graph.push(Node {
op: Op::Sub,
inputs: vec![x_id, mb_id],
shape: x_shape.clone(),
dtype,
});
let e_id = graph.push(Node {
op: Op::Exp,
inputs: vec![s_id],
shape: x_shape.clone(),
dtype,
});
let d_id = graph.push(Node {
op: Op::ReduceSumTo(keepdim_shape.clone()),
inputs: vec![e_id],
shape: keepdim_shape,
dtype,
});
let db_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![d_id],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Div,
inputs: vec![e_id, db_id],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_rope_decompose(
graph: &mut Graph,
id: NodeId,
_params: &FusedOpParams,
) -> NodeId {
let (x_id, cos_id, sin_id, x_shape, dtype) = {
let n = graph.node(id);
(
n.inputs[0],
n.inputs[1],
n.inputs[2],
n.shape.clone(),
n.dtype,
)
};
let dims = x_shape.dims().to_vec();
let rank = dims.len();
let seq = dims[rank - 2];
let d = dims[rank - 1];
let half = d / 2;
let last = rank - 1;
let mut broadcast_shape_dims: Vec<usize> = vec![1usize; rank];
broadcast_shape_dims[rank - 2] = seq;
broadcast_shape_dims[rank - 1] = d;
let broadcast_shape = Shape::from_dims(&broadcast_shape_dims);
let cos_reshaped_id = graph.push(Node {
op: Op::Reshape(broadcast_shape.clone()),
inputs: vec![cos_id],
shape: broadcast_shape.clone(),
dtype,
});
let sin_reshaped_id = graph.push(Node {
op: Op::Reshape(broadcast_shape.clone()),
inputs: vec![sin_id],
shape: broadcast_shape,
dtype,
});
let cos_bcast_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![cos_reshaped_id],
shape: x_shape.clone(),
dtype,
});
let sin_bcast_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![sin_reshaped_id],
shape: x_shape.clone(),
dtype,
});
let mut half_dims = dims.clone();
half_dims[last] = half;
let half_shape = Shape::from_dims(&half_dims);
let first_half_id = graph.push(Node {
op: Op::Slice {
dim: last,
start: 0,
len: half,
},
inputs: vec![x_id],
shape: half_shape.clone(),
dtype,
});
let second_half_id = graph.push(Node {
op: Op::Slice {
dim: last,
start: half,
len: half,
},
inputs: vec![x_id],
shape: half_shape.clone(),
dtype,
});
let neg_second_id = graph.push(Node {
op: Op::Neg,
inputs: vec![second_half_id],
shape: half_shape,
dtype,
});
let rotated_half_id = graph.push(Node {
op: Op::Concat { dim: last },
inputs: vec![neg_second_id, first_half_id],
shape: x_shape.clone(),
dtype,
});
let left_id = graph.push(Node {
op: Op::Mul,
inputs: vec![x_id, cos_bcast_id],
shape: x_shape.clone(),
dtype,
});
let right_id = graph.push(Node {
op: Op::Mul,
inputs: vec![rotated_half_id, sin_bcast_id],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Add,
inputs: vec![left_id, right_id],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_rms_norm_decompose(
graph: &mut Graph,
id: NodeId,
params: &FusedOpParams,
) -> NodeId {
let (x_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.shape.clone(), n.dtype)
};
let eps = match params {
FusedOpParams::RmsNormLastDim { eps } => *eps,
_ => return id,
};
let dims = x_shape.dims().to_vec();
let rank = dims.len();
let last = rank - 1;
let mut keepdim_dims = dims.clone();
keepdim_dims[last] = 1;
let keepdim_shape = Shape::from_dims(&keepdim_dims);
let mut reduced_dims = dims.clone();
reduced_dims.remove(last);
let reduced_shape = Shape::from_dims(&reduced_dims);
let sq_id = graph.push(Node {
op: Op::Sqr,
inputs: vec![x_id],
shape: x_shape.clone(),
dtype,
});
let mean_id = graph.push(Node {
op: Op::MeanDim(last),
inputs: vec![sq_id],
shape: reduced_shape,
dtype,
});
let mean_kd_id = graph.push(Node {
op: Op::Reshape(keepdim_shape.clone()),
inputs: vec![mean_id],
shape: keepdim_shape.clone(),
dtype,
});
let denom_sq_id = graph.push(Node {
op: Op::AddScalar(eps),
inputs: vec![mean_kd_id],
shape: keepdim_shape.clone(),
dtype,
});
let denom_id = graph.push(Node {
op: Op::Sqrt,
inputs: vec![denom_sq_id],
shape: keepdim_shape,
dtype,
});
let denom_bcast_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![denom_id],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Div,
inputs: vec![x_id, denom_bcast_id],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_layer_norm_decompose(
graph: &mut Graph,
id: NodeId,
params: &FusedOpParams,
) -> NodeId {
let (x_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.shape.clone(), n.dtype)
};
let eps = match params {
FusedOpParams::LayerNormLastDim { eps } => *eps,
_ => return id,
};
let dims = x_shape.dims().to_vec();
let rank = dims.len();
let last = rank - 1;
let mut keepdim_dims = dims.clone();
keepdim_dims[last] = 1;
let keepdim_shape = Shape::from_dims(&keepdim_dims);
let mut reduced_dims = dims.clone();
reduced_dims.remove(last);
let reduced_shape = Shape::from_dims(&reduced_dims);
let mean_id = graph.push(Node {
op: Op::MeanDim(last),
inputs: vec![x_id],
shape: reduced_shape.clone(),
dtype,
});
let mean_kd_id = graph.push(Node {
op: Op::Reshape(keepdim_shape.clone()),
inputs: vec![mean_id],
shape: keepdim_shape.clone(),
dtype,
});
let mean_bcast_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![mean_kd_id],
shape: x_shape.clone(),
dtype,
});
let centered_id = graph.push(Node {
op: Op::Sub,
inputs: vec![x_id, mean_bcast_id],
shape: x_shape.clone(),
dtype,
});
let centered_sq_id = graph.push(Node {
op: Op::Sqr,
inputs: vec![centered_id],
shape: x_shape.clone(),
dtype,
});
let var_id = graph.push(Node {
op: Op::MeanDim(last),
inputs: vec![centered_sq_id],
shape: reduced_shape,
dtype,
});
let var_kd_id = graph.push(Node {
op: Op::Reshape(keepdim_shape.clone()),
inputs: vec![var_id],
shape: keepdim_shape.clone(),
dtype,
});
let var_eps_id = graph.push(Node {
op: Op::AddScalar(eps),
inputs: vec![var_kd_id],
shape: keepdim_shape.clone(),
dtype,
});
let denom_id = graph.push(Node {
op: Op::Sqrt,
inputs: vec![var_eps_id],
shape: keepdim_shape,
dtype,
});
let denom_bcast_id = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![denom_id],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Div,
inputs: vec![centered_id, denom_bcast_id],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_layer_norm_backward_decompose(
graph: &mut Graph,
id: NodeId,
params: &FusedOpParams,
) -> NodeId {
let (x_id, g_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
};
let eps = match params {
FusedOpParams::LayerNormLastDimBackward { eps } => *eps,
_ => return id,
};
let dims = x_shape.dims().to_vec();
let last = dims.len() - 1;
let mut kd = dims.clone();
kd[last] = 1;
let keepdim = Shape::from_dims(&kd);
let mut rd = dims.clone();
rd.remove(last);
let reduced = Shape::from_dims(&rd);
let mean_b = |graph: &mut Graph, src: NodeId| -> NodeId {
let m = graph.push(Node {
op: Op::MeanDim(last),
inputs: vec![src],
shape: reduced.clone(),
dtype,
});
let m_kd = graph.push(Node {
op: Op::Reshape(keepdim.clone()),
inputs: vec![m],
shape: keepdim.clone(),
dtype,
});
graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![m_kd],
shape: x_shape.clone(),
dtype,
})
};
let mean_x = mean_b(graph, x_id);
let xc = graph.push(Node {
op: Op::Sub,
inputs: vec![x_id, mean_x],
shape: x_shape.clone(),
dtype,
});
let xc_sq = graph.push(Node {
op: Op::Sqr,
inputs: vec![xc],
shape: x_shape.clone(),
dtype,
});
let var = mean_b(graph, xc_sq);
let var_eps = graph.push(Node {
op: Op::AddScalar(eps),
inputs: vec![var],
shape: x_shape.clone(),
dtype,
});
let istd = graph.push(Node {
op: Op::Rsqrt,
inputs: vec![var_eps],
shape: x_shape.clone(),
dtype,
});
let xhat = graph.push(Node {
op: Op::Mul,
inputs: vec![xc, istd],
shape: x_shape.clone(),
dtype,
});
let mean_g = mean_b(graph, g_id);
let g_xhat = graph.push(Node {
op: Op::Mul,
inputs: vec![g_id, xhat],
shape: x_shape.clone(),
dtype,
});
let mean_gxh = mean_b(graph, g_xhat);
let t1 = graph.push(Node {
op: Op::Sub,
inputs: vec![g_id, mean_g],
shape: x_shape.clone(),
dtype,
});
let t2 = graph.push(Node {
op: Op::Mul,
inputs: vec![xhat, mean_gxh],
shape: x_shape.clone(),
dtype,
});
let inner = graph.push(Node {
op: Op::Sub,
inputs: vec![t1, t2],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Mul,
inputs: vec![istd, inner],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_rms_norm_backward_decompose(
graph: &mut Graph,
id: NodeId,
params: &FusedOpParams,
) -> NodeId {
let (x_id, up_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
};
let eps = match params {
FusedOpParams::RmsNormLastDimBackward { eps } => *eps,
_ => return id,
};
let dims = x_shape.dims().to_vec();
let last = dims.len() - 1;
let n = dims[last] as f64;
let mut kd = dims.clone();
kd[last] = 1;
let keepdim = Shape::from_dims(&kd);
let mut rd = dims.clone();
rd.remove(last);
let reduced = Shape::from_dims(&rd);
let sq = graph.push(Node {
op: Op::Sqr,
inputs: vec![x_id],
shape: x_shape.clone(),
dtype,
});
let mean = graph.push(Node {
op: Op::MeanDim(last),
inputs: vec![sq],
shape: reduced.clone(),
dtype,
});
let mean_kd = graph.push(Node {
op: Op::Reshape(keepdim.clone()),
inputs: vec![mean],
shape: keepdim.clone(),
dtype,
});
let denom_kd = graph.push(Node {
op: Op::AddScalar(eps),
inputs: vec![mean_kd],
shape: keepdim.clone(),
dtype,
});
let rrms_kd = graph.push(Node {
op: Op::Rsqrt,
inputs: vec![denom_kd],
shape: keepdim.clone(),
dtype,
});
let rrms_b = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![rrms_kd],
shape: x_shape.clone(),
dtype,
});
let gx = graph.push(Node {
op: Op::Mul,
inputs: vec![up_id, x_id],
shape: x_shape.clone(),
dtype,
});
let s = graph.push(Node {
op: Op::SumDim(last),
inputs: vec![gx],
shape: reduced,
dtype,
});
let s_kd = graph.push(Node {
op: Op::Reshape(keepdim.clone()),
inputs: vec![s],
shape: keepdim.clone(),
dtype,
});
let s_b = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![s_kd],
shape: x_shape.clone(),
dtype,
});
let ndenom_kd = graph.push(Node {
op: Op::MulScalar(n),
inputs: vec![denom_kd],
shape: keepdim.clone(),
dtype,
});
let ndenom_b = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![ndenom_kd],
shape: x_shape.clone(),
dtype,
});
let xs = graph.push(Node {
op: Op::Mul,
inputs: vec![x_id, s_b],
shape: x_shape.clone(),
dtype,
});
let term = graph.push(Node {
op: Op::Div,
inputs: vec![xs, ndenom_b],
shape: x_shape.clone(),
dtype,
});
let inner = graph.push(Node {
op: Op::Sub,
inputs: vec![up_id, term],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Mul,
inputs: vec![rrms_b, inner],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_reduce_max_to_backward_decompose(
graph: &mut Graph,
id: NodeId,
_params: &FusedOpParams,
) -> NodeId {
use fuel_ir::Scalar;
let (x_id, up_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
};
let target = graph.node(up_id).shape.clone();
let y = graph.push(Node {
op: Op::ReduceMaxTo(target.clone()),
inputs: vec![x_id],
shape: target.clone(),
dtype,
});
let y_b = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![y],
shape: x_shape.clone(),
dtype,
});
let mask_u8 = graph.push(Node {
op: Op::Equal,
inputs: vec![x_id, y_b],
shape: x_shape.clone(),
dtype: DType::U8,
});
let zeros = graph.push(Node {
op: Op::MulScalar(0.0),
inputs: vec![x_id],
shape: x_shape.clone(),
dtype,
});
let fill = match Scalar::one(dtype) {
Ok(s) => s,
Err(_) => return id,
};
let mask_f = graph.push(Node {
op: Op::MaskedFill { value: fill },
inputs: vec![zeros, mask_u8],
shape: x_shape.clone(),
dtype,
});
let ties = graph.push(Node {
op: Op::ReduceSumTo(target.clone()),
inputs: vec![mask_f],
shape: target.clone(),
dtype,
});
let share = graph.push(Node {
op: Op::Div,
inputs: vec![up_id, ties],
shape: target,
dtype,
});
let share_b = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![share],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Mul,
inputs: vec![mask_f, share_b],
shape: x_shape,
dtype,
})
}
fn frozen_legacy_powi_backward_decompose(
graph: &mut Graph,
id: NodeId,
params: &FusedOpParams,
) -> NodeId {
let (x_id, up_id, shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
};
let exp = match params {
FusedOpParams::PowIBackward { exp } => *exp,
_ => return id,
};
let pow = graph.push(Node {
op: Op::PowI(exp - 1),
inputs: vec![x_id],
shape: shape.clone(),
dtype,
});
let scaled = graph.push(Node {
op: Op::MulScalar(exp as f64),
inputs: vec![pow],
shape: shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Mul,
inputs: vec![scaled, up_id],
shape,
dtype,
})
}
fn frozen_legacy_fused_linear_decompose(
graph: &mut Graph,
id: NodeId,
_params: &FusedOpParams,
) -> NodeId {
let (a_id, b_id, bias_id, out_shape, dtype) = {
let n = graph.node(id);
(
n.inputs[0],
n.inputs[1],
n.inputs[2],
n.shape.clone(),
n.dtype,
)
};
let mm_id = graph.push(Node {
op: Op::MatMul,
inputs: vec![a_id, b_id],
shape: out_shape.clone(),
dtype,
});
let bias_bcst_id = graph.push(Node {
op: Op::BroadcastTo(out_shape.clone()),
inputs: vec![bias_id],
shape: out_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Add,
inputs: vec![mm_id, bias_bcst_id],
shape: out_shape,
dtype,
})
}
fn frozen_legacy_softmax_backward_decompose(
graph: &mut Graph,
id: NodeId,
_params: &FusedOpParams,
) -> NodeId {
let (s_id, g_id, x_shape, dtype) = {
let n = graph.node(id);
(n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
};
let mut kd = x_shape.dims().to_vec();
let last = kd.len() - 1;
kd[last] = 1;
let keepdim = Shape::from_dims(&kd);
let gs = graph.push(Node {
op: Op::Mul,
inputs: vec![g_id, s_id],
shape: x_shape.clone(),
dtype,
});
let summed = graph.push(Node {
op: Op::ReduceSumTo(keepdim.clone()),
inputs: vec![gs],
shape: keepdim,
dtype,
});
let summed_b = graph.push(Node {
op: Op::BroadcastTo(x_shape.clone()),
inputs: vec![summed],
shape: x_shape.clone(),
dtype,
});
let sub = graph.push(Node {
op: Op::Sub,
inputs: vec![g_id, summed_b],
shape: x_shape.clone(),
dtype,
});
graph.push(Node {
op: Op::Mul,
inputs: vec![s_id, sub],
shape: x_shape,
dtype,
})
}
#[test]
fn emit_matches_softmax_last_dim_decompose() {
use fuel_ir::{DType, Shape};
let mut g = Graph::new();
let sh = Shape::from_dims(&[2, 4]);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Const,
inputs: vec![x],
shape: sh.clone(),
dtype: DType::F32,
});
let oracle = frozen_legacy_softmax_decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
let kd = OpAttrs {
target_shape: vec![2, 1],
..OpAttrs::default()
};
let full = OpAttrs {
target_shape: vec![2, 4],
..OpAttrs::default()
};
let softmax_e = |kd: &OpAttrs, full: &OpAttrs| {
op_node(
OpTag::Exp,
OpAttrs::default(),
vec![op_node(
OpTag::Sub,
OpAttrs::default(),
vec![
bind(0),
op_node(
OpTag::BroadcastTo,
full.clone(),
vec![op_node(OpTag::ReduceMaxTo, kd.clone(), vec![bind(0)])],
),
],
)],
)
};
let region = op_node(
OpTag::Div,
OpAttrs::default(),
vec![
softmax_e(&kd, &full),
op_node(
OpTag::BroadcastTo,
full.clone(),
vec![op_node(
OpTag::ReduceSumTo,
kd.clone(),
vec![softmax_e(&kd, &full)],
)],
),
],
);
let emitted = emit_region(&mut g, ®ion, &[x], &[]);
assert_structural_eq(&g, oracle, emitted);
}
#[test]
fn emit_matches_rope_decompose() {
use fuel_ir::{DType, Shape};
let mut g = Graph::new();
let sh = Shape::from_dims(&[2, 4]); let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let cos = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let sin = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Const,
inputs: vec![x, cos, sin],
shape: sh.clone(),
dtype: DType::F32,
});
let oracle = frozen_legacy_rope_decompose(&mut g, fused, &FusedOpParams::Rope);
let full = OpAttrs {
target_shape: vec![2, 4],
..OpAttrs::default()
};
let sl_first = OpAttrs {
axis: Some(1),
slice_start: Some(0),
slice_len: Some(2),
..OpAttrs::default()
};
let sl_second = OpAttrs {
axis: Some(1),
slice_start: Some(2),
slice_len: Some(2),
..OpAttrs::default()
};
let cat = OpAttrs {
axis: Some(1),
..OpAttrs::default()
};
let bcast_reshape = |full: &OpAttrs, i: u8| {
op_node(
OpTag::BroadcastTo,
full.clone(),
vec![op_node(OpTag::Reshape, full.clone(), vec![bind(i)])],
)
};
let rotated = op_node(
OpTag::Concat,
cat,
vec![
op_node(
OpTag::Neg,
OpAttrs::default(),
vec![op_node(OpTag::Slice, sl_second, vec![bind(0)])],
),
op_node(OpTag::Slice, sl_first, vec![bind(0)]),
],
);
let left = op_node(
OpTag::Mul,
OpAttrs::default(),
vec![bind(0), bcast_reshape(&full, 1)],
);
let right = op_node(
OpTag::Mul,
OpAttrs::default(),
vec![rotated, bcast_reshape(&full, 2)],
);
let region = op_node(OpTag::Add, OpAttrs::default(), vec![left, right]);
let emitted = emit_region(&mut g, ®ion, &[x, cos, sin], &[]);
assert_structural_eq(&g, oracle, emitted);
}
#[test]
fn emit_matches_layer_norm_last_dim_decompose() {
use fuel_ir::{DType, Shape};
let mut g = Graph::new();
let sh = Shape::from_dims(&[2, 4]); let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Const,
inputs: vec![x],
shape: sh.clone(),
dtype: DType::F32,
});
let oracle = frozen_legacy_layer_norm_decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDim { eps: 1e-5 },
);
let kd = OpAttrs {
target_shape: vec![2, 1],
..OpAttrs::default()
};
let full = OpAttrs {
target_shape: vec![2, 4],
..OpAttrs::default()
};
let md = OpAttrs {
axis: Some(1),
..OpAttrs::default()
};
let eps_attrs = OpAttrs {
scalars: vec![1e-5],
..OpAttrs::default()
}; let centered = op_node(
OpTag::Sub,
OpAttrs::default(),
vec![
bind(0),
op_node(
OpTag::BroadcastTo,
full.clone(),
vec![op_node(
OpTag::Reshape,
kd.clone(),
vec![op_node(OpTag::MeanDim, md.clone(), vec![bind(0)])],
)],
),
],
);
let denom_bcast = op_node(
OpTag::BroadcastTo,
full.clone(),
vec![op_node(
OpTag::Sqrt,
OpAttrs::default(),
vec![op_node(
OpTag::AddScalar,
eps_attrs,
vec![op_node(
OpTag::Reshape,
kd.clone(),
vec![op_node(
OpTag::MeanDim,
md.clone(),
vec![op_node(
OpTag::Sqr,
OpAttrs::default(),
vec![centered.clone()],
)],
)],
)],
)],
)],
);
let region = op_node(OpTag::Div, OpAttrs::default(), vec![centered, denom_bcast]);
let emitted = emit_region(&mut g, ®ion, &[x], &[]);
assert_structural_eq(&g, oracle, emitted);
}
mod resolve_rel {
use super::super::{RelAttrError, resolve_rel_attrs};
use fuel_kernel_seam_types::shape_expr::{
Dim, LAST, SYMBOLIC, ShapeExpr, ShapeExprError, TAG_DIMS,
};
use fuel_kernel_seam_types::{OpAttrs, OpTag};
fn half_of_bind0_last() -> Dim {
Dim::Div(
Box::new(Dim::Extent {
operand: 0,
axis: LAST,
}),
Box::new(Dim::Const(2)),
)
}
#[test]
fn same_as_bind0_tracks_the_bind_shape() {
let attrs = OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
..OpAttrs::default()
};
let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![2, 3]], &[vec![2, 1]])
.expect("resolves");
assert_eq!(r.target_shape, vec![2, 3]);
assert!(
r.target_shape_rel.is_none(),
"resolved attrs are fully concrete"
);
let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![4, 5]], &[vec![4, 1]])
.expect("resolves");
assert_eq!(r.target_shape, vec![4, 5]);
}
#[test]
fn slice_bounds_track_the_bind_extent() {
let attrs = OpAttrs {
axis: Some(1),
slice_start_rel: Some(half_of_bind0_last()),
slice_len_rel: Some(half_of_bind0_last()),
..OpAttrs::default()
};
let r = resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]])
.expect("resolves at d=4");
assert_eq!(r.slice_start, Some(2));
assert_eq!(r.slice_len, Some(2));
assert!(r.slice_start_rel.is_none() && r.slice_len_rel.is_none());
let r = resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 8]], &[vec![2, 8]])
.expect("resolves at d=8");
assert_eq!(r.slice_start, Some(4));
assert_eq!(r.slice_len, Some(4));
}
#[test]
fn axis_last_resolves_per_tag() {
let attrs = OpAttrs {
axis_last: true,
..OpAttrs::default()
};
let r = resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![2, 4]]).expect("rank 2");
assert_eq!(r.axis, Some(1));
assert!(!r.axis_last, "resolved attrs are fully concrete");
let r =
resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![2, 3, 4]]).expect("rank 3");
assert_eq!(r.axis, Some(2));
let r =
resolve_rel_attrs(OpTag::Concat, &attrs, &[], &[vec![2, 3, 4]]).expect("concat");
assert_eq!(r.axis, Some(2));
let r =
resolve_rel_attrs(OpTag::Unsqueeze, &attrs, &[], &[vec![2, 4]]).expect("unsqueeze");
assert_eq!(r.dims, vec![2]);
assert!(!r.axis_last);
let r =
resolve_rel_attrs(OpTag::Squeeze, &attrs, &[], &[vec![2, 4, 1]]).expect("squeeze");
assert_eq!(r.dims, vec![2]);
}
#[test]
fn bind_out_of_range_is_a_typed_decline() {
let attrs = OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 3 }),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![2, 3]], &[vec![2, 3]]),
Err(RelAttrError::Expr(ShapeExprError::OperandOutOfRange {
operand: 3,
operands: 1
})),
);
let attrs = OpAttrs {
axis: Some(1),
slice_start_rel: Some(Dim::Extent {
operand: 7,
axis: LAST,
}),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
Err(RelAttrError::Expr(ShapeExprError::OperandOutOfRange {
operand: 7,
operands: 1
})),
);
}
#[test]
fn rel_and_abs_both_set_is_a_typed_conflict() {
let attrs = OpAttrs {
target_shape: vec![2, 3],
target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![2, 3]], &[vec![2, 3]]),
Err(RelAttrError::RelAbsConflict {
field: "target_shape"
}),
);
let attrs = OpAttrs {
axis: Some(1),
slice_start: Some(0),
slice_start_rel: Some(half_of_bind0_last()),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
Err(RelAttrError::RelAbsConflict {
field: "slice_start"
}),
);
let attrs = OpAttrs {
axis: Some(1),
slice_len: Some(2),
slice_len_rel: Some(half_of_bind0_last()),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
Err(RelAttrError::RelAbsConflict { field: "slice_len" }),
);
let attrs = OpAttrs {
axis: Some(0),
axis_last: true,
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![2, 4]]),
Err(RelAttrError::RelAbsConflict { field: "axis" }),
);
let attrs = OpAttrs {
dims: vec![0],
axis_last: true,
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Unsqueeze, &attrs, &[], &[vec![2, 4]]),
Err(RelAttrError::RelAbsConflict { field: "dims" }),
);
}
#[test]
fn negative_result_is_a_typed_decline() {
let neg = Dim::Sub(Box::new(Dim::Const(0)), Box::new(Dim::Const(2)));
let attrs = OpAttrs {
axis: Some(1),
slice_start_rel: Some(neg),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
Err(RelAttrError::Negative {
field: "slice_start",
value: -2
}),
);
}
#[test]
fn symbolic_extent_is_a_surfaced_gap_decline() {
let attrs = OpAttrs {
axis: Some(1),
slice_len_rel: Some(half_of_bind0_last()),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, SYMBOLIC]], &[vec![2, 4]]),
Err(RelAttrError::SymbolicGap { field: "slice_len" }),
);
let attrs = OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(
OpTag::BroadcastTo,
&attrs,
&[vec![2, SYMBOLIC]],
&[vec![2, 4]]
),
Err(RelAttrError::SymbolicGap {
field: "target_shape"
}),
);
}
#[test]
fn two_operand_max_frame_declines_instead_of_a_partial_shape() {
let binds = vec![vec![2, 1], vec![1, 3]]; for operand in [0u8, 1] {
let attrs = OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand }),
..OpAttrs::default()
};
let child = binds[operand as usize].clone();
assert_eq!(
resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &binds, &[child]),
Err(RelAttrError::FrameNotExpressible {
field: "target_shape",
frame: vec![2, 3],
missing_ctor: TAG_DIMS,
}),
"SameAs {{ operand: {operand} }} must not resolve to a PARTIAL frame",
);
}
}
#[test]
fn frame_guard_does_not_fire_when_an_operand_carries_the_frame() {
let attrs = OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
..OpAttrs::default()
};
let binds = vec![vec![2, 3, 4], vec![4]];
let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &binds, &[vec![4]])
.expect("bind0 carries the frame");
assert_eq!(r.target_shape, vec![2, 3, 4]);
let sub = OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 1 }),
..OpAttrs::default()
};
let r = resolve_rel_attrs(OpTag::BroadcastTo, &sub, &binds, &[vec![4]])
.expect("sub-frame broadcast stays expressible");
assert_eq!(r.target_shape, vec![4]);
let mm = vec![vec![8, 4096], vec![4096, 1024]];
let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &mm, &[vec![8, 1]])
.expect("no joint frame ⇒ no guard");
assert_eq!(r.target_shape, vec![8, 4096]);
let two = vec![vec![2, 1], vec![1, 3]];
let r = resolve_rel_attrs(OpTag::Reshape, &attrs, &two, &[vec![2, 1]])
.expect("Reshape's target is not a frame claim");
assert_eq!(r.target_shape, vec![2, 1]);
}
#[test]
fn axis_last_on_an_axisless_tag_or_without_a_child_declines() {
let attrs = OpAttrs {
axis_last: true,
..OpAttrs::default()
};
assert_eq!(
resolve_rel_attrs(OpTag::Add, &attrs, &[], &[vec![2, 4], vec![2, 4]]),
Err(RelAttrError::AxisLastUnsupported { tag: OpTag::Add }),
);
assert_eq!(
resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[]),
Err(RelAttrError::NoChildOperand),
);
assert_eq!(
resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![]]),
Err(RelAttrError::Expr(ShapeExprError::AxisOutOfRange {
axis: LAST,
rank: 0
})),
);
}
#[test]
fn rel_free_attrs_pass_through_unchanged() {
let attrs = OpAttrs {
axis: Some(1),
slice_start: Some(2),
slice_len: Some(3),
..OpAttrs::default()
};
let r = resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]])
.expect("no-op");
assert_eq!(r, attrs);
}
}
mod emit_rel {
use super::super::*;
use super::{assert_structural_eq, bind, op_node};
use fuel_ir::{DType, Shape};
use fuel_kernel_seam_types::shape_expr::{Dim, ShapeExpr};
fn cst(g: &mut Graph, dims: &[usize]) -> NodeId {
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(dims),
dtype: DType::F32,
})
}
fn bcast_same_as_0() -> OpAttrs {
OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
..OpAttrs::default()
}
}
#[test]
fn rel_region_emits_polymorphically_across_shapes() {
let region = op_node(
OpTag::Add,
OpAttrs::default(),
vec![
bind(0),
op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
],
);
let mut g = Graph::new();
let x1 = cst(&mut g, &[2, 3]);
let t1 = cst(&mut g, &[1, 3]);
let r1 = emit_region(&mut g, ®ion, &[x1, t1], &[]);
assert!(matches!(g.node(r1).op, Op::Add));
assert_eq!(g.node(r1).shape, Shape::from_dims(&[2, 3]));
let b1 = g.node(r1).inputs[1];
assert_eq!(g.node(b1).op, Op::BroadcastTo(Shape::from_dims(&[2, 3])));
assert_eq!(g.node(b1).shape, Shape::from_dims(&[2, 3]));
let x2 = cst(&mut g, &[4, 5]);
let t2 = cst(&mut g, &[1, 5]);
let r2 = emit_region(&mut g, ®ion, &[x2, t2], &[]);
assert_eq!(g.node(r2).shape, Shape::from_dims(&[4, 5]));
let b2 = g.node(r2).inputs[1];
assert_eq!(g.node(b2).op, Op::BroadcastTo(Shape::from_dims(&[4, 5])));
}
#[test]
fn two_operand_max_frame_region_declines_through_emit() {
let region = op_node(
OpTag::Mul,
OpAttrs::default(),
vec![
op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(0)]),
op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
],
);
let mut g = Graph::new();
let a = cst(&mut g, &[2, 1]);
let b = cst(&mut g, &[1, 3]);
assert_eq!(
try_emit_region(&mut g, ®ion, &[a, b], &[]),
Err(RelAttrError::FrameNotExpressible {
field: "target_shape",
frame: vec![2, 3],
missing_ctor: fuel_kernel_seam_types::shape_expr::TAG_DIMS,
}),
);
}
#[test]
fn broadcast_rank_raise_materializes_the_legacy_reshape_pad() {
let region = op_node(
OpTag::Mul,
OpAttrs::default(),
vec![
bind(0),
op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
],
);
let mut g = Graph::new();
let x = cst(&mut g, &[2, 3, 4]);
let t = cst(&mut g, &[4]);
let pad_shape = Shape::from_dims(&[1, 1, 4]);
let full = Shape::from_dims(&[2, 3, 4]);
let pad = g.push(Node {
op: Op::Reshape(pad_shape.clone()),
inputs: vec![t],
shape: pad_shape,
dtype: DType::F32,
});
let bc = g.push(Node {
op: Op::BroadcastTo(full.clone()),
inputs: vec![pad],
shape: full.clone(),
dtype: DType::F32,
});
let reference = g.push(Node {
op: Op::Mul,
inputs: vec![x, bc],
shape: full,
dtype: DType::F32,
});
let emitted = emit_region(&mut g, ®ion, &[x, t], &[]);
assert_structural_eq(&g, reference, emitted);
}
#[test]
fn concrete_broadcast_rank_raise_also_pads() {
let region = op_node(
OpTag::BroadcastTo,
OpAttrs {
target_shape: vec![2, 3, 4],
..OpAttrs::default()
},
vec![bind(0)],
);
let mut g = Graph::new();
let t = cst(&mut g, &[3, 4]);
let emitted = emit_region(&mut g, ®ion, &[t], &[]);
assert_eq!(
g.node(emitted).op,
Op::BroadcastTo(Shape::from_dims(&[2, 3, 4]))
);
let pad = g.node(emitted).inputs[0];
assert_eq!(
g.node(pad).op,
Op::Reshape(Shape::from_dims(&[1, 3, 4])),
"rank-raise inserts the legacy 1-padded-left Reshape",
);
assert_eq!(g.node(pad).inputs, vec![t]);
}
#[test]
fn scalar_cursor_fill_stays_pre_order_after_the_reorder() {
let region = op_node(
OpTag::AddScalar,
OpAttrs::default(),
vec![op_node(OpTag::MulScalar, OpAttrs::default(), vec![bind(0)])],
);
let mut g = Graph::new();
let x = cst(&mut g, &[4]);
let root = emit_region(&mut g, ®ion, &[x], &[10.0, 20.0]);
assert!(
matches!(g.node(root).op, Op::AddScalar(v) if v == 10.0),
"parent takes scalars[0] (pre-order), got {:?}",
g.node(root).op,
);
let child = g.node(root).inputs[0];
assert!(
matches!(g.node(child).op, Op::MulScalar(v) if v == 20.0),
"child takes scalars[1], got {:?}",
g.node(child).op,
);
}
#[test]
fn validate_accepts_a_rel_attr_region() {
let region = op_node(
OpTag::Add,
OpAttrs::default(),
vec![
bind(0),
op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
],
);
register_runtime_fused("t3::rel_bcast", region)
.expect("a rel-attr region is registrable");
}
#[test]
fn validate_rejects_rel_abs_conflict_and_bind_out_of_range() {
let conflicted = op_node(
OpTag::BroadcastTo,
OpAttrs {
target_shape: vec![2, 3],
target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
..OpAttrs::default()
},
vec![bind(0)],
);
assert_eq!(
register_runtime_fused("t3::conflict", conflicted),
Err(RuntimeFusedError::InvalidRelAttrs {
tag: OpTag::BroadcastTo,
error: RelAttrError::RelAbsConflict {
field: "target_shape"
},
}),
);
let oob = op_node(
OpTag::BroadcastTo,
OpAttrs {
target_shape_rel: Some(ShapeExpr::SameAs { operand: 7 }),
..OpAttrs::default()
},
vec![bind(0)],
);
assert_eq!(
register_runtime_fused("t3::oob", oob),
Err(RuntimeFusedError::InvalidRelAttrs {
tag: OpTag::BroadcastTo,
error: RelAttrError::Expr(
fuel_kernel_seam_types::shape_expr::ShapeExprError::OperandOutOfRange {
operand: 7,
operands: 1,
},
),
}),
);
}
#[test]
fn decompose_rel_resolution_failure_is_a_fixpoint_not_a_crash() {
let neg = Dim::Sub(Box::new(Dim::Const(0)), Box::new(Dim::Const(2)));
let region = op_node(
OpTag::Slice,
OpAttrs {
axis: Some(1),
slice_start_rel: Some(neg),
slice_len: Some(1),
..OpAttrs::default()
},
vec![bind(0)],
);
let id = register_runtime_fused("t3::neg_slice", region)
.expect("a value-dependent decline still registers");
let mut g = Graph::new();
let x = cst(&mut g, &[2, 4]);
let fused = g.push(Node {
op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![] }),
inputs: vec![x],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
assert_eq!(
decompose_region(&mut g, fused),
fused,
"resolution decline ⇒ fixpoint"
);
}
#[test]
fn try_emit_region_surfaces_typed_resolution_errors() {
let neg = Dim::Sub(Box::new(Dim::Const(0)), Box::new(Dim::Const(2)));
let region = op_node(
OpTag::Slice,
OpAttrs {
axis: Some(1),
slice_start_rel: Some(neg),
slice_len: Some(1),
..OpAttrs::default()
},
vec![bind(0)],
);
let mut g = Graph::new();
let x = cst(&mut g, &[2, 4]);
assert_eq!(
try_emit_region(&mut g, ®ion, &[x], &[]),
Err(RelAttrError::Negative {
field: "slice_start",
value: -2
}),
);
let region = op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(0)]);
let dyn_shape = Shape::from_dims(&[2, 8]).with_dynamic_axis(1, 0, fuel_ir::SymId(0));
let d = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: dyn_shape,
dtype: DType::F32,
});
assert_eq!(
try_emit_region(&mut g, ®ion, &[d], &[]),
Err(RelAttrError::SymbolicGap {
field: "target_shape"
}),
);
}
}
#[test]
fn emit_shares_repeated_slot_free_subtrees() {
let region = op_node(
OpTag::Add,
OpAttrs::default(),
vec![
op_node(OpTag::Exp, OpAttrs::default(), vec![bind(0)]),
op_node(OpTag::Exp, OpAttrs::default(), vec![bind(0)]),
],
);
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let root = emit_region(&mut g, ®ion, &[x], &[]);
let add = g.node(root);
assert!(matches!(add.op, Op::Add));
assert_eq!(
add.inputs[0], add.inputs[1],
"structurally-equal slot-free subtrees must share ONE emitted node",
);
}
#[test]
fn emit_does_not_share_subtrees_with_open_scalar_slots() {
let region = op_node(
OpTag::Add,
OpAttrs::default(),
vec![
op_node(OpTag::MulScalar, OpAttrs::default(), vec![bind(0)]),
op_node(OpTag::MulScalar, OpAttrs::default(), vec![bind(0)]),
],
);
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let root = emit_region(&mut g, ®ion, &[x], &[2.0, 3.0]);
let a = g.node(root).inputs[0];
let b = g.node(root).inputs[1];
assert_ne!(a, b, "open-slot subtrees are never shared");
assert!(matches!(g.node(a).op, Op::MulScalar(v) if v == 2.0));
assert!(matches!(g.node(b).op, Op::MulScalar(v) if v == 3.0));
}
#[test]
fn emit_matches_cast_over_add_reference() {
use fuel_ir::{DType, Shape};
let mut g = Graph::new();
let sh = Shape::from_dims(&[4]);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![a, b],
shape: sh.clone(),
dtype: DType::F32,
});
let reference = g.push(Node {
op: Op::Cast(DType::F16),
inputs: vec![add],
shape: sh.clone(),
dtype: DType::F16,
});
let region = op_node(
OpTag::Cast,
OpAttrs {
cast_dtype: Some("f16".into()),
..OpAttrs::default()
},
vec![op_node(
OpTag::Add,
OpAttrs::default(),
vec![bind(0), bind(1)],
)],
);
let emitted = emit_region(&mut g, ®ion, &[a, b], &[]);
assert_structural_eq(&g, reference, emitted);
}
mod recipe_bridge {
use super::super::*;
use super::{bind, frozen_legacy_softmax_decompose, op_node};
use crate::registry::{FusedOps, decompose_via_recipe};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
let node = g.node(id);
match &node.op {
Op::Const => leaves.get(&id).expect("leaf data provided").clone(),
Op::Exp => eval(g, node.inputs[0], leaves)
.iter()
.map(|v| v.exp())
.collect(),
Op::Sub => {
let a = eval(g, node.inputs[0], leaves);
let b = eval(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x - y).collect()
}
Op::Div => {
let a = eval(g, node.inputs[0], leaves);
let b = eval(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x / y).collect()
}
Op::MaxDim(_) | Op::ReduceMaxTo(_) => {
let input = eval(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().copied().fold(f64::NEG_INFINITY, f64::max))
.collect()
}
Op::SumDim(_) | Op::ReduceSumTo(_) => {
let input = eval(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input.chunks(last).map(|row| row.iter().sum()).collect()
}
Op::Unsqueeze { .. } | Op::Reshape(_) => eval(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval(g, node.inputs[0], leaves);
let out_n: usize = target.dims().iter().product();
let last = *target.dims().last().unwrap();
assert_eq!(
input.len() * last,
out_n,
"broadcast is a last-dim repeat in these graphs",
);
input
.iter()
.flat_map(|&v| std::iter::repeat_n(v, last))
.collect()
}
other => panic!("eval: unhandled op {other:?}"),
}
}
fn softmax_fused_node(g: &mut Graph, dims: &[usize]) -> (NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
inputs: vec![x],
shape: sh,
dtype: DType::F32,
});
(x, fused)
}
#[test]
fn softmax_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (x, fused) = softmax_fused_node(&mut g, &dims);
let sh = Shape::from_dims(&dims);
let new_root = crate::registry::softmax_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::SoftmaxLastDim,
);
assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
assert_eq!(g.node(new_root).shape, sh, "softmax is shape-preserving");
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root =
frozen_legacy_softmax_decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
let n: usize = dims.iter().product();
let data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.37).sin() * 3.0 - 0.5)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, data);
let got = eval(&g, new_root, &leaves);
let want = eval(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"softmax[{i}] at {dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn softmax_recipe_emits_the_nine_node_shared_spelling() {
let mut g = Graph::new();
let (x, fused) = softmax_fused_node(&mut g, &[2, 4]);
let sh = Shape::from_dims(&[2, 4]);
let root = crate::registry::softmax_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::SoftmaxLastDim,
);
assert!(matches!(g.node(root).op, Op::Div));
let e = g.node(root).inputs[0];
let db = g.node(root).inputs[1];
assert!(matches!(g.node(e).op, Op::Exp));
assert_eq!(g.node(db).op, Op::BroadcastTo(sh.clone()));
let u2 = g.node(db).inputs[0];
assert!(matches!(g.node(u2).op, Op::Unsqueeze { dim: 1 }));
assert_eq!(g.node(u2).shape, Shape::from_dims(&[2, 1]));
let d = g.node(u2).inputs[0];
assert!(matches!(g.node(d).op, Op::SumDim(1)));
assert_eq!(g.node(d).shape, Shape::from_dims(&[2]));
assert_eq!(
g.node(d).inputs[0],
e,
"the denominator reduces the SHARED Exp node — identity-share, not a duplicate",
);
let s = g.node(e).inputs[0];
assert!(matches!(g.node(s).op, Op::Sub));
assert_eq!(g.node(s).inputs[0], x);
let mb = g.node(s).inputs[1];
assert_eq!(g.node(mb).op, Op::BroadcastTo(sh.clone()));
let u1 = g.node(mb).inputs[0];
assert!(matches!(g.node(u1).op, Op::Unsqueeze { dim: 1 }));
let m = g.node(u1).inputs[0];
assert!(matches!(g.node(m).op, Op::MaxDim(1)));
assert_eq!(g.node(m).inputs[0], x);
assert_eq!(
crate::topo_order_multi(&g, &[root]).len(),
10,
"MaxDim/Unsqueeze/Bcast/Sub/Exp/SumDim/Unsqueeze/Bcast/Div + leaf",
);
}
#[test]
fn softmax_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, fused) = softmax_fused_node(&mut g, &[2, 4]);
let before = g.len();
let out =
crate::registry::softmax_last_dim::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
#[test]
fn decompose_via_recipe_declines_an_unknown_token_recipe() {
let fabricated = op_node(
OpTag::Clamp,
OpAttrs {
scalars: vec![0.0, 1.0],
..OpAttrs::default()
},
vec![bind(0)],
);
let mut g = Graph::new();
let (_x, fused) = softmax_fused_node(&mut g, &[2, 4]);
let before = g.len();
let out = decompose_via_recipe(&mut g, fused, &fabricated, Some(Vec::new()));
assert_eq!(out, fused, "semantics-absent token ⇒ honest-miss fixpoint");
assert_eq!(
g.len(),
before,
"declined BEFORE any emission — no partial nodes"
);
}
#[test]
fn decompose_via_recipe_bind_arity_mismatch_is_a_fixpoint() {
let recipe = op_node(OpTag::Add, OpAttrs::default(), vec![bind(0), bind(1)]);
let mut g = Graph::new();
let (_x, fused) = softmax_fused_node(&mut g, &[2, 4]);
let before = g.len();
let out = decompose_via_recipe(&mut g, fused, &recipe, Some(Vec::new()));
assert_eq!(out, fused, "bind/input arity mismatch ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
mod rope_recipe {
use super::super::*;
use super::frozen_legacy_rope_decompose;
use crate::registry::{FusedOps, rope};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval_rope(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
let node = g.node(id);
match &node.op {
Op::Const => leaves.get(&id).expect("leaf data provided").clone(),
Op::Neg => eval_rope(g, node.inputs[0], leaves)
.iter()
.map(|v| -v)
.collect(),
Op::Mul => {
let a = eval_rope(g, node.inputs[0], leaves);
let b = eval_rope(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x * y).collect()
}
Op::Add => {
let a = eval_rope(g, node.inputs[0], leaves);
let b = eval_rope(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x + y).collect()
}
Op::Reshape(_) => eval_rope(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval_rope(g, node.inputs[0], leaves);
let out_n: usize = target.dims().iter().product();
(0..out_n).map(|i| input[i % input.len()]).collect()
}
Op::Slice { dim, start, len } => {
let input = eval_rope(g, node.inputs[0], leaves);
let in_dims = g.node(node.inputs[0]).shape.dims().to_vec();
let last = in_dims.len() - 1;
assert_eq!(*dim, last, "rope slices along the last axis");
let row = in_dims[last];
input
.chunks(row)
.flat_map(|r| r[*start..*start + *len].to_vec())
.collect()
}
Op::Concat { dim } => {
let a = eval_rope(g, node.inputs[0], leaves);
let b = eval_rope(g, node.inputs[1], leaves);
let a_last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
let b_last = *g.node(node.inputs[1]).shape.dims().last().unwrap();
let last = g.node(node.inputs[0]).shape.dims().len() - 1;
assert_eq!(*dim, last, "rope concats along the last axis");
let mut out = Vec::with_capacity(a.len() + b.len());
let mut ai = a.chunks(a_last);
let mut bi = b.chunks(b_last);
loop {
match (ai.next(), bi.next()) {
(Some(ra), Some(rb)) => {
out.extend_from_slice(ra);
out.extend_from_slice(rb);
}
(None, None) => break,
_ => panic!("concat row-count mismatch"),
}
}
out
}
other => panic!("eval_rope: unhandled op {other:?}"),
}
}
fn rope_fused_node(g: &mut Graph, x_dims: &[usize]) -> (NodeId, NodeId, NodeId, NodeId) {
let rank = x_dims.len();
let table_dims = [x_dims[rank - 2], x_dims[rank - 1]];
let x_sh = Shape::from_dims(x_dims);
let t_sh = Shape::from_dims(&table_dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: x_sh.clone(),
dtype: DType::F32,
});
let cos = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: t_sh.clone(),
dtype: DType::F32,
});
let sin = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: t_sh,
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(FusedOps::ROPE, FusedOpParams::Rope),
inputs: vec![x, cos, sin],
shape: x_sh,
dtype: DType::F32,
});
(x, cos, sin, fused)
}
#[test]
fn rope_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for x_dims in [vec![2usize, 4], vec![1, 2, 3, 8]] {
let mut g = Graph::new();
let (x, cos, sin, fused) = rope_fused_node(&mut g, &x_dims);
let x_sh = Shape::from_dims(&x_dims);
let rank = x_dims.len();
let seq = x_dims[rank - 2];
let d = x_dims[rank - 1];
let new_root = rope::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_ne!(new_root, fused, "recipe decompose must fire at {x_dims:?}");
assert_eq!(g.node(new_root).shape, x_sh, "rope is shape-preserving");
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_rope_decompose(&mut g, fused, &FusedOpParams::Rope);
let x_n: usize = x_dims.iter().product();
let t_n = seq * d;
let x_data: Vec<f64> = (0..x_n)
.map(|i| ((i as f64) * 0.31).sin() * 2.0 - 0.4)
.collect();
let cos_data: Vec<f64> =
(0..t_n).map(|i| ((i as f64) * 0.17 + 0.5).cos()).collect();
let sin_data: Vec<f64> =
(0..t_n).map(|i| ((i as f64) * 0.23 - 0.2).sin()).collect();
let mut leaves = HashMap::new();
leaves.insert(x, x_data);
leaves.insert(cos, cos_data);
leaves.insert(sin, sin_data);
let got = eval_rope(&g, new_root, &leaves);
let want = eval_rope(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len(), "element count at {x_dims:?}");
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"rope[{i}] at {x_dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn rope_recipe_is_byte_identical_to_legacy_at_rank_raise() {
let mut g = Graph::new();
let (_x, _cos, _sin, fused) = rope_fused_node(&mut g, &[1, 2, 3, 8]);
let recipe_root = rope::decompose(&mut g, fused, &FusedOpParams::Rope);
let legacy_root = frozen_legacy_rope_decompose(&mut g, fused, &FusedOpParams::Rope);
super::assert_structural_eq(&g, recipe_root, legacy_root);
let reachable = crate::topo_order_multi(&g, &[recipe_root]);
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
reshapes, 2,
"rank-raise materializes both legacy prep Reshapes (D4)"
);
}
#[test]
fn rope_recipe_elides_the_noop_prep_reshape_at_equal_rank() {
let mut g = Graph::new();
let (_x, _cos, _sin, fused) = rope_fused_node(&mut g, &[2, 4]);
let root = rope::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_ne!(root, fused, "recipe decompose fires");
let reachable = crate::topo_order_multi(&g, &[root]);
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
reshapes, 0,
"equal-rank broadcast needs no prep Reshape (D4)"
);
let op_nodes = reachable
.iter()
.filter(|&&n| !matches!(g.node(n).op, Op::Const))
.count();
assert_eq!(
op_nodes, 9,
"the 9-node rope recipe (2×Bcast/2×Slice/Neg/Concat/2×Mul/Add)"
);
}
#[test]
fn rope_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, _cos, _sin, fused) = rope_fused_node(&mut g, &[2, 4]);
let before = g.len();
let out = rope::decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
mod norm_recipe {
use super::super::*;
use super::{frozen_legacy_layer_norm_decompose, frozen_legacy_rms_norm_decompose};
use crate::registry::{FusedOps, layer_norm_last_dim, rms_norm_last_dim};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval_norm(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
let node = g.node(id);
match &node.op {
Op::Const => leaves.get(&id).expect("leaf data provided").clone(),
Op::Sqr => eval_norm(g, node.inputs[0], leaves)
.iter()
.map(|v| v * v)
.collect(),
Op::Sqrt => eval_norm(g, node.inputs[0], leaves)
.iter()
.map(|v| v.sqrt())
.collect(),
Op::AddScalar(e) => eval_norm(g, node.inputs[0], leaves)
.iter()
.map(|v| v + e)
.collect(),
Op::Sub => {
let a = eval_norm(g, node.inputs[0], leaves);
let b = eval_norm(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x - y).collect()
}
Op::Div => {
let a = eval_norm(g, node.inputs[0], leaves);
let b = eval_norm(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x / y).collect()
}
Op::MeanDim(_) => {
let input = eval_norm(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().sum::<f64>() / last as f64)
.collect()
}
Op::Unsqueeze { .. } | Op::Reshape(_) => eval_norm(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval_norm(g, node.inputs[0], leaves);
let out_n: usize = target.dims().iter().product();
let last = *target.dims().last().unwrap();
assert_eq!(
input.len() * last,
out_n,
"broadcast is a last-dim repeat in these graphs",
);
input
.iter()
.flat_map(|&v| std::iter::repeat_n(v, last))
.collect()
}
other => panic!("eval_norm: unhandled op {other:?}"),
}
}
fn rms_norm_fused_node(g: &mut Graph, dims: &[usize], eps: f64) -> (NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::RMS_NORM_LAST_DIM,
FusedOpParams::RmsNormLastDim { eps },
),
inputs: vec![x],
shape: sh,
dtype: DType::F32,
});
(x, fused)
}
fn layer_norm_fused_node(g: &mut Graph, dims: &[usize], eps: f64) -> (NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::LAYER_NORM_LAST_DIM,
FusedOpParams::LayerNormLastDim { eps },
),
inputs: vec![x],
shape: sh,
dtype: DType::F32,
});
(x, fused)
}
#[test]
fn rms_norm_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (x, fused) = rms_norm_fused_node(&mut g, &dims, 1e-5);
let sh = Shape::from_dims(&dims);
let new_root = rms_norm_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDim { eps: 1e-5 },
);
assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
assert_eq!(g.node(new_root).shape, sh, "rms_norm is shape-preserving");
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_rms_norm_decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDim { eps: 1e-5 },
);
let n: usize = dims.iter().product();
let data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.37).sin() * 3.0 - 0.5)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, data);
let got = eval_norm(&g, new_root, &leaves);
let want = eval_norm(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"rms_norm[{i}] at {dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn layer_norm_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (x, fused) = layer_norm_fused_node(&mut g, &dims, 1e-5);
let sh = Shape::from_dims(&dims);
let new_root = layer_norm_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDim { eps: 1e-5 },
);
assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
assert_eq!(g.node(new_root).shape, sh, "layer_norm is shape-preserving");
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_layer_norm_decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDim { eps: 1e-5 },
);
let n: usize = dims.iter().product();
let data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.29).cos() * 2.0 + 0.3)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, data);
let got = eval_norm(&g, new_root, &leaves);
let want = eval_norm(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"layer_norm[{i}] at {dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn rms_norm_recipe_uses_the_unsqueeze_keepdim_swap() {
let mut g = Graph::new();
let (_x, fused) = rms_norm_fused_node(&mut g, &[2, 4], 1e-5);
let root = rms_norm_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDim { eps: 1e-5 },
);
assert_ne!(root, fused, "recipe decompose fires");
let reachable = crate::topo_order_multi(&g, &[root]);
let unsqueezes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
.count();
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
unsqueezes, 1,
"keepdim restored via Unsqueeze append (D3 swap)"
);
assert_eq!(reshapes, 0, "no baked keepdim Reshape after the D3 swap");
}
#[test]
fn layer_norm_recipe_uses_the_unsqueeze_keepdim_swap() {
let mut g = Graph::new();
let (_x, fused) = layer_norm_fused_node(&mut g, &[2, 4], 1e-5);
let root = layer_norm_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDim { eps: 1e-5 },
);
assert_ne!(root, fused, "recipe decompose fires");
let reachable = crate::topo_order_multi(&g, &[root]);
let unsqueezes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
.count();
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
unsqueezes, 2,
"both keepdim restores via Unsqueeze append (D3 swap)"
);
assert_eq!(
reshapes, 0,
"no baked keepdim Reshape / no D4 pad after the swap"
);
let op_nodes = reachable
.iter()
.filter(|&&n| !matches!(g.node(n).op, Op::Const))
.count();
assert_eq!(op_nodes, 11, "11 op nodes with `centered` identity-shared");
}
#[test]
fn rms_norm_recipe_eps_flows_through_the_open_slot() {
let dims = [2usize, 4];
let n: usize = dims.iter().product();
let data: Vec<f64> = (0..n)
.map(|i| 0.001 * (((i as f64) * 0.37).sin() + 1.2))
.collect();
let realize = |eps: f64| -> Vec<f64> {
let mut g = Graph::new();
let (x, fused) = rms_norm_fused_node(&mut g, &dims, eps);
let root = rms_norm_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDim { eps },
);
let mut leaves = HashMap::new();
leaves.insert(x, data.clone());
eval_norm(&g, root, &leaves)
};
let a = realize(1e-5);
let b = realize(1e-6);
assert_eq!(a.len(), b.len());
assert_ne!(
a, b,
"different eps must change the output — proves projection→slot, not a baked constant",
);
}
#[test]
fn layer_norm_recipe_eps_flows_through_the_open_slot() {
let dims = [2usize, 4];
let n: usize = dims.iter().product();
let data: Vec<f64> = (0..n)
.map(|i| 1.0 + 0.001 * ((i as f64) * 0.37).sin())
.collect();
let realize = |eps: f64| -> Vec<f64> {
let mut g = Graph::new();
let (x, fused) = layer_norm_fused_node(&mut g, &dims, eps);
let root = layer_norm_last_dim::decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDim { eps },
);
let mut leaves = HashMap::new();
leaves.insert(x, data.clone());
eval_norm(&g, root, &leaves)
};
let a = realize(1e-5);
let b = realize(1e-6);
assert_eq!(a.len(), b.len());
assert_ne!(
a, b,
"different eps must change the output — proves projection→slot, not a baked constant",
);
}
#[test]
fn rms_norm_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, fused) = rms_norm_fused_node(&mut g, &[2, 4], 1e-5);
let before = g.len();
let out = rms_norm_last_dim::decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
#[test]
fn layer_norm_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, fused) = layer_norm_fused_node(&mut g, &[2, 4], 1e-5);
let before = g.len();
let out = layer_norm_last_dim::decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
mod softmax_backward_recipe {
use super::super::*;
use super::frozen_legacy_softmax_backward_decompose;
use crate::registry::{FusedOps, softmax_last_dim_backward};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval_bwd(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
if let Some(v) = leaves.get(&id) {
return v.clone();
}
let node = g.node(id);
match &node.op {
Op::Mul => {
let a = eval_bwd(g, node.inputs[0], leaves);
let b = eval_bwd(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x * y).collect()
}
Op::Sub => {
let a = eval_bwd(g, node.inputs[0], leaves);
let b = eval_bwd(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x - y).collect()
}
Op::SumDim(_) | Op::ReduceSumTo(_) => {
let input = eval_bwd(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input.chunks(last).map(|row| row.iter().sum()).collect()
}
Op::Unsqueeze { .. } | Op::Reshape(_) => eval_bwd(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval_bwd(g, node.inputs[0], leaves);
let out_n: usize = target.dims().iter().product();
let last = *target.dims().last().unwrap();
assert_eq!(
input.len() * last,
out_n,
"broadcast is a last-dim repeat in these graphs",
);
input
.iter()
.flat_map(|&v| std::iter::repeat_n(v, last))
.collect()
}
other => panic!("eval_bwd: unhandled op {other:?}"),
}
}
fn softmax_backward_fused_node(g: &mut Graph, dims: &[usize]) -> (NodeId, NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let s = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let up = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
FusedOpParams::SoftmaxLastDimBackward,
),
inputs: vec![s, up],
shape: sh,
dtype: DType::F32,
});
(s, up, fused)
}
#[test]
fn softmax_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (s, up, fused) = softmax_backward_fused_node(&mut g, &dims);
let sh = Shape::from_dims(&dims);
let new_root = softmax_last_dim_backward::decompose(
&mut g,
fused,
&FusedOpParams::SoftmaxLastDimBackward,
);
assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
assert_eq!(
g.node(new_root).shape,
sh,
"softmax backward is shape-preserving"
);
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_softmax_backward_decompose(
&mut g,
fused,
&FusedOpParams::SoftmaxLastDimBackward,
);
let n: usize = dims.iter().product();
let s_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.29).sin() * 0.5 + 0.5)
.collect();
let g_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.53).cos() * 2.0 - 0.3)
.collect();
let mut leaves = HashMap::new();
leaves.insert(s, s_data);
leaves.insert(up, g_data);
let got = eval_bwd(&g, new_root, &leaves);
let want = eval_bwd(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"softmax_backward[{i}] at {dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn softmax_backward_recipe_uses_the_sumdim_unsqueeze_swap() {
let mut g = Graph::new();
let (_s, _up, fused) = softmax_backward_fused_node(&mut g, &[2, 4]);
let root = softmax_last_dim_backward::decompose(
&mut g,
fused,
&FusedOpParams::SoftmaxLastDimBackward,
);
assert_ne!(root, fused, "recipe decompose fires");
assert!(
matches!(g.node(root).op, Op::Mul),
"backward root is the outer Mul"
);
let reachable = crate::topo_order_multi(&g, &[root]);
let sumdims = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::SumDim(_)))
.count();
let unsqueezes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
.count();
let reduce_sum_tos = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::ReduceSumTo(_)))
.count();
assert_eq!(sumdims, 1, "the reduce is SumDim(last) — the D3 swap");
assert_eq!(unsqueezes, 1, "keepdim restored via Unsqueeze append");
assert_eq!(
reduce_sum_tos, 0,
"no baked keepdim ReduceSumTo after the swap"
);
}
#[test]
fn softmax_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_s, _up, fused) = softmax_backward_fused_node(&mut g, &[2, 4]);
let before = g.len();
let out = softmax_last_dim_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
#[test]
fn softmax_backward_reaches_the_recipe_through_autograd() {
let dev: std::sync::Arc<dyn fuel_backend_contract::DynBackendDevice> =
std::sync::Arc::new(fuel_cpu_backend::dyn_impl::CpuBackendDevice);
let x = crate::NodeHandle::from_f32(
vec![0.1f32, -0.2, 0.3, 0.4, -0.5, 0.6],
Shape::from_dims(&[2, 3]),
&dev,
)
.unwrap();
let y = x.softmax_last_dim();
let y_id = y.id();
let grads = y.backward();
let g_x = grads.get(&x).expect("softmax has an input gradient");
let handle = g_x.graph();
let bwd_id = g_x.id();
let up_id = {
let gr = handle.read().unwrap();
let node = gr.node(bwd_id).clone();
match node.op {
Op::Fused(fid, params) => {
assert_eq!(
fid,
FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
"autograd emits the registry backward fused op",
);
assert!(matches!(params, FusedOpParams::SoftmaxLastDimBackward));
}
other => panic!("expected the backward fused op, got {other:?}"),
}
assert_eq!(
node.inputs[0], y_id,
"backward input 0 = the forward softmax output"
);
node.inputs[1]
};
let (new_root, legacy_root, sh) = {
let mut gr = handle.write().unwrap();
let sh = gr.node(bwd_id).shape.clone();
let new_root = softmax_last_dim_backward::decompose(
&mut gr,
bwd_id,
&FusedOpParams::SoftmaxLastDimBackward,
);
let legacy_root = frozen_legacy_softmax_backward_decompose(
&mut gr,
bwd_id,
&FusedOpParams::SoftmaxLastDimBackward,
);
(new_root, legacy_root, sh)
};
let gr = handle.read().unwrap();
assert_ne!(
new_root, bwd_id,
"the autograd backward node decomposes via the recipe"
);
assert_eq!(gr.node(new_root).shape, sh, "shape-preserving");
let reachable = crate::topo_order_multi(&gr, &[new_root]);
assert!(
reachable
.iter()
.any(|&n| matches!(gr.node(n).op, Op::SumDim(_))),
"the autograd path reaches the D3 SumDim spelling",
);
let n: usize = sh.dims().iter().product();
let s_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.31).sin() * 0.5 + 0.5)
.collect();
let g_data: Vec<f64> = (0..n).map(|i| ((i as f64) * 0.47).cos() - 0.1).collect();
let mut leaves = HashMap::new();
leaves.insert(y_id, s_data);
leaves.insert(up_id, g_data);
let got = eval_bwd(&gr, new_root, &leaves);
let want = eval_bwd(&gr, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"autograd softmax_backward[{i}]: recipe={a} vs legacy={b}",
);
}
}
}
mod layer_norm_backward_recipe {
use super::super::*;
use super::frozen_legacy_layer_norm_backward_decompose;
use crate::registry::{FusedOps, layer_norm_last_dim_backward};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval_lnb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
if let Some(v) = leaves.get(&id) {
return v.clone();
}
let node = g.node(id);
match &node.op {
Op::Mul => {
let a = eval_lnb(g, node.inputs[0], leaves);
let b = eval_lnb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x * y).collect()
}
Op::Sub => {
let a = eval_lnb(g, node.inputs[0], leaves);
let b = eval_lnb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x - y).collect()
}
Op::Sqr => eval_lnb(g, node.inputs[0], leaves)
.iter()
.map(|v| v * v)
.collect(),
Op::Rsqrt => eval_lnb(g, node.inputs[0], leaves)
.iter()
.map(|v| 1.0 / v.sqrt())
.collect(),
Op::AddScalar(e) => eval_lnb(g, node.inputs[0], leaves)
.iter()
.map(|v| v + e)
.collect(),
Op::MeanDim(_) => {
let input = eval_lnb(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().sum::<f64>() / last as f64)
.collect()
}
Op::Unsqueeze { .. } | Op::Reshape(_) => eval_lnb(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval_lnb(g, node.inputs[0], leaves);
let out_n: usize = target.dims().iter().product();
let last = *target.dims().last().unwrap();
assert_eq!(
input.len() * last,
out_n,
"broadcast is a last-dim repeat in these graphs",
);
input
.iter()
.flat_map(|&v| std::iter::repeat_n(v, last))
.collect()
}
other => panic!("eval_lnb: unhandled op {other:?}"),
}
}
fn ln_backward_fused_node(
g: &mut Graph,
dims: &[usize],
eps: f64,
) -> (NodeId, NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let up = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::LAYER_NORM_LAST_DIM_BACKWARD,
FusedOpParams::LayerNormLastDimBackward { eps },
),
inputs: vec![x, up],
shape: sh,
dtype: DType::F32,
});
(x, up, fused)
}
#[test]
fn layer_norm_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (x, up, fused) = ln_backward_fused_node(&mut g, &dims, 1e-5);
let sh = Shape::from_dims(&dims);
let new_root = layer_norm_last_dim_backward::decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
);
assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
assert_eq!(
g.node(new_root).shape,
sh,
"layer_norm backward is shape-preserving"
);
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_layer_norm_backward_decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
);
let n: usize = dims.iter().product();
let x_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.29).sin() * 2.0 + 0.3)
.collect();
let g_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.53).cos() * 1.5 - 0.2)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, x_data);
leaves.insert(up, g_data);
let got = eval_lnb(&g, new_root, &leaves);
let want = eval_lnb(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"layer_norm_backward[{i}] at {dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn layer_norm_backward_recipe_uses_the_unsqueeze_keepdim_swap() {
let mut g = Graph::new();
let (_x, _up, fused) = ln_backward_fused_node(&mut g, &[2, 4], 1e-5);
let root = layer_norm_last_dim_backward::decompose(
&mut g,
fused,
&FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
);
assert_ne!(root, fused, "recipe decompose fires");
assert!(
matches!(g.node(root).op, Op::Mul),
"backward root is the outer Mul(istd, inner)"
);
let reachable = crate::topo_order_multi(&g, &[root]);
let unsqueezes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
.count();
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert!(
unsqueezes >= 1,
"keepdim restored via Unsqueeze append (D3 swap)"
);
assert_eq!(
reshapes, 0,
"no baked keepdim Reshape / no D4 pad after the D3 swap"
);
}
#[test]
fn layer_norm_backward_recipe_open_slot_xhat_is_recomputed_then_cse_recollapses() {
let count_ops = |g: &Graph, root: NodeId| -> usize {
crate::topo_order_multi(g, &[root])
.iter()
.filter(|&&n| !matches!(g.node(n).op, Op::Const))
.count()
};
let mut gr = Graph::new();
let (_x, _up, fused) = ln_backward_fused_node(&mut gr, &[2, 4], 1e-5);
let recipe_root = layer_norm_last_dim_backward::decompose(
&mut gr,
fused,
&FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
);
let recipe_n = count_ops(&gr, recipe_root);
let mut gl = Graph::new();
let (_x2, _up2, fused2) = ln_backward_fused_node(&mut gl, &[2, 4], 1e-5);
let legacy_root = frozen_legacy_layer_norm_backward_decompose(
&mut gl,
fused2,
&FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
);
let legacy_n = count_ops(&gl, legacy_root);
assert!(
recipe_n > legacy_n,
"RISK-A: emit's slot-free share does NOT dedup the eps-slot-carrying \
xhat/istd — recipe base map {recipe_n} > legacy {legacy_n}",
);
let shared: crate::SharedGraph = std::sync::Arc::new(std::sync::RwLock::new(gr));
let cse_roots = crate::opt::optimize(&shared, &[recipe_root]);
let gr = shared.read().unwrap();
let cse_n = count_ops(&gr, cse_roots[0]);
assert!(
cse_n < recipe_n,
"downstream CSE re-collapses the redundant recompute: {cse_n} < {recipe_n}",
);
assert_eq!(
cse_n, legacy_n,
"CSE collapses the recipe base map back to the legacy DAG node count",
);
}
#[test]
fn layer_norm_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, _up, fused) = ln_backward_fused_node(&mut g, &[2, 4], 1e-5);
let before = g.len();
let out = layer_norm_last_dim_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
mod rms_norm_backward_recipe {
use super::super::*;
use super::frozen_legacy_rms_norm_backward_decompose;
use crate::registry::{FusedOps, rms_norm_last_dim_backward};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval_rnb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
if let Some(v) = leaves.get(&id) {
return v.clone();
}
let node = g.node(id);
match &node.op {
Op::Mul => {
let a = eval_rnb(g, node.inputs[0], leaves);
let b = eval_rnb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x * y).collect()
}
Op::Sub => {
let a = eval_rnb(g, node.inputs[0], leaves);
let b = eval_rnb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x - y).collect()
}
Op::Div => {
let a = eval_rnb(g, node.inputs[0], leaves);
let b = eval_rnb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x / y).collect()
}
Op::Sqr => eval_rnb(g, node.inputs[0], leaves)
.iter()
.map(|v| v * v)
.collect(),
Op::Rsqrt => eval_rnb(g, node.inputs[0], leaves)
.iter()
.map(|v| 1.0 / v.sqrt())
.collect(),
Op::AddScalar(e) => eval_rnb(g, node.inputs[0], leaves)
.iter()
.map(|v| v + e)
.collect(),
Op::MulScalar(m) => eval_rnb(g, node.inputs[0], leaves)
.iter()
.map(|v| v * m)
.collect(),
Op::MeanDim(_) => {
let input = eval_rnb(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().sum::<f64>() / last as f64)
.collect()
}
Op::SumDim(_) => {
let input = eval_rnb(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().sum::<f64>())
.collect()
}
Op::Unsqueeze { .. } | Op::Reshape(_) => eval_rnb(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval_rnb(g, node.inputs[0], leaves);
let out_n: usize = target.dims().iter().product();
let last = *target.dims().last().unwrap();
assert_eq!(
input.len() * last,
out_n,
"broadcast is a last-dim repeat in these graphs",
);
input
.iter()
.flat_map(|&v| std::iter::repeat_n(v, last))
.collect()
}
other => panic!("eval_rnb: unhandled op {other:?}"),
}
}
fn rms_backward_fused_node(
g: &mut Graph,
dims: &[usize],
eps: f64,
) -> (NodeId, NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let up = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::RMS_NORM_LAST_DIM_BACKWARD,
FusedOpParams::RmsNormLastDimBackward { eps },
),
inputs: vec![x, up],
shape: sh,
dtype: DType::F32,
});
(x, up, fused)
}
#[test]
fn rms_norm_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (x, up, fused) = rms_backward_fused_node(&mut g, &dims, 1e-5);
let sh = Shape::from_dims(&dims);
let new_root = rms_norm_last_dim_backward::decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDimBackward { eps: 1e-5 },
);
assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
assert_eq!(
g.node(new_root).shape,
sh,
"rms_norm backward is shape-preserving"
);
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_rms_norm_backward_decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDimBackward { eps: 1e-5 },
);
let n: usize = dims.iter().product();
let x_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.41).sin() * 2.0 + 0.7)
.collect();
let g_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.61).cos() * 1.3 - 0.15)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, x_data);
leaves.insert(up, g_data);
let got = eval_rnb(&g, new_root, &leaves);
let want = eval_rnb(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"rms_norm_backward[{i}] at {dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn rms_norm_backward_recipe_uses_the_unsqueeze_keepdim_swap_and_shape_scalar() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
let mut g = Graph::new();
let (_x, _up, fused) = rms_backward_fused_node(&mut g, &dims, 1e-5);
let root = rms_norm_last_dim_backward::decompose(
&mut g,
fused,
&FusedOpParams::RmsNormLastDimBackward { eps: 1e-5 },
);
assert_ne!(root, fused, "recipe decompose fires at {dims:?}");
assert!(
matches!(g.node(root).op, Op::Mul),
"backward root is the outer Mul(rrms, inner)"
);
let reachable = crate::topo_order_multi(&g, &[root]);
let unsqueezes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
.count();
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert!(
unsqueezes >= 1,
"keepdim restored via Unsqueeze append (D3 swap) at {dims:?}"
);
assert_eq!(
reshapes, 0,
"no baked keepdim Reshape / no D4 pad after the D3 swap at {dims:?}"
);
let n = *dims.last().unwrap() as f64;
let mul_scalars: Vec<f64> = reachable
.iter()
.filter_map(|&nid| match g.node(nid).op {
Op::MulScalar(v) => Some(v),
_ => None,
})
.collect();
assert!(
mul_scalars.contains(&n),
"a MulScalar(n={n}) resolved from x's last-axis extent at {dims:?}, got {mul_scalars:?}",
);
}
}
}
mod fused_linear_recipe {
use super::super::*;
use super::frozen_legacy_fused_linear_decompose;
use crate::registry::{FusedOps, fused_linear};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn broadcast(input: &[f64], in_shape: &[usize], target: &[usize]) -> Vec<f64> {
let rank = target.len();
let pad = rank - in_shape.len();
let mut real_strides = vec![0isize; in_shape.len()];
let mut s = 1isize;
for i in (0..in_shape.len()).rev() {
real_strides[i] = s;
s *= in_shape[i] as isize;
}
let mut in_strides = vec![0isize; rank];
for (i, stride) in in_strides.iter_mut().enumerate() {
if i >= pad {
let id = i - pad;
*stride = if in_shape[id] == 1 {
0
} else {
real_strides[id]
};
}
}
let out_n: usize = target.iter().product();
let mut out = Vec::with_capacity(out_n);
let mut idx = vec![0usize; rank];
for _ in 0..out_n {
let fi: isize = (0..rank).map(|i| idx[i] as isize * in_strides[i]).sum();
out.push(input[fi as usize]);
for i in (0..rank).rev() {
idx[i] += 1;
if idx[i] < target[i] {
break;
}
idx[i] = 0;
}
}
out
}
fn matmul(a: &[f64], ash: &[usize], b: &[f64], bsh: &[usize]) -> Vec<f64> {
let r = ash.len();
let (m, k, n) = (ash[r - 2], ash[r - 1], bsh[r - 1]);
let batch: usize = ash[..r - 2].iter().product();
let (a_bs, b_bs, o_bs) = (m * k, k * n, m * n);
let mut out = vec![0.0f64; batch * o_bs];
for bi in 0..batch {
for i in 0..m {
for j in 0..n {
let mut acc = 0.0f64;
for kk in 0..k {
acc += a[bi * a_bs + i * k + kk] * b[bi * b_bs + kk * n + j];
}
out[bi * o_bs + i * n + j] = acc;
}
}
}
out
}
fn eval_fl(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
if let Some(v) = leaves.get(&id) {
return v.clone();
}
let node = g.node(id);
match &node.op {
Op::MatMul => {
let a = eval_fl(g, node.inputs[0], leaves);
let b = eval_fl(g, node.inputs[1], leaves);
let ash: Vec<usize> = g.node(node.inputs[0]).shape.dims().to_vec();
let bsh: Vec<usize> = g.node(node.inputs[1]).shape.dims().to_vec();
matmul(&a, &ash, &b, &bsh)
}
Op::Reshape(_) => eval_fl(g, node.inputs[0], leaves),
Op::BroadcastTo(target) => {
let input = eval_fl(g, node.inputs[0], leaves);
let in_shape: Vec<usize> = g.node(node.inputs[0]).shape.dims().to_vec();
broadcast(&input, &in_shape, target.dims())
}
Op::Add => {
let a = eval_fl(g, node.inputs[0], leaves);
let b = eval_fl(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x + y).collect()
}
other => panic!("eval_fl: unhandled op {other:?}"),
}
}
fn fused_linear_fused_node(
g: &mut Graph,
a_dims: &[usize],
b_dims: &[usize],
) -> (NodeId, NodeId, NodeId, NodeId) {
let ar = a_dims.len();
let n = b_dims[b_dims.len() - 1];
let mut out_dims = a_dims[..ar - 2].to_vec();
out_dims.push(a_dims[ar - 2]);
out_dims.push(n);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(a_dims),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(b_dims),
dtype: DType::F32,
});
let bias = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[n]),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(FusedOps::FUSED_LINEAR, FusedOpParams::FusedLinear),
inputs: vec![a, b, bias],
shape: Shape::from_dims(&out_dims),
dtype: DType::F32,
});
(a, b, bias, fused)
}
#[test]
fn fused_linear_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for (a_dims, b_dims) in [
(vec![2usize, 3], vec![3usize, 4]),
(vec![2, 2, 3], vec![2, 3, 4]),
] {
let mut g = Graph::new();
let (a, b, bias, fused) = fused_linear_fused_node(&mut g, &a_dims, &b_dims);
let out_sh = g.node(fused).shape.clone();
let new_root = fused_linear::decompose(&mut g, fused, &FusedOpParams::FusedLinear);
assert_ne!(new_root, fused, "recipe decompose fires at a={a_dims:?}");
assert_eq!(
g.node(new_root).shape,
out_sh,
"output = matmul output shape"
);
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_fused_linear_decompose(
&mut g,
fused,
&FusedOpParams::FusedLinear,
);
let an: usize = a_dims.iter().product();
let bn: usize = b_dims.iter().product();
let bias_n = b_dims[b_dims.len() - 1];
let a_data: Vec<f64> = (0..an)
.map(|i| ((i as f64) * 0.31).sin() * 1.5 - 0.2)
.collect();
let b_data: Vec<f64> = (0..bn)
.map(|i| ((i as f64) * 0.47).cos() * 0.8 + 0.1)
.collect();
let bias_data: Vec<f64> = (0..bias_n)
.map(|i| ((i as f64) * 0.7).sin() * 2.0)
.collect();
let mut leaves = HashMap::new();
leaves.insert(a, a_data);
leaves.insert(b, b_data);
leaves.insert(bias, bias_data);
let got = eval_fl(&g, new_root, &leaves);
let want = eval_fl(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (x, y)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
x.to_bits(),
y.to_bits(),
"fused_linear[{i}] a={a_dims:?}: recipe={x} vs legacy={y}",
);
}
}
}
#[test]
fn fused_linear_recipe_drives_withdim_broadcast_with_a_d4_pad() {
let mut g = Graph::new();
let (_a, _b, _bias, fused) = fused_linear_fused_node(&mut g, &[2, 3], &[3, 4]);
let root = fused_linear::decompose(&mut g, fused, &FusedOpParams::FusedLinear);
assert_ne!(root, fused, "recipe decompose fires");
assert!(
matches!(g.node(root).op, Op::Add),
"root is Add(mm, bias_bcst)"
);
let reachable = crate::topo_order_multi(&g, &[root]);
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
reshapes, 1,
"WithDim rank-raise materializes exactly one D4 pad Reshape on the bias",
);
let bcst = reachable
.iter()
.find(|&&n| matches!(g.node(n).op, Op::BroadcastTo(_)))
.expect("bias broadcast present");
assert_eq!(
g.node(*bcst).shape.dims(),
&[2, 4],
"bias broadcast target = the matmul output shape (WithDim resolved)",
);
}
#[test]
fn fused_linear_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_a, _b, _bias, fused) = fused_linear_fused_node(&mut g, &[2, 3], &[3, 4]);
let before = g.len();
let out = fused_linear::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
mod scan_recipe_roundtrip {
use super::super::*;
use super::{bind, op_node};
use crate::opt::base_map_hash;
use crate::{ScanEmit, ScanPredicate, ScanRole};
use fuel_ir::storage::{OutputViewSpec, compose_bundle};
use fuel_ir::{DType, Shape};
use fuel_kernel_seam_types::{SCAN_ROLE_CARRY, SCAN_ROLE_ELEM};
use std::sync::Arc;
fn placeholder(role: u8, index: u32) -> PatternNode {
PatternNode::Op {
op: OpTag::ScanPlaceholder,
attrs: OpAttrs {
scan_role: Some(role),
scan_index: Some(index),
..OpAttrs::default()
},
operands: vec![],
}
}
fn scan_recipe() -> PatternNode {
let sum = || {
op_node(
OpTag::Add,
OpAttrs::default(),
vec![
placeholder(SCAN_ROLE_CARRY, 0),
placeholder(SCAN_ROLE_ELEM, 0),
],
)
};
PatternNode::Op {
op: OpTag::Scan,
attrs: OpAttrs {
scan_n_carries: Some(1),
scan_n_xs: Some(1),
scan_bound: Some(3),
scan_emit: Some(0), scan_early_exit: None,
..OpAttrs::default()
},
operands: vec![
bind(0), bind(1), bind(2), sum(), op_node(OpTag::Mul, OpAttrs::default(), vec![sum(), bind(2)]), ],
}
}
fn reference_scan(g: &mut Graph) -> NodeId {
let cs = Shape::from_dims(&[1]);
let xs = Shape::from_dims(&[3, 1]);
let init = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: cs.clone(),
dtype: DType::F32,
});
let xs0 = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: xs,
dtype: DType::F32,
});
let konst = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: cs.clone(),
dtype: DType::F32,
});
let carry = g.push(Node {
op: Op::ScanPlaceholder {
role: ScanRole::Carry,
index: 0,
},
inputs: vec![],
shape: cs.clone(),
dtype: DType::F32,
});
let elem = g.push(Node {
op: Op::ScanPlaceholder {
role: ScanRole::Elem,
index: 0,
},
inputs: vec![],
shape: cs.clone(),
dtype: DType::F32,
});
let sum = g.push(Node {
op: Op::Add,
inputs: vec![carry, elem],
shape: cs.clone(),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::Mul,
inputs: vec![sum, konst],
shape: cs.clone(),
dtype: DType::F32,
});
g.push(Node {
op: Op::Scan {
n_carries: 1,
n_xs: 1,
bound: 3,
emit: ScanEmit::All,
early_exit: None,
},
inputs: vec![init, xs0, konst, sum, y],
shape: Shape::from_dims(&[3, 1]),
dtype: DType::F32,
})
}
#[test]
fn tag_to_op_reconstructs_scan_and_placeholder() {
let scan_attrs = OpAttrs {
scan_n_carries: Some(1),
scan_n_xs: Some(1),
scan_bound: Some(3),
scan_emit: Some(0),
scan_early_exit: None,
..OpAttrs::default()
};
assert!(matches!(
tag_to_op(OpTag::Scan, &scan_attrs),
Some(Op::Scan {
n_carries: 1,
n_xs: 1,
bound: 3,
emit: ScanEmit::All,
early_exit: None
}),
));
let fin = OpAttrs {
scan_emit: Some(1),
scan_early_exit: Some(true),
..scan_attrs.clone()
};
assert!(matches!(
tag_to_op(OpTag::Scan, &fin),
Some(Op::Scan {
emit: ScanEmit::Final,
early_exit: Some(ScanPredicate),
..
}),
));
assert!(tag_to_op(OpTag::Scan, &OpAttrs::default()).is_none());
let carry = OpAttrs {
scan_role: Some(SCAN_ROLE_CARRY),
scan_index: Some(0),
..OpAttrs::default()
};
assert!(matches!(
tag_to_op(OpTag::ScanPlaceholder, &carry),
Some(Op::ScanPlaceholder {
role: ScanRole::Carry,
index: 0
}),
));
let elem = OpAttrs {
scan_role: Some(SCAN_ROLE_ELEM),
scan_index: Some(2),
..OpAttrs::default()
};
assert!(matches!(
tag_to_op(OpTag::ScanPlaceholder, &elem),
Some(Op::ScanPlaceholder {
role: ScanRole::Elem,
index: 2
}),
));
assert!(tag_to_op(OpTag::ScanPlaceholder, &OpAttrs::default()).is_none());
}
#[test]
fn scan_recipe_validates_as_representable() {
assert_eq!(scan_recipe().bind_indices(), vec![0, 1, 2]);
assert!(
validate_recipe(&scan_recipe()).is_ok(),
"scan recipe is a total, representable region"
);
}
#[test]
fn scan_recipe_reemits_to_the_same_base_map_as_the_imperative_scan() {
let mut gr = Graph::new();
let ref_scan = reference_scan(&mut gr);
let want = base_map_hash(&gr, ref_scan);
let mut ge = Graph::new();
let cs = Shape::from_dims(&[1]);
let init = ge.push(Node {
op: Op::Const,
inputs: vec![],
shape: cs.clone(),
dtype: DType::F32,
});
let xs0 = ge.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 1]),
dtype: DType::F32,
});
let konst = ge.push(Node {
op: Op::Const,
inputs: vec![],
shape: cs,
dtype: DType::F32,
});
let root = emit_region(&mut ge, &scan_recipe(), &[init, xs0, konst], &[]);
assert!(
matches!(
ge.node(root).op,
Op::Scan {
n_carries: 1,
n_xs: 1,
bound: 3,
emit: ScanEmit::All,
early_exit: None
}
),
"recipe re-emits an Op::Scan terminal"
);
let got = base_map_hash(&ge, root);
assert_eq!(
got, want,
"scan recipe re-emit base map == imperative scan base map"
);
}
fn view_scan_recipe(slot: u32) -> PatternNode {
let sum = || {
op_node(
OpTag::Add,
OpAttrs::default(),
vec![
placeholder(SCAN_ROLE_CARRY, 0),
placeholder(SCAN_ROLE_ELEM, 0),
],
)
};
let scan = PatternNode::Op {
op: OpTag::Scan,
attrs: OpAttrs {
scan_n_carries: Some(1),
scan_n_xs: Some(1),
scan_bound: Some(3),
scan_emit: Some(0),
scan_early_exit: None,
..OpAttrs::default()
},
operands: vec![
bind(0), bind(1), bind(2), sum(), op_node(OpTag::Mul, OpAttrs::default(), vec![bind(2), sum()]), ],
};
op_node(
OpTag::View,
OpAttrs {
view_slot: Some(slot),
..OpAttrs::default()
},
vec![scan],
)
}
fn reference_view(g: &mut Graph, slot: u32) -> NodeId {
let two = Shape::from_dims(&[2]);
let init = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: two.clone(),
dtype: DType::F32,
});
let xs0 = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let konst = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: two.clone(),
dtype: DType::F32,
});
let carry = g.push(Node {
op: Op::ScanPlaceholder {
role: ScanRole::Carry,
index: 0,
},
inputs: vec![],
shape: two.clone(),
dtype: DType::F32,
});
let elem = g.push(Node {
op: Op::ScanPlaceholder {
role: ScanRole::Elem,
index: 0,
},
inputs: vec![],
shape: two.clone(),
dtype: DType::F32,
});
let sum = g.push(Node {
op: Op::Add,
inputs: vec![carry, elem],
shape: two.clone(),
dtype: DType::F32,
});
let body_y = g.push(Node {
op: Op::Mul,
inputs: vec![konst, sum],
shape: two.clone(),
dtype: DType::F32,
});
let ys = Shape::from_dims(&[3, 2]);
let scan = g.push(Node {
op: Op::Scan {
n_carries: 1,
n_xs: 1,
bound: 3,
emit: ScanEmit::All,
early_exit: None,
},
inputs: vec![init, xs0, konst, sum, body_y],
shape: ys.clone(),
dtype: DType::F32,
});
let specs = vec![
OutputViewSpec::contiguous(DType::F32, ys),
OutputViewSpec::contiguous(DType::F32, two),
];
let (_bytes, views) = compose_bundle(&specs).expect("compose_bundle");
g.set_output_views(scan, Arc::from(views.into_boxed_slice()))
.expect("set_output_views");
let (sh, dt) = {
let v = g.output_views(scan).expect("bundle set");
let s = &v[slot as usize];
(s.shape.clone(), s.dtype)
};
g.push(Node {
op: Op::View { slot },
inputs: vec![scan],
shape: sh,
dtype: dt,
})
}
fn roundtrip_view(slot: u32, want_shape: &[usize]) {
let mut gr = Graph::new();
let ref_view = reference_view(&mut gr, slot);
let want = base_map_hash(&gr, ref_view);
let mut ge = Graph::new();
let two = Shape::from_dims(&[2]);
let init = ge.push(Node {
op: Op::Const,
inputs: vec![],
shape: two.clone(),
dtype: DType::F32,
});
let xs0 = ge.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let konst = ge.push(Node {
op: Op::Const,
inputs: vec![],
shape: two,
dtype: DType::F32,
});
let root = emit_region(&mut ge, &view_scan_recipe(slot), &[init, xs0, konst], &[]);
assert!(
matches!(ge.node(root).op, Op::View { slot: s } if s == slot),
"recipe root re-emits Op::View{{slot={slot}}}"
);
assert_eq!(
ge.node(root).shape.dims(),
want_shape,
"view slot {slot} shape comes from the re-attached output_views bundle"
);
let producer = ge.node(root).inputs[0];
assert!(
matches!(ge.node(producer).op, Op::Scan { .. }),
"producer is the scan"
);
assert_eq!(
ge.output_views(producer).map(|v| v.len()),
Some(2),
"emit re-attaches the 2-slot bundle so slot>=1 views resolve"
);
let got = base_map_hash(&ge, root);
assert_eq!(
got, want,
"view+scan recipe re-emit base map == imperative scan+view base map"
);
}
#[test]
fn scan_view_slot0_reemits_to_the_same_base_map() {
roundtrip_view(0, &[3, 2]);
}
#[test]
fn scan_view_slot1_reemits_to_the_same_base_map() {
roundtrip_view(1, &[2]);
}
#[test]
fn view_scan_recipe_validates_as_representable() {
assert_eq!(view_scan_recipe(1).bind_indices(), vec![0, 1, 2]);
assert!(validate_recipe(&view_scan_recipe(1)).is_ok());
assert!(matches!(
tag_to_op(
OpTag::View,
&OpAttrs {
view_slot: Some(1),
..OpAttrs::default()
}
),
Some(Op::View { slot: 1 }),
));
assert!(tag_to_op(OpTag::View, &OpAttrs::default()).is_none());
}
#[test]
fn scan_placeholder_recipe_emits_its_declared_shape() {
use fuel_kernel_seam_types::shape_expr::{Dim, ShapeExpr};
let dims2 = || {
ShapeExpr::Dims(vec![
Dim::Extent {
operand: 0,
axis: 0,
},
Dim::Extent {
operand: 0,
axis: 1,
},
])
};
let ph = |index: u32| PatternNode::Op {
op: OpTag::ScanPlaceholder,
attrs: OpAttrs {
scan_role: Some(SCAN_ROLE_ELEM),
scan_index: Some(index),
target_shape_rel: Some(dims2()),
..OpAttrs::default()
},
operands: vec![],
};
let recipe = op_node(OpTag::Mul, OpAttrs::default(), vec![ph(0), bind(0)]);
assert!(validate_recipe(&recipe).is_ok());
let mut g = Graph::new();
let bind0 = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let root = emit_region(&mut g, &recipe, &[bind0], &[]);
let ph_id = g.node(root).inputs[0];
assert!(matches!(
g.node(ph_id).op,
Op::ScanPlaceholder {
role: ScanRole::Elem,
index: 0
}
));
assert_eq!(
g.node(ph_id).shape.dims(),
&[2, 3],
"placeholder emits its declared target_shape_rel shape, not rank-0",
);
assert_eq!(
g.node(ph_id).dtype,
DType::F32,
"placeholder dtype = bind 0's"
);
assert_eq!(
g.node(root).shape.dims(),
&[2, 3],
"Mul(placeholder,..) inherits the declared shape"
);
}
}
mod reduce_max_to_backward_recipe {
use super::super::*;
use super::frozen_legacy_reduce_max_to_backward_decompose;
use crate::registry::{FusedOps, reduce_max_to_backward};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn broadcast(input: &[f64], in_shape: &[usize], target: &[usize]) -> Vec<f64> {
let rank = target.len();
let pad = rank - in_shape.len();
let mut real_strides = vec![0isize; in_shape.len()];
let mut s = 1isize;
for i in (0..in_shape.len()).rev() {
real_strides[i] = s;
s *= in_shape[i] as isize;
}
let mut in_strides = vec![0isize; rank];
for (i, stride) in in_strides.iter_mut().enumerate() {
if i >= pad {
let id = i - pad;
*stride = if in_shape[id] == 1 {
0
} else {
real_strides[id]
};
}
}
let out_n: usize = target.iter().product();
let mut out = Vec::with_capacity(out_n);
let mut idx = vec![0usize; rank];
for _ in 0..out_n {
let fi: isize = (0..rank).map(|i| idx[i] as isize * in_strides[i]).sum();
out.push(input[fi as usize]);
for i in (0..rank).rev() {
idx[i] += 1;
if idx[i] < target[i] {
break;
}
idx[i] = 0;
}
}
out
}
fn eval_rmb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
if let Some(v) = leaves.get(&id) {
return v.clone();
}
let node = g.node(id);
match &node.op {
Op::Mul => {
let a = eval_rmb(g, node.inputs[0], leaves);
let b = eval_rmb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x * y).collect()
}
Op::Div => {
let a = eval_rmb(g, node.inputs[0], leaves);
let b = eval_rmb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x / y).collect()
}
Op::MulScalar(m) => eval_rmb(g, node.inputs[0], leaves)
.iter()
.map(|v| v * m)
.collect(),
Op::Equal => {
let a = eval_rmb(g, node.inputs[0], leaves);
let b = eval_rmb(g, node.inputs[1], leaves);
a.iter()
.zip(&b)
.map(|(x, y)| if x == y { 1.0 } else { 0.0 })
.collect()
}
Op::MaskedFill { value } => {
let input = eval_rmb(g, node.inputs[0], leaves);
let mask = eval_rmb(g, node.inputs[1], leaves);
let fill = value.to_f64();
input
.iter()
.zip(&mask)
.map(|(x, m)| if *m != 0.0 { fill } else { *x })
.collect()
}
Op::ReduceMaxTo(_) => {
let input = eval_rmb(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().cloned().fold(f64::NEG_INFINITY, f64::max))
.collect()
}
Op::ReduceSumTo(_) => {
let input = eval_rmb(g, node.inputs[0], leaves);
let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
input
.chunks(last)
.map(|row| row.iter().sum::<f64>())
.collect()
}
Op::BroadcastTo(target) => {
let input = eval_rmb(g, node.inputs[0], leaves);
let in_shape: Vec<usize> = g.node(node.inputs[0]).shape.dims().to_vec();
broadcast(&input, &in_shape, target.dims())
}
other => panic!("eval_rmb: unhandled op {other:?}"),
}
}
fn rmb_fused_node(
g: &mut Graph,
x_dims: &[usize],
up_dims: &[usize],
) -> (NodeId, NodeId, NodeId) {
let xsh = Shape::from_dims(x_dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: xsh.clone(),
dtype: DType::F32,
});
let up = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(up_dims),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::REDUCE_MAX_TO_BACKWARD,
FusedOpParams::ReduceMaxToBackward,
),
inputs: vec![x, up],
shape: xsh,
dtype: DType::F32,
});
(x, up, fused)
}
#[test]
fn reduce_max_to_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for (x_dims, up_dims) in [
(vec![2usize, 4], vec![2usize, 1]),
(vec![3, 5, 7], vec![3, 5, 1]),
] {
let mut g = Graph::new();
let (x, up, fused) = rmb_fused_node(&mut g, &x_dims, &up_dims);
let xsh = Shape::from_dims(&x_dims);
let new_root = reduce_max_to_backward::decompose(
&mut g,
fused,
&FusedOpParams::ReduceMaxToBackward,
);
assert_ne!(new_root, fused, "recipe decompose must fire at {x_dims:?}");
assert_eq!(
g.node(new_root).shape,
xsh,
"reduce_max backward is x-shaped"
);
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_reduce_max_to_backward_decompose(
&mut g,
fused,
&FusedOpParams::ReduceMaxToBackward,
);
let xn: usize = x_dims.iter().product();
let upn: usize = up_dims.iter().product();
let x_data: Vec<f64> = (0..xn)
.map(|i| ((i as f64) * 0.37).sin() * 2.0 + 0.3)
.collect();
let up_data: Vec<f64> = (0..upn)
.map(|i| ((i as f64) * 0.53).cos() * 1.1 - 0.2)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, x_data);
leaves.insert(up, up_data);
let got = eval_rmb(&g, new_root, &leaves);
let want = eval_rmb(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"reduce_max_to_backward[{i}] at {x_dims:?}: recipe={a} vs legacy={b}",
);
}
}
}
#[test]
fn reduce_max_to_backward_recipe_shares_mask_and_has_no_reshape() {
for (x_dims, up_dims) in [
(vec![2usize, 4], vec![2usize, 1]),
(vec![3, 5, 7], vec![3, 5, 1]),
] {
let mut g = Graph::new();
let (_x, _up, fused) = rmb_fused_node(&mut g, &x_dims, &up_dims);
let root = reduce_max_to_backward::decompose(
&mut g,
fused,
&FusedOpParams::ReduceMaxToBackward,
);
assert_ne!(root, fused, "recipe decompose fires at {x_dims:?}");
assert!(
matches!(g.node(root).op, Op::Mul),
"root is Mul(mask_f, share_b)"
);
let reachable = crate::topo_order_multi(&g, &[root]);
let masked = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::MaskedFill { .. }))
.count();
assert_eq!(
masked, 1,
"mask_f identity-shared across its two uses at {x_dims:?} (ONE MaskedFill)",
);
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
reshapes, 0,
"direct structural mirror — no D3/D4 Reshape at {x_dims:?}"
);
let fill_dtype = reachable
.iter()
.find_map(|&n| match &g.node(n).op {
Op::MaskedFill { value } => Some(value.dtype()),
_ => None,
})
.expect("a MaskedFill is present");
assert_eq!(
fill_dtype,
DType::F32,
"the A2 carrier re-resolved the fill Scalar to x's dtype at {x_dims:?}",
);
}
}
#[test]
fn reduce_max_to_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, _up, fused) = rmb_fused_node(&mut g, &[2, 4], &[2, 1]);
let before = g.len();
let out = reduce_max_to_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
mod powi_backward_recipe {
use super::super::*;
use super::frozen_legacy_powi_backward_decompose;
use crate::registry::{FusedOps, powi_backward};
use fuel_ir::{DType, Shape};
use std::collections::HashMap;
fn eval_pb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
if let Some(v) = leaves.get(&id) {
return v.clone();
}
let node = g.node(id);
match &node.op {
Op::Mul => {
let a = eval_pb(g, node.inputs[0], leaves);
let b = eval_pb(g, node.inputs[1], leaves);
a.iter().zip(&b).map(|(x, y)| x * y).collect()
}
Op::MulScalar(m) => eval_pb(g, node.inputs[0], leaves)
.iter()
.map(|v| v * m)
.collect(),
Op::PowI(n) => eval_pb(g, node.inputs[0], leaves)
.iter()
.map(|v| v.powi(*n))
.collect(),
other => panic!("eval_pb: unhandled op {other:?}"),
}
}
fn pb_fused_node(g: &mut Graph, dims: &[usize], exp: i32) -> (NodeId, NodeId, NodeId) {
let sh = Shape::from_dims(dims);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let up = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: sh.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(FusedOps::POWI_BACKWARD, FusedOpParams::PowIBackward { exp }),
inputs: vec![x, up],
shape: sh,
dtype: DType::F32,
});
(x, up, fused)
}
#[test]
fn powi_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
for dims in [vec![2usize, 4], vec![3, 5, 7]] {
for exp in [3i32, 2, 5, 1, 0, -2] {
let mut g = Graph::new();
let (x, up, fused) = pb_fused_node(&mut g, &dims, exp);
let sh = Shape::from_dims(&dims);
let new_root = powi_backward::decompose(
&mut g,
fused,
&FusedOpParams::PowIBackward { exp },
);
assert_ne!(
new_root, fused,
"recipe decompose must fire at {dims:?} exp={exp}"
);
assert_eq!(
g.node(new_root).shape,
sh,
"powi backward is shape-preserving"
);
assert_eq!(g.node(new_root).dtype, DType::F32);
let legacy_root = frozen_legacy_powi_backward_decompose(
&mut g,
fused,
&FusedOpParams::PowIBackward { exp },
);
let n: usize = dims.iter().product();
let x_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.41).sin() * 0.5 + 1.5)
.collect();
let up_data: Vec<f64> = (0..n)
.map(|i| ((i as f64) * 0.61).cos() * 1.3 - 0.15)
.collect();
let mut leaves = HashMap::new();
leaves.insert(x, x_data);
leaves.insert(up, up_data);
let got = eval_pb(&g, new_root, &leaves);
let want = eval_pb(&g, legacy_root, &leaves);
assert_eq!(got.len(), want.len());
for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"powi_backward[{i}] at {dims:?} exp={exp}: recipe={a} vs legacy={b}",
);
}
}
}
}
#[test]
fn powi_backward_recipe_is_a_direct_mirror_with_carrier_exponent() {
for exp in [3i32, 0, -2] {
let mut g = Graph::new();
let (_x, _up, fused) = pb_fused_node(&mut g, &[2, 4], exp);
let root =
powi_backward::decompose(&mut g, fused, &FusedOpParams::PowIBackward { exp });
assert_ne!(root, fused, "recipe decompose fires at exp={exp}");
assert!(
matches!(g.node(root).op, Op::Mul),
"root is Mul(scaled, up)"
);
let reachable = crate::topo_order_multi(&g, &[root]);
let powis: Vec<i32> = reachable
.iter()
.filter_map(|&n| match g.node(n).op {
Op::PowI(k) => Some(k),
_ => None,
})
.collect();
assert_eq!(
powis,
vec![exp - 1],
"the A3 carrier resolved Op::PowI(exp-1) at exp={exp}"
);
let mul_scalars: Vec<f64> = reachable
.iter()
.filter_map(|&n| match g.node(n).op {
Op::MulScalar(v) => Some(v),
_ => None,
})
.collect();
assert_eq!(
mul_scalars,
vec![exp as f64],
"baked MulScalar(exp) at exp={exp}"
);
let reshapes = reachable
.iter()
.filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
.count();
assert_eq!(
reshapes, 0,
"direct structural mirror — no D3/D4 Reshape at exp={exp}"
);
}
}
#[test]
fn powi_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let (_x, _up, fused) = pb_fused_node(&mut g, &[2, 4], 3);
let before = g.len();
let out = powi_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}
}