Skip to main content

CostFn

Type Alias CostFn 

Source
pub type CostFn = fn(&[Shape], &[DType], &OpParams, &BackendCapabilities) -> CostEstimate;
Expand description

Maps (OpKind, KernelDTypes, BackendId) triples to dispatch wrapper functions, where KernelDTypes lists per-operand dtypes (inputs in order, then outputs). Built once at backend registration time (typically a process-wide OnceLock though that’s not enforced here for testability), consulted at execute time via lookup.

Same-dtype ops register [T, T, ..., T] for the right operand count; mixed-precision ops (Cast, future F32×BF16→F32 matmul) register the exact combo. Variadic ops (Concat) register a canonical short shape [T, T] (one input dtype + output dtype); the lookup site for those ops collapses its dtypes vector to the same shorthand.

Backends register their wrappers via the dispatch::register_* functions in this crate (e.g. crate::dispatch::register_cpu_kernels). Cost-fn signature for primitive-op registrations stored in KernelBindingTable. Mirrors the fused-op cost-fn shape but takes OpParams (the binding-table param payload) instead of FusedOpParams.

Implementations return a crate::fused::CostEstimate computed statically from shapes + dtypes + op-specific params + backend capabilities. The architecture’s Layer-1 (FLOP-count + bandwidth) cost model lives here; Layer-2 empirical refinement composes on top via the telemetry framework, not by changing this signature.