use crate::kernel::{KernelCaps, KernelRef};
use fuel_graph::registry::{FusedOpId, FusedOpParams};
use fuel_ir::{DType, Shape, backend::BackendCapabilities, probe::BackendId};
use smallvec::SmallVec;
use std::collections::HashMap;
#[derive(Copy, Clone)]
pub struct BackendImpl {
pub kernel: KernelRef,
pub dtypes: &'static [DType],
pub cost: fn(&[Shape], &FusedOpParams, &BackendCapabilities) -> CostEstimate,
pub precision: PrecisionGuarantee,
pub caps: KernelCaps,
pub revision: KernelRevisionHash,
pub cost_expr: Option<&'static crate::fkc::CompiledCostExpr>,
}
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct CostEstimate {
pub flops: u64,
pub bytes_moved: u64,
pub kernel_overhead_ns: u32,
}
#[derive(Copy, Clone, Debug)]
pub struct PrecisionGuarantee {
pub bit_stable_on_same_hardware: bool,
pub max_ulp: Option<u32>,
pub max_relative: Option<f64>,
pub max_absolute: Option<f64>,
pub notes: &'static str,
}
impl PrecisionGuarantee {
pub const REFERENCE: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(0),
max_relative: Some(0.0),
max_absolute: Some(0.0),
notes: "Reference IEEE-754 evaluation; bit-identical re-run.",
};
pub const PRIMITIVE_DETERMINISTIC_CPU: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend primitive kernel family: deterministic \
nested-loop iteration order; F32 accumulator for \
half-precision inputs; bit-identical same-hardware \
re-run. Per-op-per-shape ULP / relative bounds land \
with the step-8 empirical calibration framework.",
};
pub const UNAUDITED: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: false,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "PrecisionGuarantee::UNAUDITED — default placeholder for \
kernels whose precision hasn't been audited yet. Coverage \
lints flag these. Sites must either annotate via \
`register_with_precision` with a real claim or use \
`PrecisionGuarantee::none(reason)` if the audit \
concluded no static bound applies.",
};
pub const fn none(reason: &'static str) -> PrecisionGuarantee {
PrecisionGuarantee {
bit_stable_on_same_hardware: false,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: reason,
}
}
pub fn makes_no_static_claim(&self) -> bool {
!self.bit_stable_on_same_hardware
&& self.max_ulp.is_none()
&& self.max_relative.is_none()
&& self.max_absolute.is_none()
}
pub fn is_audited_no_bound(&self) -> bool {
self.makes_no_static_claim()
&& !self.notes.is_empty()
&& self.notes != Self::UNAUDITED.notes
}
}
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct KernelRevisionHash(pub u64);
impl KernelRevisionHash {
pub const UNTRACKED: KernelRevisionHash = KernelRevisionHash(0);
}
type BackendImplList = SmallVec<[(BackendId, BackendImpl); 4]>;
#[derive(Default)]
pub struct FusedKernelRegistry {
by_id: HashMap<FusedOpId, BackendImplList>,
bundle_slot_names: HashMap<(FusedOpId, BackendId, Vec<DType>), Vec<String>>,
}
impl FusedKernelRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&mut self, id: FusedOpId, backend: BackendId, impl_: BackendImpl) {
self.by_id.entry(id).or_default().push((backend, impl_));
}
pub fn lookup(&self, id: FusedOpId, backend: BackendId) -> Option<BackendImpl> {
self.by_id
.get(&id)
.and_then(|impls| impls.iter().find(|(b, _)| *b == backend))
.map(|(_, impl_)| *impl_)
}
pub fn lookup_by_dtypes(
&self,
id: FusedOpId,
backend: BackendId,
dtypes: &[DType],
) -> Option<BackendImpl> {
self.by_id
.get(&id)
.and_then(|impls| {
impls
.iter()
.find(|(b, i)| *b == backend && i.dtypes == dtypes)
})
.map(|(_, impl_)| *impl_)
}
pub fn impls_for(&self, id: FusedOpId) -> &[(BackendId, BackendImpl)] {
self.by_id.get(&id).map(|v| v.as_slice()).unwrap_or(&[])
}
pub fn len(&self) -> usize {
self.by_id.len()
}
pub fn is_empty(&self) -> bool {
self.by_id.is_empty()
}
pub fn record_bundle_slot_names(
&mut self,
id: FusedOpId,
backend: BackendId,
dtypes: &[DType],
names: &[String],
) {
self.bundle_slot_names
.insert((id, backend, dtypes.to_vec()), names.to_vec());
}
pub fn bundle_slot_names(
&self,
id: FusedOpId,
backend: BackendId,
dtypes: &[DType],
) -> Option<&[String]> {
self.bundle_slot_names
.get(&(id, backend, dtypes.to_vec()))
.map(|v| v.as_slice())
}
}
#[macro_export]
macro_rules! register_fused {
(
$registry:expr,
$id:expr,
$backend:expr,
$dtypes:expr,
$kernel:expr,
cost = $cost:expr,
precision = $precision:expr
$(, caps = $caps:expr)?
$(, revision = $revision:expr)?
$(,)?
) => {{
#[allow(unused_mut, unused_assignments)]
let mut caps = $crate::kernel::KernelCaps::empty();
$( caps = $caps; )?
#[allow(unused_mut, unused_assignments)]
let mut revision = $crate::fused::KernelRevisionHash::UNTRACKED;
$( revision = $revision; )?
$registry.register(
$id,
$backend,
$crate::fused::BackendImpl {
kernel: $kernel,
dtypes: $dtypes,
cost: $cost,
precision: $precision,
caps,
revision,
cost_expr: None,
},
);
}};
}
pub fn default_kernel_registry() -> &'static FusedKernelRegistry {
use std::sync::OnceLock;
static REGISTRY: OnceLock<FusedKernelRegistry> = OnceLock::new();
REGISTRY.get_or_init(|| {
let mut r = FusedKernelRegistry::new();
register_default_kernels(&mut r);
r
})
}
pub fn register_default_kernels(r: &mut FusedKernelRegistry) {
crate::dispatch::register_default_fused_kernels(r);
}
pub fn cost_fused_linear_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(shapes.len(), 3, "FusedLinear cost: expected 3 input shapes");
let a_dims = shapes[0].dims();
let b_dims = shapes[1].dims();
let rank = a_dims.len();
if rank < 2 || b_dims.len() < 2 {
return CostEstimate::default();
}
let m = a_dims[rank - 2] as u64;
let k = a_dims[rank - 1] as u64;
let n = b_dims[b_dims.len() - 1] as u64;
let batch: u64 = a_dims[..rank - 2]
.iter()
.map(|d| *d as u64)
.product::<u64>()
.max(1);
let mm_flops = 2u64 * m * n * k * batch;
let bias_flops = m * n * batch;
let elems = batch * (m * k + k * n + m * n) + n;
let bytes_moved = elems * 4;
CostEstimate {
flops: mm_flops + bias_flops,
bytes_moved,
kernel_overhead_ns: 50, }
}
pub const FUSED_LINEAR_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend FusedLinear: deterministic matmul + bias-add; \
BF16/F16 accumulate in F32 and narrow at end; same-hardware re-run \
bit-identical.",
};
pub fn cost_fused_softmax_cross_entropy_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
2,
"FusedSoftmaxCrossEntropy cost: expected 2 input shapes (logits, targets)",
);
let logits_dims = shapes[0].dims();
if logits_dims.is_empty() {
return CostEstimate::default();
}
let vocab = *logits_dims.last().unwrap() as u64;
let n_rows: u64 = logits_dims[..logits_dims.len() - 1]
.iter()
.map(|d| *d as u64)
.product::<u64>()
.max(1);
let per_row_flops = vocab + 2 * vocab + 10 * vocab + 10;
let total_flops = n_rows * per_row_flops;
let bytes_moved = n_rows * vocab * 4 + n_rows * 8 + n_rows * 4;
CostEstimate {
flops: total_flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const FUSED_SOFTMAX_CROSS_ENTROPY_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend FusedSoftmaxCrossEntropy: stable log-sum-exp \
in F64 accumulator; gather + NLL in F64; narrow to F32 at end; \
same-hardware re-run bit-identical.",
};
pub fn cost_causal_conv1d_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
3,
"CausalConv1d cost: expected 3 input shapes (x, weight, bias)",
);
let use_silu = match params {
FusedOpParams::CausalConv1d { use_silu } => *use_silu,
_ => return CostEstimate::default(),
};
let x_dims = shapes[0].dims();
let w_dims = shapes[1].dims();
if x_dims.len() != 3 || w_dims.len() != 3 {
return CostEstimate::default();
}
let batch = x_dims[0] as u64;
let channels = x_dims[1] as u64;
let seq_in = x_dims[2] as u64;
let kernel = w_dims[2] as u64;
if seq_in < kernel - 1 {
return CostEstimate::default();
}
let seq_out = seq_in - (kernel - 1);
let per_out_flops = 2 * kernel + if use_silu { 10 } else { 0 };
let flops = batch * channels * seq_out * per_out_flops;
let bytes_moved = batch * channels * seq_in * 4
+ channels * kernel * 4
+ channels * 4
+ batch * channels * seq_out * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const CAUSAL_CONV1D_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend CausalConv1d: textbook depthwise nested-loop \
with F32 accumulator; optional SiLU via x/(1+exp(-x)); \
same-hardware re-run bit-identical.",
};
pub fn cost_selective_scan_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
5,
"SelectiveScan cost: expected 5 input shapes (u, delta, a, b, c)",
);
let u_dims = shapes[0].dims();
let a_dims = shapes[2].dims();
if u_dims.len() != 3 || a_dims.len() != 2 {
return CostEstimate::default();
}
let batch = u_dims[0] as u64;
let seqlen = u_dims[1] as u64;
let dim = u_dims[2] as u64;
let dstate = a_dims[1] as u64;
let per_iter_flops = 16;
let flops = batch * seqlen * dim * dstate * per_iter_flops;
let bytes_moved = 2 * batch * seqlen * dim * 4
+ 2 * batch * seqlen * dstate * 4
+ dim * dstate * 4
+ batch * seqlen * dim * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const SELECTIVE_SCAN_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend SelectiveScan: state recurrence in F64 \
accumulator; narrow to F32 on y store; deterministic \
iteration order; same-hardware re-run bit-identical.",
};
pub fn cost_ssd_chunk_scan_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
5,
"SsdChunkScan cost: expected 5 input shapes (x, dt, a, b, c)",
);
let x_dims = shapes[0].dims();
let b_dims = shapes[3].dims();
if x_dims.len() != 4 || b_dims.len() != 4 {
return CostEstimate::default();
}
let batch = x_dims[0] as u64;
let seqlen = x_dims[1] as u64;
let heads = x_dims[2] as u64;
let head_dim = x_dims[3] as u64;
let state_dim = b_dims[3] as u64;
let per_iter_flops = 8;
let flops = batch * seqlen * heads * head_dim * state_dim * per_iter_flops;
let bytes_moved = 2 * batch * seqlen * heads * head_dim * 4
+ batch * seqlen * heads * 4
+ 2 * batch * seqlen * heads * state_dim * 4
+ heads * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const SSD_CHUNK_SCAN_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend SsdChunkScan: per-head state recurrence \
in F64 accumulator; narrow to F32 on y store; \
deterministic iteration order; same-hardware re-run \
bit-identical. chunk_size is a GPU-parallelism knob; \
CPU runs sequential, any positive chunk_size dividing \
seqlen produces the same answer.",
};
pub fn cost_nf4_matmul_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
3,
"Nf4Matmul cost: expected 3 input shapes (activations, w_packed, absmax)",
);
let block_size = match params {
FusedOpParams::Nf4Matmul { block_size } => *block_size as u64,
_ => return CostEstimate::default(),
};
let a_dims = shapes[0].dims();
let w_dims = shapes[1].dims();
if a_dims.len() < 2 || w_dims.len() != 2 {
return CostEstimate::default();
}
let m = a_dims[a_dims.len() - 2] as u64;
let k = a_dims[a_dims.len() - 1] as u64;
let n = w_dims[0] as u64;
let batch: u64 = a_dims[..a_dims.len() - 2]
.iter()
.map(|d| *d as u64)
.product::<u64>()
.max(1);
if block_size == 0 || !k.is_multiple_of(block_size) {
return CostEstimate::default();
}
let flops = batch * m * n * k * 3;
let elem_bytes = 4;
let bytes_moved = batch * m * k * elem_bytes
+ n * (k / 2)
+ n * (k / block_size) * 4
+ batch * m * n * elem_bytes;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const NF4_MATMUL_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend Nf4Matmul: dequant via baked NF4_LUT × \
per-block absmax; F32 accumulator for the inner product \
(F16/BF16 up-cast on load, narrow on store); \
deterministic iteration order; same-hardware re-run \
bit-identical.",
};
pub fn cost_conv2d_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert!(
shapes.len() == 2 || shapes.len() == 3,
"Conv2D cost: expected 2 or 3 input shapes",
);
let (stride, padding, groups) = match params {
FusedOpParams::Conv2D {
stride,
padding,
groups,
} => (*stride, *padding, *groups),
_ => return CostEstimate::default(),
};
let x_dims = shapes[0].dims();
let w_dims = shapes[1].dims();
if x_dims.len() != 4 || w_dims.len() != 4 || groups == 0 {
return CostEstimate::default();
}
let (n, cin, h_in, w_in) = (
x_dims[0] as u64,
x_dims[1] as u64,
x_dims[2] as u64,
x_dims[3] as u64,
);
let (cout, cin_per_g, kh, kw) = (
w_dims[0] as u64,
w_dims[1] as u64,
w_dims[2] as u64,
w_dims[3] as u64,
);
let (sh, sw) = (stride.0 as u64, stride.1 as u64);
let (ph, pw) = (padding.0 as u64, padding.1 as u64);
if sh == 0 || sw == 0 || kh > h_in + 2 * ph || kw > w_in + 2 * pw {
return CostEstimate::default();
}
let h_out = (h_in + 2 * ph - kh) / sh + 1;
let w_out = (w_in + 2 * pw - kw) / sw + 1;
let conv_flops = 2u64 * n * cout * h_out * w_out * cin_per_g * kh * kw;
let bias_flops = if shapes.len() == 3 {
n * cout * h_out * w_out
} else {
0
};
let elems_in = n * cin * h_in * w_in;
let elems_w = cout * cin_per_g * kh * kw;
let elems_out = n * cout * h_out * w_out;
let elems_bias = if shapes.len() == 3 { cout } else { 0 };
let bytes_moved = (elems_in + elems_w + elems_out + elems_bias) * 4;
CostEstimate {
flops: conv_flops + bias_flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const CONV2D_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend Conv2D: textbook nested-loop \
cross-correlation; F32 accumulator (F64 for F64 input); \
BF16/F16 multiply in F32 and narrow at end; same-hardware \
re-run bit-identical.",
};
pub const NORM_FAMILY_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend Softmax/RmsNorm/LayerNorm: deterministic \
elementwise + per-row reduction; F32 accumulator (F64 for \
F64 input); BF16/F16 multiply in F32 and narrow at end; \
same-hardware re-run bit-identical.",
};
pub fn cost_norm_family_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert!(
!shapes.is_empty(),
"Norm-family cost: expected ≥1 input shape",
);
let dims = shapes[0].dims();
if dims.is_empty() {
return CostEstimate::default();
}
let elems: u64 = dims.iter().map(|&d| d as u64).product();
let flops_per_elem: u64 = match params {
FusedOpParams::SoftmaxLastDim | FusedOpParams::SoftmaxLastDimBackward => 5, FusedOpParams::RmsNormLastDim { .. } | FusedOpParams::RmsNormLastDimBackward { .. } => 4, FusedOpParams::LayerNormLastDim { .. } | FusedOpParams::LayerNormLastDimBackward { .. } => {
7
} _ => 5, };
let bandwidth_factor: u64 = if shapes.len() >= 2 { 3 } else { 2 };
CostEstimate {
flops: flops_per_elem * elems,
bytes_moved: bandwidth_factor * elems * 4,
kernel_overhead_ns: 50,
}
}
pub const ROPE_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend Rope: per-head rotation `out = x·cos + \
rotated(x)·sin` with deterministic iteration order. \
BF16/F16 multiply in F32 and narrow at end; same-hardware \
re-run bit-identical.",
};
pub fn cost_rope_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
3,
"Rope cost: expected 3 input shapes (x, cos, sin)"
);
let dims = shapes[0].dims();
if dims.is_empty() {
return CostEstimate::default();
}
let elems: u64 = dims.iter().map(|&d| d as u64).product();
let flops = 4 * elems;
let cs_elems: u64 = shapes[1].dims().iter().map(|&d| d as u64).product::<u64>() * 2;
let bytes_moved = (2 * elems + cs_elems) * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const CONV_TRANSPOSE2D_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend ConvTranspose2D: textbook scatter-with-stride \
+ nested loops; F32 accumulator (F64 for F64); BF16/F16 \
multiply in F32 and narrow at end; same-hardware bit-identical.",
};
pub fn cost_conv_transpose2d_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
2,
"ConvTranspose2D cost: expected 2 input shapes"
);
let (stride, padding, output_padding, dilation, groups) = match params {
FusedOpParams::ConvTranspose2D {
stride,
padding,
output_padding,
dilation,
groups,
} => (*stride, *padding, *output_padding, *dilation, *groups),
_ => return CostEstimate::default(),
};
let x_dims = shapes[0].dims();
let w_dims = shapes[1].dims();
if x_dims.len() != 4 || w_dims.len() != 4 || groups == 0 {
return CostEstimate::default();
}
let (n, cin, h_in, w_in) = (
x_dims[0] as u64,
x_dims[1] as u64,
x_dims[2] as u64,
x_dims[3] as u64,
);
let (_, cout_per_g, kh, kw) = (
w_dims[0] as u64,
w_dims[1] as u64,
w_dims[2] as u64,
w_dims[3] as u64,
);
let cout = cout_per_g * groups as u64;
let h_out = h_in.saturating_sub(1) * stride.0 as u64
+ dilation.0 as u64 * kh.saturating_sub(1)
+ output_padding.0 as u64
+ 1;
let h_out = h_out.saturating_sub(2 * padding.0 as u64);
let w_out = w_in.saturating_sub(1) * stride.1 as u64
+ dilation.1 as u64 * kw.saturating_sub(1)
+ output_padding.1 as u64
+ 1;
let w_out = w_out.saturating_sub(2 * padding.1 as u64);
let cin_per_g = cin / groups as u64;
let flops = 2u64 * n * cout * h_out * w_out * cin_per_g * kh * kw;
let elems_in = n * cin * h_in * w_in;
let elems_w = cin * cout_per_g * kh * kw;
let elems_out = n * cout * h_out * w_out;
let bytes_moved = (elems_in + elems_w + elems_out) * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const ATTN_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend FlashAttn/PagedAttn: naive scaled-dot- \
product reference; F32 accumulator; deterministic \
iteration order; bit-identical re-run on same hardware. \
GPU kernels produce different numerics (tiled softmax) — \
their PrecisionGuarantee is declared separately when they \
register.",
};
pub fn cost_attn_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
if shapes.is_empty() {
return CostEstimate::default();
}
let q_dims = shapes[0].dims();
if q_dims.len() != 4 {
return CostEstimate::default();
}
let (b, hq, sq, d) = (
q_dims[0] as u64,
q_dims[1] as u64,
q_dims[2] as u64,
q_dims[3] as u64,
);
let sk: u64 = if shapes.len() >= 2 {
let k_dims = shapes[1].dims();
if k_dims.len() == 4 {
(k_dims[2] as u64).max(k_dims[0] as u64 * k_dims[1] as u64)
} else {
sq
}
} else {
sq
};
let mm_flops = 4u64 * b * hq * sq * sk * d;
let sm_flops = 5u64 * b * hq * sq * sk;
let elems_qkv = b * hq * sq * d + 2 * b * hq * sk * d;
let elems_out = b * hq * sq * d;
let bytes_moved = (elems_qkv + elems_out) * 4;
CostEstimate {
flops: mm_flops + sm_flops,
bytes_moved,
kernel_overhead_ns: 200,
}
}
pub const ATTN_BACKWARD_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend FlashAttnBackward{Q,K,V}: math-definition \
gradient with recomputed softmax; sequential nested-loop \
accumulation (F32 accumulator for half types); \
deterministic iteration order; bit-identical re-run on \
same hardware. GPU backward kernels (tiled, atomics) \
declare their own PrecisionGuarantee when they register.",
};
pub fn cost_attn_backward_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
if shapes.is_empty() {
return CostEstimate::default();
}
let q_dims = shapes[0].dims();
if q_dims.len() != 4 {
return CostEstimate::default();
}
let (b, hq, sq, d) = (
q_dims[0] as u64,
q_dims[1] as u64,
q_dims[2] as u64,
q_dims[3] as u64,
);
let sk: u64 = if shapes.len() >= 2 {
let k_dims = shapes[1].dims();
if k_dims.len() == 4 {
k_dims[2] as u64
} else {
sq
}
} else {
sq
};
let mm_flops = 10u64 * b * hq * sq * sk * d;
let sm_flops = 8u64 * b * hq * sq * sk;
let elems_read = 2 * b * hq * sq * d + 2 * b * hq * sk * d;
let elems_written = b * hq * sq * d + 2 * b * hq * sk * d;
let bytes_moved = (elems_read + elems_written) * 4;
CostEstimate {
flops: mm_flops + sm_flops,
bytes_moved,
kernel_overhead_ns: 200,
}
}
pub const QMATMUL_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend QMatMul: per-block dequant + F32 matmul \
accumulate; deterministic iteration order; bit-identical \
re-run on same hardware. Inherent precision loss from \
the quantization itself is a property of the QuantType, \
not the kernel — captured separately when the calibration \
framework lands.",
};
pub const REDUCE_MAX_TO_BACKWARD_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend ReduceMaxTo backward: recomputed-max + \
tie-mask + scale + gate; deterministic iteration order; \
F32 accumulator for half-precision; bit-identical same-\
hardware re-run.",
};
pub const INPLACE_AFFINE_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend in-place affine: single-pass `mul · x + \
add` per element; deterministic iteration order; F32 pivot \
for half-precision dtypes; bit-identical same-hardware re-run.",
};
pub fn cost_inplace_affine_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
1,
"InplaceAffine cost: expected 1 input shape"
);
let n: u64 = shapes[0].dims().iter().map(|&d| d as u64).product();
let flops = n;
let bytes_moved = n * 4 * 2;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const POWI_BACKWARD_CPU_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-cpu-backend PowI backward: single-pass `exp · \
x.powi(exp-1) · upstream` per element; deterministic \
iteration order; F32 compute for half-precision; bit-\
identical same-hardware re-run.",
};
pub fn cost_powi_backward_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
2,
"PowIBackward cost: expected 2 input shapes"
);
let in_count: u64 = shapes[0].dims().iter().map(|&d| d as u64).product();
let exp_abs = match params {
FusedOpParams::PowIBackward { exp } => (*exp).unsigned_abs().max(1) as u64,
_ => 1,
};
let powi_muls = 64 - exp_abs.leading_zeros() as u64;
let flops = in_count * (powi_muls + 2);
let bytes_moved = in_count * 4 + in_count * 4 + in_count * 4 ;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 60,
}
}
pub fn cost_reduce_max_to_backward_cpu(
shapes: &[Shape],
_params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(
shapes.len(),
2,
"ReduceMaxToBackward cost: expected 2 input shapes"
);
let in_count: u64 = shapes[0].dims().iter().map(|&d| d as u64).product();
let out_count: u64 = shapes[1].dims().iter().map(|&d| d as u64).product();
let flops = 5 * in_count + 2 * out_count;
let bytes_moved = (5 * in_count + 3 * out_count) * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 80,
}
}
pub fn cost_qmatmul_cpu(
shapes: &[Shape],
params: &FusedOpParams,
_caps: &BackendCapabilities,
) -> CostEstimate {
debug_assert_eq!(shapes.len(), 2, "QMatMul cost: expected 2 input shapes");
let (k, n) = match params {
FusedOpParams::QMatMul { k, n, .. } => (*k as u64, *n as u64),
_ => return CostEstimate::default(),
};
let a_dims = shapes[0].dims();
let rank = a_dims.len();
if rank < 2 {
return CostEstimate::default();
}
let m = a_dims[rank - 2] as u64;
let batch: u64 = a_dims[..rank - 2]
.iter()
.map(|&d| d as u64)
.product::<u64>()
.max(1);
let flops = 2 * batch * m * n * k;
let a_elems = batch * m * k;
let w_u32_elems: u64 = shapes[1].dims().iter().map(|&d| d as u64).product();
let c_elems = batch * m * n;
let bytes_moved = (a_elems + c_elems) * 4 + w_u32_elems * 4;
CostEstimate {
flops,
bytes_moved,
kernel_overhead_ns: 50,
}
}
pub const VULKAN_FLOAT_POINTWISE_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(1),
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend f32/f64 pointwise: direct hardware FADD/\
FMUL/FMA per thread; no atomics; bit-stable on same hardware. \
Within 1 ULP of IEEE-754 round-to-nearest; small variation \
across implementations due to commutative-reorder + FMA fusion.",
};
pub const VULKAN_HALF_POINTWISE_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(1),
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend f16/bf16 pointwise: shaderFloat16 native or \
u32-packed bf16; bit-stable on same hardware; within 1 ULP in \
the dtype's mantissa space (10-bit f16, 7-bit bf16).",
};
pub const VULKAN_TRANSCENDENTAL_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(4),
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend transcendentals: GLSL.std.450 vendor lib \
(f32/f16/bf16) or Horner polynomial (f64); bit-stable on same \
hardware; Vulkan spec allows ≤4 ULP, real GPUs typically tighter.",
};
pub const VULKAN_MATMUL_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend f32 matmul: tiled / reg-tile / matvec \
kernels with per-thread sequential FMA accumulation; bit-stable \
on same hardware; error bounded by O(K · ULP) from K-deep \
reduction.",
};
pub const VULKAN_MATMUL_TENSORCORE_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend tensor-core matmul (f32 × bf16 → f32): \
cooperative-matrix FMAs with f32 internal accumulator; bit-stable \
on same hardware; per-tile error bounded by O(K · ULP_bf16 + \
tile_K · ULP_f32).",
};
pub const VULKAN_BYTE_LEVEL_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(0),
max_relative: Some(0.0),
max_absolute: Some(0.0),
notes: "fuel-vulkan-backend byte-level ops (Triu/Tril/Flip/Roll/\
WriteSlice/IndexSelect/Concat/Copy): pure data movement, no FP \
math; bit-identical to input bytes.",
};
pub const VULKAN_CAST_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(0),
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend Cast: pure dtype conversion via SPIR-V \
OpFConvert / OpConvertFToS / OpConvertSToF; bit-stable on same \
hardware; ULP=0 vs IEEE-754 round-to-nearest in target dtype \
(loss is inherent to the target dtype, not the implementation).",
};
pub const VULKAN_QMATMUL_PRECISION: PrecisionGuarantee = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: None,
max_relative: None,
max_absolute: None,
notes: "fuel-vulkan-backend QMatMul (Q4_0/Q4_K_M/Q8_0 × F32): per-block \
dequant + per-output-element FMA; bit-stable on same hardware. \
Quantization precision loss is a QuantType property, not the \
kernel's contribution.",
};
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::{DType, Layout, Result};
use std::sync::Arc;
use std::sync::RwLock;
fn dummy_kernel(
_inputs: &[Arc<RwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
_layouts: &[Layout],
_params: &crate::kernel::OpParams,
) -> Result<()> {
Ok(())
}
fn dummy_cost(_s: &[Shape], _p: &FusedOpParams, _c: &BackendCapabilities) -> CostEstimate {
CostEstimate::default()
}
const DUMMY_DTYPES: &[DType] = &[DType::F32];
fn make_impl() -> BackendImpl {
BackendImpl {
kernel: dummy_kernel,
dtypes: DUMMY_DTYPES,
cost: dummy_cost,
precision: PrecisionGuarantee::UNAUDITED,
caps: KernelCaps::empty(),
revision: KernelRevisionHash::UNTRACKED,
cost_expr: None,
}
}
#[test]
fn fused_kernel_registry_empty() {
let r = FusedKernelRegistry::new();
assert!(r.is_empty());
assert_eq!(r.len(), 0);
assert!(r.lookup(FusedOpId(1), BackendId::Cpu).is_none());
assert!(r.impls_for(FusedOpId(1)).is_empty());
let _ = DType::F32;
}
#[test]
fn fused_kernel_registry_register_and_lookup() {
let mut r = FusedKernelRegistry::new();
r.register(FusedOpId(1), BackendId::Cpu, make_impl());
assert!(!r.is_empty());
assert_eq!(r.len(), 1);
let got = r.lookup(FusedOpId(1), BackendId::Cpu);
assert!(got.is_some());
let impls = r.impls_for(FusedOpId(1));
assert_eq!(impls.len(), 1);
assert_eq!(impls[0].0, BackendId::Cpu);
}
#[test]
fn fused_kernel_registry_register_appends_alternatives() {
let mut r = FusedKernelRegistry::new();
r.register(FusedOpId(1), BackendId::Cpu, make_impl());
r.register(FusedOpId(1), BackendId::Cpu, make_impl());
assert_eq!(r.len(), 1, "one id");
assert_eq!(
r.impls_for(FusedOpId(1)).len(),
2,
"two alternatives registered",
);
}
#[test]
fn precision_guarantee_reference_is_strict() {
let p = PrecisionGuarantee::REFERENCE;
assert!(p.bit_stable_on_same_hardware);
assert_eq!(p.max_ulp, Some(0));
}
#[test]
fn cost_estimate_default_zero() {
let c = CostEstimate::default();
assert_eq!(c.flops, 0);
assert_eq!(c.bytes_moved, 0);
assert_eq!(c.kernel_overhead_ns, 0);
}
#[test]
fn default_kernel_registry_has_fused_linear_cpu_quartet() {
use fuel_graph::registry::FusedOps;
let r = default_kernel_registry();
let impls = r.impls_for(FusedOps::FUSED_LINEAR);
assert_eq!(
impls.len(),
4,
"expected 4 FusedLinear × Cpu × dtype registrations, got {}",
impls.len(),
);
for (backend, _) in impls {
assert_eq!(*backend, BackendId::Cpu);
}
for dtype in [DType::F32, DType::F64, DType::BF16, DType::F16] {
let want = [dtype; 4];
let got = r.lookup_by_dtypes(FusedOps::FUSED_LINEAR, BackendId::Cpu, &want);
assert!(
got.is_some(),
"lookup_by_dtypes missed FusedLinear × Cpu × {dtype:?}",
);
let impl_ = got.unwrap();
assert!(
impl_.precision.bit_stable_on_same_hardware,
"CPU FusedLinear should be bit-stable on same hardware",
);
}
}
#[test]
fn default_kernel_registry_has_conv2d_cpu_octet() {
use fuel_graph::registry::FusedOps;
let r = default_kernel_registry();
let impls = r.impls_for(FusedOps::CONV2D);
assert_eq!(
impls.len(),
8,
"expected 8 Conv2D × Cpu × dtype × (no-bias|with-bias) \
registrations, got {}",
impls.len(),
);
for (backend, _) in impls {
assert_eq!(*backend, BackendId::Cpu);
}
for dtype in [DType::F32, DType::F64, DType::BF16, DType::F16] {
let no_bias = [dtype; 3];
let with_bias = [dtype; 4];
assert!(
r.lookup_by_dtypes(FusedOps::CONV2D, BackendId::Cpu, &no_bias)
.is_some(),
"lookup_by_dtypes missed Conv2D × Cpu × {dtype:?} (no-bias)",
);
let with_bias_impl = r.lookup_by_dtypes(FusedOps::CONV2D, BackendId::Cpu, &with_bias);
assert!(
with_bias_impl.is_some(),
"lookup_by_dtypes missed Conv2D × Cpu × {dtype:?} (with-bias)",
);
let impl_ = with_bias_impl.unwrap();
assert!(
impl_.precision.bit_stable_on_same_hardware,
"CPU Conv2D should be bit-stable on same hardware",
);
}
}
#[test]
fn precision_guarantee_lint_bit_stable_cpu_coverage() {
use fuel_graph::registry::{FusedOpId, FusedOps, default_registry};
const KNOWN_GAPS: &[(FusedOpId, &str)] = &[(
FusedOps::AWQ_MATMUL,
"decompose-only in this increment; no direct CPU binding-table \
registration yet -- see fuel-graph/src/registry/awq_matmul.rs's \
module doc for the follow-on dispatch-wiring work",
)];
let meta = default_registry();
let kernels = default_kernel_registry();
let mut failures: Vec<String> = Vec::new();
let mut covered = 0usize;
let mut allowlisted = 0usize;
for entry in meta.entries_iter() {
let id = entry.id;
let name = entry.name;
if let Some((_, reason)) = KNOWN_GAPS.iter().find(|(g, _)| *g == id) {
allowlisted += 1;
let has_cpu = kernels
.impls_for(id)
.iter()
.any(|(b, _)| *b == BackendId::Cpu);
if has_cpu {
failures.push(format!(
"FusedOpId {id:?} ({name}) is on the KNOWN_GAPS allowlist \
but DOES have a CPU registration now. Reason given was: \
{reason:?}. Remove the allowlist entry to enable the \
bit-stable lint.",
));
}
continue;
}
let impls = kernels.impls_for(id);
let has_bit_stable_cpu = impls.iter().any(|(backend, impl_)| {
*backend == BackendId::Cpu && impl_.precision.bit_stable_on_same_hardware
});
if has_bit_stable_cpu {
covered += 1;
} else {
failures.push(format!(
"FusedOpId {id:?} ({name}) has no bit-stable CPU BackendImpl. \
Architecture v1.0 §05 requires the always-built backend \
(fuel-cpu-backend) to provide at least one \
`bit_stable_on_same_hardware: true` kernel per registered \
fused op. Either add a CPU registration in \
fuel-storage::dispatch::register_default_fused_kernels with \
a matching PrecisionGuarantee, or add a line to KNOWN_GAPS \
above with a documented reason.",
));
}
}
assert!(
failures.is_empty(),
"Architecture v1.0 bit-stable CPU coverage lint failed:\n{}",
failures.join("\n"),
);
assert!(
covered > 0 || allowlisted > 0,
"lint covered 0 ops — registry appears empty"
);
assert_eq!(
allowlisted, 1,
"KNOWN_GAPS allowlist has exactly 1 documented entry (AwqMatmul); \
saw {allowlisted} allowlisted",
);
assert_eq!(
covered, 24,
"expected 24 fused ops covered by bit-stable CPU impls, got {covered}",
);
}
#[test]
fn precision_guarantee_lint_detects_missing_cpu() {
use fuel_graph::registry::FusedOpId;
let empty = FusedKernelRegistry::new();
let id = FusedOpId(1);
let has_bit_stable_cpu = empty.impls_for(id).iter().any(|(backend, impl_)| {
*backend == BackendId::Cpu && impl_.precision.bit_stable_on_same_hardware
});
assert!(
!has_bit_stable_cpu,
"empty registry must not report bit-stable CPU coverage"
);
}
#[test]
fn precision_guarantee_lint_rejects_non_bit_stable_cpu() {
use fuel_graph::registry::FusedOpId;
let mut r = FusedKernelRegistry::new();
let id = FusedOpId(42);
let weak = BackendImpl {
kernel: dummy_kernel,
dtypes: DUMMY_DTYPES,
cost: dummy_cost,
precision: PrecisionGuarantee::UNAUDITED,
caps: KernelCaps::empty(),
revision: KernelRevisionHash::UNTRACKED,
cost_expr: None,
};
r.register(id, BackendId::Cpu, weak);
let has_bit_stable_cpu = r.impls_for(id).iter().any(|(backend, impl_)| {
*backend == BackendId::Cpu && impl_.precision.bit_stable_on_same_hardware
});
assert!(
!has_bit_stable_cpu,
"an UNAUDITED-precision CPU impl must not satisfy the bit-stable lint"
);
}
#[test]
fn default_kernel_registry_step6_coverage() {
use fuel_graph::registry::FusedOps;
let r = default_kernel_registry();
for (id, want) in [
(FusedOps::SOFTMAX_LAST_DIM, 4usize),
(FusedOps::RMS_NORM_LAST_DIM, 4),
(FusedOps::LAYER_NORM_LAST_DIM, 4),
(FusedOps::ROPE, 4),
(FusedOps::CONV_TRANSPOSE2D, 8),
(FusedOps::FLASH_ATTN, 8),
(FusedOps::FLASH_ATTN_BACKWARD_Q, 8),
(FusedOps::FLASH_ATTN_BACKWARD_K, 8),
(FusedOps::FLASH_ATTN_BACKWARD_V, 8),
(FusedOps::PAGED_ATTN, 8),
(FusedOps::QMATMUL, 1),
(FusedOps::SOFTMAX_LAST_DIM_BACKWARD, 4),
(FusedOps::LAYER_NORM_LAST_DIM_BACKWARD, 4),
(FusedOps::RMS_NORM_LAST_DIM_BACKWARD, 4),
(FusedOps::REDUCE_MAX_TO_BACKWARD, 4),
] {
let impls = r.impls_for(id);
assert_eq!(
impls.len(),
want,
"expected {want} CPU impls for FusedOpId {id:?}, got {}",
impls.len(),
);
for (backend, impl_) in impls {
assert_eq!(*backend, BackendId::Cpu);
assert!(
impl_.precision.bit_stable_on_same_hardware,
"CPU impl for {id:?} should be bit-stable on same hardware",
);
}
}
}
#[test]
fn cost_conv2d_cpu_flops_scale() {
use crate::fused::cost_conv2d_cpu;
use fuel_ir::DeviceLocation;
use std::collections::HashSet;
let x = Shape::from_dims(&[2, 4, 8, 8]);
let w = Shape::from_dims(&[6, 4, 3, 3]);
let bias = Shape::from_dims(&[6]);
let caps = BackendCapabilities {
backend_id: BackendId::Cpu,
device_location: DeviceLocation::Cpu,
op_dtype_support: HashSet::new(),
required_alignment: 1,
access_granularity_bits: 8,
transfer_paths: vec![],
storage_substrate: fuel_ir::backend::SubstrateClass::HostBytes,
compute_throughput_flops_per_ns: 1.0,
mem_bandwidth_bytes_per_ns: 4.0,
};
let params = FusedOpParams::Conv2D {
stride: (1, 1),
padding: (1, 1),
groups: 1,
};
let c = cost_conv2d_cpu(&[x, w, bias], ¶ms, &caps);
assert_eq!(c.flops, 55296 + 768);
assert!(c.bytes_moved > 0);
}
#[test]
fn cost_fused_linear_cpu_flops_scale() {
use fuel_ir::DeviceLocation;
use std::collections::HashSet;
let a = Shape::from_dims(&[2, 4, 8]); let b = Shape::from_dims(&[2, 8, 16]); let bias = Shape::from_dims(&[16]);
let caps = BackendCapabilities {
backend_id: BackendId::Cpu,
device_location: DeviceLocation::Cpu,
op_dtype_support: HashSet::new(),
required_alignment: 1,
access_granularity_bits: 8,
transfer_paths: vec![],
storage_substrate: fuel_ir::backend::SubstrateClass::HostBytes,
compute_throughput_flops_per_ns: 1.0,
mem_bandwidth_bytes_per_ns: 4.0,
};
let c = cost_fused_linear_cpu(&[a, b, bias], &FusedOpParams::FusedLinear, &caps);
assert_eq!(c.flops, 2048 + 128);
assert!(c.bytes_moved > 0);
}
#[test]
fn bundle_slot_names_round_trip_through_the_fused_registry() {
use fuel_graph::registry::FusedOps;
use fuel_ir::{DType, probe::BackendId};
let mut reg = FusedKernelRegistry::new();
let dtypes = &[DType::F32, DType::F32][..];
reg.record_bundle_slot_names(
FusedOps::SELECTIVE_SCAN,
BackendId::Cpu,
dtypes,
&["y".to_string(), "last_state".to_string()],
);
assert_eq!(
reg.bundle_slot_names(FusedOps::SELECTIVE_SCAN, BackendId::Cpu, dtypes),
Some(&["y".to_string(), "last_state".to_string()][..])
);
assert_eq!(
reg.bundle_slot_names(FusedOps::SELECTIVE_SCAN, BackendId::Cpu, &[DType::F16][..]),
None
);
}
}