use std::collections::HashMap;
use fuel_graph::{Graph, NodeId};
use fuel_ir::dispatch::{OpKind, SizeClass};
use fuel_ir::probe::BackendId;
use fuel_ir::{DType, DeviceLocation, Error, Result, Shape};
use smallvec::SmallVec;
use crate::kernel::KernelBindingTable;
use crate::pipelined::{build_lookup_dtypes, op_to_op_kind};
use crate::ranker::{
AlternativeSet, CapabilitiesLookup, ChainInput, DecisionContext, FilterContext, JudgeOracle,
KEEP_PER_DEVICE, PlacementDp, PrecisionRequirement, TransferEstimator, apply_filter_chain,
apply_inbound_transfer_costs, composite_ns, compute_static_costs, default_backend_rates,
default_chain, enumerate_candidates,
};
#[derive(Debug)]
pub struct ExecutionPlan {
pub order: Vec<NodeId>,
pub alternatives: HashMap<NodeId, AlternativeSet>,
pub generation: u64,
}
impl ExecutionPlan {
pub fn empty() -> Self {
Self {
order: Vec::new(),
alternatives: HashMap::new(),
generation: 0,
}
}
pub fn alternatives(&self, node: NodeId) -> Option<&AlternativeSet> {
self.alternatives.get(&node)
}
pub fn alternatives_mut(&mut self, node: NodeId) -> Option<&mut AlternativeSet> {
self.alternatives.get_mut(&node)
}
}
pub fn default_device_for(backend: BackendId) -> DeviceLocation {
match backend {
BackendId::Cpu => DeviceLocation::Cpu,
BackendId::Cuda => DeviceLocation::Cuda { gpu_id: 0 },
BackendId::Vulkan => DeviceLocation::Vulkan { gpu_id: 0 },
BackendId::Metal => DeviceLocation::Metal { gpu_id: 0 },
_ => DeviceLocation::Cpu,
}
}
pub(crate) fn decision_device_for(
graph: &Graph,
options: &PlanOptions<'_>,
id: NodeId,
) -> Option<DeviceLocation> {
graph
.placement(id)
.or(options.pinned_device)
.or_else(|| graph.target_backend(id).map(default_device_for))
}
pub(crate) fn candidate_backends_for(
graph: &Graph,
options: &PlanOptions<'_>,
id: NodeId,
node: &fuel_graph::Node,
) -> Vec<BackendId> {
let Some(target_device) = decision_device_for(graph, options, id) else {
return Vec::new();
};
let mut backends: Vec<BackendId> = match options.placements_for_device {
Some(f) => f(target_device),
None => vec![
graph
.target_backend(id)
.unwrap_or_else(|| backend_for_device(target_device)),
],
};
if node.op.destructive_input().is_none()
&& let Some(fallback) = options.fallback_placements_for
{
backends.extend(fallback(target_device).into_iter().map(|(b, _)| b));
}
let mut unique: Vec<BackendId> = Vec::with_capacity(backends.len());
for b in backends {
if !unique.contains(&b) {
unique.push(b);
}
}
unique
}
pub fn backend_for_device(loc: DeviceLocation) -> BackendId {
match loc {
DeviceLocation::Cpu => BackendId::Cpu,
DeviceLocation::Cuda { .. } => BackendId::Cuda,
DeviceLocation::Vulkan { .. } => BackendId::Vulkan,
DeviceLocation::Metal { .. } => BackendId::Metal,
}
}
pub struct PlanOptions<'env> {
pub precision_requirement: PrecisionRequirement,
pub populate_costs: bool,
pub placements_for_device: Option<&'env (dyn Fn(DeviceLocation) -> Vec<BackendId> + 'env)>,
pub pinned_device: Option<DeviceLocation>,
pub fallback_placements_for:
Option<&'env (dyn Fn(DeviceLocation) -> Vec<(BackendId, DeviceLocation)> + 'env)>,
pub capabilities_for: Option<&'env CapabilitiesLookup<'env>>,
pub judge: Option<&'env dyn JudgeOracle>,
pub transfer_estimator: Option<&'env dyn TransferEstimator>,
pub require_transfer_pricing: bool,
pub input_residency: Option<&'env (dyn Fn(NodeId) -> Option<DeviceLocation> + 'env)>,
pub reuse_plan: Option<&'env ExecutionPlan>,
pub allow_cost_placement: bool,
#[cfg(feature = "telemetry")]
pub telemetry: Option<&'env crate::telemetry::TelemetryHooks<'env>>,
}
impl Default for PlanOptions<'_> {
fn default() -> Self {
Self {
precision_requirement: PrecisionRequirement::default(),
populate_costs: true,
placements_for_device: None,
pinned_device: None,
fallback_placements_for: None,
capabilities_for: None,
judge: None,
transfer_estimator: None,
require_transfer_pricing: false,
input_residency: None,
reuse_plan: None,
allow_cost_placement: true,
#[cfg(feature = "telemetry")]
telemetry: None,
}
}
}
impl<'env> PlanOptions<'env> {
pub fn new() -> Self {
Self::default()
}
pub fn with_precision_requirement(mut self, req: PrecisionRequirement) -> Self {
self.precision_requirement = req;
self
}
pub fn without_cost_population(mut self) -> Self {
self.populate_costs = false;
self
}
pub fn with_placements_for_device(
mut self,
f: &'env (dyn Fn(DeviceLocation) -> Vec<BackendId> + 'env),
) -> Self {
self.placements_for_device = Some(f);
self
}
pub fn with_pinned_device(mut self, device: DeviceLocation) -> Self {
self.pinned_device = Some(device);
self
}
pub fn without_cost_placement(mut self) -> Self {
self.allow_cost_placement = false;
self
}
pub fn require_transfer_pricing(mut self) -> Self {
self.require_transfer_pricing = true;
self
}
pub fn with_fallback_placements_for(
mut self,
f: &'env (dyn Fn(DeviceLocation) -> Vec<(BackendId, DeviceLocation)> + 'env),
) -> Self {
self.fallback_placements_for = Some(f);
self
}
pub fn with_capabilities_for(mut self, f: &'env CapabilitiesLookup<'env>) -> Self {
self.capabilities_for = Some(f);
self
}
pub fn with_judge(mut self, judge: &'env dyn JudgeOracle) -> Self {
self.judge = Some(judge);
self
}
pub fn with_transfer_estimator(mut self, est: &'env dyn TransferEstimator) -> Self {
self.transfer_estimator = Some(est);
self
}
pub fn with_input_residency(
mut self,
f: &'env (dyn Fn(NodeId) -> Option<DeviceLocation> + 'env),
) -> Self {
self.input_residency = Some(f);
self
}
pub fn with_reuse_plan(mut self, base: &'env ExecutionPlan) -> Self {
self.reuse_plan = Some(base);
self
}
#[cfg(feature = "telemetry")]
pub fn with_telemetry(mut self, hooks: &'env crate::telemetry::TelemetryHooks<'env>) -> Self {
self.telemetry = Some(hooks);
self
}
}
pub fn compile_plan(
graph: &Graph,
order: &[NodeId],
bindings_table: &KernelBindingTable,
options: &PlanOptions<'_>,
) -> Result<ExecutionPlan> {
if options.require_transfer_pricing && options.transfer_estimator.is_none() {
return Err(Error::Msg(
"compile_plan: require_transfer_pricing is set but no transfer \
estimator is wired — placement would silently price transfers at \
zero (latency-greedy). Wire PlanOptions::with_transfer_estimator(..) \
or drop the pricing requirement."
.into(),
)
.bt());
}
let generation = crate::dispatch::topology_generation();
let mut alternatives_map = HashMap::with_capacity(order.len());
let mut residency: HashMap<NodeId, DeviceLocation> = HashMap::with_capacity(order.len());
let mut dp = PlacementDp::new();
let mut dp_drafts: HashMap<NodeId, NodeDraft> = HashMap::new();
for &id in order {
let node = graph.node(id);
if let Some(base) = options.reuse_plan
&& let Some(set) = base.alternatives(id)
{
let resident = early_residency(graph, options, id, node)
.or_else(|| set.winner().map(|w| w.device));
if let Some(loc) = resident {
residency.insert(id, loc);
}
alternatives_map.insert(id, set.clone());
continue;
}
let early = early_residency(graph, options, id, node);
let mut dp_row_opened = false;
let winner_device = match build_node_draft(graph, id, node, bindings_table, options)? {
Some(draft) => {
let devices = distinct_devices(&draft.set);
let dp_estimator = options
.transfer_estimator
.filter(|_| draft.relaxed && early.is_none() && devices.len() >= 2);
if let Some(est) = dp_estimator {
open_dp_row(
&mut dp,
graph,
id,
node,
&draft.set,
&devices,
&residency,
est,
options.capabilities_for,
);
dp_drafts.insert(id, draft);
dp_row_opened = true;
None
} else {
if let Some(est) = options.transfer_estimator
&& let [single] = devices.as_slice()
{
close_open_inputs(&mut dp, graph, node, *single, est, &mut residency);
}
let set = finalize_node_greedy(
graph,
node,
draft,
bindings_table,
options,
&residency,
)?;
let dev = set.winner().map(|w| w.device);
if devices.len() > 1
&& let (Some(est), Some(wd)) = (options.transfer_estimator, dev)
{
close_open_inputs(&mut dp, graph, node, wd, est, &mut residency);
}
alternatives_map.insert(id, set);
dev
}
}
None => {
if let fuel_graph::Op::Copy { target } | fuel_graph::Op::Move { target } = node.op
&& let Some(est) = options.transfer_estimator
{
close_open_inputs(&mut dp, graph, node, target, est, &mut residency);
}
None
}
};
if dp_row_opened {
continue;
}
let resident = early
.or(winner_device)
.or_else(|| {
graph.target_backend(id).map(|b| {
options
.pinned_device
.unwrap_or_else(|| default_device_for(b))
})
})
.or_else(|| {
if is_passthrough(node) {
node.inputs.first().and_then(|i| residency.get(i)).copied()
} else {
None
}
});
if let Some(loc) = resident {
residency.insert(id, loc);
} else if is_passthrough(node) {
if let Some(&input) = node.inputs.first() {
let root = dp.resolve(input);
if let Some(dev) = dp.committed_device(root) {
residency.insert(id, dev);
} else {
dp.add_alias(id, root);
}
}
}
}
if !dp_drafts.is_empty() {
let Some(est) = options.transfer_estimator else {
return Err(Error::Msg(
"compile_plan: placement DP rows exist without a transfer \
estimator — internal invariant violated"
.into(),
)
.bt());
};
let commits = dp.finish(options.pinned_device, |nid| node_bytes(graph, nid), est);
for (n, d) in commits {
residency.insert(n, d);
}
for &id in order {
if residency.contains_key(&id) {
continue;
}
let node = graph.node(id);
if is_passthrough(node)
&& let Some(&loc) = node.inputs.first().and_then(|i| residency.get(i))
{
residency.insert(id, loc);
}
}
for &id in order {
let Some(draft) = dp_drafts.remove(&id) else {
continue;
};
let chosen = dp.committed_device(id).ok_or_else(|| {
Error::Msg(format!(
"compile_plan: placement DP left node {:?} uncommitted — \
internal invariant violated",
id,
))
.bt()
})?;
let node = graph.node(id);
let mut set = draft.set;
let inputs = priced_inputs_for(graph, node, &residency);
apply_inbound_transfer_costs(&mut set, &inputs, est);
set.rank_by_composite_cost();
let keep: Vec<usize> = set
.alternatives()
.iter()
.enumerate()
.filter_map(|(i, c)| (c.device == chosen).then_some(i))
.collect();
if keep.len() < set.len() {
set.retain_indices(&keep);
}
set.retain_per_device_frontier(KEEP_PER_DEVICE);
if set.is_empty() {
return Err(missing_binding_error(
bindings_table,
draft.op_kind,
&draft.dtypes,
draft.diag_backend,
));
}
alternatives_map.insert(id, set);
}
}
#[cfg(feature = "telemetry")]
emit_plan_telemetry(graph, &alternatives_map, bindings_table, options);
Ok(ExecutionPlan {
order: order.to_vec(),
alternatives: alternatives_map,
generation,
})
}
#[cfg(feature = "telemetry")]
fn emit_plan_telemetry(
graph: &Graph,
alternatives: &HashMap<NodeId, AlternativeSet>,
bindings_table: &KernelBindingTable,
options: &PlanOptions<'_>,
) {
use crate::telemetry::{
DispatchRecord, FdxOperandDesc, ImplId, TELEMETRY_SCHEMA_VERSION, detect_miss_precomputed,
};
let Some(hooks) = options.telemetry else {
return;
};
if !hooks.config.is_enabled() {
return;
}
let arch = hooks.arch_tag();
let want_candidates = hooks.config.mode.wants_candidates();
for (&id, set) in alternatives {
let Some(winner) = set.winner() else {
continue;
};
let node = graph.node(id);
let Some(op_kind) = op_to_op_kind(&node.op) else {
continue;
};
let dtypes = build_lookup_dtypes(graph, node);
let winner_ptr = winner.kernel as *const () as usize;
let winner_entry = bindings_table
.lookup_alternatives(op_kind, &dtypes, winner.backend)
.iter()
.find(|e| e.kernel as *const () as usize == winner_ptr)
.copied();
let is_generic = winner_entry.is_some_and(|e| e.is_generic);
let winner_revision = winner_entry.map_or(0, |e| e.kernel_revision_hash);
let mut operands: Vec<FdxOperandDesc> = node
.inputs
.iter()
.map(|&input_id| {
FdxOperandDesc::from_layout(&graph.layout(input_id), graph.node(input_id).dtype)
})
.collect();
operands.push(FdxOperandDesc::from_layout(&graph.layout(id), node.dtype));
let op_class = format!("{op_kind:?}");
let structure_key = hooks.provider.structure_key(&op_class, &operands, &arch);
let chosen = ImplId {
backend: winner.backend,
op: op_kind,
dtypes: dtypes.to_vec(),
kernel_source: winner.kernel_source.to_string(),
kernel_revision_hash: winner_revision,
};
let candidates = if want_candidates {
build_candidates(set, op_kind, &dtypes, bindings_table, options.judge)
} else {
Vec::new()
};
let dispatch = DispatchRecord {
schema: TELEMETRY_SCHEMA_VERSION,
structure_key,
chosen: chosen.clone(),
candidates,
count: 1,
hw: hooks.hw.clone(),
};
if let Ok(mut sink) = hooks.sink.lock() {
sink.record_dispatch(dispatch);
}
if let Some(miss) = detect_miss_precomputed(
is_generic,
chosen,
&op_class,
&operands,
&arch,
hooks.provider,
hooks.hw.clone(),
) {
if let Ok(mut sink) = hooks.sink.lock() {
sink.record_miss(miss);
}
}
}
}
#[cfg(feature = "telemetry")]
fn build_candidates(
set: &AlternativeSet,
op_kind: OpKind,
dtypes: &[DType],
bindings_table: &KernelBindingTable,
judge: Option<&dyn JudgeOracle>,
) -> Vec<crate::telemetry::Candidate> {
use crate::telemetry::{Candidate as TeleCandidate, ImplId};
let ctx = set.context();
set.alternatives()
.iter()
.map(|cand| {
let ptr = cand.kernel as *const () as usize;
let revision = bindings_table
.lookup_alternatives(op_kind, dtypes, cand.backend)
.iter()
.find(|e| e.kernel as *const () as usize == ptr)
.map_or(0, |e| e.kernel_revision_hash);
let impl_id = ImplId {
backend: cand.backend,
op: op_kind,
dtypes: dtypes.to_vec(),
kernel_source: cand.kernel_source.to_string(),
kernel_revision_hash: revision,
};
let latency_ns = match (judge, ctx) {
(Some(j), Some(ctx)) => j.measured_latency_ns(
ctx.op,
ctx.principal_dtype,
ctx.size_class,
cand.backend,
cand.kernel_source,
revision,
),
_ => None,
};
TeleCandidate {
impl_id,
latency_ns,
}
})
.collect()
}
pub fn node_needs_plan_entry(op: &fuel_graph::Op) -> bool {
op_to_op_kind(op).is_some()
&& !matches!(
op,
fuel_graph::Op::Copy { .. } | fuel_graph::Op::Move { .. }
)
}
fn is_passthrough(node: &fuel_graph::Node) -> bool {
node.op.is_view_op()
|| matches!(
node.op,
fuel_graph::Op::Reshape(_) | fuel_graph::Op::Contiguize
)
}
fn distinct_devices(set: &AlternativeSet) -> SmallVec<[DeviceLocation; 4]> {
let mut out: SmallVec<[DeviceLocation; 4]> = SmallVec::new();
for c in set.alternatives() {
if !out.contains(&c.device) {
out.push(c.device);
}
}
out
}
fn node_bytes(graph: &Graph, id: NodeId) -> u64 {
let n = graph.node(id);
(n.shape.elem_count() as u64).saturating_mul(n.dtype.size_in_bytes() as u64)
}
#[allow(clippy::too_many_arguments)]
fn open_dp_row(
dp: &mut PlacementDp,
graph: &Graph,
id: NodeId,
node: &fuel_graph::Node,
set: &AlternativeSet,
devices: &[DeviceLocation],
residency: &HashMap<NodeId, DeviceLocation>,
est: &dyn TransferEstimator,
caps_for: Option<&CapabilitiesLookup<'_>>,
) {
let rates_for = |backend: BackendId| -> (f64, f64) {
caps_for
.and_then(|f| f(backend))
.map(|caps| {
(
caps.compute_throughput_flops_per_ns,
caps.mem_bandwidth_bytes_per_ns,
)
})
.unwrap_or_else(|| default_backend_rates(backend))
};
let device_costs: Vec<(DeviceLocation, u64)> = devices
.iter()
.map(|&d| {
let min = set
.alternatives()
.iter()
.filter(|c| c.device == d)
.map(|c| {
let (cr, bw) = rates_for(c.backend);
composite_ns(&c.static_cost, cr, bw)
})
.min()
.unwrap_or(u64::MAX);
(d, min)
})
.collect();
let mut fixed: Vec<(DeviceLocation, u64)> = Vec::new();
let mut chains: Vec<ChainInput> = Vec::new();
for &input in &node.inputs {
let bytes = node_bytes(graph, input);
if let Some(&src) = residency.get(&input) {
fixed.push((src, bytes));
continue;
}
let root = dp.resolve(input);
if dp.is_open(root) {
match chains.iter_mut().find(|ci| ci.producer == root) {
Some(ci) => ci.edge_bytes.push(bytes),
None => chains.push(ChainInput {
producer: root,
edge_bytes: smallvec::smallvec![bytes],
}),
}
} else if let Some(dev) = dp.committed_device(root) {
fixed.push((dev, bytes));
}
}
dp.push_row(id, &device_costs, &fixed, &chains, est);
}
fn close_open_inputs(
dp: &mut PlacementDp,
graph: &Graph,
node: &fuel_graph::Node,
toward: DeviceLocation,
est: &dyn TransferEstimator,
residency: &mut HashMap<NodeId, DeviceLocation>,
) {
for &input in &node.inputs {
if residency.contains_key(&input) {
continue;
}
let root = dp.resolve(input);
if dp.is_open(root) {
let commits = dp.close_toward(root, toward, &[node_bytes(graph, input)], est);
for (n, d) in commits {
residency.insert(n, d);
}
}
if !residency.contains_key(&input)
&& let Some(d) = dp.committed_device(root)
{
residency.insert(input, d);
}
}
}
fn early_residency(
graph: &Graph,
options: &PlanOptions<'_>,
id: NodeId,
node: &fuel_graph::Node,
) -> Option<DeviceLocation> {
match node.op {
fuel_graph::Op::Copy { target }
| fuel_graph::Op::Move { target }
| fuel_graph::Op::Alloc { target } => return Some(target),
_ => {}
}
if let Some(loc) = graph.placement(id) {
return Some(loc);
}
options.input_residency.and_then(|f| f(id))
}
fn priced_inputs_for(
graph: &Graph,
node: &fuel_graph::Node,
residency: &HashMap<NodeId, DeviceLocation>,
) -> Vec<(DeviceLocation, u64)> {
node.inputs
.iter()
.filter_map(|&input_id| {
let src = residency.get(&input_id).copied()?;
let in_node = graph.node(input_id);
let bytes = (in_node.shape.elem_count() as u64)
.saturating_mul(in_node.dtype.size_in_bytes() as u64);
Some((src, bytes))
})
.collect()
}
struct NodeDraft {
set: AlternativeSet,
relaxed: bool,
from_fallback: bool,
costed: bool,
op_kind: OpKind,
dtypes: Vec<DType>,
diag_backend: BackendId,
}
fn build_node_draft(
graph: &Graph,
id: NodeId,
node: &fuel_graph::Node,
bindings_table: &KernelBindingTable,
options: &PlanOptions<'_>,
) -> Result<Option<NodeDraft>> {
if !node_needs_plan_entry(&node.op) {
return Ok(None);
}
let Some(op_kind) = op_to_op_kind(&node.op) else {
return Ok(None);
};
let explicit_backend = graph.target_backend(id);
let target_device = decision_device_for(graph, options, id).ok_or_else(|| {
Error::Msg(format!(
"compile_plan: node {:?} ({:?}) has no device context — \
set PlanOptions::pinned_device, a graph placement, or \
the node's target_backend",
id, node.op,
))
.bt()
})?;
let diag_backend = explicit_backend.unwrap_or_else(|| backend_for_device(target_device));
let dtypes = build_lookup_dtypes(graph, node);
let placements: Vec<(BackendId, DeviceLocation)> = match options.placements_for_device {
Some(f) => f(target_device)
.into_iter()
.map(|backend| (backend, target_device))
.collect(),
None => vec![(diag_backend, target_device)],
};
let op_params = candidate_default_op_params(graph, node);
let pricing_active = options.transfer_estimator.is_some()
&& options.populate_costs
&& options.capabilities_for.is_some();
let fallback_allowed =
node.op.destructive_input().is_none() && options.fallback_placements_for.is_some();
let relaxed = options.allow_cost_placement
&& pricing_active
&& fallback_allowed
&& graph.placement(id).is_none();
let mut set;
let mut from_fallback = false;
if relaxed {
let mut merged = placements.clone();
if let Some(fallback) = options.fallback_placements_for {
merged.extend(fallback(target_device));
}
set = enumerate_candidates(op_kind, &dtypes, &merged, &op_params, bindings_table);
} else {
set = enumerate_candidates(op_kind, &dtypes, &placements, &op_params, bindings_table);
if set.is_empty()
&& fallback_allowed
&& let Some(fallback) = options.fallback_placements_for
{
let fb_placements = fallback(target_device);
if !fb_placements.is_empty() {
set = enumerate_candidates(
op_kind,
&dtypes,
&fb_placements,
&op_params,
bindings_table,
);
from_fallback = !set.is_empty();
}
}
}
if set.is_empty() {
return Err(missing_binding_error(
bindings_table,
op_kind,
&dtypes,
diag_backend,
));
}
let input_layouts: Vec<fuel_ir::Layout> = node
.inputs
.iter()
.map(|&input_id| graph.layout(input_id))
.collect();
let ctx = FilterContext::new(op_kind, &dtypes, &input_layouts);
let chain = default_chain(options.precision_requirement);
if let Err(err) = apply_filter_chain(&mut set, &chain, &ctx) {
if relaxed || from_fallback || !fallback_allowed {
return Err(err);
}
let Some(fallback) = options.fallback_placements_for else {
return Err(err);
};
let fb_placements = fallback(target_device);
if fb_placements.is_empty() {
return Err(err);
}
let mut fb_set =
enumerate_candidates(op_kind, &dtypes, &fb_placements, &op_params, bindings_table);
if fb_set.is_empty() {
return Err(err);
}
apply_filter_chain(&mut fb_set, &chain, &ctx)?;
set = fb_set;
from_fallback = true;
}
if let Some(&principal_dtype) = dtypes.first() {
let input_shapes: Vec<Shape> = node
.inputs
.iter()
.map(|&input_id| graph.node(input_id).shape.clone())
.collect();
let size_class = SizeClass::for_op(op_kind, &input_shapes);
set.set_context(DecisionContext {
op: op_kind,
principal_dtype,
size_class,
});
}
let costed = if options.populate_costs {
if let Some(caps_for) = options.capabilities_for {
let input_shapes: Vec<Shape> = node
.inputs
.iter()
.map(|&input_id| graph.node(input_id).shape.clone())
.collect();
compute_static_costs(
&mut set,
op_kind,
&dtypes,
&input_shapes,
bindings_table,
caps_for,
options.judge,
);
true
} else {
false
}
} else {
false
};
Ok(Some(NodeDraft {
set,
relaxed,
from_fallback,
costed,
op_kind,
dtypes,
diag_backend,
}))
}
fn finalize_node_greedy(
graph: &Graph,
node: &fuel_graph::Node,
draft: NodeDraft,
bindings_table: &KernelBindingTable,
options: &PlanOptions<'_>,
residency: &HashMap<NodeId, DeviceLocation>,
) -> Result<AlternativeSet> {
let NodeDraft {
mut set,
relaxed,
from_fallback,
costed,
op_kind,
dtypes,
diag_backend,
} = draft;
if costed {
if let Some(estimator) = options.transfer_estimator {
let inputs = priced_inputs_for(graph, node, residency);
apply_inbound_transfer_costs(&mut set, &inputs, estimator);
}
set.rank_by_composite_cost();
}
if relaxed && let Some(winner_device) = set.winner().map(|c| c.device) {
let keep: Vec<usize> = set
.alternatives()
.iter()
.enumerate()
.filter_map(|(i, c)| (c.device == winner_device).then_some(i))
.collect();
if keep.len() < set.len() {
set.retain_indices(&keep);
}
}
set.retain_per_device_frontier(KEEP_PER_DEVICE);
if from_fallback && set.len() > 1 {
set.retain_indices(&[0]);
}
if set.is_empty() {
return Err(missing_binding_error(
bindings_table,
op_kind,
&dtypes,
diag_backend,
));
}
Ok(set)
}
fn missing_binding_error(
table: &KernelBindingTable,
op: OpKind,
dtypes: &[DType],
backend: BackendId,
) -> Error {
let _ = backend;
let (available_backends, supported_combinations) = table.coverage_diagnostic();
Error::NoBackendForOp {
op,
dtypes: dtypes.to_vec(),
available_backends,
supported_combinations,
}
.bt()
}
fn candidate_default_op_params(
_graph: &Graph,
_node: &fuel_graph::Node,
) -> crate::kernel::OpParams {
crate::kernel::OpParams::None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fused::PrecisionGuarantee;
use crate::kernel::{KernelCaps, OpParams, unknown_cost};
use fuel_graph::{Node, Op, topo_order};
use fuel_ir::backend::{BackendCapabilities, SubstrateClass, TransferPath};
use fuel_ir::{DType, Layout, Result as FuelResult, Shape, StrideVec};
use fuel_memory::Storage;
use smallvec::smallvec;
use std::collections::HashSet;
use std::sync::{Arc, RwLock};
fn noop_kernel(
_i: &[Arc<RwLock<Storage>>],
_o: &mut [Arc<RwLock<Storage>>],
_l: &[Layout],
_p: &OpParams,
) -> FuelResult<()> {
Ok(())
}
fn noop_kernel_b(
_i: &[Arc<RwLock<Storage>>],
_o: &mut [Arc<RwLock<Storage>>],
_l: &[Layout],
_p: &OpParams,
) -> FuelResult<()> {
Ok(())
}
fn noop_kernel_c(
_i: &[Arc<RwLock<Storage>>],
_o: &mut [Arc<RwLock<Storage>>],
_l: &[Layout],
_p: &OpParams,
) -> FuelResult<()> {
Ok(())
}
fn register_add_f32(
table: &mut KernelBindingTable,
backend: BackendId,
kernel: crate::kernel::KernelRef,
precision: PrecisionGuarantee,
) {
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
backend,
kernel,
KernelCaps::empty(),
precision,
unknown_cost,
);
}
fn build_add_graph() -> (Graph, NodeId) {
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
(g, add)
}
fn cpu_caps() -> BackendCapabilities {
BackendCapabilities {
backend_id: BackendId::Cpu,
device_location: DeviceLocation::Cpu,
op_dtype_support: HashSet::new(),
required_alignment: 64,
access_granularity_bits: 8,
transfer_paths: vec![(DeviceLocation::Cpu, TransferPath::SameDevice)],
storage_substrate: SubstrateClass::HostBytes,
compute_throughput_flops_per_ns: 1.0,
mem_bandwidth_bytes_per_ns: 4.0,
}
}
#[test]
fn empty_plan_has_no_alternatives() {
let p = ExecutionPlan::empty();
assert!(p.order.is_empty());
assert!(p.alternatives.is_empty());
assert!(p.alternatives(NodeId(0)).is_none());
}
#[test]
fn default_device_per_backend_matches_router_convention() {
assert_eq!(default_device_for(BackendId::Cpu), DeviceLocation::Cpu);
assert_eq!(
default_device_for(BackendId::Cuda),
DeviceLocation::Cuda { gpu_id: 0 },
);
assert_eq!(
default_device_for(BackendId::Vulkan),
DeviceLocation::Vulkan { gpu_id: 0 },
);
}
#[test]
fn compile_plan_skips_view_and_const_nodes() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let reshape = g.push(Node {
op: Op::Reshape(Shape::from_dims(&[2, 3])),
inputs: vec![add],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
g.set_target_backend(reshape, BackendId::Cpu);
let order = topo_order(&g, reshape);
let opts = PlanOptions::new().without_cost_population();
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert_eq!(plan.order.len(), 4);
assert_eq!(
plan.alternatives.len(),
1,
"only Add carries an alternative set",
);
let alts = plan.alternatives(add).expect("Add alternatives present");
assert_eq!(alts.len(), 1);
assert!(plan.alternatives(lhs).is_none());
assert!(plan.alternatives(rhs).is_none());
assert!(plan.alternatives(reshape).is_none());
}
#[test]
fn compile_plan_fails_fast_on_missing_binding() {
let table = KernelBindingTable::new();
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
let order = topo_order(&g, mm);
let opts = PlanOptions::new().without_cost_population();
let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
match err {
Error::NoBackendForOp { op, dtypes, .. } => {
assert_eq!(op, OpKind::MatMul);
assert_eq!(dtypes, vec![DType::F32, DType::F32, DType::F32]);
}
other => panic!("expected NoBackendForOp, got {other:?}"),
}
}
#[test]
fn plan_time_miss_reports_real_coverage_not_empty() {
let bindings = crate::dispatch::global_bindings(); let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::BF16,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
let order = topo_order(&g, mm);
let opts = PlanOptions::new().without_cost_population();
let err = compile_plan(&g, &order, &bindings, &opts).unwrap_err();
match err {
Error::NoBackendForOp {
op,
dtypes,
available_backends,
supported_combinations,
} => {
assert_eq!(op, OpKind::MatMul);
assert_eq!(
dtypes,
vec![DType::F32, DType::BF16, DType::F32],
"the mixed-precision key (F32 activation, BF16 weight)",
);
assert!(
available_backends.contains(&BackendId::Cpu),
"a populated table must report Cpu, never []: {available_backends:?}",
);
assert!(
supported_combinations
.iter()
.any(|(b, o, d)| *b == BackendId::Cpu
&& *o == OpKind::MatMul
&& *d == vec![DType::F32, DType::F32, DType::F32]),
"must show uniform F32 matmul IS registered (so the mixed gap is visible): \
{supported_combinations:?}",
);
}
other => panic!("expected NoBackendForOp, got {other:?}"),
}
}
#[test]
fn compile_plan_legacy_single_backend_uses_target_backend() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let opts = PlanOptions::new().without_cost_population();
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).unwrap();
assert_eq!(alts.len(), 1, "legacy mode → single backend → single alt");
assert_eq!(alts.winner().unwrap().backend, BackendId::Cpu);
}
#[test]
fn compile_plan_cross_co_located_via_placements_callback() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu, BackendId::Cuda]
} else {
vec![]
}
};
let opts = PlanOptions::new()
.without_cost_population()
.with_placements_for_device(&placements_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).unwrap();
assert_eq!(
alts.len(),
2,
"cross-co-located mode aggregates Cpu + Aocl alternatives",
);
}
#[test]
fn compile_plan_precision_requirement_filters_via_default_chain() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee {
bit_stable_on_same_hardware: false,
..PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU
},
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let opts = PlanOptions::new()
.without_cost_population()
.with_precision_requirement(PrecisionRequirement::BIT_STABLE);
let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
match err {
Error::FilterRejected { filter, .. } => {
assert_eq!(filter, "precision-floor");
}
other => panic!("expected FilterRejected, got {other:?}"),
}
let unconstrained = PlanOptions::new().without_cost_population();
let plan = compile_plan(&g, &order, &table, &unconstrained).expect("compile");
assert!(plan.alternatives(add_id).is_some());
let _ = add_id;
}
#[test]
fn compile_plan_with_cost_ranks_cheaper_winner() {
fn cheap_cost(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 10,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn expensive_cost(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 1_000_000,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
expensive_cost,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cheap_cost,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu, BackendId::Cuda]
} else {
vec![]
}
};
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_placements_for_device(&placements_fn)
.with_capabilities_for(&caps_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).unwrap();
assert_eq!(alts.len(), 2);
assert_eq!(
alts.winner().unwrap().backend,
BackendId::Cuda,
"cheaper static cost wins after ranking",
);
}
#[test]
fn compile_plan_judge_layer2_can_flip_layer1_winner() {
use crate::ranker::HashMapJudge;
use fuel_ir::dispatch::SizeClass;
fn cpu_layer1(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 10,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn aocl_layer1(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 100_000,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cpu_layer1,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
aocl_layer1,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu, BackendId::Cuda]
} else {
vec![]
}
};
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let sc = SizeClass::from_elem_count(3);
let mut judge = HashMapJudge::new();
judge.insert(
OpKind::AddElementwise,
DType::F32,
sc,
BackendId::Cpu,
"",
5_000,
0,
);
judge.insert(
OpKind::AddElementwise,
DType::F32,
sc,
BackendId::Cuda,
"",
20,
0,
);
let opts = PlanOptions::new()
.with_placements_for_device(&placements_fn)
.with_capabilities_for(&caps_fn)
.with_judge(&judge);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).unwrap();
assert_eq!(alts.len(), 2);
assert_eq!(
alts.winner().unwrap().backend,
BackendId::Cuda,
"Layer-2 measurement reverses Layer-1 verdict via compile_plan",
);
}
#[test]
fn compile_plan_stamps_decision_context() {
use fuel_ir::dispatch::SizeClass;
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let opts = PlanOptions::new().without_cost_population();
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).expect("Add alternatives present");
let ctx = alts.context().expect("context stamped");
assert_eq!(ctx.op, OpKind::AddElementwise);
assert_eq!(ctx.principal_dtype, DType::F32);
assert_eq!(ctx.size_class, SizeClass::from_elem_count(3));
}
#[test]
fn compile_plan_retains_per_device_frontier_no_top_n() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Vulkan,
noop_kernel_c,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu, BackendId::Cuda, BackendId::Vulkan]
} else {
vec![]
}
};
let opts = PlanOptions::new()
.without_cost_population()
.with_placements_for_device(&placements_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).unwrap();
assert_eq!(
alts.len(),
3,
"all three co-located (device, backend) paths survive — no top-N truncation",
);
}
#[test]
fn compile_plan_pinned_device_matches_stamped_plan() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let (g_stamped, add_stamped) = build_add_graph();
let order_stamped = topo_order(&g_stamped, add_stamped);
let opts_stamped = PlanOptions::new()
.without_cost_population()
.with_placements_for_device(&placements_fn);
let plan_stamped = compile_plan(&g_stamped, &order_stamped, &table, &opts_stamped)
.expect("stamped compile");
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, add);
let opts = PlanOptions::new()
.without_cost_population()
.with_placements_for_device(&placements_fn)
.with_pinned_device(DeviceLocation::Cpu);
let plan = compile_plan(&g, &order, &table, &opts).expect("pinned compile");
let a = plan_stamped.alternatives(add_stamped).expect("stamped set");
let b = plan.alternatives(add).expect("pinned set");
assert_eq!(a.len(), b.len(), "same candidate count");
for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
assert_eq!(ca.backend, cb.backend);
assert_eq!(ca.device, cb.device);
assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
assert_eq!(ca.kernel_source, cb.kernel_source);
}
assert_eq!(a.context(), b.context(), "same decision context");
}
#[test]
fn compile_plan_errors_without_any_device_context() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, lhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, add);
let opts = PlanOptions::new().without_cost_population();
let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
let msg = format!("{err:?}");
assert!(
msg.contains("no device context"),
"expected the device-context diagnostic, got: {msg}",
);
}
#[test]
fn compile_plan_skips_op_copy_nodes() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
table.register_full(
OpKind::Copy,
&[DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
);
let (mut g, add_id) = build_add_graph();
let copy_id = g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cpu,
},
inputs: vec![add_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(copy_id, BackendId::Cpu);
let order = topo_order(&g, copy_id);
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(DeviceLocation::Cpu);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert!(plan.alternatives(add_id).is_some(), "Add is planned");
assert!(
plan.alternatives(copy_id).is_none(),
"Op::Copy is residency-determined — no plan entry",
);
}
#[test]
fn compile_plan_skips_op_move_nodes() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
table.register_full(
OpKind::Copy,
&[DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
);
let (mut g, add_id) = build_add_graph();
let move_id = g.push(Node {
op: Op::Move {
target: DeviceLocation::Cpu,
},
inputs: vec![add_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(move_id, BackendId::Cpu);
let order = topo_order(&g, move_id);
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(DeviceLocation::Cpu);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert!(plan.alternatives(add_id).is_some(), "Add is planned");
assert!(
plan.alternatives(move_id).is_none(),
"Op::Move is residency-determined — no plan entry",
);
}
#[test]
fn compile_plan_fallback_admits_off_device_candidate() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, lhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, add);
let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
if dev == cuda0 {
vec![BackendId::Cuda]
} else {
vec![BackendId::Cpu]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
};
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(cuda0)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add).expect("Add planned via fallback");
assert_eq!(alts.len(), 1, "fallback set frozen to one winner");
let w = alts.winner().unwrap();
assert_eq!(w.backend, BackendId::Cpu);
assert_eq!(w.device, DeviceLocation::Cpu, "off-device placement");
}
#[test]
fn compile_plan_fallback_not_consulted_when_pinned_has_impl() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
panic!(
"fallback consulted although the pinned device has an \
implementation — locality policy violated",
);
};
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add_id).unwrap();
assert_eq!(alts.len(), 1);
assert_eq!(alts.winner().unwrap().device, DeviceLocation::Cpu);
}
#[test]
fn compile_plan_fallback_denied_for_destructive_ops() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::ReluInplace,
&[DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let relu = g.push(Node {
op: Op::ReluInplace,
inputs: vec![lhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, relu);
let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
if dev == cuda0 {
vec![BackendId::Cuda]
} else {
vec![BackendId::Cpu]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
};
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(cuda0)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn);
let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
assert!(
matches!(
err,
Error::NoBackendForOp {
op: OpKind::ReluInplace,
..
}
),
"destructive op must NOT fall back off-device; got {err:?}",
);
}
#[test]
fn compile_plan_fallback_freezes_to_single_winner() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let vk0 = DeviceLocation::Vulkan { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Vulkan,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, lhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, add);
let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
if dev == cuda0 {
vec![BackendId::Cuda]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![
(BackendId::Cpu, DeviceLocation::Cpu),
(BackendId::Vulkan, vk0),
]
};
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(cuda0)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add).unwrap();
assert_eq!(
alts.len(),
1,
"fallback set must freeze to its single ranked winner",
);
}
struct FlatEstimator {
latency_ns: u64,
}
impl crate::ranker::TransferEstimator for FlatEstimator {
fn estimate_transfer_ns(
&self,
src: DeviceLocation,
dst: DeviceLocation,
_bytes: u64,
) -> u64 {
if src == dst { 0 } else { self.latency_ns }
}
}
#[test]
fn compile_plan_prices_fallback_by_input_residency() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let vk0 = DeviceLocation::Vulkan { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Vulkan,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, lhs],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let order = topo_order(&g, add);
let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
if dev == cuda0 {
vec![BackendId::Cuda]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![
(BackendId::Cpu, DeviceLocation::Cpu),
(BackendId::Vulkan, vk0),
]
};
let residency_fn =
move |id: NodeId| -> Option<DeviceLocation> { (id == lhs).then_some(vk0) };
let est = FlatEstimator { latency_ns: 1_000 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(cuda0)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add).expect("Add planned via fallback");
assert_eq!(alts.len(), 1, "fallback narrows to a single winner");
let w = alts.winner().unwrap();
assert_eq!(
w.device, vk0,
"transfer pricing moves the fallback to the inputs' residency",
);
assert_eq!(w.inbound_transfer_ns, 0, "co-resident inputs price zero");
}
fn cost_600(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 600,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_10(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 10,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_15_000(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 15_000,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_10_000_000(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 10_000_000,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn relaxed_plan_with_costs(
cpu_cost: crate::kernel::CostFn,
cuda_cost: crate::kernel::CostFn,
estimator: &FlatEstimator,
) -> AlternativeSet {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cpu_cost,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cuda_cost,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cuda, cuda0)]
};
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(estimator)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
plan.alternatives(add_id).expect("Add planned").clone()
}
#[test]
fn compile_plan_relaxed_tiny_op_stays_local_when_transfer_dominates() {
let est = FlatEstimator {
latency_ns: 1_000_000,
};
let set = relaxed_plan_with_costs(cost_600, cost_10, &est);
let w = set.winner().unwrap();
assert_eq!(
w.device,
DeviceLocation::Cpu,
"transfer dominates → stay local"
);
assert_eq!(w.inbound_transfer_ns, 0, "local winner pays no transfer");
assert!(
set.alternatives()
.iter()
.all(|c| c.device == DeviceLocation::Cpu),
"off-device siblings pruned after rank",
);
}
#[test]
fn compile_plan_uncapped_gpu_no_longer_wins_by_being_free() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
fn build_table() -> KernelBindingTable {
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_600,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_15_000,
);
table
}
fn winner_device_under(
table: &KernelBindingTable,
cpu_caps_val: &'static BackendCapabilities,
gpu_caps: Option<&'static BackendCapabilities>,
cuda0: DeviceLocation,
) -> DeviceLocation {
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cuda, cuda0)]
};
let residency_fn =
|_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let est = FlatEstimator { latency_ns: 100 };
let caps_fn = move |b: BackendId| -> Option<&BackendCapabilities> {
match b {
BackendId::Cpu => Some(cpu_caps_val),
BackendId::Cuda => gpu_caps,
_ => None,
}
};
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, table, &opts).expect("compile");
plan.alternatives(add_id).unwrap().winner().unwrap().device
}
let table = build_table();
let cpu_caps_val: &'static BackendCapabilities = Box::leak(Box::new(cpu_caps()));
let winner_bug = winner_device_under(&table, cpu_caps_val, None, cuda0);
assert_eq!(
winner_bug, cuda0,
"PRE-FIX: an uncapped CUDA candidate is skipped by \
compute_static_costs, prices at 0 + 200 transfer, and \
spuriously beats the CPU's honest 600 — this IS the bug",
);
let cuda_caps_val: &'static BackendCapabilities = Box::leak(Box::new(
crate::dispatch::derive_backend_caps(BackendId::Cuda, cuda0, &table),
));
assert!(
cuda_caps_val
.op_dtype_support
.contains(&(OpKind::AddElementwise, DType::F32)),
"derive_backend_caps must mirror the binding table's \
registered (op, output-dtype) pairs for the backend",
);
assert_eq!(cuda_caps_val.backend_id, BackendId::Cuda);
assert_eq!(cuda_caps_val.device_location, cuda0);
let winner_fixed = winner_device_under(&table, cpu_caps_val, Some(cuda_caps_val), cuda0);
assert_eq!(
winner_fixed,
DeviceLocation::Cpu,
"POST-FIX: with GPU caps registered the CUDA candidate is \
costed honestly (15000 FLOPs ÷ 30/ns = 500 ns compute + 200 \
transfer = 700 > CPU 600), so the CPU-pinned realize stays \
on CPU — no spill onto the unseeded GPU",
);
}
#[test]
fn compile_plan_places_large_compute_on_gpu_small_on_cpu_by_throughput() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
fn winner(cost: crate::kernel::CostFn, cuda0: DeviceLocation) -> DeviceLocation {
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cuda, cuda0)]
};
let residency_fn =
|_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let est = FlatEstimator { latency_ns: 100 };
let cpu_caps_val = cpu_caps();
let cuda_caps_val =
crate::dispatch::derive_backend_caps(BackendId::Cuda, cuda0, &table);
let caps_fn = |b: BackendId| -> Option<&BackendCapabilities> {
match b {
BackendId::Cpu => Some(&cpu_caps_val),
BackendId::Cuda => Some(&cuda_caps_val),
_ => None,
}
};
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
plan.alternatives(add_id).unwrap().winner().unwrap().device
}
assert_eq!(
winner(cost_100, cuda0),
DeviceLocation::Cpu,
"SMALL op: the GPU's throughput win on 100 FLOPs (~3 ns) \
can't overcome the 200 ns inbound transfer → stay on CPU",
);
assert_eq!(
winner(cost_10_000_000, cuda0),
cuda0,
"LARGE op: the GPU's throughput win on 10 M FLOPs (~333 k ns) \
dwarfs the 200 ns transfer → move to CUDA. RED under the \
fixed 1-FLOP/ns prior (equal compute-ns → the transfer \
always loses the GPU the placement)",
);
}
#[test]
fn compile_plan_without_cost_placement_keeps_large_op_on_cpu() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_10_000_000,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_10_000_000,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cuda, cuda0)]
};
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let est = FlatEstimator { latency_ns: 100 };
let cpu_caps_val = cpu_caps();
let cuda_caps_val = crate::dispatch::derive_backend_caps(BackendId::Cuda, cuda0, &table);
let caps_fn = |b: BackendId| -> Option<&BackendCapabilities> {
match b {
BackendId::Cpu => Some(&cpu_caps_val),
BackendId::Cuda => Some(&cuda_caps_val),
_ => None,
}
};
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn)
.without_cost_placement();
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let set = plan.alternatives(add_id).unwrap();
assert_eq!(
set.winner().unwrap().device,
DeviceLocation::Cpu,
"without_cost_placement: the large GPU-favorable op stays on the \
pinned CPU device (no relaxed cross-device enumeration)",
);
assert!(
set.alternatives()
.iter()
.all(|c| c.device == DeviceLocation::Cpu),
"no off-device candidate enumerated at all",
);
}
#[test]
fn compile_plan_relaxed_huge_op_moves_when_kernel_gap_dominates() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let est = FlatEstimator { latency_ns: 1_000 };
let set = relaxed_plan_with_costs(cost_10_000_000, cost_10, &est);
let w = set.winner().unwrap();
assert_eq!(w.device, cuda0, "kernel gap dominates → move");
assert_eq!(
w.inbound_transfer_ns, 2_000,
"the move's two input crossings are priced on the winner",
);
assert!(
set.alternatives().iter().all(|c| c.device == cuda0),
"set pruned to the winner's device",
);
}
#[test]
fn compile_plan_relaxed_zero_signal_ties_stay_local() {
let est = FlatEstimator { latency_ns: 0 };
let set = relaxed_plan_with_costs(unknown_cost, unknown_cost, &est);
assert_eq!(
set.winner().unwrap().device,
DeviceLocation::Cpu,
"ties break toward the decision device",
);
}
#[test]
fn compile_plan_relaxed_single_device_plan_unchanged() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> { Vec::new() };
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let base_opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn);
let base = compile_plan(&g, &order, &table, &base_opts).expect("base");
let est = FlatEstimator {
latency_ns: 1_000_000,
};
let wired_opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let wired = compile_plan(&g, &order, &table, &wired_opts).expect("wired");
let a = base.alternatives(add_id).unwrap();
let b = wired.alternatives(add_id).unwrap();
assert_eq!(a.len(), b.len(), "same candidate count");
for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
assert_eq!(ca.backend, cb.backend);
assert_eq!(ca.device, cb.device);
assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
assert_eq!(ca.static_cost, cb.static_cost, "same Layer-1 cost");
assert_eq!(cb.inbound_transfer_ns, 0, "no transfer term fires");
}
assert_eq!(a.context(), b.context(), "same decision context");
}
#[test]
fn compile_plan_relaxed_respects_explicit_placement_pin() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (mut g, add_id) = build_add_graph();
g.set_placement(add_id, DeviceLocation::Cpu);
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
panic!(
"fallback consulted for a hard-pinned node — explicit \
placements must not enter the priced relax",
);
};
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let est = FlatEstimator { latency_ns: 0 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert_eq!(
plan.alternatives(add_id).unwrap().winner().unwrap().device,
DeviceLocation::Cpu,
);
}
#[test]
fn compile_plan_relaxed_destructive_never_moves() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::ReluInplace,
&[DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_600,
);
table.register_full(
OpKind::ReluInplace,
&[DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_10,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let relu = g.push(Node {
op: Op::ReluInplace,
inputs: vec![lhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, relu);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![]
}
};
let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cuda, cuda0)]
};
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let est = FlatEstimator { latency_ns: 0 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert_eq!(
plan.alternatives(relu).unwrap().winner().unwrap().device,
DeviceLocation::Cpu,
"destructive op stays on the device owning its mutation target",
);
}
#[test]
fn compile_plan_filter_rejected_pinned_device_falls_back() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee {
bit_stable_on_same_hardware: false,
..PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU
},
);
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, lhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let order = topo_order(&g, add);
let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
if dev == cuda0 {
vec![BackendId::Cuda]
} else {
vec![]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
};
let opts = PlanOptions::new()
.without_cost_population()
.with_pinned_device(cuda0)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_precision_requirement(PrecisionRequirement::BIT_STABLE);
let plan = compile_plan(&g, &order, &table, &opts)
.expect("filter-rejected pin falls back instead of erroring");
let alts = plan.alternatives(add).expect("Add planned via fallback");
assert_eq!(alts.len(), 1, "fallback set frozen to one winner");
let w = alts.winner().unwrap();
assert_eq!(w.backend, BackendId::Cpu);
assert_eq!(w.device, DeviceLocation::Cpu);
assert!(
w.precision.bit_stable_on_same_hardware,
"the admissible off-device sibling won",
);
}
#[test]
fn compile_plan_without_estimator_prices_nothing() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let base_opts = PlanOptions::new().with_capabilities_for(&caps_fn);
let base = compile_plan(&g, &order, &table, &base_opts).expect("base");
let with_residency = PlanOptions::new()
.with_capabilities_for(&caps_fn)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &with_residency).expect("plan");
let a = base.alternatives(add_id).unwrap();
let b = plan.alternatives(add_id).unwrap();
assert_eq!(a.len(), b.len());
for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
assert_eq!(ca.backend, cb.backend);
assert_eq!(ca.device, cb.device);
assert_eq!(ca.kernel as usize, cb.kernel as usize);
assert_eq!(ca.inbound_transfer_ns, 0);
assert_eq!(cb.inbound_transfer_ns, 0);
}
}
#[test]
fn compile_plan_require_transfer_pricing_errors_without_estimator() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_capabilities_for(&caps_fn)
.require_transfer_pricing();
let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
let msg = format!("{err:?}");
assert!(
msg.contains("transfer") && (msg.contains("estimator") || msg.contains("pricing")),
"error should name the missing transfer estimator / pricing; got: {msg}",
);
}
#[test]
fn compile_plan_require_transfer_pricing_ok_with_estimator() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let est = FlatEstimator { latency_ns: 1000 };
let opts = PlanOptions::new()
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.require_transfer_pricing();
let plan = compile_plan(&g, &order, &table, &opts)
.expect("require_transfer_pricing WITH an estimator must plan fine");
assert!(plan.alternatives(add_id).is_some());
}
#[test]
fn compile_plan_passes_input_layouts_to_filter_chain() {
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel_b,
KernelCaps::strided_input(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
);
let mut g = Graph::new();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let strides: StrideVec = smallvec![2isize];
g.set_layout(lhs, Layout::new(Shape::from_dims(&[6]), strides, 0));
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
let order = topo_order(&g, add);
let opts = PlanOptions::new().without_cost_population();
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let alts = plan.alternatives(add).unwrap();
assert_eq!(alts.len(), 1);
assert!(alts.winner().unwrap().caps.strided_input);
}
struct BytesEstimator {
latency_ns: u64,
ns_per_byte: u64,
}
impl crate::ranker::TransferEstimator for BytesEstimator {
fn estimate_transfer_ns(
&self,
src: DeviceLocation,
dst: DeviceLocation,
bytes: u64,
) -> u64 {
if src == dst {
return 0;
}
self.latency_ns
.saturating_add(bytes.saturating_mul(self.ns_per_byte))
}
}
fn cost_50(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 50,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_100(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 100,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_300(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 300,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_800(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 800,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_5000(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 5000,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn cost_100_000(
_: &[Shape],
_: &[DType],
_: &OpParams,
_: &BackendCapabilities,
) -> crate::fused::CostEstimate {
crate::fused::CostEstimate {
flops: 100_000,
bytes_moved: 0,
kernel_overhead_ns: 0,
}
}
fn register_unary_f32(
table: &mut KernelBindingTable,
op: OpKind,
backend: BackendId,
kernel: crate::kernel::KernelRef,
cost: crate::kernel::CostFn,
) {
table.register_full(
op,
&[DType::F32, DType::F32],
backend,
kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost,
);
}
fn push_f32(g: &mut Graph, op: Op, inputs: Vec<NodeId>, dims: &[usize]) -> NodeId {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(dims),
dtype: DType::F32,
})
}
#[test]
fn dp_mid_sequence_gpu_segment_beats_all_cpu_and_all_gpu() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_unary_f32(
&mut table,
OpKind::SqrElementwise,
BackendId::Cpu,
noop_kernel,
cost_100,
);
register_unary_f32(
&mut table,
OpKind::SqrElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_5000,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cpu,
noop_kernel,
cost_800,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
let mut g = Graph::new();
let c = push_f32(&mut g, Op::Const, vec![], &[4]);
let n1 = push_f32(&mut g, Op::Sqr, vec![c], &[4]);
let n2 = push_f32(&mut g, Op::Neg, vec![n1], &[4]);
let n3 = push_f32(&mut g, Op::Neg, vec![n2], &[4]);
let n4 = push_f32(&mut g, Op::Neg, vec![n3], &[4]);
let n5 = push_f32(&mut g, Op::Sqr, vec![n4], &[4]);
let order = topo_order(&g, n5);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
(id == c).then_some(DeviceLocation::Cpu)
};
let est = FlatEstimator { latency_ns: 1_000 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let dev_of = |id: NodeId| plan.alternatives(id).unwrap().winner().unwrap().device;
assert_eq!(
dev_of(n1),
DeviceLocation::Cpu,
"GPU-hostile entry stays CPU"
);
assert_eq!(dev_of(n2), cuda0, "segment start migrates");
assert_eq!(dev_of(n3), cuda0, "segment interior stays GPU");
assert_eq!(dev_of(n4), cuda0, "segment end stays GPU");
assert_eq!(dev_of(n5), DeviceLocation::Cpu, "GPU-hostile exit returns");
let inbound = |id: NodeId| {
plan.alternatives(id)
.unwrap()
.winner()
.unwrap()
.inbound_transfer_ns
};
assert_eq!(
inbound(n2),
1_000,
"entry crossing priced on the first GPU op"
);
assert_eq!(inbound(n3), 0);
assert_eq!(inbound(n4), 0);
assert_eq!(inbound(n5), 1_000, "exit crossing priced on the return op");
for id in [n1, n2, n3, n4, n5] {
let d = dev_of(id);
assert!(
plan.alternatives(id)
.unwrap()
.alternatives()
.iter()
.all(|c| c.device == d),
"set pruned to the DP-committed device",
);
}
}
#[test]
fn dp_fused_op_loses_when_stranding_exceeds_local_win() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cpu,
noop_kernel,
cost_100_000,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
register_unary_f32(
&mut table,
OpKind::SoftmaxLastDim,
BackendId::Cpu,
noop_kernel,
cost_300,
);
register_unary_f32(
&mut table,
OpKind::SoftmaxLastDim,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
register_unary_f32(
&mut table,
OpKind::SqrElementwise,
BackendId::Cpu,
noop_kernel,
cost_100,
);
let mut g = Graph::new();
let c = push_f32(&mut g, Op::Const, vec![], &[1]);
let n1 = push_f32(&mut g, Op::Neg, vec![c], &[1]);
let f = push_f32(
&mut g,
Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
vec![n1],
&[256],
);
let n2 = push_f32(&mut g, Op::Sqr, vec![f], &[256]);
let order = topo_order(&g, n2);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
(id == c).then_some(DeviceLocation::Cpu)
};
let est = BytesEstimator {
latency_ns: 10,
ns_per_byte: 1,
};
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let dev_of = |id: NodeId| plan.alternatives(id).unwrap().winner().unwrap().device;
assert_eq!(dev_of(n1), cuda0, "the huge-win producer migrates");
assert_eq!(
dev_of(f),
DeviceLocation::Cpu,
"the locally-faster fused op stays where its output is needed",
);
assert_eq!(dev_of(n2), DeviceLocation::Cpu);
assert!(
plan.alternatives(f)
.unwrap()
.alternatives()
.iter()
.all(|c| c.device == DeviceLocation::Cpu),
"fused set pruned to the committed device",
);
}
#[test]
fn dp_exit_pricing_keeps_final_op_on_realize_target() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cpu,
noop_kernel,
cost_50,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
let mut g = Graph::new();
let c = push_f32(&mut g, Op::Const, vec![], &[4]);
let n = push_f32(&mut g, Op::Neg, vec![c], &[4]);
let order = topo_order(&g, n);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn =
move |id: NodeId| -> Option<DeviceLocation> { (id == c).then_some(cuda0) };
let est = FlatEstimator { latency_ns: 1_000 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let w = plan.alternatives(n).unwrap().winner().unwrap();
assert_eq!(
w.device,
DeviceLocation::Cpu,
"exit pricing flips the locally-winning GPU placement",
);
assert_eq!(
w.inbound_transfer_ns, 1_000,
"the GPU-resident input's crossing is priced on the CPU winner",
);
}
#[test]
fn dp_exit_pricing_does_not_pin_when_kernel_gap_dominates() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cpu,
noop_kernel,
cost_10_000_000,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
let mut g = Graph::new();
let c = push_f32(&mut g, Op::Const, vec![], &[4]);
let n = push_f32(&mut g, Op::Neg, vec![c], &[4]);
let order = topo_order(&g, n);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn =
move |id: NodeId| -> Option<DeviceLocation> { (id == c).then_some(cuda0) };
let est = FlatEstimator { latency_ns: 1_000 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert_eq!(
plan.alternatives(n).unwrap().winner().unwrap().device,
cuda0
);
}
#[test]
fn dp_single_device_plan_equals_greedy_plan() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
register_add_f32(
&mut table,
BackendId::Cuda,
noop_kernel_b,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu, BackendId::Cuda]
} else {
vec![]
}
};
let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> { Vec::new() };
let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let base_opts = PlanOptions::new()
.with_placements_for_device(&placements_fn)
.with_capabilities_for(&caps_fn);
let base = compile_plan(&g, &order, &table, &base_opts).expect("base");
let est = FlatEstimator {
latency_ns: 1_000_000,
};
let wired_opts = PlanOptions::new()
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let wired = compile_plan(&g, &order, &table, &wired_opts).expect("wired");
assert_eq!(base.alternatives.len(), wired.alternatives.len());
let a = base.alternatives(add_id).unwrap();
let b = wired.alternatives(add_id).unwrap();
assert_eq!(a.len(), b.len(), "same candidate count");
for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
assert_eq!(ca.backend, cb.backend);
assert_eq!(ca.device, cb.device);
assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
assert_eq!(ca.static_cost, cb.static_cost, "same Layer-1 cost");
assert_eq!(cb.inbound_transfer_ns, 0, "no transfer term fires");
}
assert_eq!(a.context(), b.context(), "same decision context");
}
fn run_diamond_regime(latency_ns: u64, expect: DeviceLocation) {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
for op in [
OpKind::NegElementwise,
OpKind::SqrElementwise,
OpKind::ReluElementwise,
] {
register_unary_f32(&mut table, op, BackendId::Cpu, noop_kernel, cost_100);
register_unary_f32(&mut table, op, BackendId::Cuda, noop_kernel_b, cost_10);
}
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_100,
);
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cuda,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
cost_10,
);
let mut g = Graph::new();
let c0 = push_f32(&mut g, Op::Const, vec![], &[4]);
let a = push_f32(&mut g, Op::Neg, vec![c0], &[4]);
let b = push_f32(&mut g, Op::Sqr, vec![a], &[4]);
let c2 = push_f32(&mut g, Op::Relu, vec![a], &[4]);
let d = push_f32(&mut g, Op::Add, vec![b, c2], &[4]);
let order = topo_order(&g, d);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
(id == c0).then_some(DeviceLocation::Cpu)
};
let est = FlatEstimator { latency_ns };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
for id in [a, b, c2, d] {
assert_eq!(
plan.alternatives(id).unwrap().winner().unwrap().device,
expect,
"diamond join must commit a consistent placement",
);
}
}
#[test]
fn dp_diamond_join_merges_consistently() {
run_diamond_regime(1_000_000, DeviceLocation::Cpu);
run_diamond_regime(0, DeviceLocation::Cuda { gpu_id: 0 });
}
#[test]
fn dp_thousand_node_chain_plans_within_sanity_bound() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cpu,
noop_kernel,
cost_100,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
let mut g = Graph::new();
let c = push_f32(&mut g, Op::Const, vec![], &[4]);
let mut prev = c;
let mut nodes = Vec::with_capacity(1000);
for _ in 0..1000 {
let n = push_f32(&mut g, Op::Neg, vec![prev], &[4]);
nodes.push(n);
prev = n;
}
let order = topo_order(&g, prev);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
(id == c).then_some(DeviceLocation::Cpu)
};
let est = FlatEstimator { latency_ns: 10 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let start = std::time::Instant::now();
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let elapsed = start.elapsed();
assert!(
elapsed.as_secs() < 5,
"1000-node DP plan took {elapsed:?} — pathological scaling",
);
assert_eq!(plan.alternatives.len(), 1000);
for &n in &nodes {
assert_eq!(
plan.alternatives(n).unwrap().winner().unwrap().device,
DeviceLocation::Cpu,
"equal-cost ties break toward the decision device end-to-end",
);
}
}
#[test]
fn dp_chain_connects_through_view_passthroughs() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut table = KernelBindingTable::new();
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cpu,
noop_kernel,
cost_800,
);
register_unary_f32(
&mut table,
OpKind::NegElementwise,
BackendId::Cuda,
noop_kernel_b,
cost_100,
);
let mut g = Graph::new();
let c = push_f32(&mut g, Op::Const, vec![], &[4]);
let n1 = push_f32(&mut g, Op::Neg, vec![c], &[4]);
let r = push_f32(
&mut g,
Op::Reshape(Shape::from_dims(&[2, 2])),
vec![n1],
&[2, 2],
);
let n2 = push_f32(&mut g, Op::Neg, vec![r], &[2, 2]);
let n3 = push_f32(&mut g, Op::Neg, vec![n2], &[2, 2]);
let order = topo_order(&g, n3);
let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
if dev == DeviceLocation::Cpu {
vec![BackendId::Cpu]
} else {
vec![BackendId::Cuda]
}
};
let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
if dev == DeviceLocation::Cpu {
vec![(BackendId::Cuda, cuda0)]
} else {
vec![(BackendId::Cpu, DeviceLocation::Cpu)]
}
};
let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
(id == c).then_some(DeviceLocation::Cpu)
};
let est = FlatEstimator { latency_ns: 1_000 };
let cpu_caps_val = cpu_caps();
let caps_fn = |_: BackendId| Some(&cpu_caps_val);
let opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_fn)
.with_fallback_placements_for(&fallback_fn)
.with_capabilities_for(&caps_fn)
.with_transfer_estimator(&est)
.with_input_residency(&residency_fn);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
for id in [n1, n2, n3] {
assert_eq!(
plan.alternatives(id).unwrap().winner().unwrap().device,
cuda0,
"the chain must not break at the Reshape pass-through",
);
}
assert!(
plan.alternatives(r).is_none(),
"Reshape carries no plan entry"
);
}
#[cfg(feature = "telemetry")]
mod telemetry_wiring {
use super::*;
use crate::telemetry::{
Contiguity, FdxOperandDesc, HwStamp, StructureKeyProvider, StructureKeyToken,
TelemetryConfig, TelemetryHooks, TelemetryMode, TelemetrySink,
};
use std::sync::Mutex;
struct StubProvider;
impl StructureKeyProvider for StubProvider {
fn structure_key(
&self,
op_class: &str,
operands: &[FdxOperandDesc],
_arch: &str,
) -> Option<StructureKeyToken> {
Some(StructureKeyToken(format!(
"{op_class}:stub:n={}",
operands.len()
)))
}
}
fn cpu_hw() -> HwStamp {
HwStamp {
compute_capability: None,
hardware_sku: "test-cpu".into(),
driver_version: "n/a".into(),
}
}
fn generic_only_add_table() -> KernelBindingTable {
let mut table = KernelBindingTable::new();
table.register_full_with_source_generic(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::strided_input(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
"baracuda-generic-strided",
true, 0, None, );
table
}
#[test]
fn generic_only_cell_emits_one_aggregated_miss_through_compile_plan() {
let table = generic_only_add_table();
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sink = Mutex::new(TelemetrySink::new());
let provider = StubProvider;
let config = TelemetryConfig {
mode: TelemetryMode::Coarse,
out_path: None,
};
let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks);
let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let sink = sink.into_inner().expect("sink lock");
assert_eq!(
sink.miss_cell_count(),
1,
"a generic-only cell must emit exactly one aggregated miss cell",
);
let recs = sink.miss_records();
assert_eq!(recs.len(), 1);
assert_eq!(
recs[0].wanted,
StructureKeyToken("AddElementwise:stub:n=3".into()),
"the wanted token must reflect the live op_class + operands",
);
assert_eq!(recs[0].fallback.kernel_source, "baracuda-generic-strided");
assert_eq!(recs[0].fallback.op, OpKind::AddElementwise);
assert_eq!(recs[0].fallback.backend, BackendId::Cpu);
assert_eq!(
recs[0].fallback.dtypes,
vec![DType::F32, DType::F32, DType::F32],
);
assert_eq!(recs[0].count, 1);
assert_eq!(recs[0].hw, cpu_hw());
}
#[test]
fn telemetry_off_emits_nothing_and_does_not_change_the_plan() {
let table = generic_only_add_table();
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sink = Mutex::new(TelemetrySink::new());
let provider = StubProvider;
let config = TelemetryConfig::default(); let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks);
let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert!(
sink.into_inner().expect("sink lock").is_empty(),
"Off mode must emit nothing",
);
let alts = plan.alternatives(add_id).expect("Add alternatives present");
assert_eq!(alts.len(), 1);
assert_eq!(alts.winner().unwrap().backend, BackendId::Cpu);
}
#[test]
fn specialized_winner_emits_no_miss() {
let mut table = KernelBindingTable::new();
register_add_f32(
&mut table,
BackendId::Cpu,
noop_kernel,
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
);
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sink = Mutex::new(TelemetrySink::new());
let provider = StubProvider;
let config = TelemetryConfig {
mode: TelemetryMode::Coarse,
out_path: None,
};
let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks);
let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let sink = sink.into_inner().expect("sink lock");
assert!(
sink.miss_records().is_empty(),
"a specialized (is_generic=false) winner must emit no miss",
);
assert_eq!(
sink.dispatch_cell_count(),
1,
"a specialized winner still emits its dispatch record",
);
let _ = add_id;
}
fn two_sibling_add_table() -> KernelBindingTable {
let mut table = KernelBindingTable::new();
table.register_full_with_source(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
"portable-cpu",
);
table.register_full_with_source(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel_b,
KernelCaps::empty(),
PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
"aocl",
);
table
}
#[test]
fn detailed_dispatch_record_emitted_with_per_candidate_judge_latency() {
use crate::ranker::HashMapJudge;
use fuel_ir::dispatch::SizeClass;
let table = two_sibling_add_table();
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sc = SizeClass::from_elem_count(3);
let mut judge = HashMapJudge::new();
judge.insert(
OpKind::AddElementwise,
DType::F32,
sc,
BackendId::Cpu,
"portable-cpu",
41_230,
0,
);
let sink = Mutex::new(TelemetrySink::new());
let provider = StubProvider;
let config = TelemetryConfig {
mode: TelemetryMode::Detailed,
out_path: None,
};
let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks)
.with_judge(&judge);
let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let sink = sink.into_inner().expect("sink lock");
assert_eq!(
sink.dispatch_cell_count(),
1,
"one planned cell ⇒ exactly one aggregated dispatch record",
);
assert_eq!(
sink.miss_cell_count(),
0,
"specialized winners emit no miss"
);
let recs = sink.dispatch_records();
assert_eq!(recs.len(), 1);
let rec = &recs[0];
assert_eq!(rec.chosen.kernel_source, "portable-cpu");
assert_eq!(rec.chosen.op, OpKind::AddElementwise);
assert_eq!(rec.chosen.backend, BackendId::Cpu);
assert_eq!(rec.chosen.dtypes, vec![DType::F32, DType::F32, DType::F32]);
assert_eq!(
rec.chosen.kernel_revision_hash, 0,
"hand-written ⇒ untracked"
);
assert_eq!(rec.count, 1);
assert_eq!(rec.hw, cpu_hw());
assert_eq!(
rec.structure_key,
Some(StructureKeyToken("AddElementwise:stub:n=3".into())),
);
assert_eq!(rec.candidates.len(), 2, "both admitted alternatives appear");
let by_src = |src: &str| {
rec.candidates
.iter()
.find(|c| c.impl_id.kernel_source == src)
.unwrap_or_else(|| panic!("candidate {src} present"))
};
assert_eq!(
by_src("portable-cpu").latency_ns,
Some(41_230),
"the measured winner carries its own ns",
);
assert_eq!(
by_src("aocl").latency_ns,
None,
"the unmeasured sibling is None, never a fabricated 0",
);
}
#[test]
fn coarse_dispatch_record_has_no_candidates() {
let table = two_sibling_add_table();
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sink = Mutex::new(TelemetrySink::new());
let provider = StubProvider;
let config = TelemetryConfig {
mode: TelemetryMode::Coarse,
out_path: None,
};
let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks);
let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let sink = sink.into_inner().expect("sink lock");
assert_eq!(
sink.dispatch_cell_count(),
1,
"Coarse still emits the chosen signal"
);
let recs = sink.dispatch_records();
assert!(
recs[0].candidates.is_empty(),
"Coarse omits candidates (no Judge reads)"
);
assert_eq!(recs[0].chosen.kernel_source, "portable-cpu");
assert_eq!(recs[0].count, 1);
}
#[test]
fn telemetry_off_emits_no_dispatch_record() {
let table = two_sibling_add_table();
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sink = Mutex::new(TelemetrySink::new());
let provider = StubProvider;
let config = TelemetryConfig::default(); let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks);
let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
assert!(
sink.into_inner().expect("sink lock").is_empty(),
"Off emits neither dispatch nor miss",
);
let _ = add_id;
}
struct RecordingProvider {
seen: Mutex<Vec<FdxOperandDesc>>,
}
impl StructureKeyProvider for RecordingProvider {
fn structure_key(
&self,
op_class: &str,
operands: &[FdxOperandDesc],
_arch: &str,
) -> Option<StructureKeyToken> {
*self.seen.lock().expect("record lock") = operands.to_vec();
Some(StructureKeyToken(format!(
"{op_class}:n={}",
operands.len()
)))
}
}
#[test]
fn structure_key_operand_slice_folds_in_the_output() {
let table = generic_only_add_table();
let (g, add_id) = build_add_graph();
let order = topo_order(&g, add_id);
let sink = Mutex::new(TelemetrySink::new());
let provider = RecordingProvider {
seen: Mutex::new(Vec::new()),
};
let config = TelemetryConfig {
mode: TelemetryMode::Coarse,
out_path: None,
};
let hooks = TelemetryHooks {
config: &config,
sink: &sink,
provider: &provider,
hw: cpu_hw(),
};
let opts = PlanOptions::new()
.without_cost_population()
.with_telemetry(&hooks);
let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
let seen = provider.seen.into_inner().expect("record lock");
assert_eq!(
seen.len(),
3,
"operand slice must be node inputs (2) then the output (1)",
);
let out = seen.last().expect("output operand present");
assert_eq!(
out.dtype,
DType::F32,
"output operand carries the node dtype"
);
assert_eq!(
out.shape,
vec![3i64],
"output operand carries the node shape"
);
assert_eq!(
out.contiguity,
Contiguity::Contiguous,
"the Add output is contiguous",
);
let _ = add_id;
}
}
}