fuel_dispatch/runtime_fused_kernels.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Runtime fused-op kernel registration — the **adopt/gate facade** over the
3//! ONE binding registry.
4//!
5//! The transitional `FusedOpId`-keyed kernel *sidecar* that used to live here
6//! is **folded** (the 2026-07-08 decisions-log end-state): a runtime kernel now
7//! registers into the global binding table
8//! ([`crate::dispatch::extend_global_bindings`]) under
9//! [`BindingKey::RuntimeFused`], exactly like a startup kernel registers under
10//! `BindingKey::Static` — same table, same entry shape, same dtype-keyed
11//! dispatch lookup. This module keeps the runtime-op-shaped API
12//! ([`adopt_runtime_fused`], [`fused_kernel_available`],
13//! [`lookup_runtime_kernel`]) so callers stay id-oriented, but there is no
14//! second registry behind it.
15//!
16//! Cost note: a runtime row is stored with the binding-native `unknown_cost`
17//! sentinel (never a lying zero). Deriving its Layer-1 estimate from the op's
18//! recipe (`decompose_region`, keyed by the `RuntimeFused` id in the binding
19//! key — no stored closure needed) is the plan-path pricing follow-up; today a
20//! runtime arm is sparse-skip unpriced there, which is safe (arm 0 stays the
21//! runnability fallback).
22
23use fuel_graph::jit::PatternNode;
24use fuel_graph::registry::{FusedOpId, FusedOps};
25use fuel_ir::DType;
26use fuel_ir::dispatch::OpKind;
27use fuel_ir::probe::BackendId;
28
29use crate::dispatch::{extend_global_bindings, global_bindings};
30use crate::fused::{PrecisionGuarantee, default_kernel_registry};
31use crate::kernel::{BindingKey, KernelBindingTable, KernelCaps, KernelRef};
32
33/// A runtime-fused binding row, as [`lookup_runtime_kernel`] returns it — a
34/// diagnostic/test view (dispatch resolves dtype-precisely through the binding
35/// table itself; see `compile_one`'s `is_runtime` arm).
36#[derive(Clone, Debug)]
37pub struct RuntimeKernelBinding {
38 pub kernel: KernelRef,
39 pub caps: KernelCaps,
40 pub precision: PrecisionGuarantee,
41 /// The per-operand dtype tuple the kernel registered for (inputs in
42 /// order, then output) — the binding key's dtype half.
43 pub dtypes: Vec<DType>,
44}
45
46/// Bind a kernel for a runtime fused op (id `>= RUNTIME_FUSED_BASE`) into the
47/// global binding table under `BindingKey::RuntimeFused(id)`. `dtypes` is the
48/// per-operand tuple (inputs in order, then output) — the same key shape every
49/// static registration uses.
50pub fn register_runtime_kernel(
51 id: FusedOpId,
52 dtypes: &[DType],
53 backend: BackendId,
54 kernel: KernelRef,
55) {
56 extend_global_bindings(|t| t.register(BindingKey::RuntimeFused(id), dtypes, backend, kernel));
57}
58
59/// **TEST-ONLY.** Table-scoped sibling of [`register_runtime_kernel`] —
60/// binds into a caller-owned `table` instead of the global binding table, so
61/// a test can construct a `KernelBindingTable` carrying a runtime-fused row
62/// without touching (or being affected by) the global sidecar every other
63/// test in this binary shares. `#[doc(hidden)] pub`, not `#[cfg(test)]`, for
64/// the same reason as [`clear_runtime_fused_for_tests`]: integration tests
65/// compile this crate without `cfg(test)`.
66#[doc(hidden)]
67pub fn register_runtime_kernel_into(
68 table: &mut KernelBindingTable,
69 id: FusedOpId,
70 dtypes: &[DType],
71 backend: BackendId,
72 kernel: KernelRef,
73) {
74 table.register(BindingKey::RuntimeFused(id), dtypes, backend, kernel);
75}
76
77/// Why a runtime-fused lookup returned no kernel. **Two distinct facts, kept
78/// distinct:** an `Option` would make them indistinguishable, and they have
79/// different fixes — *not bound* means register one, *ambiguous* means say
80/// which signature you want.
81#[derive(Clone, Debug, PartialEq, Eq)]
82pub enum FusedLookupMiss {
83 /// Nothing is registered for this `(id, backend)`.
84 NotBound,
85 /// Several dtype tuples share this id, so no single kernel answers the
86 /// question. A [`FusedOpId`] names a RECIPE, not a signature.
87 Ambiguous { signatures: Vec<Vec<DType>> },
88}
89
90/// The runtime-fused binding for `(id, backend)` when exactly one dtype tuple
91/// is registered.
92///
93/// **Refuses to guess.** One `FusedOpId` can legitimately carry several dtype
94/// tuples (`register_runtime_fused` dedups on the recipe's base-map hash), and
95/// resolving by id alone would pick between them by `HashMap` order — a coin
96/// flip per process, capable of returning a kernel of the wrong operand arity.
97/// This is the diagnostic/test path; production dispatch takes the dtype tuple
98/// and does an exact keyed lookup (`lookup_with_caps`). See GAP-213.
99pub fn lookup_runtime_kernel(
100 id: FusedOpId,
101 backend: BackendId,
102) -> Result<RuntimeKernelBinding, FusedLookupMiss> {
103 use crate::kernel::RuntimeFusedLookup;
104 let table = global_bindings();
105 match table.unique_runtime_fused(id, backend) {
106 RuntimeFusedLookup::Unique(dtypes, e) => Ok(RuntimeKernelBinding {
107 kernel: e.kernel,
108 caps: e.caps,
109 precision: e.precision,
110 dtypes: dtypes.to_vec(),
111 }),
112 RuntimeFusedLookup::NotBound => Err(FusedLookupMiss::NotBound),
113 RuntimeFusedLookup::Ambiguous { signatures } => {
114 Err(FusedLookupMiss::Ambiguous { signatures })
115 }
116 }
117}
118
119/// The capability predicate the optimizer's gate consults: is there an
120/// admissible kernel for `(id, backend)` — static **or** runtime-registered?
121/// The dispatch layer passes `|id| fused_kernel_available(id, backend)` to
122/// `capability_gated_rules`, so a runtime op fuses only once its kernel is
123/// bound. Id-level (coarse) by design; the dtype-precise check is the
124/// dispatch-time binding lookup itself.
125///
126/// **SHARED-CONSUMER COORDINATION (dd-shapes / flash-arm registry seam,
127/// 2026-07-08; carried through the binding-key fold).** Two independent
128/// callers with different invariants, and any change here must hold both:
129/// - the **Tier-2 runtime-fusion / JIT program** calls this for *runtime* ids
130/// (`id.is_runtime()`, allocated by [`adopt_runtime_fused`]) — those resolve
131/// only through their `BindingKey::RuntimeFused` rows
132/// (`KernelBindingTable::has_runtime_fused`), never the static registry and
133/// never the static-id bridge below (which returns `false` immediately for
134/// a runtime id);
135/// - the **decode flash-arm gate**
136/// (`crate::decode_flash::FlashArmCapability::production`) calls this for
137/// the *static* id `FusedOps::FLASH_ATTN` on `BackendId::Cuda`, whose CUDA
138/// `flash_decoding` kernel registers ONLY into the binding table under
139/// `(OpKind::FlashAttn, [f16|bf16;4], Cuda)` — never
140/// [`default_kernel_registry`] (a frozen CPU-only `OnceLock`). The
141/// `static_binding_table_bridge` arm makes that reachable: for a static id
142/// with a known `OpKind` mapping, consult the binding table for ANY entry on
143/// `backend` (dtype-blind — this predicate answers "is *a* kernel bound";
144/// dtype admissibility is a separate gate, e.g.
145/// `crate::decode_flash::flash_decode_admissible`).
146///
147/// **LOCK DISCIPLINE (post binding-key fold, 2026-07-08; nesting eliminated
148/// 2026-07-13 — Spec B Task 1).** This re-acquires the `global_bindings()`
149/// READ lock (once per `||` operand it actually evaluates — the checks are
150/// independent, short-lived acquisitions, not one held across the whole
151/// chain). One constraint remains:
152/// - Never call it — nor `adopt_runtime_fused` / `clear_runtime_fused_for_tests`,
153/// which take the WRITE lock — while holding a `global_bindings()` guard on the
154/// same thread. Same-thread read-then-write (or write-then-anything)
155/// self-deadlocks. (Tests scope their read guards around each optimize call and
156/// never hold one across an adopt/reset; production adopts on the G7 *background*
157/// thread, never the realize thread that holds read guards.)
158///
159/// The nested-read hazard this doc used to warn about — the optimizer's
160/// runtime-fusion pathfinder calling this while `optimize_graph` already
161/// holds a *caller-passed* `&KernelBindingTable` derived from a
162/// `global_bindings()` read guard — is now **eliminated**: the pathfinder
163/// calls [`fused_kernel_available_in`] against that already-held table
164/// (threaded in via `OptimizationContext::bindings`) instead of calling this
165/// function and re-acquiring the lock. This function stays the *global*
166/// entry point for callers with no table already in hand (the decode
167/// flash-arm gate, `jit_adopt`'s post-adopt check).
168pub fn fused_kernel_available(id: FusedOpId, backend: BackendId) -> bool {
169 fused_kernel_available_core(
170 id,
171 backend,
172 || global_bindings().has_runtime_fused(id, backend),
173 || static_binding_table_bridge(id, backend),
174 )
175}
176
177/// Table-passing variant of [`fused_kernel_available`] — reads the CALLER'S
178/// binding table (the one `optimize_graph` already holds, threaded via
179/// `OptimizationContext::bindings`) instead of re-acquiring
180/// `global_bindings()`. This is the **non-nesting** gate: a background
181/// adopt's `global_bindings()` WRITE can never queue between two reads this
182/// function takes, because it takes none — `table` is a plain reference the
183/// caller already resolved.
184pub fn fused_kernel_available_in(
185 table: &KernelBindingTable,
186 id: FusedOpId,
187 backend: BackendId,
188) -> bool {
189 fused_kernel_available_core(
190 id,
191 backend,
192 || table.has_runtime_fused(id, backend),
193 || static_binding_table_bridge_in(table, id, backend),
194 )
195}
196
197/// The shared 3-way OR behind [`fused_kernel_available`] and
198/// [`fused_kernel_available_in`]: the (lock-free) default kernel registry,
199/// then a caller-supplied runtime-fused check, then a caller-supplied
200/// static-id bridge check — the latter two as thunks so each public wrapper
201/// controls exactly how (and whether) it touches `global_bindings()`, rather
202/// than this shared core baking in one lock policy for both.
203fn fused_kernel_available_core(
204 id: FusedOpId,
205 backend: BackendId,
206 has_runtime_fused: impl FnOnce() -> bool,
207 bridge: impl FnOnce() -> bool,
208) -> bool {
209 default_kernel_registry().lookup(id, backend).is_some() || has_runtime_fused() || bridge()
210}
211
212/// The static-id half of the [`fused_kernel_available`] bridge (see its doc
213/// comment). Returns `false` immediately for a runtime id — purely additive
214/// for the static-id path; a runtime id resolves only via its
215/// `BindingKey::RuntimeFused` rows. Acquires `global_bindings()` once, for
216/// the duration of the [`static_binding_table_bridge_over`] scan — unchanged
217/// from before this function existed.
218fn static_binding_table_bridge(id: FusedOpId, backend: BackendId) -> bool {
219 static_binding_table_bridge_over(&global_bindings(), id, backend)
220}
221
222/// Table-passing sibling of [`static_binding_table_bridge`] — scans the
223/// caller's `table` instead of `global_bindings()`. See
224/// [`fused_kernel_available_in`].
225fn static_binding_table_bridge_in(
226 table: &KernelBindingTable,
227 id: FusedOpId,
228 backend: BackendId,
229) -> bool {
230 static_binding_table_bridge_over(table, id, backend)
231}
232
233/// The scan shared by [`static_binding_table_bridge`] and
234/// [`static_binding_table_bridge_in`] — the only difference between the two
235/// is which `KernelBindingTable` reference they hand in here.
236fn static_binding_table_bridge_over(
237 table: &KernelBindingTable,
238 id: FusedOpId,
239 backend: BackendId,
240) -> bool {
241 if id.is_runtime() {
242 return false;
243 }
244 let Some(op_kind) = static_fused_id_to_binding_table_op_kind(id) else {
245 return false;
246 };
247 // `iter_keys` is the static-entry view (RuntimeFused rows filtered), which
248 // is exactly the population this bridge should scan.
249 table
250 .iter_keys()
251 .any(|(op, _dtypes, b)| op == op_kind && b == backend)
252}
253
254/// The (deliberately small) set of *static* [`FusedOpId`]s whose kernel is
255/// known to register into the primitive `KernelBindingTable` under a
256/// corresponding [`OpKind`] rather than (only) [`default_kernel_registry`].
257/// See [`fused_kernel_available`]'s doc comment — widen this one id at a
258/// time as each is verified, not speculatively.
259pub(crate) fn static_fused_id_to_binding_table_op_kind(id: FusedOpId) -> Option<OpKind> {
260 if id == FusedOps::FLASH_ATTN {
261 Some(OpKind::FlashAttn)
262 } else {
263 None
264 }
265}
266
267/// **TEST-ONLY.** Reset the runtime-fused world: drop every
268/// `BindingKey::RuntimeFused` row from the global binding table AND clear the
269/// `fuel-graph` metadata sidecar — together, because clearing metadata restarts
270/// the id allocator and a reused id must never resolve a stale kernel. Bindings
271/// are cleared FIRST so no window exists where a fresh id sees an old row.
272/// Callers in one test binary share the process: serialize with every other
273/// adopting test (a bare reset races — dd-shapes coordination, 2026-07-08).
274/// `#[doc(hidden)] pub`, not `#[cfg(test)]`: integration tests compile this
275/// crate without `cfg(test)`.
276#[doc(hidden)]
277pub fn clear_runtime_fused_for_tests() {
278 extend_global_bindings(|t| t.remove_runtime_fused_for_tests());
279 fuel_graph::runtime_fused::clear_runtime_fused_for_tests();
280}
281
282/// Adopt a synthesized/imported runtime fused op: register its recipe (the
283/// `region`) in the `fuel-graph` metadata sidecar **and** bind its kernel in
284/// the global binding table, then return the freshly-allocated runtime
285/// [`FusedOpId`]. After this the capability gate sees the op as fusable on
286/// `backend`, and the executor's `is_runtime` arm resolves it dtype-precisely
287/// from the same table every static kernel lives in.
288///
289/// Takes the *resolved* parts (the region + the bound `KernelRef`), not a
290/// `JitResponse` — the `JitResponse`/`SynthArtifact` destructuring (+ the
291/// artifact-load step) happens at the seam-call site (`jit_adopt`), so the
292/// core dispatch layer stays free of the `fuel-kernel-seam` envelope crate.
293/// Returns `None` if the region is not registrable (non-decomposable / bad
294/// binds — surfaced by `register_runtime_fused`).
295///
296/// **Idempotent under region dedup (2026-07-15).** `register_runtime_fused`
297/// dedups structurally-identical regions to the SAME [`FusedOpId`] — so a
298/// second `adopt_runtime_fused` call for an already-adopted recipe returns
299/// the same `id` here too. Binding, however, is append-only
300/// ([`KernelBindingTable::register_full_with_source_generic`]): pushing the
301/// *same* `KernelRef` onto the *same* `(id, dtypes, backend)` key a second
302/// time would create a duplicate-identical sibling, which
303/// `KernelBindingTable::finalize` rejects (`Error::Msg`) — and the caller of
304/// `finalize` (`extend_global_bindings`'s `.expect()`) turns that `Err` into
305/// a panic *while holding the global bindings write lock*, poisoning it for
306/// every later caller.
307///
308/// Guard against this exactly at the idempotent-vs-genuinely-new boundary:
309/// before binding, check whether THIS `kernel` is already bound at `(id,
310/// dtypes, backend)` — the precise, dtypes-aware
311/// [`KernelBindingTable::lookup_alternatives`] view — and skip the bind if
312/// so. The check-then-bind runs inside a SINGLE [`extend_global_bindings`]
313/// write-lock critical section (not a separate read-lock check followed by
314/// a separate write-lock bind): two threads racing to adopt the same
315/// structurally-identical region + kernel + dtypes + backend (region dedup
316/// means unrelated call sites can land on the same `id`) would otherwise
317/// both observe "not yet bound" through a read-then-write race window and
318/// both bind, reproducing the duplicate-`KernelRef` panic anyway. A fresh
319/// id, or the same id with a different kernel / dtypes / backend, still
320/// binds exactly as before.
321pub fn adopt_runtime_fused(
322 name: impl Into<String>,
323 region: PatternNode,
324 kernel: KernelRef,
325 dtypes: Vec<DType>,
326 backend: BackendId,
327) -> Option<FusedOpId> {
328 let id = fuel_graph::runtime_fused::register_runtime_fused(name, region).ok()?;
329 extend_global_bindings(|table| {
330 let already_bound = table
331 .lookup_alternatives(BindingKey::RuntimeFused(id), &dtypes, backend)
332 .iter()
333 .any(|entry| std::ptr::fn_addr_eq(entry.kernel, kernel));
334 if !already_bound {
335 table.register(BindingKey::RuntimeFused(id), &dtypes, backend, kernel);
336 }
337 });
338 Some(id)
339}
340
341#[cfg(test)]
342mod tests {
343 use super::*;
344 use fuel_graph::jit::{OpAttrs, OpTag};
345 use fuel_ir::Layout;
346 use std::sync::{Arc, RwLock as StdRwLock};
347
348 /// A second, distinct `KernelRef` — `adopt_runtime_fused` skips
349 /// registration when the same fn pointer is already bound, so the two
350 /// signatures in the GAP-213 test must not share one.
351 fn other_noop_kernel(
352 _inputs: &[Arc<StdRwLock<fuel_memory::Storage>>],
353 _outputs: &mut [Arc<StdRwLock<fuel_memory::Storage>>],
354 _layouts: &[Layout],
355 _params: &crate::kernel::OpParams,
356 ) -> fuel_ir::Result<()> {
357 Ok(())
358 }
359
360 fn noop_kernel(
361 _inputs: &[Arc<StdRwLock<fuel_memory::Storage>>],
362 _outputs: &mut [Arc<StdRwLock<fuel_memory::Storage>>],
363 _layouts: &[Layout],
364 _params: &crate::kernel::OpParams,
365 ) -> fuel_ir::Result<()> {
366 Ok(())
367 }
368
369 fn relu_add() -> PatternNode {
370 PatternNode::Op {
371 op: OpTag::Relu,
372 attrs: OpAttrs::default(),
373 operands: vec![PatternNode::Op {
374 op: OpTag::Add,
375 attrs: OpAttrs::default(),
376 operands: vec![
377 PatternNode::Bind { index: 0 },
378 PatternNode::Bind { index: 1 },
379 ],
380 }],
381 }
382 }
383
384 /// sqrt(div(a, b)) — deliberately NOT the `relu(add(a, b))` shape every
385 /// other adopted-op test in this crate uses (`fused_cost`, `jit_adopt`,
386 /// `runtime_fused_arm`, and this module's own `relu_add` tests above).
387 /// `register_runtime_fused`'s dedup index is a process-global sidecar
388 /// shared by every `#[test]` in this binary (see
389 /// `runtime_fused_pathfinder`'s `tanh_mul_region` doc comment for the
390 /// full collision rationale): a test asserting on the EXACT kernel
391 /// pointer `lookup_with_caps`/`first_runtime_fused` resolves needs a
392 /// region no other adopting test also registers, or whichever test's
393 /// kernel reaches the shared binding-table key first wins the "first
394 /// alternative" — nondeterministically, by thread scheduling.
395 fn sqrt_div() -> PatternNode {
396 PatternNode::Op {
397 op: OpTag::Sqrt,
398 attrs: OpAttrs::default(),
399 operands: vec![PatternNode::Op {
400 op: OpTag::Div,
401 attrs: OpAttrs::default(),
402 operands: vec![
403 PatternNode::Bind { index: 0 },
404 PatternNode::Bind { index: 1 },
405 ],
406 }],
407 }
408 }
409
410 #[test]
411 fn adopt_registers_op_plus_kernel_and_the_gate_predicate_sees_it() {
412 let id = adopt_runtime_fused(
413 "test::adopt::relu_add",
414 relu_add(),
415 noop_kernel as KernelRef,
416 vec![DType::F32, DType::F32, DType::F32],
417 BackendId::Cpu,
418 )
419 .expect("relu(add) is a registrable region");
420
421 assert!(id.is_runtime(), "allocated a runtime id");
422 let row = lookup_runtime_kernel(id, BackendId::Cpu).expect("row bound on Cpu");
423 assert_eq!(
424 row.dtypes,
425 vec![DType::F32, DType::F32, DType::F32],
426 "the binding key carries the per-operand dtype tuple",
427 );
428 assert!(
429 fused_kernel_available(id, BackendId::Cpu),
430 "the capability predicate sees the adopted op on Cpu",
431 );
432 assert!(
433 !fused_kernel_available(id, BackendId::Cuda),
434 "but not on a backend the kernel wasn't adopted for",
435 );
436 assert!(
437 fuel_graph::runtime_fused::runtime_region(id).is_some(),
438 "the recipe (region) is registered in the fuel-graph metadata sidecar",
439 );
440 }
441
442 /// The fold invariant: a runtime kernel lives in the SAME global binding
443 /// table as static kernels — resolvable through the ordinary dtype-keyed
444 /// `lookup_with_caps` under its `RuntimeFused` key, with the exact adopted
445 /// kernel pointer. (No second registry behind the facade.)
446 #[test]
447 fn adopted_kernel_lives_in_the_one_binding_table() {
448 let id = adopt_runtime_fused(
449 "test::adopt::sqrt_div::one_registry",
450 sqrt_div(),
451 noop_kernel as KernelRef,
452 vec![DType::F32, DType::F32, DType::F32],
453 BackendId::Cpu,
454 )
455 .expect("registrable region");
456 let (kernel, _caps) = global_bindings()
457 .lookup_with_caps(
458 BindingKey::RuntimeFused(id),
459 &[DType::F32, DType::F32, DType::F32],
460 BackendId::Cpu,
461 )
462 .expect("resolves through the ordinary binding lookup");
463 assert!(
464 std::ptr::fn_addr_eq(kernel, noop_kernel as KernelRef),
465 "the exact adopted kernel, from the one registry",
466 );
467 // …and the dtype-precise lookup is an honest miss on a wrong tuple.
468 assert!(
469 global_bindings()
470 .lookup_with_caps(
471 BindingKey::RuntimeFused(id),
472 &[DType::F64, DType::F64, DType::F64],
473 BackendId::Cpu,
474 )
475 .is_err(),
476 "wrong dtype tuple ⇒ NoBackendForOp, not a wrong-kernel bind",
477 );
478 }
479
480 /// THE REGRESSION TEST (Task 7 cross-crate fix, 2026-07-15). Under
481 /// `register_runtime_fused`'s region-dedup, two `adopt_runtime_fused`
482 /// calls for the SAME region + kernel + dtypes + backend must resolve to
483 /// the SAME `FusedOpId` and must NOT panic on the second call (the
484 /// pre-fix bug: the second `register_runtime_kernel` pushed a
485 /// duplicate-identical `KernelRef` binding, and `finalize()`'s
486 /// `.expect()` in `extend_global_bindings` panicked while holding the
487 /// global bindings write lock, poisoning it for every later test in the
488 /// binary). The kernel must remain available after the re-adopt.
489 /// **GAP-213.** One `FusedOpId` can carry several dtype tuples — the id
490 /// names a RECIPE, not a signature — so resolving by id alone has no single
491 /// right answer. It must not invent one.
492 ///
493 /// Born red: the old `first_runtime_fused` scan bound `dtypes` and never
494 /// tested it, so it answered by `HashMap::iter()` order and could hand back
495 /// the 2-operand kernel for a 3-operand call. That is how GAP-001 spent
496 /// eleven days misattributed to a `baracuda-kernelgen` version.
497 ///
498 /// Built on a LOCAL table on purpose: the global-registry tests in this
499 /// module share process state and call `clear_runtime_fused_for_tests`
500 /// concurrently, so an id minted in one can be invalidated by another. This
501 /// test asserts a property of the lookup, and should not be able to fail
502 /// for a reason that has nothing to do with the lookup.
503 #[test]
504 fn one_recipe_two_signatures_refuses_to_guess() {
505 let mut table = KernelBindingTable::new();
506 let id = FusedOpId(FusedOpId::RUNTIME_FUSED_BASE);
507 register_runtime_kernel_into(
508 &mut table,
509 id,
510 &[DType::F32, DType::F32],
511 BackendId::Cpu,
512 noop_kernel as KernelRef,
513 );
514 register_runtime_kernel_into(
515 &mut table,
516 id,
517 &[DType::F32, DType::F32, DType::F32],
518 BackendId::Cpu,
519 other_noop_kernel as KernelRef,
520 );
521
522 match table.unique_runtime_fused(id, BackendId::Cpu) {
523 crate::kernel::RuntimeFusedLookup::Ambiguous { signatures } => {
524 assert_eq!(
525 signatures.len(),
526 2,
527 "both signatures reported: {signatures:?}"
528 );
529 assert!(
530 signatures.iter().any(|s| s.len() == 2)
531 && signatures.iter().any(|s| s.len() == 3),
532 "the report must name the competing ARITIES, got {signatures:?}",
533 );
534 }
535 other => {
536 panic!("an id carrying two signatures must refuse, not guess — got {other:?}",)
537 }
538 }
539 }
540
541 /// Control. With ONE signature there is nothing to be ambiguous about and
542 /// the lookup must still answer. Without this, the assertion above is
543 /// satisfied by a `unique_runtime_fused` that refuses unconditionally.
544 #[test]
545 fn one_signature_still_resolves() {
546 let mut table = KernelBindingTable::new();
547 let id = FusedOpId(FusedOpId::RUNTIME_FUSED_BASE);
548 register_runtime_kernel_into(
549 &mut table,
550 id,
551 &[DType::F32, DType::F32],
552 BackendId::Cpu,
553 noop_kernel as KernelRef,
554 );
555 match table.unique_runtime_fused(id, BackendId::Cpu) {
556 crate::kernel::RuntimeFusedLookup::Unique(dtypes, _) => {
557 assert_eq!(dtypes, &[DType::F32, DType::F32]);
558 }
559 other => panic!("one signature must resolve — got {other:?}"),
560 }
561 }
562
563 /// Second control: an unbound id must report NotBound, not Ambiguous. Keeps
564 /// the two miss reasons from collapsing back into one.
565 #[test]
566 fn unbound_id_reports_not_bound_not_ambiguous() {
567 let table = KernelBindingTable::new();
568 let id = FusedOpId(FusedOpId::RUNTIME_FUSED_BASE);
569 assert!(matches!(
570 table.unique_runtime_fused(id, BackendId::Cpu),
571 crate::kernel::RuntimeFusedLookup::NotBound
572 ));
573 }
574
575 #[test]
576 fn adopt_runtime_fused_is_idempotent_for_an_identical_region_and_kernel() {
577 let id1 = adopt_runtime_fused(
578 "test::adopt::idempotent::sqrt_div",
579 sqrt_div(),
580 noop_kernel as KernelRef,
581 vec![DType::F32, DType::F32, DType::F32],
582 BackendId::Cpu,
583 )
584 .expect("sqrt(div) is a registrable region");
585
586 // Re-adopt: identical region ⇒ dedup returns the same id; identical
587 // kernel/dtypes/backend ⇒ must be a no-op bind, not a panic.
588 let id2 = adopt_runtime_fused(
589 "test::adopt::idempotent::sqrt_div::again",
590 sqrt_div(),
591 noop_kernel as KernelRef,
592 vec![DType::F32, DType::F32, DType::F32],
593 BackendId::Cpu,
594 )
595 .expect("re-adopting the same region is still Some");
596
597 assert_eq!(
598 id1, id2,
599 "structurally-identical regions dedup to the same FusedOpId"
600 );
601 assert!(
602 fused_kernel_available(id1, BackendId::Cpu),
603 "the kernel is still available after the idempotent re-adopt",
604 );
605 let row = lookup_runtime_kernel(id1, BackendId::Cpu).expect("row still bound on Cpu");
606 assert!(
607 std::ptr::fn_addr_eq(row.kernel, noop_kernel as KernelRef),
608 "the (single, non-duplicated) kernel resolves correctly",
609 );
610 }
611
612 // ---- flash-arm registry-seam bridge (dd-shapes, 2026-07-08) -----------
613
614 /// GREEN target (cuda build only): the CUDA `flash_decoding` kernel
615 /// registers only into `KernelBindingTable`, so before the bridge this was
616 /// unconditionally `false` on every host. Forces `global_bindings()` to
617 /// initialize (which runs the CUDA registration under
618 /// `#[cfg(feature = "cuda")]`) then asserts the static-id bridge sees it.
619 #[test]
620 #[cfg(feature = "cuda")]
621 fn fused_kernel_available_bridges_static_flash_attn_from_binding_table_on_cuda() {
622 let _guard = crate::dispatch::global_bindings();
623 drop(_guard);
624 assert!(
625 fused_kernel_available(FusedOps::FLASH_ATTN, BackendId::Cuda),
626 "the CUDA flash_decoding binding registers only into KernelBindingTable; \
627 fused_kernel_available must bridge to it for the static FLASH_ATTN id",
628 );
629 }
630
631 /// GUARD (non-cuda build only): with no CUDA registration ever run, the
632 /// static-id bridge has nothing to find.
633 #[test]
634 #[cfg(not(feature = "cuda"))]
635 fn fused_kernel_available_flash_attn_false_on_cuda_without_cuda_registration() {
636 assert!(
637 !fused_kernel_available(FusedOps::FLASH_ATTN, BackendId::Cuda),
638 "no cuda registration ran in this build; the static bridge finds nothing",
639 );
640 }
641
642 // ---- lock-nesting prerequisite: table-taking variant (Spec B Task 1) --
643
644 /// [`fused_kernel_available_in`] reads ONLY the table it is handed — a
645 /// row bound into a fresh, local `KernelBindingTable` (never touching
646 /// `global_bindings()`) is visible; the same id/backend against an
647 /// unrelated empty table is not. This is the non-nesting gate the
648 /// runtime-fused pathfinder now calls (via `OptimizationContext::bindings`)
649 /// instead of re-acquiring the global read lock it may already be
650 /// nested under.
651 #[test]
652 fn fused_kernel_available_in_reads_the_passed_table_only() {
653 let mut table = KernelBindingTable::new();
654 // A runtime-fused id bound in THIS table (not the global one).
655 let id = fuel_graph::runtime_fused::register_runtime_fused(
656 "test::fused_kernel_available_in::relu_add",
657 relu_add(),
658 )
659 .expect("relu(add) is a registrable region");
660 register_runtime_kernel_into(
661 &mut table,
662 id,
663 &[DType::F32, DType::F32, DType::F32],
664 BackendId::Cuda,
665 noop_kernel as KernelRef,
666 );
667 assert!(
668 fused_kernel_available_in(&table, id, BackendId::Cuda),
669 "the row bound into THIS table is visible",
670 );
671 assert!(
672 !fused_kernel_available_in(&KernelBindingTable::new(), id, BackendId::Cuda),
673 "an unrelated empty table sees nothing — no fallback to the global table",
674 );
675 }
676
677 /// COORDINATION GUARD: the static-id bridge must be a strict no-op for
678 /// *runtime* ids — a runtime id resolves ONLY through its
679 /// `BindingKey::RuntimeFused` rows (post-fold; formerly the sidecar), and
680 /// an id adopted on one backend is never "rescued" by the bridge on
681 /// another.
682 #[test]
683 fn fused_kernel_available_runtime_id_still_routes_through_runtime_rows_untouched() {
684 let id = adopt_runtime_fused(
685 "test::adopt::flash_bridge_guard::relu_add",
686 relu_add(),
687 noop_kernel as KernelRef,
688 vec![DType::F32, DType::F32, DType::F32],
689 BackendId::Cpu,
690 )
691 .expect("relu(add) is a registrable region");
692
693 assert!(id.is_runtime(), "allocated a runtime id");
694 assert!(
695 fused_kernel_available(id, BackendId::Cpu),
696 "a runtime id still resolves via its RuntimeFused binding rows",
697 );
698 assert!(
699 !fused_kernel_available(id, BackendId::Cuda),
700 "not adopted on Cuda, and the static-id bridge must never rescue a \
701 runtime id (it returns false immediately for any id.is_runtime())",
702 );
703 }
704}