use crate::registry::{
BackwardKind, FusedOpEntry, FusedOpFamily, FusedOpParams, FusedOps, PatternMatch,
SubgraphPattern, decompose_via_recipe,
};
use crate::{Graph, NodeId};
use fuel_ir::{DType, Shape};
use fuel_kernel_seam_types::{OpAttrs, OpTag, PatternNode};
pub fn entry() -> FusedOpEntry {
FusedOpEntry {
destructive_input: None,
id: FusedOps::CAUSAL_CONV1D,
name: "CausalConv1d",
family: FusedOpFamily::Forward,
pattern: SubgraphPattern::Callable(canonical_pattern),
decompose,
backward: BackwardKind::Decompose,
shape_rule,
dtype_rule,
output_views: None,
}
}
fn shape_rule(input_shapes: &[Shape], _params: &FusedOpParams) -> Shape {
debug_assert_eq!(
input_shapes.len(),
3,
"CausalConv1d takes 3 inputs (x, weight, bias)",
);
let x_dims = input_shapes[0].dims();
let w_dims = input_shapes[1].dims();
debug_assert_eq!(
x_dims.len(),
3,
"CausalConv1d: x must be rank 3 [batch, channels, seq+pad], got {x_dims:?}"
);
debug_assert_eq!(
w_dims.len(),
3,
"CausalConv1d: weight must be rank 3 [channels, 1, kernel], got {w_dims:?}"
);
let batch = x_dims[0];
let channels = x_dims[1];
let x_seq = x_dims[2];
let kernel = w_dims[2];
debug_assert!(
x_seq >= kernel - 1,
"CausalConv1d: x time dim {x_seq} must be ≥ kernel - 1 = {} \
(caller must pre-pad with {} zeros)",
kernel - 1,
kernel - 1,
);
let out_seq = x_seq - (kernel - 1);
Shape::from_dims(&[batch, channels, out_seq])
}
fn dtype_rule(input_dtypes: &[DType], _params: &FusedOpParams) -> DType {
debug_assert_eq!(
input_dtypes.len(),
3,
"CausalConv1d takes 3 inputs (x, weight, bias)",
);
input_dtypes[0]
}
fn recipe(
out_dims: &[usize],
channels: usize,
out_seq: usize,
kernel: usize,
use_silu: bool,
) -> PatternNode {
use OpTag as T;
let op = |op, attrs, operands| PatternNode::Op {
op,
attrs,
operands,
};
let bind = |i: u8| PatternNode::Bind { index: i };
let shape_attr = |dims: &[usize]| OpAttrs {
target_shape: dims.iter().map(|&d| d as i64).collect(),
..OpAttrs::default()
};
let slice_attr = |dim: i64, start: u64, len: u64| OpAttrs {
axis: Some(dim),
slice_start: Some(start),
slice_len: Some(len),
..OpAttrs::default()
};
let per_channel = [1usize, channels, 1];
let mut acc: Option<PatternNode> = None;
for tap in 0..kernel {
let x_k = op(
T::Slice,
slice_attr(2, tap as u64, out_seq as u64),
vec![bind(0)],
);
let w_k = op(T::Slice, slice_attr(2, tap as u64, 1), vec![bind(1)]);
let w_re = op(T::Reshape, shape_attr(&per_channel), vec![w_k]);
let w_b = op(T::BroadcastTo, shape_attr(out_dims), vec![w_re]);
let term = op(T::Mul, OpAttrs::default(), vec![x_k, w_b]);
acc = Some(match acc {
None => term,
Some(a) => op(T::Add, OpAttrs::default(), vec![a, term]),
});
}
let acc = acc.expect("CausalConv1d kernel size is ≥ 1");
let b_re = op(T::Reshape, shape_attr(&per_channel), vec![bind(2)]);
let b_b = op(T::BroadcastTo, shape_attr(out_dims), vec![b_re]);
let biased = op(T::Add, OpAttrs::default(), vec![acc, b_b]);
if use_silu {
op(T::Silu, OpAttrs::default(), vec![biased])
} else {
biased
}
}
pub fn decompose(graph: &mut Graph, id: NodeId, params: &FusedOpParams) -> NodeId {
let use_silu = match params {
FusedOpParams::CausalConv1d { use_silu } => *use_silu,
_ => return id,
};
let (out_dims, kernel) = {
let n = graph.node(id);
if n.inputs.len() != 3 {
return id;
}
let out_dims = n.shape.dims().to_vec(); let w_dims = graph.node(n.inputs[1]).shape.dims().to_vec(); if out_dims.len() != 3 || w_dims.len() != 3 {
return id;
}
(out_dims, w_dims[2])
};
if kernel == 0 {
return id;
}
let channels = out_dims[1];
let out_seq = out_dims[2];
let recipe_node = recipe(&out_dims, channels, out_seq, kernel, use_silu);
decompose_via_recipe(graph, id, &recipe_node, Some(Vec::new()))
}
pub fn canonical_pattern(_graph: &Graph, _root: NodeId) -> Option<PatternMatch> {
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Node, Op};
fn frozen_legacy_causal_conv1d_decompose(
graph: &mut Graph,
id: NodeId,
params: &FusedOpParams,
) -> NodeId {
let (x_id, w_id, b_id, out_shape, dtype) = {
let n = graph.node(id);
(
n.inputs[0],
n.inputs[1],
n.inputs[2],
n.shape.clone(),
n.dtype,
)
};
let use_silu = match params {
FusedOpParams::CausalConv1d { use_silu } => *use_silu,
_ => return id,
};
let out_dims = out_shape.dims().to_vec(); let channels = out_dims[1];
let out_seq = out_dims[2];
let kernel = graph.node(w_id).shape.dims()[2]; let per_channel = Shape::from_dims(&[1, channels, 1]);
let full = out_shape.clone();
let mut acc: Option<NodeId> = None;
for tap in 0..kernel {
let x_k = graph.push(Node {
op: Op::Slice {
dim: 2,
start: tap,
len: out_seq,
},
inputs: vec![x_id],
shape: full.clone(),
dtype,
});
let w_k = graph.push(Node {
op: Op::Slice {
dim: 2,
start: tap,
len: 1,
},
inputs: vec![w_id],
shape: Shape::from_dims(&[channels, 1, 1]),
dtype,
});
let w_re = graph.push(Node {
op: Op::Reshape(per_channel.clone()),
inputs: vec![w_k],
shape: per_channel.clone(),
dtype,
});
let w_b = graph.push(Node {
op: Op::BroadcastTo(full.clone()),
inputs: vec![w_re],
shape: full.clone(),
dtype,
});
let term = graph.push(Node {
op: Op::Mul,
inputs: vec![x_k, w_b],
shape: full.clone(),
dtype,
});
acc = Some(match acc {
None => term,
Some(a) => graph.push(Node {
op: Op::Add,
inputs: vec![a, term],
shape: full.clone(),
dtype,
}),
});
}
let acc = acc.expect("CausalConv1d kernel size is ≥ 1");
let b_re = graph.push(Node {
op: Op::Reshape(per_channel.clone()),
inputs: vec![b_id],
shape: per_channel,
dtype,
});
let b_b = graph.push(Node {
op: Op::BroadcastTo(full.clone()),
inputs: vec![b_re],
shape: full.clone(),
dtype,
});
let biased = graph.push(Node {
op: Op::Add,
inputs: vec![acc, b_b],
shape: full.clone(),
dtype,
});
if use_silu {
graph.push(Node {
op: Op::Silu,
inputs: vec![biased],
shape: full,
dtype,
})
} else {
biased
}
}
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, na.shape, nb.shape
);
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 fused_node(
g: &mut Graph,
b: usize,
c: usize,
seq: usize,
kernel: usize,
use_silu: bool,
) -> NodeId {
let x_seq = seq + (kernel - 1);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[b, c, x_seq]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[c, 1, kernel]),
dtype: DType::F32,
});
let bias = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[c]),
dtype: DType::F32,
});
g.push(Node {
op: Op::Fused(
FusedOps::CAUSAL_CONV1D,
FusedOpParams::CausalConv1d { use_silu },
),
inputs: vec![x, w, bias],
shape: Shape::from_dims(&[b, c, seq]),
dtype: DType::F32,
})
}
#[test]
fn causal_conv1d_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
let shapes: [(usize, usize, usize); 4] = [
(1, 2, 3), (1, 3, 1), (2, 2, 4), (1, 1, 5), ];
let mut fired = 0usize;
for kernel in [1usize, 2, 4] {
for use_silu in [false, true] {
for (b, c, seq) in shapes {
let mut g = Graph::new();
let fused = fused_node(&mut g, b, c, seq, kernel, use_silu);
let out_shape = g.node(fused).shape.clone();
let params = FusedOpParams::CausalConv1d { use_silu };
let new_root = decompose(&mut g, fused, ¶ms);
assert_ne!(
new_root, fused,
"recipe decompose fires (kernel={kernel}, use_silu={use_silu}, \
b={b}, c={c}, seq={seq})"
);
assert_eq!(
g.node(new_root).shape,
out_shape,
"output shape matches shape_rule (kernel={kernel}, use_silu={use_silu})"
);
assert_eq!(g.node(new_root).dtype, DType::F32, "output dtype is F32");
let legacy_root = frozen_legacy_causal_conv1d_decompose(&mut g, fused, ¶ms);
assert_structural_eq(&g, new_root, legacy_root);
fired += 1;
}
}
}
assert_eq!(fired, 3 * 2 * 4, "every config in the matrix was checked");
}
#[test]
fn causal_conv1d_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let fused = fused_node(&mut g, 1, 2, 3, 4, true);
let before = g.len();
let out = 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 causal_conv1d_zero_tap_weight_is_a_fixpoint() {
let mut g = Graph::new();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 3]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 1, 0]), dtype: DType::F32,
});
let bias = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let params = FusedOpParams::CausalConv1d { use_silu: false };
let fused = g.push(Node {
op: Op::Fused(FusedOps::CAUSAL_CONV1D, params.clone()),
inputs: vec![x, w, bias],
shape: Shape::from_dims(&[1, 2, 3]),
dtype: DType::F32,
});
let before = g.len();
let out = decompose(&mut g, fused, ¶ms);
assert_eq!(out, fused, "zero-tap weight => fixpoint self-return");
assert_eq!(g.len(), before, "declined before any emission");
}
}