1use std::collections::HashMap;
31
32use fuel_graph::{Graph, NodeId};
33use fuel_ir::dispatch::{OpKind, SizeClass};
34use fuel_ir::probe::BackendId;
35use fuel_ir::{DType, DeviceLocation, Error, Result, Shape};
36
37use smallvec::SmallVec;
38
39use crate::kernel::KernelBindingTable;
40use crate::pipelined::{build_lookup_dtypes, op_to_op_kind};
41use crate::ranker::{
42 AlternativeSet, CapabilitiesLookup, ChainInput, DecisionContext, FilterContext, JudgeOracle,
43 KEEP_PER_DEVICE, PlacementDp, PrecisionRequirement, TransferEstimator, apply_filter_chain,
44 apply_inbound_transfer_costs, composite_ns, compute_static_costs, default_backend_rates,
45 default_chain, enumerate_candidates,
46};
47
48#[derive(Debug)]
57pub struct ExecutionPlan {
58 pub order: Vec<NodeId>,
60 pub alternatives: HashMap<NodeId, AlternativeSet>,
62 pub generation: u64,
73}
74
75impl ExecutionPlan {
76 pub fn empty() -> Self {
78 Self {
79 order: Vec::new(),
80 alternatives: HashMap::new(),
81 generation: 0,
82 }
83 }
84
85 pub fn alternatives(&self, node: NodeId) -> Option<&AlternativeSet> {
88 self.alternatives.get(&node)
89 }
90
91 pub fn alternatives_mut(&mut self, node: NodeId) -> Option<&mut AlternativeSet> {
95 self.alternatives.get_mut(&node)
96 }
97}
98
99pub fn default_device_for(backend: BackendId) -> DeviceLocation {
104 match backend {
105 BackendId::Cpu => DeviceLocation::Cpu,
106 BackendId::Cuda => DeviceLocation::Cuda { gpu_id: 0 },
107 BackendId::Vulkan => DeviceLocation::Vulkan { gpu_id: 0 },
108 BackendId::Metal => DeviceLocation::Metal { gpu_id: 0 },
109 _ => DeviceLocation::Cpu,
110 }
111}
112
113pub(crate) fn decision_device_for(
123 graph: &Graph,
124 options: &PlanOptions<'_>,
125 id: NodeId,
126) -> Option<DeviceLocation> {
127 graph
128 .placement(id)
129 .or(options.pinned_device)
130 .or_else(|| graph.target_backend(id).map(default_device_for))
131}
132
133pub(crate) fn candidate_backends_for(
153 graph: &Graph,
154 options: &PlanOptions<'_>,
155 id: NodeId,
156 node: &fuel_graph::Node,
157) -> Vec<BackendId> {
158 let Some(target_device) = decision_device_for(graph, options, id) else {
159 return Vec::new();
160 };
161 let mut backends: Vec<BackendId> = match options.placements_for_device {
162 Some(f) => f(target_device),
163 None => vec![
164 graph
165 .target_backend(id)
166 .unwrap_or_else(|| backend_for_device(target_device)),
167 ],
168 };
169 if node.op.destructive_input().is_none()
174 && let Some(fallback) = options.fallback_placements_for
175 {
176 backends.extend(fallback(target_device).into_iter().map(|(b, _)| b));
177 }
178 let mut unique: Vec<BackendId> = Vec::with_capacity(backends.len());
180 for b in backends {
181 if !unique.contains(&b) {
182 unique.push(b);
183 }
184 }
185 unique
186}
187
188pub fn backend_for_device(loc: DeviceLocation) -> BackendId {
193 match loc {
194 DeviceLocation::Cpu => BackendId::Cpu,
195 DeviceLocation::Cuda { .. } => BackendId::Cuda,
196 DeviceLocation::Vulkan { .. } => BackendId::Vulkan,
197 DeviceLocation::Metal { .. } => BackendId::Metal,
198 }
199}
200
201pub struct PlanOptions<'env> {
292 pub precision_requirement: PrecisionRequirement,
295 pub populate_costs: bool,
299 pub placements_for_device: Option<&'env (dyn Fn(DeviceLocation) -> Vec<BackendId> + 'env)>,
302 pub pinned_device: Option<DeviceLocation>,
305 pub fallback_placements_for:
308 Option<&'env (dyn Fn(DeviceLocation) -> Vec<(BackendId, DeviceLocation)> + 'env)>,
309 pub capabilities_for: Option<&'env CapabilitiesLookup<'env>>,
312 pub judge: Option<&'env dyn JudgeOracle>,
318 pub transfer_estimator: Option<&'env dyn TransferEstimator>,
339 pub require_transfer_pricing: bool,
356 pub input_residency: Option<&'env (dyn Fn(NodeId) -> Option<DeviceLocation> + 'env)>,
368 pub reuse_plan: Option<&'env ExecutionPlan>,
385 pub allow_cost_placement: bool,
404 #[cfg(feature = "telemetry")]
412 pub telemetry: Option<&'env crate::telemetry::TelemetryHooks<'env>>,
413}
414
415impl Default for PlanOptions<'_> {
416 fn default() -> Self {
417 Self {
418 precision_requirement: PrecisionRequirement::default(),
419 populate_costs: true,
420 placements_for_device: None,
421 pinned_device: None,
422 fallback_placements_for: None,
423 capabilities_for: None,
424 judge: None,
425 transfer_estimator: None,
426 require_transfer_pricing: false,
427 input_residency: None,
428 reuse_plan: None,
429 allow_cost_placement: true,
430 #[cfg(feature = "telemetry")]
431 telemetry: None,
432 }
433 }
434}
435
436impl<'env> PlanOptions<'env> {
437 pub fn new() -> Self {
439 Self::default()
440 }
441
442 pub fn with_precision_requirement(mut self, req: PrecisionRequirement) -> Self {
444 self.precision_requirement = req;
445 self
446 }
447
448 pub fn without_cost_population(mut self) -> Self {
452 self.populate_costs = false;
453 self
454 }
455
456 pub fn with_placements_for_device(
459 mut self,
460 f: &'env (dyn Fn(DeviceLocation) -> Vec<BackendId> + 'env),
461 ) -> Self {
462 self.placements_for_device = Some(f);
463 self
464 }
465
466 pub fn with_pinned_device(mut self, device: DeviceLocation) -> Self {
471 self.pinned_device = Some(device);
472 self
473 }
474
475 pub fn without_cost_placement(mut self) -> Self {
480 self.allow_cost_placement = false;
481 self
482 }
483
484 pub fn require_transfer_pricing(mut self) -> Self {
492 self.require_transfer_pricing = true;
493 self
494 }
495
496 pub fn with_fallback_placements_for(
505 mut self,
506 f: &'env (dyn Fn(DeviceLocation) -> Vec<(BackendId, DeviceLocation)> + 'env),
507 ) -> Self {
508 self.fallback_placements_for = Some(f);
509 self
510 }
511
512 pub fn with_capabilities_for(mut self, f: &'env CapabilitiesLookup<'env>) -> Self {
515 self.capabilities_for = Some(f);
516 self
517 }
518
519 pub fn with_judge(mut self, judge: &'env dyn JudgeOracle) -> Self {
523 self.judge = Some(judge);
524 self
525 }
526
527 pub fn with_transfer_estimator(mut self, est: &'env dyn TransferEstimator) -> Self {
533 self.transfer_estimator = Some(est);
534 self
535 }
536
537 pub fn with_input_residency(
541 mut self,
542 f: &'env (dyn Fn(NodeId) -> Option<DeviceLocation> + 'env),
543 ) -> Self {
544 self.input_residency = Some(f);
545 self
546 }
547
548 pub fn with_reuse_plan(mut self, base: &'env ExecutionPlan) -> Self {
551 self.reuse_plan = Some(base);
552 self
553 }
554
555 #[cfg(feature = "telemetry")]
561 pub fn with_telemetry(mut self, hooks: &'env crate::telemetry::TelemetryHooks<'env>) -> Self {
562 self.telemetry = Some(hooks);
563 self
564 }
565}
566
567pub fn compile_plan(
650 graph: &Graph,
651 order: &[NodeId],
652 bindings_table: &KernelBindingTable,
653 options: &PlanOptions<'_>,
654) -> Result<ExecutionPlan> {
655 if options.require_transfer_pricing && options.transfer_estimator.is_none() {
663 return Err(Error::Msg(
664 "compile_plan: require_transfer_pricing is set but no transfer \
665 estimator is wired — placement would silently price transfers at \
666 zero (latency-greedy). Wire PlanOptions::with_transfer_estimator(..) \
667 or drop the pricing requirement."
668 .into(),
669 )
670 .bt());
671 }
672
673 let generation = crate::dispatch::topology_generation();
679 let mut alternatives_map = HashMap::with_capacity(order.len());
680
681 let mut residency: HashMap<NodeId, DeviceLocation> = HashMap::with_capacity(order.len());
685
686 let mut dp = PlacementDp::new();
689 let mut dp_drafts: HashMap<NodeId, NodeDraft> = HashMap::new();
690
691 for &id in order {
692 let node = graph.node(id);
693
694 if let Some(base) = options.reuse_plan
700 && let Some(set) = base.alternatives(id)
701 {
702 let resident = early_residency(graph, options, id, node)
703 .or_else(|| set.winner().map(|w| w.device));
704 if let Some(loc) = resident {
705 residency.insert(id, loc);
706 }
707 alternatives_map.insert(id, set.clone());
708 continue;
709 }
710
711 let early = early_residency(graph, options, id, node);
714
715 let mut dp_row_opened = false;
716 let winner_device = match build_node_draft(graph, id, node, bindings_table, options)? {
717 Some(draft) => {
718 let devices = distinct_devices(&draft.set);
719 let dp_estimator = options
720 .transfer_estimator
721 .filter(|_| draft.relaxed && early.is_none() && devices.len() >= 2);
722 if let Some(est) = dp_estimator {
723 open_dp_row(
727 &mut dp,
728 graph,
729 id,
730 node,
731 &draft.set,
732 &devices,
733 &residency,
734 est,
735 options.capabilities_for,
736 );
737 dp_drafts.insert(id, draft);
738 dp_row_opened = true;
739 None
740 } else {
741 if let Some(est) = options.transfer_estimator
748 && let [single] = devices.as_slice()
749 {
750 close_open_inputs(&mut dp, graph, node, *single, est, &mut residency);
751 }
752 let set = finalize_node_greedy(
753 graph,
754 node,
755 draft,
756 bindings_table,
757 options,
758 &residency,
759 )?;
760 let dev = set.winner().map(|w| w.device);
761 if devices.len() > 1
765 && let (Some(est), Some(wd)) = (options.transfer_estimator, dev)
766 {
767 close_open_inputs(&mut dp, graph, node, wd, est, &mut residency);
768 }
769 alternatives_map.insert(id, set);
770 dev
771 }
772 }
773 None => {
774 if let fuel_graph::Op::Copy { target } | fuel_graph::Op::Move { target } = node.op
780 && let Some(est) = options.transfer_estimator
781 {
782 close_open_inputs(&mut dp, graph, node, target, est, &mut residency);
783 }
784 None
785 }
786 };
787
788 if dp_row_opened {
789 continue;
793 }
794
795 let resident = early
797 .or(winner_device)
798 .or_else(|| {
799 graph.target_backend(id).map(|b| {
804 options
805 .pinned_device
806 .unwrap_or_else(|| default_device_for(b))
807 })
808 })
809 .or_else(|| {
810 if is_passthrough(node) {
815 node.inputs.first().and_then(|i| residency.get(i)).copied()
816 } else {
817 None
818 }
819 });
820 if let Some(loc) = resident {
821 residency.insert(id, loc);
822 } else if is_passthrough(node) {
823 if let Some(&input) = node.inputs.first() {
829 let root = dp.resolve(input);
830 if let Some(dev) = dp.committed_device(root) {
831 residency.insert(id, dev);
832 } else {
833 dp.add_alias(id, root);
834 }
835 }
836 }
837 }
838
839 if !dp_drafts.is_empty() {
842 let Some(est) = options.transfer_estimator else {
843 return Err(Error::Msg(
846 "compile_plan: placement DP rows exist without a transfer \
847 estimator — internal invariant violated"
848 .into(),
849 )
850 .bt());
851 };
852 let commits = dp.finish(options.pinned_device, |nid| node_bytes(graph, nid), est);
853 for (n, d) in commits {
854 residency.insert(n, d);
855 }
856 for &id in order {
860 if residency.contains_key(&id) {
861 continue;
862 }
863 let node = graph.node(id);
864 if is_passthrough(node)
865 && let Some(&loc) = node.inputs.first().and_then(|i| residency.get(i))
866 {
867 residency.insert(id, loc);
868 }
869 }
870 for &id in order {
875 let Some(draft) = dp_drafts.remove(&id) else {
876 continue;
877 };
878 let chosen = dp.committed_device(id).ok_or_else(|| {
879 Error::Msg(format!(
880 "compile_plan: placement DP left node {:?} uncommitted — \
881 internal invariant violated",
882 id,
883 ))
884 .bt()
885 })?;
886 let node = graph.node(id);
887 let mut set = draft.set;
888 let inputs = priced_inputs_for(graph, node, &residency);
889 apply_inbound_transfer_costs(&mut set, &inputs, est);
890 set.rank_by_composite_cost();
891 let keep: Vec<usize> = set
892 .alternatives()
893 .iter()
894 .enumerate()
895 .filter_map(|(i, c)| (c.device == chosen).then_some(i))
896 .collect();
897 if keep.len() < set.len() {
898 set.retain_indices(&keep);
899 }
900 set.retain_per_device_frontier(KEEP_PER_DEVICE);
906 if set.is_empty() {
907 return Err(missing_binding_error(
910 bindings_table,
911 draft.op_kind,
912 &draft.dtypes,
913 draft.diag_backend,
914 ));
915 }
916 alternatives_map.insert(id, set);
917 }
918 }
919
920 #[cfg(feature = "telemetry")]
925 emit_plan_telemetry(graph, &alternatives_map, bindings_table, options);
926
927 Ok(ExecutionPlan {
928 order: order.to_vec(),
929 alternatives: alternatives_map,
930 generation,
931 })
932}
933
934#[cfg(feature = "telemetry")]
959fn emit_plan_telemetry(
960 graph: &Graph,
961 alternatives: &HashMap<NodeId, AlternativeSet>,
962 bindings_table: &KernelBindingTable,
963 options: &PlanOptions<'_>,
964) {
965 use crate::telemetry::{
966 DispatchRecord, FdxOperandDesc, ImplId, TELEMETRY_SCHEMA_VERSION, detect_miss_precomputed,
967 };
968
969 let Some(hooks) = options.telemetry else {
971 return;
972 };
973 if !hooks.config.is_enabled() {
974 return;
975 }
976 let arch = hooks.arch_tag();
977 let want_candidates = hooks.config.mode.wants_candidates();
978
979 for (&id, set) in alternatives {
980 let Some(winner) = set.winner() else {
981 continue;
982 };
983 let node = graph.node(id);
984 let Some(op_kind) = op_to_op_kind(&node.op) else {
985 continue;
986 };
987 let dtypes = build_lookup_dtypes(graph, node);
988
989 let winner_ptr = winner.kernel as *const () as usize;
994 let winner_entry = bindings_table
995 .lookup_alternatives(op_kind, &dtypes, winner.backend)
996 .iter()
997 .find(|e| e.kernel as *const () as usize == winner_ptr)
998 .copied();
999 let is_generic = winner_entry.is_some_and(|e| e.is_generic);
1000 let winner_revision = winner_entry.map_or(0, |e| e.kernel_revision_hash);
1001
1002 let mut operands: Vec<FdxOperandDesc> = node
1008 .inputs
1009 .iter()
1010 .map(|&input_id| {
1011 FdxOperandDesc::from_layout(&graph.layout(input_id), graph.node(input_id).dtype)
1012 })
1013 .collect();
1014 operands.push(FdxOperandDesc::from_layout(&graph.layout(id), node.dtype));
1015 let op_class = format!("{op_kind:?}");
1016
1017 let structure_key = hooks.provider.structure_key(&op_class, &operands, &arch);
1021
1022 let chosen = ImplId {
1024 backend: winner.backend,
1025 op: op_kind,
1026 dtypes: dtypes.to_vec(),
1027 kernel_source: winner.kernel_source.to_string(),
1028 kernel_revision_hash: winner_revision,
1029 };
1030
1031 let candidates = if want_candidates {
1034 build_candidates(set, op_kind, &dtypes, bindings_table, options.judge)
1035 } else {
1036 Vec::new()
1037 };
1038
1039 let dispatch = DispatchRecord {
1041 schema: TELEMETRY_SCHEMA_VERSION,
1042 structure_key,
1043 chosen: chosen.clone(),
1044 candidates,
1045 count: 1,
1046 hw: hooks.hw.clone(),
1047 };
1048 if let Ok(mut sink) = hooks.sink.lock() {
1049 sink.record_dispatch(dispatch);
1050 }
1051
1052 if let Some(miss) = detect_miss_precomputed(
1056 is_generic,
1057 chosen,
1058 &op_class,
1059 &operands,
1060 &arch,
1061 hooks.provider,
1062 hooks.hw.clone(),
1063 ) {
1064 if let Ok(mut sink) = hooks.sink.lock() {
1067 sink.record_miss(miss);
1068 }
1069 }
1070 }
1071}
1072
1073#[cfg(feature = "telemetry")]
1082fn build_candidates(
1083 set: &AlternativeSet,
1084 op_kind: OpKind,
1085 dtypes: &[DType],
1086 bindings_table: &KernelBindingTable,
1087 judge: Option<&dyn JudgeOracle>,
1088) -> Vec<crate::telemetry::Candidate> {
1089 use crate::telemetry::{Candidate as TeleCandidate, ImplId};
1090
1091 let ctx = set.context();
1094 set.alternatives()
1095 .iter()
1096 .map(|cand| {
1097 let ptr = cand.kernel as *const () as usize;
1098 let revision = bindings_table
1099 .lookup_alternatives(op_kind, dtypes, cand.backend)
1100 .iter()
1101 .find(|e| e.kernel as *const () as usize == ptr)
1102 .map_or(0, |e| e.kernel_revision_hash);
1103 let impl_id = ImplId {
1104 backend: cand.backend,
1105 op: op_kind,
1106 dtypes: dtypes.to_vec(),
1107 kernel_source: cand.kernel_source.to_string(),
1108 kernel_revision_hash: revision,
1109 };
1110 let latency_ns = match (judge, ctx) {
1111 (Some(j), Some(ctx)) => j.measured_latency_ns(
1112 ctx.op,
1113 ctx.principal_dtype,
1114 ctx.size_class,
1115 cand.backend,
1116 cand.kernel_source,
1117 revision,
1123 ),
1124 _ => None,
1125 };
1126 TeleCandidate {
1127 impl_id,
1128 latency_ns,
1129 }
1130 })
1131 .collect()
1132}
1133
1134pub fn node_needs_plan_entry(op: &fuel_graph::Op) -> bool {
1142 op_to_op_kind(op).is_some()
1143 && !matches!(
1144 op,
1145 fuel_graph::Op::Copy { .. } | fuel_graph::Op::Move { .. }
1146 )
1147}
1148
1149fn is_passthrough(node: &fuel_graph::Node) -> bool {
1153 node.op.is_view_op()
1154 || matches!(
1155 node.op,
1156 fuel_graph::Op::Reshape(_) | fuel_graph::Op::Contiguize
1157 )
1158}
1159
1160fn distinct_devices(set: &AlternativeSet) -> SmallVec<[DeviceLocation; 4]> {
1164 let mut out: SmallVec<[DeviceLocation; 4]> = SmallVec::new();
1165 for c in set.alternatives() {
1166 if !out.contains(&c.device) {
1167 out.push(c.device);
1168 }
1169 }
1170 out
1171}
1172
1173fn node_bytes(graph: &Graph, id: NodeId) -> u64 {
1177 let n = graph.node(id);
1178 (n.shape.elem_count() as u64).saturating_mul(n.dtype.size_in_bytes() as u64)
1179}
1180
1181fn open_dp_row(
1186 dp: &mut PlacementDp,
1187 graph: &Graph,
1188 id: NodeId,
1189 node: &fuel_graph::Node,
1190 set: &AlternativeSet,
1191 devices: &[DeviceLocation],
1192 residency: &HashMap<NodeId, DeviceLocation>,
1193 est: &dyn TransferEstimator,
1194 caps_for: Option<&CapabilitiesLookup<'_>>,
1195) {
1196 let rates_for = |backend: BackendId| -> (f64, f64) {
1205 caps_for
1206 .and_then(|f| f(backend))
1207 .map(|caps| {
1208 (
1209 caps.compute_throughput_flops_per_ns,
1210 caps.mem_bandwidth_bytes_per_ns,
1211 )
1212 })
1213 .unwrap_or_else(|| default_backend_rates(backend))
1214 };
1215 let device_costs: Vec<(DeviceLocation, u64)> = devices
1216 .iter()
1217 .map(|&d| {
1218 let min = set
1219 .alternatives()
1220 .iter()
1221 .filter(|c| c.device == d)
1222 .map(|c| {
1223 let (cr, bw) = rates_for(c.backend);
1224 composite_ns(&c.static_cost, cr, bw)
1225 })
1226 .min()
1227 .unwrap_or(u64::MAX);
1228 (d, min)
1229 })
1230 .collect();
1231
1232 let mut fixed: Vec<(DeviceLocation, u64)> = Vec::new();
1233 let mut chains: Vec<ChainInput> = Vec::new();
1234 for &input in &node.inputs {
1235 let bytes = node_bytes(graph, input);
1236 if let Some(&src) = residency.get(&input) {
1237 fixed.push((src, bytes));
1238 continue;
1239 }
1240 let root = dp.resolve(input);
1241 if dp.is_open(root) {
1242 match chains.iter_mut().find(|ci| ci.producer == root) {
1243 Some(ci) => ci.edge_bytes.push(bytes),
1244 None => chains.push(ChainInput {
1245 producer: root,
1246 edge_bytes: smallvec::smallvec![bytes],
1247 }),
1248 }
1249 } else if let Some(dev) = dp.committed_device(root) {
1250 fixed.push((dev, bytes));
1253 }
1254 }
1257 dp.push_row(id, &device_costs, &fixed, &chains, est);
1258}
1259
1260fn close_open_inputs(
1266 dp: &mut PlacementDp,
1267 graph: &Graph,
1268 node: &fuel_graph::Node,
1269 toward: DeviceLocation,
1270 est: &dyn TransferEstimator,
1271 residency: &mut HashMap<NodeId, DeviceLocation>,
1272) {
1273 for &input in &node.inputs {
1274 if residency.contains_key(&input) {
1275 continue;
1276 }
1277 let root = dp.resolve(input);
1278 if dp.is_open(root) {
1279 let commits = dp.close_toward(root, toward, &[node_bytes(graph, input)], est);
1280 for (n, d) in commits {
1281 residency.insert(n, d);
1282 }
1283 }
1284 if !residency.contains_key(&input)
1285 && let Some(d) = dp.committed_device(root)
1286 {
1287 residency.insert(input, d);
1288 }
1289 }
1290}
1291
1292fn early_residency(
1301 graph: &Graph,
1302 options: &PlanOptions<'_>,
1303 id: NodeId,
1304 node: &fuel_graph::Node,
1305) -> Option<DeviceLocation> {
1306 match node.op {
1307 fuel_graph::Op::Copy { target }
1308 | fuel_graph::Op::Move { target }
1309 | fuel_graph::Op::Alloc { target } => return Some(target),
1310 _ => {}
1311 }
1312 if let Some(loc) = graph.placement(id) {
1313 return Some(loc);
1314 }
1315 options.input_residency.and_then(|f| f(id))
1316}
1317
1318fn priced_inputs_for(
1324 graph: &Graph,
1325 node: &fuel_graph::Node,
1326 residency: &HashMap<NodeId, DeviceLocation>,
1327) -> Vec<(DeviceLocation, u64)> {
1328 node.inputs
1329 .iter()
1330 .filter_map(|&input_id| {
1331 let src = residency.get(&input_id).copied()?;
1332 let in_node = graph.node(input_id);
1333 let bytes = (in_node.shape.elem_count() as u64)
1334 .saturating_mul(in_node.dtype.size_in_bytes() as u64);
1335 Some((src, bytes))
1336 })
1337 .collect()
1338}
1339
1340struct NodeDraft {
1346 set: AlternativeSet,
1347 relaxed: bool,
1350 from_fallback: bool,
1353 costed: bool,
1357 op_kind: OpKind,
1358 dtypes: Vec<DType>,
1359 diag_backend: BackendId,
1360}
1361
1362fn build_node_draft(
1369 graph: &Graph,
1370 id: NodeId,
1371 node: &fuel_graph::Node,
1372 bindings_table: &KernelBindingTable,
1373 options: &PlanOptions<'_>,
1374) -> Result<Option<NodeDraft>> {
1375 if !node_needs_plan_entry(&node.op) {
1391 return Ok(None);
1392 }
1393 let Some(op_kind) = op_to_op_kind(&node.op) else {
1394 return Ok(None);
1396 };
1397
1398 let explicit_backend = graph.target_backend(id);
1402 let target_device = decision_device_for(graph, options, id).ok_or_else(|| {
1403 Error::Msg(format!(
1404 "compile_plan: node {:?} ({:?}) has no device context — \
1405 set PlanOptions::pinned_device, a graph placement, or \
1406 the node's target_backend",
1407 id, node.op,
1408 ))
1409 .bt()
1410 })?;
1411 let diag_backend = explicit_backend.unwrap_or_else(|| backend_for_device(target_device));
1415 let dtypes = build_lookup_dtypes(graph, node);
1416
1417 let placements: Vec<(BackendId, DeviceLocation)> = match options.placements_for_device {
1421 Some(f) => f(target_device)
1422 .into_iter()
1423 .map(|backend| (backend, target_device))
1424 .collect(),
1425 None => vec![(diag_backend, target_device)],
1426 };
1427
1428 let op_params = candidate_default_op_params(graph, node);
1429
1430 let pricing_active = options.transfer_estimator.is_some()
1448 && options.populate_costs
1449 && options.capabilities_for.is_some();
1450 let fallback_allowed =
1451 node.op.destructive_input().is_none() && options.fallback_placements_for.is_some();
1452 let relaxed = options.allow_cost_placement
1453 && pricing_active
1454 && fallback_allowed
1455 && graph.placement(id).is_none();
1456
1457 let mut set;
1458 let mut from_fallback = false;
1459 if relaxed {
1460 let mut merged = placements.clone();
1465 if let Some(fallback) = options.fallback_placements_for {
1466 merged.extend(fallback(target_device));
1467 }
1468 set = enumerate_candidates(op_kind, &dtypes, &merged, &op_params, bindings_table);
1469 } else {
1470 set = enumerate_candidates(op_kind, &dtypes, &placements, &op_params, bindings_table);
1471
1472 if set.is_empty()
1480 && fallback_allowed
1481 && let Some(fallback) = options.fallback_placements_for
1482 {
1483 let fb_placements = fallback(target_device);
1484 if !fb_placements.is_empty() {
1485 set = enumerate_candidates(
1486 op_kind,
1487 &dtypes,
1488 &fb_placements,
1489 &op_params,
1490 bindings_table,
1491 );
1492 from_fallback = !set.is_empty();
1493 }
1494 }
1495 }
1496
1497 if set.is_empty() {
1502 return Err(missing_binding_error(
1503 bindings_table,
1504 op_kind,
1505 &dtypes,
1506 diag_backend,
1507 ));
1508 }
1509
1510 let input_layouts: Vec<fuel_ir::Layout> = node
1512 .inputs
1513 .iter()
1514 .map(|&input_id| graph.layout(input_id))
1515 .collect();
1516 let ctx = FilterContext::new(op_kind, &dtypes, &input_layouts);
1517 let chain = default_chain(options.precision_requirement);
1518 if let Err(err) = apply_filter_chain(&mut set, &chain, &ctx) {
1519 if relaxed || from_fallback || !fallback_allowed {
1528 return Err(err);
1529 }
1530 let Some(fallback) = options.fallback_placements_for else {
1531 return Err(err);
1532 };
1533 let fb_placements = fallback(target_device);
1534 if fb_placements.is_empty() {
1535 return Err(err);
1536 }
1537 let mut fb_set =
1538 enumerate_candidates(op_kind, &dtypes, &fb_placements, &op_params, bindings_table);
1539 if fb_set.is_empty() {
1540 return Err(err);
1541 }
1542 apply_filter_chain(&mut fb_set, &chain, &ctx)?;
1545 set = fb_set;
1546 from_fallback = true;
1547 }
1548
1549 if let Some(&principal_dtype) = dtypes.first() {
1560 let input_shapes: Vec<Shape> = node
1561 .inputs
1562 .iter()
1563 .map(|&input_id| graph.node(input_id).shape.clone())
1564 .collect();
1565 let size_class = SizeClass::for_op(op_kind, &input_shapes);
1566 set.set_context(DecisionContext {
1567 op: op_kind,
1568 principal_dtype,
1569 size_class,
1570 });
1571 }
1572
1573 let costed = if options.populate_costs {
1577 if let Some(caps_for) = options.capabilities_for {
1578 let input_shapes: Vec<Shape> = node
1579 .inputs
1580 .iter()
1581 .map(|&input_id| graph.node(input_id).shape.clone())
1582 .collect();
1583 compute_static_costs(
1584 &mut set,
1585 op_kind,
1586 &dtypes,
1587 &input_shapes,
1588 bindings_table,
1589 caps_for,
1590 options.judge,
1591 );
1592 true
1593 } else {
1594 false
1595 }
1596 } else {
1597 false
1598 };
1599
1600 Ok(Some(NodeDraft {
1601 set,
1602 relaxed,
1603 from_fallback,
1604 costed,
1605 op_kind,
1606 dtypes,
1607 diag_backend,
1608 }))
1609}
1610
1611fn finalize_node_greedy(
1618 graph: &Graph,
1619 node: &fuel_graph::Node,
1620 draft: NodeDraft,
1621 bindings_table: &KernelBindingTable,
1622 options: &PlanOptions<'_>,
1623 residency: &HashMap<NodeId, DeviceLocation>,
1624) -> Result<AlternativeSet> {
1625 let NodeDraft {
1626 mut set,
1627 relaxed,
1628 from_fallback,
1629 costed,
1630 op_kind,
1631 dtypes,
1632 diag_backend,
1633 } = draft;
1634
1635 if costed {
1636 if let Some(estimator) = options.transfer_estimator {
1642 let inputs = priced_inputs_for(graph, node, residency);
1643 apply_inbound_transfer_costs(&mut set, &inputs, estimator);
1644 }
1645 set.rank_by_composite_cost();
1646 }
1647
1648 if relaxed && let Some(winner_device) = set.winner().map(|c| c.device) {
1659 let keep: Vec<usize> = set
1660 .alternatives()
1661 .iter()
1662 .enumerate()
1663 .filter_map(|(i, c)| (c.device == winner_device).then_some(i))
1664 .collect();
1665 if keep.len() < set.len() {
1666 set.retain_indices(&keep);
1667 }
1668 }
1669
1670 set.retain_per_device_frontier(KEEP_PER_DEVICE);
1678
1679 if from_fallback && set.len() > 1 {
1686 set.retain_indices(&[0]);
1687 }
1688
1689 if set.is_empty() {
1696 return Err(missing_binding_error(
1697 bindings_table,
1698 op_kind,
1699 &dtypes,
1700 diag_backend,
1701 ));
1702 }
1703
1704 Ok(set)
1705}
1706
1707fn missing_binding_error(
1713 table: &KernelBindingTable,
1714 op: OpKind,
1715 dtypes: &[DType],
1716 backend: BackendId,
1717) -> Error {
1718 let _ = backend;
1723 let (available_backends, supported_combinations) = table.coverage_diagnostic();
1724 Error::NoBackendForOp {
1725 op,
1726 dtypes: dtypes.to_vec(),
1727 available_backends,
1728 supported_combinations,
1729 }
1730 .bt()
1731}
1732
1733fn candidate_default_op_params(
1741 _graph: &Graph,
1742 _node: &fuel_graph::Node,
1743) -> crate::kernel::OpParams {
1744 crate::kernel::OpParams::None
1745}
1746
1747#[cfg(test)]
1748mod tests {
1749 use super::*;
1750 use crate::fused::PrecisionGuarantee;
1751 use crate::kernel::{KernelCaps, OpParams, unknown_cost};
1752 use fuel_graph::{Node, Op, topo_order};
1753 use fuel_ir::backend::{BackendCapabilities, SubstrateClass, TransferPath};
1754 use fuel_ir::{DType, Layout, Result as FuelResult, Shape, StrideVec};
1755 use fuel_memory::Storage;
1756 use smallvec::smallvec;
1757 use std::collections::HashSet;
1758 use std::sync::{Arc, RwLock};
1759
1760 fn noop_kernel(
1761 _i: &[Arc<RwLock<Storage>>],
1762 _o: &mut [Arc<RwLock<Storage>>],
1763 _l: &[Layout],
1764 _p: &OpParams,
1765 ) -> FuelResult<()> {
1766 Ok(())
1767 }
1768
1769 fn noop_kernel_b(
1770 _i: &[Arc<RwLock<Storage>>],
1771 _o: &mut [Arc<RwLock<Storage>>],
1772 _l: &[Layout],
1773 _p: &OpParams,
1774 ) -> FuelResult<()> {
1775 Ok(())
1776 }
1777
1778 fn noop_kernel_c(
1779 _i: &[Arc<RwLock<Storage>>],
1780 _o: &mut [Arc<RwLock<Storage>>],
1781 _l: &[Layout],
1782 _p: &OpParams,
1783 ) -> FuelResult<()> {
1784 Ok(())
1785 }
1786
1787 fn register_add_f32(
1788 table: &mut KernelBindingTable,
1789 backend: BackendId,
1790 kernel: crate::kernel::KernelRef,
1791 precision: PrecisionGuarantee,
1792 ) {
1793 table.register_full(
1794 OpKind::AddElementwise,
1795 &[DType::F32, DType::F32, DType::F32],
1796 backend,
1797 kernel,
1798 KernelCaps::empty(),
1799 precision,
1800 unknown_cost,
1801 );
1802 }
1803
1804 fn build_add_graph() -> (Graph, NodeId) {
1805 let mut g = Graph::new();
1806 let lhs = g.push(Node {
1807 op: Op::Const,
1808 inputs: vec![],
1809 shape: Shape::from_dims(&[3]),
1810 dtype: DType::F32,
1811 });
1812 let rhs = g.push(Node {
1813 op: Op::Const,
1814 inputs: vec![],
1815 shape: Shape::from_dims(&[3]),
1816 dtype: DType::F32,
1817 });
1818 let add = g.push(Node {
1819 op: Op::Add,
1820 inputs: vec![lhs, rhs],
1821 shape: Shape::from_dims(&[3]),
1822 dtype: DType::F32,
1823 });
1824 g.set_target_backend(add, BackendId::Cpu);
1825 (g, add)
1826 }
1827
1828 fn cpu_caps() -> BackendCapabilities {
1829 BackendCapabilities {
1830 backend_id: BackendId::Cpu,
1831 device_location: DeviceLocation::Cpu,
1832 op_dtype_support: HashSet::new(),
1833 required_alignment: 64,
1834 access_granularity_bits: 8,
1835 transfer_paths: vec![(DeviceLocation::Cpu, TransferPath::SameDevice)],
1836 storage_substrate: SubstrateClass::HostBytes,
1837 compute_throughput_flops_per_ns: 1.0,
1838 mem_bandwidth_bytes_per_ns: 4.0,
1839 }
1840 }
1841
1842 #[test]
1843 fn empty_plan_has_no_alternatives() {
1844 let p = ExecutionPlan::empty();
1845 assert!(p.order.is_empty());
1846 assert!(p.alternatives.is_empty());
1847 assert!(p.alternatives(NodeId(0)).is_none());
1848 }
1849
1850 #[test]
1851 fn default_device_per_backend_matches_router_convention() {
1852 assert_eq!(default_device_for(BackendId::Cpu), DeviceLocation::Cpu);
1853 assert_eq!(
1854 default_device_for(BackendId::Cuda),
1855 DeviceLocation::Cuda { gpu_id: 0 },
1856 );
1857 assert_eq!(
1858 default_device_for(BackendId::Vulkan),
1859 DeviceLocation::Vulkan { gpu_id: 0 },
1860 );
1861 }
1862
1863 #[test]
1864 fn compile_plan_skips_view_and_const_nodes() {
1865 let mut table = KernelBindingTable::new();
1866 register_add_f32(
1867 &mut table,
1868 BackendId::Cpu,
1869 noop_kernel,
1870 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
1871 );
1872
1873 let mut g = Graph::new();
1874 let lhs = g.push(Node {
1875 op: Op::Const,
1876 inputs: vec![],
1877 shape: Shape::from_dims(&[6]),
1878 dtype: DType::F32,
1879 });
1880 let rhs = g.push(Node {
1881 op: Op::Const,
1882 inputs: vec![],
1883 shape: Shape::from_dims(&[6]),
1884 dtype: DType::F32,
1885 });
1886 let add = g.push(Node {
1887 op: Op::Add,
1888 inputs: vec![lhs, rhs],
1889 shape: Shape::from_dims(&[6]),
1890 dtype: DType::F32,
1891 });
1892 let reshape = g.push(Node {
1893 op: Op::Reshape(Shape::from_dims(&[2, 3])),
1894 inputs: vec![add],
1895 shape: Shape::from_dims(&[2, 3]),
1896 dtype: DType::F32,
1897 });
1898 g.set_target_backend(add, BackendId::Cpu);
1899 g.set_target_backend(reshape, BackendId::Cpu);
1900
1901 let order = topo_order(&g, reshape);
1902 let opts = PlanOptions::new().without_cost_population();
1903 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
1904
1905 assert_eq!(plan.order.len(), 4);
1906 assert_eq!(
1907 plan.alternatives.len(),
1908 1,
1909 "only Add carries an alternative set",
1910 );
1911 let alts = plan.alternatives(add).expect("Add alternatives present");
1912 assert_eq!(alts.len(), 1);
1913 assert!(plan.alternatives(lhs).is_none());
1914 assert!(plan.alternatives(rhs).is_none());
1915 assert!(plan.alternatives(reshape).is_none());
1916 }
1917
1918 #[test]
1919 fn compile_plan_fails_fast_on_missing_binding() {
1920 let table = KernelBindingTable::new();
1921 let mut g = Graph::new();
1922 let lhs = g.push(Node {
1923 op: Op::Const,
1924 inputs: vec![],
1925 shape: Shape::from_dims(&[2, 3]),
1926 dtype: DType::F32,
1927 });
1928 let rhs = g.push(Node {
1929 op: Op::Const,
1930 inputs: vec![],
1931 shape: Shape::from_dims(&[3, 2]),
1932 dtype: DType::F32,
1933 });
1934 let mm = g.push(Node {
1935 op: Op::MatMul,
1936 inputs: vec![lhs, rhs],
1937 shape: Shape::from_dims(&[2, 2]),
1938 dtype: DType::F32,
1939 });
1940 g.set_target_backend(mm, BackendId::Cpu);
1941
1942 let order = topo_order(&g, mm);
1943 let opts = PlanOptions::new().without_cost_population();
1944 let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
1945 match err {
1946 Error::NoBackendForOp { op, dtypes, .. } => {
1947 assert_eq!(op, OpKind::MatMul);
1948 assert_eq!(dtypes, vec![DType::F32, DType::F32, DType::F32]);
1949 }
1950 other => panic!("expected NoBackendForOp, got {other:?}"),
1951 }
1952 }
1953
1954 #[test]
1965 fn plan_time_miss_reports_real_coverage_not_empty() {
1966 let bindings = crate::dispatch::global_bindings(); let mut g = Graph::new();
1968 let lhs = g.push(Node {
1969 op: Op::Const,
1970 inputs: vec![],
1971 shape: Shape::from_dims(&[2, 3]),
1972 dtype: DType::F32,
1973 });
1974 let rhs = g.push(Node {
1976 op: Op::Const,
1977 inputs: vec![],
1978 shape: Shape::from_dims(&[3, 2]),
1979 dtype: DType::BF16,
1980 });
1981 let mm = g.push(Node {
1982 op: Op::MatMul,
1983 inputs: vec![lhs, rhs],
1984 shape: Shape::from_dims(&[2, 2]),
1985 dtype: DType::F32,
1986 });
1987 g.set_target_backend(mm, BackendId::Cpu);
1988
1989 let order = topo_order(&g, mm);
1990 let opts = PlanOptions::new().without_cost_population();
1991 let err = compile_plan(&g, &order, &bindings, &opts).unwrap_err();
1992 match err {
1993 Error::NoBackendForOp {
1994 op,
1995 dtypes,
1996 available_backends,
1997 supported_combinations,
1998 } => {
1999 assert_eq!(op, OpKind::MatMul);
2000 assert_eq!(
2001 dtypes,
2002 vec![DType::F32, DType::BF16, DType::F32],
2003 "the mixed-precision key (F32 activation, BF16 weight)",
2004 );
2005 assert!(
2007 available_backends.contains(&BackendId::Cpu),
2008 "a populated table must report Cpu, never []: {available_backends:?}",
2009 );
2010 assert!(
2011 supported_combinations
2012 .iter()
2013 .any(|(b, o, d)| *b == BackendId::Cpu
2014 && *o == OpKind::MatMul
2015 && *d == vec![DType::F32, DType::F32, DType::F32]),
2016 "must show uniform F32 matmul IS registered (so the mixed gap is visible): \
2017 {supported_combinations:?}",
2018 );
2019 }
2020 other => panic!("expected NoBackendForOp, got {other:?}"),
2021 }
2022 }
2023
2024 #[test]
2025 fn compile_plan_legacy_single_backend_uses_target_backend() {
2026 let mut table = KernelBindingTable::new();
2027 register_add_f32(
2028 &mut table,
2029 BackendId::Cpu,
2030 noop_kernel,
2031 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2032 );
2033 register_add_f32(
2037 &mut table,
2038 BackendId::Cuda,
2039 noop_kernel_b,
2040 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2041 );
2042 let (g, add_id) = build_add_graph();
2043 let order = topo_order(&g, add_id);
2044 let opts = PlanOptions::new().without_cost_population();
2045 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2046 let alts = plan.alternatives(add_id).unwrap();
2047 assert_eq!(alts.len(), 1, "legacy mode → single backend → single alt");
2048 assert_eq!(alts.winner().unwrap().backend, BackendId::Cpu);
2049 }
2050
2051 #[test]
2052 fn compile_plan_cross_co_located_via_placements_callback() {
2053 let mut table = KernelBindingTable::new();
2054 register_add_f32(
2055 &mut table,
2056 BackendId::Cpu,
2057 noop_kernel,
2058 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2059 );
2060 register_add_f32(
2061 &mut table,
2062 BackendId::Cuda,
2063 noop_kernel_b,
2064 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2065 );
2066
2067 let (g, add_id) = build_add_graph();
2068 let order = topo_order(&g, add_id);
2069 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2070 if dev == DeviceLocation::Cpu {
2071 vec![BackendId::Cpu, BackendId::Cuda]
2072 } else {
2073 vec![]
2074 }
2075 };
2076 let opts = PlanOptions::new()
2077 .without_cost_population()
2078 .with_placements_for_device(&placements_fn);
2079 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2080 let alts = plan.alternatives(add_id).unwrap();
2081 assert_eq!(
2082 alts.len(),
2083 2,
2084 "cross-co-located mode aggregates Cpu + Aocl alternatives",
2085 );
2086 }
2087
2088 #[test]
2089 fn compile_plan_precision_requirement_filters_via_default_chain() {
2090 let mut table = KernelBindingTable::new();
2091 register_add_f32(
2093 &mut table,
2094 BackendId::Cpu,
2095 noop_kernel,
2096 PrecisionGuarantee {
2097 bit_stable_on_same_hardware: false,
2098 ..PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU
2099 },
2100 );
2101
2102 let (g, add_id) = build_add_graph();
2103 let order = topo_order(&g, add_id);
2104 let opts = PlanOptions::new()
2105 .without_cost_population()
2106 .with_precision_requirement(PrecisionRequirement::BIT_STABLE);
2107 let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
2108 match err {
2109 Error::FilterRejected { filter, .. } => {
2110 assert_eq!(filter, "precision-floor");
2111 }
2112 other => panic!("expected FilterRejected, got {other:?}"),
2113 }
2114 let unconstrained = PlanOptions::new().without_cost_population();
2116 let plan = compile_plan(&g, &order, &table, &unconstrained).expect("compile");
2117 assert!(plan.alternatives(add_id).is_some());
2118 let _ = add_id;
2119 }
2120
2121 #[test]
2122 fn compile_plan_with_cost_ranks_cheaper_winner() {
2123 fn cheap_cost(
2124 _: &[Shape],
2125 _: &[DType],
2126 _: &OpParams,
2127 _: &BackendCapabilities,
2128 ) -> crate::fused::CostEstimate {
2129 crate::fused::CostEstimate {
2130 flops: 10,
2131 bytes_moved: 0,
2132 kernel_overhead_ns: 0,
2133 }
2134 }
2135 fn expensive_cost(
2136 _: &[Shape],
2137 _: &[DType],
2138 _: &OpParams,
2139 _: &BackendCapabilities,
2140 ) -> crate::fused::CostEstimate {
2141 crate::fused::CostEstimate {
2142 flops: 1_000_000,
2143 bytes_moved: 0,
2144 kernel_overhead_ns: 0,
2145 }
2146 }
2147 let mut table = KernelBindingTable::new();
2148 table.register_full(
2149 OpKind::AddElementwise,
2150 &[DType::F32, DType::F32, DType::F32],
2151 BackendId::Cpu,
2152 noop_kernel,
2153 KernelCaps::empty(),
2154 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2155 expensive_cost,
2156 );
2157 table.register_full(
2158 OpKind::AddElementwise,
2159 &[DType::F32, DType::F32, DType::F32],
2160 BackendId::Cuda,
2161 noop_kernel_b,
2162 KernelCaps::empty(),
2163 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2164 cheap_cost,
2165 );
2166
2167 let (g, add_id) = build_add_graph();
2168 let order = topo_order(&g, add_id);
2169 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2170 if dev == DeviceLocation::Cpu {
2171 vec![BackendId::Cpu, BackendId::Cuda]
2172 } else {
2173 vec![]
2174 }
2175 };
2176 let cpu_caps_val = cpu_caps();
2177 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
2178 let opts = PlanOptions::new()
2179 .with_placements_for_device(&placements_fn)
2180 .with_capabilities_for(&caps_fn);
2181 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2182 let alts = plan.alternatives(add_id).unwrap();
2183 assert_eq!(alts.len(), 2);
2184 assert_eq!(
2185 alts.winner().unwrap().backend,
2186 BackendId::Cuda,
2187 "cheaper static cost wins after ranking",
2188 );
2189 }
2190
2191 #[test]
2192 fn compile_plan_judge_layer2_can_flip_layer1_winner() {
2193 use crate::ranker::HashMapJudge;
2197 use fuel_ir::dispatch::SizeClass;
2198
2199 fn cpu_layer1(
2200 _: &[Shape],
2201 _: &[DType],
2202 _: &OpParams,
2203 _: &BackendCapabilities,
2204 ) -> crate::fused::CostEstimate {
2205 crate::fused::CostEstimate {
2206 flops: 10,
2207 bytes_moved: 0,
2208 kernel_overhead_ns: 0,
2209 }
2210 }
2211 fn aocl_layer1(
2212 _: &[Shape],
2213 _: &[DType],
2214 _: &OpParams,
2215 _: &BackendCapabilities,
2216 ) -> crate::fused::CostEstimate {
2217 crate::fused::CostEstimate {
2218 flops: 100_000,
2219 bytes_moved: 0,
2220 kernel_overhead_ns: 0,
2221 }
2222 }
2223 let mut table = KernelBindingTable::new();
2224 table.register_full(
2225 OpKind::AddElementwise,
2226 &[DType::F32, DType::F32, DType::F32],
2227 BackendId::Cpu,
2228 noop_kernel,
2229 KernelCaps::empty(),
2230 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2231 cpu_layer1,
2232 );
2233 table.register_full(
2234 OpKind::AddElementwise,
2235 &[DType::F32, DType::F32, DType::F32],
2236 BackendId::Cuda,
2237 noop_kernel_b,
2238 KernelCaps::empty(),
2239 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2240 aocl_layer1,
2241 );
2242
2243 let (g, add_id) = build_add_graph();
2244 let order = topo_order(&g, add_id);
2245 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2246 if dev == DeviceLocation::Cpu {
2247 vec![BackendId::Cpu, BackendId::Cuda]
2248 } else {
2249 vec![]
2250 }
2251 };
2252 let cpu_caps_val = cpu_caps();
2253 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
2254
2255 let sc = SizeClass::from_elem_count(3);
2258 let mut judge = HashMapJudge::new();
2259 judge.insert(
2260 OpKind::AddElementwise,
2261 DType::F32,
2262 sc,
2263 BackendId::Cpu,
2264 "",
2265 5_000,
2266 0,
2267 );
2268 judge.insert(
2269 OpKind::AddElementwise,
2270 DType::F32,
2271 sc,
2272 BackendId::Cuda,
2273 "",
2274 20,
2275 0,
2276 );
2277
2278 let opts = PlanOptions::new()
2279 .with_placements_for_device(&placements_fn)
2280 .with_capabilities_for(&caps_fn)
2281 .with_judge(&judge);
2282
2283 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2284 let alts = plan.alternatives(add_id).unwrap();
2285 assert_eq!(alts.len(), 2);
2286 assert_eq!(
2287 alts.winner().unwrap().backend,
2288 BackendId::Cuda,
2289 "Layer-2 measurement reverses Layer-1 verdict via compile_plan",
2290 );
2291 }
2292
2293 #[test]
2298 fn compile_plan_stamps_decision_context() {
2299 use fuel_ir::dispatch::SizeClass;
2300
2301 let mut table = KernelBindingTable::new();
2302 register_add_f32(
2303 &mut table,
2304 BackendId::Cpu,
2305 noop_kernel,
2306 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2307 );
2308 let (g, add_id) = build_add_graph();
2309 let order = topo_order(&g, add_id);
2310 let opts = PlanOptions::new().without_cost_population();
2311 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2312 let alts = plan.alternatives(add_id).expect("Add alternatives present");
2313 let ctx = alts.context().expect("context stamped");
2314 assert_eq!(ctx.op, OpKind::AddElementwise);
2315 assert_eq!(ctx.principal_dtype, DType::F32);
2316 assert_eq!(ctx.size_class, SizeClass::from_elem_count(3));
2318 }
2319
2320 #[test]
2329 fn compile_plan_retains_per_device_frontier_no_top_n() {
2330 let mut table = KernelBindingTable::new();
2331 register_add_f32(
2334 &mut table,
2335 BackendId::Cpu,
2336 noop_kernel,
2337 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2338 );
2339 register_add_f32(
2340 &mut table,
2341 BackendId::Cuda,
2342 noop_kernel_b,
2343 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2344 );
2345 register_add_f32(
2346 &mut table,
2347 BackendId::Vulkan,
2348 noop_kernel_c,
2349 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2350 );
2351
2352 let (g, add_id) = build_add_graph();
2353 let order = topo_order(&g, add_id);
2354 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2355 if dev == DeviceLocation::Cpu {
2356 vec![BackendId::Cpu, BackendId::Cuda, BackendId::Vulkan]
2357 } else {
2358 vec![]
2359 }
2360 };
2361 let opts = PlanOptions::new()
2362 .without_cost_population()
2363 .with_placements_for_device(&placements_fn);
2364 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2365 let alts = plan.alternatives(add_id).unwrap();
2366 assert_eq!(
2367 alts.len(),
2368 3,
2369 "all three co-located (device, backend) paths survive — no top-N truncation",
2370 );
2371 }
2372
2373 #[test]
2378 fn compile_plan_pinned_device_matches_stamped_plan() {
2379 let mut table = KernelBindingTable::new();
2380 register_add_f32(
2381 &mut table,
2382 BackendId::Cpu,
2383 noop_kernel,
2384 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2385 );
2386
2387 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2388 if dev == DeviceLocation::Cpu {
2389 vec![BackendId::Cpu]
2390 } else {
2391 vec![]
2392 }
2393 };
2394
2395 let (g_stamped, add_stamped) = build_add_graph();
2397 let order_stamped = topo_order(&g_stamped, add_stamped);
2398 let opts_stamped = PlanOptions::new()
2399 .without_cost_population()
2400 .with_placements_for_device(&placements_fn);
2401 let plan_stamped = compile_plan(&g_stamped, &order_stamped, &table, &opts_stamped)
2402 .expect("stamped compile");
2403
2404 let mut g = Graph::new();
2407 let lhs = g.push(Node {
2408 op: Op::Const,
2409 inputs: vec![],
2410 shape: Shape::from_dims(&[3]),
2411 dtype: DType::F32,
2412 });
2413 let rhs = g.push(Node {
2414 op: Op::Const,
2415 inputs: vec![],
2416 shape: Shape::from_dims(&[3]),
2417 dtype: DType::F32,
2418 });
2419 let add = g.push(Node {
2420 op: Op::Add,
2421 inputs: vec![lhs, rhs],
2422 shape: Shape::from_dims(&[3]),
2423 dtype: DType::F32,
2424 });
2425 let order = topo_order(&g, add);
2426 let opts = PlanOptions::new()
2427 .without_cost_population()
2428 .with_placements_for_device(&placements_fn)
2429 .with_pinned_device(DeviceLocation::Cpu);
2430 let plan = compile_plan(&g, &order, &table, &opts).expect("pinned compile");
2431
2432 let a = plan_stamped.alternatives(add_stamped).expect("stamped set");
2433 let b = plan.alternatives(add).expect("pinned set");
2434 assert_eq!(a.len(), b.len(), "same candidate count");
2435 for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
2436 assert_eq!(ca.backend, cb.backend);
2437 assert_eq!(ca.device, cb.device);
2438 assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
2439 assert_eq!(ca.kernel_source, cb.kernel_source);
2440 }
2441 assert_eq!(a.context(), b.context(), "same decision context");
2442 }
2443
2444 #[test]
2448 fn compile_plan_errors_without_any_device_context() {
2449 let mut table = KernelBindingTable::new();
2450 register_add_f32(
2451 &mut table,
2452 BackendId::Cpu,
2453 noop_kernel,
2454 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2455 );
2456 let mut g = Graph::new();
2457 let lhs = g.push(Node {
2458 op: Op::Const,
2459 inputs: vec![],
2460 shape: Shape::from_dims(&[3]),
2461 dtype: DType::F32,
2462 });
2463 let add = g.push(Node {
2464 op: Op::Add,
2465 inputs: vec![lhs, lhs],
2466 shape: Shape::from_dims(&[3]),
2467 dtype: DType::F32,
2468 });
2469 let order = topo_order(&g, add);
2470 let opts = PlanOptions::new().without_cost_population();
2471 let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
2472 let msg = format!("{err:?}");
2473 assert!(
2474 msg.contains("no device context"),
2475 "expected the device-context diagnostic, got: {msg}",
2476 );
2477 }
2478
2479 #[test]
2484 fn compile_plan_skips_op_copy_nodes() {
2485 let mut table = KernelBindingTable::new();
2486 register_add_f32(
2487 &mut table,
2488 BackendId::Cpu,
2489 noop_kernel,
2490 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2491 );
2492 table.register_full(
2495 OpKind::Copy,
2496 &[DType::F32, DType::F32],
2497 BackendId::Cpu,
2498 noop_kernel_b,
2499 KernelCaps::empty(),
2500 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2501 unknown_cost,
2502 );
2503
2504 let (mut g, add_id) = build_add_graph();
2505 let copy_id = g.push(Node {
2506 op: Op::Copy {
2507 target: DeviceLocation::Cpu,
2508 },
2509 inputs: vec![add_id],
2510 shape: Shape::from_dims(&[3]),
2511 dtype: DType::F32,
2512 });
2513 g.set_target_backend(copy_id, BackendId::Cpu);
2514 let order = topo_order(&g, copy_id);
2515 let opts = PlanOptions::new()
2516 .without_cost_population()
2517 .with_pinned_device(DeviceLocation::Cpu);
2518 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2519 assert!(plan.alternatives(add_id).is_some(), "Add is planned");
2520 assert!(
2521 plan.alternatives(copy_id).is_none(),
2522 "Op::Copy is residency-determined — no plan entry",
2523 );
2524 }
2525
2526 #[test]
2530 fn compile_plan_skips_op_move_nodes() {
2531 let mut table = KernelBindingTable::new();
2532 register_add_f32(
2533 &mut table,
2534 BackendId::Cpu,
2535 noop_kernel,
2536 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2537 );
2538 table.register_full(
2541 OpKind::Copy,
2542 &[DType::F32, DType::F32],
2543 BackendId::Cpu,
2544 noop_kernel_b,
2545 KernelCaps::empty(),
2546 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2547 unknown_cost,
2548 );
2549
2550 let (mut g, add_id) = build_add_graph();
2551 let move_id = g.push(Node {
2552 op: Op::Move {
2553 target: DeviceLocation::Cpu,
2554 },
2555 inputs: vec![add_id],
2556 shape: Shape::from_dims(&[3]),
2557 dtype: DType::F32,
2558 });
2559 g.set_target_backend(move_id, BackendId::Cpu);
2560 let order = topo_order(&g, move_id);
2561 let opts = PlanOptions::new()
2562 .without_cost_population()
2563 .with_pinned_device(DeviceLocation::Cpu);
2564 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2565 assert!(plan.alternatives(add_id).is_some(), "Add is planned");
2566 assert!(
2567 plan.alternatives(move_id).is_none(),
2568 "Op::Move is residency-determined — no plan entry",
2569 );
2570 }
2571
2572 #[test]
2576 fn compile_plan_fallback_admits_off_device_candidate() {
2577 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
2578 let mut table = KernelBindingTable::new();
2579 register_add_f32(
2582 &mut table,
2583 BackendId::Cpu,
2584 noop_kernel,
2585 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2586 );
2587
2588 let mut g = Graph::new();
2589 let lhs = g.push(Node {
2590 op: Op::Const,
2591 inputs: vec![],
2592 shape: Shape::from_dims(&[3]),
2593 dtype: DType::F32,
2594 });
2595 let add = g.push(Node {
2596 op: Op::Add,
2597 inputs: vec![lhs, lhs],
2598 shape: Shape::from_dims(&[3]),
2599 dtype: DType::F32,
2600 });
2601 let order = topo_order(&g, add);
2602
2603 let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
2604 if dev == cuda0 {
2605 vec![BackendId::Cuda]
2606 } else {
2607 vec![BackendId::Cpu]
2608 }
2609 };
2610 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
2611 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
2612 };
2613 let opts = PlanOptions::new()
2614 .without_cost_population()
2615 .with_pinned_device(cuda0)
2616 .with_placements_for_device(&placements_fn)
2617 .with_fallback_placements_for(&fallback_fn);
2618 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2619 let alts = plan.alternatives(add).expect("Add planned via fallback");
2620 assert_eq!(alts.len(), 1, "fallback set frozen to one winner");
2621 let w = alts.winner().unwrap();
2622 assert_eq!(w.backend, BackendId::Cpu);
2623 assert_eq!(w.device, DeviceLocation::Cpu, "off-device placement");
2624 }
2625
2626 #[test]
2632 fn compile_plan_fallback_not_consulted_when_pinned_has_impl() {
2633 let mut table = KernelBindingTable::new();
2634 register_add_f32(
2635 &mut table,
2636 BackendId::Cpu,
2637 noop_kernel,
2638 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2639 );
2640 register_add_f32(
2641 &mut table,
2642 BackendId::Cuda,
2643 noop_kernel_b,
2644 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2645 );
2646 let (g, add_id) = build_add_graph();
2647 let order = topo_order(&g, add_id);
2648 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2649 if dev == DeviceLocation::Cpu {
2650 vec![BackendId::Cpu]
2651 } else {
2652 vec![]
2653 }
2654 };
2655 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
2656 panic!(
2657 "fallback consulted although the pinned device has an \
2658 implementation — locality policy violated",
2659 );
2660 };
2661 let opts = PlanOptions::new()
2662 .without_cost_population()
2663 .with_pinned_device(DeviceLocation::Cpu)
2664 .with_placements_for_device(&placements_fn)
2665 .with_fallback_placements_for(&fallback_fn);
2666 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2667 let alts = plan.alternatives(add_id).unwrap();
2668 assert_eq!(alts.len(), 1);
2669 assert_eq!(alts.winner().unwrap().device, DeviceLocation::Cpu);
2670 }
2671
2672 #[test]
2675 fn compile_plan_fallback_denied_for_destructive_ops() {
2676 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
2677 let mut table = KernelBindingTable::new();
2678 table.register_full(
2680 OpKind::ReluInplace,
2681 &[DType::F32, DType::F32],
2682 BackendId::Cpu,
2683 noop_kernel,
2684 KernelCaps::empty(),
2685 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2686 unknown_cost,
2687 );
2688 let mut g = Graph::new();
2689 let lhs = g.push(Node {
2690 op: Op::Const,
2691 inputs: vec![],
2692 shape: Shape::from_dims(&[3]),
2693 dtype: DType::F32,
2694 });
2695 let relu = g.push(Node {
2696 op: Op::ReluInplace,
2697 inputs: vec![lhs],
2698 shape: Shape::from_dims(&[3]),
2699 dtype: DType::F32,
2700 });
2701 let order = topo_order(&g, relu);
2702 let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
2703 if dev == cuda0 {
2704 vec![BackendId::Cuda]
2705 } else {
2706 vec![BackendId::Cpu]
2707 }
2708 };
2709 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
2710 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
2711 };
2712 let opts = PlanOptions::new()
2713 .without_cost_population()
2714 .with_pinned_device(cuda0)
2715 .with_placements_for_device(&placements_fn)
2716 .with_fallback_placements_for(&fallback_fn);
2717 let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
2718 assert!(
2719 matches!(
2720 err,
2721 Error::NoBackendForOp {
2722 op: OpKind::ReluInplace,
2723 ..
2724 }
2725 ),
2726 "destructive op must NOT fall back off-device; got {err:?}",
2727 );
2728 }
2729
2730 #[test]
2735 fn compile_plan_fallback_freezes_to_single_winner() {
2736 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
2737 let vk0 = DeviceLocation::Vulkan { gpu_id: 0 };
2738 let mut table = KernelBindingTable::new();
2739 register_add_f32(
2740 &mut table,
2741 BackendId::Cpu,
2742 noop_kernel,
2743 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2744 );
2745 register_add_f32(
2746 &mut table,
2747 BackendId::Vulkan,
2748 noop_kernel_b,
2749 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2750 );
2751
2752 let mut g = Graph::new();
2753 let lhs = g.push(Node {
2754 op: Op::Const,
2755 inputs: vec![],
2756 shape: Shape::from_dims(&[3]),
2757 dtype: DType::F32,
2758 });
2759 let add = g.push(Node {
2760 op: Op::Add,
2761 inputs: vec![lhs, lhs],
2762 shape: Shape::from_dims(&[3]),
2763 dtype: DType::F32,
2764 });
2765 let order = topo_order(&g, add);
2766 let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
2767 if dev == cuda0 {
2768 vec![BackendId::Cuda]
2769 } else {
2770 vec![]
2771 }
2772 };
2773 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
2774 vec![
2775 (BackendId::Cpu, DeviceLocation::Cpu),
2776 (BackendId::Vulkan, vk0),
2777 ]
2778 };
2779 let opts = PlanOptions::new()
2780 .without_cost_population()
2781 .with_pinned_device(cuda0)
2782 .with_placements_for_device(&placements_fn)
2783 .with_fallback_placements_for(&fallback_fn);
2784 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2785 let alts = plan.alternatives(add).unwrap();
2786 assert_eq!(
2787 alts.len(),
2788 1,
2789 "fallback set must freeze to its single ranked winner",
2790 );
2791 }
2792
2793 struct FlatEstimator {
2796 latency_ns: u64,
2797 }
2798
2799 impl crate::ranker::TransferEstimator for FlatEstimator {
2800 fn estimate_transfer_ns(
2801 &self,
2802 src: DeviceLocation,
2803 dst: DeviceLocation,
2804 _bytes: u64,
2805 ) -> u64 {
2806 if src == dst { 0 } else { self.latency_ns }
2807 }
2808 }
2809
2810 #[test]
2816 fn compile_plan_prices_fallback_by_input_residency() {
2817 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
2818 let vk0 = DeviceLocation::Vulkan { gpu_id: 0 };
2819 let mut table = KernelBindingTable::new();
2820 register_add_f32(
2822 &mut table,
2823 BackendId::Cpu,
2824 noop_kernel,
2825 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2826 );
2827 register_add_f32(
2828 &mut table,
2829 BackendId::Vulkan,
2830 noop_kernel_b,
2831 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2832 );
2833
2834 let mut g = Graph::new();
2835 let lhs = g.push(Node {
2836 op: Op::Const,
2837 inputs: vec![],
2838 shape: Shape::from_dims(&[4]),
2839 dtype: DType::F32,
2840 });
2841 let add = g.push(Node {
2842 op: Op::Add,
2843 inputs: vec![lhs, lhs],
2844 shape: Shape::from_dims(&[4]),
2845 dtype: DType::F32,
2846 });
2847 let order = topo_order(&g, add);
2848
2849 let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
2850 if dev == cuda0 {
2851 vec![BackendId::Cuda]
2852 } else {
2853 vec![]
2854 }
2855 };
2856 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
2857 vec![
2858 (BackendId::Cpu, DeviceLocation::Cpu),
2859 (BackendId::Vulkan, vk0),
2860 ]
2861 };
2862 let residency_fn =
2866 move |id: NodeId| -> Option<DeviceLocation> { (id == lhs).then_some(vk0) };
2867 let est = FlatEstimator { latency_ns: 1_000 };
2868 let cpu_caps_val = cpu_caps();
2869 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
2870 let opts = PlanOptions::new()
2871 .with_pinned_device(cuda0)
2872 .with_placements_for_device(&placements_fn)
2873 .with_fallback_placements_for(&fallback_fn)
2874 .with_capabilities_for(&caps_fn)
2875 .with_transfer_estimator(&est)
2876 .with_input_residency(&residency_fn);
2877 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2878 let alts = plan.alternatives(add).expect("Add planned via fallback");
2879 assert_eq!(alts.len(), 1, "fallback narrows to a single winner");
2880 let w = alts.winner().unwrap();
2881 assert_eq!(
2882 w.device, vk0,
2883 "transfer pricing moves the fallback to the inputs' residency",
2884 );
2885 assert_eq!(w.inbound_transfer_ns, 0, "co-resident inputs price zero");
2886 }
2887
2888 fn cost_600(
2891 _: &[Shape],
2892 _: &[DType],
2893 _: &OpParams,
2894 _: &BackendCapabilities,
2895 ) -> crate::fused::CostEstimate {
2896 crate::fused::CostEstimate {
2897 flops: 600,
2898 bytes_moved: 0,
2899 kernel_overhead_ns: 0,
2900 }
2901 }
2902 fn cost_10(
2903 _: &[Shape],
2904 _: &[DType],
2905 _: &OpParams,
2906 _: &BackendCapabilities,
2907 ) -> crate::fused::CostEstimate {
2908 crate::fused::CostEstimate {
2909 flops: 10,
2910 bytes_moved: 0,
2911 kernel_overhead_ns: 0,
2912 }
2913 }
2914 fn cost_15_000(
2918 _: &[Shape],
2919 _: &[DType],
2920 _: &OpParams,
2921 _: &BackendCapabilities,
2922 ) -> crate::fused::CostEstimate {
2923 crate::fused::CostEstimate {
2924 flops: 15_000,
2925 bytes_moved: 0,
2926 kernel_overhead_ns: 0,
2927 }
2928 }
2929 fn cost_10_000_000(
2930 _: &[Shape],
2931 _: &[DType],
2932 _: &OpParams,
2933 _: &BackendCapabilities,
2934 ) -> crate::fused::CostEstimate {
2935 crate::fused::CostEstimate {
2936 flops: 10_000_000,
2937 bytes_moved: 0,
2938 kernel_overhead_ns: 0,
2939 }
2940 }
2941
2942 fn relaxed_plan_with_costs(
2948 cpu_cost: crate::kernel::CostFn,
2949 cuda_cost: crate::kernel::CostFn,
2950 estimator: &FlatEstimator,
2951 ) -> AlternativeSet {
2952 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
2953 let mut table = KernelBindingTable::new();
2954 table.register_full(
2955 OpKind::AddElementwise,
2956 &[DType::F32, DType::F32, DType::F32],
2957 BackendId::Cpu,
2958 noop_kernel,
2959 KernelCaps::empty(),
2960 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2961 cpu_cost,
2962 );
2963 table.register_full(
2964 OpKind::AddElementwise,
2965 &[DType::F32, DType::F32, DType::F32],
2966 BackendId::Cuda,
2967 noop_kernel_b,
2968 KernelCaps::empty(),
2969 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
2970 cuda_cost,
2971 );
2972
2973 let (g, add_id) = build_add_graph();
2974 let order = topo_order(&g, add_id);
2975 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
2976 if dev == DeviceLocation::Cpu {
2977 vec![BackendId::Cpu]
2978 } else {
2979 vec![]
2980 }
2981 };
2982 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
2983 vec![(BackendId::Cuda, cuda0)]
2984 };
2985 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
2988 let cpu_caps_val = cpu_caps();
2989 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
2990 let opts = PlanOptions::new()
2991 .with_pinned_device(DeviceLocation::Cpu)
2992 .with_placements_for_device(&placements_fn)
2993 .with_fallback_placements_for(&fallback_fn)
2994 .with_capabilities_for(&caps_fn)
2995 .with_transfer_estimator(estimator)
2996 .with_input_residency(&residency_fn);
2997 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
2998 plan.alternatives(add_id).expect("Add planned").clone()
2999 }
3000
3001 #[test]
3005 fn compile_plan_relaxed_tiny_op_stays_local_when_transfer_dominates() {
3006 let est = FlatEstimator {
3009 latency_ns: 1_000_000,
3010 };
3011 let set = relaxed_plan_with_costs(cost_600, cost_10, &est);
3012 let w = set.winner().unwrap();
3013 assert_eq!(
3014 w.device,
3015 DeviceLocation::Cpu,
3016 "transfer dominates → stay local"
3017 );
3018 assert_eq!(w.inbound_transfer_ns, 0, "local winner pays no transfer");
3019 assert!(
3020 set.alternatives()
3021 .iter()
3022 .all(|c| c.device == DeviceLocation::Cpu),
3023 "off-device siblings pruned after rank",
3024 );
3025 }
3026
3027 #[test]
3060 fn compile_plan_uncapped_gpu_no_longer_wins_by_being_free() {
3061 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3062
3063 fn build_table() -> KernelBindingTable {
3070 let mut table = KernelBindingTable::new();
3071 table.register_full(
3072 OpKind::AddElementwise,
3073 &[DType::F32, DType::F32, DType::F32],
3074 BackendId::Cpu,
3075 noop_kernel,
3076 KernelCaps::empty(),
3077 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3078 cost_600,
3079 );
3080 table.register_full(
3081 OpKind::AddElementwise,
3082 &[DType::F32, DType::F32, DType::F32],
3083 BackendId::Cuda,
3084 noop_kernel_b,
3085 KernelCaps::empty(),
3086 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3087 cost_15_000,
3088 );
3089 table
3090 }
3091
3092 fn winner_device_under(
3100 table: &KernelBindingTable,
3101 cpu_caps_val: &'static BackendCapabilities,
3102 gpu_caps: Option<&'static BackendCapabilities>,
3103 cuda0: DeviceLocation,
3104 ) -> DeviceLocation {
3105 let (g, add_id) = build_add_graph();
3106 let order = topo_order(&g, add_id);
3107 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3108 if dev == DeviceLocation::Cpu {
3109 vec![BackendId::Cpu]
3110 } else {
3111 vec![]
3112 }
3113 };
3114 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3115 vec![(BackendId::Cuda, cuda0)]
3116 };
3117 let residency_fn =
3121 |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3122 let est = FlatEstimator { latency_ns: 100 };
3125 let caps_fn = move |b: BackendId| -> Option<&BackendCapabilities> {
3126 match b {
3127 BackendId::Cpu => Some(cpu_caps_val),
3128 BackendId::Cuda => gpu_caps,
3129 _ => None,
3130 }
3131 };
3132 let opts = PlanOptions::new()
3133 .with_pinned_device(DeviceLocation::Cpu)
3134 .with_placements_for_device(&placements_fn)
3135 .with_fallback_placements_for(&fallback_fn)
3136 .with_capabilities_for(&caps_fn)
3137 .with_transfer_estimator(&est)
3138 .with_input_residency(&residency_fn);
3139 let plan = compile_plan(&g, &order, table, &opts).expect("compile");
3140 plan.alternatives(add_id).unwrap().winner().unwrap().device
3141 }
3142
3143 let table = build_table();
3144 let cpu_caps_val: &'static BackendCapabilities = Box::leak(Box::new(cpu_caps()));
3148
3149 let winner_bug = winner_device_under(&table, cpu_caps_val, None, cuda0);
3153 assert_eq!(
3154 winner_bug, cuda0,
3155 "PRE-FIX: an uncapped CUDA candidate is skipped by \
3156 compute_static_costs, prices at 0 + 200 transfer, and \
3157 spuriously beats the CPU's honest 600 — this IS the bug",
3158 );
3159
3160 let cuda_caps_val: &'static BackendCapabilities = Box::leak(Box::new(
3163 crate::dispatch::derive_backend_caps(BackendId::Cuda, cuda0, &table),
3164 ));
3165 assert!(
3168 cuda_caps_val
3169 .op_dtype_support
3170 .contains(&(OpKind::AddElementwise, DType::F32)),
3171 "derive_backend_caps must mirror the binding table's \
3172 registered (op, output-dtype) pairs for the backend",
3173 );
3174 assert_eq!(cuda_caps_val.backend_id, BackendId::Cuda);
3175 assert_eq!(cuda_caps_val.device_location, cuda0);
3176
3177 let winner_fixed = winner_device_under(&table, cpu_caps_val, Some(cuda_caps_val), cuda0);
3178 assert_eq!(
3179 winner_fixed,
3180 DeviceLocation::Cpu,
3181 "POST-FIX: with GPU caps registered the CUDA candidate is \
3182 costed honestly (15000 FLOPs ÷ 30/ns = 500 ns compute + 200 \
3183 transfer = 700 > CPU 600), so the CPU-pinned realize stays \
3184 on CPU — no spill onto the unseeded GPU",
3185 );
3186 }
3187
3188 #[test]
3205 fn compile_plan_places_large_compute_on_gpu_small_on_cpu_by_throughput() {
3206 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3207
3208 fn winner(cost: crate::kernel::CostFn, cuda0: DeviceLocation) -> DeviceLocation {
3209 let mut table = KernelBindingTable::new();
3210 table.register_full(
3211 OpKind::AddElementwise,
3212 &[DType::F32, DType::F32, DType::F32],
3213 BackendId::Cpu,
3214 noop_kernel,
3215 KernelCaps::empty(),
3216 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3217 cost,
3218 );
3219 table.register_full(
3220 OpKind::AddElementwise,
3221 &[DType::F32, DType::F32, DType::F32],
3222 BackendId::Cuda,
3223 noop_kernel_b,
3224 KernelCaps::empty(),
3225 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3226 cost,
3227 );
3228
3229 let (g, add_id) = build_add_graph();
3230 let order = topo_order(&g, add_id);
3231 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3232 if dev == DeviceLocation::Cpu {
3233 vec![BackendId::Cpu]
3234 } else {
3235 vec![]
3236 }
3237 };
3238 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3239 vec![(BackendId::Cuda, cuda0)]
3240 };
3241 let residency_fn =
3242 |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3243 let est = FlatEstimator { latency_ns: 100 };
3246 let cpu_caps_val = cpu_caps();
3247 let cuda_caps_val =
3248 crate::dispatch::derive_backend_caps(BackendId::Cuda, cuda0, &table);
3249 let caps_fn = |b: BackendId| -> Option<&BackendCapabilities> {
3250 match b {
3251 BackendId::Cpu => Some(&cpu_caps_val),
3252 BackendId::Cuda => Some(&cuda_caps_val),
3253 _ => None,
3254 }
3255 };
3256 let opts = PlanOptions::new()
3257 .with_pinned_device(DeviceLocation::Cpu)
3258 .with_placements_for_device(&placements_fn)
3259 .with_fallback_placements_for(&fallback_fn)
3260 .with_capabilities_for(&caps_fn)
3261 .with_transfer_estimator(&est)
3262 .with_input_residency(&residency_fn);
3263 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
3264 plan.alternatives(add_id).unwrap().winner().unwrap().device
3265 }
3266
3267 assert_eq!(
3268 winner(cost_100, cuda0),
3269 DeviceLocation::Cpu,
3270 "SMALL op: the GPU's throughput win on 100 FLOPs (~3 ns) \
3271 can't overcome the 200 ns inbound transfer → stay on CPU",
3272 );
3273 assert_eq!(
3274 winner(cost_10_000_000, cuda0),
3275 cuda0,
3276 "LARGE op: the GPU's throughput win on 10 M FLOPs (~333 k ns) \
3277 dwarfs the 200 ns transfer → move to CUDA. RED under the \
3278 fixed 1-FLOP/ns prior (equal compute-ns → the transfer \
3279 always loses the GPU the placement)",
3280 );
3281 }
3282
3283 #[test]
3294 fn compile_plan_without_cost_placement_keeps_large_op_on_cpu() {
3295 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3296 let mut table = KernelBindingTable::new();
3297 table.register_full(
3298 OpKind::AddElementwise,
3299 &[DType::F32, DType::F32, DType::F32],
3300 BackendId::Cpu,
3301 noop_kernel,
3302 KernelCaps::empty(),
3303 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3304 cost_10_000_000,
3305 );
3306 table.register_full(
3307 OpKind::AddElementwise,
3308 &[DType::F32, DType::F32, DType::F32],
3309 BackendId::Cuda,
3310 noop_kernel_b,
3311 KernelCaps::empty(),
3312 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3313 cost_10_000_000,
3314 );
3315
3316 let (g, add_id) = build_add_graph();
3317 let order = topo_order(&g, add_id);
3318 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3319 if dev == DeviceLocation::Cpu {
3320 vec![BackendId::Cpu]
3321 } else {
3322 vec![]
3323 }
3324 };
3325 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3326 vec![(BackendId::Cuda, cuda0)]
3327 };
3328 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3329 let est = FlatEstimator { latency_ns: 100 };
3330 let cpu_caps_val = cpu_caps();
3331 let cuda_caps_val = crate::dispatch::derive_backend_caps(BackendId::Cuda, cuda0, &table);
3332 let caps_fn = |b: BackendId| -> Option<&BackendCapabilities> {
3333 match b {
3334 BackendId::Cpu => Some(&cpu_caps_val),
3335 BackendId::Cuda => Some(&cuda_caps_val),
3336 _ => None,
3337 }
3338 };
3339 let opts = PlanOptions::new()
3340 .with_pinned_device(DeviceLocation::Cpu)
3341 .with_placements_for_device(&placements_fn)
3342 .with_fallback_placements_for(&fallback_fn)
3343 .with_capabilities_for(&caps_fn)
3344 .with_transfer_estimator(&est)
3345 .with_input_residency(&residency_fn)
3346 .without_cost_placement();
3347 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
3348 let set = plan.alternatives(add_id).unwrap();
3349 assert_eq!(
3350 set.winner().unwrap().device,
3351 DeviceLocation::Cpu,
3352 "without_cost_placement: the large GPU-favorable op stays on the \
3353 pinned CPU device (no relaxed cross-device enumeration)",
3354 );
3355 assert!(
3356 set.alternatives()
3357 .iter()
3358 .all(|c| c.device == DeviceLocation::Cpu),
3359 "no off-device candidate enumerated at all",
3360 );
3361 }
3362
3363 #[test]
3368 fn compile_plan_relaxed_huge_op_moves_when_kernel_gap_dominates() {
3369 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3370 let est = FlatEstimator { latency_ns: 1_000 };
3372 let set = relaxed_plan_with_costs(cost_10_000_000, cost_10, &est);
3373 let w = set.winner().unwrap();
3374 assert_eq!(w.device, cuda0, "kernel gap dominates → move");
3375 assert_eq!(
3376 w.inbound_transfer_ns, 2_000,
3377 "the move's two input crossings are priced on the winner",
3378 );
3379 assert!(
3380 set.alternatives().iter().all(|c| c.device == cuda0),
3381 "set pruned to the winner's device",
3382 );
3383 }
3384
3385 #[test]
3390 fn compile_plan_relaxed_zero_signal_ties_stay_local() {
3391 let est = FlatEstimator { latency_ns: 0 };
3392 let set = relaxed_plan_with_costs(unknown_cost, unknown_cost, &est);
3393 assert_eq!(
3394 set.winner().unwrap().device,
3395 DeviceLocation::Cpu,
3396 "ties break toward the decision device",
3397 );
3398 }
3399
3400 #[test]
3404 fn compile_plan_relaxed_single_device_plan_unchanged() {
3405 let mut table = KernelBindingTable::new();
3406 register_add_f32(
3407 &mut table,
3408 BackendId::Cpu,
3409 noop_kernel,
3410 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3411 );
3412 let (g, add_id) = build_add_graph();
3413 let order = topo_order(&g, add_id);
3414 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3415 if dev == DeviceLocation::Cpu {
3416 vec![BackendId::Cpu]
3417 } else {
3418 vec![]
3419 }
3420 };
3421 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> { Vec::new() };
3423 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3424 let cpu_caps_val = cpu_caps();
3425 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
3426
3427 let base_opts = PlanOptions::new()
3428 .with_pinned_device(DeviceLocation::Cpu)
3429 .with_placements_for_device(&placements_fn)
3430 .with_fallback_placements_for(&fallback_fn)
3431 .with_capabilities_for(&caps_fn);
3432 let base = compile_plan(&g, &order, &table, &base_opts).expect("base");
3433
3434 let est = FlatEstimator {
3435 latency_ns: 1_000_000,
3436 };
3437 let wired_opts = PlanOptions::new()
3438 .with_pinned_device(DeviceLocation::Cpu)
3439 .with_placements_for_device(&placements_fn)
3440 .with_fallback_placements_for(&fallback_fn)
3441 .with_capabilities_for(&caps_fn)
3442 .with_transfer_estimator(&est)
3443 .with_input_residency(&residency_fn);
3444 let wired = compile_plan(&g, &order, &table, &wired_opts).expect("wired");
3445
3446 let a = base.alternatives(add_id).unwrap();
3447 let b = wired.alternatives(add_id).unwrap();
3448 assert_eq!(a.len(), b.len(), "same candidate count");
3449 for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
3450 assert_eq!(ca.backend, cb.backend);
3451 assert_eq!(ca.device, cb.device);
3452 assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
3453 assert_eq!(ca.static_cost, cb.static_cost, "same Layer-1 cost");
3454 assert_eq!(cb.inbound_transfer_ns, 0, "no transfer term fires");
3455 }
3456 assert_eq!(a.context(), b.context(), "same decision context");
3457 }
3458
3459 #[test]
3464 fn compile_plan_relaxed_respects_explicit_placement_pin() {
3465 let mut table = KernelBindingTable::new();
3466 register_add_f32(
3467 &mut table,
3468 BackendId::Cpu,
3469 noop_kernel,
3470 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3471 );
3472 register_add_f32(
3473 &mut table,
3474 BackendId::Cuda,
3475 noop_kernel_b,
3476 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3477 );
3478 let (mut g, add_id) = build_add_graph();
3479 g.set_placement(add_id, DeviceLocation::Cpu);
3481 let order = topo_order(&g, add_id);
3482 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3483 if dev == DeviceLocation::Cpu {
3484 vec![BackendId::Cpu]
3485 } else {
3486 vec![]
3487 }
3488 };
3489 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3490 panic!(
3491 "fallback consulted for a hard-pinned node — explicit \
3492 placements must not enter the priced relax",
3493 );
3494 };
3495 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3496 let est = FlatEstimator { latency_ns: 0 };
3497 let cpu_caps_val = cpu_caps();
3498 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
3499 let opts = PlanOptions::new()
3500 .with_pinned_device(DeviceLocation::Cpu)
3501 .with_placements_for_device(&placements_fn)
3502 .with_fallback_placements_for(&fallback_fn)
3503 .with_capabilities_for(&caps_fn)
3504 .with_transfer_estimator(&est)
3505 .with_input_residency(&residency_fn);
3506 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
3507 assert_eq!(
3508 plan.alternatives(add_id).unwrap().winner().unwrap().device,
3509 DeviceLocation::Cpu,
3510 );
3511 }
3512
3513 #[test]
3516 fn compile_plan_relaxed_destructive_never_moves() {
3517 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3518 let mut table = KernelBindingTable::new();
3519 table.register_full(
3520 OpKind::ReluInplace,
3521 &[DType::F32, DType::F32],
3522 BackendId::Cpu,
3523 noop_kernel,
3524 KernelCaps::empty(),
3525 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3526 cost_600,
3527 );
3528 table.register_full(
3529 OpKind::ReluInplace,
3530 &[DType::F32, DType::F32],
3531 BackendId::Cuda,
3532 noop_kernel_b,
3533 KernelCaps::empty(),
3534 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3535 cost_10,
3536 );
3537 let mut g = Graph::new();
3538 let lhs = g.push(Node {
3539 op: Op::Const,
3540 inputs: vec![],
3541 shape: Shape::from_dims(&[3]),
3542 dtype: DType::F32,
3543 });
3544 let relu = g.push(Node {
3545 op: Op::ReluInplace,
3546 inputs: vec![lhs],
3547 shape: Shape::from_dims(&[3]),
3548 dtype: DType::F32,
3549 });
3550 let order = topo_order(&g, relu);
3551 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3552 if dev == DeviceLocation::Cpu {
3553 vec![BackendId::Cpu]
3554 } else {
3555 vec![]
3556 }
3557 };
3558 let fallback_fn = move |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3559 vec![(BackendId::Cuda, cuda0)]
3560 };
3561 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3562 let est = FlatEstimator { latency_ns: 0 };
3563 let cpu_caps_val = cpu_caps();
3564 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
3565 let opts = PlanOptions::new()
3566 .with_pinned_device(DeviceLocation::Cpu)
3567 .with_placements_for_device(&placements_fn)
3568 .with_fallback_placements_for(&fallback_fn)
3569 .with_capabilities_for(&caps_fn)
3570 .with_transfer_estimator(&est)
3571 .with_input_residency(&residency_fn);
3572 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
3573 assert_eq!(
3574 plan.alternatives(relu).unwrap().winner().unwrap().device,
3575 DeviceLocation::Cpu,
3576 "destructive op stays on the device owning its mutation target",
3577 );
3578 }
3579
3580 #[test]
3585 fn compile_plan_filter_rejected_pinned_device_falls_back() {
3586 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3587 let mut table = KernelBindingTable::new();
3588 register_add_f32(
3591 &mut table,
3592 BackendId::Cuda,
3593 noop_kernel_b,
3594 PrecisionGuarantee {
3595 bit_stable_on_same_hardware: false,
3596 ..PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU
3597 },
3598 );
3599 register_add_f32(
3600 &mut table,
3601 BackendId::Cpu,
3602 noop_kernel,
3603 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3604 );
3605 let mut g = Graph::new();
3606 let lhs = g.push(Node {
3607 op: Op::Const,
3608 inputs: vec![],
3609 shape: Shape::from_dims(&[3]),
3610 dtype: DType::F32,
3611 });
3612 let add = g.push(Node {
3613 op: Op::Add,
3614 inputs: vec![lhs, lhs],
3615 shape: Shape::from_dims(&[3]),
3616 dtype: DType::F32,
3617 });
3618 let order = topo_order(&g, add);
3619 let placements_fn = move |dev: DeviceLocation| -> Vec<BackendId> {
3620 if dev == cuda0 {
3621 vec![BackendId::Cuda]
3622 } else {
3623 vec![]
3624 }
3625 };
3626 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3627 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
3628 };
3629 let opts = PlanOptions::new()
3630 .without_cost_population()
3631 .with_pinned_device(cuda0)
3632 .with_placements_for_device(&placements_fn)
3633 .with_fallback_placements_for(&fallback_fn)
3634 .with_precision_requirement(PrecisionRequirement::BIT_STABLE);
3635 let plan = compile_plan(&g, &order, &table, &opts)
3636 .expect("filter-rejected pin falls back instead of erroring");
3637 let alts = plan.alternatives(add).expect("Add planned via fallback");
3638 assert_eq!(alts.len(), 1, "fallback set frozen to one winner");
3639 let w = alts.winner().unwrap();
3640 assert_eq!(w.backend, BackendId::Cpu);
3641 assert_eq!(w.device, DeviceLocation::Cpu);
3642 assert!(
3643 w.precision.bit_stable_on_same_hardware,
3644 "the admissible off-device sibling won",
3645 );
3646 }
3647
3648 #[test]
3653 fn compile_plan_without_estimator_prices_nothing() {
3654 let mut table = KernelBindingTable::new();
3655 register_add_f32(
3656 &mut table,
3657 BackendId::Cpu,
3658 noop_kernel,
3659 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3660 );
3661 let (g, add_id) = build_add_graph();
3662 let order = topo_order(&g, add_id);
3663 let cpu_caps_val = cpu_caps();
3664 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
3665 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
3666
3667 let base_opts = PlanOptions::new().with_capabilities_for(&caps_fn);
3668 let base = compile_plan(&g, &order, &table, &base_opts).expect("base");
3669
3670 let with_residency = PlanOptions::new()
3671 .with_capabilities_for(&caps_fn)
3672 .with_input_residency(&residency_fn);
3673 let plan = compile_plan(&g, &order, &table, &with_residency).expect("plan");
3674
3675 let a = base.alternatives(add_id).unwrap();
3676 let b = plan.alternatives(add_id).unwrap();
3677 assert_eq!(a.len(), b.len());
3678 for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
3679 assert_eq!(ca.backend, cb.backend);
3680 assert_eq!(ca.device, cb.device);
3681 assert_eq!(ca.kernel as usize, cb.kernel as usize);
3682 assert_eq!(ca.inbound_transfer_ns, 0);
3683 assert_eq!(cb.inbound_transfer_ns, 0);
3684 }
3685 }
3686
3687 #[test]
3693 fn compile_plan_require_transfer_pricing_errors_without_estimator() {
3694 let mut table = KernelBindingTable::new();
3695 register_add_f32(
3696 &mut table,
3697 BackendId::Cpu,
3698 noop_kernel,
3699 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3700 );
3701 let (g, add_id) = build_add_graph();
3702 let order = topo_order(&g, add_id);
3703 let cpu_caps_val = cpu_caps();
3704 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
3705
3706 let opts = PlanOptions::new()
3709 .with_capabilities_for(&caps_fn)
3710 .require_transfer_pricing();
3711 let err = compile_plan(&g, &order, &table, &opts).unwrap_err();
3712 let msg = format!("{err:?}");
3713 assert!(
3714 msg.contains("transfer") && (msg.contains("estimator") || msg.contains("pricing")),
3715 "error should name the missing transfer estimator / pricing; got: {msg}",
3716 );
3717 }
3718
3719 #[test]
3722 fn compile_plan_require_transfer_pricing_ok_with_estimator() {
3723 let mut table = KernelBindingTable::new();
3724 register_add_f32(
3725 &mut table,
3726 BackendId::Cpu,
3727 noop_kernel,
3728 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3729 );
3730 let (g, add_id) = build_add_graph();
3731 let order = topo_order(&g, add_id);
3732 let cpu_caps_val = cpu_caps();
3733 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
3734 let est = FlatEstimator { latency_ns: 1000 };
3735
3736 let opts = PlanOptions::new()
3737 .with_capabilities_for(&caps_fn)
3738 .with_transfer_estimator(&est)
3739 .require_transfer_pricing();
3740 let plan = compile_plan(&g, &order, &table, &opts)
3741 .expect("require_transfer_pricing WITH an estimator must plan fine");
3742 assert!(plan.alternatives(add_id).is_some());
3743 }
3744
3745 #[test]
3746 fn compile_plan_passes_input_layouts_to_filter_chain() {
3747 let mut table = KernelBindingTable::new();
3751 table.register_full(
3752 OpKind::AddElementwise,
3753 &[DType::F32, DType::F32, DType::F32],
3754 BackendId::Cpu,
3755 noop_kernel,
3756 KernelCaps::empty(),
3757 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3758 unknown_cost,
3759 );
3760 table.register_full(
3761 OpKind::AddElementwise,
3762 &[DType::F32, DType::F32, DType::F32],
3763 BackendId::Cpu,
3764 noop_kernel_b,
3765 KernelCaps::strided_input(),
3766 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3767 unknown_cost,
3768 );
3769
3770 let mut g = Graph::new();
3772 let lhs = g.push(Node {
3773 op: Op::Const,
3774 inputs: vec![],
3775 shape: Shape::from_dims(&[6]),
3776 dtype: DType::F32,
3777 });
3778 let strides: StrideVec = smallvec![2isize];
3781 g.set_layout(lhs, Layout::new(Shape::from_dims(&[6]), strides, 0));
3782 let rhs = g.push(Node {
3783 op: Op::Const,
3784 inputs: vec![],
3785 shape: Shape::from_dims(&[6]),
3786 dtype: DType::F32,
3787 });
3788 let add = g.push(Node {
3789 op: Op::Add,
3790 inputs: vec![lhs, rhs],
3791 shape: Shape::from_dims(&[6]),
3792 dtype: DType::F32,
3793 });
3794 g.set_target_backend(add, BackendId::Cpu);
3795 let order = topo_order(&g, add);
3796 let opts = PlanOptions::new().without_cost_population();
3797 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
3798 let alts = plan.alternatives(add).unwrap();
3799 assert_eq!(alts.len(), 1);
3802 assert!(alts.winner().unwrap().caps.strided_input);
3803 }
3804
3805 struct BytesEstimator {
3812 latency_ns: u64,
3813 ns_per_byte: u64,
3814 }
3815
3816 impl crate::ranker::TransferEstimator for BytesEstimator {
3817 fn estimate_transfer_ns(
3818 &self,
3819 src: DeviceLocation,
3820 dst: DeviceLocation,
3821 bytes: u64,
3822 ) -> u64 {
3823 if src == dst {
3824 return 0;
3825 }
3826 self.latency_ns
3827 .saturating_add(bytes.saturating_mul(self.ns_per_byte))
3828 }
3829 }
3830
3831 fn cost_50(
3832 _: &[Shape],
3833 _: &[DType],
3834 _: &OpParams,
3835 _: &BackendCapabilities,
3836 ) -> crate::fused::CostEstimate {
3837 crate::fused::CostEstimate {
3838 flops: 50,
3839 bytes_moved: 0,
3840 kernel_overhead_ns: 0,
3841 }
3842 }
3843 fn cost_100(
3844 _: &[Shape],
3845 _: &[DType],
3846 _: &OpParams,
3847 _: &BackendCapabilities,
3848 ) -> crate::fused::CostEstimate {
3849 crate::fused::CostEstimate {
3850 flops: 100,
3851 bytes_moved: 0,
3852 kernel_overhead_ns: 0,
3853 }
3854 }
3855 fn cost_300(
3856 _: &[Shape],
3857 _: &[DType],
3858 _: &OpParams,
3859 _: &BackendCapabilities,
3860 ) -> crate::fused::CostEstimate {
3861 crate::fused::CostEstimate {
3862 flops: 300,
3863 bytes_moved: 0,
3864 kernel_overhead_ns: 0,
3865 }
3866 }
3867 fn cost_800(
3868 _: &[Shape],
3869 _: &[DType],
3870 _: &OpParams,
3871 _: &BackendCapabilities,
3872 ) -> crate::fused::CostEstimate {
3873 crate::fused::CostEstimate {
3874 flops: 800,
3875 bytes_moved: 0,
3876 kernel_overhead_ns: 0,
3877 }
3878 }
3879 fn cost_5000(
3880 _: &[Shape],
3881 _: &[DType],
3882 _: &OpParams,
3883 _: &BackendCapabilities,
3884 ) -> crate::fused::CostEstimate {
3885 crate::fused::CostEstimate {
3886 flops: 5000,
3887 bytes_moved: 0,
3888 kernel_overhead_ns: 0,
3889 }
3890 }
3891 fn cost_100_000(
3892 _: &[Shape],
3893 _: &[DType],
3894 _: &OpParams,
3895 _: &BackendCapabilities,
3896 ) -> crate::fused::CostEstimate {
3897 crate::fused::CostEstimate {
3898 flops: 100_000,
3899 bytes_moved: 0,
3900 kernel_overhead_ns: 0,
3901 }
3902 }
3903
3904 fn register_unary_f32(
3906 table: &mut KernelBindingTable,
3907 op: OpKind,
3908 backend: BackendId,
3909 kernel: crate::kernel::KernelRef,
3910 cost: crate::kernel::CostFn,
3911 ) {
3912 table.register_full(
3913 op,
3914 &[DType::F32, DType::F32],
3915 backend,
3916 kernel,
3917 KernelCaps::empty(),
3918 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
3919 cost,
3920 );
3921 }
3922
3923 fn push_f32(g: &mut Graph, op: Op, inputs: Vec<NodeId>, dims: &[usize]) -> NodeId {
3924 g.push(Node {
3925 op,
3926 inputs,
3927 shape: Shape::from_dims(dims),
3928 dtype: DType::F32,
3929 })
3930 }
3931
3932 #[test]
3947 fn dp_mid_sequence_gpu_segment_beats_all_cpu_and_all_gpu() {
3948 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
3949 let mut table = KernelBindingTable::new();
3950 register_unary_f32(
3951 &mut table,
3952 OpKind::SqrElementwise,
3953 BackendId::Cpu,
3954 noop_kernel,
3955 cost_100,
3956 );
3957 register_unary_f32(
3958 &mut table,
3959 OpKind::SqrElementwise,
3960 BackendId::Cuda,
3961 noop_kernel_b,
3962 cost_5000,
3963 );
3964 register_unary_f32(
3965 &mut table,
3966 OpKind::NegElementwise,
3967 BackendId::Cpu,
3968 noop_kernel,
3969 cost_800,
3970 );
3971 register_unary_f32(
3972 &mut table,
3973 OpKind::NegElementwise,
3974 BackendId::Cuda,
3975 noop_kernel_b,
3976 cost_100,
3977 );
3978
3979 let mut g = Graph::new();
3980 let c = push_f32(&mut g, Op::Const, vec![], &[4]);
3981 let n1 = push_f32(&mut g, Op::Sqr, vec![c], &[4]);
3982 let n2 = push_f32(&mut g, Op::Neg, vec![n1], &[4]);
3983 let n3 = push_f32(&mut g, Op::Neg, vec![n2], &[4]);
3984 let n4 = push_f32(&mut g, Op::Neg, vec![n3], &[4]);
3985 let n5 = push_f32(&mut g, Op::Sqr, vec![n4], &[4]);
3986 let order = topo_order(&g, n5);
3987
3988 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
3989 if dev == DeviceLocation::Cpu {
3990 vec![BackendId::Cpu]
3991 } else {
3992 vec![BackendId::Cuda]
3993 }
3994 };
3995 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
3996 if dev == DeviceLocation::Cpu {
3997 vec![(BackendId::Cuda, cuda0)]
3998 } else {
3999 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4000 }
4001 };
4002 let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
4003 (id == c).then_some(DeviceLocation::Cpu)
4004 };
4005 let est = FlatEstimator { latency_ns: 1_000 };
4006 let cpu_caps_val = cpu_caps();
4007 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4008 let opts = PlanOptions::new()
4009 .with_pinned_device(DeviceLocation::Cpu)
4010 .with_placements_for_device(&placements_fn)
4011 .with_fallback_placements_for(&fallback_fn)
4012 .with_capabilities_for(&caps_fn)
4013 .with_transfer_estimator(&est)
4014 .with_input_residency(&residency_fn);
4015 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4016
4017 let dev_of = |id: NodeId| plan.alternatives(id).unwrap().winner().unwrap().device;
4018 assert_eq!(
4019 dev_of(n1),
4020 DeviceLocation::Cpu,
4021 "GPU-hostile entry stays CPU"
4022 );
4023 assert_eq!(dev_of(n2), cuda0, "segment start migrates");
4024 assert_eq!(dev_of(n3), cuda0, "segment interior stays GPU");
4025 assert_eq!(dev_of(n4), cuda0, "segment end stays GPU");
4026 assert_eq!(dev_of(n5), DeviceLocation::Cpu, "GPU-hostile exit returns");
4027
4028 let inbound = |id: NodeId| {
4032 plan.alternatives(id)
4033 .unwrap()
4034 .winner()
4035 .unwrap()
4036 .inbound_transfer_ns
4037 };
4038 assert_eq!(
4039 inbound(n2),
4040 1_000,
4041 "entry crossing priced on the first GPU op"
4042 );
4043 assert_eq!(inbound(n3), 0);
4044 assert_eq!(inbound(n4), 0);
4045 assert_eq!(inbound(n5), 1_000, "exit crossing priced on the return op");
4046
4047 for id in [n1, n2, n3, n4, n5] {
4050 let d = dev_of(id);
4051 assert!(
4052 plan.alternatives(id)
4053 .unwrap()
4054 .alternatives()
4055 .iter()
4056 .all(|c| c.device == d),
4057 "set pruned to the DP-committed device",
4058 );
4059 }
4060 }
4061
4062 #[test]
4078 fn dp_fused_op_loses_when_stranding_exceeds_local_win() {
4079 use fuel_graph::registry::{FusedOpParams, FusedOps};
4080
4081 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
4082 let mut table = KernelBindingTable::new();
4083 register_unary_f32(
4084 &mut table,
4085 OpKind::NegElementwise,
4086 BackendId::Cpu,
4087 noop_kernel,
4088 cost_100_000,
4089 );
4090 register_unary_f32(
4091 &mut table,
4092 OpKind::NegElementwise,
4093 BackendId::Cuda,
4094 noop_kernel_b,
4095 cost_100,
4096 );
4097 register_unary_f32(
4098 &mut table,
4099 OpKind::SoftmaxLastDim,
4100 BackendId::Cpu,
4101 noop_kernel,
4102 cost_300,
4103 );
4104 register_unary_f32(
4105 &mut table,
4106 OpKind::SoftmaxLastDim,
4107 BackendId::Cuda,
4108 noop_kernel_b,
4109 cost_100,
4110 );
4111 register_unary_f32(
4113 &mut table,
4114 OpKind::SqrElementwise,
4115 BackendId::Cpu,
4116 noop_kernel,
4117 cost_100,
4118 );
4119
4120 let mut g = Graph::new();
4121 let c = push_f32(&mut g, Op::Const, vec![], &[1]);
4122 let n1 = push_f32(&mut g, Op::Neg, vec![c], &[1]);
4123 let f = push_f32(
4124 &mut g,
4125 Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
4126 vec![n1],
4127 &[256],
4128 );
4129 let n2 = push_f32(&mut g, Op::Sqr, vec![f], &[256]);
4130 let order = topo_order(&g, n2);
4131
4132 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4133 if dev == DeviceLocation::Cpu {
4134 vec![BackendId::Cpu]
4135 } else {
4136 vec![BackendId::Cuda]
4137 }
4138 };
4139 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
4140 if dev == DeviceLocation::Cpu {
4141 vec![(BackendId::Cuda, cuda0)]
4142 } else {
4143 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4144 }
4145 };
4146 let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
4147 (id == c).then_some(DeviceLocation::Cpu)
4148 };
4149 let est = BytesEstimator {
4150 latency_ns: 10,
4151 ns_per_byte: 1,
4152 };
4153 let cpu_caps_val = cpu_caps();
4154 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4155 let opts = PlanOptions::new()
4156 .with_pinned_device(DeviceLocation::Cpu)
4157 .with_placements_for_device(&placements_fn)
4158 .with_fallback_placements_for(&fallback_fn)
4159 .with_capabilities_for(&caps_fn)
4160 .with_transfer_estimator(&est)
4161 .with_input_residency(&residency_fn);
4162 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4163
4164 let dev_of = |id: NodeId| plan.alternatives(id).unwrap().winner().unwrap().device;
4165 assert_eq!(dev_of(n1), cuda0, "the huge-win producer migrates");
4166 assert_eq!(
4167 dev_of(f),
4168 DeviceLocation::Cpu,
4169 "the locally-faster fused op stays where its output is needed",
4170 );
4171 assert_eq!(dev_of(n2), DeviceLocation::Cpu);
4172 assert!(
4173 plan.alternatives(f)
4174 .unwrap()
4175 .alternatives()
4176 .iter()
4177 .all(|c| c.device == DeviceLocation::Cpu),
4178 "fused set pruned to the committed device",
4179 );
4180 }
4181
4182 #[test]
4190 fn dp_exit_pricing_keeps_final_op_on_realize_target() {
4191 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
4192 let mut table = KernelBindingTable::new();
4193 register_unary_f32(
4194 &mut table,
4195 OpKind::NegElementwise,
4196 BackendId::Cpu,
4197 noop_kernel,
4198 cost_50,
4199 );
4200 register_unary_f32(
4201 &mut table,
4202 OpKind::NegElementwise,
4203 BackendId::Cuda,
4204 noop_kernel_b,
4205 cost_100,
4206 );
4207
4208 let mut g = Graph::new();
4209 let c = push_f32(&mut g, Op::Const, vec![], &[4]);
4210 let n = push_f32(&mut g, Op::Neg, vec![c], &[4]);
4211 let order = topo_order(&g, n);
4212
4213 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4214 if dev == DeviceLocation::Cpu {
4215 vec![BackendId::Cpu]
4216 } else {
4217 vec![BackendId::Cuda]
4218 }
4219 };
4220 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
4221 if dev == DeviceLocation::Cpu {
4222 vec![(BackendId::Cuda, cuda0)]
4223 } else {
4224 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4225 }
4226 };
4227 let residency_fn =
4228 move |id: NodeId| -> Option<DeviceLocation> { (id == c).then_some(cuda0) };
4229 let est = FlatEstimator { latency_ns: 1_000 };
4230 let cpu_caps_val = cpu_caps();
4231 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4232 let opts = PlanOptions::new()
4233 .with_pinned_device(DeviceLocation::Cpu)
4234 .with_placements_for_device(&placements_fn)
4235 .with_fallback_placements_for(&fallback_fn)
4236 .with_capabilities_for(&caps_fn)
4237 .with_transfer_estimator(&est)
4238 .with_input_residency(&residency_fn);
4239 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4240 let w = plan.alternatives(n).unwrap().winner().unwrap();
4241 assert_eq!(
4242 w.device,
4243 DeviceLocation::Cpu,
4244 "exit pricing flips the locally-winning GPU placement",
4245 );
4246 assert_eq!(
4247 w.inbound_transfer_ns, 1_000,
4248 "the GPU-resident input's crossing is priced on the CPU winner",
4249 );
4250 }
4251
4252 #[test]
4255 fn dp_exit_pricing_does_not_pin_when_kernel_gap_dominates() {
4256 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
4257 let mut table = KernelBindingTable::new();
4258 register_unary_f32(
4259 &mut table,
4260 OpKind::NegElementwise,
4261 BackendId::Cpu,
4262 noop_kernel,
4263 cost_10_000_000,
4264 );
4265 register_unary_f32(
4266 &mut table,
4267 OpKind::NegElementwise,
4268 BackendId::Cuda,
4269 noop_kernel_b,
4270 cost_100,
4271 );
4272
4273 let mut g = Graph::new();
4274 let c = push_f32(&mut g, Op::Const, vec![], &[4]);
4275 let n = push_f32(&mut g, Op::Neg, vec![c], &[4]);
4276 let order = topo_order(&g, n);
4277
4278 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4279 if dev == DeviceLocation::Cpu {
4280 vec![BackendId::Cpu]
4281 } else {
4282 vec![BackendId::Cuda]
4283 }
4284 };
4285 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
4286 if dev == DeviceLocation::Cpu {
4287 vec![(BackendId::Cuda, cuda0)]
4288 } else {
4289 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4290 }
4291 };
4292 let residency_fn =
4293 move |id: NodeId| -> Option<DeviceLocation> { (id == c).then_some(cuda0) };
4294 let est = FlatEstimator { latency_ns: 1_000 };
4295 let cpu_caps_val = cpu_caps();
4296 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4297 let opts = PlanOptions::new()
4298 .with_pinned_device(DeviceLocation::Cpu)
4299 .with_placements_for_device(&placements_fn)
4300 .with_fallback_placements_for(&fallback_fn)
4301 .with_capabilities_for(&caps_fn)
4302 .with_transfer_estimator(&est)
4303 .with_input_residency(&residency_fn);
4304 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4305 assert_eq!(
4306 plan.alternatives(n).unwrap().winner().unwrap().device,
4307 cuda0
4308 );
4309 }
4310
4311 #[test]
4317 fn dp_single_device_plan_equals_greedy_plan() {
4318 let mut table = KernelBindingTable::new();
4319 register_add_f32(
4320 &mut table,
4321 BackendId::Cpu,
4322 noop_kernel,
4323 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4324 );
4325 register_add_f32(
4326 &mut table,
4327 BackendId::Cuda,
4328 noop_kernel_b,
4329 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4330 );
4331 let (g, add_id) = build_add_graph();
4332 let order = topo_order(&g, add_id);
4333 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4336 if dev == DeviceLocation::Cpu {
4337 vec![BackendId::Cpu, BackendId::Cuda]
4338 } else {
4339 vec![]
4340 }
4341 };
4342 let fallback_fn = |_dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> { Vec::new() };
4343 let residency_fn = |_id: NodeId| -> Option<DeviceLocation> { Some(DeviceLocation::Cpu) };
4344 let cpu_caps_val = cpu_caps();
4345 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4346
4347 let base_opts = PlanOptions::new()
4348 .with_placements_for_device(&placements_fn)
4349 .with_capabilities_for(&caps_fn);
4350 let base = compile_plan(&g, &order, &table, &base_opts).expect("base");
4351
4352 let est = FlatEstimator {
4353 latency_ns: 1_000_000,
4354 };
4355 let wired_opts = PlanOptions::new()
4356 .with_placements_for_device(&placements_fn)
4357 .with_fallback_placements_for(&fallback_fn)
4358 .with_capabilities_for(&caps_fn)
4359 .with_transfer_estimator(&est)
4360 .with_input_residency(&residency_fn);
4361 let wired = compile_plan(&g, &order, &table, &wired_opts).expect("wired");
4362
4363 assert_eq!(base.alternatives.len(), wired.alternatives.len());
4364 let a = base.alternatives(add_id).unwrap();
4365 let b = wired.alternatives(add_id).unwrap();
4366 assert_eq!(a.len(), b.len(), "same candidate count");
4367 for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
4368 assert_eq!(ca.backend, cb.backend);
4369 assert_eq!(ca.device, cb.device);
4370 assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
4371 assert_eq!(ca.static_cost, cb.static_cost, "same Layer-1 cost");
4372 assert_eq!(cb.inbound_transfer_ns, 0, "no transfer term fires");
4373 }
4374 assert_eq!(a.context(), b.context(), "same decision context");
4375 }
4376
4377 fn run_diamond_regime(latency_ns: u64, expect: DeviceLocation) {
4381 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
4382 let mut table = KernelBindingTable::new();
4383 for op in [
4384 OpKind::NegElementwise,
4385 OpKind::SqrElementwise,
4386 OpKind::ReluElementwise,
4387 ] {
4388 register_unary_f32(&mut table, op, BackendId::Cpu, noop_kernel, cost_100);
4389 register_unary_f32(&mut table, op, BackendId::Cuda, noop_kernel_b, cost_10);
4390 }
4391 table.register_full(
4392 OpKind::AddElementwise,
4393 &[DType::F32, DType::F32, DType::F32],
4394 BackendId::Cpu,
4395 noop_kernel,
4396 KernelCaps::empty(),
4397 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4398 cost_100,
4399 );
4400 table.register_full(
4401 OpKind::AddElementwise,
4402 &[DType::F32, DType::F32, DType::F32],
4403 BackendId::Cuda,
4404 noop_kernel_b,
4405 KernelCaps::empty(),
4406 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4407 cost_10,
4408 );
4409
4410 let mut g = Graph::new();
4411 let c0 = push_f32(&mut g, Op::Const, vec![], &[4]);
4412 let a = push_f32(&mut g, Op::Neg, vec![c0], &[4]);
4413 let b = push_f32(&mut g, Op::Sqr, vec![a], &[4]);
4414 let c2 = push_f32(&mut g, Op::Relu, vec![a], &[4]);
4415 let d = push_f32(&mut g, Op::Add, vec![b, c2], &[4]);
4416 let order = topo_order(&g, d);
4417
4418 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4419 if dev == DeviceLocation::Cpu {
4420 vec![BackendId::Cpu]
4421 } else {
4422 vec![BackendId::Cuda]
4423 }
4424 };
4425 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
4426 if dev == DeviceLocation::Cpu {
4427 vec![(BackendId::Cuda, cuda0)]
4428 } else {
4429 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4430 }
4431 };
4432 let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
4433 (id == c0).then_some(DeviceLocation::Cpu)
4434 };
4435 let est = FlatEstimator { latency_ns };
4436 let cpu_caps_val = cpu_caps();
4437 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4438 let opts = PlanOptions::new()
4439 .with_pinned_device(DeviceLocation::Cpu)
4440 .with_placements_for_device(&placements_fn)
4441 .with_fallback_placements_for(&fallback_fn)
4442 .with_capabilities_for(&caps_fn)
4443 .with_transfer_estimator(&est)
4444 .with_input_residency(&residency_fn);
4445 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4446 for id in [a, b, c2, d] {
4447 assert_eq!(
4448 plan.alternatives(id).unwrap().winner().unwrap().device,
4449 expect,
4450 "diamond join must commit a consistent placement",
4451 );
4452 }
4453 }
4454
4455 #[test]
4460 fn dp_diamond_join_merges_consistently() {
4461 run_diamond_regime(1_000_000, DeviceLocation::Cpu);
4463 run_diamond_regime(0, DeviceLocation::Cuda { gpu_id: 0 });
4465 }
4466
4467 #[test]
4473 fn dp_thousand_node_chain_plans_within_sanity_bound() {
4474 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
4475 let mut table = KernelBindingTable::new();
4476 register_unary_f32(
4479 &mut table,
4480 OpKind::NegElementwise,
4481 BackendId::Cpu,
4482 noop_kernel,
4483 cost_100,
4484 );
4485 register_unary_f32(
4486 &mut table,
4487 OpKind::NegElementwise,
4488 BackendId::Cuda,
4489 noop_kernel_b,
4490 cost_100,
4491 );
4492
4493 let mut g = Graph::new();
4494 let c = push_f32(&mut g, Op::Const, vec![], &[4]);
4495 let mut prev = c;
4496 let mut nodes = Vec::with_capacity(1000);
4497 for _ in 0..1000 {
4498 let n = push_f32(&mut g, Op::Neg, vec![prev], &[4]);
4499 nodes.push(n);
4500 prev = n;
4501 }
4502 let order = topo_order(&g, prev);
4503
4504 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4505 if dev == DeviceLocation::Cpu {
4506 vec![BackendId::Cpu]
4507 } else {
4508 vec![BackendId::Cuda]
4509 }
4510 };
4511 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
4512 if dev == DeviceLocation::Cpu {
4513 vec![(BackendId::Cuda, cuda0)]
4514 } else {
4515 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4516 }
4517 };
4518 let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
4519 (id == c).then_some(DeviceLocation::Cpu)
4520 };
4521 let est = FlatEstimator { latency_ns: 10 };
4522 let cpu_caps_val = cpu_caps();
4523 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4524 let opts = PlanOptions::new()
4525 .with_pinned_device(DeviceLocation::Cpu)
4526 .with_placements_for_device(&placements_fn)
4527 .with_fallback_placements_for(&fallback_fn)
4528 .with_capabilities_for(&caps_fn)
4529 .with_transfer_estimator(&est)
4530 .with_input_residency(&residency_fn);
4531
4532 let start = std::time::Instant::now();
4533 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4534 let elapsed = start.elapsed();
4535 assert!(
4536 elapsed.as_secs() < 5,
4537 "1000-node DP plan took {elapsed:?} — pathological scaling",
4538 );
4539
4540 assert_eq!(plan.alternatives.len(), 1000);
4541 for &n in &nodes {
4542 assert_eq!(
4543 plan.alternatives(n).unwrap().winner().unwrap().device,
4544 DeviceLocation::Cpu,
4545 "equal-cost ties break toward the decision device end-to-end",
4546 );
4547 }
4548 }
4549
4550 #[test]
4554 fn dp_chain_connects_through_view_passthroughs() {
4555 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
4556 let mut table = KernelBindingTable::new();
4557 register_unary_f32(
4561 &mut table,
4562 OpKind::NegElementwise,
4563 BackendId::Cpu,
4564 noop_kernel,
4565 cost_800,
4566 );
4567 register_unary_f32(
4568 &mut table,
4569 OpKind::NegElementwise,
4570 BackendId::Cuda,
4571 noop_kernel_b,
4572 cost_100,
4573 );
4574
4575 let mut g = Graph::new();
4576 let c = push_f32(&mut g, Op::Const, vec![], &[4]);
4577 let n1 = push_f32(&mut g, Op::Neg, vec![c], &[4]);
4578 let r = push_f32(
4579 &mut g,
4580 Op::Reshape(Shape::from_dims(&[2, 2])),
4581 vec![n1],
4582 &[2, 2],
4583 );
4584 let n2 = push_f32(&mut g, Op::Neg, vec![r], &[2, 2]);
4585 let n3 = push_f32(&mut g, Op::Neg, vec![n2], &[2, 2]);
4586 let order = topo_order(&g, n3);
4587
4588 let placements_fn = |dev: DeviceLocation| -> Vec<BackendId> {
4589 if dev == DeviceLocation::Cpu {
4590 vec![BackendId::Cpu]
4591 } else {
4592 vec![BackendId::Cuda]
4593 }
4594 };
4595 let fallback_fn = move |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
4596 if dev == DeviceLocation::Cpu {
4597 vec![(BackendId::Cuda, cuda0)]
4598 } else {
4599 vec![(BackendId::Cpu, DeviceLocation::Cpu)]
4600 }
4601 };
4602 let residency_fn = move |id: NodeId| -> Option<DeviceLocation> {
4603 (id == c).then_some(DeviceLocation::Cpu)
4604 };
4605 let est = FlatEstimator { latency_ns: 1_000 };
4606 let cpu_caps_val = cpu_caps();
4607 let caps_fn = |_: BackendId| Some(&cpu_caps_val);
4608 let opts = PlanOptions::new()
4609 .with_pinned_device(DeviceLocation::Cpu)
4610 .with_placements_for_device(&placements_fn)
4611 .with_fallback_placements_for(&fallback_fn)
4612 .with_capabilities_for(&caps_fn)
4613 .with_transfer_estimator(&est)
4614 .with_input_residency(&residency_fn);
4615 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4616
4617 for id in [n1, n2, n3] {
4621 assert_eq!(
4622 plan.alternatives(id).unwrap().winner().unwrap().device,
4623 cuda0,
4624 "the chain must not break at the Reshape pass-through",
4625 );
4626 }
4627 assert!(
4628 plan.alternatives(r).is_none(),
4629 "Reshape carries no plan entry"
4630 );
4631 }
4632
4633 #[cfg(feature = "telemetry")]
4638 mod telemetry_wiring {
4639 use super::*;
4640 use crate::telemetry::{
4641 Contiguity, FdxOperandDesc, HwStamp, StructureKeyProvider, StructureKeyToken,
4642 TelemetryConfig, TelemetryHooks, TelemetryMode, TelemetrySink,
4643 };
4644 use std::sync::Mutex;
4645
4646 struct StubProvider;
4650 impl StructureKeyProvider for StubProvider {
4651 fn structure_key(
4652 &self,
4653 op_class: &str,
4654 operands: &[FdxOperandDesc],
4655 _arch: &str,
4656 ) -> Option<StructureKeyToken> {
4657 Some(StructureKeyToken(format!(
4658 "{op_class}:stub:n={}",
4659 operands.len()
4660 )))
4661 }
4662 }
4663
4664 fn cpu_hw() -> HwStamp {
4665 HwStamp {
4666 compute_capability: None,
4667 hardware_sku: "test-cpu".into(),
4668 driver_version: "n/a".into(),
4669 }
4670 }
4671
4672 fn generic_only_add_table() -> KernelBindingTable {
4677 let mut table = KernelBindingTable::new();
4678 table.register_full_with_source_generic(
4679 OpKind::AddElementwise,
4680 &[DType::F32, DType::F32, DType::F32],
4681 BackendId::Cpu,
4682 noop_kernel,
4683 KernelCaps::strided_input(),
4684 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4685 unknown_cost,
4686 "baracuda-generic-strided",
4687 true, 0, None, );
4691 table
4692 }
4693
4694 #[test]
4706 fn generic_only_cell_emits_one_aggregated_miss_through_compile_plan() {
4707 let table = generic_only_add_table();
4708 let (g, add_id) = build_add_graph();
4709 let order = topo_order(&g, add_id);
4710
4711 let sink = Mutex::new(TelemetrySink::new());
4712 let provider = StubProvider;
4713 let config = TelemetryConfig {
4714 mode: TelemetryMode::Coarse,
4715 out_path: None,
4716 };
4717 let hooks = TelemetryHooks {
4718 config: &config,
4719 sink: &sink,
4720 provider: &provider,
4721 hw: cpu_hw(),
4722 };
4723 let opts = PlanOptions::new()
4724 .without_cost_population()
4725 .with_telemetry(&hooks);
4726 let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4727
4728 let sink = sink.into_inner().expect("sink lock");
4729 assert_eq!(
4730 sink.miss_cell_count(),
4731 1,
4732 "a generic-only cell must emit exactly one aggregated miss cell",
4733 );
4734 let recs = sink.miss_records();
4735 assert_eq!(recs.len(), 1);
4736 assert_eq!(
4740 recs[0].wanted,
4741 StructureKeyToken("AddElementwise:stub:n=3".into()),
4742 "the wanted token must reflect the live op_class + operands",
4743 );
4744 assert_eq!(recs[0].fallback.kernel_source, "baracuda-generic-strided");
4745 assert_eq!(recs[0].fallback.op, OpKind::AddElementwise);
4746 assert_eq!(recs[0].fallback.backend, BackendId::Cpu);
4747 assert_eq!(
4748 recs[0].fallback.dtypes,
4749 vec![DType::F32, DType::F32, DType::F32],
4750 );
4751 assert_eq!(recs[0].count, 1);
4752 assert_eq!(recs[0].hw, cpu_hw());
4753 }
4754
4755 #[test]
4761 fn telemetry_off_emits_nothing_and_does_not_change_the_plan() {
4762 let table = generic_only_add_table();
4763 let (g, add_id) = build_add_graph();
4764 let order = topo_order(&g, add_id);
4765
4766 let sink = Mutex::new(TelemetrySink::new());
4767 let provider = StubProvider;
4768 let config = TelemetryConfig::default(); let hooks = TelemetryHooks {
4770 config: &config,
4771 sink: &sink,
4772 provider: &provider,
4773 hw: cpu_hw(),
4774 };
4775 let opts = PlanOptions::new()
4776 .without_cost_population()
4777 .with_telemetry(&hooks);
4778 let plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4779
4780 assert!(
4781 sink.into_inner().expect("sink lock").is_empty(),
4782 "Off mode must emit nothing",
4783 );
4784 let alts = plan.alternatives(add_id).expect("Add alternatives present");
4786 assert_eq!(alts.len(), 1);
4787 assert_eq!(alts.winner().unwrap().backend, BackendId::Cpu);
4788 }
4789
4790 #[test]
4795 fn specialized_winner_emits_no_miss() {
4796 let mut table = KernelBindingTable::new();
4797 register_add_f32(
4799 &mut table,
4800 BackendId::Cpu,
4801 noop_kernel,
4802 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4803 );
4804 let (g, add_id) = build_add_graph();
4805 let order = topo_order(&g, add_id);
4806
4807 let sink = Mutex::new(TelemetrySink::new());
4808 let provider = StubProvider;
4809 let config = TelemetryConfig {
4810 mode: TelemetryMode::Coarse,
4811 out_path: None,
4812 };
4813 let hooks = TelemetryHooks {
4814 config: &config,
4815 sink: &sink,
4816 provider: &provider,
4817 hw: cpu_hw(),
4818 };
4819 let opts = PlanOptions::new()
4820 .without_cost_population()
4821 .with_telemetry(&hooks);
4822 let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4823
4824 let sink = sink.into_inner().expect("sink lock");
4825 assert!(
4826 sink.miss_records().is_empty(),
4827 "a specialized (is_generic=false) winner must emit no miss",
4828 );
4829 assert_eq!(
4833 sink.dispatch_cell_count(),
4834 1,
4835 "a specialized winner still emits its dispatch record",
4836 );
4837 let _ = add_id;
4838 }
4839
4840 fn two_sibling_add_table() -> KernelBindingTable {
4848 let mut table = KernelBindingTable::new();
4849 table.register_full_with_source(
4850 OpKind::AddElementwise,
4851 &[DType::F32, DType::F32, DType::F32],
4852 BackendId::Cpu,
4853 noop_kernel,
4854 KernelCaps::empty(),
4855 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4856 unknown_cost,
4857 "portable-cpu",
4858 );
4859 table.register_full_with_source(
4860 OpKind::AddElementwise,
4861 &[DType::F32, DType::F32, DType::F32],
4862 BackendId::Cpu,
4863 noop_kernel_b,
4864 KernelCaps::empty(),
4865 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
4866 unknown_cost,
4867 "aocl",
4868 );
4869 table
4870 }
4871
4872 #[test]
4884 fn detailed_dispatch_record_emitted_with_per_candidate_judge_latency() {
4885 use crate::ranker::HashMapJudge;
4886 use fuel_ir::dispatch::SizeClass;
4887
4888 let table = two_sibling_add_table();
4889 let (g, add_id) = build_add_graph();
4890 let order = topo_order(&g, add_id);
4891
4892 let sc = SizeClass::from_elem_count(3);
4896 let mut judge = HashMapJudge::new();
4897 judge.insert(
4898 OpKind::AddElementwise,
4899 DType::F32,
4900 sc,
4901 BackendId::Cpu,
4902 "portable-cpu",
4903 41_230,
4904 0,
4905 );
4906
4907 let sink = Mutex::new(TelemetrySink::new());
4908 let provider = StubProvider;
4909 let config = TelemetryConfig {
4910 mode: TelemetryMode::Detailed,
4911 out_path: None,
4912 };
4913 let hooks = TelemetryHooks {
4914 config: &config,
4915 sink: &sink,
4916 provider: &provider,
4917 hw: cpu_hw(),
4918 };
4919 let opts = PlanOptions::new()
4920 .without_cost_population()
4921 .with_telemetry(&hooks)
4922 .with_judge(&judge);
4923 let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
4924
4925 let sink = sink.into_inner().expect("sink lock");
4926 assert_eq!(
4927 sink.dispatch_cell_count(),
4928 1,
4929 "one planned cell ⇒ exactly one aggregated dispatch record",
4930 );
4931 assert_eq!(
4932 sink.miss_cell_count(),
4933 0,
4934 "specialized winners emit no miss"
4935 );
4936 let recs = sink.dispatch_records();
4937 assert_eq!(recs.len(), 1);
4938 let rec = &recs[0];
4939 assert_eq!(rec.chosen.kernel_source, "portable-cpu");
4941 assert_eq!(rec.chosen.op, OpKind::AddElementwise);
4942 assert_eq!(rec.chosen.backend, BackendId::Cpu);
4943 assert_eq!(rec.chosen.dtypes, vec![DType::F32, DType::F32, DType::F32]);
4944 assert_eq!(
4945 rec.chosen.kernel_revision_hash, 0,
4946 "hand-written ⇒ untracked"
4947 );
4948 assert_eq!(rec.count, 1);
4949 assert_eq!(rec.hw, cpu_hw());
4950 assert_eq!(
4953 rec.structure_key,
4954 Some(StructureKeyToken("AddElementwise:stub:n=3".into())),
4955 );
4956 assert_eq!(rec.candidates.len(), 2, "both admitted alternatives appear");
4959 let by_src = |src: &str| {
4960 rec.candidates
4961 .iter()
4962 .find(|c| c.impl_id.kernel_source == src)
4963 .unwrap_or_else(|| panic!("candidate {src} present"))
4964 };
4965 assert_eq!(
4966 by_src("portable-cpu").latency_ns,
4967 Some(41_230),
4968 "the measured winner carries its own ns",
4969 );
4970 assert_eq!(
4971 by_src("aocl").latency_ns,
4972 None,
4973 "the unmeasured sibling is None, never a fabricated 0",
4974 );
4975 }
4976
4977 #[test]
4981 fn coarse_dispatch_record_has_no_candidates() {
4982 let table = two_sibling_add_table();
4983 let (g, add_id) = build_add_graph();
4984 let order = topo_order(&g, add_id);
4985
4986 let sink = Mutex::new(TelemetrySink::new());
4987 let provider = StubProvider;
4988 let config = TelemetryConfig {
4989 mode: TelemetryMode::Coarse,
4990 out_path: None,
4991 };
4992 let hooks = TelemetryHooks {
4993 config: &config,
4994 sink: &sink,
4995 provider: &provider,
4996 hw: cpu_hw(),
4997 };
4998 let opts = PlanOptions::new()
4999 .without_cost_population()
5000 .with_telemetry(&hooks);
5001 let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
5002
5003 let sink = sink.into_inner().expect("sink lock");
5004 assert_eq!(
5005 sink.dispatch_cell_count(),
5006 1,
5007 "Coarse still emits the chosen signal"
5008 );
5009 let recs = sink.dispatch_records();
5010 assert!(
5011 recs[0].candidates.is_empty(),
5012 "Coarse omits candidates (no Judge reads)"
5013 );
5014 assert_eq!(recs[0].chosen.kernel_source, "portable-cpu");
5015 assert_eq!(recs[0].count, 1);
5016 }
5017
5018 #[test]
5021 fn telemetry_off_emits_no_dispatch_record() {
5022 let table = two_sibling_add_table();
5023 let (g, add_id) = build_add_graph();
5024 let order = topo_order(&g, add_id);
5025
5026 let sink = Mutex::new(TelemetrySink::new());
5027 let provider = StubProvider;
5028 let config = TelemetryConfig::default(); let hooks = TelemetryHooks {
5030 config: &config,
5031 sink: &sink,
5032 provider: &provider,
5033 hw: cpu_hw(),
5034 };
5035 let opts = PlanOptions::new()
5036 .without_cost_population()
5037 .with_telemetry(&hooks);
5038 let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
5039
5040 assert!(
5041 sink.into_inner().expect("sink lock").is_empty(),
5042 "Off emits neither dispatch nor miss",
5043 );
5044 let _ = add_id;
5045 }
5046
5047 struct RecordingProvider {
5056 seen: Mutex<Vec<FdxOperandDesc>>,
5057 }
5058 impl StructureKeyProvider for RecordingProvider {
5059 fn structure_key(
5060 &self,
5061 op_class: &str,
5062 operands: &[FdxOperandDesc],
5063 _arch: &str,
5064 ) -> Option<StructureKeyToken> {
5065 *self.seen.lock().expect("record lock") = operands.to_vec();
5066 Some(StructureKeyToken(format!(
5067 "{op_class}:n={}",
5068 operands.len()
5069 )))
5070 }
5071 }
5072
5073 #[test]
5083 fn structure_key_operand_slice_folds_in_the_output() {
5084 let table = generic_only_add_table();
5085 let (g, add_id) = build_add_graph();
5086 let order = topo_order(&g, add_id);
5087
5088 let sink = Mutex::new(TelemetrySink::new());
5089 let provider = RecordingProvider {
5090 seen: Mutex::new(Vec::new()),
5091 };
5092 let config = TelemetryConfig {
5093 mode: TelemetryMode::Coarse,
5094 out_path: None,
5095 };
5096 let hooks = TelemetryHooks {
5097 config: &config,
5098 sink: &sink,
5099 provider: &provider,
5100 hw: cpu_hw(),
5101 };
5102 let opts = PlanOptions::new()
5103 .without_cost_population()
5104 .with_telemetry(&hooks);
5105 let _plan = compile_plan(&g, &order, &table, &opts).expect("compile");
5106
5107 let seen = provider.seen.into_inner().expect("record lock");
5108 assert_eq!(
5110 seen.len(),
5111 3,
5112 "operand slice must be node inputs (2) then the output (1)",
5113 );
5114 let out = seen.last().expect("output operand present");
5115 assert_eq!(
5117 out.dtype,
5118 DType::F32,
5119 "output operand carries the node dtype"
5120 );
5121 assert_eq!(
5122 out.shape,
5123 vec![3i64],
5124 "output operand carries the node shape"
5125 );
5126 assert_eq!(
5127 out.contiguity,
5128 Contiguity::Contiguous,
5129 "the Add output is contiguous",
5130 );
5131 let _ = add_id;
5132 }
5133 }
5134}