use std::collections::{HashMap, HashSet};
use crate::fkc::ImportWarning;
use crate::fkc::cost_expr::{CostNode, parse_expr as parse_cost_expr};
use crate::fkc::error::FkcError;
use crate::fkc::schema::TensorDesc;
use fuel_ir::{DType, Shape};
pub type ProbeCombo = Vec<(String, fuel_ir::Shape, fuel_ir::DType)>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RankSpec {
Exact(usize),
Any,
Range { lo: usize, hi: Option<usize> },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AxisIndex {
FromStart(usize),
FromEnd(usize),
}
#[derive(Debug, Clone, PartialEq)]
pub enum ShapeAtom {
SameAs(String),
SameRank(String),
Rank(RankSpec),
BroadcastTo(String),
LastDimEq(String),
DimEq { axis: AxisIndex, expr: CostNode },
Divisible { lhs: CostNode, rhs: CostNode },
CapacityGe { axis: AxisIndex, sym: String },
}
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ShapeConstraint {
pub atoms: Vec<ShapeAtom>,
pub freetext: Vec<String>,
}
fn parse_axis(s: &str) -> Option<AxisIndex> {
let s = s.trim();
if let Some(n) = s.strip_prefix('-') {
n.trim().parse::<usize>().ok().map(AxisIndex::FromEnd)
} else {
s.parse::<usize>().ok().map(AxisIndex::FromStart)
}
}
pub fn parse_rank_spec(s: &str) -> Option<RankSpec> {
let s = s.trim();
if s == "any" {
return Some(RankSpec::Any);
}
if let Ok(n) = s.parse::<usize>() {
return Some(RankSpec::Exact(n));
}
if let Some((lo, hi)) = s.split_once("..=") {
return Some(RankSpec::Range {
lo: lo.trim().parse().ok()?,
hi: Some(hi.trim().parse().ok()?),
});
}
if let Some(lo) = s.strip_suffix("..") {
return Some(RankSpec::Range {
lo: lo.trim().parse().ok()?,
hi: None,
});
}
None
}
fn split_two_args(inner: &str) -> Option<(&str, &str)> {
let mut depth = 0i32;
for (i, c) in inner.char_indices() {
match c {
'(' | '[' => depth += 1,
')' | ']' => depth -= 1,
',' if depth == 0 => return Some((&inner[..i], &inner[i + 1..])),
_ => {}
}
}
None
}
pub fn parse_shape_constraint(
raw: &str,
section: &str,
operand: &str,
) -> Result<ShapeConstraint, FkcError> {
let mut out = ShapeConstraint::default();
for seg_raw in raw.split(';') {
let seg = seg_raw.trim();
if seg.is_empty() {
continue;
}
let unparse = || FkcError::UnparseableShapeConstraint {
section: section.into(),
operand: operand.into(),
raw: seg.to_string(),
};
if let Some(r) = seg.strip_prefix("same_as=") {
out.atoms.push(ShapeAtom::SameAs(r.trim().into()));
continue;
}
if let Some(r) = seg.strip_prefix("same_rank=") {
out.atoms.push(ShapeAtom::SameRank(r.trim().into()));
continue;
}
if let Some(r) = seg.strip_prefix("broadcast_to=") {
out.atoms.push(ShapeAtom::BroadcastTo(r.trim().into()));
continue;
}
if let Some(r) = seg.strip_prefix("last_dim_eq=") {
out.atoms.push(ShapeAtom::LastDimEq(r.trim().into()));
continue;
}
if let Some(r) = seg.strip_prefix("rank=") {
out.atoms
.push(ShapeAtom::Rank(parse_rank_spec(r).ok_or_else(unparse)?));
continue;
}
if seg.starts_with("divisible(") {
let inner = seg
.strip_prefix("divisible(")
.and_then(|s| s.strip_suffix(')'))
.ok_or_else(unparse)?;
let (a, b) = split_two_args(inner).ok_or_else(unparse)?;
let lhs = parse_cost_expr(a.trim()).map_err(|_| unparse())?;
let rhs = parse_cost_expr(b.trim()).map_err(|_| unparse())?;
out.atoms.push(ShapeAtom::Divisible { lhs, rhs });
continue;
}
if seg.starts_with("capacity_ge(") {
let inner = seg
.strip_prefix("capacity_ge(")
.and_then(|s| s.strip_suffix(')'))
.ok_or_else(unparse)?;
let (a, b) = split_two_args(inner).ok_or_else(unparse)?;
let axis = a
.trim()
.strip_prefix("dim[")
.and_then(|s| s.strip_suffix(']'))
.and_then(parse_axis)
.ok_or_else(unparse)?;
out.atoms.push(ShapeAtom::CapacityGe {
axis,
sym: b.trim().to_string(),
});
continue;
}
if seg.starts_with("dim[")
&& let Some(close) = seg.find(']')
{
let idx = &seg["dim[".len()..close];
let after = seg[close + 1..].trim_start();
match (parse_axis(idx), after.strip_prefix('=')) {
(Some(axis), Some(rhs)) if !rhs.starts_with('=') => {
let rhs = rhs.trim();
if rhs.is_empty() {
return Err(unparse());
}
out.atoms.push(ShapeAtom::DimEq {
axis,
expr: parse_cost_expr(rhs).map_err(|_| unparse())?,
});
continue;
}
_ => {
out.freetext.push(seg.to_string());
continue;
}
}
}
out.freetext.push(seg.to_string()); }
Ok(out)
}
#[derive(Clone, Copy)]
struct SeedProfile {
base: i64,
odd_last: bool,
}
const PROFILES: [SeedProfile; 3] = [
SeedProfile {
base: 2,
odd_last: false,
}, SeedProfile {
base: 4,
odd_last: true,
}, SeedProfile {
base: 8,
odd_last: false,
}, ];
fn resolve_rank_spec_field(v: Option<&serde_yaml_ng::Value>) -> Option<RankSpec> {
match v {
Some(serde_yaml_ng::Value::Number(n)) => n.as_u64().map(|u| RankSpec::Exact(u as usize)),
Some(serde_yaml_ng::Value::String(s)) => parse_rank_spec(s),
_ => None,
}
}
fn rank_for_probe(spec: Option<RankSpec>) -> usize {
match spec {
Some(RankSpec::Exact(n)) => n,
Some(RankSpec::Range { lo, .. }) => lo,
Some(RankSpec::Any) | None => 4, }
}
fn seed_axis(profile: SeedProfile, axis: usize, rank: usize) -> i64 {
if profile.odd_last && rank > 0 && axis == rank - 1 {
3
} else {
profile.base
}
}
fn first_probe_dtype(d: &TensorDesc) -> DType {
if let Some(tok) = d.dtypes.first()
&& let Ok(dt) = crate::fkc::lower::lower_dtype(tok, "", "")
{
return dt;
}
match d.dtype_class.as_deref() {
Some(class) => crate::fkc::lower::expand_dtype_class(class, &[], "", "")
.ok()
.and_then(|v| v.first().copied())
.unwrap_or(DType::F32),
None => DType::F32,
}
}
struct Solve {
dims: HashMap<String, Vec<i64>>,
sym: HashMap<String, i64>,
base: i64,
}
fn as_dim_ref(node: &CostNode) -> Option<(Option<String>, AxisIndex)> {
if let CostNode::Index { base, index } = node {
let axis = match &**index {
CostNode::Lit(v) => AxisIndex::FromStart(*v as usize),
CostNode::Neg(inner) => {
if let CostNode::Lit(v) = &**inner {
AxisIndex::FromEnd(*v as usize)
} else {
return None;
}
}
_ => return None,
};
if let CostNode::Sym(s) = &**base {
return Some(if let Some(r) = s.strip_suffix(".dim") {
(Some(r.to_string()), axis)
} else if s == "dim" {
(None, axis)
} else {
return None;
});
}
}
None
}
fn axis_to_index(axis: AxisIndex, rank: usize) -> Option<usize> {
match axis {
AxisIndex::FromStart(i) => Some(i),
AxisIndex::FromEnd(n) => rank.checked_sub(n),
}
}
fn eval_dim_expr(
node: &CostNode,
s: &mut Solve,
ranks: &HashMap<String, usize>,
self_role: &str,
) -> Option<i64> {
use crate::fkc::cost_expr::BinOp::*;
match node {
CostNode::Lit(v) => Some(*v as i64),
CostNode::Neg(i) => eval_dim_expr(i, s, ranks, self_role)?.checked_neg(),
CostNode::Bin { op, lhs, rhs } => {
let (l, r) = (
eval_dim_expr(lhs, s, ranks, self_role)?,
eval_dim_expr(rhs, s, ranks, self_role)?,
);
Some(match op {
Add => l.checked_add(r)?,
Sub => l.checked_sub(r)?,
Mul => l.checked_mul(r)?,
Div => l.checked_div(r)?,
Rem => l.checked_rem(r)?,
})
}
CostNode::Index { .. } => {
let (role, axis) = as_dim_ref(node)?;
let rrole = role.as_deref().unwrap_or(self_role);
let idx = axis_to_index(axis, *ranks.get(rrole)?)?;
s.dims.get(rrole)?.get(idx).copied()
}
CostNode::Sym(name) => Some(*s.sym.entry(name.clone()).or_insert(s.base)), CostNode::Call { .. } => None,
}
}
fn warn(section: &str, message: String) -> ImportWarning {
ImportWarning {
section: section.into(),
message,
}
}
fn dep_sources(atoms: &[ShapeAtom], input_roles: &HashSet<String>) -> Vec<String> {
fn collect(n: &CostNode, roles: &HashSet<String>, out: &mut Vec<String>) {
if let Some((Some(r), _)) = as_dim_ref(n)
&& roles.contains(&r)
{
out.push(r);
}
match n {
CostNode::Bin { lhs, rhs, .. } => {
collect(lhs, roles, out);
collect(rhs, roles, out);
}
CostNode::Neg(i) => collect(i, roles, out),
CostNode::Index { base, index } => {
collect(base, roles, out);
collect(index, roles, out);
}
CostNode::Call { args, .. } => {
for a in args {
collect(a, roles, out);
}
}
_ => {}
}
}
let mut deps = Vec::new();
for a in atoms {
match a {
ShapeAtom::SameAs(r) | ShapeAtom::BroadcastTo(r) | ShapeAtom::LastDimEq(r)
if input_roles.contains(r) =>
{
deps.push(r.clone())
}
ShapeAtom::DimEq { expr, .. } => collect(expr, input_roles, &mut deps),
ShapeAtom::Divisible { lhs, rhs } => {
collect(lhs, input_roles, &mut deps);
collect(rhs, input_roles, &mut deps);
}
_ => {}
}
}
deps.sort();
deps.dedup();
deps
}
fn topo_order(
order_in: &[String],
edges: &HashMap<String, Vec<String>>,
section: &str,
w: &mut Vec<ImportWarning>,
) -> Vec<String> {
let set: HashSet<&String> = order_in.iter().collect();
let mut indeg: HashMap<&String, usize> = order_in.iter().map(|r| (r, 0usize)).collect();
for (r, deps) in edges {
if set.contains(r) {
for d in deps {
if set.contains(d) {
*indeg.get_mut(r).unwrap() += 1;
}
}
}
}
let mut queue: Vec<&String> = order_in.iter().filter(|r| indeg[r] == 0).collect();
let mut out = Vec::new();
let mut i = 0;
while i < queue.len() {
let cur = queue[i];
i += 1;
out.push(cur.clone());
for r in order_in {
if let Some(deps) = edges.get(r)
&& deps.contains(cur)
{
let e = indeg.get_mut(r).unwrap();
*e = e.saturating_sub(1);
if *e == 0 && !out.contains(r) && !queue.contains(&r) {
queue.push(r);
}
}
}
}
if out.len() < order_in.len() {
let residual: Vec<&String> = order_in.iter().filter(|r| !out.contains(r)).collect();
w.push(warn(
section,
format!("shape_constraint dependency cycle among {residual:?}; using seed shapes"),
));
for r in order_in {
if !out.contains(r) {
out.push(r.clone());
}
}
}
out
}
fn set_axis(s: &mut Solve, role: &str, axis: AxisIndex, rank: usize, v: i64) {
if let Some(idx) = axis_to_index(axis, rank)
&& let Some(d) = s.dims.get_mut(role)
&& idx < d.len()
{
d[idx] = v.max(1);
}
}
fn apply_atom(
atom: &ShapeAtom,
self_role: &str,
s: &mut Solve,
ranks: &HashMap<String, usize>,
w: &mut Vec<ImportWarning>,
section: &str,
) {
let self_rank = *ranks.get(self_role).unwrap_or(&0);
match atom {
ShapeAtom::Rank(_) | ShapeAtom::SameRank(_) | ShapeAtom::CapacityGe { .. } => {} ShapeAtom::SameAs(src) | ShapeAtom::BroadcastTo(src) => match s.dims.get(src).cloned() {
Some(src_dims) => {
let n = self_rank.min(src_dims.len());
if let Some(d) = s.dims.get_mut(self_role) {
d[..n].copy_from_slice(&src_dims[..n]);
}
}
None => w.push(warn(
section,
format!("operand `{self_role}` references unknown role `{src}`; using seed shape"),
)),
},
ShapeAtom::LastDimEq(src) => match s.dims.get(src).and_then(|d| d.last().copied()) {
Some(sr) => set_axis(s, self_role, AxisIndex::FromEnd(1), self_rank, sr),
None => w.push(warn(
section,
format!(
"operand `{self_role}` last_dim_eq references unknown role `{src}`; using seed"
),
)),
},
ShapeAtom::DimEq { axis, expr } => match eval_dim_expr(expr, s, ranks, self_role) {
Some(v) => set_axis(s, self_role, *axis, self_rank, v),
None => w.push(warn(
section,
format!("operand `{self_role}` dim rule unresolved; using seed"),
)),
},
ShapeAtom::Divisible { lhs, rhs } => {
if let (Some((role, axis)), Some(v)) =
(as_dim_ref(lhs), eval_dim_expr(rhs, s, ranks, self_role))
{
if v > 0 {
let target = role.as_deref().unwrap_or(self_role).to_string();
let trank = *ranks.get(&target).unwrap_or(&0);
if let Some(idx) = axis_to_index(axis, trank)
&& let Some(cur) = s.dims.get(&target).and_then(|d| d.get(idx).copied())
{
if let Some(rounded) = cur
.checked_add(v - 1)
.map(|x| x / v)
.and_then(|q| q.checked_mul(v))
{
let final_val = if target != self_role && rounded == v {
rounded.checked_mul(2).unwrap_or(rounded)
} else {
rounded
};
set_axis(s, &target, axis, trank, final_val);
}
}
}
} else if let CostNode::Sym(k) = lhs
&& let Some(v) = eval_dim_expr(rhs, s, ranks, self_role)
&& v > 0
{
let e = s.sym.entry(k.clone()).or_insert(s.base);
if let Some(rounded) = e
.checked_add(v - 1)
.map(|x| x / v)
.and_then(|q| q.checked_mul(v))
{
*e = rounded;
}
}
}
}
}
pub fn solve_probe_shapes(
inputs: &[TensorDesc],
section: &str,
warnings: &mut Vec<ImportWarning>,
) -> Result<Vec<ProbeCombo>, FkcError> {
let mut parsed = Vec::with_capacity(inputs.len());
let mut roles = Vec::with_capacity(inputs.len());
let mut keys = Vec::with_capacity(inputs.len());
let mut ranks: HashMap<String, usize> = HashMap::with_capacity(inputs.len());
let mut seen_names: HashSet<String> = HashSet::with_capacity(inputs.len());
let mut warned_dup_names: HashSet<String> = HashSet::new();
for (i, d) in inputs.iter().enumerate() {
let operand = d.name.as_deref().unwrap_or("<unnamed>");
let sc = match &d.shape_constraint {
Some(raw) => parse_shape_constraint(raw, section, operand)?,
None => ShapeConstraint::default(),
};
for seg in &sc.freetext {
warnings.push(warn(
section,
format!("operand `{operand}` shape_constraint free text: {seg}"),
));
}
let role = d.name.clone().unwrap_or_default();
let key = d
.name
.clone()
.filter(|n| !n.is_empty())
.unwrap_or_else(|| format!("#unnamed{i}"));
if let Some(name) = d.name.as_deref().filter(|n| !n.is_empty())
&& !seen_names.insert(name.to_string())
&& warned_dup_names.insert(name.to_string())
{
warnings.push(warn(section, format!(
"operand name `{name}` is declared more than once; references to it are ambiguous and only one operand's shape constraints are applied"
)));
}
let rank_spec = resolve_rank_spec_field(d.rank.as_ref());
if d.rank.is_some() && rank_spec.is_none() {
warnings.push(warn(
section,
format!("operand `{operand}` rank field is malformed; defaulting to rank 4"),
));
}
let rank = rank_for_probe(rank_spec);
ranks.insert(key.clone(), rank);
roles.push(role);
keys.push(key);
parsed.push(sc);
}
let key_set: HashSet<String> = keys.iter().cloned().collect();
let mut edges: HashMap<String, Vec<String>> = HashMap::with_capacity(keys.len());
for (i, key) in keys.iter().enumerate() {
edges.insert(key.clone(), dep_sources(&parsed[i].atoms, &key_set));
}
let topo = topo_order(&keys, &edges, section, warnings);
let key_to_idx: HashMap<&String, usize> =
keys.iter().enumerate().map(|(i, k)| (k, i)).collect();
let mut combos = Vec::with_capacity(PROFILES.len());
for profile in PROFILES {
let mut s = Solve {
dims: HashMap::with_capacity(inputs.len()),
sym: HashMap::new(),
base: profile.base,
};
for key in &keys {
let rank = *ranks.get(key).unwrap_or(&0);
let dims: Vec<i64> = (0..rank).map(|a| seed_axis(profile, a, rank)).collect();
s.dims.insert(key.clone(), dims);
}
for key in &topo {
if let Some(&i) = key_to_idx.get(key) {
for atom in &parsed[i].atoms {
apply_atom(atom, key, &mut s, &ranks, warnings, section);
}
}
}
let mut combo: ProbeCombo = Vec::with_capacity(inputs.len());
for ((d, role), key) in inputs.iter().zip(roles.iter()).zip(keys.iter()) {
let dims: Vec<usize> = s
.dims
.get(key)
.map(|v| v.iter().map(|&x| x.max(1) as usize).collect())
.unwrap_or_default();
combo.push((role.clone(), Shape::from_dims(&dims), first_probe_dtype(d)));
}
combos.push(combo);
}
Ok(combos)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fkc::cost_expr::CostNode;
fn desc(name: &str, dtypes: &[&str], rank: Option<u64>) -> crate::fkc::schema::TensorDesc {
crate::fkc::schema::TensorDesc {
name: Some(name.into()),
optional: false,
dtypes: dtypes.iter().map(|s| s.to_string()).collect(),
dtype_class: None,
layout: None,
rank: rank.map(|r| serde_yaml_ng::Value::Number(r.into())),
shape_constraint: None,
fdx: None,
device: None,
substrate: None,
}
}
#[test]
fn seed_unconstrained_operands_over_three_profiles() {
use fuel_ir::Shape;
let inputs = vec![
desc("lhs", &["F32"], Some(2)),
desc("rhs", &["F32"], Some(2)),
];
let mut w = Vec::new();
let combos = solve_probe_shapes(&inputs, "s", &mut w).unwrap();
assert_eq!(combos.len(), 3);
assert_eq!(combos[0][0].1, Shape::from_dims(&[2, 2])); assert_eq!(combos[1][0].1, Shape::from_dims(&[4, 3])); assert_eq!(combos[2][0].1, Shape::from_dims(&[8, 8])); assert!(w.is_empty());
}
#[test]
fn rank_any_defaults_to_4_and_open_range_uses_lo() {
let any = desc("a", &["F32"], None); assert_eq!(
solve_probe_shapes(&[any], "s", &mut Vec::new()).unwrap()[0][0]
.1
.rank(),
4
);
let mut open = desc("b", &["F32"], None);
open.rank = Some(serde_yaml_ng::Value::String("2..".into()));
assert_eq!(
solve_probe_shapes(&[open], "s", &mut Vec::new()).unwrap()[0][0]
.1
.rank(),
2
);
}
#[test]
fn parse_covers_vocab_freetext_and_rejects_malformed_vocab() {
assert_eq!(
parse_shape_constraint("same_as=k", "s", "v").unwrap().atoms,
vec![ShapeAtom::SameAs("k".into())]
);
assert_eq!(
parse_shape_constraint("same_rank=k", "s", "v")
.unwrap()
.atoms,
vec![ShapeAtom::SameRank("k".into())]
);
assert_eq!(
parse_shape_constraint("broadcast_to=x", "s", "v")
.unwrap()
.atoms,
vec![ShapeAtom::BroadcastTo("x".into())]
);
assert_eq!(
parse_shape_constraint("last_dim_eq=x", "s", "v")
.unwrap()
.atoms,
vec![ShapeAtom::LastDimEq("x".into())]
);
assert_eq!(
parse_shape_constraint("rank=4", "s", "x").unwrap().atoms,
vec![ShapeAtom::Rank(RankSpec::Exact(4))]
);
assert_eq!(
parse_shape_constraint("rank=2..=4", "s", "x")
.unwrap()
.atoms,
vec![ShapeAtom::Rank(RankSpec::Range { lo: 2, hi: Some(4) })]
);
assert_eq!(
parse_shape_constraint("rank=any", "s", "x").unwrap().atoms,
vec![ShapeAtom::Rank(RankSpec::Any)]
);
assert_eq!(
parse_shape_constraint("rank=2..", "s", "x").unwrap().atoms,
vec![ShapeAtom::Rank(RankSpec::Range { lo: 2, hi: None })]
);
let a = parse_shape_constraint("dim[-1]=k; same_rank=b", "linear", "a").unwrap();
assert_eq!(a.atoms.len(), 2);
match &a.atoms[0] {
ShapeAtom::DimEq { axis, expr } => {
assert_eq!(*axis, AxisIndex::FromEnd(1));
assert_eq!(*expr, CostNode::Sym("k".into()));
}
other => panic!("got {other:?}"),
}
assert_eq!(a.atoms[1], ShapeAtom::SameRank("b".into()));
assert!(matches!(
parse_shape_constraint("divisible(q.dim[2], k.dim[2])", "f", "k")
.unwrap()
.atoms[0],
ShapeAtom::Divisible { .. }
));
assert!(matches!(
parse_shape_constraint("capacity_ge(dim[0], seqlen)", "f", "kv")
.unwrap()
.atoms[0],
ShapeAtom::CapacityGe { .. }
));
let nested = parse_shape_constraint("divisible(foo(a, b), c)", "f", "k").unwrap();
assert_eq!(
nested.atoms,
vec![ShapeAtom::Divisible {
lhs: CostNode::Call {
name: "foo".into(),
args: vec![CostNode::Sym("a".into()), CostNode::Sym("b".into())],
},
rhs: CostNode::Sym("c".into()),
}]
);
let mixed = parse_shape_constraint(
"same_as=out; read-modify-written in place (this operand IS the output)",
"shape-ops",
"dst",
)
.unwrap();
assert_eq!(mixed.atoms, vec![ShapeAtom::SameAs("out".into())]);
assert_eq!(mixed.freetext.len(), 1);
let ft = parse_shape_constraint("byte length >= 4 (one u32)", "shape-ops", "seed").unwrap();
assert!(ft.atoms.is_empty());
assert_eq!(ft.freetext.len(), 1);
let sym_i = parse_shape_constraint("dim[i] == in_shape[i]", "shape-ops", "out").unwrap();
assert!(sym_i.atoms.is_empty());
assert_eq!(sym_i.freetext.len(), 1);
assert!(matches!(
parse_shape_constraint("rank=banana", "s", "x").unwrap_err(),
FkcError::UnparseableShapeConstraint { .. }
));
assert!(parse_shape_constraint("divisible(x.dim[0]", "s", "x").is_err()); assert!(parse_shape_constraint("dim[0]=", "s", "x").is_err()); }
#[test]
fn solve_same_as_copies_dims_and_divisible_rounds_up() {
use fuel_ir::Shape;
let mut k = desc("k", &["F32"], Some(3));
k.shape_constraint = Some("divisible(dim[0], 8)".into());
let mut v = desc("v", &["F32"], Some(3));
v.shape_constraint = Some("same_as=k".into());
let combos = solve_probe_shapes(&[k, v], "s", &mut Vec::new()).unwrap();
let a = &combos[0]; let ks = &a.iter().find(|(r, _, _)| r == "k").unwrap().1;
let vs = &a.iter().find(|(r, _, _)| r == "v").unwrap().1;
assert_eq!(ks, &Shape::from_dims(&[8, 2, 2]));
assert_eq!(vs, ks);
}
#[test]
fn bare_symbol_is_shared_mutably_across_operands() {
let mut a = desc("a", &["F32"], Some(2));
a.shape_constraint = Some("divisible(k, 8)".into()); let mut b = desc("b", &["F32"], Some(2));
b.shape_constraint = Some("dim[-1]=k".into()); let combos = solve_probe_shapes(&[a, b], "s", &mut Vec::new()).unwrap();
let a0 = &combos[0]; let bk = a0.iter().find(|(r, _, _)| r == "b").unwrap().1.dims()[1];
assert_eq!(
bk, 8,
"b's last axis must observe the SHARED, bumped k (=8), not a fresh k=2"
);
}
#[test]
fn dim_expr_overflow_degrades_without_panic() {
let mut a = desc("a", &["F32"], Some(2));
a.shape_constraint = Some("dim[0]=1000000*1000000*1000000*1000000".into());
let mut w = Vec::new();
let result = solve_probe_shapes(&[a], "s", &mut w);
assert!(
result.is_ok(),
"overflow must degrade gracefully, not panic or error"
);
assert!(
!w.is_empty(),
"the unresolved (overflowed) atom should surface an ImportWarning"
);
}
#[test]
fn pathological_int_min_div_and_neg_degrade_without_panic() {
use fuel_ir::Shape;
let mut a = desc("a", &["F32"], Some(1));
a.shape_constraint = Some("dim[0]=(-9223372036854775807-1)/-1".into()); let mut wa = Vec::new();
let ra = solve_probe_shapes(&[a], "s", &mut wa);
assert!(ra.is_ok(), "MIN/-1 must degrade, not panic");
assert!(
!wa.is_empty(),
"the unresolved (MIN/-1) atom should surface an ImportWarning"
);
assert_eq!(ra.unwrap()[0][0].1, Shape::from_dims(&[2]));
let mut b = desc("b", &["F32"], Some(1));
b.shape_constraint = Some("dim[0]=-(-9223372036854775807-1)".into()); let mut wb = Vec::new();
let rb = solve_probe_shapes(&[b], "s", &mut wb);
assert!(rb.is_ok(), "-(i64::MIN) must degrade, not panic");
assert!(
!wb.is_empty(),
"the unresolved (Neg MIN) atom should surface an ImportWarning"
);
assert_eq!(rb.unwrap()[0][0].1, Shape::from_dims(&[2]));
}
#[test]
fn two_unnamed_operands_are_not_aliased() {
let a = crate::fkc::schema::TensorDesc {
name: None,
optional: false,
dtypes: vec!["F32".into()],
dtype_class: None,
layout: None,
rank: Some(serde_yaml_ng::Value::Number(2u64.into())),
shape_constraint: None,
fdx: None,
device: None,
substrate: None,
};
let mut b = a.clone();
b.rank = Some(serde_yaml_ng::Value::Number(3u64.into()));
let combos = solve_probe_shapes(&[a, b], "s", &mut Vec::new()).unwrap();
assert_eq!(
combos[0][0].1.rank(),
2,
"first unnamed operand keeps rank 2"
);
assert_eq!(
combos[0][1].1.rank(),
3,
"second unnamed operand keeps rank 3, not aliased to 2"
);
}
#[test]
fn dependency_ordering_resolves_and_cycles_fall_back_with_warning() {
use fuel_ir::Shape;
let mut k = desc("k", &["F32"], Some(2));
k.shape_constraint = Some("divisible(dim[0], 8)".into());
let mut v = desc("v", &["F32"], Some(2));
v.shape_constraint = Some("same_as=k".into());
let combos = solve_probe_shapes(&[v, k], "s", &mut Vec::new()).unwrap();
assert_eq!(
combos[0].iter().find(|(r, _, _)| r == "v").unwrap().1,
Shape::from_dims(&[8, 2])
);
let mut ca = desc("a", &["F32"], Some(2));
ca.shape_constraint = Some("same_as=b".into());
let mut cb = desc("b", &["F32"], Some(2));
cb.shape_constraint = Some("same_as=a".into());
let mut w = Vec::new();
let combos = solve_probe_shapes(&[ca, cb], "cyc", &mut w).unwrap();
assert_eq!(combos.len(), 3);
assert_eq!(combos[0][0].1, Shape::from_dims(&[2, 2]));
assert!(
w.iter().any(|x| x.message.to_lowercase().contains("cycle")),
"warns: {w:?}"
);
}
#[test]
fn double_reference_to_one_predecessor_is_not_a_false_cycle() {
let mut x = desc("x", &["F32"], Some(2));
x.shape_constraint = Some("dim[0]=k.dim[0]+k.dim[1]".into());
let mut k = desc("k", &["F32"], Some(2));
k.shape_constraint = Some("divisible(dim[0], 8)".into()); let mut w = Vec::new();
let combos = solve_probe_shapes(&[x, k], "s", &mut w).unwrap(); assert!(
!w.iter().any(|m| m.message.to_lowercase().contains("cycle")),
"a double reference to one predecessor is NOT a cycle: {w:?}"
);
let a0 = &combos[0];
let x0 = a0.iter().find(|(r, _, _)| r == "x").unwrap().1.dims()[0];
assert_eq!(
x0, 10,
"x.dim0 = rounded k.dim0(8) + k.dim1(2); proves k-before-x ordering"
);
}
#[test]
fn duplicate_operand_name_emits_a_warning() {
let a = desc("dup", &["F32"], Some(2));
let b = desc("dup", &["F32"], Some(2));
let mut w = Vec::new();
let _ = solve_probe_shapes(&[a, b], "s", &mut w).unwrap(); assert!(
w.iter().any(|m| m.message.contains("dup")
&& m.message.to_lowercase().contains("more than once")),
"duplicate operand name must be surfaced as a warning: {w:?}"
);
}
#[test]
fn dtype_pick_and_soft_fallback_warnings() {
use fuel_ir::{DType, Shape};
assert_eq!(
solve_probe_shapes(
&[desc("a", &["BF16", "F16", "F32"], Some(1))],
"s",
&mut Vec::new()
)
.unwrap()[0][0]
.2,
DType::BF16
);
let mut c = desc("a", &[], Some(1));
c.dtype_class = Some("float".into());
assert_eq!(
solve_probe_shapes(&[c], "s", &mut Vec::new()).unwrap()[0][0].2,
DType::BF16
);
let mut r = desc("residual", &["F32"], Some(2));
r.shape_constraint = Some("same_as=out".into());
let mut w = Vec::new();
let combos = solve_probe_shapes(&[r], "norm-softmax", &mut w).unwrap();
assert_eq!(combos[0][0].1, Shape::from_dims(&[2, 2]));
assert!(w.iter().any(|x| x.message.contains("out")));
let mut f = desc("seed", &["U8"], Some(1));
f.shape_constraint = Some("byte length >= 4 (one u32)".into());
let mut w2 = Vec::new();
solve_probe_shapes(&[f], "shape-ops", &mut w2).unwrap();
assert!(!w2.is_empty());
}
#[test]
fn malformed_rank_field_warns_and_defaults_to_rank_4() {
let mut bad = desc("weird", &["F32"], None);
bad.rank = Some(serde_yaml_ng::Value::String("banana".into()));
let mut w = Vec::new();
let combos = solve_probe_shapes(&[bad], "s", &mut w).unwrap();
assert_eq!(
combos[0][0].1.rank(),
4,
"malformed rank field still defaults to rank 4"
);
assert!(
w.iter()
.any(|x| x.message.contains("weird") && x.message.to_lowercase().contains("rank")),
"malformed rank field must surface a warning naming the operand: {w:?}"
);
}
}