use fuel_graph::jit::PatternNode;
use fuel_graph::registry::{FusedOpId, FusedOps};
use fuel_ir::DType;
use fuel_ir::dispatch::OpKind;
use fuel_ir::probe::BackendId;
use crate::dispatch::{extend_global_bindings, global_bindings};
use crate::fused::{PrecisionGuarantee, default_kernel_registry};
use crate::kernel::{BindingKey, KernelBindingTable, KernelCaps, KernelRef};
#[derive(Clone, Debug)]
pub struct RuntimeKernelBinding {
pub kernel: KernelRef,
pub caps: KernelCaps,
pub precision: PrecisionGuarantee,
pub dtypes: Vec<DType>,
}
pub fn register_runtime_kernel(
id: FusedOpId,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
) {
extend_global_bindings(|t| t.register(BindingKey::RuntimeFused(id), dtypes, backend, kernel));
}
#[doc(hidden)]
pub fn register_runtime_kernel_into(
table: &mut KernelBindingTable,
id: FusedOpId,
dtypes: &[DType],
backend: BackendId,
kernel: KernelRef,
) {
table.register(BindingKey::RuntimeFused(id), dtypes, backend, kernel);
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FusedLookupMiss {
NotBound,
Ambiguous { signatures: Vec<Vec<DType>> },
}
pub fn lookup_runtime_kernel(
id: FusedOpId,
backend: BackendId,
) -> Result<RuntimeKernelBinding, FusedLookupMiss> {
use crate::kernel::RuntimeFusedLookup;
let table = global_bindings();
match table.unique_runtime_fused(id, backend) {
RuntimeFusedLookup::Unique(dtypes, e) => Ok(RuntimeKernelBinding {
kernel: e.kernel,
caps: e.caps,
precision: e.precision,
dtypes: dtypes.to_vec(),
}),
RuntimeFusedLookup::NotBound => Err(FusedLookupMiss::NotBound),
RuntimeFusedLookup::Ambiguous { signatures } => {
Err(FusedLookupMiss::Ambiguous { signatures })
}
}
}
pub fn fused_kernel_available(id: FusedOpId, backend: BackendId) -> bool {
fused_kernel_available_core(
id,
backend,
|| global_bindings().has_runtime_fused(id, backend),
|| static_binding_table_bridge(id, backend),
)
}
pub fn fused_kernel_available_in(
table: &KernelBindingTable,
id: FusedOpId,
backend: BackendId,
) -> bool {
fused_kernel_available_core(
id,
backend,
|| table.has_runtime_fused(id, backend),
|| static_binding_table_bridge_in(table, id, backend),
)
}
fn fused_kernel_available_core(
id: FusedOpId,
backend: BackendId,
has_runtime_fused: impl FnOnce() -> bool,
bridge: impl FnOnce() -> bool,
) -> bool {
default_kernel_registry().lookup(id, backend).is_some() || has_runtime_fused() || bridge()
}
fn static_binding_table_bridge(id: FusedOpId, backend: BackendId) -> bool {
static_binding_table_bridge_over(&global_bindings(), id, backend)
}
fn static_binding_table_bridge_in(
table: &KernelBindingTable,
id: FusedOpId,
backend: BackendId,
) -> bool {
static_binding_table_bridge_over(table, id, backend)
}
fn static_binding_table_bridge_over(
table: &KernelBindingTable,
id: FusedOpId,
backend: BackendId,
) -> bool {
if id.is_runtime() {
return false;
}
let Some(op_kind) = static_fused_id_to_binding_table_op_kind(id) else {
return false;
};
table
.iter_keys()
.any(|(op, _dtypes, b)| op == op_kind && b == backend)
}
pub(crate) fn static_fused_id_to_binding_table_op_kind(id: FusedOpId) -> Option<OpKind> {
if id == FusedOps::FLASH_ATTN {
Some(OpKind::FlashAttn)
} else {
None
}
}
#[doc(hidden)]
pub fn clear_runtime_fused_for_tests() {
extend_global_bindings(|t| t.remove_runtime_fused_for_tests());
fuel_graph::runtime_fused::clear_runtime_fused_for_tests();
}
pub fn adopt_runtime_fused(
name: impl Into<String>,
region: PatternNode,
kernel: KernelRef,
dtypes: Vec<DType>,
backend: BackendId,
) -> Option<FusedOpId> {
let id = fuel_graph::runtime_fused::register_runtime_fused(name, region).ok()?;
extend_global_bindings(|table| {
let already_bound = table
.lookup_alternatives(BindingKey::RuntimeFused(id), &dtypes, backend)
.iter()
.any(|entry| std::ptr::fn_addr_eq(entry.kernel, kernel));
if !already_bound {
table.register(BindingKey::RuntimeFused(id), &dtypes, backend, kernel);
}
});
Some(id)
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_graph::jit::{OpAttrs, OpTag};
use fuel_ir::Layout;
use std::sync::{Arc, RwLock as StdRwLock};
fn other_noop_kernel(
_inputs: &[Arc<StdRwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<StdRwLock<fuel_memory::Storage>>],
_layouts: &[Layout],
_params: &crate::kernel::OpParams,
) -> fuel_ir::Result<()> {
Ok(())
}
fn noop_kernel(
_inputs: &[Arc<StdRwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<StdRwLock<fuel_memory::Storage>>],
_layouts: &[Layout],
_params: &crate::kernel::OpParams,
) -> fuel_ir::Result<()> {
Ok(())
}
fn relu_add() -> PatternNode {
PatternNode::Op {
op: OpTag::Relu,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::Add,
attrs: OpAttrs::default(),
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 1 },
],
}],
}
}
fn sqrt_div() -> PatternNode {
PatternNode::Op {
op: OpTag::Sqrt,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::Div,
attrs: OpAttrs::default(),
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 1 },
],
}],
}
}
#[test]
fn adopt_registers_op_plus_kernel_and_the_gate_predicate_sees_it() {
let id = adopt_runtime_fused(
"test::adopt::relu_add",
relu_add(),
noop_kernel as KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("relu(add) is a registrable region");
assert!(id.is_runtime(), "allocated a runtime id");
let row = lookup_runtime_kernel(id, BackendId::Cpu).expect("row bound on Cpu");
assert_eq!(
row.dtypes,
vec![DType::F32, DType::F32, DType::F32],
"the binding key carries the per-operand dtype tuple",
);
assert!(
fused_kernel_available(id, BackendId::Cpu),
"the capability predicate sees the adopted op on Cpu",
);
assert!(
!fused_kernel_available(id, BackendId::Cuda),
"but not on a backend the kernel wasn't adopted for",
);
assert!(
fuel_graph::runtime_fused::runtime_region(id).is_some(),
"the recipe (region) is registered in the fuel-graph metadata sidecar",
);
}
#[test]
fn adopted_kernel_lives_in_the_one_binding_table() {
let id = adopt_runtime_fused(
"test::adopt::sqrt_div::one_registry",
sqrt_div(),
noop_kernel as KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("registrable region");
let (kernel, _caps) = global_bindings()
.lookup_with_caps(
BindingKey::RuntimeFused(id),
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("resolves through the ordinary binding lookup");
assert!(
std::ptr::fn_addr_eq(kernel, noop_kernel as KernelRef),
"the exact adopted kernel, from the one registry",
);
assert!(
global_bindings()
.lookup_with_caps(
BindingKey::RuntimeFused(id),
&[DType::F64, DType::F64, DType::F64],
BackendId::Cpu,
)
.is_err(),
"wrong dtype tuple ⇒ NoBackendForOp, not a wrong-kernel bind",
);
}
#[test]
fn one_recipe_two_signatures_refuses_to_guess() {
let mut table = KernelBindingTable::new();
let id = FusedOpId(FusedOpId::RUNTIME_FUSED_BASE);
register_runtime_kernel_into(
&mut table,
id,
&[DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel as KernelRef,
);
register_runtime_kernel_into(
&mut table,
id,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
other_noop_kernel as KernelRef,
);
match table.unique_runtime_fused(id, BackendId::Cpu) {
crate::kernel::RuntimeFusedLookup::Ambiguous { signatures } => {
assert_eq!(
signatures.len(),
2,
"both signatures reported: {signatures:?}"
);
assert!(
signatures.iter().any(|s| s.len() == 2)
&& signatures.iter().any(|s| s.len() == 3),
"the report must name the competing ARITIES, got {signatures:?}",
);
}
other => {
panic!("an id carrying two signatures must refuse, not guess — got {other:?}",)
}
}
}
#[test]
fn one_signature_still_resolves() {
let mut table = KernelBindingTable::new();
let id = FusedOpId(FusedOpId::RUNTIME_FUSED_BASE);
register_runtime_kernel_into(
&mut table,
id,
&[DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel as KernelRef,
);
match table.unique_runtime_fused(id, BackendId::Cpu) {
crate::kernel::RuntimeFusedLookup::Unique(dtypes, _) => {
assert_eq!(dtypes, &[DType::F32, DType::F32]);
}
other => panic!("one signature must resolve — got {other:?}"),
}
}
#[test]
fn unbound_id_reports_not_bound_not_ambiguous() {
let table = KernelBindingTable::new();
let id = FusedOpId(FusedOpId::RUNTIME_FUSED_BASE);
assert!(matches!(
table.unique_runtime_fused(id, BackendId::Cpu),
crate::kernel::RuntimeFusedLookup::NotBound
));
}
#[test]
fn adopt_runtime_fused_is_idempotent_for_an_identical_region_and_kernel() {
let id1 = adopt_runtime_fused(
"test::adopt::idempotent::sqrt_div",
sqrt_div(),
noop_kernel as KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("sqrt(div) is a registrable region");
let id2 = adopt_runtime_fused(
"test::adopt::idempotent::sqrt_div::again",
sqrt_div(),
noop_kernel as KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("re-adopting the same region is still Some");
assert_eq!(
id1, id2,
"structurally-identical regions dedup to the same FusedOpId"
);
assert!(
fused_kernel_available(id1, BackendId::Cpu),
"the kernel is still available after the idempotent re-adopt",
);
let row = lookup_runtime_kernel(id1, BackendId::Cpu).expect("row still bound on Cpu");
assert!(
std::ptr::fn_addr_eq(row.kernel, noop_kernel as KernelRef),
"the (single, non-duplicated) kernel resolves correctly",
);
}
#[test]
#[cfg(feature = "cuda")]
fn fused_kernel_available_bridges_static_flash_attn_from_binding_table_on_cuda() {
let _guard = crate::dispatch::global_bindings();
drop(_guard);
assert!(
fused_kernel_available(FusedOps::FLASH_ATTN, BackendId::Cuda),
"the CUDA flash_decoding binding registers only into KernelBindingTable; \
fused_kernel_available must bridge to it for the static FLASH_ATTN id",
);
}
#[test]
#[cfg(not(feature = "cuda"))]
fn fused_kernel_available_flash_attn_false_on_cuda_without_cuda_registration() {
assert!(
!fused_kernel_available(FusedOps::FLASH_ATTN, BackendId::Cuda),
"no cuda registration ran in this build; the static bridge finds nothing",
);
}
#[test]
fn fused_kernel_available_in_reads_the_passed_table_only() {
let mut table = KernelBindingTable::new();
let id = fuel_graph::runtime_fused::register_runtime_fused(
"test::fused_kernel_available_in::relu_add",
relu_add(),
)
.expect("relu(add) is a registrable region");
register_runtime_kernel_into(
&mut table,
id,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel as KernelRef,
);
assert!(
fused_kernel_available_in(&table, id, BackendId::Cuda),
"the row bound into THIS table is visible",
);
assert!(
!fused_kernel_available_in(&KernelBindingTable::new(), id, BackendId::Cuda),
"an unrelated empty table sees nothing — no fallback to the global table",
);
}
#[test]
fn fused_kernel_available_runtime_id_still_routes_through_runtime_rows_untouched() {
let id = adopt_runtime_fused(
"test::adopt::flash_bridge_guard::relu_add",
relu_add(),
noop_kernel as KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("relu(add) is a registrable region");
assert!(id.is_runtime(), "allocated a runtime id");
assert!(
fused_kernel_available(id, BackendId::Cpu),
"a runtime id still resolves via its RuntimeFused binding rows",
);
assert!(
!fused_kernel_available(id, BackendId::Cuda),
"not adopted on Cuda, and the static-id bridge must never rescue a \
runtime id (it returns false immediately for any id.is_runtime())",
);
}
}