use std::collections::HashMap;
use std::sync::Arc;
use std::sync::RwLock;
use fuel_graph::QuantType;
use fuel_graph::registry::FusedOpId;
use fuel_ir::conv::{ParamsConv1D, ParamsConvTranspose1D};
use fuel_ir::dispatch::OpKind;
use fuel_ir::probe::BackendId;
use fuel_ir::{DType, Error, Layout, Result};
use smallvec::SmallVec;
pub type KernelDTypes = SmallVec<[DType; 8]>;
pub type CoverageBinding = (BackendId, OpKind, Vec<DType>);
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct KernelCaps {
pub strided_input: bool,
pub requires_broadcast: bool,
}
impl KernelCaps {
pub const fn empty() -> Self {
Self {
strided_input: false,
requires_broadcast: false,
}
}
pub const fn strided_input() -> Self {
Self {
strided_input: true,
requires_broadcast: false,
}
}
}
use fuel_memory::Storage;
pub type KernelRef = fn(
inputs: &[Arc<RwLock<Storage>>],
outputs: &mut [Arc<RwLock<Storage>>],
layouts: &[Layout],
params: &OpParams,
) -> Result<()>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MatmulM {
All,
Rows(usize),
Deferred(fuel_ir::DynScalar),
}
#[derive(Debug, Clone)]
pub enum OpParams {
None,
Reduce { dims: Vec<usize>, keepdim: bool },
Matmul {
lhs_batch_dims: Vec<usize>,
rhs_batch_dims: Vec<usize>,
m: usize,
n: usize,
k: usize,
m_compute: MatmulM,
},
Conv1D(ParamsConv1D),
Conv2D {
x_shape: [usize; 4],
w_shape: [usize; 4],
out_shape: [usize; 4],
stride: (usize, usize),
padding: (usize, usize),
dilation: (usize, usize),
groups: usize,
},
ConvTranspose1D(ParamsConvTranspose1D),
ConvTranspose2D {
x_shape: [usize; 4],
w_shape: [usize; 4],
out_shape: [usize; 4],
stride: (usize, usize),
padding: (usize, usize),
output_padding: (usize, usize),
dilation: (usize, usize),
groups: usize,
},
ReduceSumTo {
input_shape: Vec<usize>,
output_shape: Vec<usize>,
},
ReduceMaxTo {
input_shape: Vec<usize>,
output_shape: Vec<usize>,
},
ReduceMaxToBackward {
input_shape: Vec<usize>,
output_shape: Vec<usize>,
},
FlashAttn {
b: usize,
hq: usize,
hkv: usize,
sq: usize,
sk: usize,
d: usize,
k_len: usize,
softmax_scale: f32,
causal: bool,
window_size_left: Option<usize>,
window_size_right: Option<usize>,
softcap: Option<f32>,
},
PagedAttn {
b: usize,
hq: usize,
hkv: usize,
sq: usize,
d: usize,
block_size: usize,
max_blocks_per_seq: usize,
num_blocks: usize,
softmax_scale: f32,
softcap: Option<f32>,
},
Slice {
dim: usize,
start: usize,
end: usize,
step: usize,
},
Cast,
Affine { mul: f64, add: f64 },
Clamp { min: f64, max: f64 },
PowI { exp: i32 },
Concat {
outer_count: usize,
input_dim_sizes: Vec<usize>,
inner_count: usize,
axis: usize,
},
SoftmaxLastDim { outer_count: usize, last_dim: usize },
NormLastDim {
outer_count: usize,
last_dim: usize,
eps: f64,
},
IndexSelect {
outer_count: usize,
source_dim_size: usize,
n_indices: usize,
inner_count: usize,
},
Gather {
source_shape: Vec<usize>,
output_shape: Vec<usize>,
dim: usize,
},
Rope {
outer_count: usize,
seq: usize,
head_dim: usize,
},
QMatMul {
quant_type: QuantType,
batch_count: usize,
m: usize,
n: usize,
k: usize,
},
IndexAdd {
outer_count: usize,
base_dim_size: usize,
n_indices: usize,
inner_count: usize,
},
ScatterAdd {
base_shape: Vec<usize>,
src_shape: Vec<usize>,
dim: usize,
},
Flip {
outer_count: usize,
dim_size: usize,
inner_count: usize,
axis: usize,
},
Roll {
outer_count: usize,
dim_size: usize,
inner_count: usize,
shift: i64,
axis: usize,
},
CumSum {
outer_count: usize,
dim_size: usize,
inner_count: usize,
axis: usize,
},
Triangular {
batch_count: usize,
rows: usize,
cols: usize,
diagonal: i64,
},
LogSoftmaxLastDim { outer_count: usize, last_dim: usize },
MaskedFill { fill_bytes: Vec<u8> },
PadBackward {
in_shape: Vec<usize>,
out_shape: Vec<usize>,
padding: Vec<(usize, usize)>,
mode_tag: u8,
},
Pad {
in_shape: Vec<usize>,
out_shape: Vec<usize>,
padding: Vec<(usize, usize)>,
mode_tag: u8,
fill_bytes: Vec<u8>,
},
WriteSlice {
dest_shape: Vec<usize>,
ranges: Vec<(usize, usize)>,
deferred_dyn_offset: Option<(usize, fuel_ir::DynScalar)>,
},
WriteSliceRotating {
dest_shape: Vec<usize>,
axis: usize,
modulus: usize,
ranges: Vec<(usize, usize)>,
},
WriteSliceDoff {
dest_shape: Vec<usize>,
axis: usize,
ranges: Vec<(usize, usize)>,
},
FusedSoftmaxCrossEntropy {
n_rows: usize,
vocab: usize,
reduction: fuel_graph::registry::Reduction,
ignore_index: i64,
},
CausalConv1d {
batch: usize,
channels: usize,
seq_in: usize,
seq_out: usize,
kernel: usize,
use_silu: bool,
},
SelectiveScan {
batch: usize,
seqlen: usize,
dim: usize,
dstate: usize,
delta_softplus: bool,
},
SsdChunkScan {
batch: usize,
seqlen: usize,
heads: usize,
head_dim: usize,
state_dim: usize,
chunk_size: usize,
},
Nf4Matmul {
batch: usize,
m: usize,
n: usize,
k: usize,
block_size: usize,
},
JitScalars {
scalars: Vec<f64>,
},
NonZeroIndices {
capacity: usize,
count_sym: fuel_ir::SymId,
},
}
pub type CostFn = fn(
&[fuel_ir::Shape],
&[DType],
&OpParams,
&fuel_ir::backend::BackendCapabilities,
) -> crate::fused::CostEstimate;
pub fn unknown_cost(
_shapes: &[fuel_ir::Shape],
_dtypes: &[DType],
_params: &OpParams,
_caps: &fuel_ir::backend::BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate::default()
}
#[derive(Clone, Copy, Debug)]
pub struct BindingEntry {
pub kernel: KernelRef,
pub caps: KernelCaps,
pub precision: crate::fused::PrecisionGuarantee,
pub cost: CostFn,
pub kernel_source: &'static str,
pub is_generic: bool,
pub kernel_revision_hash: u64,
pub cost_expr: Option<&'static crate::fkc::CompiledCostExpr>,
}
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BindingKey {
Static(OpKind),
RuntimeFused(FusedOpId),
}
impl From<OpKind> for BindingKey {
fn from(op: OpKind) -> Self {
BindingKey::Static(op)
}
}
impl BindingKey {
pub fn static_op(self) -> Option<OpKind> {
match self {
BindingKey::Static(op) => Some(op),
BindingKey::RuntimeFused(_) => None,
}
}
}
#[derive(Default, Clone)]
pub struct KernelBindingTable {
bindings: HashMap<(BindingKey, KernelDTypes, BackendId), SmallVec<[BindingEntry; 2]>>,
}
#[derive(Debug)]
pub enum RuntimeFusedLookup<'a> {
NotBound,
Unique(&'a [DType], &'a BindingEntry),
Ambiguous { signatures: Vec<Vec<DType>> },
}
impl KernelBindingTable {
pub fn new() -> Self {
Self {
bindings: HashMap::new(),
}
}
pub fn register(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
) {
self.register_full(
op,
dtypes,
backend,
kernel,
KernelCaps::empty(),
crate::fused::PrecisionGuarantee::UNAUDITED,
unknown_cost,
);
}
pub fn register_with_caps(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
caps: KernelCaps,
) {
self.register_full(
op,
dtypes,
backend,
kernel,
caps,
crate::fused::PrecisionGuarantee::UNAUDITED,
unknown_cost,
);
}
pub fn register_with_precision(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
precision: crate::fused::PrecisionGuarantee,
) {
self.register_full(
op,
dtypes,
backend,
kernel,
KernelCaps::empty(),
precision,
unknown_cost,
);
}
pub fn register_with_caps_and_precision(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
caps: KernelCaps,
precision: crate::fused::PrecisionGuarantee,
) {
self.register_full(op, dtypes, backend, kernel, caps, precision, unknown_cost);
}
pub fn register_full(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
caps: KernelCaps,
precision: crate::fused::PrecisionGuarantee,
cost: CostFn,
) {
self.register_full_with_source(op, dtypes, backend, kernel, caps, precision, cost, "");
}
pub fn register_full_with_source(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
caps: KernelCaps,
precision: crate::fused::PrecisionGuarantee,
cost: CostFn,
kernel_source: &'static str,
) {
self.register_full_with_source_generic(
op,
dtypes,
backend,
kernel,
caps,
precision,
cost,
kernel_source,
false,
crate::fused::KernelRevisionHash::UNTRACKED.0,
None,
);
}
#[allow(clippy::too_many_arguments)]
pub fn register_full_with_source_generic(
&mut self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
caps: KernelCaps,
precision: crate::fused::PrecisionGuarantee,
cost: CostFn,
kernel_source: &'static str,
is_generic: bool,
kernel_revision_hash: u64,
cost_expr: Option<&'static crate::fkc::CompiledCostExpr>,
) {
let key = (op.into(), SmallVec::from_slice(dtypes), backend);
let entry = BindingEntry {
kernel,
caps,
precision,
cost,
kernel_source,
is_generic,
kernel_revision_hash,
cost_expr,
};
self.bindings.entry(key).or_default().push(entry);
}
pub fn finalize(&self) -> Result<()> {
for ((op, dtypes, backend), alts) in self.bindings.iter() {
for i in 0..alts.len() {
let ptr_i = alts[i].kernel as *const () as usize;
for entry_j in &alts[i + 1..] {
if entry_j.kernel as *const () as usize == ptr_i {
return Err(Error::Msg(format!(
"KernelBindingTable: duplicate KernelRef registered \
for (op={op:?}, dtypes={dtypes:?}, backend={backend:?}). \
The same function pointer was registered twice; \
distinct alternatives at one decision point must be \
distinct `fn` items."
)));
}
}
}
}
Ok(())
}
pub fn fill_unset_cpu_precision(&mut self, default: crate::fused::PrecisionGuarantee) {
for ((_, _, backend), alts) in self.bindings.iter_mut() {
if *backend != BackendId::Cpu {
continue;
}
for entry in alts.iter_mut() {
let p = &mut entry.precision;
if p.makes_no_static_claim() && !p.is_audited_no_bound() {
*p = default;
}
}
}
}
pub fn fill_unset_cpu_cost(&mut self, dispatcher: fn(OpKind) -> CostFn) {
self.fill_unset_cost_for_backend(BackendId::Cpu, dispatcher);
}
pub fn fill_unset_cost_for_backend(
&mut self,
backend: BackendId,
dispatcher: fn(OpKind) -> CostFn,
) {
let sentinel = unknown_cost as *const () as usize;
for ((key, _, this_backend), alts) in self.bindings.iter_mut() {
if *this_backend != backend {
continue;
}
let BindingKey::Static(op) = key else {
continue;
};
for entry in alts.iter_mut() {
if (entry.cost as usize) == sentinel {
entry.cost = dispatcher(*op);
}
}
}
}
pub fn lookup_precision(
&self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
) -> crate::fused::PrecisionGuarantee {
let key = (op.into(), SmallVec::from_slice(dtypes), backend);
self.bindings
.get(&key)
.and_then(|alts| alts.first())
.map(|e| e.precision)
.unwrap_or(crate::fused::PrecisionGuarantee::UNAUDITED)
}
pub fn lookup_cost(
&self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
) -> CostFn {
let key = (op.into(), SmallVec::from_slice(dtypes), backend);
self.bindings
.get(&key)
.and_then(|alts| alts.first())
.map(|e| e.cost)
.unwrap_or(unknown_cost)
}
pub fn iter_precision(
&self,
) -> impl Iterator<
Item = (
OpKind,
&[DType],
BackendId,
crate::fused::PrecisionGuarantee,
),
> {
self.bindings
.iter()
.filter_map(|((key, dtypes, backend), alts)| {
key.static_op().map(|op| (op, dtypes, backend, alts))
})
.flat_map(|(op, dtypes, backend, alts)| {
alts.iter()
.map(move |e| (op, dtypes.as_slice(), *backend, e.precision))
})
}
pub fn iter_entries(
&self,
) -> impl Iterator<Item = (OpKind, &[DType], BackendId, &BindingEntry)> {
self.bindings
.iter()
.filter_map(|((key, dtypes, backend), alts)| {
key.static_op().map(|op| (op, dtypes, backend, alts))
})
.flat_map(|(op, dtypes, backend, alts)| {
alts.iter()
.map(move |e| (op, dtypes.as_slice(), *backend, e))
})
}
pub fn iter_cost(&self) -> impl Iterator<Item = (OpKind, &[DType], BackendId, CostFn)> {
self.bindings
.iter()
.filter_map(|((key, dtypes, backend), alts)| {
key.static_op().map(|op| (op, dtypes, backend, alts))
})
.flat_map(|(op, dtypes, backend, alts)| {
alts.iter()
.map(move |e| (op, dtypes.as_slice(), *backend, e.cost))
})
}
pub fn iter_keys(&self) -> impl Iterator<Item = (OpKind, &[DType], BackendId)> {
self.bindings.keys().filter_map(|(key, dtypes, backend)| {
key.static_op().map(|op| (op, dtypes.as_slice(), *backend))
})
}
pub fn has_runtime_fused(&self, fid: FusedOpId, backend: BackendId) -> bool {
self.bindings
.keys()
.any(|(k, _, b)| *b == backend && matches!(k, BindingKey::RuntimeFused(f) if *f == fid))
}
pub fn unique_runtime_fused(
&self,
fid: FusedOpId,
backend: BackendId,
) -> RuntimeFusedLookup<'_> {
let mut matches = self.bindings.iter().filter(|((k, _, b), _)| {
*b == backend && matches!(k, BindingKey::RuntimeFused(f) if *f == fid)
});
let Some(((_, dtypes, _), alts)) = matches.next() else {
return RuntimeFusedLookup::NotBound;
};
if matches.next().is_some() {
let mut signatures: Vec<Vec<DType>> = self
.bindings
.iter()
.filter(|((k, _, b), _)| {
*b == backend && matches!(k, BindingKey::RuntimeFused(f) if *f == fid)
})
.map(|((_, d, _), _)| d.to_vec())
.collect();
signatures.sort_by_key(|d| (d.len(), format!("{d:?}")));
return RuntimeFusedLookup::Ambiguous { signatures };
}
match alts.first() {
Some(e) => RuntimeFusedLookup::Unique(dtypes.as_slice(), e),
None => RuntimeFusedLookup::NotBound,
}
}
#[doc(hidden)]
pub fn remove_runtime_fused_for_tests(&mut self) {
self.bindings
.retain(|(k, _, _), _| matches!(k, BindingKey::Static(_)));
}
pub fn lookup(
&self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
) -> Result<KernelRef> {
self.lookup_with_caps(op, dtypes, backend).map(|(k, _)| k)
}
pub fn lookup_with_caps(
&self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
) -> Result<(KernelRef, KernelCaps)> {
let key = (op.into(), SmallVec::from_slice(dtypes), backend);
self.bindings
.get(&key)
.and_then(|alts| alts.iter().find(|e| !e.caps.requires_broadcast))
.map(|e| (e.kernel, e.caps))
.ok_or_else(|| {
let (available_backends, supported_combinations) = self.coverage_diagnostic();
Error::NoBackendForOp {
op: key.0.static_op().unwrap_or(OpKind::RuntimeFused),
dtypes: dtypes.to_vec(),
available_backends,
supported_combinations,
}
.bt()
})
}
pub fn coverage_diagnostic(&self) -> (Vec<BackendId>, Vec<CoverageBinding>) {
let available_backends: Vec<BackendId> = self
.bindings
.keys()
.map(|(_, _, b)| *b)
.collect::<std::collections::HashSet<_>>()
.into_iter()
.collect();
let supported_combinations: Vec<(BackendId, OpKind, Vec<DType>)> = self
.bindings
.keys()
.filter_map(|(k, d, b)| k.static_op().map(|o| (*b, o, d.to_vec())))
.collect();
(available_backends, supported_combinations)
}
pub fn backend_supports(
&self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
) -> bool {
let key = (op.into(), SmallVec::from_slice(dtypes), backend);
self.bindings
.get(&key)
.is_some_and(|alts| alts.iter().any(|e| !e.caps.requires_broadcast))
}
pub fn lookup_alternatives(
&self,
op: impl Into<BindingKey>,
dtypes: &[DType],
backend: BackendId,
) -> &[BindingEntry] {
let key = (op.into(), SmallVec::from_slice(dtypes), backend);
self.bindings
.get(&key)
.map(|alts| alts.as_slice())
.unwrap_or(&[])
}
pub fn len(&self) -> usize {
self.bindings.values().map(|alts| alts.len()).sum()
}
pub fn is_empty(&self) -> bool {
self.bindings.is_empty()
}
pub fn key_count(&self) -> usize {
self.bindings.len()
}
}
impl std::fmt::Debug for KernelBindingTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("KernelBindingTable")
.field("keys", &self.bindings.len())
.field("total_alternatives", &self.len())
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::DType;
#[test]
fn kernel_ref_stores_function_pointer() {
fn dummy_kernel(
_inputs: &[Arc<RwLock<Storage>>],
_outputs: &mut [Arc<RwLock<Storage>>],
_layouts: &[Layout],
_params: &OpParams,
) -> Result<()> {
Ok(())
}
let k: KernelRef = dummy_kernel;
let _k2 = k;
let _k3: KernelRef = k;
}
#[test]
fn op_params_variants_construct() {
let _ = OpParams::None;
let _ = OpParams::Reduce {
dims: vec![0, 1],
keepdim: false,
};
let _ = OpParams::Matmul {
lhs_batch_dims: vec![],
rhs_batch_dims: vec![],
m: 4,
n: 8,
k: 16,
m_compute: MatmulM::All,
};
let _ = OpParams::Slice {
dim: 0,
start: 0,
end: 10,
step: 1,
};
let _ = OpParams::Cast;
let _ = OpParams::Affine { mul: 2.0, add: 1.0 };
let _ = format!("{:?}", OpParams::None);
}
#[test]
fn kernel_ref_can_be_called() {
fn copy_kernel(
inputs: &[Arc<RwLock<Storage>>],
outputs: &mut [Arc<RwLock<Storage>>],
_layouts: &[Layout],
_params: &OpParams,
) -> Result<()> {
let in_arc = &inputs[0];
let out_arc = &outputs[0];
let in_guard = in_arc.read().unwrap();
let mut out_guard = out_arc.write().unwrap();
let in_bytes = match &in_guard.inner {
fuel_memory::BackendStorage::Cpu(s) => s.bytes(),
#[allow(unreachable_patterns)]
_ => return Ok(()),
};
if let fuel_memory::BackendStorage::Cpu(s) = &mut out_guard.inner {
s.bytes_mut().copy_from_slice(in_bytes);
}
Ok(())
}
let input = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let output = fuel_memory::alloc_cpu_zeroed(DType::F32, 4).unwrap();
let inputs = vec![Arc::new(RwLock::new(input))];
let mut outputs = vec![Arc::new(RwLock::new(output))];
let k: KernelRef = copy_kernel;
k(&inputs, &mut outputs, &[], &OpParams::None).unwrap();
let out_guard = outputs[0].read().unwrap();
if let fuel_memory::BackendStorage::Cpu(s) = &out_guard.inner {
let typed: &[f32] = s.as_slice().unwrap();
assert_eq!(typed, &[1.0, 2.0, 3.0, 4.0]);
}
}
fn ok_kernel(
_inputs: &[Arc<RwLock<Storage>>],
_outputs: &mut [Arc<RwLock<Storage>>],
_layouts: &[Layout],
_params: &OpParams,
) -> Result<()> {
Ok(())
}
fn ok_kernel_alt(
_inputs: &[Arc<RwLock<Storage>>],
_outputs: &mut [Arc<RwLock<Storage>>],
_layouts: &[Layout],
_params: &OpParams,
) -> Result<()> {
Ok(())
}
#[test]
fn step_9a_two_alternatives_append_then_first_wins_on_legacy_lookup() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::BF16, DType::BF16, DType::BF16];
table.register(OpKind::MatMul, &dts, BackendId::Cuda, ok_kernel);
table.register(OpKind::MatMul, &dts, BackendId::Cuda, ok_kernel_alt);
let (k1, _caps) = table
.lookup_with_caps(OpKind::MatMul, &dts, BackendId::Cuda)
.unwrap();
assert_eq!(k1 as *const () as usize, ok_kernel as *const () as usize);
let alts = table.lookup_alternatives(OpKind::MatMul, &dts, BackendId::Cuda);
assert_eq!(alts.len(), 2);
assert_eq!(
alts[0].kernel as *const () as usize,
ok_kernel as *const () as usize
);
assert_eq!(
alts[1].kernel as *const () as usize,
ok_kernel_alt as *const () as usize
);
assert_eq!(table.key_count(), 1, "one decision point");
assert_eq!(table.len(), 2, "two alternatives total");
}
#[test]
fn cost_expr_field_defaults_none_for_hand_written_registration() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
let alts = table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu);
assert_eq!(alts.len(), 1);
assert!(
alts[0].cost_expr.is_none(),
"hand-written registrations carry no declared cost AST"
);
}
#[test]
fn iter_entries_yields_full_entries_and_filters_runtime_fused() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel_alt);
table.register(
BindingKey::RuntimeFused(FusedOpId(FusedOpId::RUNTIME_FUSED_BASE)),
&dts,
BackendId::Cpu,
ok_kernel,
);
let entries: Vec<_> = table.iter_entries().collect();
assert_eq!(
entries.len(),
2,
"both AddElementwise alternatives, RuntimeFused excluded"
);
for (op, dtypes, backend, entry) in &entries {
assert_eq!(*op, OpKind::AddElementwise);
assert_eq!(*dtypes, dts.as_slice());
assert_eq!(*backend, BackendId::Cpu);
let _ = entry.caps;
let _ = entry.kernel_revision_hash;
}
let fns: Vec<usize> = entries
.iter()
.map(|(_, _, _, e)| e.kernel as *const () as usize)
.collect();
assert!(fns.contains(&(ok_kernel as *const () as usize)));
assert!(fns.contains(&(ok_kernel_alt as *const () as usize)));
}
#[test]
fn step_9a_duplicate_kernel_ref_detected_by_finalize_not_panic() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel);
table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel); let err = table
.finalize()
.expect_err("finalize must report the duplicate KernelRef");
assert!(
format!("{err}").contains("duplicate KernelRef"),
"unexpected error: {err}"
);
}
#[test]
fn finalize_ok_for_distinct_sibling_kernels() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel);
table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel_alt);
table
.finalize()
.expect("distinct sibling kernels must finalize OK");
}
#[test]
fn lookup_with_caps_skips_broadcast_baked_sibling_regardless_of_order() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register_with_caps(
OpKind::AddElementwise,
&dts,
BackendId::Cpu,
ok_kernel,
KernelCaps {
strided_input: true,
requires_broadcast: true,
},
);
table.register_with_caps(
OpKind::AddElementwise,
&dts,
BackendId::Cpu,
ok_kernel_alt,
KernelCaps {
strided_input: true,
requires_broadcast: false,
},
);
let (picked, caps) = table
.lookup_with_caps(OpKind::AddElementwise, &dts, BackendId::Cpu)
.expect("a dense-safe sibling exists");
assert!(
!caps.requires_broadcast,
"the realize pick must never be a broadcast-baked kernel",
);
assert_eq!(
picked as *const () as usize, ok_kernel_alt as *const () as usize,
"lookup_with_caps must skip the broadcast-baked sibling and return the generic \
one regardless of registration order",
);
let alts = table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu);
assert_eq!(
alts.len(),
2,
"both siblings registered; neither is dropped"
);
}
#[test]
fn lookup_with_caps_misses_when_only_broadcast_baked_siblings_exist() {
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register_with_caps(
OpKind::AddElementwise,
&dts,
BackendId::Cpu,
ok_kernel,
KernelCaps {
strided_input: true,
requires_broadcast: true,
},
);
assert!(
table
.lookup_with_caps(OpKind::AddElementwise, &dts, BackendId::Cpu)
.is_err(),
"no dense-safe sibling ⇒ honest miss, never a wrong read",
);
}
#[test]
fn step_9a_fill_passes_touch_every_alternative() {
use crate::fused::PrecisionGuarantee;
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel_alt);
for e in table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu) {
assert!(!e.precision.bit_stable_on_same_hardware);
}
let bit_stable = PrecisionGuarantee {
bit_stable_on_same_hardware: true,
max_ulp: Some(0),
max_relative: None,
max_absolute: None,
notes: "test default",
};
table.fill_unset_cpu_precision(bit_stable);
let alts = table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu);
assert_eq!(alts.len(), 2);
for e in alts {
assert!(e.precision.bit_stable_on_same_hardware);
assert_eq!(e.precision.max_ulp, Some(0));
}
}
#[test]
fn gap_095_fill_preserves_audited_none_but_still_upgrades_unaudited() {
use crate::fused::PrecisionGuarantee;
use fuel_ir::probe::BackendId;
let mut table = KernelBindingTable::new();
let dts = [DType::F32, DType::F32, DType::F32];
const REASON: &str = "audited: accumulation order is scheduler-dependent; \
no static bound applies";
table.register_with_precision(
OpKind::MatMul,
&dts,
BackendId::Cpu,
ok_kernel,
PrecisionGuarantee::none(REASON),
);
table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
table.fill_unset_cpu_precision(PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU);
let audited = table
.lookup_alternatives(OpKind::MatMul, &dts, BackendId::Cpu)
.first()
.expect("audited-none entry registered")
.precision;
assert!(
!audited.bit_stable_on_same_hardware,
"the fill asserted bit-stability on a kernel whose audit \
explicitly concluded that no static bound applies",
);
assert_eq!(
audited.notes, REASON,
"the audit reason was overwritten — the reasoning that lived at \
the registration site is gone",
);
let placeholder = table
.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu)
.first()
.expect("unaudited entry registered")
.precision;
assert!(
placeholder.bit_stable_on_same_hardware,
"control FAILED: the UNAUDITED placeholder was not upgraded, so \
this test would pass with the fill disabled entirely and proves \
nothing about the discrimination it is meant to check",
);
}
}