fuel_dispatch/kernel.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Kernel reference + op parameters + binding table. Phase 7.5 B1+B5.
3//!
4//! [`KernelRef`] is the uniform function-pointer type that every
5//! dispatch wrapper matches. Backend-specific typed kernels live in
6//! their backend crates (e.g. `fuel_cpu_backend::byte_kernels`); the
7//! *wrapper* functions here in fuel-storage bridge the dispatch-
8//! erased `Storage` to those typed kernels by matching on
9//! `BackendStorage::Cpu(...)` etc.
10//!
11//! [`OpParams`] is the typed extras bag — one variant per op family
12//! that needs auxiliary data beyond inputs and outputs. Most
13//! elementwise ops use `OpParams::None`; reductions carry their
14//! reduce dims; conv2d carries kernel/stride/padding; etc.
15//!
16//! [`KernelBindingTable`] (B5) maps
17//! `(OpKind, SmallVec<[DType; N]>, BackendId) -> KernelRef`. The dtype
18//! list carries per-operand types — inputs in order, then outputs —
19//! so mixed-precision ops (e.g. `Cast: src→dst`) and same-dtype ops
20//! (e.g. `Add: [T, T, T]`) share the same key shape. Backends register
21//! their dispatch wrappers via the `dispatch::register_*` functions;
22//! op-builder methods consult the table at DAG construction time using
23//! `Graph::target_backend(id)` as the BackendId key.
24//!
25//! ## Architecture (cycle-avoidance)
26//!
27//! - Backend crates (`fuel-cpu-backend`, …): typed kernels on their
28//! concrete storage types (`CpuStorageBytes`, …). No fuel-storage
29//! dependency.
30//! - fuel-storage (this crate): dispatch wrappers that match
31//! `BackendStorage::Cpu(...)`, extract the typed storage, and
32//! call the backend's typed kernel. KernelBindingTable lives here.
33//! - Backend crates depend on fuel-storage? No — only the wrappers
34//! do. fuel-storage already depends on backend crates for variant
35//! types; this round-trip closes naturally.
36
37use std::collections::HashMap;
38use std::sync::Arc;
39use std::sync::RwLock;
40
41use fuel_graph::QuantType;
42use fuel_graph::registry::FusedOpId;
43use fuel_ir::conv::{ParamsConv1D, ParamsConvTranspose1D};
44use fuel_ir::dispatch::OpKind;
45use fuel_ir::probe::BackendId;
46use fuel_ir::{DType, Error, Layout, Result};
47use smallvec::SmallVec;
48
49/// Inline capacity for the per-operand dtype list in the binding-table
50/// key. 8 covers every op currently in flight without spilling to heap:
51/// PagedAttn (q + 2 caches + block_table + context_lens + alibi + out)
52/// at 7 entries is the worst case in inference; mixed-precision matmul
53/// is ≤ 4. Bumping later is one constant change.
54pub type KernelDTypes = SmallVec<[DType; 8]>;
55
56/// One entry of the coverage diagnostic: `(backend, op, operand dtypes)`.
57///
58/// Named rather than spelled out because the bare 3-tuple appears in a return
59/// position where it was unreadable, and `#[allow(type_complexity)]` would
60/// have asserted the tuple is fine as-is. It is not — the alias is the fix and
61/// the lint was right.
62///
63/// NOT `BindingKey`: that name is already a distinct `pub enum` in this
64/// module (`Static` / …). The compiler caught the collision; the lesson is
65/// that a type alias introduced for a lint still needs a name check.
66pub type CoverageBinding = (BackendId, OpKind, Vec<DType>);
67
68/// Per-binding capability flags. Today carries one flag — `strided_input`
69/// — that signals "this kernel walks input strides explicitly and so
70/// can consume non-contiguous input layouts (including stride-0
71/// broadcast axes) without auto-Contiguize materializing them first."
72/// The executor's contiguize gate consults this to skip the materialize
73/// step for capable kernels, so broadcast/transpose/slice layouts can
74/// reach the kernel as metadata-only views.
75///
76/// Default is the conservative all-false; binding sites opt in via
77/// [`KernelBindingTable::register_with_caps`]. Forward-extensible by
78/// adding fields (no enum/bitflags churn).
79#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
80pub struct KernelCaps {
81 /// Kernel handles non-contiguous input layouts directly (consumes
82 /// the strides side-channel from `KernelRef::layouts`). When true,
83 /// the executor passes non-contiguous inputs through unchanged
84 /// instead of running auto-Contiguize first. Inputs with non-zero
85 /// `start_offset` still go through auto-Contiguize today (slicing
86 /// the device buffer to honor offset is a separate concern).
87 pub strided_input: bool,
88 /// This kernel BAKES a specific broadcast pattern into its body — a
89 /// `broadcast_stride0: required` operand it reads as stride-0 (hoisting
90 /// `in[0]` / dropping the bcast-axis stride terms). It is therefore NOT
91 /// a generic dense-safe kernel: a dense operand routed here would be
92 /// read WRONG (element 0 repeatedly). The realize pick
93 /// ([`KernelBindingTable::lookup_with_caps`]) SKIPS such a sibling so a dense operand
94 /// is never routed into a broadcast-baked kernel — the exclusion is
95 /// order-independent (it does not rest on the generic contract being
96 /// registered first). Selecting it CORRECTLY, for an operand genuinely
97 /// stride-0 on exactly the required axes, is the layout-aware follow-on
98 /// (path 1b) that also teaches the contiguize gate; today the safe
99 /// behavior is to exclude it, so a `required`-broadcast contract is a
100 /// present-but-never-mis-selected sibling. Defaults `false` — every
101 /// existing binding is unaffected (the pick stays first-registered).
102 pub requires_broadcast: bool,
103}
104
105impl KernelCaps {
106 /// All flags off. Equivalent to `Default::default()`; provided as a
107 /// const for use in const-context registration tables.
108 pub const fn empty() -> Self {
109 Self {
110 strided_input: false,
111 requires_broadcast: false,
112 }
113 }
114
115 /// Just `strided_input` on. Ergonomic for binary/unary registrations
116 /// that opt in to the wrapper-side broadcast path.
117 pub const fn strided_input() -> Self {
118 Self {
119 strided_input: true,
120 requires_broadcast: false,
121 }
122 }
123}
124
125use fuel_memory::Storage;
126
127/// Uniform function-pointer signature for per-backend op kernels.
128///
129/// Inputs are passed as a slice of `&Storage` to handle multi-input
130/// ops (binary, ternary, custom). Outputs are passed as a slice of
131/// `&mut Storage` to handle multi-output ops (topk, var_mean,
132/// custom). Most ops use single-element slices.
133///
134/// `layouts` is a side-channel that parallels `inputs.append(outputs)`
135/// — i.e. `layouts[0..inputs.len()]` are the input layouts in order,
136/// followed by `layouts[inputs.len()..]` for the output layouts. This
137/// is the load-bearing primitive for stride-aware kernels (broadcast,
138/// transpose, gather-from-strided). Kernels that only support
139/// contiguous-equal-shape inputs (today's default) can ignore
140/// `layouts` because the executor's auto-Contiguize pass guarantees
141/// every input is contiguous before this is called; the layout slice
142/// then carries `Layout::contiguous(shape)` entries that are useful
143/// for shape inference but redundant with `Storage.len_bytes / dtype`.
144///
145/// `OpParams` carries op-family-specific extras (reduce dims,
146/// conv2d geometry, etc.). Most kernels match a specific variant;
147/// mismatches are programming bugs that the dispatch resolver
148/// must prevent. Pure layout/shape duplication that used to live in
149/// `OpParams` (e.g. the old `Reduce::input_layout`) now flows through
150/// `layouts` instead.
151///
152/// **Output Storage is pre-allocated** by the executor before the
153/// kernel is called. Kernels write into the pre-allocated bytes;
154/// they never allocate.
155///
156/// ## Multi-output kernels (Option C, Session 5)
157///
158/// Multi-output ops (e.g. SelectiveScan returning `(y, last_state)`)
159/// emit ONE `KernelRef`. The contract:
160///
161/// - `outputs.len() == 1`. The single `Arc<RwLock<Storage>>` is the
162/// producer's *bundled* Storage — its `bundle()` returns
163/// `Some(&[OutputView; N])` describing each logical slot's
164/// `dtype`/`shape`/`layout`/`byte_offset` inside the bundle's
165/// underlying byte buffer.
166/// - The kernel writes each logical output into its slot's byte
167/// range by acquiring a single write lock on the output and
168/// striding by the slot's `byte_offset`. The bundle metadata is
169/// the authoritative per-slot spec — `outputs[0].read().unwrap()
170/// .bundle().expect("bundled storage").get(slot_idx)` is the
171/// canonical access path.
172/// - Per-slot dtype tags do NOT travel through the kernel's
173/// `KernelDTypes` key (which describes inputs + the bundle's
174/// primary dtype only); the bundle metadata IS the per-slot
175/// dispatch info.
176/// - The bundle is pre-allocated by the executor via
177/// `allocate_bundled_storage(device, &output_views_spec)` (see
178/// `fuel_backend_contract::storage::allocate_bundled_storage`). Kernels
179/// never allocate; they only fill bytes.
180///
181/// Consumers of multi-output producers are NOT multi-output kernels —
182/// they're `Op::View` (zero-copy slot projection) or `Op::ViewOwned`
183/// (independent slot buffer), which the executor handles directly
184/// without invoking a kernel.
185///
186/// **Production-correct**: kernels return `Result`, never panic.
187pub type KernelRef = fn(
188 inputs: &[Arc<RwLock<Storage>>],
189 outputs: &mut [Arc<RwLock<Storage>>],
190 layouts: &[Layout],
191 params: &OpParams,
192) -> Result<()>;
193
194/// How many rows a [`OpParams::Matmul`] computes — the data-determined-M
195/// state for sparse-MoE / capacity-buffer matmuls. `m` in `OpParams::Matmul`
196/// is always the row **capacity** (buffer/stride/size); this says how many
197/// of those rows to actually compute (`<= m`), the rest left zeroed.
198#[derive(Debug, Clone, PartialEq, Eq)]
199pub enum MatmulM {
200 /// Compute every one of the `m` rows — the universal static default.
201 All,
202 /// Compute exactly this many rows of the `m`-row capacity buffer. Set
203 /// when the count is known at compile (an input-determined extent) or
204 /// after a [`Self::Deferred`] is resolved at execute.
205 Rows(usize),
206 /// The row count is **data-determined** (a producer binds it mid-pass,
207 /// e.g. `Op::NonZeroIndices`'s per-expert token count) — resolve at
208 /// execute from `produced_syms`, then patch to [`Self::Rows`].
209 Deferred(fuel_ir::DynScalar),
210}
211
212/// Typed extras bag passed to kernels via [`KernelRef`]. One variant
213/// per op family that needs auxiliary parameters; `None` for
214/// parameter-less ops (most elementwise unary and binary).
215#[derive(Debug, Clone)]
216pub enum OpParams {
217 /// Op needs no auxiliary parameters. Used by elementwise
218 /// unary (relu, neg, sqr, …), elementwise binary (add, mul,
219 /// sub, div, …), shape-only ops (reshape, transpose), etc.
220 None,
221
222 /// Reduction (sum, max, mean, …) along specific dims. The input
223 /// tensor's [`Layout`] flows through the new `layouts` side-channel
224 /// (`layouts[0]`) on `KernelRef`, so this variant carries only the
225 /// op-specific extras: which dims to reduce and the keepdim flag.
226 ///
227 /// `dims` is the sorted list of dims to reduce; `keepdim`
228 /// controls whether reduced dims are retained as size-1 in
229 /// the output (today fuel-graph never asks for keepdim, but
230 /// the field is reserved for the future).
231 Reduce { dims: Vec<usize>, keepdim: bool },
232
233 /// Matrix multiplication. Carries the dimensions explicitly
234 /// because [`Storage`](fuel_memory::Storage) only holds bytes + dtype;
235 /// the kernel needs the batch shape and `(m, n, k)` to walk
236 /// inputs and outputs.
237 ///
238 /// Shape contract: lhs `[..lhs_batch.., m, k]` @
239 /// rhs `[..rhs_batch.., k, n]` → out `[..lhs_batch.., m, n]`.
240 /// Per-axis the batch dims must either match or follow GQA-style
241 /// divisibility (`lhs_dim > rhs_dim && lhs_dim % rhs_dim == 0`)
242 /// — the kernel maps each lhs batch slot to the corresponding
243 /// rhs slot via `rhs_axis_idx = lhs_axis_idx / n_rep_axis`.
244 ///
245 /// Equal-batch case: `lhs_batch_dims == rhs_batch_dims`. Rank-2
246 /// case: both batch vectors are empty. Both work uniformly.
247 ///
248 /// Inputs are guaranteed contiguous by the executor's
249 /// auto-Contiguize pass; transpose flags don't appear here
250 /// because `Op::Transpose` is its own metadata-only op in
251 /// fuel-graph.
252 Matmul {
253 lhs_batch_dims: Vec<usize>,
254 rhs_batch_dims: Vec<usize>,
255 /// The output/LHS row **capacity** (drives size checks + strides).
256 m: usize,
257 n: usize,
258 k: usize,
259 /// How many of the `m` rows to actually compute (data-determined-M
260 /// sparse-MoE support). `MatmulM::All` for every ordinary matmul.
261 m_compute: MatmulM,
262 },
263
264 /// 1D convolution geometry (forward path).
265 Conv1D(ParamsConv1D),
266
267 /// 2D convolution geometry (forward path). Carries the tuple-shaped
268 /// stride/padding that fuel-graph's `Op::Conv2D` uses (asymmetric
269 /// supported), plus the input/weight/output shapes the kernel
270 /// needs to walk the multi-index. Storage holds only bytes +
271 /// dtype, so spatial shapes flow through OpParams.
272 ///
273 /// Inputs: `x` shape `[N, Cin, Hin, Win]`, `weight` shape
274 /// `[Cout, Cin/groups, Kh, Kw]`, optional `bias` shape `[Cout]`.
275 /// Output: `[N, Cout, Hout, Wout]`.
276 Conv2D {
277 x_shape: [usize; 4],
278 w_shape: [usize; 4],
279 out_shape: [usize; 4],
280 stride: (usize, usize),
281 padding: (usize, usize),
282 dilation: (usize, usize),
283 groups: usize,
284 },
285
286 /// 1D transposed-convolution geometry.
287 ConvTranspose1D(ParamsConvTranspose1D),
288
289 /// 2D transposed-convolution geometry. Mirrors `OpParams::Conv2D`
290 /// in shape (inline fields, asymmetric stride/padding/dilation,
291 /// `groups`), with the additional `output_padding` parameter that
292 /// transposed conv needs to disambiguate output spatial size.
293 ///
294 /// Inputs: `x` shape `[N, Cin, Hin, Win]`, `weight` shape
295 /// `[Cin, Cout/groups, Kh, Kw]` (note transposed channel order
296 /// vs Conv2D), optional `bias` shape `[Cout]`.
297 /// Output: `[N, Cout, Hout, Wout]`.
298 ConvTranspose2D {
299 x_shape: [usize; 4],
300 w_shape: [usize; 4],
301 out_shape: [usize; 4],
302 stride: (usize, usize),
303 padding: (usize, usize),
304 output_padding: (usize, usize),
305 dilation: (usize, usize),
306 groups: usize,
307 },
308
309 /// Sum-reduce a tensor to a smaller broadcast-compatible target
310 /// shape. The kernel left-pads `output_shape` with 1s to match
311 /// `input_shape`'s rank, then for each axis: if the padded output
312 /// dim equals the input dim that axis carries through; if it's 1
313 /// the axis is summed away.
314 ReduceSumTo {
315 input_shape: Vec<usize>,
316 output_shape: Vec<usize>,
317 },
318
319 /// Max-reduce a tensor to a smaller broadcast-compatible target
320 /// shape — the max-symmetric counterpart of `ReduceSumTo`. Same
321 /// axis-alignment rules; per-axis reduction is `max` instead of
322 /// `+`.
323 ReduceMaxTo {
324 input_shape: Vec<usize>,
325 output_shape: Vec<usize>,
326 },
327
328 /// Backward of `ReduceMaxTo`. Inputs: `(x, upstream)` where
329 /// `x.shape == input_shape` and `upstream.shape == output_shape`.
330 /// Output: `grad_x` of `input_shape`. The kernel recomputes the
331 /// forward max via `input_shape → output_shape` axis alignment,
332 /// builds a tie-count mask, and routes upstream back to argmax
333 /// positions (fair-share on ties).
334 ReduceMaxToBackward {
335 input_shape: Vec<usize>,
336 output_shape: Vec<usize>,
337 },
338
339 /// Multi-head scaled-dot-product attention shape + math params.
340 /// `q` is `[B, Hq, Sq, D]`, `k` and `v` are `[B, Hkv, Sk, D]`
341 /// (Hkv ≤ Hq, GQA-divisible). Optional 4th input `alibi_slopes`
342 /// has shape `[Hq]` (presence is implicit in `inputs.len() == 4`).
343 ///
344 /// `sk` is the **physical** K/V length axis (the allocated extent —
345 /// used for strides + byte-length checks). `k_len` is the
346 /// **logical** attended length (`k_len <= sk`): a kernel attends only
347 /// the first `k_len` rows of the K/V buffers and bottom-right-aligns
348 /// the causal mask at offset `k_len - sq` (Phase D symbolic extents:
349 /// flash over a fixed-capacity KV-cache with a runtime live prefix).
350 /// For the static path the lowering sets `k_len == sk`, so a kernel
351 /// that loops `0..k_len` with the `k_len - sq` causal offset is
352 /// byte-identical to the old `0..sk` / `kj > qi` form.
353 FlashAttn {
354 b: usize,
355 hq: usize,
356 hkv: usize,
357 sq: usize,
358 sk: usize,
359 d: usize,
360 k_len: usize,
361 softmax_scale: f32,
362 causal: bool,
363 window_size_left: Option<usize>,
364 window_size_right: Option<usize>,
365 softcap: Option<f32>,
366 },
367
368 /// Paged-cache attention. `q` is `[B, Hq, Sq, D]`; `k_cache` and
369 /// `v_cache` are `[num_blocks, block_size, Hkv, D]`. The 4th input
370 /// is `block_table: [B, max_blocks_per_seq]` (U32 — physical block
371 /// index per logical position). The 5th is `context_lens: [B]`
372 /// (U32 — true context length per sequence). Optional 6th is
373 /// `alibi_slopes: [Hq]`.
374 PagedAttn {
375 b: usize,
376 hq: usize,
377 hkv: usize,
378 sq: usize,
379 d: usize,
380 block_size: usize,
381 max_blocks_per_seq: usize,
382 num_blocks: usize,
383 softmax_scale: f32,
384 softcap: Option<f32>,
385 },
386
387 /// Slice along a single dim with explicit start/end/step. The
388 /// dim is implicit in the input Layout's relabeling for
389 /// multi-dim slice; this variant covers the simple case.
390 Slice {
391 dim: usize,
392 start: usize,
393 end: usize,
394 step: usize,
395 },
396
397 // (Earlier `OpParams::Pad { padding: Vec<(usize, usize)>, fill_bytes: Vec<u8> }`
398 // was a speculative multi-dim shape with no consumers. The single-
399 // dim shape that Op::Pad actually emits lives at the bottom of this
400 // enum; if multi-dim padding lands later, it can extend either
401 // shape additively.)
402 /// Cast input dtype → target dtype. The target lives on the
403 /// output Storage's `dtype` field; this variant signals the
404 /// op family without requiring a re-read.
405 Cast,
406
407 /// Affine transformation `y = mul * x + add`. Used by
408 /// `Tensor::affine` / `scale_and_shift`.
409 Affine { mul: f64, add: f64 },
410
411 /// Element-wise clamp: `y = clamp(x, min, max)`.
412 Clamp { min: f64, max: f64 },
413
414 /// Element-wise integer power: `y = x.powi(exp)`.
415 PowI { exp: i32 },
416
417 /// Concatenate N inputs along one dim. The kernel needs the
418 /// outer/inner element counts (product of dims before/after
419 /// the concat dim) plus each input's size along the concat
420 /// dim — that's all that distinguishes a concat from a
421 /// sequence of slab-copies. Order matches `Node::inputs`.
422 ///
423 /// `axis` is the original concat dim in the output's shape;
424 /// kernels that want to walk strided rank-N inputs need it
425 /// to build the per-axis stride mask (the outer/inner factoring
426 /// alone loses the position info).
427 Concat {
428 /// Product of output dims before the concat dim.
429 outer_count: usize,
430 /// Per-input size along the concat dim (length = N inputs).
431 input_dim_sizes: Vec<usize>,
432 /// Product of output dims after the concat dim.
433 inner_count: usize,
434 /// Original concat axis index in the output's rank-N shape.
435 axis: usize,
436 },
437
438 /// Softmax along the last dim. The kernel walks
439 /// `outer_count` rows of `last_dim` elements each.
440 SoftmaxLastDim { outer_count: usize, last_dim: usize },
441
442 /// Last-dim norm parameters shared by RMS-norm and LayerNorm
443 /// (no affine flavor for both today). The OpKind selects which
444 /// kernel reads this variant.
445 NormLastDim {
446 outer_count: usize,
447 last_dim: usize,
448 eps: f64,
449 },
450
451 /// Pick slices along a single dim using a rank-1 U32 index
452 /// tensor. The kernel needs the outer/inner counts (dims
453 /// before/after the selected axis), the source's selected-dim
454 /// size for index bounds checking, and the index count.
455 IndexSelect {
456 outer_count: usize,
457 source_dim_size: usize,
458 n_indices: usize,
459 inner_count: usize,
460 },
461
462 /// N-dimensional gather along `dim`. Source and output shapes
463 /// agree on every dim except `dim`; the indices tensor (U32)
464 /// has output_shape and supplies the source coord for `dim`
465 /// at every output position.
466 Gather {
467 source_shape: Vec<usize>,
468 output_shape: Vec<usize>,
469 dim: usize,
470 },
471
472 /// Rotary position embedding parameters. The kernel walks
473 /// `outer_count` × `seq` × `head_dim` elements; `cos`/`sin`
474 /// have `[seq, head_dim]` shape and broadcast across the
475 /// outer dims.
476 Rope {
477 outer_count: usize,
478 seq: usize,
479 head_dim: usize,
480 },
481
482 /// Quantized matmul shape: `A [batch, m, k] @ dequant(W) [n, k] →
483 /// out [batch, m, n]`. The weight tensor is a flat U32-typed
484 /// byte stream representing `n * k / elements_per_block` blocks
485 /// of `quant_type`. Activations are F32 today; output is F32.
486 QMatMul {
487 quant_type: QuantType,
488 batch_count: usize,
489 m: usize,
490 n: usize,
491 k: usize,
492 },
493
494 /// Index-add along a single dim with rank-1 U32 indices.
495 /// Output is `base` with `src[..., i, ...]` accumulated into
496 /// `base[..., indices[i], ...]` for every i ∈ 0..n_indices.
497 IndexAdd {
498 outer_count: usize,
499 base_dim_size: usize,
500 n_indices: usize,
501 inner_count: usize,
502 },
503
504 /// N-dimensional scatter-add. Indices and src share the same
505 /// shape; base may differ only along `dim`. The kernel walks
506 /// every src/indices position, reads `indices[p]` for the
507 /// destination's `dim` coord, and accumulates `src[p]` into
508 /// `base` at that destination.
509 ScatterAdd {
510 base_shape: Vec<usize>,
511 src_shape: Vec<usize>,
512 dim: usize,
513 },
514
515 /// Flip the order of elements along one dim. The flat-3-axis
516 /// view (outer × dim × inner) lets the kernel walk a tight loop
517 /// without re-deriving the axis split per call. `axis` is the
518 /// original dim index in the input's rank-N shape — needed by
519 /// stride-aware kernels to build a per-axis flip mask.
520 Flip {
521 outer_count: usize,
522 dim_size: usize,
523 inner_count: usize,
524 axis: usize,
525 },
526
527 /// Cyclic shift along one dim by `shift` positions (positive
528 /// shifts move elements to higher indices, wrapping around).
529 /// Same flat-3-axis view as `Flip`. `shift` is signed: negative
530 /// shifts move elements to lower indices.
531 Roll {
532 outer_count: usize,
533 dim_size: usize,
534 inner_count: usize,
535 shift: i64,
536 axis: usize,
537 },
538
539 /// Running cumulative sum along one dim. Same flat-3-axis view
540 /// as `Flip`/`Roll`. Output is always the same dtype as input;
541 /// kernel needs typed addition so it's per-dtype (unlike Flip/Roll).
542 CumSum {
543 outer_count: usize,
544 dim_size: usize,
545 inner_count: usize,
546 axis: usize,
547 },
548
549 /// Triangular mask parameters (used by both Triu and Tril — the
550 /// op-kind picks the direction). `batch_count` is the product of
551 /// leading dims; `rows`/`cols` are the last two dims. `diagonal`
552 /// is signed (0 = main diagonal, positive shifts up, negative
553 /// shifts down).
554 Triangular {
555 batch_count: usize,
556 rows: usize,
557 cols: usize,
558 diagonal: i64,
559 },
560
561 /// LogSoftmax along the last dim. Walks `outer_count` rows of
562 /// `last_dim` elements each. Per-dtype kernel (uses log/exp).
563 LogSoftmaxLastDim { outer_count: usize, last_dim: usize },
564
565 /// MaskedFill: per-element fill where mask is nonzero. The kernel
566 /// reads the element count from the layout; `fill_bytes` is
567 /// pre-encoded in the output's dtype (one element's worth).
568 MaskedFill { fill_bytes: Vec<u8> },
569
570 /// Backward helper for Pad. Carries the input shape so the kernel
571 /// can size its scatter-add buffer; the output shape is implicit
572 /// from the input shape + padding.
573 PadBackward {
574 in_shape: Vec<usize>,
575 out_shape: Vec<usize>,
576 padding: Vec<(usize, usize)>,
577 mode_tag: u8,
578 },
579
580 /// Multi-dim Pad: per-axis (before, after) plus a `mode_tag`
581 /// (0=Constant, 1=Reflect, 2=Replicate) and pre-encoded
582 /// `fill_bytes` for Constant fill. Dtype-agnostic at the byte
583 /// level: the kernel just copies bytes per the input/output shapes,
584 /// which is why fill is bytes (already encoded in the output's
585 /// dtype) rather than `f64` (which would force the kernel to
586 /// know its dtype).
587 ///
588 /// `in_shape.len() == out_shape.len() == padding.len()`, and
589 /// `out_shape[i] == in_shape[i] + padding[i].0 + padding[i].1`.
590 /// `fill_bytes.len() == dtype_size_in_bytes` (one element's worth).
591 Pad {
592 in_shape: Vec<usize>,
593 out_shape: Vec<usize>,
594 padding: Vec<(usize, usize)>,
595 mode_tag: u8,
596 fill_bytes: Vec<u8>,
597 },
598
599 /// In-place scatter write parameters. The destination shape +
600 /// source shape together determine the kernel walk; `ranges`
601 /// gives the per-axis (start, end) slab inside the destination.
602 /// `dest_shape[i] >= ranges[i].1`, and source shape is implicitly
603 /// `ranges[i].1 - ranges[i].0` along axis `i`.
604 ///
605 /// Phase E.3.2: backs `Op::WriteSlice` for persistent KV-cache
606 /// writes.
607 ///
608 /// `deferred_dyn_offset` is the **data-determined dynamic-shapes** seam
609 /// (increment 2b). Normally `None` — `ranges` is already resolved at
610 /// compile (static, or an *input-determined* `dyn_offset` bound in the
611 /// per-pass `SymEnv`). It is `Some((axis, sym))` only when the offset is
612 /// *data-determined* — its `SymId` is UNBOUND at compile because a
613 /// producer (e.g. `Op::NonZeroIndices`) computes it mid-pass. The
614 /// `WorkItemKind::WriteSlice` executor arm resolves it from
615 /// `produced_syms` and overrides `ranges[axis]` to `(base, base+width)`
616 /// just before the kernel runs; the kernel wrappers ignore this field.
617 WriteSlice {
618 dest_shape: Vec<usize>,
619 ranges: Vec<(usize, usize)>,
620 deferred_dyn_offset: Option<(usize, fuel_ir::DynScalar)>,
621 },
622
623 /// In-place ring-buffer scatter write parameters. Like
624 /// `WriteSlice` except `axis` wraps modulo `modulus`; the
625 /// dynamic start on that axis is read from an extra input (rank-0
626 /// U32). On the rotating axis, `ranges[axis].0` is ignored
627 /// (dynamic) and `ranges[axis].1 - ranges[axis].0` is the slab
628 /// width (must equal the source dim on that axis, and must not
629 /// exceed `modulus`).
630 ///
631 /// Phase C: backs `Op::WriteSliceRotating` for sliding-window
632 /// KV caches.
633 WriteSliceRotating {
634 dest_shape: Vec<usize>,
635 axis: usize,
636 modulus: usize,
637 ranges: Vec<(usize, usize)>,
638 },
639
640 /// In-place scatter write whose start on `axis` is read from a
641 /// **device-resident** rank-0 `I64` input (`inputs[2]`) at kernel
642 /// launch — the CUDA-graph-capturable KV-cache append. Like
643 /// `WriteSlice` except `ranges[axis].0` is a placeholder overridden
644 /// device-side by the offset; no modulo wrap. The CUDA wrapper
645 /// passes the offset buffer's device pointer straight through as
646 /// `dyn_start_dev` (NO D2H — that would break capture); the CPU
647 /// wrapper reads it host-side and delegates to `write_slice_cpu`.
648 /// Bounds (`offset + width <= dest_shape[axis]`) are the caller's
649 /// contract; the kernel does not clamp.
650 ///
651 /// CapturedRun (2026-07): backs `Op::WriteSliceDoff` for
652 /// `DecodeSession` decode-step capture/replay.
653 WriteSliceDoff {
654 dest_shape: Vec<usize>,
655 axis: usize,
656 ranges: Vec<(usize, usize)>,
657 },
658
659 /// FusedSoftmaxCrossEntropy execution parameters. Inputs:
660 /// `logits [n_rows, vocab]` (F32, flattened from the original
661 /// `[..., V]` shape) and `targets [n_rows]` (I64). The kernel
662 /// walks row-by-row, computing stable log-softmax + NLL +
663 /// `ignore_index` masking + the requested reduction in a single
664 /// pass, allocating only an `[n_rows]` per-row accumulator (and
665 /// a scalar for Mean/Sum). Output is F32 scalar `[]` for
666 /// Mean/Sum, F32 `[n_rows]` for None.
667 ///
668 /// `n_rows` and `vocab` are derived at translate time from the
669 /// logits layout; the kernel uses them to iterate without
670 /// re-parsing shapes.
671 FusedSoftmaxCrossEntropy {
672 n_rows: usize,
673 vocab: usize,
674 reduction: fuel_graph::registry::Reduction,
675 ignore_index: i64,
676 },
677
678 /// CausalConv1d execution parameters. Inputs:
679 /// `x [batch, channels, seq_in]` (pre-padded by caller with
680 /// `kernel - 1` left zeros), `weight [channels, 1, kernel]`,
681 /// `bias [channels]`. Output `[batch, channels, seq_out]` where
682 /// `seq_out = seq_in - (kernel - 1)`. All tensors share dtype.
683 ///
684 /// Carries both `seq_in` and `seq_out` so the kernel walks the
685 /// time axis without re-deriving from shapes. `use_silu` toggles
686 /// the fused SiLU activation on the output store (matches
687 /// baracuda's `causal_conv1d_*_run` signature flag).
688 CausalConv1d {
689 batch: usize,
690 channels: usize,
691 seq_in: usize,
692 seq_out: usize,
693 kernel: usize,
694 use_silu: bool,
695 },
696
697 /// SelectiveScan execution parameters. 5 required inputs:
698 /// `u [batch, seqlen, dim]`, `delta [batch, seqlen, dim]`,
699 /// `a [dim, dstate]`, `b [batch, seqlen, dstate]`,
700 /// `c [batch, seqlen, dstate]`. Output `y [batch, seqlen, dim]`.
701 /// All tensors share dtype.
702 ///
703 /// `delta_softplus` toggles applying softplus(delta) before use
704 /// (matches baracuda's `selective_scan_*_run` flag).
705 SelectiveScan {
706 batch: usize,
707 seqlen: usize,
708 dim: usize,
709 dstate: usize,
710 delta_softplus: bool,
711 },
712
713 /// SsdChunkScan execution parameters. 5 inputs:
714 /// `x [batch, seqlen, heads, head_dim]`,
715 /// `dt [batch, seqlen, heads]`, `a [heads]`,
716 /// `b [batch, seqlen, heads, state_dim]`,
717 /// `c [batch, seqlen, heads, state_dim]`. Output `y` matches `x`'s
718 /// shape. All tensors share dtype.
719 ///
720 /// `chunk_size` is the SSD block size (GPU-parallelism knob).
721 /// On CPU the kernel runs a sequential scan regardless;
722 /// validation requires `chunk_size > 0` and
723 /// `seqlen % chunk_size == 0`.
724 SsdChunkScan {
725 batch: usize,
726 seqlen: usize,
727 heads: usize,
728 head_dim: usize,
729 state_dim: usize,
730 chunk_size: usize,
731 },
732
733 /// Nf4Matmul execution parameters. 3 inputs:
734 /// `activations [batch, m, k]` (rank ≥ 2; leading dims flattened
735 /// into `batch`), `w_packed [n, k/2]` U8, `absmax [n, k/block_size]`
736 /// F32. Output `[batch, m, n]` matching activations' dtype.
737 ///
738 /// `block_size` is the per-output-row, per-block scale granularity
739 /// (typically 64 in bitsandbytes). `k` must be even (w_packed
740 /// layout requirement) and a multiple of `block_size`.
741 Nf4Matmul {
742 batch: usize,
743 m: usize,
744 n: usize,
745 k: usize,
746 block_size: usize,
747 },
748
749 /// AwqMatmul execution parameters. 4 inputs: `activations [..., m, k]`
750 /// (rank ≥ 2; every leading dim is flattened into `batch`, so the CUDA
751 /// wrapper runs a GEMM of `batch * m` rows), `weight [n, k/8]` U32,
752 /// `zeros [n, zeros_w]` U32, `scales [n, 8 * zeros_w]` F16 (baracuda W4A16
753 /// kernel layout). Output `[batch, m, n]` matching activations' dtype.
754 ///
755 /// `group_size` ∈ {64, 128}; `split_k_iters` is the CUDA launch-tuning
756 /// parameter (typically 8).
757 AwqMatmul {
758 batch: usize,
759 m: usize,
760 n: usize,
761 k: usize,
762 group_size: usize,
763 split_k_iters: i32,
764 },
765
766 /// A **runtime fused op's** live scalar args (the `extract:` values from
767 /// `FusedOpParams::Runtime { scalars }`), in the region's pattern
768 /// pre-order — the same canonical slot order the synthesized kernel's
769 /// trailing `, float p0, float p1, …` launch parameters are declared in
770 /// (the synthesizer's scalar ABI appends them after `long long n`, always
771 /// `float`-typed regardless of the kernel's element dtype). Only produced
772 /// by `compile_one`'s `is_runtime` arm; a param-free runtime op carries
773 /// `OpParams::None` instead.
774 JitScalars {
775 /// Slot values, pattern pre-order. `f64` at rest (the graph-side
776 /// carrier); the CUDA launch narrows each to `f32` per the ABI.
777 scalars: Vec<f64>,
778 },
779
780 /// NonZeroIndices execution parameters. 1 input `x` (the value
781 /// tensor); one bundled output `[indices [capacity] U32 ; count
782 /// [1] U32]` where `capacity == x.elem_count()`. The kernel writes
783 /// the ascending flat indices of `x`'s nonzero elements into slot 0
784 /// (first `count` entries) and the runtime `count` into slot 1.
785 ///
786 /// `count_sym` is the [`fuel_ir::symbol::SymId`] the executor binds to the realized
787 /// `count` in the per-pass `SymEnv` *after* this op runs — the
788 /// data-determined dynamic-shape seam. The kernel itself ignores it.
789 NonZeroIndices {
790 capacity: usize,
791 count_sym: fuel_ir::SymId,
792 },
793}
794
795// =============================================================================
796// Phase 7.5 B5 — kernel binding table
797// =============================================================================
798
799/// Maps `(OpKind, KernelDTypes, BackendId)` triples to dispatch wrapper
800/// functions, where `KernelDTypes` lists per-operand dtypes (inputs in
801/// order, then outputs). Built once at backend registration time
802/// (typically a process-wide `OnceLock` though that's not enforced
803/// here for testability), consulted at execute time via lookup.
804///
805/// Same-dtype ops register `[T, T, ..., T]` for the right operand
806/// count; mixed-precision ops (Cast, future F32×BF16→F32 matmul)
807/// register the exact combo. Variadic ops (Concat) register a
808/// canonical short shape `[T, T]` (one input dtype + output dtype);
809/// the lookup site for those ops collapses its dtypes vector to the
810/// same shorthand.
811///
812/// Backends register their wrappers via the `dispatch::register_*`
813/// functions in this crate (e.g.
814/// [`crate::dispatch::register_cpu_kernels`]).
815/// Cost-fn signature for primitive-op registrations stored in
816/// [`KernelBindingTable`]. Mirrors the fused-op cost-fn shape but
817/// takes [`OpParams`] (the binding-table param payload) instead of
818/// `FusedOpParams`.
819///
820/// Implementations return a [`crate::fused::CostEstimate`] computed
821/// statically from shapes + dtypes + op-specific params + backend
822/// capabilities. The architecture's Layer-1 (FLOP-count + bandwidth)
823/// cost model lives here; Layer-2 empirical refinement composes on
824/// top via the telemetry framework, not by changing this signature.
825pub type CostFn = fn(
826 &[fuel_ir::Shape],
827 &[DType],
828 &OpParams,
829 &fuel_ir::backend::BackendCapabilities,
830) -> crate::fused::CostEstimate;
831
832/// Sentinel "no cost claim" function — returns
833/// [`crate::fused::CostEstimate::default`] (all-zero). The fill pass
834/// (`fill_unset_cpu_cost`) recognizes this exact function pointer
835/// to decide which entries get the OpKind-family default.
836pub fn unknown_cost(
837 _shapes: &[fuel_ir::Shape],
838 _dtypes: &[DType],
839 _params: &OpParams,
840 _caps: &fuel_ir::backend::BackendCapabilities,
841) -> crate::fused::CostEstimate {
842 crate::fused::CostEstimate::default()
843}
844
845/// One concrete kernel registered against a `(OpKind, dtypes, backend)`
846/// decision-point key. Phase 7.6 step 9a: a single key now carries
847/// `SmallVec<[BindingEntry; 2]>` — multiple alternatives compete at
848/// the same decision point (e.g. cuBLAS bf16 matmul + CUTLASS bf16
849/// matmul at `(MatMul, [BF16, BF16, BF16], Cuda)`). The route picker
850/// (step 9b) selects among them at plan time; today's single-impl
851/// callers (`lookup`, `lookup_with_caps`, `lookup_precision`,
852/// `lookup_cost`) return the first entry, preserving the pre-9a
853/// behavior for sites that haven't migrated.
854#[derive(Clone, Copy, Debug)]
855pub struct BindingEntry {
856 pub kernel: KernelRef,
857 pub caps: KernelCaps,
858 pub precision: crate::fused::PrecisionGuarantee,
859 pub cost: CostFn,
860 /// Diagnostic tag identifying the kernel's implementation source
861 /// when multiple alternatives register at the same `(op, dtypes,
862 /// backend)` key. Used to distinguish e.g. AOCL's matmul from
863 /// MKL's matmul from the portable scalar matmul — all three live
864 /// at `(MatMul, [F32×3], Cpu)` as siblings.
865 ///
866 /// `""` is the default for kernels that don't need to distinguish.
867 /// Conventional values: `"portable-cpu"` for fuel-cpu-backend
868 /// scalar/gemm kernels, `"aocl"` for AOCL-BLAS, `"mkl"` for oneMKL,
869 /// `"cublas"` / `"cutlass"` for CUDA variants. The string is
870 /// diagnostic-only — never used for dispatch routing (the kernel
871 /// function pointer is the dispatch identifier).
872 pub kernel_source: &'static str,
873 /// Whether this kernel's FKC accept-contract is the **GENERIC**
874 /// (fully-permissive strided) one — admissible anywhere because it
875 /// imposes no structure tightness (FKC §4.12). This is the
876 /// structural-**miss** fallback bit: when the winning contract at a
877 /// dispatch cell is generic, a structure-specialized kernel would have
878 /// fit but none is registered, and the Baracuda miss telemetry records
879 /// the demand.
880 ///
881 /// Computed **once at registration** by the FKC importer from the
882 /// kernel's retained `ResolvedLayout` set
883 /// ([`crate::fkc::is_generic_contract`]) and stamped here, so
884 /// the live dispatch pick site reads a precomputed bool rather than
885 /// re-deriving genericity from the lossy single-bool [`KernelCaps`].
886 ///
887 /// **Hand-written registrations default `false`.** Every non-FKC
888 /// `register*` path routes through
889 /// [`KernelBindingTable::register_full_with_source`], which sets this
890 /// `false`: genericity is a claim about being the FKC generic-fallback
891 /// contract in Baracuda's structure-key specialization space, and a
892 /// hand-written kernel makes no such claim. This holds even for a
893 /// hand-written strided kernel (`register_with_caps(strided_input=true)`):
894 /// strided-capability is not the same as being the generic fallback, and
895 /// conflating them would fabricate a miss demand signal for cells that are
896 /// not part of Baracuda's matrix. Only the FKC importer, which computed
897 /// genericity from the declared contract, ever sets this `true`.
898 pub is_generic: bool,
899 /// Stable per-implementation-version hash — the fifth axis of the FKC
900 /// kernel-identity basis tuple (FKC §4.11), pinning the exact kernel build
901 /// so a persisted plan / telemetry record re-resolves to it. `0`
902 /// ([`crate::fused::KernelRevisionHash::UNTRACKED`]) for hand-written
903 /// registrations (which carry no revision); the FKC importer threads the
904 /// contract-computed revision from `ResolvedPrimitive.revision` here. Read
905 /// by the Baracuda dispatch-telemetry `ImplId` at the plan pick site so the
906 /// `chosen`/`candidates` ids carry a real revision instead of a stamped `0`.
907 pub kernel_revision_hash: u64,
908 /// The contract's parsed cost-expression AST (FKC §4.4 / adoption plan
909 /// §2.3), when the kernel came from an FKC-imported contract that
910 /// declares one. `&'static` so the entry stays `Copy` (`CompiledCostExpr`
911 /// derives `Debug`, which is all `&'static T` needs to stay `Copy` /
912 /// `Clone` / `Debug` itself). `None` for every hand-written registration
913 /// (`register`/`register_full*`) — those carry no declared cost AST, only
914 /// a `CostFn` fn-pointer. Only the FKC importer, once it lowers a
915 /// contract's cost block to a live `CostFn` (a later Group-2 task),
916 /// stamps `Some(&CompiledCostExpr)` here.
917 pub cost_expr: Option<&'static crate::fkc::CompiledCostExpr>,
918}
919
920/// The binding table's op key — generalized so **runtime-registered** fused
921/// ops live in the SAME registry as static kernels (the 2026-07-08
922/// decisions-log end-state; retires the transitional runtime-kernel sidecar).
923///
924/// `Static` is the compile-time `OpKind` family every existing registration
925/// uses (a bare `OpKind` converts via `From`, so those call sites are
926/// unchanged). `RuntimeFused` names a Tier-2 (JIT-synthesized / import-time)
927/// fused op by its runtime `FusedOpId` (`>= RUNTIME_FUSED_BASE`) — the id IS
928/// the op identity, so no fake `OpKind` ever enters the table.
929#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
930pub enum BindingKey {
931 /// A compile-time op family (every static kernel registration).
932 Static(OpKind),
933 /// A runtime-registered fused op, keyed by its runtime id.
934 RuntimeFused(FusedOpId),
935}
936
937impl From<OpKind> for BindingKey {
938 fn from(op: OpKind) -> Self {
939 BindingKey::Static(op)
940 }
941}
942
943impl BindingKey {
944 /// The `Static` op kind, if this is a static-family key.
945 pub fn static_op(self) -> Option<OpKind> {
946 match self {
947 BindingKey::Static(op) => Some(op),
948 BindingKey::RuntimeFused(_) => None,
949 }
950 }
951}
952
953/// `Clone` is load-bearing, not a convenience: the process-wide table is
954/// published as a copy-on-write snapshot (`RwLock<Arc<KernelBindingTable>>`
955/// in `dispatch.rs`), so a writer clones the current table, mutates the
956/// clone, and swaps the pointer. Readers therefore never hold a lock guard.
957/// The clone is shallow-ish and cheap — [`BindingEntry`] is `Copy` and the
958/// keys are `Copy`/`SmallVec`, so it is a `HashMap` bulk copy of POD, on a
959/// write-rare path (backend init + runtime-fusion adoption).
960#[derive(Default, Clone)]
961pub struct KernelBindingTable {
962 bindings: HashMap<(BindingKey, KernelDTypes, BackendId), SmallVec<[BindingEntry; 2]>>,
963}
964
965/// The outcome of [`KernelBindingTable::unique_runtime_fused`] — three
966/// distinct facts, deliberately not collapsed into an `Option`.
967#[derive(Debug)]
968pub enum RuntimeFusedLookup<'a> {
969 /// No runtime-fused binding for this `(fid, backend)`.
970 NotBound,
971 /// Exactly one dtype tuple is registered; this is its first alternative.
972 Unique(&'a [DType], &'a BindingEntry),
973 /// Several dtype tuples share this id. **No kernel is returned**, because
974 /// any choice among them would be arbitrary — and one of them may have a
975 /// different operand arity than the caller expects.
976 Ambiguous { signatures: Vec<Vec<DType>> },
977}
978
979impl KernelBindingTable {
980 pub fn new() -> Self {
981 Self {
982 bindings: HashMap::new(),
983 }
984 }
985
986 /// Register a dispatch wrapper for `(op, dtypes, backend)` with the
987 /// default (all-false) capabilities. `dtypes` is the per-operand
988 /// dtype list — inputs in order, then outputs.
989 ///
990 /// Phase 7.6 step 9a: multiple distinct kernels may register
991 /// against the same `(op, dtypes, backend)` key — they become
992 /// sibling alternatives at one decision point. Registration is
993 /// **append-only and never panics**; the **same** `KernelRef`
994 /// function pointer registered twice is detected by
995 /// [`Self::finalize`] (run once after bulk registration) and
996 /// returned as a `Result`, not raised as an inline panic.
997 ///
998 /// PrecisionGuarantee defaults to
999 /// [`PrecisionGuarantee::UNAUDITED`](crate::fused::PrecisionGuarantee::UNAUDITED).
1000 /// Step-7b convention: the always-built backend
1001 /// (fuel-cpu-backend) runs [`Self::fill_unset_cpu_precision`] at
1002 /// the end of its bulk registration pass to upgrade every UNAUDITED
1003 /// CPU entry to `PRIMITIVE_DETERMINISTIC_CPU` (bit-stable per
1004 /// hardware). Kernels with weaker guarantees should call
1005 /// [`Self::register_with_precision`] explicitly *before* the fill
1006 /// pass to opt out of the default.
1007 pub fn register(
1008 &mut self,
1009 op: impl Into<BindingKey>,
1010 dtypes: &[DType],
1011 backend: BackendId,
1012 kernel: KernelRef,
1013 ) {
1014 self.register_full(
1015 op,
1016 dtypes,
1017 backend,
1018 kernel,
1019 KernelCaps::empty(),
1020 crate::fused::PrecisionGuarantee::UNAUDITED,
1021 unknown_cost,
1022 );
1023 }
1024
1025 /// Register a dispatch wrapper with explicit capability flags.
1026 /// Used by binding sites that opt into kernel-side broadcast or
1027 /// other non-default behavior — the executor consults the caps to
1028 /// decide whether to auto-Contiguize inputs before kernel call.
1029 pub fn register_with_caps(
1030 &mut self,
1031 op: impl Into<BindingKey>,
1032 dtypes: &[DType],
1033 backend: BackendId,
1034 kernel: KernelRef,
1035 caps: KernelCaps,
1036 ) {
1037 self.register_full(
1038 op,
1039 dtypes,
1040 backend,
1041 kernel,
1042 caps,
1043 crate::fused::PrecisionGuarantee::UNAUDITED,
1044 unknown_cost,
1045 );
1046 }
1047
1048 /// Register with an explicit [`crate::fused::PrecisionGuarantee`].
1049 /// Use this when a kernel has a precision claim that differs from
1050 /// the bulk-fill default (e.g., a multi-threaded reduction with
1051 /// non-deterministic accumulation order — bit_stable_on_same_hardware
1052 /// must be false).
1053 pub fn register_with_precision(
1054 &mut self,
1055 op: impl Into<BindingKey>,
1056 dtypes: &[DType],
1057 backend: BackendId,
1058 kernel: KernelRef,
1059 precision: crate::fused::PrecisionGuarantee,
1060 ) {
1061 self.register_full(
1062 op,
1063 dtypes,
1064 backend,
1065 kernel,
1066 KernelCaps::empty(),
1067 precision,
1068 unknown_cost,
1069 );
1070 }
1071
1072 /// Backwards-compatible alias for full-form registration without
1073 /// an explicit cost function. New code should prefer
1074 /// [`Self::register_full`].
1075 pub fn register_with_caps_and_precision(
1076 &mut self,
1077 op: impl Into<BindingKey>,
1078 dtypes: &[DType],
1079 backend: BackendId,
1080 kernel: KernelRef,
1081 caps: KernelCaps,
1082 precision: crate::fused::PrecisionGuarantee,
1083 ) {
1084 self.register_full(op, dtypes, backend, kernel, caps, precision, unknown_cost);
1085 }
1086
1087 /// Phase 7.6 step 8: full-form register with caps + precision +
1088 /// cost. Cost defaults to [`unknown_cost`] in the other
1089 /// signatures; sites with non-default cost claims must use this
1090 /// form (and not be overwritten by the bulk
1091 /// [`Self::fill_unset_cpu_cost`] pass, which only touches
1092 /// entries whose cost is still `unknown_cost`).
1093 ///
1094 /// Phase 7.6 step 9a: appends to the alternative set for `(op,
1095 /// dtypes, backend)`. Distinct `KernelRef` function pointers
1096 /// compose as siblings. Append-only — a duplicate `KernelRef` is
1097 /// detected by [`Self::finalize`] after bulk registration, not by
1098 /// an inline panic (see [`Self::register`]).
1099 pub fn register_full(
1100 &mut self,
1101 op: impl Into<BindingKey>,
1102 dtypes: &[DType],
1103 backend: BackendId,
1104 kernel: KernelRef,
1105 caps: KernelCaps,
1106 precision: crate::fused::PrecisionGuarantee,
1107 cost: CostFn,
1108 ) {
1109 self.register_full_with_source(op, dtypes, backend, kernel, caps, precision, cost, "");
1110 }
1111
1112 /// Phase: backend-extensions refactor (2026-06-07). Same as
1113 /// [`Self::register_full`] but tags the binding with a
1114 /// `kernel_source` identifier used to distinguish kernels that
1115 /// register at the same `(op, dtypes, backend)` key from
1116 /// different implementation sources (AOCL vs MKL vs portable
1117 /// CPU, etc.). The tag is diagnostic-only — the picker still
1118 /// distinguishes alternatives by function-pointer identity.
1119 pub fn register_full_with_source(
1120 &mut self,
1121 op: impl Into<BindingKey>,
1122 dtypes: &[DType],
1123 backend: BackendId,
1124 kernel: KernelRef,
1125 caps: KernelCaps,
1126 precision: crate::fused::PrecisionGuarantee,
1127 cost: CostFn,
1128 kernel_source: &'static str,
1129 ) {
1130 // Hand-written registrations make no FKC generic-fallback claim —
1131 // see [`BindingEntry::is_generic`]. The FKC importer uses
1132 // [`Self::register_full_with_source_generic`] to stamp the computed
1133 // bit. Hand-written kernels carry no revision, so the revision axis is
1134 // `UNTRACKED` (`0`); only the FKC importer threads a real revision.
1135 self.register_full_with_source_generic(
1136 op,
1137 dtypes,
1138 backend,
1139 kernel,
1140 caps,
1141 precision,
1142 cost,
1143 kernel_source,
1144 false,
1145 crate::fused::KernelRevisionHash::UNTRACKED.0,
1146 None,
1147 );
1148 }
1149
1150 /// Full-form register with the FKC structural-fallback `is_generic` bit
1151 /// (see [`BindingEntry::is_generic`]). The FKC importer computes the bit
1152 /// once from a kernel's retained `ResolvedLayout` set
1153 /// ([`crate::fkc::is_generic_contract`]) and stamps it here;
1154 /// every hand-written path routes through
1155 /// [`Self::register_full_with_source`], which passes `false`.
1156 ///
1157 /// Otherwise identical to [`Self::register_full_with_source`]:
1158 /// append-only, never-panic; the same `KernelRef` on one key is caught by
1159 /// [`Self::finalize`], not an inline panic.
1160 pub fn register_full_with_source_generic(
1161 &mut self,
1162 op: impl Into<BindingKey>,
1163 dtypes: &[DType],
1164 backend: BackendId,
1165 kernel: KernelRef,
1166 caps: KernelCaps,
1167 precision: crate::fused::PrecisionGuarantee,
1168 cost: CostFn,
1169 kernel_source: &'static str,
1170 is_generic: bool,
1171 kernel_revision_hash: u64,
1172 cost_expr: Option<&'static crate::fkc::CompiledCostExpr>,
1173 ) {
1174 let key = (op.into(), SmallVec::from_slice(dtypes), backend);
1175 let entry = BindingEntry {
1176 kernel,
1177 caps,
1178 precision,
1179 cost,
1180 kernel_source,
1181 is_generic,
1182 kernel_revision_hash,
1183 cost_expr,
1184 };
1185 // Append-only (never-panic): distinct `KernelRef`s compose as
1186 // sibling alternatives at one decision point; the *same*
1187 // function pointer registered twice is NOT rejected inline.
1188 // [`Self::finalize`] detects duplicates after bulk registration
1189 // and returns a `Result`, so the static init boundary can
1190 // fail-fast (`.expect()`) while a dynamic importer surfaces the
1191 // error — honoring never-panic on the import path.
1192 self.bindings.entry(key).or_default().push(entry);
1193 }
1194
1195 /// Validate the fully-assembled table: within the alternative set
1196 /// for any one `(op, dtypes, backend)` key, each `KernelRef`
1197 /// function pointer must appear **at most once**. Distinct
1198 /// alternatives compose as siblings; the same pointer twice is a
1199 /// registration bug (a copy-paste in a hand-written table, or a
1200 /// contract importer wiring one kernel onto two keys that collapse).
1201 ///
1202 /// `register*` is append-only and never panics; this single pass —
1203 /// run once after bulk registration — surfaces a duplicate as a
1204 /// typed `Result`. The static init boundary
1205 /// ([`crate::dispatch::global_bindings`]) `.expect()`s it (fail-fast
1206 /// at startup for hand-written tables); the FKC contract importer
1207 /// will call it with `?`. This is the never-panic replacement for
1208 /// the former inline duplicate `panic!`.
1209 pub fn finalize(&self) -> Result<()> {
1210 for ((op, dtypes, backend), alts) in self.bindings.iter() {
1211 for i in 0..alts.len() {
1212 let ptr_i = alts[i].kernel as *const () as usize;
1213 for entry_j in &alts[i + 1..] {
1214 if entry_j.kernel as *const () as usize == ptr_i {
1215 return Err(Error::Msg(format!(
1216 "KernelBindingTable: duplicate KernelRef registered \
1217 for (op={op:?}, dtypes={dtypes:?}, backend={backend:?}). \
1218 The same function pointer was registered twice; \
1219 distinct alternatives at one decision point must be \
1220 distinct `fn` items."
1221 )));
1222 }
1223 }
1224 }
1225 }
1226 Ok(())
1227 }
1228
1229 /// Phase 7.6 step 7b: upgrade every UNAUDITED-precision CPU
1230 /// registration to the supplied `default`. Convention is to call
1231 /// this at the *end* of bulk registration, so any
1232 /// `register_with_precision(...)` calls that explicitly claimed a
1233 /// non-default value are preserved. The architecture-target
1234 /// shape would be precision-per-call-site, but for the ~335 CPU
1235 /// primitive registrations that all share the deterministic
1236 /// F32-accumulator property, a fill pass keeps the call sites
1237 /// concise without sacrificing the architectural commitment
1238 /// (every entry ends up with an explicit, non-UNAUDITED precision
1239 /// claim before lookup is ever exercised).
1240 ///
1241 /// Only entries with `backend == BackendId::Cpu` and the current
1242 /// `PrecisionGuarantee::UNAUDITED` sentinel are touched. Non-CPU
1243 /// backends register their own precision claims; non-UNAUDITED
1244 /// entries are preserved.
1245 ///
1246 /// Step 9a: applies to **every alternative** registered under a
1247 /// CPU key, not just the first — so a key with N siblings all
1248 /// starting at UNAUDITED ends with N entries upgraded to `default`.
1249 pub fn fill_unset_cpu_precision(&mut self, default: crate::fused::PrecisionGuarantee) {
1250 for ((_, _, backend), alts) in self.bindings.iter_mut() {
1251 if *backend != BackendId::Cpu {
1252 continue;
1253 }
1254 for entry in alts.iter_mut() {
1255 let p = &mut entry.precision;
1256 // Upgrade entries that claim nothing — EXCEPT an audited
1257 // "no static bound applies" conclusion.
1258 //
1259 // GAP-095: the old predicate was purely structural ("all four
1260 // value fields at defaults"), and the comment here asserted
1261 // that "a real weaker claim would set at least one of them."
1262 // That is false for `PrecisionGuarantee::none(reason)`, which
1263 // is field-identical to UNAUDITED and sets none of them. The
1264 // fill therefore overwrote an explicit audit conclusion with
1265 // `bit_stable_on_same_hardware: true` — asserting on the
1266 // kernel's behalf exactly what its audit had declined to
1267 // claim. FKC-4.8-0001 sanctions this fill for blocks left
1268 // absent or `audited: false`; it does not sanction that.
1269 //
1270 // Note what must KEEP being upgraded: a contract claim that
1271 // the V-FKC-9 ledger gate downgraded collapses to the
1272 // UNAUDITED sentinel exactly, so it stays eligible here. That
1273 // is load-bearing — every one of the CPU table's bit-stable
1274 // entries arrives through this path (GAP-077), so excluding
1275 // them would break the FKC-4.8-0001 coverage commitment
1276 // outright.
1277 if p.makes_no_static_claim() && !p.is_audited_no_bound() {
1278 *p = default;
1279 }
1280 }
1281 }
1282 }
1283
1284 /// Phase 7.6 step 8: upgrade every still-`unknown_cost` CPU
1285 /// registration to the cost function returned by `dispatcher` for
1286 /// its OpKind. The dispatcher's contract: given an `OpKind`,
1287 /// return the appropriate cost-family function — typically
1288 /// [`crate::cost::default_cost_for_op_kind`] in production code,
1289 /// but tests may pass their own dispatcher to exercise specific
1290 /// cost shapes.
1291 ///
1292 /// Convention is to call this at the *end* of bulk registration,
1293 /// so any `register_full(...)` calls that explicitly claimed a
1294 /// non-default cost are preserved. The function-pointer equality
1295 /// check identifies UNKNOWN entries — any caller that wants a
1296 /// "weaker" claim with non-zero values must use a distinct
1297 /// function pointer (which won't compare equal to `unknown_cost`
1298 /// and thus won't be overwritten).
1299 ///
1300 /// Step 9a: applies to **every alternative** with a CPU backend —
1301 /// each sibling entry at a key starting at `unknown_cost` is
1302 /// upgraded independently.
1303 pub fn fill_unset_cpu_cost(&mut self, dispatcher: fn(OpKind) -> CostFn) {
1304 self.fill_unset_cost_for_backend(BackendId::Cpu, dispatcher);
1305 }
1306
1307 /// Per-backend variant of [`Self::fill_unset_cpu_cost`]. Walks
1308 /// every alternative whose backend matches `backend` and replaces
1309 /// `unknown_cost` entries with the dispatcher's choice for the
1310 /// op. Used by [`crate::vulkan_dispatch::register_vulkan_kernels`]
1311 /// to bulk-fill cost functions after per-kernel precision
1312 /// registration; future MKL/AOCL/Metal backends register through
1313 /// this same helper.
1314 ///
1315 /// Backend-specific cost adjustments (e.g. higher
1316 /// `kernel_overhead_ns` for Vulkan command-buffer submission vs
1317 /// CPU's nested-loop call overhead) are a Layer-2 refinement —
1318 /// today's `default_cost_for_op_kind` returns the same CPU-flavored
1319 /// functions regardless of backend, which over-estimates Vulkan
1320 /// throughput on small tensors. Either pass a Vulkan-flavored
1321 /// dispatcher to this method, or rely on the empirical calibration
1322 /// framework's per-(op, dtype, backend) refinement.
1323 pub fn fill_unset_cost_for_backend(
1324 &mut self,
1325 backend: BackendId,
1326 dispatcher: fn(OpKind) -> CostFn,
1327 ) {
1328 let sentinel = unknown_cost as *const () as usize;
1329 for ((key, _, this_backend), alts) in self.bindings.iter_mut() {
1330 if *this_backend != backend {
1331 continue;
1332 }
1333 // Runtime-fused entries keep the sentinel: their cost derives from
1334 // the op's recipe (keyed by the runtime id), not an OpKind table.
1335 let BindingKey::Static(op) = key else {
1336 continue;
1337 };
1338 for entry in alts.iter_mut() {
1339 if (entry.cost as usize) == sentinel {
1340 entry.cost = dispatcher(*op);
1341 }
1342 }
1343 }
1344 }
1345
1346 /// Look up the [`crate::fused::PrecisionGuarantee`] for a
1347 /// registered `(op, dtypes, backend)` triple. Returns the **first
1348 /// alternative**'s precision; multi-impl callers wanting all
1349 /// alternatives use [`Self::lookup_alternatives`]. Returns
1350 /// `PrecisionGuarantee::UNAUDITED` if the binding is missing —
1351 /// callers (notably the step-7b lint) treat that as "no claim"
1352 /// rather than panicking.
1353 pub fn lookup_precision(
1354 &self,
1355 op: impl Into<BindingKey>,
1356 dtypes: &[DType],
1357 backend: BackendId,
1358 ) -> crate::fused::PrecisionGuarantee {
1359 let key = (op.into(), SmallVec::from_slice(dtypes), backend);
1360 self.bindings
1361 .get(&key)
1362 .and_then(|alts| alts.first())
1363 .map(|e| e.precision)
1364 .unwrap_or(crate::fused::PrecisionGuarantee::UNAUDITED)
1365 }
1366
1367 /// Look up the [`CostFn`] for a registered `(op, dtypes, backend)`
1368 /// triple. Returns the **first alternative**'s cost; multi-impl
1369 /// callers wanting all alternatives use
1370 /// [`Self::lookup_alternatives`]. Returns [`unknown_cost`] for
1371 /// missing bindings — the step-8 lint treats that as "no cost
1372 /// claim."
1373 pub fn lookup_cost(
1374 &self,
1375 op: impl Into<BindingKey>,
1376 dtypes: &[DType],
1377 backend: BackendId,
1378 ) -> CostFn {
1379 let key = (op.into(), SmallVec::from_slice(dtypes), backend);
1380 self.bindings
1381 .get(&key)
1382 .and_then(|alts| alts.first())
1383 .map(|e| e.cost)
1384 .unwrap_or(unknown_cost)
1385 }
1386
1387 /// Iterate `(op, dtypes, backend, precision)` over every
1388 /// registered alternative — one tuple per `BindingEntry`, not per
1389 /// key. Used by the step-7b coverage lint, which groups by
1390 /// OpKind and checks the bit-stable commitment per group; the
1391 /// grouping handles N-alternatives-per-key naturally.
1392 pub fn iter_precision(
1393 &self,
1394 ) -> impl Iterator<
1395 Item = (
1396 OpKind,
1397 &[DType],
1398 BackendId,
1399 crate::fused::PrecisionGuarantee,
1400 ),
1401 > {
1402 self.bindings
1403 .iter()
1404 .filter_map(|((key, dtypes, backend), alts)| {
1405 // Static-kernel audit view: runtime-fused entries are outside the
1406 // per-OpKind commitments these iterators exist to check.
1407 key.static_op().map(|op| (op, dtypes, backend, alts))
1408 })
1409 .flat_map(|(op, dtypes, backend, alts)| {
1410 alts.iter()
1411 .map(move |e| (op, dtypes.as_slice(), *backend, e.precision))
1412 })
1413 }
1414
1415 /// Iterate `(op, dtypes, backend, &entry)` over every registered
1416 /// alternative — one tuple per `BindingEntry`, mirroring
1417 /// [`Self::iter_precision`] but yielding the FULL entry (by reference)
1418 /// instead of projecting out just the `precision` field. Used by the
1419 /// FKC empirical verification harness (Task 4.4+), which needs the
1420 /// whole `BindingEntry` (the `kernel` fn pointer to invoke, `caps`,
1421 /// `kernel_revision_hash` for ledger keys, ...), not just one field.
1422 /// Runtime-fused/JIT entries are filtered out (`static_op()` returns
1423 /// `None` for them) — same "static-kernel audit view" as
1424 /// `iter_precision`/`iter_cost`.
1425 pub fn iter_entries(
1426 &self,
1427 ) -> impl Iterator<Item = (OpKind, &[DType], BackendId, &BindingEntry)> {
1428 self.bindings
1429 .iter()
1430 .filter_map(|((key, dtypes, backend), alts)| {
1431 key.static_op().map(|op| (op, dtypes, backend, alts))
1432 })
1433 .flat_map(|(op, dtypes, backend, alts)| {
1434 alts.iter()
1435 .map(move |e| (op, dtypes.as_slice(), *backend, e))
1436 })
1437 }
1438
1439 /// Iterate `(op, dtypes, backend, cost_fn)` over every registered
1440 /// alternative — one tuple per `BindingEntry`. Used by the step-8
1441 /// coverage lint to check the "every primitive op has a
1442 /// non-default cost function" commitment per OpKind.
1443 pub fn iter_cost(&self) -> impl Iterator<Item = (OpKind, &[DType], BackendId, CostFn)> {
1444 self.bindings
1445 .iter()
1446 .filter_map(|((key, dtypes, backend), alts)| {
1447 key.static_op().map(|op| (op, dtypes, backend, alts))
1448 })
1449 .flat_map(|(op, dtypes, backend, alts)| {
1450 alts.iter()
1451 .map(move |e| (op, dtypes.as_slice(), *backend, e.cost))
1452 })
1453 }
1454
1455 /// Iterate every registered `(op, dtypes, backend)` decision-point
1456 /// key — one item per key, regardless of how many alternatives
1457 /// share it. Consumed by `SystemTopology::build_at` to derive
1458 /// which backends have kernels registered + which (op, dtype)
1459 /// pairs each backend covers.
1460 pub fn iter_keys(&self) -> impl Iterator<Item = (OpKind, &[DType], BackendId)> {
1461 self.bindings.keys().filter_map(|(key, dtypes, backend)| {
1462 key.static_op().map(|op| (op, dtypes.as_slice(), *backend))
1463 })
1464 }
1465
1466 /// Whether ANY runtime-fused entry is bound for `(fid, backend)`,
1467 /// regardless of dtype tuple — the coarse capability-gate query (the
1468 /// optimizer's fuse/don't-fuse predicate is id-level; the dtype-precise
1469 /// check is the dispatch-time [`Self::lookup_with_caps`] itself).
1470 pub fn has_runtime_fused(&self, fid: FusedOpId, backend: BackendId) -> bool {
1471 self.bindings
1472 .keys()
1473 .any(|(k, _, b)| *b == backend && matches!(k, BindingKey::RuntimeFused(f) if *f == fid))
1474 }
1475
1476 /// Resolve `(fid, backend)` to a runtime-fused binding, **refusing to pick
1477 /// when the answer is ambiguous**.
1478 ///
1479 /// A [`FusedOpId`] names a RECIPE, not a signature: `register_runtime_fused`
1480 /// deduplicates on the region's base-map hash, so one id can legitimately
1481 /// carry several dtype tuples — `[F32, F32]` and `[F32, F32, F32]` are
1482 /// different arities under the same id. The pre-GAP-213 `first_runtime_fused` scanned
1483 /// for `(fid, backend)` and **binds `dtypes` without testing it**, so it
1484 /// answers such a question by `HashMap::iter()` order: a coin flip per
1485 /// process, and capable of returning a kernel of the wrong operand arity.
1486 ///
1487 /// The three outcomes are kept as three values rather than folded into an
1488 /// `Option`, because *not bound* and *ambiguous* are different facts with
1489 /// different fixes, and an `Option` would make them indistinguishable to the
1490 /// caller — the same collapse that let GAP-001 read as a version regression
1491 /// for eleven days.
1492 ///
1493 /// **Scope: this is the diagnostic/test resolution path.** Production
1494 /// dispatch never came through here — it uses [`Self::lookup_with_caps`],
1495 /// which takes the dtype tuple and does an exact keyed `get`, so it can
1496 /// neither be ambiguous nor return a wrong arity. See GAP-213.
1497 pub fn unique_runtime_fused(
1498 &self,
1499 fid: FusedOpId,
1500 backend: BackendId,
1501 ) -> RuntimeFusedLookup<'_> {
1502 let mut matches = self.bindings.iter().filter(|((k, _, b), _)| {
1503 *b == backend && matches!(k, BindingKey::RuntimeFused(f) if *f == fid)
1504 });
1505 let Some(((_, dtypes, _), alts)) = matches.next() else {
1506 return RuntimeFusedLookup::NotBound;
1507 };
1508 if matches.next().is_some() {
1509 let mut signatures: Vec<Vec<DType>> = self
1510 .bindings
1511 .iter()
1512 .filter(|((k, _, b), _)| {
1513 *b == backend && matches!(k, BindingKey::RuntimeFused(f) if *f == fid)
1514 })
1515 .map(|((_, d, _), _)| d.to_vec())
1516 .collect();
1517 // HashMap order is per-process, so the report would otherwise vary run to
1518 // run. `DType` is not `Ord`, so key on arity then the debug spelling —
1519 // a presentation order only; nothing depends on which comes first.
1520 signatures.sort_by_key(|d| (d.len(), format!("{d:?}")));
1521 return RuntimeFusedLookup::Ambiguous { signatures };
1522 }
1523 match alts.first() {
1524 Some(e) => RuntimeFusedLookup::Unique(dtypes.as_slice(), e),
1525 None => RuntimeFusedLookup::NotBound,
1526 }
1527 }
1528
1529 /// **TEST-ONLY.** Drop every runtime-fused entry (static entries retained).
1530 /// The runtime-fused reset story: the metadata sidecar's Vec length is the
1531 /// id allocator, so clearing it MUST clear these bindings in the same
1532 /// breath or a reused id resolves a stale kernel
1533 /// (`runtime_fused_kernels::clear_runtime_fused_for_tests` does both).
1534 #[doc(hidden)]
1535 pub fn remove_runtime_fused_for_tests(&mut self) {
1536 self.bindings
1537 .retain(|(k, _, _), _| matches!(k, BindingKey::Static(_)));
1538 }
1539
1540 /// Look up the wrapper for `(op, dtypes, backend)`. Returns the
1541 /// **first alternative**'s [`KernelRef`]; for capability-aware
1542 /// lookup use [`Self::lookup_with_caps`]; for the full alternative
1543 /// set use [`Self::lookup_alternatives`].
1544 pub fn lookup(
1545 &self,
1546 op: impl Into<BindingKey>,
1547 dtypes: &[DType],
1548 backend: BackendId,
1549 ) -> Result<KernelRef> {
1550 self.lookup_with_caps(op, dtypes, backend).map(|(k, _)| k)
1551 }
1552
1553 /// Capability-aware lookup. Returns the first **dense-safe** alternative
1554 /// — the first sibling whose caps do NOT set
1555 /// [`KernelCaps::requires_broadcast`] — paired with its registered
1556 /// [`KernelCaps`]. A baked-broadcast sibling (`broadcast_stride0:
1557 /// required`) is SKIPPED, so a dense operand is never routed into a
1558 /// kernel that reads `in[0]`; the exclusion is order-independent (path
1559 /// 1a). For every existing binding — none set `requires_broadcast` —
1560 /// this is byte-identical to returning the first-registered alternative.
1561 /// Surfaces [`Error::NoBackendForOp`] on a missing binding OR when every
1562 /// sibling is broadcast-baked (production-correct: a surfaced miss, never
1563 /// a wrong read, no panic).
1564 pub fn lookup_with_caps(
1565 &self,
1566 op: impl Into<BindingKey>,
1567 dtypes: &[DType],
1568 backend: BackendId,
1569 ) -> Result<(KernelRef, KernelCaps)> {
1570 let key = (op.into(), SmallVec::from_slice(dtypes), backend);
1571 self.bindings
1572 .get(&key)
1573 .and_then(|alts| alts.iter().find(|e| !e.caps.requires_broadcast))
1574 .map(|e| (e.kernel, e.caps))
1575 .ok_or_else(|| {
1576 let (available_backends, supported_combinations) = self.coverage_diagnostic();
1577 Error::NoBackendForOp {
1578 op: key.0.static_op().unwrap_or(OpKind::RuntimeFused),
1579 dtypes: dtypes.to_vec(),
1580 available_backends,
1581 supported_combinations,
1582 }
1583 .bt()
1584 })
1585 }
1586
1587 /// Build the truthful `(available_backends, supported_combinations)`
1588 /// diagnostic pair for a [`Error::NoBackendForOp`] miss: the DISTINCT
1589 /// registered backends, and the full static `(backend, op, dtypes)` coverage
1590 /// (the useful "what exists" context). Shared by
1591 /// [`lookup_with_caps`](Self::lookup_with_caps) AND the plan-time miss
1592 /// (`plan::missing_binding_error`) so BOTH surface the same real context. The
1593 /// plan-time path formerly hardcoded `[]` here, which falsely reads as "no
1594 /// backend is registered" and misdirects diagnosis — a mixed-precision
1595 /// `(op, dtypes)` gap is not an empty registry (observed in the wild:
1596 /// `[F32, BF16, F32]` matmul reported as "available backends: []").
1597 pub fn coverage_diagnostic(&self) -> (Vec<BackendId>, Vec<CoverageBinding>) {
1598 let available_backends: Vec<BackendId> = self
1599 .bindings
1600 .keys()
1601 .map(|(_, _, b)| *b)
1602 .collect::<std::collections::HashSet<_>>()
1603 .into_iter()
1604 .collect();
1605 let supported_combinations: Vec<(BackendId, OpKind, Vec<DType>)> = self
1606 .bindings
1607 .keys()
1608 .filter_map(|(k, d, b)| k.static_op().map(|o| (*b, o, d.to_vec())))
1609 .collect();
1610 (available_backends, supported_combinations)
1611 }
1612
1613 /// Does **this specific backend** serve `(op, dtypes)` with a usable
1614 /// (non-broadcast-baked) binding?
1615 ///
1616 /// Mirrors [`lookup_with_caps`](Self::lookup_with_caps)'s success condition
1617 /// exactly — same key, same `!requires_broadcast` admission — but skips
1618 /// constructing the miss diagnostic, so it is cheap enough to call in a
1619 /// per-node loop (`coverage_diagnostic` walks the entire table).
1620 ///
1621 /// Ask this once per *candidate* backend. There is deliberately no
1622 /// "does any backend anywhere serve it?" sibling: a predicate that ignores
1623 /// the key's backend component cannot answer a dispatch question, because a
1624 /// node executes on one backend rather than on the union of them. Such a
1625 /// sibling did exist (`any_backend_supports`) and it caused a real defect —
1626 /// the dtype-reconciliation pass consulted it, so a backend a node could
1627 /// never run on suppressed the `Op::Cast` that node needed and planning
1628 /// dead-ended at `NoBackendForOp`. See `crate::plan::candidate_backends_for`.
1629 pub fn backend_supports(
1630 &self,
1631 op: impl Into<BindingKey>,
1632 dtypes: &[DType],
1633 backend: BackendId,
1634 ) -> bool {
1635 let key = (op.into(), SmallVec::from_slice(dtypes), backend);
1636 self.bindings
1637 .get(&key)
1638 .is_some_and(|alts| alts.iter().any(|e| !e.caps.requires_broadcast))
1639 }
1640
1641 /// Phase 7.6 step 9a: return the full set of registered
1642 /// alternatives at the `(op, dtypes, backend)` decision point.
1643 /// Order is registration order — append-on-register is the step-9a
1644 /// contract. The empty slice means "no kernel registered at this
1645 /// decision point."
1646 ///
1647 /// The route picker (step 9b) consumes this to rank alternatives
1648 /// by cost + precision at plan time. Single-impl callers continue
1649 /// to use [`Self::lookup_with_caps`] (which returns the first
1650 /// alternative — the route picker's job is the longer-term home
1651 /// for selection logic).
1652 pub fn lookup_alternatives(
1653 &self,
1654 op: impl Into<BindingKey>,
1655 dtypes: &[DType],
1656 backend: BackendId,
1657 ) -> &[BindingEntry] {
1658 let key = (op.into(), SmallVec::from_slice(dtypes), backend);
1659 self.bindings
1660 .get(&key)
1661 .map(|alts| alts.as_slice())
1662 .unwrap_or(&[])
1663 }
1664
1665 /// Total number of registered alternatives across all keys (not
1666 /// just unique keys). Step 9a: a key with N siblings counts N.
1667 pub fn len(&self) -> usize {
1668 self.bindings.values().map(|alts| alts.len()).sum()
1669 }
1670
1671 /// Empty binding table.
1672 pub fn is_empty(&self) -> bool {
1673 self.bindings.is_empty()
1674 }
1675
1676 /// Total number of unique `(op, dtypes, backend)` keys. Step 9a:
1677 /// distinct from [`Self::len`], which counts alternatives.
1678 pub fn key_count(&self) -> usize {
1679 self.bindings.len()
1680 }
1681}
1682
1683impl std::fmt::Debug for KernelBindingTable {
1684 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1685 f.debug_struct("KernelBindingTable")
1686 .field("keys", &self.bindings.len())
1687 .field("total_alternatives", &self.len())
1688 .finish_non_exhaustive()
1689 }
1690}
1691
1692#[cfg(test)]
1693mod tests {
1694 use super::*;
1695 use fuel_ir::DType;
1696
1697 /// Smoke: KernelRef can be constructed and stored.
1698 #[test]
1699 fn kernel_ref_stores_function_pointer() {
1700 fn dummy_kernel(
1701 _inputs: &[Arc<RwLock<Storage>>],
1702 _outputs: &mut [Arc<RwLock<Storage>>],
1703 _layouts: &[Layout],
1704 _params: &OpParams,
1705 ) -> Result<()> {
1706 Ok(())
1707 }
1708 let k: KernelRef = dummy_kernel;
1709 // Function pointer is Copy + Clone.
1710 let _k2 = k;
1711 let _k3: KernelRef = k;
1712 }
1713
1714 /// Smoke: OpParams variants construct cleanly and Debug-format.
1715 #[test]
1716 fn op_params_variants_construct() {
1717 let _ = OpParams::None;
1718 let _ = OpParams::Reduce {
1719 dims: vec![0, 1],
1720 keepdim: false,
1721 };
1722 let _ = OpParams::Matmul {
1723 lhs_batch_dims: vec![],
1724 rhs_batch_dims: vec![],
1725 m: 4,
1726 n: 8,
1727 k: 16,
1728 m_compute: MatmulM::All,
1729 };
1730 let _ = OpParams::Slice {
1731 dim: 0,
1732 start: 0,
1733 end: 10,
1734 step: 1,
1735 };
1736 let _ = OpParams::Cast;
1737 let _ = OpParams::Affine { mul: 2.0, add: 1.0 };
1738 // Debug format compiles.
1739 let _ = format!("{:?}", OpParams::None);
1740 }
1741
1742 /// Smoke: a hand-constructed kernel that allocates the output
1743 /// Storage outside this crate would be:
1744 /// 1. allocate output via fuel_memory::alloc_cpu_zeroed
1745 /// 2. wrap inputs as &[Arc<RwLock<Storage>>]
1746 /// 3. call the kernel
1747 ///
1748 /// Phase B5 lands the first such migration; B1 just type-checks
1749 /// the surface.
1750 #[test]
1751 fn kernel_ref_can_be_called() {
1752 fn copy_kernel(
1753 inputs: &[Arc<RwLock<Storage>>],
1754 outputs: &mut [Arc<RwLock<Storage>>],
1755 _layouts: &[Layout],
1756 _params: &OpParams,
1757 ) -> Result<()> {
1758 // Simplest real kernel: copy bytes from inputs[0] to outputs[0].
1759 let in_arc = &inputs[0];
1760 let out_arc = &outputs[0];
1761 let in_guard = in_arc.read().unwrap();
1762 let mut out_guard = out_arc.write().unwrap();
1763 let in_bytes = match &in_guard.inner {
1764 fuel_memory::BackendStorage::Cpu(s) => s.bytes(),
1765 #[allow(unreachable_patterns)]
1766 _ => return Ok(()),
1767 };
1768 if let fuel_memory::BackendStorage::Cpu(s) = &mut out_guard.inner {
1769 s.bytes_mut().copy_from_slice(in_bytes);
1770 }
1771 Ok(())
1772 }
1773
1774 let input = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
1775 let output = fuel_memory::alloc_cpu_zeroed(DType::F32, 4).unwrap();
1776 let inputs = vec![Arc::new(RwLock::new(input))];
1777 let mut outputs = vec![Arc::new(RwLock::new(output))];
1778
1779 let k: KernelRef = copy_kernel;
1780 k(&inputs, &mut outputs, &[], &OpParams::None).unwrap();
1781
1782 // Output bytes match input.
1783 let out_guard = outputs[0].read().unwrap();
1784 if let fuel_memory::BackendStorage::Cpu(s) = &out_guard.inner {
1785 let typed: &[f32] = s.as_slice().unwrap();
1786 assert_eq!(typed, &[1.0, 2.0, 3.0, 4.0]);
1787 }
1788 }
1789
1790 fn ok_kernel(
1791 _inputs: &[Arc<RwLock<Storage>>],
1792 _outputs: &mut [Arc<RwLock<Storage>>],
1793 _layouts: &[Layout],
1794 _params: &OpParams,
1795 ) -> Result<()> {
1796 Ok(())
1797 }
1798
1799 fn ok_kernel_alt(
1800 _inputs: &[Arc<RwLock<Storage>>],
1801 _outputs: &mut [Arc<RwLock<Storage>>],
1802 _layouts: &[Layout],
1803 _params: &OpParams,
1804 ) -> Result<()> {
1805 Ok(())
1806 }
1807
1808 /// Phase 7.6 step 9a: two distinct `KernelRef`s register against
1809 /// the same `(op, dtypes, backend)` key as sibling alternatives.
1810 /// `lookup_with_caps` returns the first; `lookup_alternatives`
1811 /// returns both in registration order.
1812 #[test]
1813 fn step_9a_two_alternatives_append_then_first_wins_on_legacy_lookup() {
1814 use fuel_ir::probe::BackendId;
1815 let mut table = KernelBindingTable::new();
1816 let dts = [DType::BF16, DType::BF16, DType::BF16];
1817 table.register(OpKind::MatMul, &dts, BackendId::Cuda, ok_kernel);
1818 table.register(OpKind::MatMul, &dts, BackendId::Cuda, ok_kernel_alt);
1819
1820 // Legacy single-impl lookup returns the first-registered.
1821 let (k1, _caps) = table
1822 .lookup_with_caps(OpKind::MatMul, &dts, BackendId::Cuda)
1823 .unwrap();
1824 assert_eq!(k1 as *const () as usize, ok_kernel as *const () as usize);
1825
1826 // Multi-impl lookup returns both in registration order.
1827 let alts = table.lookup_alternatives(OpKind::MatMul, &dts, BackendId::Cuda);
1828 assert_eq!(alts.len(), 2);
1829 assert_eq!(
1830 alts[0].kernel as *const () as usize,
1831 ok_kernel as *const () as usize
1832 );
1833 assert_eq!(
1834 alts[1].kernel as *const () as usize,
1835 ok_kernel_alt as *const () as usize
1836 );
1837
1838 // Step-9a accounting helpers stay consistent.
1839 assert_eq!(table.key_count(), 1, "one decision point");
1840 assert_eq!(table.len(), 2, "two alternatives total");
1841 }
1842
1843 /// FKC gap-closure Task 2.1: hand-written registrations (the
1844 /// `register`/`register_full*` family that never touches the FKC
1845 /// importer) carry no declared cost-expression AST — `cost_expr`
1846 /// defaults to `None`. Only the FKC importer (a later Group-2 task)
1847 /// stamps a real `Some(&CompiledCostExpr)`.
1848 #[test]
1849 fn cost_expr_field_defaults_none_for_hand_written_registration() {
1850 use fuel_ir::probe::BackendId;
1851 let mut table = KernelBindingTable::new();
1852 let dts = [DType::F32, DType::F32, DType::F32];
1853 table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
1854 let alts = table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu);
1855 assert_eq!(alts.len(), 1);
1856 assert!(
1857 alts[0].cost_expr.is_none(),
1858 "hand-written registrations carry no declared cost AST"
1859 );
1860 }
1861
1862 /// Task 4.4: `iter_entries` mirrors `iter_precision`/`iter_cost` —
1863 /// yields one `(op, dtypes, backend, &entry)` tuple per registered
1864 /// alternative (both siblings at a shared key show up), and filters
1865 /// out `RuntimeFused` keys just like the precision/cost audit
1866 /// iterators (a "static-kernel audit view").
1867 #[test]
1868 fn iter_entries_yields_full_entries_and_filters_runtime_fused() {
1869 use fuel_ir::probe::BackendId;
1870 let mut table = KernelBindingTable::new();
1871 let dts = [DType::F32, DType::F32, DType::F32];
1872 table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
1873 table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel_alt);
1874 // A runtime-fused key — must be excluded from `iter_entries`.
1875 table.register(
1876 BindingKey::RuntimeFused(FusedOpId(FusedOpId::RUNTIME_FUSED_BASE)),
1877 &dts,
1878 BackendId::Cpu,
1879 ok_kernel,
1880 );
1881
1882 let entries: Vec<_> = table.iter_entries().collect();
1883 assert_eq!(
1884 entries.len(),
1885 2,
1886 "both AddElementwise alternatives, RuntimeFused excluded"
1887 );
1888 for (op, dtypes, backend, entry) in &entries {
1889 assert_eq!(*op, OpKind::AddElementwise);
1890 assert_eq!(*dtypes, dts.as_slice());
1891 assert_eq!(*backend, BackendId::Cpu);
1892 // Full entry is reachable, not just one projected field.
1893 let _ = entry.caps;
1894 let _ = entry.kernel_revision_hash;
1895 }
1896 let fns: Vec<usize> = entries
1897 .iter()
1898 .map(|(_, _, _, e)| e.kernel as *const () as usize)
1899 .collect();
1900 assert!(fns.contains(&(ok_kernel as *const () as usize)));
1901 assert!(fns.contains(&(ok_kernel_alt as *const () as usize)));
1902 }
1903
1904 /// Never-panic (option B): registering the same `KernelRef` twice
1905 /// against one key is append-only — it does NOT panic at
1906 /// registration. `finalize()` detects the duplicate and returns
1907 /// `Err`. (Replaces the former `#[should_panic]` guard.)
1908 #[test]
1909 /// FKC-CLAUSE: FKC-10.10-0001 — the boundary itself: registration stays
1910 /// infallible and `finalize()` RETURNS the duplicate as an error rather than
1911 /// panicking. This is why the clause does not require a `Result` signature on
1912 /// `register_full_with_source`.
1913 fn step_9a_duplicate_kernel_ref_detected_by_finalize_not_panic() {
1914 use fuel_ir::probe::BackendId;
1915 let mut table = KernelBindingTable::new();
1916 let dts = [DType::F32, DType::F32, DType::F32];
1917 table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel);
1918 table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel); // dup — must NOT panic
1919 let err = table
1920 .finalize()
1921 .expect_err("finalize must report the duplicate KernelRef");
1922 assert!(
1923 format!("{err}").contains("duplicate KernelRef"),
1924 "unexpected error: {err}"
1925 );
1926 }
1927
1928 /// Distinct sibling kernels under one key finalize cleanly.
1929 #[test]
1930 fn finalize_ok_for_distinct_sibling_kernels() {
1931 use fuel_ir::probe::BackendId;
1932 let mut table = KernelBindingTable::new();
1933 let dts = [DType::F32, DType::F32, DType::F32];
1934 table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel);
1935 table.register(OpKind::MatMul, &dts, BackendId::Cpu, ok_kernel_alt);
1936 table
1937 .finalize()
1938 .expect("distinct sibling kernels must finalize OK");
1939 }
1940
1941 /// Path 1a: a `broadcast_stride0: required` (baked-broadcast) kernel
1942 /// reads its operand as stride-0; a dense operand routed there
1943 /// miscomputes (reads `in[0]` repeatedly). The realize pick MUST skip
1944 /// such a sibling when a generic one exists, and — the load-bearing
1945 /// part — that exclusion must NOT rest on the generic being registered
1946 /// first. Register the broadcast-baked one FIRST (adversarial: `first()`
1947 /// would return it) and assert the pick returns the GENERIC sibling.
1948 #[test]
1949 fn lookup_with_caps_skips_broadcast_baked_sibling_regardless_of_order() {
1950 use fuel_ir::probe::BackendId;
1951 let mut table = KernelBindingTable::new();
1952 let dts = [DType::F32, DType::F32, DType::F32];
1953 // Baked-broadcast sibling registered FIRST.
1954 table.register_with_caps(
1955 OpKind::AddElementwise,
1956 &dts,
1957 BackendId::Cpu,
1958 ok_kernel,
1959 KernelCaps {
1960 strided_input: true,
1961 requires_broadcast: true,
1962 },
1963 );
1964 // Generic dense-safe sibling registered SECOND.
1965 table.register_with_caps(
1966 OpKind::AddElementwise,
1967 &dts,
1968 BackendId::Cpu,
1969 ok_kernel_alt,
1970 KernelCaps {
1971 strided_input: true,
1972 requires_broadcast: false,
1973 },
1974 );
1975
1976 let (picked, caps) = table
1977 .lookup_with_caps(OpKind::AddElementwise, &dts, BackendId::Cpu)
1978 .expect("a dense-safe sibling exists");
1979 assert!(
1980 !caps.requires_broadcast,
1981 "the realize pick must never be a broadcast-baked kernel",
1982 );
1983 assert_eq!(
1984 picked as *const () as usize, ok_kernel_alt as *const () as usize,
1985 "lookup_with_caps must skip the broadcast-baked sibling and return the generic \
1986 one regardless of registration order",
1987 );
1988 // Both siblings coexist — the exclusion is at the pick, not at
1989 // registration (the baked-broadcast contract is still PRESENT).
1990 let alts = table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu);
1991 assert_eq!(
1992 alts.len(),
1993 2,
1994 "both siblings registered; neither is dropped"
1995 );
1996 }
1997
1998 /// The corner that keeps the exclusion honest: if EVERY sibling is
1999 /// broadcast-baked (no dense-safe fallback), the pick is a surfaced
2000 /// miss (`NoBackendForOp`), NOT a silent wrong dense read.
2001 #[test]
2002 fn lookup_with_caps_misses_when_only_broadcast_baked_siblings_exist() {
2003 use fuel_ir::probe::BackendId;
2004 let mut table = KernelBindingTable::new();
2005 let dts = [DType::F32, DType::F32, DType::F32];
2006 table.register_with_caps(
2007 OpKind::AddElementwise,
2008 &dts,
2009 BackendId::Cpu,
2010 ok_kernel,
2011 KernelCaps {
2012 strided_input: true,
2013 requires_broadcast: true,
2014 },
2015 );
2016 assert!(
2017 table
2018 .lookup_with_caps(OpKind::AddElementwise, &dts, BackendId::Cpu)
2019 .is_err(),
2020 "no dense-safe sibling ⇒ honest miss, never a wrong read",
2021 );
2022 }
2023
2024 /// Phase 7.6 step 9a: `fill_unset_cpu_precision` upgrades every
2025 /// alternative under a CPU key, not just the first.
2026 #[test]
2027 fn step_9a_fill_passes_touch_every_alternative() {
2028 use crate::fused::PrecisionGuarantee;
2029 use fuel_ir::probe::BackendId;
2030 let mut table = KernelBindingTable::new();
2031 let dts = [DType::F32, DType::F32, DType::F32];
2032 table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
2033 table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel_alt);
2034
2035 // Both start at UNAUDITED.
2036 for e in table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu) {
2037 assert!(!e.precision.bit_stable_on_same_hardware);
2038 }
2039
2040 let bit_stable = PrecisionGuarantee {
2041 bit_stable_on_same_hardware: true,
2042 max_ulp: Some(0),
2043 max_relative: None,
2044 max_absolute: None,
2045 notes: "test default",
2046 };
2047 table.fill_unset_cpu_precision(bit_stable);
2048
2049 // Both got upgraded.
2050 let alts = table.lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu);
2051 assert_eq!(alts.len(), 2);
2052 for e in alts {
2053 assert!(e.precision.bit_stable_on_same_hardware);
2054 assert_eq!(e.precision.max_ulp, Some(0));
2055 }
2056 }
2057
2058 /// GAP-095: the CPU bulk fill must NOT overwrite an audited
2059 /// "no static bound applies" conclusion.
2060 ///
2061 /// `PrecisionGuarantee::none(reason)` and `PrecisionGuarantee::UNAUDITED`
2062 /// carry identical *values* — all four claim fields at their defaults —
2063 /// and differ only in `notes`. The fill's predicate was purely
2064 /// structural, so it conflated them and rewrote an explicit audit
2065 /// conclusion into `PRIMITIVE_DETERMINISTIC_CPU`: it asserted
2066 /// `bit_stable_on_same_hardware: true` on a kernel whose audit had
2067 /// specifically declined to claim it.
2068 ///
2069 /// FKC-4.8-0001 sanctions the fill for precision blocks left *absent* or
2070 /// `audited: false`. It does not sanction overwriting an `audited: true`
2071 /// + `none(reason)` conclusion.
2072 ///
2073 /// **Both directions asserted.** A test that only checked that
2074 /// `none(reason)` survives would also pass with the fill disabled
2075 /// outright, so the UNAUDITED sibling must still be upgraded in the same
2076 /// run — otherwise the test proves nothing about discrimination.
2077 #[test]
2078 fn gap_095_fill_preserves_audited_none_but_still_upgrades_unaudited() {
2079 use crate::fused::PrecisionGuarantee;
2080 use fuel_ir::probe::BackendId;
2081 let mut table = KernelBindingTable::new();
2082 let dts = [DType::F32, DType::F32, DType::F32];
2083
2084 const REASON: &str = "audited: accumulation order is scheduler-dependent; \
2085 no static bound applies";
2086 // The audited conclusion — must survive untouched.
2087 table.register_with_precision(
2088 OpKind::MatMul,
2089 &dts,
2090 BackendId::Cpu,
2091 ok_kernel,
2092 PrecisionGuarantee::none(REASON),
2093 );
2094 // The never-audited placeholder — must still be upgraded.
2095 table.register(OpKind::AddElementwise, &dts, BackendId::Cpu, ok_kernel);
2096
2097 table.fill_unset_cpu_precision(PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU);
2098
2099 let audited = table
2100 .lookup_alternatives(OpKind::MatMul, &dts, BackendId::Cpu)
2101 .first()
2102 .expect("audited-none entry registered")
2103 .precision;
2104 assert!(
2105 !audited.bit_stable_on_same_hardware,
2106 "the fill asserted bit-stability on a kernel whose audit \
2107 explicitly concluded that no static bound applies",
2108 );
2109 assert_eq!(
2110 audited.notes, REASON,
2111 "the audit reason was overwritten — the reasoning that lived at \
2112 the registration site is gone",
2113 );
2114
2115 // Opposite direction: the fill must still do its job.
2116 let placeholder = table
2117 .lookup_alternatives(OpKind::AddElementwise, &dts, BackendId::Cpu)
2118 .first()
2119 .expect("unaudited entry registered")
2120 .precision;
2121 assert!(
2122 placeholder.bit_stable_on_same_hardware,
2123 "control FAILED: the UNAUDITED placeholder was not upgraded, so \
2124 this test would pass with the fill disabled entirely and proves \
2125 nothing about the discrimination it is meant to check",
2126 );
2127 }
2128}