1use std::collections::{HashMap, HashSet};
39use std::sync::mpsc::{Sender, channel};
40use std::sync::{Arc, RwLock};
41use std::thread;
42
43use fuel_graph::opt::{execution_plan, insert_safety_copies};
44use fuel_graph::registry::FusedOpParams;
45use fuel_graph::{Graph, Node, NodeId, Op, PickedRoute};
46use fuel_ir::dispatch::OpKind;
47use fuel_ir::probe::BackendId;
48use fuel_ir::{DType, DeviceLocation, Error, Layout, Result, SymEnv};
49
50use crate::compiled::{
51 CompiledNode, CompletionHandle, compile_node, execute_compiled_with_wait_hint,
52};
53use crate::dispatch::global_bindings;
54use crate::kernel::{KernelBindingTable, KernelDTypes, MatmulM, OpParams};
55use crate::optimize::OptimizedGraph;
56use crate::ranker::{BackendRuntimeLookup, RuntimeSelector, resolve_branch};
57#[cfg(test)]
60use crate::ranker::pick_route;
61use fuel_memory::Storage;
62
63enum OrderSource<'a> {
90 Default,
91 Optimized {
92 optimized: &'a OptimizedGraph,
93 route: Option<&'a PickedRoute>,
96 },
97 Streaming {
108 optimized: &'a OptimizedGraph,
109 pick: &'a StreamingPick,
110 },
111 Cached {
128 optimized: &'a OptimizedGraph,
129 order_cache: &'a std::sync::OnceLock<Vec<NodeId>>,
130 },
131}
132
133#[derive(Clone)]
149pub(crate) struct StreamingPick {
150 pub selector: Arc<dyn RuntimeSelector>,
155}
156
157pub type StorageCache = HashMap<NodeId, Arc<RwLock<Storage>>>;
161
162pub struct PersistentOutputs {
178 map: HashMap<NodeId, Arc<RwLock<Storage>>>,
179 mode: PersistentMode,
180}
181
182#[derive(Clone, Copy, PartialEq, Eq)]
183enum PersistentMode {
184 Record,
186 Reuse,
188}
189
190impl PersistentOutputs {
191 pub fn recording() -> Self {
193 Self {
194 map: HashMap::new(),
195 mode: PersistentMode::Record,
196 }
197 }
198
199 pub fn reuse(map: HashMap<NodeId, Arc<RwLock<Storage>>>) -> Self {
202 Self {
203 map,
204 mode: PersistentMode::Reuse,
205 }
206 }
207
208 pub fn into_map(self) -> HashMap<NodeId, Arc<RwLock<Storage>>> {
211 self.map
212 }
213}
214
215#[cfg(feature = "cuda")]
292fn op_kind_is_capture_writeinto(op: OpKind) -> bool {
293 matches!(
294 op,
295 OpKind::AddElementwise
297 | OpKind::SubElementwise
298 | OpKind::MulElementwise
299 | OpKind::DivElementwise
300 | OpKind::MaximumElementwise
301 | OpKind::MinimumElementwise
302 | OpKind::PowElementwise
303 | OpKind::RemElementwise
304 | OpKind::SoftmaxLastDim
306 | OpKind::LogSoftmaxLastDim
307 | OpKind::Affine
309 | OpKind::MatMul
312 | OpKind::IndexSelect
314 | OpKind::RmsNormLastDim
317 | OpKind::SiluElementwise
319 | OpKind::GeluTanhElementwise
320 | OpKind::GeluElementwise
321 | OpKind::NegElementwise
325 | OpKind::Concat
329 | OpKind::Rope
332 )
333}
334
335#[cfg(feature = "cuda")]
346pub struct CapturedDecode {
347 pub run: fuel_cuda_backend::CapturedRun,
349 pub output: Arc<RwLock<Storage>>,
351 pub output_layout: Layout,
353 pub persistent: HashMap<NodeId, Arc<RwLock<Storage>>>,
356 pub retained_inputs: Vec<Arc<RwLock<Storage>>>,
364}
365
366#[cfg(feature = "cuda")]
386pub struct CapturedDecodeSession {
387 device: fuel_cuda_backend::CudaDevice,
388 captured: CapturedDecode,
389 per_token_inputs: HashMap<NodeId, Arc<RwLock<Storage>>>,
393 _retained_inputs: Vec<Arc<RwLock<Storage>>>,
399}
400
401#[cfg(feature = "cuda")]
402impl CapturedDecodeSession {
403 pub fn capture(
413 graph: Arc<RwLock<Graph>>,
414 target: NodeId,
415 inputs: StorageCache,
416 per_token_input_ids: &[NodeId],
417 sym_env: SymEnv,
418 ) -> Result<Self> {
419 use fuel_memory::BackendStorage;
420 let retained_inputs: Vec<Arc<RwLock<Storage>>> = inputs.values().cloned().collect();
424 let mut per_token_inputs = HashMap::with_capacity(per_token_input_ids.len());
425 for &id in per_token_input_ids {
426 let arc = inputs.get(&id).cloned().ok_or_else(|| {
427 Error::Msg(format!(
428 "CapturedDecodeSession::capture: per-token input {:?} not in inputs",
429 id,
430 ))
431 .bt()
432 })?;
433 per_token_inputs.insert(id, arc);
434 }
435 let device = {
437 let any = per_token_inputs
438 .values()
439 .next()
440 .or_else(|| inputs.values().next())
441 .ok_or_else(|| {
442 Error::Msg("CapturedDecodeSession::capture: no inputs".into()).bt()
443 })?;
444 let guard = any.read().map_err(|_| poisoned("input storage"))?;
445 match &guard.inner {
446 BackendStorage::Cuda(c) => c.device().clone(),
447 other => {
448 return Err(Error::Msg(format!(
449 "CapturedDecodeSession::capture: input is BackendStorage::{:?}, not Cuda",
450 std::mem::discriminant(other),
451 ))
452 .bt());
453 }
454 }
455 };
456 let captured = PipelinedExecutor::capture_decode(graph, target, inputs, sym_env)?;
457 Ok(Self {
458 device,
459 captured,
460 per_token_inputs,
461 _retained_inputs: retained_inputs,
462 })
463 }
464
465 pub fn node_outputs(&self) -> &HashMap<NodeId, Arc<RwLock<Storage>>> {
496 &self.captured.persistent
497 }
498
499 pub fn replay_token(&self, updates: &[(NodeId, &[u8])]) -> Result<Arc<RwLock<Storage>>> {
500 use fuel_memory::BackendStorage;
501 for (id, bytes) in updates {
502 let arc = self.per_token_inputs.get(id).ok_or_else(|| {
503 Error::Msg(format!(
504 "CapturedDecodeSession::replay_token: {:?} is not a per-token input",
505 id,
506 ))
507 .bt()
508 })?;
509 let guard = arc.read().map_err(|_| poisoned("per-token input"))?;
510 match &guard.inner {
511 BackendStorage::Cuda(c) => c.write_from_host(bytes)?,
512 other => {
513 return Err(Error::Msg(format!(
514 "CapturedDecodeSession::replay_token: input {:?} is \
515 BackendStorage::{:?}, not Cuda",
516 id,
517 std::mem::discriminant(other),
518 ))
519 .bt());
520 }
521 }
522 }
523 self.captured.run.replay(&self.device)?;
524 self.device.synchronize()?;
525 Ok(Arc::clone(&self.captured.output))
526 }
527
528 pub fn output(&self) -> &Arc<RwLock<Storage>> {
532 &self.captured.output
533 }
534
535 pub fn output_layout(&self) -> &Layout {
537 &self.captured.output_layout
538 }
539}
540
541enum WorkItemKind {
544 ConstAdopt,
547 ViewOf { input: NodeId },
552 ContiguizeOf { input: NodeId },
559 Kernel,
562 Iota { len: usize },
568 ReleaseMarker,
579 WriteSlice { dest: NodeId, source: NodeId },
597 WriteSliceRotating {
603 dest: NodeId,
604 source: NodeId,
605 position: NodeId,
606 },
607 WriteSliceDoff {
614 dest: NodeId,
615 source: NodeId,
616 offset: NodeId,
617 },
618 Copy { target_location: DeviceLocation },
636 Move { target_location: DeviceLocation },
660 Alloc { target_location: DeviceLocation },
678 ZeroFill,
697 InplaceKernel {
717 target_idx: usize,
721 },
722 SlotView { producer: NodeId },
732 SlotOwn { producer: NodeId, slot: u32 },
741}
742
743fn work_item_kind_name(kind: &WorkItemKind) -> &'static str {
749 match kind {
750 WorkItemKind::ConstAdopt => "item.ConstAdopt",
751 WorkItemKind::ViewOf { .. } => "item.ViewOf",
752 WorkItemKind::ContiguizeOf { .. } => "item.ContiguizeOf",
753 WorkItemKind::Kernel => "item.Kernel",
754 WorkItemKind::Iota { .. } => "item.Iota",
755 WorkItemKind::ReleaseMarker => "item.ReleaseMarker",
756 WorkItemKind::WriteSlice { .. } => "item.WriteSlice",
757 WorkItemKind::WriteSliceRotating { .. } => "item.WriteSliceRotating",
758 WorkItemKind::WriteSliceDoff { .. } => "item.WriteSliceDoff",
759 WorkItemKind::Copy { .. } => "item.Copy",
760 WorkItemKind::Move { .. } => "item.Move",
761 WorkItemKind::Alloc { .. } => "item.Alloc",
762 WorkItemKind::ZeroFill => "item.ZeroFill",
763 WorkItemKind::InplaceKernel { .. } => "item.InplaceKernel",
764 WorkItemKind::SlotView { .. } => "item.SlotView",
765 WorkItemKind::SlotOwn { .. } => "item.SlotOwn",
766 }
767}
768
769fn is_backend_neutral(kind: &WorkItemKind) -> bool {
782 matches!(
783 kind,
784 WorkItemKind::ConstAdopt
785 | WorkItemKind::ViewOf { .. }
786 | WorkItemKind::SlotView { .. }
787 | WorkItemKind::ContiguizeOf { .. }
788 )
789}
790
791struct WorkItem {
792 node_id: NodeId,
793 inputs: Vec<NodeId>,
794 elem_count: usize,
797 dtype: DType,
798 target_backend: BackendId,
799 kind: WorkItemKind,
801 compiled: Option<CompiledNode>,
804 output_layout: Layout,
811 destructive_input: Option<usize>,
818 output_bundle: Option<Arc<[fuel_ir::storage::OutputView]>>,
826}
827
828fn extend_with_side_effect_roots(graph: &Graph, user_roots: &[NodeId]) -> Vec<NodeId> {
837 let side = graph.side_effect_roots();
838 if side.is_empty() {
839 return user_roots.to_vec();
840 }
841 let mut out = Vec::with_capacity(user_roots.len() + side.len());
842 out.extend_from_slice(user_roots);
843 for &s in side {
844 if !out.contains(&s) {
845 out.push(s);
846 }
847 }
848 out
849}
850
851pub struct PipelinedExecutor;
909
910impl PipelinedExecutor {
911 pub fn realize(
934 graph: Arc<RwLock<Graph>>,
935 target: NodeId,
936 inputs: StorageCache,
937 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
938 Self::realize_inner(
939 graph,
940 target,
941 inputs,
942 OrderSource::Default,
943 SymEnv::default(),
944 )
945 .map(|(s, l, _produced)| (s, l))
946 }
947
948 pub fn realize_with_env(
961 graph: Arc<RwLock<Graph>>,
962 target: NodeId,
963 inputs: StorageCache,
964 sym_env: SymEnv,
965 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
966 Self::realize_inner(graph, target, inputs, OrderSource::Default, sym_env)
967 .map(|(s, l, _produced)| (s, l))
968 }
969
970 pub fn realize_with_env_producing(
979 graph: Arc<RwLock<Graph>>,
980 target: NodeId,
981 inputs: StorageCache,
982 sym_env: SymEnv,
983 ) -> Result<(Arc<RwLock<Storage>>, Layout, SymEnv)> {
984 Self::realize_inner(graph, target, inputs, OrderSource::Default, sym_env)
985 }
986
987 pub fn realize_with_optimized(
999 graph: Arc<RwLock<Graph>>,
1000 target: NodeId,
1001 inputs: StorageCache,
1002 optimized: &OptimizedGraph,
1003 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
1004 Self::realize_inner(
1005 graph,
1006 target,
1007 inputs,
1008 OrderSource::Optimized {
1009 optimized,
1010 route: None,
1011 },
1012 SymEnv::default(),
1013 )
1014 .map(|(s, l, _produced)| (s, l))
1015 }
1016
1017 pub fn realize_with_optimized_env(
1024 graph: Arc<RwLock<Graph>>,
1025 target: NodeId,
1026 inputs: StorageCache,
1027 optimized: &OptimizedGraph,
1028 sym_env: SymEnv,
1029 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
1030 Self::realize_inner(
1031 graph,
1032 target,
1033 inputs,
1034 OrderSource::Optimized {
1035 optimized,
1036 route: None,
1037 },
1038 sym_env,
1039 )
1040 .map(|(s, l, _produced)| (s, l))
1041 }
1042
1043 pub fn realize_with_optimized_route(
1053 graph: Arc<RwLock<Graph>>,
1054 target: NodeId,
1055 inputs: StorageCache,
1056 optimized: &OptimizedGraph,
1057 route: &PickedRoute,
1058 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
1059 Self::realize_inner(
1060 graph,
1061 target,
1062 inputs,
1063 OrderSource::Optimized {
1064 optimized,
1065 route: Some(route),
1066 },
1067 SymEnv::default(),
1068 )
1069 .map(|(s, l, _produced)| (s, l))
1070 }
1071
1072 pub fn realize_with_optimized_route_env(
1076 graph: Arc<RwLock<Graph>>,
1077 target: NodeId,
1078 inputs: StorageCache,
1079 optimized: &OptimizedGraph,
1080 route: &PickedRoute,
1081 sym_env: SymEnv,
1082 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
1083 Self::realize_inner(
1084 graph,
1085 target,
1086 inputs,
1087 OrderSource::Optimized {
1088 optimized,
1089 route: Some(route),
1090 },
1091 sym_env,
1092 )
1093 .map(|(s, l, _produced)| (s, l))
1094 }
1095
1096 pub fn realize_with_optimized_picking_env(
1116 graph: Arc<RwLock<Graph>>,
1117 target: NodeId,
1118 inputs: StorageCache,
1119 optimized: &OptimizedGraph,
1120 selector: Option<Arc<dyn RuntimeSelector>>,
1121 lookup: Option<BackendRuntimeLookup>,
1122 sym_env: SymEnv,
1123 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
1124 match Self::streaming_pick_for(&graph, selector, lookup)? {
1125 Some(pick) => Self::realize_inner(
1126 graph,
1127 target,
1128 inputs,
1129 OrderSource::Streaming {
1130 optimized,
1131 pick: &pick,
1132 },
1133 sym_env,
1134 )
1135 .map(|(s, l, _produced)| (s, l)),
1136 None => Self::realize_with_optimized_env(graph, target, inputs, optimized, sym_env),
1137 }
1138 }
1139
1140 pub fn realize_with_optimized_picking_cached_env(
1153 graph: Arc<RwLock<Graph>>,
1154 target: NodeId,
1155 inputs: StorageCache,
1156 optimized: &OptimizedGraph,
1157 order_cache: &std::sync::OnceLock<Vec<NodeId>>,
1158 selector: Option<Arc<dyn RuntimeSelector>>,
1159 lookup: Option<BackendRuntimeLookup>,
1160 sym_env: SymEnv,
1161 ) -> Result<(Arc<RwLock<Storage>>, Layout)> {
1162 match Self::streaming_pick_for(&graph, selector, lookup)? {
1163 Some(pick) => Self::realize_inner(
1164 graph,
1165 target,
1166 inputs,
1167 OrderSource::Streaming {
1168 optimized,
1169 pick: &pick,
1170 },
1171 sym_env,
1172 )
1173 .map(|(s, l, _produced)| (s, l)),
1174 None => Self::realize_inner(
1175 graph,
1176 target,
1177 inputs,
1178 OrderSource::Cached {
1179 optimized,
1180 order_cache,
1181 },
1182 sym_env,
1183 )
1184 .map(|(s, l, _produced)| (s, l)),
1185 }
1186 }
1187
1188 fn streaming_pick_for(
1201 graph: &Arc<RwLock<Graph>>,
1202 selector: Option<Arc<dyn RuntimeSelector>>,
1203 _lookup: Option<BackendRuntimeLookup>,
1207 ) -> Result<Option<StreamingPick>> {
1208 let Some(selector) = selector else {
1209 return Ok(None);
1210 };
1211 let g = graph
1212 .read()
1213 .map_err(|_| Error::Msg("graph lock poisoned".into()).bt())?;
1214 let has_branch = (0..g.len()).any(|i| matches!(g.node(NodeId(i)).op, Op::Branch { .. }));
1215 if !has_branch {
1216 return Ok(None);
1217 }
1218 Ok(Some(StreamingPick { selector }))
1219 }
1220
1221 #[cfg(feature = "cuda")]
1247 pub fn capture_decode(
1248 graph: Arc<RwLock<Graph>>,
1249 target: NodeId,
1250 inputs: StorageCache,
1251 sym_env: SymEnv,
1252 ) -> Result<CapturedDecode> {
1253 use fuel_memory::BackendStorage;
1254
1255 {
1257 let mut g = graph.write().map_err(|_| poisoned("graph lock"))?;
1258 let effective_roots = extend_with_side_effect_roots(&g, &[target]);
1259 insert_safety_copies(&mut g, &effective_roots);
1260 }
1261 let (compiler_work, _wait_set, seed_layouts) = {
1262 let g = graph.read().map_err(|_| poisoned("graph lock"))?;
1263 let effective_roots = extend_with_side_effect_roots(&g, &[target]);
1264 let (work, wait_set) =
1265 compiler_work_and_wait_set_for(&g, &effective_roots, &OrderSource::Default);
1266 let mut layouts: HashMap<NodeId, Layout> = HashMap::with_capacity(inputs.len());
1267 for &id in inputs.keys() {
1268 layouts.insert(id, g.layout(id));
1269 }
1270 (work, wait_set, layouts)
1271 };
1272
1273 let (tx, rx) = channel::<Result<WorkItem>>();
1276 let graph_for_compiler = Arc::clone(&graph);
1277 let compiler = thread::spawn(move || {
1278 compiler_thread_body(graph_for_compiler, compiler_work, sym_env, tx);
1279 });
1280 let mut items: Vec<WorkItem> = Vec::new();
1281 for item in rx {
1282 items.push(item?);
1283 }
1284 compiler
1285 .join()
1286 .map_err(|_| Error::Msg("compiler thread panicked".to_string()).bt())?;
1287
1288 {
1294 let off: Vec<String> = items
1295 .iter()
1296 .filter_map(|it| match &it.kind {
1297 WorkItemKind::Kernel if it.target_backend != BackendId::Cuda => {
1298 let desc = it
1299 .compiled
1300 .as_ref()
1301 .map(|c| format!("{:?} dtypes={:?}", c.op, c.dtypes))
1302 .unwrap_or_else(|| "?".to_string());
1303 Some(format!(
1304 "node {:?}: {desc} on {:?}",
1305 it.node_id, it.target_backend
1306 ))
1307 }
1308 _ => None,
1309 })
1310 .collect();
1311 if !off.is_empty() {
1312 eprintln!(
1313 "[capture-diag] {} kernel(s) placed OFF Cuda (these block capture): {}",
1314 off.len(),
1315 off.join("; "),
1316 );
1317 }
1318 }
1319
1320 for item in &items {
1323 match &item.kind {
1324 WorkItemKind::Move { .. } => {
1333 return Err(Error::Msg(format!(
1334 "capture_decode: graph has an Op::Move at node {:?}; \
1335 Move's destructive source-release half is driven by \
1336 per-realize cleanup, not per-replay, so capture is \
1337 not supported for Move (out of scope — see the \
1338 `capture_decode` validation loop's Move arm).",
1339 item.node_id,
1340 ))
1341 .bt());
1342 }
1343 WorkItemKind::Copy { target_location } => {
1355 if !matches!(target_location, DeviceLocation::Cuda { .. }) {
1356 return Err(Error::Msg(format!(
1357 "capture_decode: graph has a cross-device Op::Copy at \
1358 node {:?} (target {:?}); capture requires a \
1359 single-device CUDA decode step — only same-device \
1360 CUDA→CUDA copies are capture-safe.",
1361 item.node_id, target_location,
1362 ))
1363 .bt());
1364 }
1365 }
1366 WorkItemKind::Kernel => {
1367 if item.target_backend != BackendId::Cuda {
1368 return Err(Error::Msg(format!(
1369 "capture_decode: kernel node {:?} targets {:?}, not Cuda",
1370 item.node_id, item.target_backend,
1371 ))
1372 .bt());
1373 }
1374 let compiled = item.compiled.as_ref().ok_or_else(|| {
1375 Error::Msg(format!(
1376 "capture_decode: kernel node {:?} has no compiled op",
1377 item.node_id,
1378 ))
1379 .bt()
1380 })?;
1381 let op = compiled.op;
1382 if !op_kind_is_capture_writeinto(op) {
1383 return Err(Error::Msg(format!(
1384 "capture_decode: op {:?} (node {:?}) has no write-into CUDA \
1385 wrapper yet — capture unsupported for this op. Extend \
1386 op_kind_is_capture_writeinto once its write-into variant lands.",
1387 op, item.node_id,
1388 ))
1389 .bt());
1390 }
1391 if op == OpKind::Concat {
1399 if let OpParams::Concat {
1400 input_dim_sizes, ..
1401 } = &compiled.op_params
1402 {
1403 if input_dim_sizes.len() != 2 {
1404 return Err(Error::Msg(format!(
1405 "capture_decode: Concat node {:?} has {} inputs; only \
1406 2-input concat (the rope rotate-half shape) has a \
1407 write-into CUDA path today — N-ary (N>2) concat would \
1408 allocate mid-capture.",
1409 item.node_id,
1410 input_dim_sizes.len(),
1411 ))
1412 .bt());
1413 }
1414 }
1415 }
1416 }
1417 _ => {}
1419 }
1420 }
1421
1422 let mut persistent = PersistentOutputs::recording();
1424 let mut warm_cache: StorageCache = inputs.clone();
1425 let mut warm_layouts = seed_layouts.clone();
1426 let mut warm_syms = SymEnv::new();
1427 for item in &items {
1428 let _handle = execute_work_item(
1429 item,
1430 &mut warm_cache,
1431 &mut warm_layouts,
1432 &mut warm_syms,
1433 false,
1434 Some(&mut persistent),
1435 )
1436 .map_err(|e| with_node_location(&graph, item.node_id, e))?;
1437 }
1438 sync_active_cuda_devices(&warm_cache)?;
1440
1441 let output = warm_cache.get(&target).cloned().ok_or_else(|| {
1442 Error::Msg(format!(
1443 "capture_decode: target {:?} not realized in warm run",
1444 target,
1445 ))
1446 .bt()
1447 })?;
1448 let output_layout = warm_layouts.get(&target).cloned().ok_or_else(|| {
1449 Error::Msg(format!("capture_decode: target {:?} has no layout", target)).bt()
1450 })?;
1451
1452 let device = {
1454 let guard = output.read().map_err(|_| poisoned("output storage"))?;
1455 match &guard.inner {
1456 BackendStorage::Cuda(c) => c.device().clone(),
1457 other => {
1458 return Err(Error::Msg(format!(
1459 "capture_decode: target {:?} is BackendStorage::{:?}, not Cuda",
1460 target,
1461 std::mem::discriminant(other),
1462 ))
1463 .bt());
1464 }
1465 }
1466 };
1467
1468 let mut reuse = PersistentOutputs::reuse(persistent.into_map());
1470 let retained_inputs: Vec<Arc<RwLock<Storage>>> = inputs.values().cloned().collect();
1474 let mut cap_cache: StorageCache = inputs;
1475 let mut cap_layouts = seed_layouts;
1476 let mut cap_syms = SymEnv::new();
1477 let mut inner_err: Option<Error> = None;
1478 let items_ref = &items;
1479 let graph_ref = &graph;
1480 let run = device.capture_run(|_stream| {
1481 for item in items_ref {
1485 match execute_work_item(
1486 item,
1487 &mut cap_cache,
1488 &mut cap_layouts,
1489 &mut cap_syms,
1490 false,
1491 Some(&mut reuse),
1492 ) {
1493 Ok(_handle) => {}
1494 Err(e) => {
1495 inner_err = Some(with_node_location(graph_ref, item.node_id, e));
1496 break;
1497 }
1498 }
1499 }
1500 Ok(())
1501 })?;
1502 if let Some(e) = inner_err {
1503 return Err(e);
1504 }
1505
1506 Ok(CapturedDecode {
1507 run,
1508 output,
1509 output_layout,
1510 retained_inputs,
1511 persistent: reuse.into_map(),
1512 })
1513 }
1514
1515 fn realize_inner(
1516 graph: Arc<RwLock<Graph>>,
1517 target: NodeId,
1518 inputs: StorageCache,
1519 order_source: OrderSource<'_>,
1520 sym_env: SymEnv,
1521 ) -> Result<(Arc<RwLock<Storage>>, Layout, SymEnv)> {
1522 {
1527 let _t = crate::compiled::op_profile::PhaseTimer::start("realize.safety_copies");
1528 let mut g = graph.write().map_err(|_| poisoned("graph lock"))?;
1529 let effective_roots = extend_with_side_effect_roots(&g, &[target]);
1530 insert_safety_copies(&mut g, &effective_roots);
1531 }
1532
1533 let (compiler_work, wait_set, mut layout_cache): (
1543 CompilerWork,
1544 WaitSet,
1545 HashMap<NodeId, Layout>,
1546 ) = {
1547 let _t = crate::compiled::op_profile::PhaseTimer::start("realize.plan_order");
1548 let g = graph.read().map_err(|_| poisoned("graph lock"))?;
1549 let effective_roots = extend_with_side_effect_roots(&g, &[target]);
1550 let (work, wait_set) =
1566 compiler_work_and_wait_set_for(&g, &effective_roots, &order_source);
1567 let mut layouts = HashMap::with_capacity(inputs.len());
1568 for &id in inputs.keys() {
1569 layouts.insert(id, g.layout(id));
1570 }
1571 (work, wait_set, layouts)
1572 };
1573
1574 let (tx, rx) = channel::<Result<WorkItem>>();
1575 let graph_for_compiler = Arc::clone(&graph);
1576
1577 let compiler = {
1583 let _t = crate::compiled::op_profile::PhaseTimer::start("realize.spawn_compiler");
1584 thread::spawn(move || {
1585 let _t = crate::compiled::op_profile::PhaseTimer::start(
1586 "compiler_thread (runs concurrently)",
1587 );
1588 compiler_thread_body(graph_for_compiler, compiler_work, sym_env, tx);
1589 })
1590 };
1591
1592 let plan_generation: Option<u64> = generation_for(&order_source);
1605 let mut current_chunk_backend: Option<BackendId> = None;
1606 let mut cache: StorageCache = inputs;
1607 let mut produced_syms = SymEnv::new();
1613 let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
1620 let mut inflight_vulkan: Vec<InflightVulkan> = Vec::new();
1630 let mut multi_backend = false;
1636 let loop_total = crate::compiled::op_profile::PhaseTimer::start("realize.loop_total");
1637 loop {
1638 let t_wait =
1639 crate::compiled::op_profile::PhaseTimer::start("realize.wait_for_compiler");
1640 let Ok(item) = rx.recv() else {
1641 break;
1642 };
1643 drop(t_wait);
1644 let item = item?;
1645 if !is_backend_neutral(&item.kind) && current_chunk_backend != Some(item.target_backend)
1652 {
1653 if let Some(plan_gen) = plan_generation {
1654 let live = crate::dispatch::topology_generation();
1655 if live != plan_gen {
1656 return Err(Error::TopologyChanged {
1657 plan_generation: plan_gen,
1658 current_generation: live,
1659 }
1660 .bt());
1661 }
1662 }
1663 let leaving = current_chunk_backend;
1664 if let Some(prev) = leaving
1665 && prev != item.target_backend
1666 {
1667 multi_backend = true;
1670 if prev == BackendId::Vulkan {
1676 {
1677 let _pt = crate::compiled::op_profile::PhaseTimer::start(
1678 "realize.eager_submit_vulkan",
1679 );
1680 eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
1681 }
1682 }
1683 }
1684 current_chunk_backend = Some(item.target_backend);
1685 }
1686 if item.target_backend == BackendId::Vulkan && matches!(item.kind, WorkItemKind::Kernel)
1690 {
1691 multi_backend = true;
1692 }
1693 if matches!(
1705 item.kind,
1706 WorkItemKind::Copy { .. } | WorkItemKind::Move { .. }
1707 ) {
1708 if multi_backend {
1713 {
1714 let _pt = crate::compiled::op_profile::PhaseTimer::start(
1715 "realize.eager_submit_vulkan",
1716 );
1717 eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
1718 }
1719 }
1720 if let Some(&producer) = item.inputs.first() {
1721 if multi_backend && copy_source_is_vulkan(&cache, producer)? {
1730 {
1731 let _pt = crate::compiled::op_profile::PhaseTimer::start(
1732 "realize.drain_inflight_vulkan",
1733 );
1734 drain_inflight_vulkan(&mut inflight_vulkan)?;
1735 }
1736 }
1737 wait_producer_handle(&mut handles, producer)?;
1742 }
1743 }
1744 let t_exec =
1745 crate::compiled::op_profile::PhaseTimer::start(work_item_kind_name(&item.kind));
1746 let handle = execute_work_item(
1747 &item,
1748 &mut cache,
1749 &mut layout_cache,
1750 &mut produced_syms,
1751 will_be_waited(&wait_set, item.node_id),
1752 None,
1755 )
1756 .map_err(|e| with_node_location(&graph, item.node_id, e))?;
1757 drop(t_exec);
1758 store_handle(&mut handles, item.node_id, handle);
1759 if let Some(d_idx) = item.destructive_input
1760 && let Some(&destroyed) = item.inputs.get(d_idx)
1761 && destroyed != target
1762 {
1763 let _t =
1787 crate::compiled::op_profile::PhaseTimer::start("realize.evict_destroyed_input");
1788 if multi_backend {
1789 defer_evicted_vulkan_buffer(&cache, destroyed, &mut inflight_vulkan)?;
1790 } else if let Some(d_arc) = cache.get(&destroyed) {
1791 let _f = crate::compiled::op_profile::PhaseTimer::start(
1792 "realize.evict_force_flush_vulkan",
1793 );
1794 force_flush_vulkan(d_arc)?;
1795 }
1796 cache.remove(&destroyed);
1797 layout_cache.remove(&destroyed);
1798 handles.remove(&destroyed);
1802 }
1803 }
1804
1805 drop(loop_total);
1806 {
1807 let _t = crate::compiled::op_profile::PhaseTimer::start("realize.join_compiler");
1808 compiler
1809 .join()
1810 .map_err(|_| Error::Msg("compiler thread panicked".to_string()).bt())?;
1811 }
1812
1813 let post_loop = crate::compiled::op_profile::PhaseTimer::start("realize.post_loop_drains");
1818 drain_handles(&mut handles)?;
1819 sync_active_cuda_devices(&cache)?;
1826 {
1830 let _pt =
1831 crate::compiled::op_profile::PhaseTimer::start("realize.drain_inflight_vulkan");
1832 drain_inflight_vulkan(&mut inflight_vulkan)?;
1833 }
1834 drain_vulkan_pending(&cache)?;
1840 debug_assert!(
1841 inflight_vulkan.is_empty(),
1842 "A4b-4: in-flight Vulkan batch list must be empty after realize-end drain",
1843 );
1844
1845 drop(post_loop);
1846 let storage = cache.remove(&target).ok_or_else(|| {
1847 Error::Msg(format!(
1848 "PipelinedExecutor::realize: target slot {:?} not populated after execution",
1849 target
1850 ))
1851 .bt()
1852 })?;
1853 let layout = layout_cache.remove(&target).ok_or_else(|| {
1854 Error::Msg(format!(
1855 "PipelinedExecutor::realize: target layout {:?} not populated after execution",
1856 target
1857 ))
1858 .bt()
1859 })?;
1860 {
1861 let _t = crate::compiled::op_profile::PhaseTimer::start("realize.drop_cache");
1862 drop(cache);
1863 drop(handles);
1864 drop(layout_cache);
1865 }
1866 Ok((storage, layout, produced_syms))
1867 }
1868
1869 pub fn realize_many(
1884 graph: Arc<RwLock<Graph>>,
1885 targets: &[NodeId],
1886 inputs: StorageCache,
1887 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
1888 Self::realize_many_inner(
1889 graph,
1890 targets,
1891 inputs,
1892 OrderSource::Default,
1893 SymEnv::default(),
1894 )
1895 }
1896
1897 pub fn realize_many_with_optimized(
1903 graph: Arc<RwLock<Graph>>,
1904 targets: &[NodeId],
1905 inputs: StorageCache,
1906 optimized: &OptimizedGraph,
1907 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
1908 Self::realize_many_inner(
1909 graph,
1910 targets,
1911 inputs,
1912 OrderSource::Optimized {
1913 optimized,
1914 route: None,
1915 },
1916 SymEnv::default(),
1917 )
1918 }
1919
1920 pub fn realize_many_with_optimized_env(
1923 graph: Arc<RwLock<Graph>>,
1924 targets: &[NodeId],
1925 inputs: StorageCache,
1926 optimized: &OptimizedGraph,
1927 sym_env: SymEnv,
1928 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
1929 Self::realize_many_inner(
1930 graph,
1931 targets,
1932 inputs,
1933 OrderSource::Optimized {
1934 optimized,
1935 route: None,
1936 },
1937 sym_env,
1938 )
1939 }
1940
1941 pub fn realize_many_with_optimized_route(
1946 graph: Arc<RwLock<Graph>>,
1947 targets: &[NodeId],
1948 inputs: StorageCache,
1949 optimized: &OptimizedGraph,
1950 route: &PickedRoute,
1951 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
1952 Self::realize_many_inner(
1953 graph,
1954 targets,
1955 inputs,
1956 OrderSource::Optimized {
1957 optimized,
1958 route: Some(route),
1959 },
1960 SymEnv::default(),
1961 )
1962 }
1963
1964 pub fn realize_many_with_optimized_route_env(
1967 graph: Arc<RwLock<Graph>>,
1968 targets: &[NodeId],
1969 inputs: StorageCache,
1970 optimized: &OptimizedGraph,
1971 route: &PickedRoute,
1972 sym_env: SymEnv,
1973 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
1974 Self::realize_many_inner(
1975 graph,
1976 targets,
1977 inputs,
1978 OrderSource::Optimized {
1979 optimized,
1980 route: Some(route),
1981 },
1982 sym_env,
1983 )
1984 }
1985
1986 pub fn realize_many_with_optimized_picking_env(
1992 graph: Arc<RwLock<Graph>>,
1993 targets: &[NodeId],
1994 inputs: StorageCache,
1995 optimized: &OptimizedGraph,
1996 selector: Option<Arc<dyn RuntimeSelector>>,
1997 lookup: Option<BackendRuntimeLookup>,
1998 sym_env: SymEnv,
1999 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
2000 match Self::streaming_pick_for(&graph, selector, lookup)? {
2001 Some(pick) => Self::realize_many_inner(
2002 graph,
2003 targets,
2004 inputs,
2005 OrderSource::Streaming {
2006 optimized,
2007 pick: &pick,
2008 },
2009 sym_env,
2010 ),
2011 None => {
2012 Self::realize_many_with_optimized_env(graph, targets, inputs, optimized, sym_env)
2013 }
2014 }
2015 }
2016
2017 fn realize_many_inner(
2018 graph: Arc<RwLock<Graph>>,
2019 targets: &[NodeId],
2020 inputs: StorageCache,
2021 order_source: OrderSource<'_>,
2022 sym_env: SymEnv,
2023 ) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
2024 if targets.is_empty() {
2025 return Ok(Vec::new());
2026 }
2027
2028 {
2033 let mut g = graph.write().map_err(|_| poisoned("graph lock"))?;
2034 let effective_roots = extend_with_side_effect_roots(&g, targets);
2035 insert_safety_copies(&mut g, &effective_roots);
2036 }
2037
2038 let (compiler_work, wait_set, mut layout_cache): (
2045 CompilerWork,
2046 WaitSet,
2047 HashMap<NodeId, Layout>,
2048 ) = {
2049 let g = graph.read().map_err(|_| poisoned("graph lock"))?;
2050 let effective_roots = extend_with_side_effect_roots(&g, targets);
2051 let (work, wait_set) =
2061 compiler_work_and_wait_set_for(&g, &effective_roots, &order_source);
2062 let mut layouts = HashMap::with_capacity(inputs.len());
2063 for &id in inputs.keys() {
2064 layouts.insert(id, g.layout(id));
2065 }
2066 (work, wait_set, layouts)
2067 };
2068
2069 let target_set: std::collections::HashSet<NodeId> = targets.iter().copied().collect();
2072
2073 let (tx, rx) = channel::<Result<WorkItem>>();
2074 let graph_for_compiler = Arc::clone(&graph);
2075
2076 let compiler = thread::spawn(move || {
2077 compiler_thread_body(graph_for_compiler, compiler_work, sym_env, tx);
2078 });
2079
2080 let plan_generation: Option<u64> = generation_for(&order_source);
2084 let mut current_chunk_backend: Option<BackendId> = None;
2085 let mut cache: StorageCache = inputs;
2086 let mut produced_syms = SymEnv::new();
2091 let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
2095 let mut inflight_vulkan: Vec<InflightVulkan> = Vec::new();
2102 let mut multi_backend = false;
2103 for item in rx {
2104 let item = item?;
2105 if !is_backend_neutral(&item.kind) && current_chunk_backend != Some(item.target_backend)
2108 {
2109 if let Some(plan_gen) = plan_generation {
2110 let live = crate::dispatch::topology_generation();
2111 if live != plan_gen {
2112 return Err(Error::TopologyChanged {
2113 plan_generation: plan_gen,
2114 current_generation: live,
2115 }
2116 .bt());
2117 }
2118 }
2119 let leaving = current_chunk_backend;
2120 if let Some(prev) = leaving
2121 && prev != item.target_backend
2122 {
2123 multi_backend = true;
2124 if prev == BackendId::Vulkan {
2125 eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
2126 }
2127 }
2128 current_chunk_backend = Some(item.target_backend);
2129 }
2130 if item.target_backend == BackendId::Vulkan && matches!(item.kind, WorkItemKind::Kernel)
2134 {
2135 multi_backend = true;
2136 }
2137 if matches!(
2142 item.kind,
2143 WorkItemKind::Copy { .. } | WorkItemKind::Move { .. }
2144 ) {
2145 if multi_backend {
2148 eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
2149 }
2150 if let Some(&producer) = item.inputs.first() {
2151 if multi_backend && copy_source_is_vulkan(&cache, producer)? {
2155 drain_inflight_vulkan(&mut inflight_vulkan)?;
2156 }
2157 wait_producer_handle(&mut handles, producer)?;
2159 }
2160 }
2161 let handle = execute_work_item(
2162 &item,
2163 &mut cache,
2164 &mut layout_cache,
2165 &mut produced_syms,
2166 will_be_waited(&wait_set, item.node_id),
2167 None,
2170 )
2171 .map_err(|e| with_node_location(&graph, item.node_id, e))?;
2172 store_handle(&mut handles, item.node_id, handle);
2173 if let Some(d_idx) = item.destructive_input
2174 && let Some(&destroyed) = item.inputs.get(d_idx)
2175 && !target_set.contains(&destroyed)
2176 {
2177 if multi_backend {
2192 defer_evicted_vulkan_buffer(&cache, destroyed, &mut inflight_vulkan)?;
2193 } else if let Some(d_arc) = cache.get(&destroyed) {
2194 force_flush_vulkan(d_arc)?;
2195 }
2196 cache.remove(&destroyed);
2197 layout_cache.remove(&destroyed);
2198 handles.remove(&destroyed);
2201 }
2202 }
2203
2204 compiler
2205 .join()
2206 .map_err(|_| Error::Msg("compiler thread panicked".to_string()).bt())?;
2207
2208 drain_handles(&mut handles)?;
2211 sync_active_cuda_devices(&cache)?;
2215 drain_inflight_vulkan(&mut inflight_vulkan)?;
2218 drain_vulkan_pending(&cache)?;
2222 debug_assert!(
2223 inflight_vulkan.is_empty(),
2224 "A4b-4: in-flight Vulkan batch list must be empty after realize-end drain",
2225 );
2226
2227 let mut out = Vec::with_capacity(targets.len());
2232 for &target in targets {
2233 let storage = cache.get(&target).cloned().ok_or_else(|| {
2234 Error::Msg(format!(
2235 "PipelinedExecutor::realize_many: target {:?} not populated after execution",
2236 target
2237 ))
2238 .bt()
2239 })?;
2240 let layout = layout_cache.get(&target).cloned().ok_or_else(|| {
2241 Error::Msg(format!(
2242 "PipelinedExecutor::realize_many: target layout {:?} not populated after execution",
2243 target
2244 ))
2245 .bt()
2246 })?;
2247 out.push((storage, layout));
2248 }
2249 Ok(out)
2250 }
2251}
2252
2253fn order_for(
2261 graph: &Graph,
2262 effective_roots: &[NodeId],
2263 order_source: &OrderSource<'_>,
2264) -> Vec<NodeId> {
2265 match order_source {
2266 OrderSource::Default => execution_plan(graph, effective_roots),
2267 OrderSource::Streaming {
2275 optimized: _optimized,
2276 pick,
2277 } => {
2278 let bindings = global_bindings();
2279 fuel_graph::lower_picked_route_streaming(graph, effective_roots, |branch| {
2280 resolve_branch(graph, branch, &bindings, pick.selector.as_ref())
2281 })
2282 }
2283 OrderSource::Optimized { optimized, route } => {
2284 match route {
2296 Some(route) => fuel_graph::lower_picked_route(graph, effective_roots, route),
2297 None => {
2298 let view = OptimizedGraph {
2302 roots: effective_roots.to_vec(),
2303 generation: optimized.generation,
2304 placements: None,
2305 };
2306 view.dispatch_order(graph)
2307 }
2308 }
2309 }
2310 OrderSource::Cached { .. } => {
2315 unreachable!("OrderSource::Cached is resolved before order_for is called")
2316 }
2317 }
2318}
2319
2320fn build_wait_set(graph: &Graph, order: &[NodeId]) -> HashSet<NodeId> {
2336 let mut set = HashSet::new();
2337 for &id in order {
2338 let node = graph.node(id);
2339 if matches!(node.op, Op::Copy { .. } | Op::Move { .. })
2340 && let Some(&producer) = node.inputs.first()
2341 {
2342 set.insert(producer);
2343 }
2344 }
2345 set
2346}
2347
2348type WaitSet = Option<HashSet<NodeId>>;
2354
2355#[inline]
2359fn will_be_waited(wait_set: &WaitSet, node_id: NodeId) -> bool {
2360 match wait_set {
2361 None => true,
2362 Some(set) => set.contains(&node_id),
2363 }
2364}
2365
2366fn compiler_work_and_wait_set_for(
2380 graph: &Graph,
2381 effective_roots: &[NodeId],
2382 order_source: &OrderSource<'_>,
2383) -> (CompilerWork, WaitSet) {
2384 match order_source {
2385 OrderSource::Streaming { pick, .. } => (
2386 CompilerWork::Streaming {
2387 roots: effective_roots.to_vec(),
2388 pick: (*pick).clone(),
2389 },
2390 None,
2391 ),
2392 OrderSource::Cached {
2393 optimized,
2394 order_cache,
2395 } => {
2396 let order = order_cache.get_or_init(|| {
2397 order_for(
2398 graph,
2399 effective_roots,
2400 &OrderSource::Optimized {
2401 optimized,
2402 route: None,
2403 },
2404 )
2405 });
2406 let wait_set = build_wait_set(graph, order);
2407 (CompilerWork::Order(order.clone()), Some(wait_set))
2408 }
2409 other => {
2410 let order = order_for(graph, effective_roots, other);
2411 let wait_set = build_wait_set(graph, &order);
2412 (CompilerWork::Order(order), Some(wait_set))
2413 }
2414 }
2415}
2416
2417fn generation_for(order_source: &OrderSource<'_>) -> Option<u64> {
2422 match order_source {
2423 OrderSource::Optimized { optimized, .. } => Some(optimized.generation),
2424 OrderSource::Streaming { optimized, .. } => Some(optimized.generation),
2425 OrderSource::Cached { optimized, .. } => Some(optimized.generation),
2426 OrderSource::Default => None,
2427 }
2428}
2429
2430enum CompilerWork {
2435 Order(Vec<NodeId>),
2438 Streaming {
2446 roots: Vec<NodeId>,
2447 pick: StreamingPick,
2448 },
2449}
2450
2451fn compiler_thread_body(
2458 graph: Arc<RwLock<Graph>>,
2459 work: CompilerWork,
2460 sym_env: SymEnv,
2461 tx: Sender<Result<WorkItem>>,
2462) {
2463 let bindings = global_bindings();
2464
2465 let g = match graph.read() {
2466 Ok(g) => g,
2467 Err(_) => {
2468 let _ = tx.send(Err(poisoned("graph lock in compiler")));
2469 return;
2470 }
2471 };
2472
2473 let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
2479 let compile_and_send = |id: NodeId, layout_cache: &mut HashMap<NodeId, Layout>| -> bool {
2483 let item = compile_one(&g, id, layout_cache, &bindings, &sym_env);
2484 let stop_on_err = item.is_err();
2485 if tx.send(item).is_err() {
2486 return false;
2487 }
2488 !stop_on_err
2489 };
2490
2491 match work {
2492 CompilerWork::Order(order) => {
2493 for id in order {
2494 if !compile_and_send(id, &mut layout_cache) {
2495 return;
2496 }
2497 }
2498 }
2499 CompilerWork::Streaming { roots, pick } => {
2500 let mut stopped = false;
2512 fuel_graph::walk_picked_route_streaming(
2513 &g,
2514 &roots,
2515 |branch| resolve_branch(&g, branch, &bindings, pick.selector.as_ref()),
2516 |members| {
2517 if stopped {
2518 return;
2519 }
2520 for &id in members {
2521 if !compile_and_send(id, &mut layout_cache) {
2522 stopped = true;
2523 return;
2524 }
2525 }
2526 },
2527 );
2528 }
2529 }
2530}
2531
2532fn resolve_compiled(
2541 op_kind: OpKind,
2542 dtypes: &[DType],
2543 target_backend: BackendId,
2544 op_params: OpParams,
2545 bindings: &KernelBindingTable,
2546) -> Result<CompiledNode> {
2547 compile_node(op_kind, dtypes, target_backend, op_params, bindings)
2548}
2549
2550fn compile_one(
2565 graph: &Graph,
2566 id: NodeId,
2567 layout_cache: &mut HashMap<NodeId, Layout>,
2568 bindings: &KernelBindingTable,
2569 sym_env: &SymEnv,
2570) -> Result<WorkItem> {
2571 let node = graph.node(id);
2572 let elem_count = node.shape.elem_count();
2573 let inputs = node.inputs.clone();
2574 let destructive_input = node.op.destructive_input();
2578
2579 if matches!(node.op, Op::Const) {
2580 let output_layout = graph.layout(id);
2581 layout_cache.insert(id, output_layout.clone());
2582 return Ok(WorkItem {
2583 node_id: id,
2584 inputs,
2585 elem_count,
2586 dtype: node.dtype,
2587 target_backend: BackendId::Cpu,
2588 kind: WorkItemKind::ConstAdopt,
2589 compiled: None,
2590 output_layout,
2591 destructive_input,
2592 output_bundle: None,
2593 });
2594 }
2595
2596 if node.op.is_view_op() {
2597 if inputs.len() != 1 {
2598 return Err(Error::Msg(format!(
2599 "view op {:?} expects 1 input, got {}",
2600 node.op,
2601 inputs.len(),
2602 ))
2603 .bt());
2604 }
2605 let output_layout = graph.layout(id);
2611 layout_cache.insert(id, output_layout.clone());
2612 let target_backend = graph
2617 .target_backend(id)
2618 .or_else(|| graph.target_backend(inputs[0]))
2619 .unwrap_or(BackendId::Cpu);
2620 return Ok(WorkItem {
2621 node_id: id,
2622 inputs,
2623 elem_count,
2624 dtype: node.dtype,
2625 target_backend,
2626 kind: WorkItemKind::ViewOf {
2627 input: node.inputs[0],
2628 },
2629 compiled: None,
2630 output_layout,
2631 destructive_input,
2632 output_bundle: None,
2633 });
2634 }
2635
2636 if let Op::View { .. } = node.op {
2644 if inputs.len() != 1 {
2645 return Err(Error::Msg(format!(
2646 "Op::View expects 1 input (the producer), got {}",
2647 inputs.len(),
2648 ))
2649 .bt());
2650 }
2651 let output_layout = graph.layout(id);
2652 layout_cache.insert(id, output_layout.clone());
2653 let target_backend = graph
2654 .target_backend(id)
2655 .or_else(|| graph.target_backend(inputs[0]))
2656 .unwrap_or(BackendId::Cpu);
2657 return Ok(WorkItem {
2658 node_id: id,
2659 inputs: inputs.clone(),
2660 elem_count,
2661 dtype: node.dtype,
2662 target_backend,
2663 kind: WorkItemKind::SlotView {
2664 producer: inputs[0],
2665 },
2666 compiled: None,
2667 output_layout,
2668 destructive_input,
2669 output_bundle: None,
2670 });
2671 }
2672
2673 if let Op::ViewOwned { slot } = node.op {
2678 if inputs.len() != 1 {
2679 return Err(Error::Msg(format!(
2680 "Op::ViewOwned expects 1 input (the producer), got {}",
2681 inputs.len(),
2682 ))
2683 .bt());
2684 }
2685 let output_layout = Layout::contiguous(node.shape.clone());
2686 layout_cache.insert(id, output_layout.clone());
2687 let target_backend = graph
2688 .target_backend(id)
2689 .or_else(|| graph.target_backend(inputs[0]))
2690 .unwrap_or(BackendId::Cpu);
2691 return Ok(WorkItem {
2692 node_id: id,
2693 inputs: inputs.clone(),
2694 elem_count,
2695 dtype: node.dtype,
2696 target_backend,
2697 kind: WorkItemKind::SlotOwn {
2698 producer: inputs[0],
2699 slot,
2700 },
2701 compiled: None,
2702 output_layout,
2703 destructive_input,
2704 output_bundle: None,
2705 });
2706 }
2707
2708 if matches!(node.op, Op::Release) {
2709 if inputs.len() != 1 {
2715 return Err(
2716 Error::Msg(format!("Op::Release expects 1 input, got {}", inputs.len(),)).bt(),
2717 );
2718 }
2719 let output_layout = Layout::contiguous(node.shape.clone());
2720 layout_cache.insert(id, output_layout.clone());
2721 let target_backend = graph
2725 .target_backend(id)
2726 .or_else(|| graph.target_backend(inputs[0]))
2727 .unwrap_or(BackendId::Cpu);
2728 return Ok(WorkItem {
2729 node_id: id,
2730 inputs,
2731 elem_count,
2732 dtype: node.dtype,
2733 target_backend,
2734 kind: WorkItemKind::ReleaseMarker,
2735 compiled: None,
2736 output_layout,
2737 destructive_input,
2738 output_bundle: None,
2739 });
2740 }
2741
2742 if !matches!(
2750 node.op,
2751 Op::WriteSlice { .. }
2752 | Op::WriteSliceRotating { .. }
2753 | Op::WriteSliceDoff { .. }
2754 | Op::ZeroFill
2755 | Op::Release
2756 | Op::Move { .. },
2757 ) && let Some(target_idx) = node.op.destructive_input()
2758 {
2759 if inputs.len() <= target_idx {
2760 return Err(Error::Msg(format!(
2761 "in-place op {:?} declares destructive_input={target_idx} \
2762 but has only {} input(s)",
2763 node.op,
2764 inputs.len(),
2765 ))
2766 .bt());
2767 }
2768 let target_backend = graph.target_backend(id).ok_or_else(|| {
2769 Error::Msg(format!(
2770 "PipelinedExecutor: in-place op node {id:?} has no target_backend set",
2771 ))
2772 .bt()
2773 })?;
2774 let op_kind = op_to_op_kind(&node.op).ok_or_else(|| {
2775 Error::Msg(format!(
2776 "PipelinedExecutor: in-place op {:?} has no op_to_op_kind mapping",
2777 node.op,
2778 ))
2779 .bt()
2780 })?;
2781 let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
2782 let dtypes = build_lookup_dtypes(graph, node);
2783 let compiled = resolve_compiled(op_kind, &dtypes, target_backend, op_params, bindings)?;
2784 let output_layout = graph.layout(inputs[target_idx]);
2786 layout_cache.insert(id, output_layout.clone());
2787 return Ok(WorkItem {
2788 node_id: id,
2789 inputs: inputs.clone(),
2790 elem_count,
2791 dtype: node.dtype,
2792 target_backend,
2793 kind: WorkItemKind::InplaceKernel { target_idx },
2794 compiled: Some(compiled),
2795 output_layout,
2796 destructive_input,
2797 output_bundle: None,
2798 });
2799 }
2800
2801 if matches!(node.op, Op::WriteSlice { .. }) {
2802 if inputs.len() != 2 {
2808 return Err(Error::Msg(format!(
2809 "Op::WriteSlice expects 2 inputs (destination, source), got {}",
2810 inputs.len(),
2811 ))
2812 .bt());
2813 }
2814 let target_backend = graph.target_backend(id).ok_or_else(|| {
2815 Error::Msg(format!(
2816 "PipelinedExecutor: WriteSlice node {:?} has no target_backend set",
2817 id
2818 ))
2819 .bt()
2820 })?;
2821 let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
2822 let dtypes = build_lookup_dtypes(graph, node);
2823 let compiled = resolve_compiled(
2824 OpKind::WriteSlice,
2825 &dtypes,
2826 target_backend,
2827 op_params,
2828 bindings,
2829 )?;
2830 let output_layout = graph.layout(inputs[0]);
2835 layout_cache.insert(id, output_layout.clone());
2836 return Ok(WorkItem {
2837 node_id: id,
2838 inputs: inputs.clone(),
2839 elem_count,
2840 dtype: node.dtype,
2841 target_backend,
2842 kind: WorkItemKind::WriteSlice {
2843 dest: inputs[0],
2844 source: inputs[1],
2845 },
2846 compiled: Some(compiled),
2847 output_layout,
2848 destructive_input,
2849 output_bundle: None,
2850 });
2851 }
2852
2853 if matches!(node.op, Op::WriteSliceRotating { .. }) {
2854 if inputs.len() != 3 {
2860 return Err(Error::Msg(format!(
2861 "Op::WriteSliceRotating expects 3 inputs (destination, source, position), got {}",
2862 inputs.len(),
2863 ))
2864 .bt());
2865 }
2866 let target_backend = graph.target_backend(id).ok_or_else(|| {
2867 Error::Msg(format!(
2868 "PipelinedExecutor: WriteSliceRotating node {:?} has no target_backend set",
2869 id
2870 ))
2871 .bt()
2872 })?;
2873 let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
2874 let dtypes = build_lookup_dtypes(graph, node);
2875 let compiled = resolve_compiled(
2876 OpKind::WriteSliceRotating,
2877 &dtypes,
2878 target_backend,
2879 op_params,
2880 bindings,
2881 )?;
2882 let output_layout = graph.layout(inputs[0]);
2883 layout_cache.insert(id, output_layout.clone());
2884 return Ok(WorkItem {
2885 node_id: id,
2886 inputs: inputs.clone(),
2887 elem_count,
2888 dtype: node.dtype,
2889 target_backend,
2890 kind: WorkItemKind::WriteSliceRotating {
2891 dest: inputs[0],
2892 source: inputs[1],
2893 position: inputs[2],
2894 },
2895 compiled: Some(compiled),
2896 output_layout,
2897 destructive_input,
2898 output_bundle: None,
2899 });
2900 }
2901
2902 if matches!(node.op, Op::WriteSliceDoff { .. }) {
2903 if inputs.len() != 3 {
2909 return Err(Error::Msg(format!(
2910 "Op::WriteSliceDoff expects 3 inputs (destination, source, offset), got {}",
2911 inputs.len(),
2912 ))
2913 .bt());
2914 }
2915 let target_backend = graph.target_backend(id).ok_or_else(|| {
2916 Error::Msg(format!(
2917 "PipelinedExecutor: WriteSliceDoff node {:?} has no target_backend set",
2918 id
2919 ))
2920 .bt()
2921 })?;
2922 let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
2923 let dtypes = build_lookup_dtypes(graph, node);
2924 let compiled = resolve_compiled(
2925 OpKind::WriteSliceDoff,
2926 &dtypes,
2927 target_backend,
2928 op_params,
2929 bindings,
2930 )?;
2931 let output_layout = graph.layout(inputs[0]);
2932 layout_cache.insert(id, output_layout.clone());
2933 return Ok(WorkItem {
2934 node_id: id,
2935 inputs: inputs.clone(),
2936 elem_count,
2937 dtype: node.dtype,
2938 target_backend,
2939 kind: WorkItemKind::WriteSliceDoff {
2940 dest: inputs[0],
2941 source: inputs[1],
2942 offset: inputs[2],
2943 },
2944 compiled: Some(compiled),
2945 output_layout,
2946 destructive_input,
2947 output_bundle: None,
2948 });
2949 }
2950
2951 if let Op::Alloc { target } = node.op {
2952 if !inputs.is_empty() {
2958 return Err(
2959 Error::Msg(format!("Op::Alloc expects 0 inputs, got {}", inputs.len(),)).bt(),
2960 );
2961 }
2962 let output_layout = Layout::contiguous(node.shape.clone());
2963 layout_cache.insert(id, output_layout.clone());
2964 let target_backend = match target {
2969 DeviceLocation::Cpu => BackendId::Cpu,
2970 DeviceLocation::Cuda { .. } => BackendId::Cuda,
2971 DeviceLocation::Vulkan { .. } => BackendId::Vulkan,
2972 DeviceLocation::Metal { .. } => BackendId::Metal,
2973 };
2974 return Ok(WorkItem {
2975 node_id: id,
2976 inputs,
2977 elem_count,
2978 dtype: node.dtype,
2979 target_backend,
2980 kind: WorkItemKind::Alloc {
2981 target_location: target,
2982 },
2983 compiled: None,
2984 output_layout,
2985 destructive_input,
2986 output_bundle: None,
2987 });
2988 }
2989
2990 if let Op::Iota { len } = node.op {
2991 if !inputs.is_empty() {
2996 return Err(
2997 Error::Msg(format!("Op::Iota expects 0 inputs, got {}", inputs.len(),)).bt(),
2998 );
2999 }
3000 let output_layout = Layout::contiguous(node.shape.clone());
3001 layout_cache.insert(id, output_layout.clone());
3002 let target_backend = graph.target_backend(id).unwrap_or(BackendId::Cpu);
3006 return Ok(WorkItem {
3007 node_id: id,
3008 inputs,
3009 elem_count,
3010 dtype: node.dtype,
3011 target_backend,
3012 kind: WorkItemKind::Iota { len },
3013 compiled: None,
3014 output_layout,
3015 destructive_input,
3016 output_bundle: None,
3017 });
3018 }
3019
3020 if matches!(node.op, Op::ZeroFill) {
3021 if inputs.len() != 1 {
3025 return Err(Error::Msg(format!(
3026 "Op::ZeroFill expects 1 input, got {}",
3027 inputs.len(),
3028 ))
3029 .bt());
3030 }
3031 let output_layout = Layout::contiguous(node.shape.clone());
3037 layout_cache.insert(id, output_layout.clone());
3038 let target_backend = graph
3039 .target_backend(id)
3040 .or_else(|| graph.target_backend(inputs[0]))
3041 .unwrap_or(BackendId::Cpu);
3042 return Ok(WorkItem {
3043 node_id: id,
3044 inputs,
3045 elem_count,
3046 dtype: node.dtype,
3047 target_backend,
3048 kind: WorkItemKind::ZeroFill,
3049 compiled: None,
3050 output_layout,
3051 destructive_input,
3052 output_bundle: None,
3053 });
3054 }
3055
3056 if let Op::Copy { target } = node.op {
3057 if inputs.len() != 1 {
3066 return Err(
3067 Error::Msg(format!("Op::Copy expects 1 input, got {}", inputs.len(),)).bt(),
3068 );
3069 }
3070 let target_backend = graph.target_backend(id).ok_or_else(|| {
3071 Error::Msg(format!(
3072 "PipelinedExecutor: Op::Copy node {:?} has no target_backend \
3073 set (= source backend, the one whose kernel runs the \
3074 download)",
3075 id
3076 ))
3077 .bt()
3078 })?;
3079 let op_params = OpParams::None;
3080 let dtypes = build_lookup_dtypes(graph, node);
3081 let compiled =
3082 resolve_compiled(OpKind::Copy, &dtypes, target_backend, op_params, bindings)?;
3083 let output_layout = Layout::contiguous(node.shape.clone());
3089 layout_cache.insert(id, output_layout.clone());
3090 return Ok(WorkItem {
3091 node_id: id,
3092 inputs,
3093 elem_count,
3094 dtype: node.dtype,
3095 target_backend,
3096 kind: WorkItemKind::Copy {
3097 target_location: target,
3098 },
3099 compiled: Some(compiled),
3100 output_layout,
3101 destructive_input,
3102 output_bundle: None,
3103 });
3104 }
3105
3106 if let Op::Move { target } = node.op {
3107 if inputs.len() != 1 {
3119 return Err(
3120 Error::Msg(format!("Op::Move expects 1 input, got {}", inputs.len(),)).bt(),
3121 );
3122 }
3123 let target_backend = graph.target_backend(id).ok_or_else(|| {
3124 Error::Msg(format!(
3125 "PipelinedExecutor: Op::Move node {:?} has no target_backend \
3126 set (= source backend, the one whose kernel runs the \
3127 transfer)",
3128 id
3129 ))
3130 .bt()
3131 })?;
3132 let op_params = OpParams::None;
3133 let dtypes = build_lookup_dtypes(graph, node);
3134 let compiled =
3135 resolve_compiled(OpKind::Copy, &dtypes, target_backend, op_params, bindings)?;
3136 let output_layout = Layout::contiguous(node.shape.clone());
3140 layout_cache.insert(id, output_layout.clone());
3141 return Ok(WorkItem {
3142 node_id: id,
3143 inputs,
3144 elem_count,
3145 dtype: node.dtype,
3146 target_backend,
3147 kind: WorkItemKind::Move {
3148 target_location: target,
3149 },
3150 compiled: Some(compiled),
3151 output_layout,
3152 destructive_input,
3153 output_bundle: None,
3154 });
3155 }
3156
3157 if matches!(node.op, Op::Reshape(_)) {
3158 if inputs.len() != 1 {
3159 return Err(
3160 Error::Msg(format!("Op::Reshape expects 1 input, got {}", inputs.len(),)).bt(),
3161 );
3162 }
3163 let input_layout = graph.layout(inputs[0]);
3164 let output_layout = Layout::contiguous(node.shape.clone());
3169 layout_cache.insert(id, output_layout.clone());
3170 let target_backend = graph
3171 .target_backend(id)
3172 .or_else(|| graph.target_backend(inputs[0]))
3173 .unwrap_or(BackendId::Cpu);
3174 let in_elem_count = input_layout.shape().elem_count();
3176 if in_elem_count != elem_count {
3177 return Err(Error::Msg(format!(
3178 "Op::Reshape changes element count: input {} → output {}",
3179 in_elem_count, elem_count,
3180 ))
3181 .bt());
3182 }
3183 return Ok(WorkItem {
3184 node_id: id,
3185 inputs,
3186 elem_count,
3187 dtype: node.dtype,
3188 target_backend,
3189 kind: WorkItemKind::ContiguizeOf {
3190 input: node.inputs[0],
3191 },
3192 compiled: None,
3193 output_layout,
3194 destructive_input,
3195 output_bundle: None,
3196 });
3197 }
3198
3199 if matches!(node.op, Op::Contiguize) {
3200 if inputs.len() != 1 {
3207 return Err(Error::Msg(format!(
3208 "Op::Contiguize expects 1 input, got {}",
3209 inputs.len(),
3210 ))
3211 .bt());
3212 }
3213 let output_layout = Layout::contiguous(node.shape.clone());
3214 layout_cache.insert(id, output_layout.clone());
3215 let target_backend = graph
3216 .target_backend(id)
3217 .or_else(|| graph.target_backend(inputs[0]))
3218 .unwrap_or(BackendId::Cpu);
3219 return Ok(WorkItem {
3220 node_id: id,
3221 inputs,
3222 elem_count,
3223 dtype: node.dtype,
3224 target_backend,
3225 kind: WorkItemKind::ContiguizeOf {
3226 input: node.inputs[0],
3227 },
3228 compiled: None,
3229 output_layout,
3230 destructive_input,
3231 output_bundle: None,
3232 });
3233 }
3234
3235 if let Op::Fused(fid, _) = &node.op
3245 && fid.is_runtime()
3246 {
3247 let target_backend = graph.target_backend(id).ok_or_else(|| {
3248 Error::Msg(format!(
3249 "PipelinedExecutor: runtime fused op {:?} (node {:?}) has no \
3250 target_backend — the arm emitter pins it at emission",
3251 fid, id,
3252 ))
3253 .bt()
3254 })?;
3255 if graph.output_views_arc(id).is_some() {
3258 return Err(Error::Msg(format!(
3259 "PipelinedExecutor: runtime fused op {:?} (node {:?}) carries a \
3260 multi-output bundle — synthesized kernels are single-output in v1",
3261 fid, id,
3262 ))
3263 .bt());
3264 }
3265 let dtypes = build_lookup_dtypes(graph, node);
3266 let (kernel, caps) = bindings.lookup_with_caps(
3267 crate::kernel::BindingKey::RuntimeFused(*fid),
3268 &dtypes,
3269 target_backend,
3270 )?;
3271 let output_layout = Layout::contiguous(node.shape.clone());
3272 layout_cache.insert(id, output_layout.clone());
3273 let op_params = match &node.op {
3280 Op::Fused(_, FusedOpParams::Runtime { scalars }) if !scalars.is_empty() => {
3281 OpParams::JitScalars {
3282 scalars: scalars.clone(),
3283 }
3284 }
3285 _ => OpParams::None,
3286 };
3287 let compiled = CompiledNode {
3288 op: OpKind::RuntimeFused,
3289 dtypes: KernelDTypes::from_slice(&dtypes),
3290 backend: target_backend,
3291 kernel,
3292 caps,
3293 op_params,
3294 };
3295 return Ok(WorkItem {
3296 node_id: id,
3297 inputs,
3298 elem_count,
3299 dtype: node.dtype,
3300 target_backend,
3301 kind: WorkItemKind::Kernel,
3302 compiled: Some(compiled),
3303 output_layout,
3304 destructive_input,
3305 output_bundle: None,
3306 });
3307 }
3308
3309 let target_backend = graph.target_backend(id).ok_or_else(|| {
3310 Error::Msg(format!(
3311 "PipelinedExecutor: node {:?} ({:?}) has no target_backend set",
3312 id, node.op
3313 ))
3314 .bt()
3315 })?;
3316
3317 let op_kind = op_to_op_kind(&node.op).ok_or_else(|| {
3318 Error::Msg(format!(
3319 "PipelinedExecutor: op {:?} not yet mapped to an OpKind \
3320 (Phase C migrates more ops as they're registered)",
3321 node.op,
3322 ))
3323 .bt()
3324 })?;
3325
3326 let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
3327 let dtypes = build_lookup_dtypes(graph, node);
3332 let compiled = resolve_compiled(op_kind, &dtypes, target_backend, op_params, bindings)?;
3333 let output_layout = Layout::contiguous(node.shape.clone());
3334 layout_cache.insert(id, output_layout.clone());
3335 let output_bundle = graph.output_views_arc(id);
3341 Ok(WorkItem {
3342 node_id: id,
3343 inputs,
3344 elem_count,
3345 dtype: node.dtype,
3346 target_backend,
3347 kind: WorkItemKind::Kernel,
3348 compiled: Some(compiled),
3349 output_layout,
3350 destructive_input,
3351 output_bundle,
3352 })
3353}
3354
3355pub(crate) fn build_lookup_dtypes(graph: &Graph, node: &Node) -> Vec<DType> {
3365 if matches!(node.op, Op::Concat { .. }) {
3366 let in_dt = node
3368 .inputs
3369 .first()
3370 .map(|&id| graph.node(id).dtype)
3371 .unwrap_or(node.dtype);
3372 return vec![in_dt, node.dtype];
3373 }
3374 if matches!(node.op, Op::WriteSlice { .. }) {
3375 let src_dt = node
3381 .inputs
3382 .get(1)
3383 .map(|&id| graph.node(id).dtype)
3384 .unwrap_or(node.dtype);
3385 return vec![src_dt, node.dtype];
3386 }
3387 if matches!(node.op, Op::WriteSliceRotating { .. }) {
3388 let src_dt = node
3392 .inputs
3393 .get(1)
3394 .map(|&id| graph.node(id).dtype)
3395 .unwrap_or(node.dtype);
3396 return vec![src_dt, node.dtype];
3397 }
3398 if matches!(node.op, Op::WriteSliceDoff { .. }) {
3399 let src_dt = node
3403 .inputs
3404 .get(1)
3405 .map(|&id| graph.node(id).dtype)
3406 .unwrap_or(node.dtype);
3407 return vec![src_dt, node.dtype];
3408 }
3409 let mut dts: Vec<DType> = node.inputs.iter().map(|&id| graph.node(id).dtype).collect();
3410 dts.push(node.dtype);
3411 dts
3412}
3413
3414fn op_is_softmax_last_dim(op: &Op) -> bool {
3420 matches!(op, Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM)
3421}
3422
3423fn op_is_fused_linear(op: &Op) -> bool {
3426 matches!(op, Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::FUSED_LINEAR)
3427}
3428
3429pub(crate) fn op_to_op_kind(op: &Op) -> Option<OpKind> {
3437 match op {
3438 Op::Add => Some(OpKind::AddElementwise),
3439 Op::Sub => Some(OpKind::SubElementwise),
3440 Op::Mul => Some(OpKind::MulElementwise),
3441 Op::Div => Some(OpKind::DivElementwise),
3442 Op::Relu => Some(OpKind::ReluElementwise),
3443 Op::Neg => Some(OpKind::NegElementwise),
3444 Op::Sqr => Some(OpKind::SqrElementwise),
3445 Op::Sqrt => Some(OpKind::SqrtElementwise),
3446 Op::Tanh => Some(OpKind::TanhElementwise),
3447 Op::Exp => Some(OpKind::ExpElementwise),
3448 Op::Log => Some(OpKind::LogElementwise),
3449 Op::Sin => Some(OpKind::SinElementwise),
3450 Op::Cos => Some(OpKind::CosElementwise),
3451 Op::Sigmoid => Some(OpKind::SigmoidElementwise),
3452 Op::Silu => Some(OpKind::SiluElementwise),
3453 Op::GeluTanh => Some(OpKind::GeluTanhElementwise),
3454 Op::Step => Some(OpKind::StepElementwise),
3455 Op::Recip => Some(OpKind::RecipElementwise),
3456 Op::Abs => Some(OpKind::AbsElementwise),
3457 Op::Equal => Some(OpKind::EqualElementwise),
3458 Op::Ne => Some(OpKind::NotEqualElementwise),
3459 Op::Lt => Some(OpKind::LessElementwise),
3460 Op::Le => Some(OpKind::LessEqualElementwise),
3461 Op::Gt => Some(OpKind::GreaterElementwise),
3462 Op::Ge => Some(OpKind::GreaterEqualElementwise),
3463 Op::Where => Some(OpKind::Where),
3464 Op::Floor => Some(OpKind::FloorElementwise),
3465 Op::Ceil => Some(OpKind::CeilElementwise),
3466 Op::Round => Some(OpKind::RoundElementwise),
3467 Op::Sign => Some(OpKind::SignElementwise),
3468 Op::Erf => Some(OpKind::ErfElementwise),
3469 Op::Gelu => Some(OpKind::GeluElementwise),
3470 Op::Pow => Some(OpKind::PowElementwise),
3471 Op::Rsqrt => Some(OpKind::RsqrtElementwise),
3472 Op::Rem => Some(OpKind::RemElementwise),
3473 Op::Flip { .. } => Some(OpKind::Flip),
3474 Op::Roll { .. } => Some(OpKind::Roll),
3475 Op::CumSum { .. } => Some(OpKind::CumSum),
3476 Op::Triu { .. } => Some(OpKind::Triu),
3477 Op::Tril { .. } => Some(OpKind::Tril),
3478 Op::LogSoftmaxLastDim => Some(OpKind::LogSoftmaxLastDim),
3479 Op::LogSoftmaxLastDimBackward => Some(OpKind::LogSoftmaxLastDimBackward),
3480 Op::MaskedFill { .. } => Some(OpKind::MaskedFill),
3481 Op::Pad { .. } => Some(OpKind::Pad),
3482 Op::PadBackward { .. } => Some(OpKind::PadBackward),
3483 Op::SumDim(_) => Some(OpKind::SumReduce),
3484 Op::MaxDim(_) => Some(OpKind::MaxReduce),
3485 Op::MinDim(_) => Some(OpKind::MinReduce),
3486 Op::MeanDim(_) => Some(OpKind::MeanReduce),
3487 Op::SumAll => Some(OpKind::SumReduce),
3488 Op::MaxAll => Some(OpKind::MaxReduce),
3489 Op::MinAll => Some(OpKind::MinReduce),
3490 Op::MeanAll => Some(OpKind::MeanReduce),
3491 Op::MatMul => Some(OpKind::MatMul),
3492 Op::Cast(_) => Some(OpKind::Cast),
3493 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::CONV2D => Some(OpKind::Conv2D),
3499 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D => {
3503 Some(OpKind::ConvTranspose2D)
3504 }
3505 Op::ReduceSumTo(_) => Some(OpKind::ReduceSumTo),
3506 Op::ReduceMaxTo(_) => Some(OpKind::ReduceMaxTo),
3507 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::FLASH_ATTN => {
3508 Some(OpKind::FlashAttn)
3509 }
3510 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::PAGED_ATTN => {
3511 Some(OpKind::PagedAttn)
3512 }
3513 Op::AddScalar(_) => Some(OpKind::Affine),
3514 Op::MulScalar(_) => Some(OpKind::Affine),
3515 Op::Clamp { .. } => Some(OpKind::ClampElementwise),
3516 Op::PowI(_) => Some(OpKind::PowIElementwise),
3517 Op::Maximum => Some(OpKind::MaximumElementwise),
3518 Op::Minimum => Some(OpKind::MinimumElementwise),
3519 Op::Concat { .. } => Some(OpKind::Concat),
3520 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM => {
3521 Some(OpKind::SoftmaxLastDim)
3522 }
3523 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM_BACKWARD => {
3526 Some(OpKind::SoftmaxLastDimBackward)
3527 }
3528 Op::Fused(fid, _)
3529 if *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM_BACKWARD =>
3530 {
3531 Some(OpKind::LayerNormLastDimBackward)
3532 }
3533 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM_BACKWARD => {
3534 Some(OpKind::RmsNormLastDimBackward)
3535 }
3536 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::REDUCE_MAX_TO_BACKWARD => {
3537 Some(OpKind::ReduceMaxToBackward)
3538 }
3539 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::POWI_BACKWARD => {
3540 Some(OpKind::PowIElementwiseBackward)
3541 }
3542 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::FUSED_LINEAR => {
3543 Some(OpKind::FusedLinear)
3544 }
3545 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM => {
3546 Some(OpKind::RmsNormLastDim)
3547 }
3548 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM => {
3549 Some(OpKind::LayerNormLastDim)
3550 }
3551 Op::IndexSelect { .. } => Some(OpKind::IndexSelect),
3552 Op::Gather { .. } => Some(OpKind::Gather),
3553 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::ROPE => Some(OpKind::Rope),
3554 Op::IndexAdd { .. } => Some(OpKind::IndexAdd),
3555 Op::ScatterAdd { .. } => Some(OpKind::ScatterAdd),
3556 Op::NonZeroIndices { .. } => Some(OpKind::NonZeroIndices),
3557 Op::ArgMaxDim(_) => Some(OpKind::ArgMaxDim),
3558 Op::ArgMinDim(_) => Some(OpKind::ArgMinDim),
3559 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::QMATMUL => {
3561 Some(OpKind::QMatMul)
3562 }
3563 Op::WriteSlice { .. } => Some(OpKind::WriteSlice),
3564 Op::WriteSliceRotating { .. } => Some(OpKind::WriteSliceRotating),
3565 Op::WriteSliceDoff { .. } => Some(OpKind::WriteSliceDoff),
3566 Op::Copy { .. } => Some(OpKind::Copy),
3574 Op::Move { .. } => Some(OpKind::Copy),
3581 Op::ReluInplace => Some(OpKind::ReluInplace),
3586 Op::SiluInplace => Some(OpKind::SiluInplace),
3587 Op::GeluTanhInplace => Some(OpKind::GeluTanhInplace),
3588 Op::TanhInplace => Some(OpKind::TanhInplace),
3589 Op::SigmoidInplace => Some(OpKind::SigmoidInplace),
3590 Op::NegInplace => Some(OpKind::NegInplace),
3591 Op::AbsInplace => Some(OpKind::AbsInplace),
3592 Op::SqrInplace => Some(OpKind::SqrInplace),
3593 Op::SqrtInplace => Some(OpKind::SqrtInplace),
3594 Op::RsqrtInplace => Some(OpKind::RsqrtInplace),
3595 Op::RecipInplace => Some(OpKind::RecipInplace),
3596 Op::ExpInplace => Some(OpKind::ExpInplace),
3597 Op::LogInplace => Some(OpKind::LogInplace),
3598 Op::SinInplace => Some(OpKind::SinInplace),
3599 Op::CosInplace => Some(OpKind::CosInplace),
3600 Op::SignInplace => Some(OpKind::SignInplace),
3601 Op::FloorInplace => Some(OpKind::FloorInplace),
3602 Op::CeilInplace => Some(OpKind::CeilInplace),
3603 Op::RoundInplace => Some(OpKind::RoundInplace),
3604 Op::ErfInplace => Some(OpKind::ErfInplace),
3605 Op::GeluInplace => Some(OpKind::GeluInplace),
3606 Op::ClampInplace { .. } => Some(OpKind::ClampInplace),
3607 Op::PowIInplace(_) => Some(OpKind::PowIInplace),
3608 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::INPLACE_AFFINE => {
3614 Some(OpKind::InplaceAffine)
3615 }
3616 Op::Fused(fid, _)
3617 if *fid == fuel_graph::registry::FusedOps::FUSED_SOFTMAX_CROSS_ENTROPY =>
3618 {
3619 Some(OpKind::FusedSoftmaxCrossEntropy)
3620 }
3621 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::CAUSAL_CONV1D => {
3622 Some(OpKind::CausalConv1d)
3623 }
3624 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SELECTIVE_SCAN => {
3625 Some(OpKind::SelectiveScan)
3626 }
3627 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SSD_CHUNK_SCAN => {
3628 Some(OpKind::SsdChunkScan)
3629 }
3630 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::NF4_MATMUL => {
3631 Some(OpKind::Nf4Matmul)
3632 }
3633 Op::Alloc { .. } | Op::ZeroFill => None,
3638 _ => None,
3639 }
3640}
3641
3642fn encode_value_to_bytes(dtype: DType, value: f64) -> Result<Vec<u8>> {
3655 match dtype {
3656 DType::F32 => Ok((value as f32).to_le_bytes().to_vec()),
3657 DType::F64 => Ok(value.to_le_bytes().to_vec()),
3658 DType::BF16 => Ok(half::bf16::from_f32(value as f32).to_le_bytes().to_vec()),
3659 DType::F16 => Ok(half::f16::from_f32(value as f32).to_le_bytes().to_vec()),
3660 DType::U8 => Ok(vec![value as u8]),
3661 DType::U32 => Ok((value as u32).to_le_bytes().to_vec()),
3662 other => Err(Error::Msg(format!(
3663 "encode_value_to_bytes: dtype {other:?} not yet supported for Pad fill",
3664 ))
3665 .bt()),
3666 }
3667}
3668
3669fn scalar_to_bytes(s: fuel_ir::Scalar) -> Vec<u8> {
3673 use fuel_ir::Scalar;
3674 match s {
3675 Scalar::U8(v) => vec![v],
3676 Scalar::I8(v) => vec![v as u8],
3677 Scalar::U32(v) => v.to_le_bytes().to_vec(),
3678 Scalar::I16(v) => v.to_le_bytes().to_vec(),
3679 Scalar::I32(v) => v.to_le_bytes().to_vec(),
3680 Scalar::I64(v) => v.to_le_bytes().to_vec(),
3681 Scalar::BF16(v) => v.to_le_bytes().to_vec(),
3682 Scalar::F16(v) => v.to_le_bytes().to_vec(),
3683 Scalar::F32(v) => v.to_le_bytes().to_vec(),
3684 Scalar::F64(v) => v.to_le_bytes().to_vec(),
3685 Scalar::Bool(b) => vec![b as u8],
3686 Scalar::F8E4M3(v) => vec![v.to_bits()],
3687 Scalar::F8E5M2(v) => vec![v.to_bits()],
3688 Scalar::F8E8M0(v) => vec![v],
3691 }
3692}
3693
3694fn op_to_op_params(
3695 graph: &Graph,
3696 node: &Node,
3697 node_id: NodeId,
3698 layout_cache: &HashMap<NodeId, Layout>,
3699 sym_env: &SymEnv,
3700) -> Result<OpParams> {
3701 let input_layout = |input_id: NodeId| -> Layout {
3707 layout_cache
3708 .get(&input_id)
3709 .cloned()
3710 .unwrap_or_else(|| graph.layout(input_id))
3711 };
3712 Ok(match &node.op {
3713 Op::Fused(_, fuel_graph::registry::FusedOpParams::QMatMul { quant_type, k, n }) => {
3715 if node.inputs.len() != 2 {
3719 return Err(Error::Msg(format!(
3720 "Op::QMatMul expects 2 inputs (activations, weight_bytes), got {}",
3721 node.inputs.len(),
3722 ))
3723 .bt());
3724 }
3725 let act_layout = input_layout(node.inputs[0]);
3726 let act_dims = act_layout.shape().dims();
3727 if act_dims.len() < 2 {
3728 return Err(Error::Msg(format!(
3729 "Op::QMatMul: activations must be rank ≥ 2, got {act_dims:?}",
3730 ))
3731 .bt());
3732 }
3733 let m = act_dims[act_dims.len() - 2];
3734 let k_act = act_dims[act_dims.len() - 1];
3735 if k_act != *k {
3736 return Err(Error::Msg(format!(
3737 "Op::QMatMul: activation last dim ({k_act}) must equal Op's k ({k})",
3738 ))
3739 .bt());
3740 }
3741 let batch_count: usize = act_dims[..act_dims.len() - 2].iter().product();
3742 OpParams::QMatMul {
3743 quant_type: *quant_type,
3744 batch_count,
3745 m,
3746 n: *n,
3747 k: *k,
3748 }
3749 }
3750 Op::ArgMaxDim(d) | Op::ArgMinDim(d) => {
3751 OpParams::Reduce {
3755 dims: vec![*d],
3756 keepdim: false,
3757 }
3758 }
3759 Op::SumDim(d) | Op::MaxDim(d) | Op::MinDim(d) | Op::MeanDim(d) => OpParams::Reduce {
3760 dims: vec![*d],
3761 keepdim: false,
3762 },
3763 Op::SumAll | Op::MaxAll | Op::MinAll | Op::MeanAll => {
3764 let il = input_layout(node.inputs[0]);
3765 let rank = il.shape().rank();
3766 OpParams::Reduce {
3767 dims: (0..rank).collect(),
3768 keepdim: false,
3769 }
3770 }
3771 Op::MatMul => {
3772 if node.inputs.len() != 2 {
3778 return Err(Error::Msg(format!(
3779 "Op::MatMul expects 2 inputs, got {}",
3780 node.inputs.len(),
3781 ))
3782 .bt());
3783 }
3784 let lhs = input_layout(node.inputs[0]);
3785 let rhs = input_layout(node.inputs[1]);
3786 let lhs_dims = lhs.shape().dims();
3787 let rhs_dims = rhs.shape().dims();
3788 if lhs_dims.len() < 2 || rhs_dims.len() < 2 {
3789 return Err(Error::Msg(format!(
3790 "Op::MatMul requires both inputs rank ≥ 2; got lhs={:?} rhs={:?}",
3791 lhs_dims, rhs_dims,
3792 ))
3793 .bt());
3794 }
3795 if lhs_dims.len() != rhs_dims.len() {
3796 return Err(Error::Msg(format!(
3797 "Op::MatMul: ranks must match (auto-broadcast happens at \
3798 graph construction time); got lhs rank {} vs rhs rank {}",
3799 lhs_dims.len(),
3800 rhs_dims.len(),
3801 ))
3802 .bt());
3803 }
3804 let rank = lhs_dims.len();
3805 let batch_rank = rank - 2;
3806 for i in 0..batch_rank {
3808 let la = lhs_dims[i];
3809 let ra = rhs_dims[i];
3810 let ok = la == ra || (ra > 0 && la > ra && la % ra == 0);
3811 if !ok {
3812 return Err(Error::Msg(format!(
3813 "Op::MatMul: batch dim {i} disallowed combination \
3814 (lhs={la}, rhs={ra}); must be equal or \
3815 GQA-divisible (lhs > rhs && lhs % rhs == 0)",
3816 ))
3817 .bt());
3818 }
3819 }
3820 let lhs_batch_dims: Vec<usize> = lhs_dims[..batch_rank].to_vec();
3821 let rhs_batch_dims: Vec<usize> = rhs_dims[..batch_rank].to_vec();
3822 let (m, k_lhs) = (lhs_dims[rank - 2], lhs_dims[rank - 1]);
3823 let (k_rhs, n) = (rhs_dims[rank - 2], rhs_dims[rank - 1]);
3824 if k_lhs != k_rhs {
3825 return Err(Error::Msg(format!(
3826 "Op::MatMul: contracting dims disagree — lhs trailing is \
3827 [{m}, {k_lhs}], rhs trailing is [{k_rhs}, {n}]",
3828 ))
3829 .bt());
3830 }
3831 let m_compute = match graph.matmul_row_count(node_id) {
3837 None => MatmulM::All,
3838 Some(dyn_m) => match dyn_m.resolve(sym_env) {
3839 Some(c) => {
3840 if c > m {
3841 return Err(Error::Msg(format!(
3842 "Op::MatMul: resolved row count {c} exceeds capacity {m}",
3843 ))
3844 .bt());
3845 }
3846 MatmulM::Rows(c)
3847 }
3848 None => match dyn_m {
3849 fuel_ir::DynScalar::Sym(s) if graph.binds_data_determined_sym(s) => {
3850 MatmulM::Deferred(dyn_m)
3851 }
3852 _ => {
3853 return Err(Error::Msg(format!(
3854 "Op::MatMul: dynamic row count {dyn_m:?} is unbound in the \
3855 SymEnv at compile and no producer binds it (not \
3856 data-determined)",
3857 ))
3858 .bt());
3859 }
3860 },
3861 },
3862 };
3863 OpParams::Matmul {
3864 lhs_batch_dims,
3865 rhs_batch_dims,
3866 m,
3867 n,
3868 k: k_lhs,
3869 m_compute,
3870 }
3871 }
3872 op if op_is_fused_linear(op) => {
3876 if node.inputs.len() != 3 {
3881 return Err(Error::Msg(format!(
3882 "Op::FusedLinear expects 3 inputs (a, b, bias), got {}",
3883 node.inputs.len(),
3884 ))
3885 .bt());
3886 }
3887 let lhs = input_layout(node.inputs[0]);
3888 let rhs = input_layout(node.inputs[1]);
3889 let bias = input_layout(node.inputs[2]);
3890 let lhs_dims = lhs.shape().dims();
3891 let rhs_dims = rhs.shape().dims();
3892 let bias_dims = bias.shape().dims();
3893 if lhs_dims.len() < 2 || rhs_dims.len() < 2 {
3894 return Err(Error::Msg(format!(
3895 "Op::FusedLinear requires both a, b rank ≥ 2; got a={:?} b={:?}",
3896 lhs_dims, rhs_dims,
3897 ))
3898 .bt());
3899 }
3900 if lhs_dims.len() != rhs_dims.len() {
3901 return Err(Error::Msg(format!(
3902 "Op::FusedLinear: ranks must match (auto-broadcast happens at \
3903 graph construction time); got a rank {} vs b rank {}",
3904 lhs_dims.len(),
3905 rhs_dims.len(),
3906 ))
3907 .bt());
3908 }
3909 let rank = lhs_dims.len();
3910 let batch_rank = rank - 2;
3911 for i in 0..batch_rank {
3912 let la = lhs_dims[i];
3913 let ra = rhs_dims[i];
3914 let ok = la == ra || (ra > 0 && la > ra && la % ra == 0);
3915 if !ok {
3916 return Err(Error::Msg(format!(
3917 "Op::FusedLinear: batch dim {i} disallowed (a={la}, b={ra})",
3918 ))
3919 .bt());
3920 }
3921 }
3922 let lhs_batch_dims: Vec<usize> = lhs_dims[..batch_rank].to_vec();
3923 let rhs_batch_dims: Vec<usize> = rhs_dims[..batch_rank].to_vec();
3924 let (m, k_lhs) = (lhs_dims[rank - 2], lhs_dims[rank - 1]);
3925 let (k_rhs, n) = (rhs_dims[rank - 2], rhs_dims[rank - 1]);
3926 if k_lhs != k_rhs {
3927 return Err(Error::Msg(format!(
3928 "Op::FusedLinear: contracting dims disagree — a trailing is \
3929 [{m}, {k_lhs}], b trailing is [{k_rhs}, {n}]",
3930 ))
3931 .bt());
3932 }
3933 if bias_dims.len() != 1 || bias_dims[0] != n {
3934 return Err(Error::Msg(format!(
3935 "Op::FusedLinear: bias must be rank-1 [{n}], got {bias_dims:?}",
3936 ))
3937 .bt());
3938 }
3939 OpParams::Matmul {
3940 lhs_batch_dims,
3941 rhs_batch_dims,
3942 m,
3943 n,
3944 k: k_lhs,
3945 m_compute: MatmulM::All,
3947 }
3948 }
3949 Op::Fused(
3954 _,
3955 fuel_graph::registry::FusedOpParams::FusedSoftmaxCrossEntropy {
3956 reduction,
3957 ignore_index,
3958 },
3959 ) => {
3960 if node.inputs.len() != 2 {
3961 return Err(Error::Msg(format!(
3962 "FusedSoftmaxCrossEntropy expects 2 inputs (logits, targets), got {}",
3963 node.inputs.len(),
3964 ))
3965 .bt());
3966 }
3967 let logits_layout = input_layout(node.inputs[0]);
3968 let logits_dims = logits_layout.shape().dims();
3969 if logits_dims.is_empty() {
3970 return Err(Error::Msg(
3971 "FusedSoftmaxCrossEntropy: logits must have rank ≥ 1".to_string(),
3972 )
3973 .bt());
3974 }
3975 let vocab = *logits_dims.last().unwrap();
3976 let n_rows: usize = logits_dims[..logits_dims.len() - 1]
3977 .iter()
3978 .product::<usize>()
3979 .max(1);
3980 let targets_layout = input_layout(node.inputs[1]);
3981 let target_count: usize = targets_layout
3982 .shape()
3983 .dims()
3984 .iter()
3985 .product::<usize>()
3986 .max(1);
3987 if target_count != n_rows {
3988 return Err(Error::Msg(format!(
3989 "FusedSoftmaxCrossEntropy: targets element count {target_count} must \
3990 equal logits row count {n_rows} (logits {logits_dims:?}, targets {:?})",
3991 targets_layout.shape().dims(),
3992 ))
3993 .bt());
3994 }
3995 OpParams::FusedSoftmaxCrossEntropy {
3996 n_rows,
3997 vocab,
3998 reduction: *reduction,
3999 ignore_index: *ignore_index,
4000 }
4001 }
4002 Op::Fused(_, fuel_graph::registry::FusedOpParams::CausalConv1d { use_silu }) => {
4007 if node.inputs.len() != 3 {
4008 return Err(Error::Msg(format!(
4009 "CausalConv1d expects 3 inputs (x, weight, bias), got {}",
4010 node.inputs.len(),
4011 ))
4012 .bt());
4013 }
4014 let x_layout = input_layout(node.inputs[0]);
4015 let w_layout = input_layout(node.inputs[1]);
4016 let x_dims = x_layout.shape().dims();
4017 let w_dims = w_layout.shape().dims();
4018 if x_dims.len() != 3 {
4019 return Err(Error::Msg(format!(
4020 "CausalConv1d: x must be rank 3 [batch, channels, seq+pad], got {x_dims:?}",
4021 ))
4022 .bt());
4023 }
4024 if w_dims.len() != 3 {
4025 return Err(Error::Msg(format!(
4026 "CausalConv1d: weight must be rank 3 [channels, 1, kernel], got {w_dims:?}",
4027 ))
4028 .bt());
4029 }
4030 let batch = x_dims[0];
4031 let channels = x_dims[1];
4032 let seq_in = x_dims[2];
4033 let kernel = w_dims[2];
4034 if seq_in < kernel - 1 {
4035 return Err(Error::Msg(format!(
4036 "CausalConv1d: x time dim {seq_in} must be ≥ kernel-1 = {} \
4037 (caller must pre-pad with kernel-1 zeros)",
4038 kernel - 1,
4039 ))
4040 .bt());
4041 }
4042 let seq_out = seq_in - (kernel - 1);
4043 OpParams::CausalConv1d {
4044 batch,
4045 channels,
4046 seq_in,
4047 seq_out,
4048 kernel,
4049 use_silu: *use_silu,
4050 }
4051 }
4052 Op::Fused(_, fuel_graph::registry::FusedOpParams::SelectiveScan { delta_softplus }) => {
4055 if node.inputs.len() != 5 {
4056 return Err(Error::Msg(format!(
4057 "SelectiveScan expects 5 inputs (u, delta, a, b, c), got {}",
4058 node.inputs.len(),
4059 ))
4060 .bt());
4061 }
4062 let u_layout = input_layout(node.inputs[0]);
4063 let a_layout = input_layout(node.inputs[2]);
4064 let u_dims = u_layout.shape().dims();
4065 let a_dims = a_layout.shape().dims();
4066 if u_dims.len() != 3 {
4067 return Err(Error::Msg(format!(
4068 "SelectiveScan: u must be rank 3 [batch, seqlen, dim], got {u_dims:?}",
4069 ))
4070 .bt());
4071 }
4072 if a_dims.len() != 2 {
4073 return Err(Error::Msg(format!(
4074 "SelectiveScan: a must be rank 2 [dim, dstate], got {a_dims:?}",
4075 ))
4076 .bt());
4077 }
4078 let batch = u_dims[0];
4079 let seqlen = u_dims[1];
4080 let dim = u_dims[2];
4081 let dstate = a_dims[1];
4082 if a_dims[0] != dim {
4083 return Err(Error::Msg(format!(
4084 "SelectiveScan: a's first dim {} must equal dim {dim}",
4085 a_dims[0],
4086 ))
4087 .bt());
4088 }
4089 OpParams::SelectiveScan {
4090 batch,
4091 seqlen,
4092 dim,
4093 dstate,
4094 delta_softplus: *delta_softplus,
4095 }
4096 }
4097 Op::Fused(_, fuel_graph::registry::FusedOpParams::SsdChunkScan { chunk_size }) => {
4101 if node.inputs.len() != 5 {
4102 return Err(Error::Msg(format!(
4103 "SsdChunkScan expects 5 inputs (x, dt, a, b, c), got {}",
4104 node.inputs.len(),
4105 ))
4106 .bt());
4107 }
4108 let x_layout = input_layout(node.inputs[0]);
4109 let b_layout = input_layout(node.inputs[3]);
4110 let x_dims = x_layout.shape().dims();
4111 let b_dims = b_layout.shape().dims();
4112 if x_dims.len() != 4 {
4113 return Err(Error::Msg(format!(
4114 "SsdChunkScan: x must be rank 4 [batch, seqlen, heads, head_dim], got {x_dims:?}",
4115 ))
4116 .bt());
4117 }
4118 if b_dims.len() != 4 {
4119 return Err(Error::Msg(format!(
4120 "SsdChunkScan: b must be rank 4 [batch, seqlen, heads, state_dim], got {b_dims:?}",
4121 ))
4122 .bt());
4123 }
4124 let batch = x_dims[0];
4125 let seqlen = x_dims[1];
4126 let heads = x_dims[2];
4127 let head_dim = x_dims[3];
4128 let state_dim = b_dims[3];
4129 OpParams::SsdChunkScan {
4130 batch,
4131 seqlen,
4132 heads,
4133 head_dim,
4134 state_dim,
4135 chunk_size: *chunk_size,
4136 }
4137 }
4138 Op::Fused(_, fuel_graph::registry::FusedOpParams::Nf4Matmul { block_size }) => {
4142 if node.inputs.len() != 3 {
4143 return Err(Error::Msg(format!(
4144 "Nf4Matmul expects 3 inputs (activations, w_packed, absmax), got {}",
4145 node.inputs.len(),
4146 ))
4147 .bt());
4148 }
4149 let a_layout = input_layout(node.inputs[0]);
4150 let w_layout = input_layout(node.inputs[1]);
4151 let a_dims = a_layout.shape().dims();
4152 let w_dims = w_layout.shape().dims();
4153 if a_dims.len() < 2 {
4154 return Err(Error::Msg(format!(
4155 "Nf4Matmul: activations must be rank ≥ 2, got {a_dims:?}",
4156 ))
4157 .bt());
4158 }
4159 if w_dims.len() != 2 {
4160 return Err(Error::Msg(format!(
4161 "Nf4Matmul: w_packed must be rank 2 [n, k/2], got {w_dims:?}",
4162 ))
4163 .bt());
4164 }
4165 let m = a_dims[a_dims.len() - 2];
4166 let k = a_dims[a_dims.len() - 1];
4167 let n = w_dims[0];
4168 let batch: usize = a_dims[..a_dims.len() - 2].iter().product::<usize>().max(1);
4169 if k % 2 != 0 {
4170 return Err(Error::Msg(format!(
4171 "Nf4Matmul: k={k} must be even (w_packed holds 2 nibbles per byte along k)",
4172 ))
4173 .bt());
4174 }
4175 if k % *block_size != 0 {
4176 return Err(Error::Msg(format!(
4177 "Nf4Matmul: k={k} must be a multiple of block_size={block_size}",
4178 ))
4179 .bt());
4180 }
4181 if w_dims[1] != k / 2 {
4182 return Err(Error::Msg(format!(
4183 "Nf4Matmul: w_packed second dim {} must equal k/2 = {}",
4184 w_dims[1],
4185 k / 2,
4186 ))
4187 .bt());
4188 }
4189 OpParams::Nf4Matmul {
4190 batch,
4191 m,
4192 n,
4193 k,
4194 block_size: *block_size,
4195 }
4196 }
4197 op if op_is_softmax_last_dim(op)
4205 || matches!(op, Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM_BACKWARD) =>
4206 {
4207 let il = input_layout(node.inputs[0]);
4208 let dims = il.shape().dims();
4209 if dims.is_empty() {
4210 return Err(Error::Msg("Op::SoftmaxLastDim requires rank ≥ 1".to_string()).bt());
4211 }
4212 let last_dim = *dims.last().unwrap();
4213 let outer_count: usize = dims[..dims.len() - 1].iter().product();
4214 OpParams::SoftmaxLastDim {
4215 outer_count,
4216 last_dim,
4217 }
4218 }
4219 Op::Flip { dim } => {
4220 if node.inputs.len() != 1 {
4223 return Err(Error::Msg(format!(
4224 "Op::Flip expects 1 input, got {}",
4225 node.inputs.len(),
4226 ))
4227 .bt());
4228 }
4229 let in_layout = input_layout(node.inputs[0]);
4230 let in_dims = in_layout.shape().dims();
4231 if *dim >= in_dims.len() {
4232 return Err(Error::Msg(format!(
4233 "Op::Flip: dim {dim} out of range for rank {}",
4234 in_dims.len(),
4235 ))
4236 .bt());
4237 }
4238 let outer_count: usize = in_dims[..*dim].iter().product();
4239 let dim_size = in_dims[*dim];
4240 let inner_count: usize = in_dims[*dim + 1..].iter().product();
4241 OpParams::Flip {
4242 outer_count,
4243 dim_size,
4244 inner_count,
4245 axis: *dim,
4246 }
4247 }
4248 Op::Roll { dim, shift } => {
4249 if node.inputs.len() != 1 {
4250 return Err(Error::Msg(format!(
4251 "Op::Roll expects 1 input, got {}",
4252 node.inputs.len(),
4253 ))
4254 .bt());
4255 }
4256 let in_layout = input_layout(node.inputs[0]);
4257 let in_dims = in_layout.shape().dims();
4258 if *dim >= in_dims.len() {
4259 return Err(Error::Msg(format!(
4260 "Op::Roll: dim {dim} out of range for rank {}",
4261 in_dims.len(),
4262 ))
4263 .bt());
4264 }
4265 let outer_count: usize = in_dims[..*dim].iter().product();
4266 let dim_size = in_dims[*dim];
4267 let inner_count: usize = in_dims[*dim + 1..].iter().product();
4268 OpParams::Roll {
4269 outer_count,
4270 dim_size,
4271 inner_count,
4272 shift: *shift,
4273 axis: *dim,
4274 }
4275 }
4276 Op::CumSum { dim } => {
4277 if node.inputs.len() != 1 {
4278 return Err(Error::Msg(format!(
4279 "Op::CumSum expects 1 input, got {}",
4280 node.inputs.len(),
4281 ))
4282 .bt());
4283 }
4284 let in_layout = input_layout(node.inputs[0]);
4285 let in_dims = in_layout.shape().dims();
4286 if *dim >= in_dims.len() {
4287 return Err(Error::Msg(format!(
4288 "Op::CumSum: dim {dim} out of range for rank {}",
4289 in_dims.len(),
4290 ))
4291 .bt());
4292 }
4293 let outer_count: usize = in_dims[..*dim].iter().product();
4294 let dim_size = in_dims[*dim];
4295 let inner_count: usize = in_dims[*dim + 1..].iter().product();
4296 OpParams::CumSum {
4297 outer_count,
4298 dim_size,
4299 inner_count,
4300 axis: *dim,
4301 }
4302 }
4303 Op::Triu { diagonal } | Op::Tril { diagonal } => {
4304 if node.inputs.len() != 1 {
4305 return Err(Error::Msg(format!(
4306 "Op::Triu/Tril expects 1 input, got {}",
4307 node.inputs.len(),
4308 ))
4309 .bt());
4310 }
4311 let in_layout = input_layout(node.inputs[0]);
4312 let in_dims = in_layout.shape().dims();
4313 if in_dims.len() < 2 {
4314 return Err(Error::Msg(format!(
4315 "Op::Triu/Tril requires rank >= 2, got {}",
4316 in_dims.len(),
4317 ))
4318 .bt());
4319 }
4320 let rows = in_dims[in_dims.len() - 2];
4321 let cols = in_dims[in_dims.len() - 1];
4322 let batch_count: usize = in_dims[..in_dims.len() - 2]
4323 .iter()
4324 .product::<usize>()
4325 .max(1);
4326 OpParams::Triangular {
4327 batch_count,
4328 rows,
4329 cols,
4330 diagonal: *diagonal,
4331 }
4332 }
4333 Op::LogSoftmaxLastDim => {
4334 if node.inputs.len() != 1 {
4335 return Err(Error::Msg(format!(
4336 "Op::LogSoftmaxLastDim expects 1 input, got {}",
4337 node.inputs.len(),
4338 ))
4339 .bt());
4340 }
4341 let in_layout = input_layout(node.inputs[0]);
4342 let in_dims = in_layout.shape().dims();
4343 if in_dims.is_empty() {
4344 return Err(Error::Msg("Op::LogSoftmaxLastDim requires rank ≥ 1".to_string()).bt());
4345 }
4346 let last_dim = *in_dims.last().unwrap();
4347 let outer_count: usize = in_dims[..in_dims.len() - 1].iter().product();
4348 OpParams::LogSoftmaxLastDim {
4349 outer_count,
4350 last_dim,
4351 }
4352 }
4353 Op::LogSoftmaxLastDimBackward => {
4354 if node.inputs.len() != 2 {
4357 return Err(Error::Msg(format!(
4358 "Op::LogSoftmaxLastDimBackward expects 2 inputs, got {}",
4359 node.inputs.len(),
4360 ))
4361 .bt());
4362 }
4363 let in_layout = input_layout(node.inputs[0]);
4364 let in_dims = in_layout.shape().dims();
4365 if in_dims.is_empty() {
4366 return Err(Error::Msg(
4367 "Op::LogSoftmaxLastDimBackward requires rank ≥ 1".to_string(),
4368 )
4369 .bt());
4370 }
4371 let last_dim = *in_dims.last().unwrap();
4372 let outer_count: usize = in_dims[..in_dims.len() - 1].iter().product();
4373 OpParams::LogSoftmaxLastDim {
4374 outer_count,
4375 last_dim,
4376 }
4377 }
4378 Op::MaskedFill { value } => {
4379 if node.inputs.len() != 2 {
4380 return Err(Error::Msg(format!(
4381 "Op::MaskedFill expects 2 inputs (x, mask), got {}",
4382 node.inputs.len(),
4383 ))
4384 .bt());
4385 }
4386 let fill_bytes = scalar_to_bytes(*value);
4387 OpParams::MaskedFill { fill_bytes }
4388 }
4389 Op::PadBackward {
4392 in_shape,
4393 padding,
4394 mode,
4395 } => {
4396 if node.inputs.len() != 1 {
4399 return Err(Error::Msg(format!(
4400 "Op::PadBackward expects 1 input, got {}",
4401 node.inputs.len(),
4402 ))
4403 .bt());
4404 }
4405 let in_dims: Vec<usize> = in_shape.dims().to_vec();
4406 let out_dims: Vec<usize> = in_dims
4409 .iter()
4410 .zip(padding.iter())
4411 .map(|(&d, &(b, a))| d + b + a)
4412 .collect();
4413 let mode_tag: u8 = match mode {
4414 fuel_graph::PadMode::Constant => 0,
4415 fuel_graph::PadMode::Reflect => 1,
4416 fuel_graph::PadMode::Replicate => 2,
4417 };
4418 OpParams::PadBackward {
4419 in_shape: in_dims,
4420 out_shape: out_dims,
4421 padding: padding.clone(),
4422 mode_tag,
4423 }
4424 }
4425 Op::Pad {
4426 padding,
4427 mode,
4428 value,
4429 } => {
4430 if node.inputs.len() != 1 {
4431 return Err(Error::Msg(format!(
4432 "Op::Pad expects 1 input, got {}",
4433 node.inputs.len(),
4434 ))
4435 .bt());
4436 }
4437 let in_layout = input_layout(node.inputs[0]);
4438 let in_dims: Vec<usize> = in_layout.shape().dims().to_vec();
4439 if padding.len() != in_dims.len() {
4440 return Err(Error::Msg(format!(
4441 "Op::Pad: padding.len() ({}) != input rank ({})",
4442 padding.len(),
4443 in_dims.len(),
4444 ))
4445 .bt());
4446 }
4447 let out_dims: Vec<usize> = in_dims
4448 .iter()
4449 .zip(padding.iter())
4450 .map(|(&d, &(b, a))| d + b + a)
4451 .collect();
4452 let mode_tag: u8 = match mode {
4453 fuel_graph::PadMode::Constant => 0,
4454 fuel_graph::PadMode::Reflect => 1,
4455 fuel_graph::PadMode::Replicate => 2,
4456 };
4457 let fill_bytes = encode_value_to_bytes(node.dtype, *value)?;
4462 OpParams::Pad {
4463 in_shape: in_dims,
4464 out_shape: out_dims,
4465 padding: padding.clone(),
4466 mode_tag,
4467 fill_bytes,
4468 }
4469 }
4470 Op::IndexAdd { dim } => {
4471 if node.inputs.len() != 3 {
4474 return Err(Error::Msg(format!(
4475 "Op::IndexAdd expects 3 inputs (base, indices, src), got {}",
4476 node.inputs.len(),
4477 ))
4478 .bt());
4479 }
4480 let base_layout = input_layout(node.inputs[0]);
4481 let idx_layout = input_layout(node.inputs[1]);
4482 let src_layout = input_layout(node.inputs[2]);
4483 let base_dims = base_layout.shape().dims();
4484 let idx_dims = idx_layout.shape().dims();
4485 let src_dims = src_layout.shape().dims();
4486 if *dim >= base_dims.len() {
4487 return Err(Error::Msg(format!(
4488 "Op::IndexAdd: dim {dim} out of range for base rank {}",
4489 base_dims.len(),
4490 ))
4491 .bt());
4492 }
4493 if idx_dims.len() != 1 {
4494 return Err(Error::Msg(format!(
4495 "Op::IndexAdd: indices must be rank 1, got {idx_dims:?}",
4496 ))
4497 .bt());
4498 }
4499 if base_dims.len() != src_dims.len() {
4500 return Err(Error::Msg(format!(
4501 "Op::IndexAdd: base rank ({}) != src rank ({})",
4502 base_dims.len(),
4503 src_dims.len(),
4504 ))
4505 .bt());
4506 }
4507 let outer_count: usize = base_dims[..*dim].iter().product();
4508 let base_dim_size = base_dims[*dim];
4509 let inner_count: usize = base_dims[*dim + 1..].iter().product();
4510 let n_indices = idx_dims[0];
4511 OpParams::IndexAdd {
4512 outer_count,
4513 base_dim_size,
4514 n_indices,
4515 inner_count,
4516 }
4517 }
4518 Op::ScatterAdd { dim } => {
4519 if node.inputs.len() != 3 {
4520 return Err(Error::Msg(format!(
4521 "Op::ScatterAdd expects 3 inputs (base, indices, src), got {}",
4522 node.inputs.len(),
4523 ))
4524 .bt());
4525 }
4526 let base_layout = input_layout(node.inputs[0]);
4527 let src_layout = input_layout(node.inputs[2]);
4528 let base_shape: Vec<usize> = base_layout.shape().dims().to_vec();
4529 let src_shape: Vec<usize> = src_layout.shape().dims().to_vec();
4530 if base_shape.len() != src_shape.len() {
4531 return Err(Error::Msg(format!(
4532 "Op::ScatterAdd: base rank ({}) != src rank ({})",
4533 base_shape.len(),
4534 src_shape.len(),
4535 ))
4536 .bt());
4537 }
4538 if *dim >= base_shape.len() {
4539 return Err(Error::Msg(format!(
4540 "Op::ScatterAdd: dim {dim} out of range for rank {}",
4541 base_shape.len(),
4542 ))
4543 .bt());
4544 }
4545 OpParams::ScatterAdd {
4546 base_shape,
4547 src_shape,
4548 dim: *dim,
4549 }
4550 }
4551 Op::Fused(_, fuel_graph::registry::FusedOpParams::Rope) => {
4554 if node.inputs.len() != 3 {
4557 return Err(Error::Msg(format!(
4558 "Op::Rope expects 3 inputs (x, cos, sin), got {}",
4559 node.inputs.len(),
4560 ))
4561 .bt());
4562 }
4563 let x_layout = input_layout(node.inputs[0]);
4564 let x_dims = x_layout.shape().dims();
4565 if x_dims.len() < 2 {
4566 return Err(
4567 Error::Msg(format!("Op::Rope: x must have rank ≥ 2, got {x_dims:?}",)).bt(),
4568 );
4569 }
4570 let head_dim = *x_dims.last().unwrap();
4571 let seq = x_dims[x_dims.len() - 2];
4572 let outer_count: usize = x_dims[..x_dims.len() - 2].iter().product();
4573 OpParams::Rope {
4574 outer_count,
4575 seq,
4576 head_dim,
4577 }
4578 }
4579 Op::Gather { dim } => {
4580 if node.inputs.len() != 2 {
4582 return Err(Error::Msg(format!(
4583 "Op::Gather expects 2 inputs (source, indices), got {}",
4584 node.inputs.len(),
4585 ))
4586 .bt());
4587 }
4588 let src_layout = input_layout(node.inputs[0]);
4589 let idx_layout = input_layout(node.inputs[1]);
4590 let source_shape: Vec<usize> = src_layout.shape().dims().to_vec();
4591 let output_shape: Vec<usize> = idx_layout.shape().dims().to_vec();
4592 if source_shape.len() != output_shape.len() {
4593 return Err(Error::Msg(format!(
4594 "Op::Gather: source rank ({}) != indices rank ({})",
4595 source_shape.len(),
4596 output_shape.len(),
4597 ))
4598 .bt());
4599 }
4600 if *dim >= source_shape.len() {
4601 return Err(Error::Msg(format!(
4602 "Op::Gather: dim {dim} out of range for rank {}",
4603 source_shape.len(),
4604 ))
4605 .bt());
4606 }
4607 OpParams::Gather {
4608 source_shape,
4609 output_shape,
4610 dim: *dim,
4611 }
4612 }
4613 Op::NonZeroIndices { count_sym } => {
4614 if node.inputs.len() != 1 {
4617 return Err(Error::Msg(format!(
4618 "Op::NonZeroIndices expects 1 input, got {}",
4619 node.inputs.len(),
4620 ))
4621 .bt());
4622 }
4623 let x_layout = input_layout(node.inputs[0]);
4624 let capacity = x_layout.shape().elem_count();
4625 OpParams::NonZeroIndices {
4626 capacity,
4627 count_sym: *count_sym,
4628 }
4629 }
4630 Op::IndexSelect { dim } => {
4631 if node.inputs.len() != 2 {
4633 return Err(Error::Msg(format!(
4634 "Op::IndexSelect expects 2 inputs (source, indices), got {}",
4635 node.inputs.len(),
4636 ))
4637 .bt());
4638 }
4639 let src_layout = input_layout(node.inputs[0]);
4640 let idx_layout = input_layout(node.inputs[1]);
4641 let src_dims = src_layout.shape().dims();
4642 let idx_dims = idx_layout.shape().dims();
4643 if *dim >= src_dims.len() {
4644 return Err(Error::Msg(format!(
4645 "Op::IndexSelect: dim {dim} out of range for source rank {}",
4646 src_dims.len(),
4647 ))
4648 .bt());
4649 }
4650 if idx_dims.len() != 1 {
4651 return Err(Error::Msg(format!(
4652 "Op::IndexSelect: indices must be rank 1, got shape {idx_dims:?}",
4653 ))
4654 .bt());
4655 }
4656 let outer_count: usize = src_dims[..*dim].iter().product();
4657 let source_dim_size = src_dims[*dim];
4658 let inner_count: usize = src_dims[*dim + 1..].iter().product();
4659 let n_indices = idx_dims[0];
4660 OpParams::IndexSelect {
4661 outer_count,
4662 source_dim_size,
4663 n_indices,
4664 inner_count,
4665 }
4666 }
4667 Op::Fused(fid, params)
4673 if *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM
4674 || *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM
4675 || *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM_BACKWARD
4676 || *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM_BACKWARD =>
4677 {
4678 let eps = match params {
4679 fuel_graph::registry::FusedOpParams::RmsNormLastDim { eps } => *eps,
4680 fuel_graph::registry::FusedOpParams::LayerNormLastDim { eps } => *eps,
4681 fuel_graph::registry::FusedOpParams::RmsNormLastDimBackward { eps } => *eps,
4682 fuel_graph::registry::FusedOpParams::LayerNormLastDimBackward { eps } => *eps,
4683 _ => {
4684 return Err(Error::Msg(format!(
4685 "Op::Fused(NORM_LAST_DIM[_BACKWARD], _) expected \
4686 RmsNormLastDim[Backward] or LayerNormLastDim[Backward] params, got {params:?}",
4687 ))
4688 .bt());
4689 }
4690 };
4691 let il = input_layout(node.inputs[0]);
4692 let dims = il.shape().dims();
4693 if dims.is_empty() {
4694 return Err(Error::Msg(
4695 "Op::Fused(NORM_LAST_DIM[_BACKWARD], _) requires rank ≥ 1".to_string(),
4696 )
4697 .bt());
4698 }
4699 let last_dim = *dims.last().unwrap();
4700 let outer_count: usize = dims[..dims.len() - 1].iter().product();
4701 OpParams::NormLastDim {
4702 outer_count,
4703 last_dim,
4704 eps,
4705 }
4706 }
4707 Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::REDUCE_MAX_TO_BACKWARD => {
4711 if node.inputs.len() != 2 {
4712 return Err(Error::Msg(format!(
4713 "Op::Fused(REDUCE_MAX_TO_BACKWARD) expects 2 inputs, got {}",
4714 node.inputs.len(),
4715 ))
4716 .bt());
4717 }
4718 let x_layout = input_layout(node.inputs[0]);
4719 let up_layout = input_layout(node.inputs[1]);
4720 let input_shape: Vec<usize> = x_layout.shape().dims().to_vec();
4721 let output_shape: Vec<usize> = up_layout.shape().dims().to_vec();
4722 OpParams::ReduceMaxToBackward {
4723 input_shape,
4724 output_shape,
4725 }
4726 }
4727 Op::Concat { dim } => {
4728 if node.inputs.is_empty() {
4733 return Err(Error::Msg("Op::Concat requires at least 1 input".to_string()).bt());
4734 }
4735 let out_dims = node.shape.dims();
4736 if *dim >= out_dims.len() {
4737 return Err(Error::Msg(format!(
4738 "Op::Concat: dim {dim} out of range for output rank {}",
4739 out_dims.len(),
4740 ))
4741 .bt());
4742 }
4743 let outer_count: usize = out_dims[..*dim].iter().product();
4744 let inner_count: usize = out_dims[*dim + 1..].iter().product();
4745 let mut input_dim_sizes: Vec<usize> = Vec::with_capacity(node.inputs.len());
4746 for in_id in &node.inputs {
4747 let il = input_layout(*in_id);
4748 let il_dims = il.shape().dims();
4749 if *dim >= il_dims.len() {
4750 return Err(Error::Msg(format!(
4751 "Op::Concat: input {in_id:?} has rank {} but concat dim is {dim}",
4752 il_dims.len(),
4753 ))
4754 .bt());
4755 }
4756 input_dim_sizes.push(il_dims[*dim]);
4757 }
4758 OpParams::Concat {
4759 outer_count,
4760 input_dim_sizes,
4761 inner_count,
4762 axis: *dim,
4763 }
4764 }
4765 Op::AddScalar(c) => OpParams::Affine { mul: 1.0, add: *c },
4766 Op::MulScalar(c) => OpParams::Affine { mul: *c, add: 0.0 },
4767 Op::Clamp { min, max } => OpParams::Clamp {
4768 min: *min,
4769 max: *max,
4770 },
4771 Op::PowI(exp) => OpParams::PowI { exp: *exp },
4772 Op::ClampInplace { min, max } => OpParams::Clamp {
4773 min: *min,
4774 max: *max,
4775 },
4776 Op::PowIInplace(exp) => OpParams::PowI { exp: *exp },
4777 Op::Fused(_, fuel_graph::registry::FusedOpParams::PowIBackward { exp }) => {
4781 OpParams::PowI { exp: *exp }
4782 }
4783 Op::Fused(_, fuel_graph::registry::FusedOpParams::InplaceAffine { mul, add }) => {
4789 OpParams::Affine {
4790 mul: *mul,
4791 add: *add,
4792 }
4793 }
4794 Op::Fused(
4797 _,
4798 fuel_graph::registry::FusedOpParams::Conv2D {
4799 stride,
4800 padding,
4801 groups,
4802 },
4803 ) => {
4804 if node.inputs.len() != 2 && node.inputs.len() != 3 {
4808 return Err(Error::Msg(format!(
4809 "Op::Conv2D expects 2 or 3 inputs, got {}",
4810 node.inputs.len(),
4811 ))
4812 .bt());
4813 }
4814 let x_layout = input_layout(node.inputs[0]);
4815 let w_layout = input_layout(node.inputs[1]);
4816 let x_dims = x_layout.shape().dims();
4817 let w_dims = w_layout.shape().dims();
4818 if x_dims.len() != 4 || w_dims.len() != 4 {
4819 return Err(Error::Msg(format!(
4820 "Op::Conv2D requires rank-4 x and weight; got x={x_dims:?} w={w_dims:?}",
4821 ))
4822 .bt());
4823 }
4824 let out_dims = node.shape.dims();
4825 if out_dims.len() != 4 {
4826 return Err(Error::Msg(format!(
4827 "Op::Conv2D output must be rank 4, got {out_dims:?}",
4828 ))
4829 .bt());
4830 }
4831 let x_shape = [x_dims[0], x_dims[1], x_dims[2], x_dims[3]];
4832 let w_shape = [w_dims[0], w_dims[1], w_dims[2], w_dims[3]];
4833 let out_shape = [out_dims[0], out_dims[1], out_dims[2], out_dims[3]];
4834 OpParams::Conv2D {
4835 x_shape,
4836 w_shape,
4837 out_shape,
4838 stride: *stride,
4839 padding: *padding,
4840 dilation: (1, 1),
4841 groups: *groups,
4842 }
4843 }
4844 Op::Fused(
4846 _,
4847 fuel_graph::registry::FusedOpParams::PagedAttn {
4848 softmax_scale,
4849 block_size,
4850 softcap,
4851 },
4852 ) => {
4853 if node.inputs.len() != 5 && node.inputs.len() != 6 {
4858 return Err(Error::Msg(format!(
4859 "Op::PagedAttn expects 5 or 6 inputs, got {}",
4860 node.inputs.len(),
4861 ))
4862 .bt());
4863 }
4864 let q_layout = input_layout(node.inputs[0]);
4865 let kc_layout = input_layout(node.inputs[1]);
4866 let vc_layout = input_layout(node.inputs[2]);
4867 let bt_layout = input_layout(node.inputs[3]);
4868 let q_dims = q_layout.shape().dims();
4869 let kc_dims = kc_layout.shape().dims();
4870 let vc_dims = vc_layout.shape().dims();
4871 let bt_dims = bt_layout.shape().dims();
4872 if q_dims.len() != 4 || kc_dims.len() != 4 || vc_dims.len() != 4 {
4873 return Err(Error::Msg(format!(
4874 "Op::PagedAttn requires rank-4 q/k_cache/v_cache; \
4875 got q={q_dims:?} k_cache={kc_dims:?} v_cache={vc_dims:?}",
4876 ))
4877 .bt());
4878 }
4879 if kc_dims != vc_dims {
4880 return Err(Error::Msg(format!(
4881 "Op::PagedAttn: k_cache {kc_dims:?} and v_cache {vc_dims:?} must match",
4882 ))
4883 .bt());
4884 }
4885 if bt_dims.len() != 2 {
4886 return Err(Error::Msg(format!(
4887 "Op::PagedAttn: block_table must be rank 2 [B, max_blocks_per_seq], got {bt_dims:?}",
4888 ))
4889 .bt());
4890 }
4891 if kc_dims[1] != *block_size {
4892 return Err(Error::Msg(format!(
4893 "Op::PagedAttn: k_cache block_size dim ({}) must equal Op's block_size ({})",
4894 kc_dims[1], block_size,
4895 ))
4896 .bt());
4897 }
4898 OpParams::PagedAttn {
4899 b: q_dims[0],
4900 hq: q_dims[1],
4901 hkv: kc_dims[2],
4902 sq: q_dims[2],
4903 d: q_dims[3],
4904 block_size: *block_size,
4905 max_blocks_per_seq: bt_dims[1],
4906 num_blocks: kc_dims[0],
4907 softmax_scale: *softmax_scale,
4908 softcap: *softcap,
4909 }
4910 }
4911 Op::Fused(
4913 _,
4914 fuel_graph::registry::FusedOpParams::FlashAttn {
4915 softmax_scale,
4916 causal,
4917 window_size_left,
4918 window_size_right,
4919 softcap,
4920 k_len,
4921 },
4922 ) => {
4923 if node.inputs.len() != 3 && node.inputs.len() != 4 {
4928 return Err(Error::Msg(format!(
4929 "Op::FlashAttn expects 3 or 4 inputs, got {}",
4930 node.inputs.len(),
4931 ))
4932 .bt());
4933 }
4934 let q_layout = input_layout(node.inputs[0]);
4935 let k_layout = input_layout(node.inputs[1]);
4936 let v_layout = input_layout(node.inputs[2]);
4937 let q_dims = q_layout.shape().dims();
4938 let k_dims = k_layout.shape().dims();
4939 let v_dims = v_layout.shape().dims();
4940 if q_dims.len() != 4 || k_dims.len() != 4 || v_dims.len() != 4 {
4941 return Err(Error::Msg(format!(
4942 "Op::FlashAttn requires rank-4 q/k/v; got q={q_dims:?} k={k_dims:?} v={v_dims:?}",
4943 ))
4944 .bt());
4945 }
4946 if k_dims != v_dims {
4947 return Err(Error::Msg(format!(
4948 "Op::FlashAttn: k {k_dims:?} and v {v_dims:?} must share shape",
4949 ))
4950 .bt());
4951 }
4952 if q_dims[0] != k_dims[0] || q_dims[3] != k_dims[3] {
4953 return Err(Error::Msg(format!(
4954 "Op::FlashAttn: q {q_dims:?} and k {k_dims:?} must share B and D",
4955 ))
4956 .bt());
4957 }
4958 let sk = k_dims[2];
4959 let attended_k_len = match k_len {
4964 None => sk,
4965 Some(dyn_k) => {
4966 let kl = dyn_k.resolve(sym_env).ok_or_else(|| {
4967 Error::Msg(format!(
4968 "Op::FlashAttn: dynamic k_len {dyn_k:?} is unbound in the SymEnv at realize",
4969 ))
4970 .bt()
4971 })?;
4972 if kl > sk {
4973 return Err(Error::Msg(format!(
4974 "Op::FlashAttn: resolved k_len ({kl}) exceeds K capacity ({sk})",
4975 ))
4976 .bt());
4977 }
4978 if kl < q_dims[2] {
4979 return Err(Error::Msg(format!(
4980 "Op::FlashAttn: resolved k_len ({kl}) < Sq ({}) — not a valid causal prefix",
4981 q_dims[2],
4982 ))
4983 .bt());
4984 }
4985 kl
4986 }
4987 };
4988 OpParams::FlashAttn {
4989 b: q_dims[0],
4990 hq: q_dims[1],
4991 hkv: k_dims[1],
4992 sq: q_dims[2],
4993 sk,
4994 d: q_dims[3],
4995 k_len: attended_k_len,
4996 softmax_scale: *softmax_scale,
4997 causal: *causal,
4998 window_size_left: *window_size_left,
4999 window_size_right: *window_size_right,
5000 softcap: *softcap,
5001 }
5002 }
5003 Op::ReduceSumTo(target_shape) => {
5004 if node.inputs.len() != 1 {
5005 return Err(Error::Msg(format!(
5006 "Op::ReduceSumTo expects 1 input, got {}",
5007 node.inputs.len(),
5008 ))
5009 .bt());
5010 }
5011 let in_layout = input_layout(node.inputs[0]);
5012 let input_shape: Vec<usize> = in_layout.shape().dims().to_vec();
5013 let output_shape: Vec<usize> = target_shape.dims().to_vec();
5014 OpParams::ReduceSumTo {
5015 input_shape,
5016 output_shape,
5017 }
5018 }
5019 Op::ReduceMaxTo(target_shape) => {
5020 if node.inputs.len() != 1 {
5021 return Err(Error::Msg(format!(
5022 "Op::ReduceMaxTo expects 1 input, got {}",
5023 node.inputs.len(),
5024 ))
5025 .bt());
5026 }
5027 let in_layout = input_layout(node.inputs[0]);
5028 let input_shape: Vec<usize> = in_layout.shape().dims().to_vec();
5029 let output_shape: Vec<usize> = target_shape.dims().to_vec();
5030 OpParams::ReduceMaxTo {
5031 input_shape,
5032 output_shape,
5033 }
5034 }
5035 Op::Fused(
5037 _,
5038 fuel_graph::registry::FusedOpParams::ConvTranspose2D {
5039 stride,
5040 padding,
5041 output_padding,
5042 dilation,
5043 groups,
5044 },
5045 ) => {
5046 if node.inputs.len() != 2 && node.inputs.len() != 3 {
5050 return Err(Error::Msg(format!(
5051 "Op::ConvTranspose2D expects 2 or 3 inputs, got {}",
5052 node.inputs.len(),
5053 ))
5054 .bt());
5055 }
5056 let x_layout = input_layout(node.inputs[0]);
5057 let w_layout = input_layout(node.inputs[1]);
5058 let x_dims = x_layout.shape().dims();
5059 let w_dims = w_layout.shape().dims();
5060 if x_dims.len() != 4 || w_dims.len() != 4 {
5061 return Err(Error::Msg(format!(
5062 "Op::ConvTranspose2D requires rank-4 x and weight; got x={x_dims:?} w={w_dims:?}",
5063 ))
5064 .bt());
5065 }
5066 let out_dims = node.shape.dims();
5067 if out_dims.len() != 4 {
5068 return Err(Error::Msg(format!(
5069 "Op::ConvTranspose2D output must be rank 4, got {out_dims:?}",
5070 ))
5071 .bt());
5072 }
5073 let x_shape = [x_dims[0], x_dims[1], x_dims[2], x_dims[3]];
5074 let w_shape = [w_dims[0], w_dims[1], w_dims[2], w_dims[3]];
5075 let out_shape = [out_dims[0], out_dims[1], out_dims[2], out_dims[3]];
5076 OpParams::ConvTranspose2D {
5077 x_shape,
5078 w_shape,
5079 out_shape,
5080 stride: *stride,
5081 padding: *padding,
5082 output_padding: *output_padding,
5083 dilation: *dilation,
5084 groups: *groups,
5085 }
5086 }
5087 Op::WriteSlice { ranges, dyn_offset } => {
5088 if node.inputs.len() != 2 {
5094 return Err(Error::Msg(format!(
5095 "Op::WriteSlice expects 2 inputs (destination, source), got {}",
5096 node.inputs.len(),
5097 ))
5098 .bt());
5099 }
5100 let dest_dims = node.shape.dims().to_vec();
5101 if ranges.len() != dest_dims.len() {
5102 return Err(Error::Msg(format!(
5103 "Op::WriteSlice: ranges.len() ({}) must equal destination rank ({})",
5104 ranges.len(),
5105 dest_dims.len(),
5106 ))
5107 .bt());
5108 }
5109 let mut ranges = ranges.clone();
5130 let mut deferred_dyn_offset: Option<(usize, fuel_ir::DynScalar)> = None;
5131 if let Some((axis, off)) = dyn_offset {
5132 if *axis >= ranges.len() {
5133 return Err(Error::Msg(format!(
5134 "Op::WriteSlice: dyn_offset axis ({axis}) out of bounds for rank {}",
5135 ranges.len(),
5136 ))
5137 .bt());
5138 }
5139 match off.resolve(sym_env) {
5140 Some(base) => {
5141 let (start, end) = ranges[*axis];
5142 let width = end - start;
5143 ranges[*axis] = (base, base + width);
5144 }
5145 None => {
5146 let is_data_determined = matches!(off, fuel_ir::DynScalar::Sym(s)
5147 if graph.binds_data_determined_sym(*s));
5148 if is_data_determined {
5149 deferred_dyn_offset = Some((*axis, *off));
5150 } else {
5151 return Err(Error::Msg(format!(
5152 "Op::WriteSlice: dynamic offset {off:?} on axis {axis} is \
5153 unbound in the SymEnv at realize (and no producer in the \
5154 graph binds it, so it is not data-determined)",
5155 ))
5156 .bt());
5157 }
5158 }
5159 }
5160 }
5161 for (i, &(start, end)) in ranges.iter().enumerate() {
5162 if deferred_dyn_offset.is_some_and(|(a, _)| a == i) {
5166 continue;
5167 }
5168 if end < start || end > dest_dims[i] {
5169 return Err(Error::Msg(format!(
5170 "Op::WriteSlice: ranges[{i}] = ({start}, {end}) invalid \
5171 for destination dim {i} = {}",
5172 dest_dims[i],
5173 ))
5174 .bt());
5175 }
5176 }
5177 OpParams::WriteSlice {
5178 dest_shape: dest_dims,
5179 ranges,
5180 deferred_dyn_offset,
5181 }
5182 }
5183 Op::WriteSliceRotating {
5184 axis,
5185 modulus,
5186 ranges,
5187 } => {
5188 if node.inputs.len() != 3 {
5194 return Err(Error::Msg(format!(
5195 "Op::WriteSliceRotating expects 3 inputs (destination, source, position), got {}",
5196 node.inputs.len(),
5197 ))
5198 .bt());
5199 }
5200 let dest_dims = node.shape.dims().to_vec();
5201 if ranges.len() != dest_dims.len() {
5202 return Err(Error::Msg(format!(
5203 "Op::WriteSliceRotating: ranges.len() ({}) must equal destination rank ({})",
5204 ranges.len(),
5205 dest_dims.len(),
5206 ))
5207 .bt());
5208 }
5209 if *axis >= dest_dims.len() {
5210 return Err(Error::Msg(format!(
5211 "Op::WriteSliceRotating: axis {axis} out of bounds for rank {}",
5212 dest_dims.len(),
5213 ))
5214 .bt());
5215 }
5216 if *modulus == 0 || *modulus > dest_dims[*axis] {
5217 return Err(Error::Msg(format!(
5218 "Op::WriteSliceRotating: modulus {modulus} invalid for dest_shape[{axis}] = {}",
5219 dest_dims[*axis],
5220 ))
5221 .bt());
5222 }
5223 OpParams::WriteSliceRotating {
5224 dest_shape: dest_dims,
5225 axis: *axis,
5226 modulus: *modulus,
5227 ranges: ranges.clone(),
5228 }
5229 }
5230 Op::WriteSliceDoff { axis, ranges } => {
5231 if node.inputs.len() != 3 {
5238 return Err(Error::Msg(format!(
5239 "Op::WriteSliceDoff expects 3 inputs (destination, source, offset), got {}",
5240 node.inputs.len(),
5241 ))
5242 .bt());
5243 }
5244 let dest_dims = node.shape.dims().to_vec();
5245 if ranges.len() != dest_dims.len() {
5246 return Err(Error::Msg(format!(
5247 "Op::WriteSliceDoff: ranges.len() ({}) must equal destination rank ({})",
5248 ranges.len(),
5249 dest_dims.len(),
5250 ))
5251 .bt());
5252 }
5253 if *axis >= dest_dims.len() {
5254 return Err(Error::Msg(format!(
5255 "Op::WriteSliceDoff: axis {axis} out of bounds for rank {}",
5256 dest_dims.len(),
5257 ))
5258 .bt());
5259 }
5260 OpParams::WriteSliceDoff {
5261 dest_shape: dest_dims,
5262 axis: *axis,
5263 ranges: ranges.clone(),
5264 }
5265 }
5266 _ => OpParams::None,
5267 })
5268}
5269
5270fn with_node_location(graph: &Arc<RwLock<Graph>>, node_id: NodeId, err: Error) -> Error {
5292 match graph.read() {
5293 Ok(g) => err.context(g.describe_node(node_id)),
5294 Err(_) => err,
5295 }
5296}
5297
5298fn bind_data_determined_count(
5318 arc: &Arc<RwLock<Storage>>,
5319 capacity: usize,
5320 count_sym: fuel_ir::SymId,
5321 produced_syms: &mut SymEnv,
5322) -> Result<()> {
5323 let guard = arc
5324 .read()
5325 .map_err(|_| poisoned("NonZeroIndices output storage"))?;
5326 let count_byte_off = capacity * std::mem::size_of::<u32>();
5327 let end = count_byte_off + std::mem::size_of::<u32>();
5328 #[allow(unreachable_patterns)]
5332 let count = match &guard.inner {
5333 fuel_memory::BackendStorage::Cpu(cpu_bytes) => {
5334 let bytes = cpu_bytes.bytes();
5335 if end > bytes.len() {
5336 return Err(Error::Msg(format!(
5337 "NonZeroIndices: count slot byte range [{count_byte_off}..{end}) \
5338 exceeds bundle byte length {}",
5339 bytes.len(),
5340 ))
5341 .bt());
5342 }
5343 u32::from_le_bytes(
5344 bytes[count_byte_off..end]
5345 .try_into()
5346 .expect("4 bytes → u32"),
5347 )
5348 }
5349 other => {
5350 return Err(Error::Msg(format!(
5351 "NonZeroIndices: data-determined count read is CPU-only today; got \
5352 BackendStorage::{:?} (device backends need a D2H after the completion \
5353 handle is waited — a documented follow-up)",
5354 std::mem::discriminant(other),
5355 ))
5356 .bt());
5357 }
5358 };
5359 produced_syms.bind(count_sym, count as usize)?;
5360 Ok(())
5361}
5362
5363fn resolve_deferred_write_slice(
5374 compiled: &CompiledNode,
5375 produced_syms: &SymEnv,
5376 node_id: NodeId,
5377) -> Result<Option<CompiledNode>> {
5378 let OpParams::WriteSlice {
5379 dest_shape,
5380 ranges,
5381 deferred_dyn_offset,
5382 } = &compiled.op_params
5383 else {
5384 return Ok(None);
5385 };
5386 let Some((axis, off)) = deferred_dyn_offset else {
5387 return Ok(None);
5388 };
5389 let base = off.resolve(produced_syms).ok_or_else(|| {
5390 Error::Msg(format!(
5391 "Op::WriteSlice {node_id:?}: data-determined offset {off:?} on axis {axis} is \
5392 unbound in produced_syms at execute (its producer did not run / bind it)",
5393 ))
5394 .bt()
5395 })?;
5396 let mut ranges = ranges.clone();
5397 let (start, end) = ranges[*axis];
5398 let width = end - start;
5399 let resolved = (base, base + width);
5400 if resolved.1 > dest_shape[*axis] {
5401 return Err(Error::Msg(format!(
5402 "Op::WriteSlice {node_id:?}: resolved data-determined range {resolved:?} on axis \
5403 {axis} exceeds destination dim = {}",
5404 dest_shape[*axis],
5405 ))
5406 .bt());
5407 }
5408 ranges[*axis] = resolved;
5409 let mut patched = compiled.clone();
5410 patched.op_params = OpParams::WriteSlice {
5411 dest_shape: dest_shape.clone(),
5412 ranges,
5413 deferred_dyn_offset: None,
5414 };
5415 Ok(Some(patched))
5416}
5417
5418fn resolve_deferred_matmul(
5426 compiled: &CompiledNode,
5427 produced_syms: &SymEnv,
5428 node_id: NodeId,
5429) -> Result<Option<CompiledNode>> {
5430 let OpParams::Matmul {
5431 lhs_batch_dims,
5432 rhs_batch_dims,
5433 m,
5434 n,
5435 k,
5436 m_compute,
5437 } = &compiled.op_params
5438 else {
5439 return Ok(None);
5440 };
5441 let MatmulM::Deferred(dyn_m) = m_compute else {
5442 return Ok(None);
5443 };
5444 let rows = dyn_m.resolve(produced_syms).ok_or_else(|| {
5445 Error::Msg(format!(
5446 "Op::MatMul {node_id:?}: data-determined row count {dyn_m:?} is unbound in \
5447 produced_syms at execute (its producer did not run / bind it)",
5448 ))
5449 .bt()
5450 })?;
5451 if rows > *m {
5452 return Err(Error::Msg(format!(
5453 "Op::MatMul {node_id:?}: resolved row count {rows} exceeds capacity {m}",
5454 ))
5455 .bt());
5456 }
5457 let mut patched = compiled.clone();
5458 patched.op_params = OpParams::Matmul {
5459 lhs_batch_dims: lhs_batch_dims.clone(),
5460 rhs_batch_dims: rhs_batch_dims.clone(),
5461 m: *m,
5462 n: *n,
5463 k: *k,
5464 m_compute: MatmulM::Rows(rows),
5465 };
5466 Ok(Some(patched))
5467}
5468
5469fn execute_work_item(
5470 item: &WorkItem,
5471 cache: &mut StorageCache,
5472 layout_cache: &mut HashMap<NodeId, Layout>,
5473 produced_syms: &mut SymEnv,
5474 will_be_waited: bool,
5480 mut persistent: Option<&mut PersistentOutputs>,
5487) -> Result<CompletionHandle> {
5488 match &item.kind {
5489 WorkItemKind::ConstAdopt => {
5490 if !cache.contains_key(&item.node_id) {
5491 return Err(Error::Msg(format!(
5492 "PipelinedExecutor: Const node {:?} not in input cache",
5493 item.node_id
5494 ))
5495 .bt());
5496 }
5497 layout_cache.insert(item.node_id, item.output_layout.clone());
5501 Ok(CompletionHandle::Ready)
5502 }
5503 WorkItemKind::ViewOf { input } => {
5504 let input_arc = cache.get(input).cloned().ok_or_else(|| {
5505 Error::Msg(format!(
5506 "PipelinedExecutor: view-op input {:?} of {:?} not realized",
5507 input, item.node_id,
5508 ))
5509 .bt()
5510 })?;
5511 cache.insert(item.node_id, input_arc);
5512 layout_cache.insert(item.node_id, item.output_layout.clone());
5513 Ok(CompletionHandle::Ready)
5514 }
5515 WorkItemKind::SlotView { producer } => {
5516 let producer_arc = cache.get(producer).cloned().ok_or_else(|| {
5525 Error::Msg(format!(
5526 "PipelinedExecutor: Op::View producer {:?} of {:?} not realized",
5527 producer, item.node_id,
5528 ))
5529 .bt()
5530 })?;
5531 cache.insert(item.node_id, producer_arc);
5532 layout_cache.insert(item.node_id, item.output_layout.clone());
5533 Ok(CompletionHandle::Ready)
5534 }
5535 WorkItemKind::SlotOwn { producer, slot } => {
5536 let producer_arc = cache.get(producer).cloned().ok_or_else(|| {
5550 Error::Msg(format!(
5551 "PipelinedExecutor: Op::ViewOwned producer {:?} of {:?} not realized",
5552 producer, item.node_id,
5553 ))
5554 .bt()
5555 })?;
5556 let new_storage = {
5557 let guard = producer_arc
5558 .read()
5559 .map_err(|_| poisoned("ViewOwned producer storage"))?;
5560 let bundle = guard.bundle().ok_or_else(|| {
5561 Error::Msg(format!(
5562 "Op::ViewOwned: producer {:?} has no bundle side-table \
5563 — the multi-output authoring contract requires \
5564 Storage::with_bundle / new_bundled at allocation \
5565 time. Node {:?} slot {} cannot resolve its source.",
5566 producer, item.node_id, slot,
5567 ))
5568 .bt()
5569 })?;
5570 let sv = bundle.get(*slot as usize).cloned().ok_or_else(|| {
5571 Error::Msg(format!(
5572 "Op::ViewOwned: slot {} out of range for producer {:?} \
5573 (bundle has {} slots)",
5574 slot,
5575 producer,
5576 bundle.len(),
5577 ))
5578 .bt()
5579 })?;
5580 use fuel_memory::BackendStorage;
5581 match &guard.inner {
5582 BackendStorage::Cpu(cpu_bytes) => {
5583 let src_bytes = cpu_bytes.bytes();
5584 let end = sv.byte_offset + sv.len_bytes();
5585 if end > src_bytes.len() {
5586 return Err(Error::Msg(format!(
5587 "Op::ViewOwned: slot {} byte range [{}..{}) \
5588 exceeds producer's CPU byte length {}",
5589 slot,
5590 sv.byte_offset,
5591 end,
5592 src_bytes.len(),
5593 ))
5594 .bt());
5595 }
5596 let owned = fuel_cpu_backend::CpuStorageBytes::from_bytes(
5597 &src_bytes[sv.byte_offset..end],
5598 );
5599 fuel_memory::Storage::new(BackendStorage::Cpu(owned), sv.dtype)
5600 }
5601 #[cfg(feature = "cuda")]
5602 BackendStorage::Cuda(cuda_bytes) => {
5603 let dst = cuda_bytes.slot_copy_to_new(sv.byte_offset, sv.len_bytes())?;
5604 fuel_memory::Storage::new(BackendStorage::Cuda(dst), sv.dtype)
5605 }
5606 #[cfg(feature = "vulkan")]
5607 BackendStorage::Vulkan(vk_bytes) => {
5608 let backend = vk_bytes.backend().ok_or_else(|| {
5609 Error::Msg(format!(
5610 "Op::ViewOwned: producer {:?} is on Vulkan but \
5611 its storage has no VulkanBackend handle. \
5612 Storages flowing through the pipelined executor \
5613 must be constructed via alloc_bytes_handle / \
5614 upload_bytes_handle so the bundle's slot extraction \
5615 path can reach the backend.",
5616 producer,
5617 ))
5618 .bt()
5619 })?;
5620 let dst = backend.slot_copy_to_new_handle(
5621 vk_bytes,
5622 sv.byte_offset,
5623 sv.len_bytes(),
5624 )?;
5625 fuel_memory::Storage::new(BackendStorage::Vulkan(dst), sv.dtype)
5626 }
5627 #[allow(unreachable_patterns)]
5628 _ => {
5629 return Err(Error::Msg(format!(
5630 "Op::ViewOwned: producer {:?} is on a backend \
5631 without a slot-copy-with-offset hook \
5632 (CPU + CUDA + Vulkan wired). Add a per-backend \
5633 `slot_copy_to_new`-equivalent and extend the \
5634 match arm.",
5635 producer,
5636 ))
5637 .bt());
5638 }
5639 }
5640 };
5641 cache.insert(item.node_id, Arc::new(RwLock::new(new_storage)));
5642 layout_cache.insert(item.node_id, item.output_layout.clone());
5643 Ok(CompletionHandle::Ready)
5647 }
5648 WorkItemKind::WriteSlice { dest, source } => {
5649 let compiled = item.compiled.as_ref().ok_or_else(|| {
5650 Error::Msg(format!(
5651 "PipelinedExecutor: WriteSlice work item {:?} has no compiled node",
5652 item.node_id,
5653 ))
5654 .bt()
5655 })?;
5656 let patched_ws = resolve_deferred_write_slice(compiled, produced_syms, item.node_id)?;
5660 let compiled = patched_ws.as_ref().unwrap_or(compiled);
5661 let dest_arc = cache.get(dest).cloned().ok_or_else(|| {
5662 Error::Msg(format!(
5663 "PipelinedExecutor: WriteSlice destination {:?} of {:?} not realized",
5664 dest, item.node_id,
5665 ))
5666 .bt()
5667 })?;
5668 let source_arc = cache.get(source).cloned().ok_or_else(|| {
5669 Error::Msg(format!(
5670 "PipelinedExecutor: WriteSlice source {:?} of {:?} not realized",
5671 source, item.node_id,
5672 ))
5673 .bt()
5674 })?;
5675 let dest_layout = layout_cache.get(dest).cloned().ok_or_else(|| {
5676 Error::Msg(format!(
5677 "PipelinedExecutor: WriteSlice destination {:?} of {:?} has no cached layout",
5678 dest, item.node_id,
5679 ))
5680 .bt()
5681 })?;
5682 let source_layout = layout_cache.get(source).cloned().ok_or_else(|| {
5683 Error::Msg(format!(
5684 "PipelinedExecutor: WriteSlice source {:?} of {:?} has no cached layout",
5685 source, item.node_id,
5686 ))
5687 .bt()
5688 })?;
5689 if !dest_layout.is_contiguous() || dest_layout.start_offset() != 0 {
5696 return Err(Error::Msg(format!(
5697 "Op::WriteSlice (Phase E.3.2 v1): destination {:?} must be \
5698 contiguous + zero-offset; got Layout {:?}",
5699 dest, dest_layout,
5700 ))
5701 .bt());
5702 }
5703 let source_arc_contig = {
5704 let s_dtype = source_arc
5705 .read()
5706 .map_err(|_| poisoned("WriteSlice source storage"))?
5707 .dtype;
5708 let s_len_bytes = source_arc
5709 .read()
5710 .map_err(|_| poisoned("WriteSlice source storage"))?
5711 .inner
5712 .len_bytes();
5713 let layout_bytes = source_layout.shape().elem_count() * s_dtype.size_in_bytes();
5714 let bytes_match_shape = s_len_bytes == layout_bytes;
5715 let already_contig = source_layout.is_contiguous()
5716 && source_layout.start_offset() == 0
5717 && bytes_match_shape;
5718 if already_contig {
5719 source_arc
5720 } else {
5721 auto_contiguize(&source_arc, &source_layout)?
5722 }
5723 };
5724 let source_layout_kernel = fuel_ir::Layout::contiguous(source_layout.shape().clone());
5725 let input_arcs = vec![source_arc_contig];
5730 let mut output_arcs = vec![dest_arc.clone()];
5731 let kernel_layouts = vec![source_layout_kernel, dest_layout.clone()];
5732 let handle = execute_compiled_with_wait_hint(
5734 compiled,
5735 &input_arcs,
5736 &mut output_arcs,
5737 &kernel_layouts,
5738 will_be_waited,
5739 )?;
5740 cache.insert(item.node_id, dest_arc);
5745 layout_cache.insert(item.node_id, item.output_layout.clone());
5746 Ok(handle)
5747 }
5748 WorkItemKind::WriteSliceRotating {
5749 dest,
5750 source,
5751 position,
5752 } => {
5753 let compiled = item.compiled.as_ref().ok_or_else(|| {
5757 Error::Msg(format!(
5758 "PipelinedExecutor: WriteSliceRotating work item {:?} has no compiled node",
5759 item.node_id,
5760 ))
5761 .bt()
5762 })?;
5763 let dest_arc = cache.get(dest).cloned().ok_or_else(|| {
5764 Error::Msg(format!(
5765 "PipelinedExecutor: WriteSliceRotating destination {:?} of {:?} not realized",
5766 dest, item.node_id,
5767 ))
5768 .bt()
5769 })?;
5770 let source_arc = cache.get(source).cloned().ok_or_else(|| {
5771 Error::Msg(format!(
5772 "PipelinedExecutor: WriteSliceRotating source {:?} of {:?} not realized",
5773 source, item.node_id,
5774 ))
5775 .bt()
5776 })?;
5777 let position_arc = cache.get(position).cloned().ok_or_else(|| {
5778 Error::Msg(format!(
5779 "PipelinedExecutor: WriteSliceRotating position {:?} of {:?} not realized",
5780 position, item.node_id,
5781 ))
5782 .bt()
5783 })?;
5784 let dest_layout = layout_cache.get(dest).cloned().ok_or_else(|| {
5785 Error::Msg(format!(
5786 "PipelinedExecutor: WriteSliceRotating destination {:?} of {:?} has no cached layout",
5787 dest, item.node_id,
5788 ))
5789 .bt()
5790 })?;
5791 let source_layout = layout_cache.get(source).cloned().ok_or_else(|| {
5792 Error::Msg(format!(
5793 "PipelinedExecutor: WriteSliceRotating source {:?} of {:?} has no cached layout",
5794 source, item.node_id,
5795 ))
5796 .bt()
5797 })?;
5798 if !dest_layout.is_contiguous() || dest_layout.start_offset() != 0 {
5800 return Err(Error::Msg(format!(
5801 "Op::WriteSliceRotating (v1): destination {:?} must be \
5802 contiguous + zero-offset; got Layout {:?}",
5803 dest, dest_layout,
5804 ))
5805 .bt());
5806 }
5807 let source_arc_contig = {
5808 let s_dtype = source_arc
5809 .read()
5810 .map_err(|_| poisoned("WriteSliceRotating source storage"))?
5811 .dtype;
5812 let s_len_bytes = source_arc
5813 .read()
5814 .map_err(|_| poisoned("WriteSliceRotating source storage"))?
5815 .inner
5816 .len_bytes();
5817 let layout_bytes = source_layout.shape().elem_count() * s_dtype.size_in_bytes();
5818 let bytes_match_shape = s_len_bytes == layout_bytes;
5819 let already_contig = source_layout.is_contiguous()
5820 && source_layout.start_offset() == 0
5821 && bytes_match_shape;
5822 if already_contig {
5823 source_arc
5824 } else {
5825 auto_contiguize(&source_arc, &source_layout)?
5826 }
5827 };
5828 let source_layout_kernel = fuel_ir::Layout::contiguous(source_layout.shape().clone());
5829 let position_layout = fuel_ir::Layout::contiguous(fuel_ir::Shape::from_dims(&[]));
5830 let input_arcs = vec![source_arc_contig, position_arc];
5831 let mut output_arcs = vec![dest_arc.clone()];
5832 let kernel_layouts = vec![source_layout_kernel, position_layout, dest_layout.clone()];
5833 let handle = execute_compiled_with_wait_hint(
5835 compiled,
5836 &input_arcs,
5837 &mut output_arcs,
5838 &kernel_layouts,
5839 will_be_waited,
5840 )?;
5841 cache.insert(item.node_id, dest_arc);
5842 layout_cache.insert(item.node_id, item.output_layout.clone());
5843 Ok(handle)
5844 }
5845 WorkItemKind::WriteSliceDoff {
5846 dest,
5847 source,
5848 offset,
5849 } => {
5850 let compiled = item.compiled.as_ref().ok_or_else(|| {
5858 Error::Msg(format!(
5859 "PipelinedExecutor: WriteSliceDoff work item {:?} has no compiled node",
5860 item.node_id,
5861 ))
5862 .bt()
5863 })?;
5864 let dest_arc = cache.get(dest).cloned().ok_or_else(|| {
5865 Error::Msg(format!(
5866 "PipelinedExecutor: WriteSliceDoff destination {:?} of {:?} not realized",
5867 dest, item.node_id,
5868 ))
5869 .bt()
5870 })?;
5871 let source_arc = cache.get(source).cloned().ok_or_else(|| {
5872 Error::Msg(format!(
5873 "PipelinedExecutor: WriteSliceDoff source {:?} of {:?} not realized",
5874 source, item.node_id,
5875 ))
5876 .bt()
5877 })?;
5878 let offset_arc = cache.get(offset).cloned().ok_or_else(|| {
5879 Error::Msg(format!(
5880 "PipelinedExecutor: WriteSliceDoff offset {:?} of {:?} not realized",
5881 offset, item.node_id,
5882 ))
5883 .bt()
5884 })?;
5885 let dest_layout = layout_cache.get(dest).cloned().ok_or_else(|| {
5886 Error::Msg(format!(
5887 "PipelinedExecutor: WriteSliceDoff destination {:?} of {:?} has no cached layout",
5888 dest, item.node_id,
5889 ))
5890 .bt()
5891 })?;
5892 let source_layout = layout_cache.get(source).cloned().ok_or_else(|| {
5893 Error::Msg(format!(
5894 "PipelinedExecutor: WriteSliceDoff source {:?} of {:?} has no cached layout",
5895 source, item.node_id,
5896 ))
5897 .bt()
5898 })?;
5899 if !dest_layout.is_contiguous() || dest_layout.start_offset() != 0 {
5901 return Err(Error::Msg(format!(
5902 "Op::WriteSliceDoff (v1): destination {:?} must be \
5903 contiguous + zero-offset; got Layout {:?}",
5904 dest, dest_layout,
5905 ))
5906 .bt());
5907 }
5908 let source_arc_contig = {
5909 let s_dtype = source_arc
5910 .read()
5911 .map_err(|_| poisoned("WriteSliceDoff source storage"))?
5912 .dtype;
5913 let s_len_bytes = source_arc
5914 .read()
5915 .map_err(|_| poisoned("WriteSliceDoff source storage"))?
5916 .inner
5917 .len_bytes();
5918 let layout_bytes = source_layout.shape().elem_count() * s_dtype.size_in_bytes();
5919 let bytes_match_shape = s_len_bytes == layout_bytes;
5920 let already_contig = source_layout.is_contiguous()
5921 && source_layout.start_offset() == 0
5922 && bytes_match_shape;
5923 if already_contig {
5924 source_arc
5925 } else {
5926 auto_contiguize(&source_arc, &source_layout)?
5927 }
5928 };
5929 let source_layout_kernel = fuel_ir::Layout::contiguous(source_layout.shape().clone());
5930 let offset_layout = fuel_ir::Layout::contiguous(fuel_ir::Shape::from_dims(&[]));
5931 let input_arcs = vec![source_arc_contig, offset_arc];
5932 let mut output_arcs = vec![dest_arc.clone()];
5933 let kernel_layouts = vec![source_layout_kernel, offset_layout, dest_layout.clone()];
5934 let handle = execute_compiled_with_wait_hint(
5936 compiled,
5937 &input_arcs,
5938 &mut output_arcs,
5939 &kernel_layouts,
5940 will_be_waited,
5941 )?;
5942 cache.insert(item.node_id, dest_arc);
5943 layout_cache.insert(item.node_id, item.output_layout.clone());
5944 Ok(handle)
5945 }
5946 WorkItemKind::Alloc { target_location } => {
5947 #[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
5959 let n_bytes = item.elem_count * item.dtype.size_in_bytes();
5960 let alloced: Storage = match target_location {
5961 DeviceLocation::Cpu => fuel_memory::alloc_cpu_zeroed(item.dtype, item.elem_count)?,
5967 #[cfg(feature = "cuda")]
5968 DeviceLocation::Cuda { gpu_id } => {
5969 let cuda_dev = find_cuda_device_in_cache(cache, *gpu_id).ok_or_else(|| {
5970 Error::Msg(format!(
5971 "Op::Alloc on Cuda {{ gpu_id: {} }}: no CUDA \
5972 storage in input cache to derive the device \
5973 handle from. The caller must seed the cache \
5974 (e.g. via `fuel-core::pipelined_bridge::\
5975 device_seed_storage`) before realizing.",
5976 gpu_id,
5977 ))
5978 .bt()
5979 })?;
5980 let cuda_bytes =
5987 fuel_cuda_backend::CudaStorageBytes::alloc_uninit(&cuda_dev, n_bytes)?;
5988 Storage::new(fuel_memory::BackendStorage::Cuda(cuda_bytes), item.dtype)
5989 }
5990 #[cfg(not(feature = "cuda"))]
5991 DeviceLocation::Cuda { .. } => {
5992 return Err(Error::Msg(
5993 "Op::Alloc on Cuda but fuel-storage wasn't built \
5994 with --features cuda"
5995 .to_string(),
5996 )
5997 .bt());
5998 }
5999 #[cfg(feature = "vulkan")]
6000 DeviceLocation::Vulkan { gpu_id } => {
6001 let backend =
6002 find_vulkan_backend_in_cache(cache, *gpu_id).ok_or_else(|| {
6003 Error::Msg(format!(
6004 "Op::Alloc on Vulkan {{ gpu_id: {} }}: no Vulkan \
6005 storage in input cache to derive the backend \
6006 handle from. The caller must seed the cache \
6007 (e.g. via `fuel-core::pipelined_bridge::\
6008 device_seed_storage`) before realizing.",
6009 gpu_id,
6010 ))
6011 .bt()
6012 })?;
6013 let vk_bytes = backend.alloc_bytes_handle(n_bytes)?;
6020 Storage::new(fuel_memory::BackendStorage::Vulkan(vk_bytes), item.dtype)
6021 }
6022 #[cfg(not(feature = "vulkan"))]
6023 DeviceLocation::Vulkan { .. } => {
6024 return Err(Error::Msg(
6025 "Op::Alloc on Vulkan but fuel-storage wasn't built \
6026 with --features vulkan"
6027 .to_string(),
6028 )
6029 .bt());
6030 }
6031 other => {
6032 return Err(Error::Msg(format!(
6033 "Op::Alloc on {:?}: target not yet wired (CPU + \
6034 CUDA + Vulkan land in Phase 3a of bridge-retirement; \
6035 Metal extends when its byte-storage substrate is \
6036 ready).",
6037 other,
6038 ))
6039 .bt());
6040 }
6041 };
6042 cache.insert(item.node_id, Arc::new(RwLock::new(alloced)));
6043 layout_cache.insert(item.node_id, item.output_layout.clone());
6044 Ok(CompletionHandle::Ready)
6045 }
6046 WorkItemKind::Iota { len } => {
6047 if item.target_backend != BackendId::Cpu {
6053 return Err(Error::Msg(format!(
6054 "WorkItemKind::Iota: only CPU is wired today, got {:?}",
6055 item.target_backend,
6056 ))
6057 .bt());
6058 }
6059 let v: Vec<f32> = (0..*len).map(|i| i as f32).collect();
6060 let iota = fuel_memory::from_slice_cpu(&v);
6061 cache.insert(item.node_id, Arc::new(RwLock::new(iota)));
6062 layout_cache.insert(item.node_id, item.output_layout.clone());
6063 Ok(CompletionHandle::Ready)
6064 }
6065 WorkItemKind::ZeroFill => {
6066 if item.inputs.len() != 1 {
6071 return Err(Error::Msg(format!(
6072 "PipelinedExecutor: ZeroFill work item {:?} expects \
6073 1 input, got {}",
6074 item.node_id,
6075 item.inputs.len(),
6076 ))
6077 .bt());
6078 }
6079 let src_id = item.inputs[0];
6080 let src_arc = cache.get(&src_id).cloned().ok_or_else(|| {
6081 Error::Msg(format!(
6082 "PipelinedExecutor: ZeroFill input {:?} of {:?} \
6083 not realized",
6084 src_id, item.node_id,
6085 ))
6086 .bt()
6087 })?;
6088 {
6094 let mut guard = src_arc
6095 .write()
6096 .map_err(|_| poisoned("ZeroFill destination storage"))?;
6097 match &mut guard.inner {
6098 fuel_memory::BackendStorage::Cpu(c) => {
6099 let bytes = c.bytes_mut();
6100 bytes.fill(0);
6101 }
6102 #[cfg(feature = "cuda")]
6103 fuel_memory::BackendStorage::Cuda(c) => {
6104 c.zero_async()?;
6105 }
6106 #[cfg(feature = "vulkan")]
6107 fuel_memory::BackendStorage::Vulkan(v) => {
6108 let backend = v
6109 .backend()
6110 .ok_or_else(|| {
6111 Error::Msg(
6112 "Op::ZeroFill on Vulkan: input has no \
6113 attached backend handle. Storage must \
6114 come from VulkanBackend::alloc_bytes_handle \
6115 / upload_bytes_handle."
6116 .to_string(),
6117 )
6118 .bt()
6119 })?
6120 .clone();
6121 backend.fill_bytes_zero(v)?;
6122 }
6123 #[allow(unreachable_patterns)]
6124 other => {
6125 return Err(Error::Msg(format!(
6126 "Op::ZeroFill: backend not wired ({other:?}); \
6127 CPU + CUDA + Vulkan covered, Metal extends \
6128 when its byte-storage substrate is ready.",
6129 ))
6130 .bt());
6131 }
6132 }
6133 }
6134 cache.insert(item.node_id, src_arc);
6139 layout_cache.insert(item.node_id, item.output_layout.clone());
6140 Ok(CompletionHandle::Ready)
6145 }
6146 WorkItemKind::Copy { target_location } | WorkItemKind::Move { target_location } => {
6147 let op_label = if matches!(item.kind, WorkItemKind::Move { .. }) {
6160 "Move"
6161 } else {
6162 "Copy"
6163 };
6164 let compiled = item.compiled.as_ref().ok_or_else(|| {
6165 Error::Msg(format!(
6166 "PipelinedExecutor: {op_label} work item {:?} has no compiled node",
6167 item.node_id,
6168 ))
6169 .bt()
6170 })?;
6171 if item.inputs.len() != 1 {
6172 return Err(Error::Msg(format!(
6173 "PipelinedExecutor: {op_label} work item {:?} expects 1 input, got {}",
6174 item.node_id,
6175 item.inputs.len(),
6176 ))
6177 .bt());
6178 }
6179 let src_id = item.inputs[0];
6180 let src_arc = cache.get(&src_id).cloned().ok_or_else(|| {
6181 Error::Msg(format!(
6182 "PipelinedExecutor: {op_label} input {:?} of {:?} not realized",
6183 src_id, item.node_id,
6184 ))
6185 .bt()
6186 })?;
6187 {
6205 let guard = src_arc
6206 .read()
6207 .map_err(|_| poisoned("Copy source storage"))?;
6208 if guard.is_bundled() {
6209 let layout_bytes_for_guard = layout_cache
6210 .get(&src_id)
6211 .map(|l| l.shape().elem_count() * guard.dtype.size_in_bytes())
6212 .unwrap_or(usize::MAX);
6213 let bundle_bytes = guard.inner.len_bytes();
6214 let is_whole_bundle = layout_bytes_for_guard >= bundle_bytes;
6215 if is_whole_bundle {
6216 return Err(Error::Msg(format!(
6217 "PipelinedExecutor: Op::{op_label} on a bundled producer \
6218 ({:?}) is not yet supported (Option C followup). \
6219 Use Op::View → Op::Copy → Op::ViewOwned for \
6220 per-slot cross-device moves, or wait for the \
6221 per-backend whole-bundle-Copy hook.",
6222 src_id,
6223 ))
6224 .bt());
6225 }
6226 }
6227 }
6228 let src_layout = layout_cache.get(&src_id).cloned().ok_or_else(|| {
6229 Error::Msg(format!(
6230 "PipelinedExecutor: {op_label} input {:?} of {:?} has no cached layout",
6231 src_id, item.node_id,
6232 ))
6233 .bt()
6234 })?;
6235 let src_dtype = src_arc
6236 .read()
6237 .map_err(|_| poisoned("Copy source storage"))?
6238 .dtype;
6239 let src_len_bytes = src_arc
6240 .read()
6241 .map_err(|_| poisoned("Copy source storage"))?
6242 .inner
6243 .len_bytes();
6244 let layout_bytes = src_layout.shape().elem_count() * src_dtype.size_in_bytes();
6245 let bytes_match_shape = src_len_bytes == layout_bytes;
6246 let already_contig =
6247 src_layout.is_contiguous() && src_layout.start_offset() == 0 && bytes_match_shape;
6248 let src_input = if already_contig {
6249 src_arc
6250 } else {
6251 auto_contiguize(&src_arc, &src_layout)?
6252 };
6253 let kernel_input_layout = fuel_ir::Layout::contiguous(src_layout.shape().clone());
6254 let reuse_arc = match &persistent {
6265 Some(p) if p.mode == PersistentMode::Reuse => match target_location {
6266 DeviceLocation::Cuda { .. } => p.map.get(&item.node_id).cloned(),
6267 _ => None,
6268 },
6269 _ => None,
6270 };
6271 #[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
6285 let n_bytes = item.elem_count * item.dtype.size_in_bytes();
6286 let input_arcs = vec![src_input];
6287 let mut output_arcs = if let Some(arc) = reuse_arc {
6288 vec![arc]
6289 } else {
6290 let output = match target_location {
6291 DeviceLocation::Cpu => {
6292 fuel_memory::alloc_cpu_zeroed(item.dtype, item.elem_count)?
6293 }
6294 #[cfg(feature = "cuda")]
6295 DeviceLocation::Cuda { gpu_id } => {
6296 let cuda_dev =
6297 find_cuda_device_in_cache(cache, *gpu_id).ok_or_else(|| {
6298 Error::Msg(format!(
6299 "Op::{op_label} on Cuda {{ gpu_id: {} }}: no CUDA \
6300 storage in input cache to derive the device \
6301 handle from. The caller must seed the cache \
6302 (e.g. via `fuel-core::pipelined_bridge::\
6303 device_seed_storage`) before realizing an \
6304 H2D Op::{op_label}.",
6305 gpu_id,
6306 ))
6307 .bt()
6308 })?;
6309 let cuda_bytes =
6310 fuel_cuda_backend::CudaStorageBytes::alloc_uninit(&cuda_dev, n_bytes)?;
6311 Storage::new(fuel_memory::BackendStorage::Cuda(cuda_bytes), item.dtype)
6312 }
6313 #[cfg(not(feature = "cuda"))]
6314 DeviceLocation::Cuda { .. } => {
6315 return Err(Error::Msg(format!(
6316 "Op::{op_label} target Cuda but fuel-storage wasn't built \
6317 with --features cuda",
6318 ))
6319 .bt());
6320 }
6321 #[cfg(feature = "vulkan")]
6322 DeviceLocation::Vulkan { gpu_id } => {
6323 let backend =
6324 find_vulkan_backend_in_cache(cache, *gpu_id).ok_or_else(|| {
6325 Error::Msg(format!(
6326 "Op::{op_label} on Vulkan {{ gpu_id: {} }}: no Vulkan \
6327 storage in input cache to derive the backend \
6328 handle from. The caller must seed the cache \
6329 (e.g. via `fuel-core::pipelined_bridge::\
6330 device_seed_storage`) before realizing an \
6331 H2D Op::{op_label}.",
6332 gpu_id,
6333 ))
6334 .bt()
6335 })?;
6336 let vk_bytes = backend.alloc_bytes_handle(n_bytes)?;
6337 Storage::new(fuel_memory::BackendStorage::Vulkan(vk_bytes), item.dtype)
6338 }
6339 #[cfg(not(feature = "vulkan"))]
6340 DeviceLocation::Vulkan { .. } => {
6341 return Err(Error::Msg(format!(
6342 "Op::{op_label} target Vulkan but fuel-storage wasn't built \
6343 with --features vulkan",
6344 ))
6345 .bt());
6346 }
6347 other => {
6348 return Err(Error::Msg(format!(
6349 "PipelinedExecutor: Op::{op_label} target_location {:?} not yet \
6350 wired (CPU + CUDA + Vulkan covered; Metal extends when \
6351 its byte-storage substrate is ready).",
6352 other
6353 ))
6354 .bt());
6355 }
6356 };
6357 vec![Arc::new(RwLock::new(output))]
6358 };
6359 let kernel_layouts = vec![kernel_input_layout, item.output_layout.clone()];
6360 let handle = execute_compiled_with_wait_hint(
6362 compiled,
6363 &input_arcs,
6364 &mut output_arcs,
6365 &kernel_layouts,
6366 will_be_waited,
6367 )?;
6368 let arc = output_arcs.into_iter().next().expect("one output");
6369 if let Some(p) = persistent.as_deref_mut()
6375 && p.mode == PersistentMode::Record
6376 {
6377 p.map.insert(item.node_id, Arc::clone(&arc));
6378 }
6379 cache.insert(item.node_id, arc);
6380 layout_cache.insert(item.node_id, item.output_layout.clone());
6381 Ok(handle)
6382 }
6383 WorkItemKind::ReleaseMarker => {
6384 let marker = fuel_memory::alloc_cpu_zeroed(item.dtype, 0)?;
6391 cache.insert(item.node_id, Arc::new(RwLock::new(marker)));
6392 layout_cache.insert(item.node_id, item.output_layout.clone());
6393 Ok(CompletionHandle::Ready)
6394 }
6395 WorkItemKind::ContiguizeOf { input } => {
6396 let input_arc = cache.get(input).cloned().ok_or_else(|| {
6397 Error::Msg(format!(
6398 "PipelinedExecutor: reshape input {:?} of {:?} not realized",
6399 input, item.node_id,
6400 ))
6401 .bt()
6402 })?;
6403 let input_layout = layout_cache.get(input).cloned().ok_or_else(|| {
6404 Error::Msg(format!(
6405 "PipelinedExecutor: reshape input {:?} of {:?} has no cached layout",
6406 input, item.node_id,
6407 ))
6408 .bt()
6409 })?;
6410 let reuse_arc = match &persistent {
6421 Some(p) if p.mode == PersistentMode::Reuse => p.map.get(&item.node_id).cloned(),
6422 _ => None,
6423 };
6424 let out_arc = if let Some(dest_arc) = reuse_arc {
6425 contiguize_into(&dest_arc, &input_arc, &input_layout)?;
6426 dest_arc
6427 } else if input_layout.is_contiguous() && input_layout.start_offset() == 0 {
6428 input_arc
6432 } else {
6433 let arc = auto_contiguize(&input_arc, &input_layout)?;
6435 if let Some(p) = persistent.as_deref_mut()
6436 && p.mode == PersistentMode::Record
6437 {
6438 p.map.insert(item.node_id, Arc::clone(&arc));
6439 }
6440 arc
6441 };
6442 cache.insert(item.node_id, out_arc);
6443 layout_cache.insert(item.node_id, item.output_layout.clone());
6444 Ok(CompletionHandle::Ready)
6448 }
6449 WorkItemKind::Kernel => {
6450 let compiled = item.compiled.as_ref().ok_or_else(|| {
6451 Error::Msg(format!(
6452 "PipelinedExecutor: Kernel work item {:?} has no compiled node",
6453 item.node_id,
6454 ))
6455 .bt()
6456 })?;
6457 let patched_mm = resolve_deferred_matmul(compiled, produced_syms, item.node_id)?;
6461 let compiled = patched_mm.as_ref().unwrap_or(compiled);
6462 let mut input_arcs: Vec<Arc<RwLock<Storage>>> = Vec::with_capacity(item.inputs.len());
6474 let mut kernel_layouts: Vec<fuel_ir::Layout> =
6475 Vec::with_capacity(item.inputs.len() + 1);
6476 let kernel_handles_strided = compiled.caps.strided_input;
6483 let t_gather = crate::compiled::op_profile::PhaseTimer::start("kernel.gather_inputs");
6484 for in_id in &item.inputs {
6485 let in_arc = cache.get(in_id).cloned().ok_or_else(|| {
6486 Error::Msg(format!(
6487 "PipelinedExecutor: input {:?} of {:?} not realized",
6488 in_id, item.node_id,
6489 ))
6490 .bt()
6491 })?;
6492 let in_layout = layout_cache.get(in_id).cloned().ok_or_else(|| {
6493 Error::Msg(format!(
6494 "PipelinedExecutor: input {:?} of {:?} has no cached layout",
6495 in_id, item.node_id,
6496 ))
6497 .bt()
6498 })?;
6499 let in_dtype = in_arc.read().map_err(|_| poisoned("input storage"))?.dtype;
6516 let in_len_bytes = in_arc
6517 .read()
6518 .map_err(|_| poisoned("input storage"))?
6519 .inner
6520 .len_bytes();
6521 let layout_bytes = in_layout.shape().elem_count() * in_dtype.size_in_bytes();
6522 let bytes_match_shape = in_len_bytes == layout_bytes;
6523 let already_contig =
6524 in_layout.is_contiguous() && in_layout.start_offset() == 0 && bytes_match_shape;
6525 let strided_ok =
6526 kernel_handles_strided && in_layout.start_offset() == 0 && bytes_match_shape;
6527 if already_contig || strided_ok {
6528 input_arcs.push(in_arc);
6529 kernel_layouts.push(in_layout);
6530 } else {
6531 let contig_arc = auto_contiguize(&in_arc, &in_layout)?;
6532 input_arcs.push(contig_arc);
6533 kernel_layouts.push(fuel_ir::Layout::contiguous(in_layout.shape().clone()));
6534 }
6535 }
6536 drop(t_gather);
6537 kernel_layouts.push(item.output_layout.clone());
6538
6539 let (alloc_elem_count, bundle_to_attach) = match &item.output_bundle {
6552 Some(bundle) => {
6553 let dtype_bytes = item.dtype.size_in_bytes().max(1);
6554 let total_bytes = bundle.iter().fold(0usize, |acc, s| {
6555 let end = s.byte_offset.saturating_add(s.len_bytes());
6556 acc.max(end)
6557 });
6558 let elems = total_bytes.div_ceil(dtype_bytes);
6559 (elems, Some(Arc::clone(bundle)))
6560 }
6561 None => (item.elem_count, None),
6562 };
6563 let reuse_arc = match &persistent {
6570 Some(p) if p.mode == PersistentMode::Reuse => {
6571 Some(p.map.get(&item.node_id).cloned().ok_or_else(|| {
6572 Error::Msg(format!(
6573 "PipelinedExecutor: capture (reuse) mode has no persistent \
6574 output buffer for kernel node {:?}; the warm (record) run \
6575 must realize every kernel node before capture",
6576 item.node_id,
6577 ))
6578 .bt()
6579 })?)
6580 }
6581 _ => None,
6582 };
6583 let t_alloc = crate::compiled::op_profile::PhaseTimer::start("kernel.alloc_output");
6584 let mut output_arcs = if let Some(arc) = reuse_arc {
6585 vec![arc]
6586 } else {
6587 let output = match item.target_backend {
6588 BackendId::Cpu => fuel_memory::alloc_cpu_zeroed(item.dtype, alloc_elem_count)?,
6589 #[cfg(feature = "cuda")]
6590 BackendId::Cuda => {
6591 let first_in = input_arcs.first().ok_or_else(|| {
6592 Error::Msg(format!(
6593 "PipelinedExecutor: kernel {:?} on Cuda has no inputs; \
6594 cannot derive device for output allocation",
6595 item.node_id,
6596 ))
6597 .bt()
6598 })?;
6599 let guard = first_in.read().map_err(|_| poisoned("input storage"))?;
6600 let cuda_in = match &guard.inner {
6601 fuel_memory::BackendStorage::Cuda(c) => c,
6602 other => {
6603 return Err(Error::Msg(format!(
6604 "PipelinedExecutor: kernel {:?} target_backend=Cuda but \
6605 input has BackendStorage::{:?}; mixed-backend kernels \
6606 require an explicit Op::Copy first",
6607 item.node_id,
6608 std::mem::discriminant(other),
6609 ))
6610 .bt());
6611 }
6612 };
6613 let n_bytes = alloc_elem_count * item.dtype.size_in_bytes();
6614 let cuda_bytes =
6615 fuel_cuda_backend::CudaStorageBytes::alloc(cuda_in.device(), n_bytes)?;
6616 fuel_memory::Storage::new(
6617 fuel_memory::BackendStorage::Cuda(cuda_bytes),
6618 item.dtype,
6619 )
6620 }
6621 #[cfg(feature = "vulkan")]
6622 BackendId::Vulkan => {
6623 let first_in = input_arcs.first().ok_or_else(|| {
6631 Error::Msg(format!(
6632 "PipelinedExecutor: kernel {:?} on Vulkan has no inputs; \
6633 cannot derive backend handle for output allocation",
6634 item.node_id,
6635 ))
6636 .bt()
6637 })?;
6638 let guard = first_in.read().map_err(|_| poisoned("input storage"))?;
6639 let vk_in = match &guard.inner {
6640 fuel_memory::BackendStorage::Vulkan(v) => v,
6641 other => {
6642 return Err(Error::Msg(format!(
6643 "PipelinedExecutor: kernel {:?} target_backend=Vulkan but \
6644 input has BackendStorage::{:?}; mixed-backend kernels \
6645 require an explicit Op::Copy first",
6646 item.node_id,
6647 std::mem::discriminant(other),
6648 ))
6649 .bt());
6650 }
6651 };
6652 let backend = vk_in.backend().ok_or_else(|| {
6653 Error::Msg(format!(
6654 "PipelinedExecutor: Vulkan kernel {:?} input has no backend \
6655 handle. Storages flowing through the pipelined executor's \
6656 Vulkan path must be constructed via \
6657 VulkanBackend::alloc_bytes_handle / upload_bytes_handle \
6658 (not the legacy alloc_bytes / upload_bytes which leave \
6659 VulkanStorageBytes::backend = None).",
6660 item.node_id,
6661 ))
6662 .bt()
6663 })?;
6664 let n_bytes = alloc_elem_count * item.dtype.size_in_bytes();
6665 let vk_bytes = backend.alloc_bytes_handle(n_bytes)?;
6666 fuel_memory::Storage::new(
6667 fuel_memory::BackendStorage::Vulkan(vk_bytes),
6668 item.dtype,
6669 )
6670 }
6671 other => {
6672 return Err(Error::Msg(format!(
6673 "PipelinedExecutor: target_backend {:?} output allocation \
6674 not yet implemented (CPU + CUDA + Vulkan wired; Metal extends later)",
6675 other
6676 ))
6677 .bt());
6678 }
6679 };
6680 let output = match bundle_to_attach {
6685 Some(bundle) => output.with_bundle(bundle)?,
6686 None => output,
6687 };
6688 vec![Arc::new(RwLock::new(output))]
6689 };
6690 drop(t_alloc);
6691
6692 let t_launch =
6694 crate::compiled::op_profile::PhaseTimer::start("kernel.execute_compiled");
6695 let handle = execute_compiled_with_wait_hint(
6696 compiled,
6697 &input_arcs,
6698 &mut output_arcs,
6699 &kernel_layouts,
6700 will_be_waited,
6701 )?;
6702 drop(t_launch);
6703
6704 let arc = output_arcs.into_iter().next().expect("one output");
6705
6706 if let Some(p) = persistent
6711 && p.mode == PersistentMode::Record
6712 {
6713 p.map.insert(item.node_id, Arc::clone(&arc));
6714 }
6715
6716 if let OpParams::NonZeroIndices {
6726 capacity,
6727 count_sym,
6728 } = &compiled.op_params
6729 {
6730 bind_data_determined_count(&arc, *capacity, *count_sym, produced_syms)?;
6731 }
6732
6733 cache.insert(item.node_id, arc);
6734 layout_cache.insert(item.node_id, item.output_layout.clone());
6735 Ok(handle)
6736 }
6737 WorkItemKind::InplaceKernel { target_idx } => {
6738 let compiled = item.compiled.as_ref().ok_or_else(|| {
6739 Error::Msg(format!(
6740 "PipelinedExecutor: InplaceKernel work item {:?} has no compiled node",
6741 item.node_id,
6742 ))
6743 .bt()
6744 })?;
6745 let target_id = *item.inputs.get(*target_idx).ok_or_else(|| {
6746 Error::Msg(format!(
6747 "PipelinedExecutor: InplaceKernel work item {:?} target_idx={} \
6748 out of bounds (inputs.len()={})",
6749 item.node_id,
6750 target_idx,
6751 item.inputs.len(),
6752 ))
6753 .bt()
6754 })?;
6755 let target_arc = cache.get(&target_id).cloned().ok_or_else(|| {
6756 Error::Msg(format!(
6757 "PipelinedExecutor: InplaceKernel target {:?} of {:?} not realized",
6758 target_id, item.node_id,
6759 ))
6760 .bt()
6761 })?;
6762 let target_layout = layout_cache.get(&target_id).cloned().ok_or_else(|| {
6763 Error::Msg(format!(
6764 "PipelinedExecutor: InplaceKernel target {:?} of {:?} has no cached layout",
6765 target_id, item.node_id,
6766 ))
6767 .bt()
6768 })?;
6769 if !target_layout.is_contiguous() || target_layout.start_offset() != 0 {
6774 return Err(Error::Msg(format!(
6775 "InplaceKernel (Phase 3 v1): target {:?} must be contiguous + \
6776 zero-offset; got Layout {:?}",
6777 target_id, target_layout,
6778 ))
6779 .bt());
6780 }
6781 let mut input_arcs: Vec<Arc<RwLock<Storage>>> = Vec::new();
6791 let mut kernel_layouts: Vec<Layout> = Vec::new();
6792 for (i, &input_id) in item.inputs.iter().enumerate() {
6793 if i == *target_idx {
6794 continue;
6795 }
6796 let arc = cache.get(&input_id).cloned().ok_or_else(|| {
6797 Error::Msg(format!(
6798 "PipelinedExecutor: InplaceKernel non-target input {:?} of {:?} \
6799 not realized",
6800 input_id, item.node_id,
6801 ))
6802 .bt()
6803 })?;
6804 let lay = layout_cache.get(&input_id).cloned().ok_or_else(|| {
6805 Error::Msg(format!(
6806 "PipelinedExecutor: InplaceKernel non-target input {:?} of {:?} \
6807 has no cached layout",
6808 input_id, item.node_id,
6809 ))
6810 .bt()
6811 })?;
6812 input_arcs.push(arc);
6813 kernel_layouts.push(lay);
6814 }
6815 let mut output_arcs = vec![target_arc.clone()];
6816 kernel_layouts.push(target_layout.clone());
6817 let handle = execute_compiled_with_wait_hint(
6819 compiled,
6820 &input_arcs,
6821 &mut output_arcs,
6822 &kernel_layouts,
6823 will_be_waited,
6824 )?;
6825 cache.insert(item.node_id, target_arc);
6829 layout_cache.insert(item.node_id, item.output_layout.clone());
6830 Ok(handle)
6831 }
6832 }
6833}
6834
6835fn poisoned(what: &'static str) -> Error {
6836 Error::Msg(format!("PipelinedExecutor: {} poisoned", what)).bt()
6837}
6838
6839#[inline]
6848fn store_handle(
6849 handles: &mut HashMap<NodeId, CompletionHandle>,
6850 node_id: NodeId,
6851 handle: CompletionHandle,
6852) {
6853 if matches!(handle, CompletionHandle::Ready) {
6854 if let Some(prev) = handles.remove(&node_id) {
6858 let _ = prev.wait();
6859 }
6860 return;
6861 }
6862 if let Some(prev) = handles.insert(node_id, handle) {
6863 let _ = prev.wait();
6864 }
6865}
6866
6867fn wait_producer_handle(
6894 handles: &mut HashMap<NodeId, CompletionHandle>,
6895 producer: NodeId,
6896) -> Result<()> {
6897 if let Some(handle) = handles.remove(&producer) {
6898 handle.wait()?;
6899 }
6900 Ok(())
6901}
6902
6903fn drain_handles(handles: &mut HashMap<NodeId, CompletionHandle>) -> Result<()> {
6930 for (_node, handle) in handles.drain() {
6931 handle.wait()?;
6932 }
6933 debug_assert!(
6939 handles.is_empty(),
6940 "A4b-1: handle map must be empty after drain"
6941 );
6942 Ok(())
6943}
6944
6945#[cfg(feature = "cuda")]
6974fn sync_active_cuda_devices(cache: &StorageCache) -> Result<()> {
6975 let mut devices: HashMap<fuel_ir::DeviceLocation, fuel_cuda_backend::CudaDevice> =
6976 HashMap::new();
6977 for arc in cache.values() {
6978 let guard = arc
6979 .read()
6980 .map_err(|_| poisoned("cache storage (cuda realize-end sync scan)"))?;
6981 if let fuel_memory::BackendStorage::Cuda(c) = &guard.inner {
6982 devices
6983 .entry(c.device().location())
6984 .or_insert_with(|| c.device().clone());
6985 }
6986 }
6987 for dev in devices.values() {
6988 dev.synchronize()?;
6989 }
6990 Ok(())
6991}
6992
6993#[cfg(not(feature = "cuda"))]
6996fn sync_active_cuda_devices(_cache: &StorageCache) -> Result<()> {
6997 Ok(())
6998}
6999
7000#[cfg(feature = "cuda")]
7011fn find_cuda_device_in_cache(
7012 cache: &StorageCache,
7013 gpu_id: usize,
7014) -> Option<fuel_cuda_backend::CudaDevice> {
7015 for arc in cache.values() {
7016 let guard = arc.read().ok()?;
7017 if let fuel_memory::BackendStorage::Cuda(c) = &guard.inner {
7018 if matches!(c.device().location(), DeviceLocation::Cuda { gpu_id: g } if g == gpu_id) {
7023 return Some(c.device().clone());
7024 }
7025 }
7026 }
7027 None
7028}
7029
7030#[cfg(feature = "vulkan")]
7055struct VulkanCompletion {
7056 backend: Arc<fuel_vulkan_backend::VulkanBackend>,
7057 batch: Option<fuel_vulkan_backend::SubmittedBatch>,
7058 loc: fuel_ir::DeviceLocation,
7061}
7062
7063#[cfg(feature = "vulkan")]
7069type InflightVulkan = VulkanCompletion;
7070#[cfg(not(feature = "vulkan"))]
7071type InflightVulkan = ();
7072
7073#[cfg(feature = "vulkan")]
7086fn copy_source_is_vulkan(cache: &StorageCache, producer: NodeId) -> Result<bool> {
7087 if let Some(arc) = cache.get(&producer) {
7088 let guard = arc
7089 .read()
7090 .map_err(|_| poisoned("storage lock probing copy-source backend"))?;
7091 return Ok(matches!(
7092 guard.inner,
7093 fuel_memory::BackendStorage::Vulkan(_)
7094 ));
7095 }
7096 Ok(false)
7097}
7098
7099#[cfg(not(feature = "vulkan"))]
7100#[inline]
7101fn copy_source_is_vulkan(_cache: &StorageCache, _producer: NodeId) -> Result<bool> {
7102 Ok(false)
7103}
7104
7105#[cfg(feature = "vulkan")]
7106impl crate::compiled::Completion for VulkanCompletion {
7107 fn wait(self: Box<Self>) -> Result<()> {
7108 (*self).into_wait()
7111 }
7112}
7113
7114#[cfg(feature = "vulkan")]
7115impl VulkanCompletion {
7116 fn new(
7120 backend: Arc<fuel_vulkan_backend::VulkanBackend>,
7121 batch: fuel_vulkan_backend::SubmittedBatch,
7122 loc: fuel_ir::DeviceLocation,
7123 ) -> Self {
7124 crate::dispatch::inflight_inc(loc);
7125 Self {
7126 backend,
7127 batch: Some(batch),
7128 loc,
7129 }
7130 }
7131
7132 fn into_wait(mut self) -> Result<()> {
7143 if let Some(batch) = self.batch.take() {
7148 self.backend.wait_submitted(batch)?;
7149 }
7150 Ok(())
7151 }
7153
7154 fn retain_data(&mut self, buf: std::sync::Arc<fuel_vulkan_backend::VulkanBuffer>) {
7165 if let Some(batch) = self.batch.as_mut() {
7166 batch.retain_buffer(buf);
7167 }
7168 }
7169}
7170
7171#[cfg(feature = "vulkan")]
7172impl Drop for VulkanCompletion {
7173 fn drop(&mut self) {
7174 crate::dispatch::inflight_dec(self.loc);
7179 }
7180}
7181
7182#[cfg(feature = "vulkan")]
7209fn eager_submit_all_vulkan(
7210 cache: &StorageCache,
7211 inflight: &mut Vec<VulkanCompletion>,
7212) -> Result<()> {
7213 let mut seen: Vec<usize> = Vec::new();
7215 let mut backends: Vec<Arc<fuel_vulkan_backend::VulkanBackend>> = Vec::new();
7216 for arc in cache.values() {
7217 let guard = arc
7218 .read()
7219 .map_err(|_| poisoned("storage lock during eager vulkan submit"))?;
7220 if let fuel_memory::BackendStorage::Vulkan(v) = &guard.inner
7221 && let Some(backend) = v.backend()
7222 && !seen.contains(&backend.gpu_id)
7223 {
7224 seen.push(backend.gpu_id);
7225 backends.push(backend.clone());
7226 }
7227 }
7228 for backend in backends {
7229 let loc = DeviceLocation::Vulkan {
7233 gpu_id: backend.gpu_id,
7234 };
7235 if let Some(batch) = backend.submit_pending()? {
7236 inflight.push(VulkanCompletion::new(backend, batch, loc));
7237 }
7238 }
7239 Ok(())
7240}
7241
7242#[cfg(not(feature = "vulkan"))]
7243#[inline]
7244fn eager_submit_all_vulkan(
7245 _cache: &StorageCache,
7246 _inflight: &mut Vec<InflightVulkan>,
7247) -> Result<()> {
7248 Ok(())
7249}
7250
7251#[cfg(feature = "vulkan")]
7268fn drain_inflight_vulkan(inflight: &mut Vec<VulkanCompletion>) -> Result<()> {
7269 for vc in inflight.drain(..) {
7270 vc.into_wait()?;
7271 }
7272 Ok(())
7273}
7274
7275#[cfg(not(feature = "vulkan"))]
7276#[inline]
7277fn drain_inflight_vulkan(_inflight: &mut Vec<InflightVulkan>) -> Result<()> {
7278 Ok(())
7279}
7280
7281#[cfg(feature = "vulkan")]
7297fn drain_vulkan_pending(cache: &StorageCache) -> Result<()> {
7298 use crate::compiled::CompletionHandle;
7299 for arc in cache.values() {
7300 let backend = {
7301 let guard = arc
7302 .read()
7303 .map_err(|_| poisoned("storage lock during vulkan submit_pending"))?;
7304 match &guard.inner {
7305 fuel_memory::BackendStorage::Vulkan(v) => v.backend().cloned(),
7306 _ => None,
7307 }
7308 };
7309 let Some(backend) = backend else { continue };
7310 let loc = DeviceLocation::Vulkan {
7316 gpu_id: backend.gpu_id,
7317 };
7318 if let Some(batch) = backend.submit_pending()? {
7319 let handle: CompletionHandle =
7320 CompletionHandle::Pending(Box::new(VulkanCompletion::new(backend, batch, loc)));
7321 handle.wait()?;
7322 }
7323 }
7324 Ok(())
7325}
7326
7327#[cfg(not(feature = "vulkan"))]
7328fn drain_vulkan_pending(_cache: &StorageCache) -> Result<()> {
7329 Ok(())
7330}
7331
7332#[cfg(feature = "vulkan")]
7349fn force_flush_vulkan(arc: &Arc<RwLock<Storage>>) -> Result<()> {
7350 let guard = arc
7351 .read()
7352 .map_err(|_| poisoned("storage lock during vulkan force_flush"))?;
7353 if let fuel_memory::BackendStorage::Vulkan(v) = &guard.inner
7354 && let Some(backend) = v.backend()
7355 {
7356 backend.force_flush()?;
7357 }
7358 Ok(())
7359}
7360
7361#[cfg(not(feature = "vulkan"))]
7362fn force_flush_vulkan(_arc: &Arc<RwLock<Storage>>) -> Result<()> {
7363 Ok(())
7364}
7365
7366#[cfg(feature = "vulkan")]
7409const EVICT_SUBMIT_MIN_DISPATCHES: usize = 48;
7410
7411#[cfg(feature = "vulkan")]
7412fn defer_evicted_vulkan_buffer(
7413 cache: &StorageCache,
7414 destroyed: NodeId,
7415 inflight: &mut Vec<VulkanCompletion>,
7416) -> Result<()> {
7417 let (buf, backend) = {
7422 let Some(arc) = cache.get(&destroyed) else {
7423 return Ok(());
7424 };
7425 let guard = arc
7426 .read()
7427 .map_err(|_| poisoned("storage lock during deferred-eviction retain"))?;
7428 match &guard.inner {
7429 fuel_memory::BackendStorage::Vulkan(v) => (v.device_buffer_arc(), v.backend().cloned()),
7430 _ => (None, None),
7431 }
7432 };
7433 let Some(buf) = buf else { return Ok(()) };
7434
7435 if let Some(backend) = &backend {
7440 backend.retain_in_open_batch(buf.clone());
7441 }
7442 for vc in inflight.iter_mut() {
7447 vc.retain_data(buf.clone());
7448 }
7449 if let Some(backend) = backend
7451 && backend.open_batch_len() >= EVICT_SUBMIT_MIN_DISPATCHES
7452 {
7453 eager_submit_all_vulkan(cache, inflight)?;
7454 }
7455 if !inflight.is_empty() {
7465 DEFERRED_EVICTION_RETAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7466 }
7467 Ok(())
7468}
7469
7470#[cfg(feature = "vulkan")]
7476pub static DEFERRED_EVICTION_RETAINS: std::sync::atomic::AtomicUsize =
7477 std::sync::atomic::AtomicUsize::new(0);
7478
7479#[cfg(feature = "vulkan")]
7483pub fn deferred_eviction_retains() -> usize {
7484 DEFERRED_EVICTION_RETAINS.load(std::sync::atomic::Ordering::Relaxed)
7485}
7486
7487#[cfg(not(feature = "vulkan"))]
7488#[inline]
7489fn defer_evicted_vulkan_buffer(
7490 _cache: &StorageCache,
7491 _destroyed: NodeId,
7492 _inflight: &mut Vec<InflightVulkan>,
7493) -> Result<()> {
7494 Ok(())
7495}
7496
7497#[cfg(feature = "vulkan")]
7501fn find_vulkan_backend_in_cache(
7502 cache: &StorageCache,
7503 gpu_id: usize,
7504) -> Option<std::sync::Arc<fuel_vulkan_backend::VulkanBackend>> {
7505 for arc in cache.values() {
7506 let guard = arc.read().ok()?;
7507 if let fuel_memory::BackendStorage::Vulkan(v) = &guard.inner
7508 && let Some(backend) = v.backend()
7509 && backend.gpu_id == gpu_id
7510 {
7511 return Some(backend.clone());
7512 }
7513 }
7514 None
7515}
7516
7517fn auto_contiguize(arc: &Arc<RwLock<Storage>>, layout: &Layout) -> Result<Arc<RwLock<Storage>>> {
7527 let in_guard = arc
7528 .read()
7529 .map_err(|_| poisoned("input storage lock during auto_contiguize"))?;
7530 let dtype = in_guard.dtype;
7531 let dtype_size = dtype.size_in_bytes();
7532 let new_storage = match &in_guard.inner {
7533 fuel_memory::BackendStorage::Cpu(c) => {
7534 let new_bytes = fuel_cpu_backend::byte_kernels::contiguize_cpu(c, layout, dtype_size)?;
7535 Storage::new(fuel_memory::BackendStorage::Cpu(new_bytes), dtype)
7536 }
7537 #[cfg(feature = "cuda")]
7538 fuel_memory::BackendStorage::Cuda(c) => {
7539 let contig = fuel_cuda_backend::baracuda::contiguize::contiguize_to_fresh(
7551 c, layout, dtype_size,
7552 )?;
7553 Storage::new(fuel_memory::BackendStorage::Cuda(contig), dtype)
7554 }
7555 #[cfg(feature = "vulkan")]
7556 fuel_memory::BackendStorage::Vulkan(v) => {
7557 let backend = v.backend().ok_or_else(|| {
7562 Error::Msg(
7563 "auto_contiguize: Vulkan input has no backend handle. \
7564 Storages flowing through the pipelined executor's Vulkan \
7565 path must be constructed via VulkanBackend::alloc_bytes_handle \
7566 / upload_bytes_handle (not the legacy alloc_bytes / \
7567 upload_bytes which leave VulkanStorageBytes::backend = None)."
7568 .to_string(),
7569 )
7570 .bt()
7571 })?;
7572 if let Some(native) = crate::vulkan_dispatch::contiguize_vulkan(v, layout, dtype_size)?
7573 {
7574 return Ok(Arc::new(RwLock::new(Storage::new(
7575 fuel_memory::BackendStorage::Vulkan(native),
7576 dtype,
7577 ))));
7578 }
7579 let _t = crate::compiled::op_profile::PhaseTimer::start(
7580 "auto_contiguize.vulkan_host_roundtrip",
7581 );
7582 let host_bytes = backend.download_bytes(v)?;
7583 let host_cpu = fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(&host_bytes);
7584 let contig_cpu =
7585 fuel_cpu_backend::byte_kernels::contiguize_cpu(&host_cpu, layout, dtype_size)?;
7586 let host_contig_bytes = contig_cpu.bytes().to_vec();
7587 let vk_bytes = backend.upload_bytes_handle(&host_contig_bytes)?;
7588 Storage::new(fuel_memory::BackendStorage::Vulkan(vk_bytes), dtype)
7589 }
7590 #[allow(unreachable_patterns)]
7591 _ => {
7592 return Err(Error::Msg(
7593 "auto_contiguize: backend not wired (CPU + CUDA + Vulkan are \
7594 wired today; Metal extends this match when its byte-storage \
7595 substrate lands)"
7596 .to_string(),
7597 )
7598 .bt());
7599 }
7600 };
7601 Ok(Arc::new(RwLock::new(new_storage)))
7602}
7603
7604#[cfg_attr(not(feature = "cuda"), allow(unused_variables))]
7621fn contiguize_into(
7622 dest: &Arc<RwLock<Storage>>,
7623 arc: &Arc<RwLock<Storage>>,
7624 layout: &Layout,
7625) -> Result<()> {
7626 let in_guard = arc
7627 .read()
7628 .map_err(|_| poisoned("input storage lock during contiguize_into"))?;
7629 match &in_guard.inner {
7630 #[cfg(feature = "cuda")]
7631 fuel_memory::BackendStorage::Cuda(src_cuda) => {
7632 let dtype_size = in_guard.dtype.size_in_bytes();
7633 let mut dest_guard = dest
7634 .write()
7635 .map_err(|_| poisoned("dest storage lock during contiguize_into"))?;
7636 let dest_cuda = match &mut dest_guard.inner {
7637 fuel_memory::BackendStorage::Cuda(d) => d,
7638 other => {
7639 return Err(Error::Msg(format!(
7640 "contiguize_into: dest is BackendStorage::{:?}, not Cuda",
7641 std::mem::discriminant(other),
7642 ))
7643 .bt());
7644 }
7645 };
7646 fuel_cuda_backend::baracuda::contiguize::contiguize_into(
7647 src_cuda, layout, dest_cuda, dtype_size,
7648 )
7649 }
7650 #[allow(unreachable_patterns)]
7651 _ => Err(Error::Msg(
7652 "contiguize_into: only CUDA is wired (PersistentOutputs::reuse / \
7653 capture_decode is CUDA-only today — no CPU/Vulkan consumer exists \
7654 yet; extend this match if one appears)"
7655 .to_string(),
7656 )
7657 .bt()),
7658 }
7659}
7660
7661#[cfg(test)]
7662mod backend_neutral_tests {
7663 use super::*;
7664
7665 #[test]
7666 fn only_deviceless_items_are_backend_neutral() {
7667 let n = NodeId(0);
7668 assert!(is_backend_neutral(&WorkItemKind::ConstAdopt));
7670 assert!(is_backend_neutral(&WorkItemKind::ViewOf { input: n }));
7671 assert!(is_backend_neutral(&WorkItemKind::SlotView { producer: n }));
7672 assert!(is_backend_neutral(&WorkItemKind::ContiguizeOf { input: n }));
7675 assert!(!is_backend_neutral(&WorkItemKind::Kernel));
7679 assert!(!is_backend_neutral(&WorkItemKind::Copy {
7680 target_location: DeviceLocation::Cpu
7681 }));
7682 assert!(!is_backend_neutral(&WorkItemKind::WriteSlice {
7683 dest: n,
7684 source: n
7685 }));
7686 }
7687}
7688
7689#[cfg(test)]
7690mod tests {
7691 use super::*;
7692 use fuel_graph::Node;
7693 use fuel_ir::Shape;
7694
7695 #[test]
7701 fn scan_has_no_dispatch_op_kind() {
7702 use fuel_graph::{Op, ScanEmit};
7703 assert_eq!(
7704 crate::pipelined::op_to_op_kind(&Op::Scan {
7705 n_carries: 1,
7706 n_xs: 0,
7707 bound: 2,
7708 emit: ScanEmit::All,
7709 early_exit: None
7710 }),
7711 None,
7712 );
7713 }
7714
7715 #[test]
7721 fn build_wait_set_collects_only_cross_device_copy_producers() {
7722 let mut g = Graph::new();
7723 let const_id = g.push(Node {
7724 op: Op::Const,
7725 inputs: vec![],
7726 shape: Shape::from_dims(&[4]),
7727 dtype: DType::F32,
7728 });
7729 let same_device_id = g.push(Node {
7731 op: Op::Contiguize,
7732 inputs: vec![const_id],
7733 shape: Shape::from_dims(&[4]),
7734 dtype: DType::F32,
7735 });
7736 let copy_id = g.push(Node {
7738 op: Op::Copy {
7739 target: DeviceLocation::Cuda { gpu_id: 0 },
7740 },
7741 inputs: vec![same_device_id],
7742 shape: Shape::from_dims(&[4]),
7743 dtype: DType::F32,
7744 });
7745 let order = vec![const_id, same_device_id, copy_id];
7746
7747 let set = build_wait_set(&g, &order);
7748 assert!(
7749 set.contains(&same_device_id),
7750 "the Copy's producer must be in the wait set"
7751 );
7752 assert!(
7753 !set.contains(&const_id),
7754 "a node with no Copy/Move consumer must be absent from the wait set"
7755 );
7756 assert!(
7757 !set.contains(©_id),
7758 "the Copy node itself is not its own producer"
7759 );
7760 assert_eq!(
7761 set.len(),
7762 1,
7763 "only the one true cross-device producer is in the set"
7764 );
7765 }
7766
7767 #[test]
7769 fn build_wait_set_collects_move_producers_too() {
7770 let mut g = Graph::new();
7771 let const_id = g.push(Node {
7772 op: Op::Const,
7773 inputs: vec![],
7774 shape: Shape::from_dims(&[4]),
7775 dtype: DType::F32,
7776 });
7777 let move_id = g.push(Node {
7778 op: Op::Move {
7779 target: DeviceLocation::Cpu,
7780 },
7781 inputs: vec![const_id],
7782 shape: Shape::from_dims(&[4]),
7783 dtype: DType::F32,
7784 });
7785 let order = vec![const_id, move_id];
7786
7787 let set = build_wait_set(&g, &order);
7788 assert!(
7789 set.contains(&const_id),
7790 "the Move's producer must be in the wait set"
7791 );
7792 assert_eq!(set.len(), 1);
7793 }
7794
7795 #[test]
7798 fn build_wait_set_is_empty_for_cpu_only_order() {
7799 let mut g = Graph::new();
7800 let const_id = g.push(Node {
7801 op: Op::Const,
7802 inputs: vec![],
7803 shape: Shape::from_dims(&[4]),
7804 dtype: DType::F32,
7805 });
7806 let contig_id = g.push(Node {
7807 op: Op::Contiguize,
7808 inputs: vec![const_id],
7809 shape: Shape::from_dims(&[4]),
7810 dtype: DType::F32,
7811 });
7812 let order = vec![const_id, contig_id];
7813
7814 let set = build_wait_set(&g, &order);
7815 assert!(
7816 set.is_empty(),
7817 "a CPU-only order (no Copy/Move) yields an empty wait set"
7818 );
7819 }
7820
7821 #[test]
7825 fn will_be_waited_none_is_conservative_true_for_any_node() {
7826 let ws: WaitSet = None;
7827 assert!(will_be_waited(&ws, NodeId(0)));
7828 assert!(will_be_waited(&ws, NodeId(12345)));
7829 }
7830
7831 #[test]
7832 fn will_be_waited_some_checks_membership() {
7833 let mut set = HashSet::new();
7834 set.insert(NodeId(3));
7835 let ws: WaitSet = Some(set);
7836 assert!(will_be_waited(&ws, NodeId(3)), "member must be waited");
7837 assert!(!will_be_waited(&ws, NodeId(4)), "non-member must be elided");
7838 }
7839
7840 struct FlagCompletion(std::sync::Arc<std::sync::atomic::AtomicUsize>);
7843 impl crate::compiled::Completion for FlagCompletion {
7844 fn wait(self: Box<Self>) -> Result<()> {
7845 self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
7846 Ok(())
7847 }
7848 }
7849
7850 #[test]
7854 fn handle_map_stores_pending_and_drains_once() {
7855 use std::sync::Arc as StdArc;
7856 use std::sync::atomic::{AtomicUsize, Ordering};
7857
7858 let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
7859
7860 store_handle(&mut handles, NodeId(1), CompletionHandle::Ready);
7862 assert!(handles.is_empty(), "Ready handles must not be stored");
7863
7864 let waited = StdArc::new(AtomicUsize::new(0));
7866 store_handle(
7867 &mut handles,
7868 NodeId(2),
7869 CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
7870 );
7871 store_handle(
7872 &mut handles,
7873 NodeId(3),
7874 CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
7875 );
7876 assert_eq!(handles.len(), 2);
7877
7878 drain_handles(&mut handles).expect("drain");
7879 assert_eq!(
7880 waited.load(Ordering::SeqCst),
7881 2,
7882 "each pending handle waited once"
7883 );
7884 assert!(handles.is_empty(), "map empty after drain");
7885 }
7886
7887 #[test]
7891 fn handle_map_rekey_waits_previous() {
7892 use std::sync::Arc as StdArc;
7893 use std::sync::atomic::{AtomicUsize, Ordering};
7894
7895 let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
7896 let waited = StdArc::new(AtomicUsize::new(0));
7897
7898 store_handle(
7899 &mut handles,
7900 NodeId(7),
7901 CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
7902 );
7903 store_handle(
7905 &mut handles,
7906 NodeId(7),
7907 CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
7908 );
7909 assert_eq!(
7910 waited.load(Ordering::SeqCst),
7911 1,
7912 "rekey waited the prior pending"
7913 );
7914
7915 store_handle(&mut handles, NodeId(7), CompletionHandle::Ready);
7917 assert_eq!(
7918 waited.load(Ordering::SeqCst),
7919 2,
7920 "Ready-over-Pending waited it"
7921 );
7922 assert!(handles.is_empty());
7923
7924 drain_handles(&mut handles).expect("drain empty");
7925 }
7926
7927 #[test]
7934 fn wait_producer_handle_consumes_named_handle_once() {
7935 use std::sync::Arc as StdArc;
7936 use std::sync::atomic::{AtomicUsize, Ordering};
7937
7938 let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
7939 let waited = StdArc::new(AtomicUsize::new(0));
7940
7941 store_handle(
7943 &mut handles,
7944 NodeId(2),
7945 CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
7946 );
7947 store_handle(
7948 &mut handles,
7949 NodeId(5),
7950 CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
7951 );
7952
7953 wait_producer_handle(&mut handles, NodeId(99)).expect("absent is no-op");
7955 assert_eq!(
7956 waited.load(Ordering::SeqCst),
7957 0,
7958 "absent producer waited nothing"
7959 );
7960 assert_eq!(handles.len(), 2);
7961
7962 wait_producer_handle(&mut handles, NodeId(2)).expect("wait source");
7964 assert_eq!(
7965 waited.load(Ordering::SeqCst),
7966 1,
7967 "source producer waited once"
7968 );
7969 assert!(!handles.contains_key(&NodeId(2)), "source handle consumed");
7970 assert!(
7971 handles.contains_key(&NodeId(5)),
7972 "unrelated handle untouched"
7973 );
7974
7975 wait_producer_handle(&mut handles, NodeId(2)).expect("re-wait is no-op");
7977 assert_eq!(waited.load(Ordering::SeqCst), 1, "no double-wait");
7978
7979 drain_handles(&mut handles).expect("drain");
7981 assert_eq!(
7982 waited.load(Ordering::SeqCst),
7983 2,
7984 "drain waited the remaining one"
7985 );
7986 assert!(handles.is_empty());
7987 }
7988
7989 #[test]
7992 fn op_contiguize_zero_copy_on_contiguous_input() {
7993 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
7994
7995 let graph = Arc::new(RwLock::new(Graph::new()));
7996 let (const_id, contig_id) = {
7997 let mut g = graph.write().unwrap();
7998 let const_id = g.push(Node {
7999 op: Op::Const,
8000 inputs: vec![],
8001 shape: Shape::from_dims(&[2, 2]),
8002 dtype: DType::F32,
8003 });
8004 let contig_id = g.push(Node {
8005 op: Op::Contiguize,
8006 inputs: vec![const_id],
8007 shape: Shape::from_dims(&[2, 2]),
8008 dtype: DType::F32,
8009 });
8010 (const_id, contig_id)
8013 };
8014
8015 let mut inputs = StorageCache::new();
8016 let input_arc = Arc::new(RwLock::new(storage));
8017 let input_arc_for_compare = Arc::clone(&input_arc);
8018 inputs.insert(const_id, input_arc);
8019
8020 let (result_arc, _) =
8021 PipelinedExecutor::realize(graph, contig_id, inputs).expect("realize");
8022
8023 assert!(
8025 Arc::ptr_eq(&result_arc, &input_arc_for_compare),
8026 "Op::Contiguize on contiguous input must adopt the input Arc (zero copy)",
8027 );
8028 }
8029
8030 #[test]
8035 fn op_contiguize_materializes_on_strided_input() {
8036 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
8042
8043 let graph = Arc::new(RwLock::new(Graph::new()));
8044 let (const_id, _transpose_id, contig_id) = {
8045 let mut g = graph.write().unwrap();
8046 let const_id = g.push(Node {
8047 op: Op::Const,
8048 inputs: vec![],
8049 shape: Shape::from_dims(&[2, 3]),
8050 dtype: DType::F32,
8051 });
8052 let transpose_id = g.push(Node {
8053 op: Op::Transpose,
8054 inputs: vec![const_id],
8055 shape: Shape::from_dims(&[3, 2]),
8056 dtype: DType::F32,
8057 });
8058 let contig_id = g.push(Node {
8059 op: Op::Contiguize,
8060 inputs: vec![transpose_id],
8061 shape: Shape::from_dims(&[3, 2]),
8062 dtype: DType::F32,
8063 });
8064 (const_id, transpose_id, contig_id)
8065 };
8066
8067 let mut inputs = StorageCache::new();
8068 inputs.insert(const_id, Arc::new(RwLock::new(storage)));
8069
8070 let (result_arc, _) =
8071 PipelinedExecutor::realize(graph, contig_id, inputs).expect("realize");
8072
8073 let guard = result_arc.read().unwrap();
8074 let cpu = if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
8075 c
8076 } else {
8077 panic!("expected Cpu storage");
8078 };
8079 let typed: &[f32] = cpu.as_slice().expect("f32 cast");
8080 assert_eq!(
8082 typed,
8083 &[1.0, 4.0, 2.0, 5.0, 3.0, 6.0],
8084 "contiguized transpose must reflect logical (transposed) byte order",
8085 );
8086 }
8087
8088 #[test]
8092 fn pipelined_realize_const_const_add() {
8093 let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
8094 let rhs_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
8095
8096 let graph = Arc::new(RwLock::new(Graph::new()));
8097 let (lhs_id, rhs_id, add_id) = {
8098 let mut g = graph.write().unwrap();
8099 let lhs_id = g.push(Node {
8100 op: Op::Const,
8101 inputs: vec![],
8102 shape: Shape::from_dims(&[3]),
8103 dtype: DType::F32,
8104 });
8105 let rhs_id = g.push(Node {
8106 op: Op::Const,
8107 inputs: vec![],
8108 shape: Shape::from_dims(&[3]),
8109 dtype: DType::F32,
8110 });
8111 let add_id = g.push(Node {
8112 op: Op::Add,
8113 inputs: vec![lhs_id, rhs_id],
8114 shape: Shape::from_dims(&[3]),
8115 dtype: DType::F32,
8116 });
8117 g.set_target_backend(add_id, BackendId::Cpu);
8118 (lhs_id, rhs_id, add_id)
8119 };
8120
8121 let mut inputs = StorageCache::new();
8122 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
8123 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
8124
8125 let (result_arc, _result_layout) =
8126 PipelinedExecutor::realize(graph, add_id, inputs).expect("realize");
8127
8128 let guard = result_arc.read().unwrap();
8129 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
8130 let typed: &[f32] = c.as_slice().expect("f32 cast");
8131 assert_eq!(typed, &[11.0, 22.0, 33.0]);
8132 } else {
8133 panic!("expected CPU output");
8134 }
8135 }
8136
8137 #[test]
8150 #[ignore = "requires a live CUDA device"]
8151 #[cfg(feature = "cuda")]
8152 fn capture_decode_binary_chain_replays_bit_exact_cuda() {
8153 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8154 use fuel_ir::Shape;
8155 use fuel_memory::BackendStorage;
8156
8157 let Ok(dev) = CudaDevice::new(0) else {
8158 return fuel_test_support::hardware::skip(
8159 fuel_test_support::hardware::Hardware::Cuda,
8160 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
8161 );
8162 };
8163
8164 let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId, NodeId) {
8167 let graph = Arc::new(RwLock::new(Graph::new()));
8168 let (a, b, c, z) = {
8169 let mut g = graph.write().unwrap();
8170 let a = g.push(Node {
8171 op: Op::Const,
8172 inputs: vec![],
8173 shape: Shape::from_dims(&[4]),
8174 dtype: DType::F32,
8175 });
8176 let b = g.push(Node {
8177 op: Op::Const,
8178 inputs: vec![],
8179 shape: Shape::from_dims(&[4]),
8180 dtype: DType::F32,
8181 });
8182 let c = g.push(Node {
8183 op: Op::Const,
8184 inputs: vec![],
8185 shape: Shape::from_dims(&[4]),
8186 dtype: DType::F32,
8187 });
8188 let y = g.push(Node {
8189 op: Op::Add,
8190 inputs: vec![a, b],
8191 shape: Shape::from_dims(&[4]),
8192 dtype: DType::F32,
8193 });
8194 let z = g.push(Node {
8195 op: Op::Mul,
8196 inputs: vec![y, c],
8197 shape: Shape::from_dims(&[4]),
8198 dtype: DType::F32,
8199 });
8200 g.set_target_backend(y, BackendId::Cuda);
8201 g.set_target_backend(z, BackendId::Cuda);
8202 (a, b, c, z)
8203 };
8204 (graph, a, b, c, z)
8205 };
8206
8207 let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
8208 let bytes: &[u8] = bytemuck::cast_slice(v);
8209 let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytes).expect("h2d");
8210 Arc::new(RwLock::new(Storage::new(
8211 BackendStorage::Cuda(cb),
8212 DType::F32,
8213 )))
8214 };
8215 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8216 let g = arc.read().unwrap();
8217 let BackendStorage::Cuda(c) = &g.inner else {
8218 panic!("not cuda")
8219 };
8220 let host = c.to_cpu_bytes().expect("d2h");
8221 bytemuck::cast_slice(&host).to_vec()
8222 };
8223
8224 let a_data = [1.0_f32, 2.0, 3.0, 4.0];
8225 let b_data = [10.0_f32, 20.0, 30.0, 40.0];
8226 let c_data = [2.0_f32, 2.0, 2.0, 2.0];
8227 let expect0 = [22.0_f32, 44.0, 66.0, 88.0]; let a_arc = up(&a_data);
8231 let b_arc = up(&b_data);
8232 let c_arc = up(&c_data);
8233
8234 let (graph, a, b, c, z) = build();
8235 let mut inputs = StorageCache::new();
8236 inputs.insert(a, a_arc.clone());
8237 inputs.insert(b, b_arc.clone());
8238 inputs.insert(c, c_arc.clone());
8239
8240 let captured =
8241 PipelinedExecutor::capture_decode(Arc::clone(&graph), z, inputs, SymEnv::default())
8242 .expect("capture_decode");
8243
8244 assert_eq!(
8246 read(&captured.output),
8247 &expect0,
8248 "warm-run result after capture"
8249 );
8250
8251 captured.run.replay(&dev).expect("replay");
8253 dev.synchronize().expect("sync");
8254 assert_eq!(read(&captured.output), &expect0, "replay reproduces result");
8255
8256 let a2 = [100.0_f32, 200.0, 300.0, 400.0];
8260 {
8261 let g = a_arc.read().unwrap();
8262 let BackendStorage::Cuda(c) = &g.inner else {
8263 panic!("a not cuda")
8264 };
8265 c.write_from_host(bytemuck::cast_slice(&a2))
8266 .expect("in-place H2D");
8267 }
8268 captured.run.replay(&dev).expect("replay 2");
8269 dev.synchronize().expect("sync 2");
8270 let got = read(&captured.output);
8271
8272 let (graph_ref, ar, br, cr, zr) = build();
8273 let mut inputs_ref = StorageCache::new();
8274 inputs_ref.insert(ar, up(&a2));
8275 inputs_ref.insert(br, up(&b_data));
8276 inputs_ref.insert(cr, up(&c_data));
8277 let (ref_arc, _) =
8278 PipelinedExecutor::realize(graph_ref, zr, inputs_ref).expect("ref realize");
8279 let expect1 = read(&ref_arc); assert_eq!(
8281 got, expect1,
8282 "replay after in-place input update matches uncaptured realize"
8283 );
8284 assert_eq!(expect1, vec![220.0, 440.0, 660.0, 880.0]);
8285 }
8286
8287 #[test]
8300 #[ignore = "requires a live CUDA device"]
8301 #[cfg(feature = "cuda")]
8302 fn capture_decode_matmul_replays_bit_exact_cuda() {
8303 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8304 use fuel_ir::Shape;
8305 use fuel_memory::BackendStorage;
8306 let Ok(dev) = CudaDevice::new(0) else { return };
8307 const H: usize = 4;
8308 const F: usize = 8;
8309 let a_data: Vec<f32> = (0..H).map(|i| 0.5 + i as f32).collect();
8310 let w_data: Vec<f32> = (0..H * F).map(|i| 0.01 * (i as f32 + 1.0)).collect();
8311 let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId) {
8312 let graph = Arc::new(RwLock::new(Graph::new()));
8313 let (a, w, y) = {
8314 let mut g = graph.write().unwrap();
8315 let a = g.push(Node {
8316 op: Op::Const,
8317 inputs: vec![],
8318 shape: Shape::from_dims(&[1, H]),
8319 dtype: DType::F32,
8320 });
8321 let w = g.push(Node {
8322 op: Op::Const,
8323 inputs: vec![],
8324 shape: Shape::from_dims(&[H, F]),
8325 dtype: DType::F32,
8326 });
8327 let y = g.push(Node {
8328 op: Op::MatMul,
8329 inputs: vec![a, w],
8330 shape: Shape::from_dims(&[1, F]),
8331 dtype: DType::F32,
8332 });
8333 g.set_target_backend(y, BackendId::Cuda);
8334 (a, w, y)
8335 };
8336 (graph, a, w, y)
8337 };
8338 let up = |v: &[f32]| {
8339 let b: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
8340 Arc::new(RwLock::new(Storage::new(
8341 BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
8342 DType::F32,
8343 )))
8344 };
8345 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8346 let g = arc.read().unwrap();
8347 let BackendStorage::Cuda(c) = &g.inner else {
8348 panic!("not cuda")
8349 };
8350 c.to_cpu_bytes()
8351 .unwrap()
8352 .chunks_exact(4)
8353 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
8354 .collect()
8355 };
8356 let (gr, ar, wr, yr) = build();
8358 let mut insr = StorageCache::new();
8359 insr.insert(ar, up(&a_data));
8360 insr.insert(wr, up(&w_data));
8361 let (refarc, _) = PipelinedExecutor::realize(gr, yr, insr).unwrap();
8362 let expect = read(&refarc);
8363 let (graph, a, w, y) = build();
8365 let mut inputs = StorageCache::new();
8366 inputs.insert(a, up(&a_data));
8367 inputs.insert(w, up(&w_data));
8368 let captured =
8369 PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
8370 .expect("capture");
8371 assert_eq!(read(&captured.output), expect, "warm result after capture");
8372 let mut squatters = Vec::new();
8375 for _ in 0..8 {
8376 let g: Vec<u8> = (0..H * F).flat_map(|_| 999.0f32.to_le_bytes()).collect();
8377 squatters.push(CudaStorageBytes::from_cpu_bytes(&dev, &g).unwrap());
8378 }
8379 dev.synchronize().unwrap();
8380 let mut bad = 0usize;
8381 for _ in 0..50 {
8382 captured.run.replay(&dev).unwrap();
8383 dev.synchronize().unwrap();
8384 if read(&captured.output) != expect {
8385 bad += 1;
8386 }
8387 }
8388 assert_eq!(
8389 bad, 0,
8390 "cuBLAS matmul capture replays bit-exact ({bad}/50 wrong)"
8391 );
8392 }
8393
8394 #[test]
8397 #[ignore = "requires a live CUDA device"]
8398 #[cfg(feature = "cuda")]
8399 fn capture_decode_affine_replays_bit_exact_cuda() {
8400 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8401 use fuel_ir::Shape;
8402 use fuel_memory::BackendStorage;
8403 let Ok(dev) = CudaDevice::new(0) else { return };
8404 const N: usize = 6;
8405 let x_data: Vec<f32> = (0..N).map(|i| 1.0 + i as f32).collect();
8406 let expect: Vec<f32> = x_data.iter().map(|v| v * 0.5).collect();
8407 let graph = Arc::new(RwLock::new(Graph::new()));
8408 let (x, y) = {
8409 let mut g = graph.write().unwrap();
8410 let x = g.push(Node {
8411 op: Op::Const,
8412 inputs: vec![],
8413 shape: Shape::from_dims(&[1, N]),
8414 dtype: DType::F32,
8415 });
8416 let y = g.push(Node {
8417 op: Op::MulScalar(0.5),
8418 inputs: vec![x],
8419 shape: Shape::from_dims(&[1, N]),
8420 dtype: DType::F32,
8421 });
8422 g.set_target_backend(y, BackendId::Cuda);
8423 (x, y)
8424 };
8425 let up = |v: &[f32]| {
8426 let b: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
8427 Arc::new(RwLock::new(Storage::new(
8428 BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
8429 DType::F32,
8430 )))
8431 };
8432 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8433 let g = arc.read().unwrap();
8434 let BackendStorage::Cuda(c) = &g.inner else {
8435 panic!("not cuda")
8436 };
8437 c.to_cpu_bytes()
8438 .unwrap()
8439 .chunks_exact(4)
8440 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
8441 .collect()
8442 };
8443 let mut inputs = StorageCache::new();
8444 inputs.insert(x, up(&x_data));
8445 let captured =
8446 PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
8447 .expect("capture");
8448 assert_eq!(read(&captured.output), expect, "warm result after capture");
8449 let mut sq = Vec::new();
8450 for _ in 0..8 {
8451 let g: Vec<u8> = (0..N).flat_map(|_| 999.0f32.to_le_bytes()).collect();
8452 sq.push(CudaStorageBytes::from_cpu_bytes(&dev, &g).unwrap());
8453 }
8454 dev.synchronize().unwrap();
8455 let mut bad = 0usize;
8456 for _ in 0..50 {
8457 captured.run.replay(&dev).unwrap();
8458 dev.synchronize().unwrap();
8459 if read(&captured.output) != expect {
8460 bad += 1;
8461 }
8462 }
8463 assert_eq!(bad, 0, "affine capture replays bit-exact ({bad}/50 wrong)");
8464 }
8465
8466 #[test]
8473 #[ignore = "requires a live CUDA device"]
8474 #[cfg(feature = "cuda")]
8475 fn capture_decode_index_select_replays_bit_exact_cuda() {
8476 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8477 use fuel_ir::Shape;
8478 use fuel_memory::BackendStorage;
8479 let Ok(dev) = CudaDevice::new(0) else { return };
8480 const V: usize = 3;
8481 const H: usize = 4;
8482 let wte_data: Vec<f32> = (0..V)
8483 .flat_map(|v| (0..H).map(move |c| 1.0 + 10.0 * v as f32 + c as f32))
8484 .collect();
8485 let graph = Arc::new(RwLock::new(Graph::new()));
8486 let (tok, wte, emb) = {
8487 let mut g = graph.write().unwrap();
8488 let tok = g.push(Node {
8489 op: Op::Const,
8490 inputs: vec![],
8491 shape: Shape::from_dims(&[1]),
8492 dtype: DType::U32,
8493 });
8494 let wte = g.push(Node {
8495 op: Op::Const,
8496 inputs: vec![],
8497 shape: Shape::from_dims(&[V, H]),
8498 dtype: DType::F32,
8499 });
8500 let emb = g.push(Node {
8501 op: Op::IndexSelect { dim: 0 },
8502 inputs: vec![wte, tok],
8503 shape: Shape::from_dims(&[1, H]),
8504 dtype: DType::F32,
8505 });
8506 g.set_target_backend(emb, BackendId::Cuda);
8507 (tok, wte, emb)
8508 };
8509 let up_f32 = |v: &[f32]| {
8510 let b: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
8511 Arc::new(RwLock::new(Storage::new(
8512 BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
8513 DType::F32,
8514 )))
8515 };
8516 let up_u32 = |v: &[u32]| {
8517 let b: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
8518 Arc::new(RwLock::new(Storage::new(
8519 BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
8520 DType::U32,
8521 )))
8522 };
8523 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8524 let g = arc.read().unwrap();
8525 let BackendStorage::Cuda(c) = &g.inner else {
8526 panic!("not cuda")
8527 };
8528 c.to_cpu_bytes()
8529 .unwrap()
8530 .chunks_exact(4)
8531 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
8532 .collect()
8533 };
8534 let tok_arc = up_u32(&[0]);
8535 let mut inputs = StorageCache::new();
8536 inputs.insert(tok, tok_arc.clone());
8537 inputs.insert(wte, up_f32(&wte_data));
8538 let captured =
8539 PipelinedExecutor::capture_decode(Arc::clone(&graph), emb, inputs, SymEnv::default())
8540 .expect("capture");
8541 assert_eq!(
8542 read(&captured.output),
8543 vec![1.0, 2.0, 3.0, 4.0],
8544 "warm result after capture"
8545 );
8546 let mut squatters = Vec::new();
8549 for _ in 0..8 {
8550 let g: Vec<u8> = (0..V * H).flat_map(|_| 999.0f32.to_le_bytes()).collect();
8551 squatters.push(CudaStorageBytes::from_cpu_bytes(&dev, &g).unwrap());
8552 }
8553 dev.synchronize().unwrap();
8554 let mut bad = 0usize;
8555 for _ in 0..50 {
8556 captured.run.replay(&dev).unwrap();
8557 dev.synchronize().unwrap();
8558 if read(&captured.output) != vec![1.0, 2.0, 3.0, 4.0] {
8559 bad += 1;
8560 }
8561 }
8562 assert_eq!(
8563 bad, 0,
8564 "index_select capture replays bit-exact ({bad}/50 wrong)"
8565 );
8566 }
8567
8568 #[test]
8580 #[ignore = "requires a live CUDA device"]
8581 #[cfg(feature = "cuda")]
8582 fn capture_decode_norm_silu_softmax_replays_bit_exact_cuda() {
8583 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8584 use fuel_graph::registry::{FusedOpParams, FusedOps};
8585 use fuel_ir::Shape;
8586 use fuel_memory::BackendStorage;
8587
8588 let Ok(dev) = CudaDevice::new(0) else {
8589 return fuel_test_support::hardware::skip(
8590 fuel_test_support::hardware::Hardware::Cuda,
8591 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
8592 );
8593 };
8594 const H: usize = 6;
8595 let eps = 1e-5f64;
8596
8597 let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId) {
8598 let graph = Arc::new(RwLock::new(Graph::new()));
8599 let (x, y) = {
8600 let mut g = graph.write().unwrap();
8601 let mk = |g: &mut Graph, op: Op, ins: Vec<NodeId>, dims: &[usize]| {
8602 g.push(Node {
8603 op,
8604 inputs: ins,
8605 shape: Shape::from_dims(dims),
8606 dtype: DType::F32,
8607 })
8608 };
8609 let rms = || {
8610 Op::Fused(
8611 FusedOps::RMS_NORM_LAST_DIM,
8612 FusedOpParams::RmsNormLastDim { eps },
8613 )
8614 };
8615 let sm = || Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim);
8616 let x = mk(&mut g, Op::Const, vec![], &[1, H]);
8617 let n1 = mk(&mut g, rms(), vec![x], &[1, H]);
8618 let s = mk(&mut g, Op::Silu, vec![n1], &[1, H]);
8619 let p = mk(&mut g, sm(), vec![s], &[1, H]);
8620 let r = mk(&mut g, Op::Add, vec![p, x], &[1, H]);
8621 let n2 = mk(&mut g, rms(), vec![r], &[1, H]);
8622 let y = mk(&mut g, Op::Add, vec![n2, x], &[1, H]);
8623 for id in [n1, s, p, r, n2, y] {
8624 g.set_target_backend(id, BackendId::Cuda);
8625 }
8626 (x, y)
8627 };
8628 (graph, x, y)
8629 };
8630
8631 let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
8632 let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
8633 Arc::new(RwLock::new(Storage::new(
8634 BackendStorage::Cuda(cb),
8635 DType::F32,
8636 )))
8637 };
8638 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8639 let g = arc.read().unwrap();
8640 let BackendStorage::Cuda(c) = &g.inner else {
8641 panic!("not cuda")
8642 };
8643 bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
8644 };
8645
8646 let x_data: Vec<f32> = (0..H).map(|c| 0.1 * (c as f32 + 1.0)).collect();
8647 let x_arc = up(&x_data);
8648 let (graph, x, y) = build();
8649 let mut inputs = StorageCache::new();
8650 inputs.insert(x, x_arc.clone());
8651
8652 let captured =
8653 PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
8654 .expect("capture_decode norm/silu/softmax");
8655
8656 let realize_ref = |xv: &[f32]| -> Vec<f32> {
8657 let (g, xr, yr) = build();
8658 let mut ins = StorageCache::new();
8659 ins.insert(xr, up(xv));
8660 let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
8661 read(&arc)
8662 };
8663 let expect0 = realize_ref(&x_data);
8664
8665 assert_eq!(read(&captured.output), expect0, "warm result after capture");
8667 captured.run.replay(&dev).expect("replay");
8669 dev.synchronize().expect("sync");
8670 assert_eq!(read(&captured.output), expect0, "replay reproduces result");
8671
8672 let x2: Vec<f32> = (0..H).map(|c| 0.3 * (c as f32 + 1.0)).collect();
8675 {
8676 let g = x_arc.read().unwrap();
8677 let BackendStorage::Cuda(c) = &g.inner else {
8678 panic!("x not cuda")
8679 };
8680 c.write_from_host(bytemuck::cast_slice(&x2))
8681 .expect("in-place H2D");
8682 }
8683 captured.run.replay(&dev).expect("replay 2");
8684 dev.synchronize().expect("sync 2");
8685 assert_eq!(
8686 read(&captured.output),
8687 realize_ref(&x2),
8688 "replay after x update matches realize"
8689 );
8690 }
8691
8692 #[test]
8704 #[ignore = "requires a live CUDA device"]
8705 #[cfg(feature = "cuda")]
8706 fn capture_decode_contiguize_then_concat_replays_bit_exact_cuda() {
8707 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8708 use fuel_ir::Shape;
8709 use fuel_memory::BackendStorage;
8710
8711 let Ok(dev) = CudaDevice::new(0) else {
8712 return fuel_test_support::hardware::skip(
8713 fuel_test_support::hardware::Hardware::Cuda,
8714 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
8715 );
8716 };
8717
8718 let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId) {
8719 let graph = Arc::new(RwLock::new(Graph::new()));
8720 let (x, b, y) = {
8721 let mut g = graph.write().unwrap();
8722 let x = g.push(Node {
8723 op: Op::Const,
8724 inputs: vec![],
8725 shape: Shape::from_dims(&[2, 3, 4]),
8726 dtype: DType::F32,
8727 });
8728 let p = g.push(Node {
8729 op: Op::Permute(vec![1, 0, 2]),
8730 inputs: vec![x],
8731 shape: Shape::from_dims(&[3, 2, 4]),
8732 dtype: DType::F32,
8733 });
8734 let r = g.push(Node {
8735 op: Op::Reshape(Shape::from_dims(&[3, 8])),
8736 inputs: vec![p],
8737 shape: Shape::from_dims(&[3, 8]),
8738 dtype: DType::F32,
8739 });
8740 let b = g.push(Node {
8741 op: Op::Const,
8742 inputs: vec![],
8743 shape: Shape::from_dims(&[3, 8]),
8744 dtype: DType::F32,
8745 });
8746 let y = g.push(Node {
8747 op: Op::Concat { dim: 1 },
8748 inputs: vec![r, b],
8749 shape: Shape::from_dims(&[3, 16]),
8750 dtype: DType::F32,
8751 });
8752 g.set_target_backend(p, BackendId::Cuda);
8753 g.set_target_backend(r, BackendId::Cuda);
8754 g.set_target_backend(y, BackendId::Cuda);
8755 (x, b, y)
8756 };
8757 (graph, x, b, y)
8758 };
8759
8760 let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
8761 let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
8762 Arc::new(RwLock::new(Storage::new(
8763 BackendStorage::Cuda(cb),
8764 DType::F32,
8765 )))
8766 };
8767 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8768 let g = arc.read().unwrap();
8769 let BackendStorage::Cuda(c) = &g.inner else {
8770 panic!("not cuda")
8771 };
8772 bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
8773 };
8774
8775 let x_data: Vec<f32> = (0..24).map(|i| i as f32).collect();
8776 let b_data: Vec<f32> = (0..24).map(|i| 100.0 + i as f32).collect();
8777 let x_arc = up(&x_data);
8778 let (graph, x, b, y) = build();
8779 let mut inputs = StorageCache::new();
8780 inputs.insert(x, x_arc.clone());
8781 inputs.insert(b, up(&b_data));
8785
8786 let captured =
8787 PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
8788 .expect("capture_decode contiguize+concat");
8789
8790 let realize_ref = |xv: &[f32], bv: &[f32]| -> Vec<f32> {
8791 let (g, xr, br, yr) = build();
8792 let mut ins = StorageCache::new();
8793 ins.insert(xr, up(xv));
8794 ins.insert(br, up(bv));
8795 let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
8796 read(&arc)
8797 };
8798 let expect0 = realize_ref(&x_data, &b_data);
8799
8800 assert_eq!(read(&captured.output), expect0, "warm result after capture");
8803
8804 for i in 0..2 {
8808 captured
8809 .run
8810 .replay(&dev)
8811 .unwrap_or_else(|e| panic!("replay {i}: {e:?}"));
8812 dev.synchronize().expect("sync");
8813 assert_eq!(
8814 read(&captured.output),
8815 expect0,
8816 "replay {i} reproduces result"
8817 );
8818 }
8819
8820 let x2: Vec<f32> = (0..24).map(|i| 1000.0 + i as f32).collect();
8825 {
8826 let g = x_arc.read().unwrap();
8827 let BackendStorage::Cuda(c) = &g.inner else {
8828 panic!("x not cuda")
8829 };
8830 c.write_from_host(bytemuck::cast_slice(&x2))
8831 .expect("in-place H2D");
8832 }
8833 captured.run.replay(&dev).expect("replay 3");
8834 dev.synchronize().expect("sync 3");
8835 assert_eq!(
8836 read(&captured.output),
8837 realize_ref(&x2, &b_data),
8838 "replay after x update matches independent uncaptured realize"
8839 );
8840 }
8841
8842 #[test]
8853 #[ignore = "requires a live CUDA device"]
8854 #[cfg(feature = "cuda")]
8855 fn capture_decode_same_device_copy_replays_bit_exact_cuda() {
8856 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8857 use fuel_ir::Shape;
8858 use fuel_memory::BackendStorage;
8859
8860 let Ok(dev) = CudaDevice::new(0) else {
8861 return fuel_test_support::hardware::skip(
8862 fuel_test_support::hardware::Hardware::Cuda,
8863 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
8864 );
8865 };
8866
8867 let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId) {
8868 let graph = Arc::new(RwLock::new(Graph::new()));
8869 let (x, b, y) = {
8870 let mut g = graph.write().unwrap();
8871 let x = g.push(Node {
8872 op: Op::Const,
8873 inputs: vec![],
8874 shape: Shape::from_dims(&[2, 4]),
8875 dtype: DType::F32,
8876 });
8877 let b = g.push(Node {
8878 op: Op::Const,
8879 inputs: vec![],
8880 shape: Shape::from_dims(&[2, 4]),
8881 dtype: DType::F32,
8882 });
8883 let c = g.push(Node {
8885 op: Op::Copy {
8886 target: DeviceLocation::Cuda { gpu_id: 0 },
8887 },
8888 inputs: vec![x],
8889 shape: Shape::from_dims(&[2, 4]),
8890 dtype: DType::F32,
8891 });
8892 let y = g.push(Node {
8893 op: Op::Add,
8894 inputs: vec![c, b],
8895 shape: Shape::from_dims(&[2, 4]),
8896 dtype: DType::F32,
8897 });
8898 g.set_target_backend(c, BackendId::Cuda);
8900 g.set_target_backend(y, BackendId::Cuda);
8901 (x, b, y)
8902 };
8903 (graph, x, b, y)
8904 };
8905
8906 let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
8907 let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
8908 Arc::new(RwLock::new(Storage::new(
8909 BackendStorage::Cuda(cb),
8910 DType::F32,
8911 )))
8912 };
8913 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
8914 let g = arc.read().unwrap();
8915 let BackendStorage::Cuda(c) = &g.inner else {
8916 panic!("not cuda")
8917 };
8918 bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
8919 };
8920
8921 let x_data: Vec<f32> = (0..8).map(|i| i as f32).collect();
8922 let b_data: Vec<f32> = (0..8).map(|i| 100.0 + i as f32).collect();
8923 let x_arc = up(&x_data);
8924 let (graph, x, b, y) = build();
8925 let mut inputs = StorageCache::new();
8926 inputs.insert(x, x_arc.clone());
8927 inputs.insert(b, up(&b_data));
8930
8931 let captured =
8932 PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
8933 .expect("capture_decode same-device copy");
8934
8935 let realize_ref = |xv: &[f32], bv: &[f32]| -> Vec<f32> {
8936 let (g, xr, br, yr) = build();
8937 let mut ins = StorageCache::new();
8938 ins.insert(xr, up(xv));
8939 ins.insert(br, up(bv));
8940 let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
8941 read(&arc)
8942 };
8943 let expect0 = realize_ref(&x_data, &b_data);
8944
8945 assert_eq!(read(&captured.output), expect0, "warm result after capture");
8947
8948 for i in 0..2 {
8951 captured
8952 .run
8953 .replay(&dev)
8954 .unwrap_or_else(|e| panic!("replay {i}: {e:?}"));
8955 dev.synchronize().expect("sync");
8956 assert_eq!(
8957 read(&captured.output),
8958 expect0,
8959 "replay {i} reproduces result"
8960 );
8961 }
8962
8963 let x2: Vec<f32> = (0..8).map(|i| 1000.0 + i as f32).collect();
8967 {
8968 let g = x_arc.read().unwrap();
8969 let BackendStorage::Cuda(c) = &g.inner else {
8970 panic!("x not cuda")
8971 };
8972 c.write_from_host(bytemuck::cast_slice(&x2))
8973 .expect("in-place H2D");
8974 }
8975 captured.run.replay(&dev).expect("replay 3");
8976 dev.synchronize().expect("sync 3");
8977 assert_eq!(
8978 read(&captured.output),
8979 realize_ref(&x2, &b_data),
8980 "replay after x update matches independent uncaptured realize"
8981 );
8982 }
8983
8984 #[test]
8990 #[ignore = "requires a live CUDA device"]
8991 #[cfg(feature = "cuda")]
8992 fn capture_decode_rejects_cross_device_copy_cuda() {
8993 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
8994 use fuel_ir::Shape;
8995 use fuel_memory::BackendStorage;
8996
8997 let Ok(dev) = CudaDevice::new(0) else {
8998 return fuel_test_support::hardware::skip(
8999 fuel_test_support::hardware::Hardware::Cuda,
9000 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
9001 );
9002 };
9003
9004 let graph = Arc::new(RwLock::new(Graph::new()));
9005 let (x, y) = {
9006 let mut g = graph.write().unwrap();
9007 let x = g.push(Node {
9008 op: Op::Const,
9009 inputs: vec![],
9010 shape: Shape::from_dims(&[2, 4]),
9011 dtype: DType::F32,
9012 });
9013 let y = g.push(Node {
9015 op: Op::Copy {
9016 target: DeviceLocation::Cpu,
9017 },
9018 inputs: vec![x],
9019 shape: Shape::from_dims(&[2, 4]),
9020 dtype: DType::F32,
9021 });
9022 g.set_target_backend(y, BackendId::Cuda); (x, y)
9024 };
9025
9026 let x_data: Vec<f32> = (0..8).map(|i| i as f32).collect();
9027 let cb =
9028 CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(&x_data)).expect("h2d");
9029 let x_arc = Arc::new(RwLock::new(Storage::new(
9030 BackendStorage::Cuda(cb),
9031 DType::F32,
9032 )));
9033 let mut inputs = StorageCache::new();
9034 inputs.insert(x, x_arc);
9035
9036 let result = PipelinedExecutor::capture_decode(graph, y, inputs, SymEnv::default());
9037 let msg = match result {
9038 Ok(_) => panic!("D2H Op::Copy must be rejected by capture_decode"),
9039 Err(e) => e.to_string(),
9040 };
9041 assert!(
9042 msg.contains("cross-device"),
9043 "error should describe the cross-device rejection: {msg}"
9044 );
9045 }
9046
9047 #[test]
9055 #[ignore = "requires a live CUDA device"]
9056 #[cfg(feature = "cuda")]
9057 fn captured_decode_session_replays_many_tokens_cuda() {
9058 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
9059 use fuel_graph::registry::{FusedOpParams, FusedOps};
9060 use fuel_ir::Shape;
9061 use fuel_memory::BackendStorage;
9062
9063 let Ok(dev) = CudaDevice::new(0) else {
9064 return fuel_test_support::hardware::skip(
9065 fuel_test_support::hardware::Hardware::Cuda,
9066 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
9067 );
9068 };
9069 const H: usize = 6;
9070 let eps = 1e-5f64;
9071
9072 let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId) {
9073 let graph = Arc::new(RwLock::new(Graph::new()));
9074 let (x, y) = {
9075 let mut g = graph.write().unwrap();
9076 let mk = |g: &mut Graph, op: Op, ins: Vec<NodeId>, dims: &[usize]| {
9077 g.push(Node {
9078 op,
9079 inputs: ins,
9080 shape: Shape::from_dims(dims),
9081 dtype: DType::F32,
9082 })
9083 };
9084 let rms = || {
9085 Op::Fused(
9086 FusedOps::RMS_NORM_LAST_DIM,
9087 FusedOpParams::RmsNormLastDim { eps },
9088 )
9089 };
9090 let sm = || Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim);
9091 let x = mk(&mut g, Op::Const, vec![], &[1, H]);
9092 let n1 = mk(&mut g, rms(), vec![x], &[1, H]);
9093 let s = mk(&mut g, Op::Silu, vec![n1], &[1, H]);
9094 let p = mk(&mut g, sm(), vec![s], &[1, H]);
9095 let r = mk(&mut g, Op::Add, vec![p, x], &[1, H]);
9096 let n2 = mk(&mut g, rms(), vec![r], &[1, H]);
9097 let y = mk(&mut g, Op::Add, vec![n2, x], &[1, H]);
9098 for id in [n1, s, p, r, n2, y] {
9099 g.set_target_backend(id, BackendId::Cuda);
9100 }
9101 (x, y)
9102 };
9103 (graph, x, y)
9104 };
9105 let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
9106 let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
9107 Arc::new(RwLock::new(Storage::new(
9108 BackendStorage::Cuda(cb),
9109 DType::F32,
9110 )))
9111 };
9112 let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
9113 let g = arc.read().unwrap();
9114 let BackendStorage::Cuda(c) = &g.inner else {
9115 panic!("not cuda")
9116 };
9117 bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
9118 };
9119 let realize_ref = |xv: &[f32]| -> Vec<f32> {
9120 let (g, xr, yr) = build();
9121 let mut ins = StorageCache::new();
9122 ins.insert(xr, up(xv));
9123 let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
9124 read(&arc)
9125 };
9126
9127 let tokens: Vec<Vec<f32>> = (0..5)
9129 .map(|t| {
9130 (0..H)
9131 .map(|c| 0.1 * (t as f32 + 1.0) * (c as f32 + 1.0))
9132 .collect()
9133 })
9134 .collect();
9135
9136 let x_arc = up(&tokens[0]);
9138 let (graph, x, y) = build();
9139 let mut inputs = StorageCache::new();
9140 inputs.insert(x, x_arc.clone());
9141 let session =
9142 CapturedDecodeSession::capture(Arc::clone(&graph), y, inputs, &[x], SymEnv::default())
9143 .expect("capture session");
9144
9145 assert_eq!(
9147 read(session.output()),
9148 realize_ref(&tokens[0]),
9149 "warm token 0"
9150 );
9151
9152 for (t, tok) in tokens.iter().enumerate() {
9155 let out = session
9156 .replay_token(&[(x, bytemuck::cast_slice(tok))])
9157 .expect("replay_token");
9158 assert_eq!(
9159 read(&out),
9160 realize_ref(tok),
9161 "replayed token {t} matches uncaptured realize"
9162 );
9163 }
9164 assert_ne!(
9166 realize_ref(&tokens[0]),
9167 realize_ref(&tokens[1]),
9168 "tokens must differ"
9169 );
9170 }
9171
9172 #[test]
9178 fn realize_with_optimized_route_empty_equals_arm0() {
9179 use crate::optimize::OptimizedGraph;
9180
9181 let build = || {
9182 let graph = Arc::new(RwLock::new(Graph::new()));
9183 let (lhs_id, rhs_id, add_id) = {
9184 let mut g = graph.write().unwrap();
9185 let lhs_id = g.push(Node {
9186 op: Op::Const,
9187 inputs: vec![],
9188 shape: Shape::from_dims(&[3]),
9189 dtype: DType::F32,
9190 });
9191 let rhs_id = g.push(Node {
9192 op: Op::Const,
9193 inputs: vec![],
9194 shape: Shape::from_dims(&[3]),
9195 dtype: DType::F32,
9196 });
9197 let add_id = g.push(Node {
9198 op: Op::Add,
9199 inputs: vec![lhs_id, rhs_id],
9200 shape: Shape::from_dims(&[3]),
9201 dtype: DType::F32,
9202 });
9203 g.set_target_backend(add_id, BackendId::Cpu);
9204 (lhs_id, rhs_id, add_id)
9205 };
9206 let mut inputs = StorageCache::new();
9207 inputs.insert(
9208 lhs_id,
9209 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
9210 1.0_f32, 2.0, 3.0,
9211 ]))),
9212 );
9213 inputs.insert(
9214 rhs_id,
9215 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
9216 10.0_f32, 20.0, 30.0,
9217 ]))),
9218 );
9219 (graph, add_id, inputs)
9220 };
9221
9222 let read_bytes = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
9223 let guard = arc.read().unwrap();
9224 match &guard.inner {
9225 fuel_memory::BackendStorage::Cpu(c) => {
9226 c.as_slice::<f32>().expect("f32 cast").to_vec()
9227 }
9228 _ => panic!("expected CPU output"),
9229 }
9230 };
9231
9232 drop(global_bindings());
9244 let live_gen = crate::dispatch::topology_generation();
9245
9246 let (g_a, add_a, in_a) = build();
9248 let opt_a = OptimizedGraph {
9249 roots: vec![add_a],
9250 generation: live_gen,
9251 placements: None,
9252 };
9253 let (arm0_arc, _) = PipelinedExecutor::realize_with_optimized(g_a, add_a, in_a, &opt_a)
9254 .expect("arm-0 realize");
9255 let arm0 = read_bytes(&arm0_arc);
9256
9257 let (g_r, add_r, in_r) = build();
9259 let opt_r = OptimizedGraph {
9260 roots: vec![add_r],
9261 generation: live_gen,
9262 placements: None,
9263 };
9264 let empty_route = PickedRoute::new();
9265 let (route_arc, _) =
9266 PipelinedExecutor::realize_with_optimized_route(g_r, add_r, in_r, &opt_r, &empty_route)
9267 .expect("empty-route realize");
9268 let routed = read_bytes(&route_arc);
9269
9270 assert_eq!(arm0, vec![11.0, 22.0, 33.0]);
9271 assert_eq!(
9272 routed, arm0,
9273 "an empty route realizes byte-identically to arm-0 (Phase B contract)",
9274 );
9275 }
9276
9277 #[test]
9284 fn realize_with_optimized_picking_cached_env_matches_arm0_cold_and_warm() {
9285 use crate::optimize::OptimizedGraph;
9286 use std::sync::OnceLock;
9287
9288 let build = || {
9289 let graph = Arc::new(RwLock::new(Graph::new()));
9290 let (lhs_id, rhs_id, add_id) = {
9291 let mut g = graph.write().unwrap();
9292 let lhs_id = g.push(Node {
9293 op: Op::Const,
9294 inputs: vec![],
9295 shape: Shape::from_dims(&[3]),
9296 dtype: DType::F32,
9297 });
9298 let rhs_id = g.push(Node {
9299 op: Op::Const,
9300 inputs: vec![],
9301 shape: Shape::from_dims(&[3]),
9302 dtype: DType::F32,
9303 });
9304 let add_id = g.push(Node {
9305 op: Op::Add,
9306 inputs: vec![lhs_id, rhs_id],
9307 shape: Shape::from_dims(&[3]),
9308 dtype: DType::F32,
9309 });
9310 g.set_target_backend(add_id, BackendId::Cpu);
9311 (lhs_id, rhs_id, add_id)
9312 };
9313 let mut inputs = StorageCache::new();
9314 inputs.insert(
9315 lhs_id,
9316 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
9317 1.0_f32, 2.0, 3.0,
9318 ]))),
9319 );
9320 inputs.insert(
9321 rhs_id,
9322 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
9323 10.0_f32, 20.0, 30.0,
9324 ]))),
9325 );
9326 (graph, add_id, inputs)
9327 };
9328 let read_bytes = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
9329 let guard = arc.read().unwrap();
9330 match &guard.inner {
9331 fuel_memory::BackendStorage::Cpu(c) => {
9332 c.as_slice::<f32>().expect("f32 cast").to_vec()
9333 }
9334 _ => panic!("expected CPU output"),
9335 }
9336 };
9337
9338 drop(global_bindings());
9339 let live_gen = crate::dispatch::topology_generation();
9340
9341 let (g, add, inputs) = build();
9342 let optimized = OptimizedGraph {
9343 roots: vec![add],
9344 generation: live_gen,
9345 placements: None,
9346 };
9347 let order_cache: OnceLock<Vec<NodeId>> = OnceLock::new();
9348 assert!(order_cache.get().is_none(), "sanity: cache starts empty");
9349
9350 let (cold_arc, _) = PipelinedExecutor::realize_with_optimized_picking_cached_env(
9352 g.clone(),
9353 add,
9354 inputs,
9355 &optimized,
9356 &order_cache,
9357 None,
9358 None,
9359 SymEnv::default(),
9360 )
9361 .expect("cold cached realize");
9362 assert_eq!(read_bytes(&cold_arc), vec![11.0, 22.0, 33.0]);
9363 assert!(
9364 order_cache.get().is_some(),
9365 "the cold call must populate order_cache, not leave it empty"
9366 );
9367 let cached_order_after_cold = order_cache.get().unwrap().clone();
9368
9369 let mut warm_inputs = StorageCache::new();
9373 warm_inputs.insert(
9374 g.read().unwrap().node(add).inputs[0],
9375 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
9376 100.0_f32, 200.0, 300.0,
9377 ]))),
9378 );
9379 warm_inputs.insert(
9380 g.read().unwrap().node(add).inputs[1],
9381 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
9382 1.0_f32, 2.0, 3.0,
9383 ]))),
9384 );
9385 let (warm_arc, _) = PipelinedExecutor::realize_with_optimized_picking_cached_env(
9386 g,
9387 add,
9388 warm_inputs,
9389 &optimized,
9390 &order_cache,
9391 None,
9392 None,
9393 SymEnv::default(),
9394 )
9395 .expect("warm cached realize");
9396 assert_eq!(
9397 read_bytes(&warm_arc),
9398 vec![101.0, 202.0, 303.0],
9399 "warm call must realize its OWN rebound inputs, not replay stale output",
9400 );
9401 assert_eq!(
9402 order_cache.get().unwrap(),
9403 &cached_order_after_cold,
9404 "the warm call must reuse the SAME cached order, not recompute a different one",
9405 );
9406 }
9407
9408 #[test]
9424 fn cached_generation_for_matches_optimized_generation_for() {
9425 let optimized = OptimizedGraph {
9426 roots: vec![NodeId(0)],
9427 generation: 42,
9428 placements: None,
9429 };
9430 let order_cache: std::sync::OnceLock<Vec<NodeId>> = std::sync::OnceLock::new();
9431
9432 let optimized_gen = generation_for(&OrderSource::Optimized {
9433 optimized: &optimized,
9434 route: None,
9435 });
9436 let cached_gen = generation_for(&OrderSource::Cached {
9437 optimized: &optimized,
9438 order_cache: &order_cache,
9439 });
9440 assert_eq!(
9441 optimized_gen,
9442 Some(42),
9443 "sanity: Optimized reports the OptimizedGraph's own generation",
9444 );
9445 assert_eq!(
9446 cached_gen, optimized_gen,
9447 "Cached must report the SAME generation as Optimized for the same \
9448 OptimizedGraph -- the TopologyChanged check cannot tell the two apart, \
9449 by construction, so caching cannot silently weaken it",
9450 );
9451 }
9452
9453 #[test]
9470 fn c3_reorder_is_a_valid_topo_permutation() {
9471 let n = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
9482 g.push(Node {
9483 op,
9484 inputs,
9485 shape: Shape::from_dims(&[3]),
9486 dtype: DType::F32,
9487 })
9488 };
9489 let mut g = Graph::new();
9490 let ca = n(&mut g, Op::Const, vec![]);
9491 let cb = n(&mut g, Op::Const, vec![]);
9492 let mut cprev = ca;
9494 for _ in 0..4 {
9495 cprev = n(&mut g, Op::Relu, vec![cprev]);
9496 g.set_target_backend(cprev, BackendId::Cuda);
9497 }
9498 let c2 = cprev; let mut vprev = cb;
9501 for _ in 0..2 {
9502 vprev = n(&mut g, Op::Relu, vec![vprev]);
9503 g.set_target_backend(vprev, BackendId::Vulkan);
9504 }
9505 let v2 = vprev; let copy = g.push(Node {
9509 op: Op::Copy {
9510 target: DeviceLocation::Cuda { gpu_id: 0 },
9511 },
9512 inputs: vec![v2],
9513 shape: Shape::from_dims(&[3]),
9514 dtype: DType::F32,
9515 });
9516 g.set_target_backend(copy, BackendId::Cuda);
9517 let out = n(&mut g, Op::Add, vec![c2, copy]);
9518 g.set_target_backend(out, BackendId::Cuda);
9519
9520 let unreordered = execution_plan(&g, &[out]);
9521 let reordered = fuel_graph::lower_runs_arm0(&g, &[out]);
9522
9523 assert_ne!(
9526 reordered, unreordered,
9527 "the C3 reorder must change the order on a two-device graph; \
9528 reordered={reordered:?} unreordered={unreordered:?}",
9529 );
9530
9531 let mut a = unreordered.clone();
9533 let mut b = reordered.clone();
9534 a.sort();
9535 b.sort();
9536 assert_eq!(a, b, "the reorder is a permutation of the same node set");
9537
9538 let pos: std::collections::HashMap<NodeId, usize> = reordered
9541 .iter()
9542 .enumerate()
9543 .map(|(i, &node)| (node, i))
9544 .collect();
9545 for (i, &node) in reordered.iter().enumerate() {
9546 for &inp in &g.node(node).inputs {
9547 if let Some(&j) = pos.get(&inp) {
9548 assert!(
9549 j < i,
9550 "input {inp:?} of {node:?} must precede it in the reordered \
9551 order; reordered={reordered:?}",
9552 );
9553 }
9554 }
9555 }
9556
9557 let p = |node: NodeId| pos[&node];
9564 assert!(p(c2) < p(copy), "CUDA chunk before the cross-device copy");
9565 assert!(p(v2) < p(copy), "Vulkan chunk before the cross-device copy");
9566 assert!(
9567 p(c2) < p(v2),
9568 "the heavier CUDA chunk is emitted before the lighter Vulkan chunk \
9569 (critical-path heavy-first); reordered={reordered:?}",
9570 );
9571 }
9572
9573 #[test]
9594 fn c3_reorder_is_operand_order_invariant() {
9595 let build = |cuda_first: bool| -> (Graph, NodeId) {
9599 let n = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
9600 g.push(Node {
9601 op,
9602 inputs,
9603 shape: Shape::from_dims(&[3]),
9604 dtype: DType::F32,
9605 })
9606 };
9607 let mut g = Graph::new();
9608 let ca = n(&mut g, Op::Const, vec![]);
9609 let cb = n(&mut g, Op::Const, vec![]);
9610 let mut cprev = ca;
9611 for _ in 0..4 {
9612 cprev = n(&mut g, Op::Relu, vec![cprev]);
9613 g.set_target_backend(cprev, BackendId::Cuda);
9614 }
9615 let c2 = cprev;
9616 let mut vprev = cb;
9617 for _ in 0..2 {
9618 vprev = n(&mut g, Op::Relu, vec![vprev]);
9619 g.set_target_backend(vprev, BackendId::Vulkan);
9620 }
9621 let v2 = vprev;
9622 let copy = g.push(Node {
9623 op: Op::Copy {
9624 target: DeviceLocation::Cuda { gpu_id: 0 },
9625 },
9626 inputs: vec![v2],
9627 shape: Shape::from_dims(&[3]),
9628 dtype: DType::F32,
9629 });
9630 g.set_target_backend(copy, BackendId::Cuda);
9631 let inputs = if cuda_first {
9633 vec![c2, copy]
9634 } else {
9635 vec![copy, c2]
9636 };
9637 let out = n(&mut g, Op::Add, inputs);
9638 g.set_target_backend(out, BackendId::Cuda);
9639 (g, out)
9640 };
9641
9642 let (g_cf, out_cf) = build(true);
9643 let (g_vf, out_vf) = build(false);
9644
9645 let order_cf = fuel_graph::lower_runs_arm0(&g_cf, &[out_cf]);
9648 let order_vf = fuel_graph::lower_runs_arm0(&g_vf, &[out_vf]);
9649
9650 let tag = |g: &Graph, ids: &[NodeId]| -> Vec<(String, Option<BackendId>)> {
9656 ids.iter()
9657 .map(|&id| (format!("{:?}", g.node(id).op), g.target_backend(id)))
9658 .collect()
9659 };
9660 let seq_cf = tag(&g_cf, &order_cf);
9661 let seq_vf = tag(&g_vf, &order_vf);
9662 assert_eq!(
9663 seq_cf, seq_vf,
9664 "the C3-reordered dispatch order MUST be identical for both \
9665 reconverge operand orders (operand-order-invariance — the \
9666 executor then sees the same WorkItem stream regardless of \
9667 operand order, so overlap is operand-insensitive).\n\
9668 cuda-first = {seq_cf:?}\nvulkan-first = {seq_vf:?}",
9669 );
9670 }
9671
9672 #[cfg(test)]
9675 #[derive(Debug)]
9676 struct PickLast;
9677 #[cfg(test)]
9678 impl RuntimeSelector for PickLast {
9679 fn select<'a>(
9680 &self,
9681 set: &'a crate::ranker::AlternativeSet,
9682 ) -> Option<&'a crate::ranker::Candidate> {
9683 set.alternatives().last()
9684 }
9685 }
9686
9687 #[cfg(test)]
9692 fn picking_diamond() -> (Arc<RwLock<Graph>>, NodeId, NodeId) {
9693 let node = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
9694 g.push(Node {
9695 op,
9696 inputs,
9697 shape: Shape::from_dims(&[2]),
9698 dtype: DType::F32,
9699 })
9700 };
9701 let mut g = Graph::new();
9702 let pre = node(&mut g, Op::Const, vec![]);
9703 let diverge = node(&mut g, Op::Relu, vec![pre]);
9704 let arm0 = node(&mut g, Op::Silu, vec![diverge]);
9705 let arm1 = node(&mut g, Op::GeluTanh, vec![diverge]);
9706 g.set_target_backend(arm0, BackendId::Cuda);
9707 g.set_target_backend(arm1, BackendId::Cpu);
9708 let reconverge = node(&mut g, Op::Relu, vec![arm0]);
9709 let mut b = g.open_branch(diverge);
9710 b.add_arm(arm0);
9711 b.add_arm(arm1);
9712 let branch = b
9713 .finalize_branches(&mut g, reconverge)
9714 .expect("well-formed 2-arm branch")
9715 .expect("2 arms survive");
9716 let post = node(&mut g, Op::Tanh, vec![reconverge]);
9717 (Arc::new(RwLock::new(g)), branch, post)
9718 }
9719
9720 #[test]
9728 fn streaming_pick_for_gates() {
9729 let (graph, _branch, post) = picking_diamond();
9730 let _ = post;
9731
9732 let sel: Arc<dyn RuntimeSelector> = Arc::new(PickLast);
9734 assert!(
9735 PipelinedExecutor::streaming_pick_for(&graph, Some(sel.clone()), None)
9736 .expect("ok")
9737 .is_some(),
9738 "selector + branch ⇒ a streaming pick",
9739 );
9740
9741 assert!(
9743 PipelinedExecutor::streaming_pick_for(&graph, None, None)
9744 .expect("ok")
9745 .is_none(),
9746 "no selector ⇒ no streaming pick (arm-0 lowering)",
9747 );
9748
9749 let node = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
9751 g.push(Node {
9752 op,
9753 inputs,
9754 shape: Shape::from_dims(&[2]),
9755 dtype: DType::F32,
9756 })
9757 };
9758 let bl_graph = {
9759 let mut g = Graph::new();
9760 let c1 = node(&mut g, Op::Const, vec![]);
9761 let c2 = node(&mut g, Op::Const, vec![]);
9762 let add = node(&mut g, Op::Add, vec![c1, c2]);
9763 g.set_target_backend(add, BackendId::Cpu);
9764 Arc::new(RwLock::new(g))
9765 };
9766 assert!(
9767 PipelinedExecutor::streaming_pick_for(&bl_graph, Some(sel), None)
9768 .expect("ok")
9769 .is_none(),
9770 "branchless graph ⇒ no branches ⇒ no streaming pick (fast-path)",
9771 );
9772 }
9773
9774 #[test]
9784 fn streamed_order_equals_one_shot() {
9785 let (graph, _branch, post) = picking_diamond();
9786 let g = graph.read().unwrap();
9787 let bindings = global_bindings();
9788
9789 for sel in [
9790 Arc::new(crate::ranker::WinnerSelector) as Arc<dyn RuntimeSelector>,
9791 Arc::new(PickLast) as Arc<dyn RuntimeSelector>,
9792 ] {
9793 let one_shot_route = pick_route(&g, &[post], &bindings, sel.as_ref(), None);
9795 let one_shot = fuel_graph::lower_picked_route(&g, &[post], &one_shot_route);
9796
9797 let streamed = fuel_graph::lower_picked_route_streaming(&g, &[post], |branch| {
9800 resolve_branch(&g, branch, &bindings, sel.as_ref())
9801 });
9802
9803 assert_eq!(
9804 streamed, one_shot,
9805 "streamed order must equal one-shot for selector {sel:?}",
9806 );
9807 }
9808 }
9809
9810 #[test]
9813 fn pipelined_errors_on_unset_target_backend() {
9814 let graph = Arc::new(RwLock::new(Graph::new()));
9815 let (lhs_id, rhs_id, add_id) = {
9816 let mut g = graph.write().unwrap();
9817 let lhs = g.push(Node {
9818 op: Op::Const,
9819 inputs: vec![],
9820 shape: Shape::from_dims(&[2]),
9821 dtype: DType::F32,
9822 });
9823 let rhs = g.push(Node {
9824 op: Op::Const,
9825 inputs: vec![],
9826 shape: Shape::from_dims(&[2]),
9827 dtype: DType::F32,
9828 });
9829 let add = g.push(Node {
9831 op: Op::Add,
9832 inputs: vec![lhs, rhs],
9833 shape: Shape::from_dims(&[2]),
9834 dtype: DType::F32,
9835 });
9836 (lhs, rhs, add)
9837 };
9838
9839 let mut inputs = StorageCache::new();
9840 inputs.insert(
9841 lhs_id,
9842 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]))),
9843 );
9844 inputs.insert(
9845 rhs_id,
9846 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[3.0_f32, 4.0]))),
9847 );
9848
9849 let result = PipelinedExecutor::realize(graph, add_id, inputs);
9850 assert!(result.is_err(), "missing target_backend must error");
9851 let msg = result.unwrap_err().to_string();
9852 assert!(
9853 msg.contains("target_backend"),
9854 "error names the unmet precondition: {msg}"
9855 );
9856 }
9857
9858 #[test]
9861 fn pipelined_realize_const_only() {
9862 let storage = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0, 7.0]);
9863
9864 let graph = Arc::new(RwLock::new(Graph::new()));
9865 let const_id = {
9866 let mut g = graph.write().unwrap();
9867 g.push(Node {
9868 op: Op::Const,
9869 inputs: vec![],
9870 shape: Shape::from_dims(&[3]),
9871 dtype: DType::F32,
9872 })
9873 };
9874
9875 let mut inputs = StorageCache::new();
9876 inputs.insert(const_id, Arc::new(RwLock::new(storage)));
9877
9878 let (result_arc, _) =
9879 PipelinedExecutor::realize(graph, const_id, inputs).expect("realize const");
9880 let guard = result_arc.read().unwrap();
9881 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
9882 let typed: &[f32] = c.as_slice().unwrap();
9883 assert_eq!(typed, &[5.0, 6.0, 7.0]);
9884 }
9885 }
9886
9887 #[test]
9889 fn pipelined_errors_on_missing_const_input() {
9890 let graph = Arc::new(RwLock::new(Graph::new()));
9891 let const_id = {
9892 let mut g = graph.write().unwrap();
9893 g.push(Node {
9894 op: Op::Const,
9895 inputs: vec![],
9896 shape: Shape::from_dims(&[2]),
9897 dtype: DType::F32,
9898 })
9899 };
9900 let inputs = StorageCache::new();
9902 let result = PipelinedExecutor::realize(graph, const_id, inputs);
9903 assert!(result.is_err());
9904 }
9905
9906 #[test]
9909 fn pipelined_realize_const_relu() {
9910 let storage = fuel_memory::from_slice_cpu(&[-1.0_f32, 0.0, 0.5, -3.5, 7.25]);
9911 let graph = Arc::new(RwLock::new(Graph::new()));
9912 let (in_id, relu_id) = {
9913 let mut g = graph.write().unwrap();
9914 let in_id = g.push(Node {
9915 op: Op::Const,
9916 inputs: vec![],
9917 shape: Shape::from_dims(&[5]),
9918 dtype: DType::F32,
9919 });
9920 let relu_id = g.push(Node {
9921 op: Op::Relu,
9922 inputs: vec![in_id],
9923 shape: Shape::from_dims(&[5]),
9924 dtype: DType::F32,
9925 });
9926 g.set_target_backend(relu_id, BackendId::Cpu);
9927 (in_id, relu_id)
9928 };
9929 let mut inputs = StorageCache::new();
9930 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
9931
9932 let (result_arc, _) = PipelinedExecutor::realize(graph, relu_id, inputs).expect("realize");
9933 let guard = result_arc.read().unwrap();
9934 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
9935 let typed: &[f32] = c.as_slice().expect("f32 cast");
9936 assert_eq!(typed, &[0.0, 0.0, 0.5, 0.0, 7.25]);
9937 }
9938 }
9939
9940 #[test]
9944 fn pipelined_realize_chained_binary_ops() {
9945 let a_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
9946 let b_storage = fuel_memory::from_slice_cpu(&[3.0_f32, 5.0]);
9947 let c_storage = fuel_memory::from_slice_cpu(&[2.0_f32, 4.0]);
9948
9949 let graph = Arc::new(RwLock::new(Graph::new()));
9950 let (a_id, b_id, c_id, sub_id, mul_id, div_id) = {
9951 let mut g = graph.write().unwrap();
9952 let a = g.push(Node {
9953 op: Op::Const,
9954 inputs: vec![],
9955 shape: Shape::from_dims(&[2]),
9956 dtype: DType::F32,
9957 });
9958 let b = g.push(Node {
9959 op: Op::Const,
9960 inputs: vec![],
9961 shape: Shape::from_dims(&[2]),
9962 dtype: DType::F32,
9963 });
9964 let c = g.push(Node {
9965 op: Op::Const,
9966 inputs: vec![],
9967 shape: Shape::from_dims(&[2]),
9968 dtype: DType::F32,
9969 });
9970 let sub = g.push(Node {
9972 op: Op::Sub,
9973 inputs: vec![a, b],
9974 shape: Shape::from_dims(&[2]),
9975 dtype: DType::F32,
9976 });
9977 let mul = g.push(Node {
9979 op: Op::Mul,
9980 inputs: vec![sub, c],
9981 shape: Shape::from_dims(&[2]),
9982 dtype: DType::F32,
9983 });
9984 let div = g.push(Node {
9986 op: Op::Div,
9987 inputs: vec![mul, b],
9988 shape: Shape::from_dims(&[2]),
9989 dtype: DType::F32,
9990 });
9991 g.set_target_backend(sub, BackendId::Cpu);
9992 g.set_target_backend(mul, BackendId::Cpu);
9993 g.set_target_backend(div, BackendId::Cpu);
9994 (a, b, c, sub, mul, div)
9995 };
9996
9997 let _ = (sub_id, mul_id);
9998 let mut inputs = StorageCache::new();
9999 inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
10000 inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
10001 inputs.insert(c_id, Arc::new(RwLock::new(c_storage)));
10002
10003 let (result_arc, _) = PipelinedExecutor::realize(graph, div_id, inputs).expect("realize");
10004 let guard = result_arc.read().unwrap();
10005 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
10006 let typed: &[f32] = c.as_slice().unwrap();
10007 assert!((typed[0] - (14.0_f32 / 3.0)).abs() < 1e-6);
10008 assert!((typed[1] - 12.0).abs() < 1e-6);
10009 }
10010 }
10011
10012 #[test]
10016 fn pipelined_realize_transpose_is_metadata_only() {
10017 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
10019 let graph = Arc::new(RwLock::new(Graph::new()));
10020 let (in_id, t_id) = {
10021 let mut g = graph.write().unwrap();
10022 let in_id = g.push(Node {
10023 op: Op::Const,
10024 inputs: vec![],
10025 shape: Shape::from_dims(&[2, 3]),
10026 dtype: DType::F32,
10027 });
10028 let t_id = g.push(Node {
10029 op: Op::Transpose,
10030 inputs: vec![in_id],
10031 shape: Shape::from_dims(&[3, 2]),
10032 dtype: DType::F32,
10033 });
10034 (in_id, t_id)
10037 };
10038 let mut inputs = StorageCache::new();
10039 let in_arc = Arc::new(RwLock::new(storage));
10040 inputs.insert(in_id, Arc::clone(&in_arc));
10041
10042 let (result_arc, result_layout) =
10043 PipelinedExecutor::realize(graph, t_id, inputs).expect("realize");
10044
10045 assert!(
10048 Arc::ptr_eq(&result_arc, &in_arc),
10049 "transpose must share input bytes"
10050 );
10051
10052 assert_eq!(result_layout.shape().dims(), &[3, 2]);
10054 assert_eq!(result_layout.stride(), &[1, 3]);
10055 assert!(!result_layout.is_contiguous());
10056 }
10057
10058 #[test]
10063 fn pipelined_realize_flip_is_metadata_only() {
10064 let data: Vec<f32> = (1..=12).map(|x| x as f32).collect();
10066 let storage = fuel_memory::from_slice_cpu(&data);
10067 let graph = Arc::new(RwLock::new(Graph::new()));
10068 let (in_id, f_id) = {
10069 let mut g = graph.write().unwrap();
10070 let in_id = g.push(Node {
10071 op: Op::Const,
10072 inputs: vec![],
10073 shape: Shape::from_dims(&[3, 4]),
10074 dtype: DType::F32,
10075 });
10076 let f_id = g.push(Node {
10077 op: Op::Flip { dim: 0 },
10078 inputs: vec![in_id],
10079 shape: Shape::from_dims(&[3, 4]),
10080 dtype: DType::F32,
10081 });
10082 (in_id, f_id)
10083 };
10084 let mut inputs = StorageCache::new();
10085 let in_arc = Arc::new(RwLock::new(storage));
10086 inputs.insert(in_id, Arc::clone(&in_arc));
10087
10088 let (result_arc, result_layout) =
10089 PipelinedExecutor::realize(graph, f_id, inputs).expect("realize");
10090
10091 assert!(
10092 Arc::ptr_eq(&result_arc, &in_arc),
10093 "flip must share input bytes"
10094 );
10095 assert_eq!(result_layout.shape().dims(), &[3, 4]);
10096 assert_eq!(result_layout.stride(), &[-4_isize, 1]);
10097 assert_eq!(result_layout.start_offset(), 8);
10098 assert!(!result_layout.is_contiguous());
10099 }
10100
10101 #[test]
10104 fn pipelined_realize_permute_is_metadata_only() {
10105 let data: Vec<f32> = (1..=24).map(|x| x as f32).collect();
10107 let storage = fuel_memory::from_slice_cpu(&data);
10108 let graph = Arc::new(RwLock::new(Graph::new()));
10109 let (in_id, p_id) = {
10110 let mut g = graph.write().unwrap();
10111 let in_id = g.push(Node {
10112 op: Op::Const,
10113 inputs: vec![],
10114 shape: Shape::from_dims(&[2, 3, 4]),
10115 dtype: DType::F32,
10116 });
10117 let p_id = g.push(Node {
10118 op: Op::Permute(vec![2, 0, 1]),
10119 inputs: vec![in_id],
10120 shape: Shape::from_dims(&[4, 2, 3]),
10121 dtype: DType::F32,
10122 });
10123 (in_id, p_id)
10124 };
10125 let mut inputs = StorageCache::new();
10126 let in_arc = Arc::new(RwLock::new(storage));
10127 inputs.insert(in_id, Arc::clone(&in_arc));
10128
10129 let (result_arc, result_layout) =
10130 PipelinedExecutor::realize(graph, p_id, inputs).expect("realize");
10131
10132 assert!(Arc::ptr_eq(&result_arc, &in_arc));
10133 assert_eq!(result_layout.shape().dims(), &[4, 2, 3]);
10134 assert_eq!(result_layout.stride(), &[1, 12, 4]);
10137 }
10138
10139 #[test]
10143 fn pipelined_realize_cast_f32_to_f64() {
10144 let storage = fuel_memory::from_slice_cpu(&[1.5_f32, -2.25, 100.0]);
10145 let graph = Arc::new(RwLock::new(Graph::new()));
10146 let (in_id, c_id) = {
10147 let mut g = graph.write().unwrap();
10148 let in_id = g.push(Node {
10149 op: Op::Const,
10150 inputs: vec![],
10151 shape: Shape::from_dims(&[3]),
10152 dtype: DType::F32,
10153 });
10154 let c_id = g.push(Node {
10155 op: Op::Cast(DType::F64),
10156 inputs: vec![in_id],
10157 shape: Shape::from_dims(&[3]),
10158 dtype: DType::F64,
10159 });
10160 g.set_target_backend(c_id, BackendId::Cpu);
10161 (in_id, c_id)
10162 };
10163 let mut inputs = StorageCache::new();
10164 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
10165
10166 let (result_arc, result_layout) =
10167 PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
10168 assert_eq!(result_layout.shape().dims(), &[3]);
10169
10170 let guard = result_arc.read().unwrap();
10171 assert_eq!(guard.dtype, DType::F64);
10172 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10173 panic!("expected Cpu storage");
10174 };
10175 assert_eq!(c.as_slice::<f64>().unwrap(), &[1.5_f64, -2.25, 100.0]);
10176 }
10177
10178 #[test]
10183 fn pipelined_realize_cast_round_trip_via_bf16() {
10184 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, -3.0, 0.5]);
10185 let graph = Arc::new(RwLock::new(Graph::new()));
10186 let (in_id, c1_id, c2_id) = {
10187 let mut g = graph.write().unwrap();
10188 let in_id = g.push(Node {
10189 op: Op::Const,
10190 inputs: vec![],
10191 shape: Shape::from_dims(&[4]),
10192 dtype: DType::F32,
10193 });
10194 let c1_id = g.push(Node {
10195 op: Op::Cast(DType::BF16),
10196 inputs: vec![in_id],
10197 shape: Shape::from_dims(&[4]),
10198 dtype: DType::BF16,
10199 });
10200 let c2_id = g.push(Node {
10201 op: Op::Cast(DType::F32),
10202 inputs: vec![c1_id],
10203 shape: Shape::from_dims(&[4]),
10204 dtype: DType::F32,
10205 });
10206 g.set_target_backend(c1_id, BackendId::Cpu);
10207 g.set_target_backend(c2_id, BackendId::Cpu);
10208 (in_id, c1_id, c2_id)
10209 };
10210 let _ = c1_id;
10211 let mut inputs = StorageCache::new();
10212 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
10213
10214 let (result_arc, _) = PipelinedExecutor::realize(graph, c2_id, inputs).expect("realize");
10215 let guard = result_arc.read().unwrap();
10216 assert_eq!(guard.dtype, DType::F32);
10217 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10218 panic!("expected Cpu storage");
10219 };
10220 assert_eq!(c.as_slice::<f32>().unwrap(), &[1.0_f32, 2.0, -3.0, 0.5]);
10221 }
10222
10223 #[test]
10229 fn pipelined_realize_slice_is_metadata_only() {
10230 let storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0]);
10233 let graph = Arc::new(RwLock::new(Graph::new()));
10234 let (in_id, s_id) = {
10235 let mut g = graph.write().unwrap();
10236 let in_id = g.push(Node {
10237 op: Op::Const,
10238 inputs: vec![],
10239 shape: Shape::from_dims(&[5]),
10240 dtype: DType::F32,
10241 });
10242 let s_id = g.push(Node {
10243 op: Op::Slice {
10244 dim: 0,
10245 start: 1,
10246 len: 3,
10247 },
10248 inputs: vec![in_id],
10249 shape: Shape::from_dims(&[3]),
10250 dtype: DType::F32,
10251 });
10252 (in_id, s_id)
10253 };
10254 let mut inputs = StorageCache::new();
10255 let in_arc = Arc::new(RwLock::new(storage));
10256 inputs.insert(in_id, Arc::clone(&in_arc));
10257
10258 let (result_arc, result_layout) =
10259 PipelinedExecutor::realize(graph, s_id, inputs).expect("realize");
10260
10261 assert!(Arc::ptr_eq(&result_arc, &in_arc));
10263 assert_eq!(result_layout.shape().dims(), &[3]);
10264 assert_eq!(result_layout.start_offset(), 1);
10268 assert_eq!(result_layout.stride(), &[1]);
10269 }
10270
10271 #[test]
10276 fn pipelined_realize_slice_then_sum_all() {
10277 let storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0]);
10278 let graph = Arc::new(RwLock::new(Graph::new()));
10279 let (in_id, s_id, sum_id) = {
10280 let mut g = graph.write().unwrap();
10281 let in_id = g.push(Node {
10282 op: Op::Const,
10283 inputs: vec![],
10284 shape: Shape::from_dims(&[5]),
10285 dtype: DType::F32,
10286 });
10287 let s_id = g.push(Node {
10288 op: Op::Slice {
10289 dim: 0,
10290 start: 1,
10291 len: 3,
10292 },
10293 inputs: vec![in_id],
10294 shape: Shape::from_dims(&[3]),
10295 dtype: DType::F32,
10296 });
10297 let sum_id = g.push(Node {
10298 op: Op::SumAll,
10299 inputs: vec![s_id],
10300 shape: Shape::from_dims(&[]),
10301 dtype: DType::F32,
10302 });
10303 g.set_target_backend(sum_id, BackendId::Cpu);
10304 (in_id, s_id, sum_id)
10305 };
10306 let _ = s_id;
10307 let mut inputs = StorageCache::new();
10308 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
10309
10310 let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
10311 let guard = result_arc.read().unwrap();
10312 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10313 panic!("expected Cpu storage");
10314 };
10315 assert_eq!(c.as_slice::<f32>().unwrap(), &[90.0]);
10316 }
10317
10318 #[test]
10324 fn pipelined_realize_flip_then_relu_materializes_in_reverse() {
10325 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
10326 let graph = Arc::new(RwLock::new(Graph::new()));
10327 let (in_id, _f_id, r_id) = {
10328 let mut g = graph.write().unwrap();
10329 let in_id = g.push(Node {
10330 op: Op::Const,
10331 inputs: vec![],
10332 shape: Shape::from_dims(&[4]),
10333 dtype: DType::F32,
10334 });
10335 let f_id = g.push(Node {
10336 op: Op::Flip { dim: 0 },
10337 inputs: vec![in_id],
10338 shape: Shape::from_dims(&[4]),
10339 dtype: DType::F32,
10340 });
10341 let r_id = g.push(Node {
10342 op: Op::Relu,
10343 inputs: vec![f_id],
10344 shape: Shape::from_dims(&[4]),
10345 dtype: DType::F32,
10346 });
10347 g.set_target_backend(r_id, BackendId::Cpu);
10348 (in_id, f_id, r_id)
10349 };
10350 let mut inputs = StorageCache::new();
10351 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
10352
10353 let (result_arc, result_layout) =
10354 PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
10355 let guard = result_arc.read().unwrap();
10356 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10357 panic!("expected Cpu storage");
10358 };
10359 assert_eq!(c.as_slice::<f32>().unwrap(), &[4.0, 3.0, 2.0, 1.0]);
10360 assert_eq!(result_layout.shape().dims(), &[4]);
10361 assert!(result_layout.is_contiguous(), "relu output is contiguous");
10362 }
10363
10364 #[test]
10368 fn pipelined_realize_broadcast_is_metadata_only() {
10369 let storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
10371 let graph = Arc::new(RwLock::new(Graph::new()));
10372 let (in_id, b_id) = {
10373 let mut g = graph.write().unwrap();
10374 let in_id = g.push(Node {
10375 op: Op::Const,
10376 inputs: vec![],
10377 shape: Shape::from_dims(&[3]),
10378 dtype: DType::F32,
10379 });
10380 let b_id = g.push(Node {
10381 op: Op::BroadcastTo(Shape::from_dims(&[4, 3])),
10382 inputs: vec![in_id],
10383 shape: Shape::from_dims(&[4, 3]),
10384 dtype: DType::F32,
10385 });
10386 (in_id, b_id)
10387 };
10388 let mut inputs = StorageCache::new();
10389 let in_arc = Arc::new(RwLock::new(storage));
10390 inputs.insert(in_id, Arc::clone(&in_arc));
10391
10392 let (result_arc, result_layout) =
10393 PipelinedExecutor::realize(graph, b_id, inputs).expect("realize");
10394
10395 assert!(Arc::ptr_eq(&result_arc, &in_arc));
10396 assert_eq!(result_layout.shape().dims(), &[4, 3]);
10397 assert_eq!(result_layout.stride(), &[0, 1]);
10399 }
10400
10401 #[test]
10406 fn pipelined_realize_matmul_2x3_times_3x2() {
10407 let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
10410 let rhs_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
10411
10412 let graph = Arc::new(RwLock::new(Graph::new()));
10413 let (lhs_id, rhs_id, mm_id) = {
10414 let mut g = graph.write().unwrap();
10415 let lhs = g.push(Node {
10416 op: Op::Const,
10417 inputs: vec![],
10418 shape: Shape::from_dims(&[2, 3]),
10419 dtype: DType::F32,
10420 });
10421 let rhs = g.push(Node {
10422 op: Op::Const,
10423 inputs: vec![],
10424 shape: Shape::from_dims(&[3, 2]),
10425 dtype: DType::F32,
10426 });
10427 let mm = g.push(Node {
10428 op: Op::MatMul,
10429 inputs: vec![lhs, rhs],
10430 shape: Shape::from_dims(&[2, 2]),
10431 dtype: DType::F32,
10432 });
10433 g.set_target_backend(mm, BackendId::Cpu);
10434 (lhs, rhs, mm)
10435 };
10436 let mut inputs = StorageCache::new();
10437 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
10438 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
10439
10440 let (result_arc, result_layout) =
10441 PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
10442 assert_eq!(result_layout.shape().dims(), &[2, 2]);
10443 assert!(result_layout.is_contiguous());
10444
10445 let guard = result_arc.read().unwrap();
10446 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10447 panic!("expected Cpu storage");
10448 };
10449 assert_eq!(c.as_slice::<f32>().unwrap(), &[58.0, 64.0, 139.0, 154.0]);
10450 }
10451
10452 #[test]
10456 fn pipelined_realize_matmul_batched_2x_2x2_times_2x2() {
10457 let lhs_storage = fuel_memory::from_slice_cpu(&[
10458 1.0_f32, 2.0, 3.0, 4.0, 1.0, 0.0, 0.0, 1.0, ]);
10461 let rhs_storage = fuel_memory::from_slice_cpu(&[
10462 5.0_f32, 6.0, 7.0, 8.0, 10.0, 20.0, 30.0, 40.0, ]);
10465
10466 let graph = Arc::new(RwLock::new(Graph::new()));
10467 let (lhs_id, rhs_id, mm_id) = {
10468 let mut g = graph.write().unwrap();
10469 let lhs = g.push(Node {
10470 op: Op::Const,
10471 inputs: vec![],
10472 shape: Shape::from_dims(&[2, 2, 2]),
10473 dtype: DType::F32,
10474 });
10475 let rhs = g.push(Node {
10476 op: Op::Const,
10477 inputs: vec![],
10478 shape: Shape::from_dims(&[2, 2, 2]),
10479 dtype: DType::F32,
10480 });
10481 let mm = g.push(Node {
10482 op: Op::MatMul,
10483 inputs: vec![lhs, rhs],
10484 shape: Shape::from_dims(&[2, 2, 2]),
10485 dtype: DType::F32,
10486 });
10487 g.set_target_backend(mm, BackendId::Cpu);
10488 (lhs, rhs, mm)
10489 };
10490 let mut inputs = StorageCache::new();
10491 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
10492 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
10493
10494 let (result_arc, result_layout) =
10495 PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
10496 assert_eq!(result_layout.shape().dims(), &[2, 2, 2]);
10497 assert!(result_layout.is_contiguous());
10498
10499 let guard = result_arc.read().unwrap();
10500 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10501 panic!("expected Cpu storage");
10502 };
10503 assert_eq!(
10506 c.as_slice::<f32>().unwrap(),
10507 &[19.0, 22.0, 43.0, 50.0, 10.0, 20.0, 30.0, 40.0]
10508 );
10509 }
10510
10511 #[test]
10515 fn pipelined_realize_add_f64() {
10516 let lhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0]);
10517 let rhs = fuel_memory::from_slice_cpu(&[10.0_f64, 20.0, 30.0]);
10518 let graph = Arc::new(RwLock::new(Graph::new()));
10519 let (l_id, r_id, op_id) = {
10520 let mut g = graph.write().unwrap();
10521 let l = g.push(Node {
10522 op: Op::Const,
10523 inputs: vec![],
10524 shape: Shape::from_dims(&[3]),
10525 dtype: DType::F64,
10526 });
10527 let r = g.push(Node {
10528 op: Op::Const,
10529 inputs: vec![],
10530 shape: Shape::from_dims(&[3]),
10531 dtype: DType::F64,
10532 });
10533 let op = g.push(Node {
10534 op: Op::Add,
10535 inputs: vec![l, r],
10536 shape: Shape::from_dims(&[3]),
10537 dtype: DType::F64,
10538 });
10539 g.set_target_backend(op, BackendId::Cpu);
10540 (l, r, op)
10541 };
10542 let mut inputs = StorageCache::new();
10543 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10544 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10545 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
10546 let guard = result_arc.read().unwrap();
10547 assert_eq!(guard.dtype, DType::F64);
10548 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10549 panic!("expected Cpu storage");
10550 };
10551 assert_eq!(c.as_slice::<f64>().unwrap(), &[11.0_f64, 22.0, 33.0]);
10552 }
10553
10554 #[test]
10561 fn pipelined_realize_eq_f32_to_bool_mask() {
10562 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, f32::NAN, 0.0]);
10563 let rhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -0.0]);
10564 let graph = Arc::new(RwLock::new(Graph::new()));
10565 let (l_id, r_id, eq_id) = {
10566 let mut g = graph.write().unwrap();
10567 let l = g.push(Node {
10568 op: Op::Const,
10569 inputs: vec![],
10570 shape: Shape::from_dims(&[5]),
10571 dtype: DType::F32,
10572 });
10573 let r = g.push(Node {
10574 op: Op::Const,
10575 inputs: vec![],
10576 shape: Shape::from_dims(&[5]),
10577 dtype: DType::F32,
10578 });
10579 let eq = g.push(Node {
10580 op: Op::Equal,
10581 inputs: vec![l, r],
10582 shape: Shape::from_dims(&[5]),
10583 dtype: DType::Bool,
10584 });
10585 g.set_target_backend(eq, BackendId::Cpu);
10586 (l, r, eq)
10587 };
10588 let mut inputs = StorageCache::new();
10589 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10590 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10591 let (result_arc, _) = PipelinedExecutor::realize(graph, eq_id, inputs).expect("realize");
10592 let guard = result_arc.read().unwrap();
10593 assert_eq!(guard.dtype, DType::Bool);
10594 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10595 panic!("expected Cpu storage");
10596 };
10597 let mask: &[u8] = c.as_slice().expect("u8 view");
10598 assert_eq!(mask, &[1, 0, 1, 0, 1]);
10604 }
10605
10606 #[test]
10610 fn pipelined_realize_ne_f32_to_bool_mask() {
10611 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, f32::NAN, 0.0]);
10612 let rhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -0.0]);
10613 let graph = Arc::new(RwLock::new(Graph::new()));
10614 let (l_id, r_id, ne_id) = {
10615 let mut g = graph.write().unwrap();
10616 let l = g.push(Node {
10617 op: Op::Const,
10618 inputs: vec![],
10619 shape: Shape::from_dims(&[5]),
10620 dtype: DType::F32,
10621 });
10622 let r = g.push(Node {
10623 op: Op::Const,
10624 inputs: vec![],
10625 shape: Shape::from_dims(&[5]),
10626 dtype: DType::F32,
10627 });
10628 let ne = g.push(Node {
10629 op: Op::Ne,
10630 inputs: vec![l, r],
10631 shape: Shape::from_dims(&[5]),
10632 dtype: DType::Bool,
10633 });
10634 g.set_target_backend(ne, BackendId::Cpu);
10635 (l, r, ne)
10636 };
10637 let mut inputs = StorageCache::new();
10638 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10639 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10640 let (result_arc, _) = PipelinedExecutor::realize(graph, ne_id, inputs).expect("realize");
10641 let guard = result_arc.read().unwrap();
10642 assert_eq!(guard.dtype, DType::Bool);
10643 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10644 panic!("expected Cpu storage");
10645 };
10646 let mask: &[u8] = c.as_slice().expect("u8 view");
10647 assert_eq!(mask, &[0, 1, 0, 1, 0]);
10651 }
10652
10653 #[test]
10658 fn pipelined_realize_lt_f32_to_bool_mask() {
10659 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -1.0]);
10660 let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 3.0, 0.0, 0.0]);
10661 let graph = Arc::new(RwLock::new(Graph::new()));
10662 let (l_id, r_id, lt_id) = {
10663 let mut g = graph.write().unwrap();
10664 let l = g.push(Node {
10665 op: Op::Const,
10666 inputs: vec![],
10667 shape: Shape::from_dims(&[5]),
10668 dtype: DType::F32,
10669 });
10670 let r = g.push(Node {
10671 op: Op::Const,
10672 inputs: vec![],
10673 shape: Shape::from_dims(&[5]),
10674 dtype: DType::F32,
10675 });
10676 let lt = g.push(Node {
10677 op: Op::Lt,
10678 inputs: vec![l, r],
10679 shape: Shape::from_dims(&[5]),
10680 dtype: DType::Bool,
10681 });
10682 g.set_target_backend(lt, BackendId::Cpu);
10683 (l, r, lt)
10684 };
10685 let mut inputs = StorageCache::new();
10686 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10687 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10688 let (result_arc, _) = PipelinedExecutor::realize(graph, lt_id, inputs).expect("realize");
10689 let guard = result_arc.read().unwrap();
10690 assert_eq!(guard.dtype, DType::Bool);
10691 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10692 panic!("expected Cpu storage");
10693 };
10694 let mask: &[u8] = c.as_slice().expect("u8 view");
10695 assert_eq!(mask, &[1, 0, 0, 0, 1]);
10698 }
10699
10700 #[test]
10703 fn pipelined_realize_le_f32_to_bool_mask() {
10704 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -1.0]);
10705 let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 2.0, 0.0, 0.0]);
10706 let graph = Arc::new(RwLock::new(Graph::new()));
10707 let (l_id, r_id, le_id) = {
10708 let mut g = graph.write().unwrap();
10709 let l = g.push(Node {
10710 op: Op::Const,
10711 inputs: vec![],
10712 shape: Shape::from_dims(&[5]),
10713 dtype: DType::F32,
10714 });
10715 let r = g.push(Node {
10716 op: Op::Const,
10717 inputs: vec![],
10718 shape: Shape::from_dims(&[5]),
10719 dtype: DType::F32,
10720 });
10721 let le = g.push(Node {
10722 op: Op::Le,
10723 inputs: vec![l, r],
10724 shape: Shape::from_dims(&[5]),
10725 dtype: DType::Bool,
10726 });
10727 g.set_target_backend(le, BackendId::Cpu);
10728 (l, r, le)
10729 };
10730 let mut inputs = StorageCache::new();
10731 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10732 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10733 let (result_arc, _) = PipelinedExecutor::realize(graph, le_id, inputs).expect("realize");
10734 let guard = result_arc.read().unwrap();
10735 assert_eq!(guard.dtype, DType::Bool);
10736 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10737 panic!("expected Cpu storage");
10738 };
10739 let mask: &[u8] = c.as_slice().expect("u8 view");
10740 assert_eq!(mask, &[1, 1, 0, 0, 1]);
10743 }
10744
10745 #[test]
10748 fn pipelined_realize_gt_f32_to_bool_mask() {
10749 let lhs = fuel_memory::from_slice_cpu(&[3.0_f32, 5.0, 2.0, f32::NAN, 1.0]);
10750 let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 3.0, 0.0, 0.0]);
10751 let graph = Arc::new(RwLock::new(Graph::new()));
10752 let (l_id, r_id, gt_id) = {
10753 let mut g = graph.write().unwrap();
10754 let l = g.push(Node {
10755 op: Op::Const,
10756 inputs: vec![],
10757 shape: Shape::from_dims(&[5]),
10758 dtype: DType::F32,
10759 });
10760 let r = g.push(Node {
10761 op: Op::Const,
10762 inputs: vec![],
10763 shape: Shape::from_dims(&[5]),
10764 dtype: DType::F32,
10765 });
10766 let gt = g.push(Node {
10767 op: Op::Gt,
10768 inputs: vec![l, r],
10769 shape: Shape::from_dims(&[5]),
10770 dtype: DType::Bool,
10771 });
10772 g.set_target_backend(gt, BackendId::Cpu);
10773 (l, r, gt)
10774 };
10775 let mut inputs = StorageCache::new();
10776 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10777 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10778 let (result_arc, _) = PipelinedExecutor::realize(graph, gt_id, inputs).expect("realize");
10779 let guard = result_arc.read().unwrap();
10780 assert_eq!(guard.dtype, DType::Bool);
10781 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10782 panic!("expected Cpu storage");
10783 };
10784 let mask: &[u8] = c.as_slice().expect("u8 view");
10785 assert_eq!(mask, &[1, 0, 0, 0, 1]);
10788 }
10789
10790 #[test]
10795 fn pipelined_realize_ge_f32_to_bool_mask() {
10796 let lhs = fuel_memory::from_slice_cpu(&[3.0_f32, 5.0, 2.0, f32::NAN, 0.0]);
10797 let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 3.0, 0.0, 0.0]);
10798 let graph = Arc::new(RwLock::new(Graph::new()));
10799 let (l_id, r_id, ge_id) = {
10800 let mut g = graph.write().unwrap();
10801 let l = g.push(Node {
10802 op: Op::Const,
10803 inputs: vec![],
10804 shape: Shape::from_dims(&[5]),
10805 dtype: DType::F32,
10806 });
10807 let r = g.push(Node {
10808 op: Op::Const,
10809 inputs: vec![],
10810 shape: Shape::from_dims(&[5]),
10811 dtype: DType::F32,
10812 });
10813 let ge = g.push(Node {
10814 op: Op::Ge,
10815 inputs: vec![l, r],
10816 shape: Shape::from_dims(&[5]),
10817 dtype: DType::Bool,
10818 });
10819 g.set_target_backend(ge, BackendId::Cpu);
10820 (l, r, ge)
10821 };
10822 let mut inputs = StorageCache::new();
10823 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10824 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10825 let (result_arc, _) = PipelinedExecutor::realize(graph, ge_id, inputs).expect("realize");
10826 let guard = result_arc.read().unwrap();
10827 assert_eq!(guard.dtype, DType::Bool);
10828 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10829 panic!("expected Cpu storage");
10830 };
10831 let mask: &[u8] = c.as_slice().expect("u8 view");
10832 assert_eq!(mask, &[1, 1, 0, 0, 1]);
10835 }
10836
10837 #[test]
10841 fn pipelined_realize_eq_f64_to_bool_mask() {
10842 let lhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0]);
10843 let rhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 4.0]);
10844 let graph = Arc::new(RwLock::new(Graph::new()));
10845 let (l_id, r_id, eq_id) = {
10846 let mut g = graph.write().unwrap();
10847 let l = g.push(Node {
10848 op: Op::Const,
10849 inputs: vec![],
10850 shape: Shape::from_dims(&[3]),
10851 dtype: DType::F64,
10852 });
10853 let r = g.push(Node {
10854 op: Op::Const,
10855 inputs: vec![],
10856 shape: Shape::from_dims(&[3]),
10857 dtype: DType::F64,
10858 });
10859 let eq = g.push(Node {
10860 op: Op::Equal,
10861 inputs: vec![l, r],
10862 shape: Shape::from_dims(&[3]),
10863 dtype: DType::Bool,
10864 });
10865 g.set_target_backend(eq, BackendId::Cpu);
10866 (l, r, eq)
10867 };
10868 let mut inputs = StorageCache::new();
10869 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
10870 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
10871 let (result_arc, _) = PipelinedExecutor::realize(graph, eq_id, inputs).expect("realize");
10872 let guard = result_arc.read().unwrap();
10873 assert_eq!(guard.dtype, DType::Bool);
10874 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10875 panic!("expected Cpu storage");
10876 };
10877 let mask: &[u8] = c.as_slice().expect("u8 view");
10878 assert_eq!(mask, &[1, 1, 0]);
10879 }
10880
10881 #[test]
10890 fn pipelined_realize_where_f32_picks_per_slot_from_bool_mask() {
10891 let mut cond_storage = fuel_memory::from_slice_cpu(&[1u8, 0, 1, 0, 1]);
10894 cond_storage.dtype = DType::Bool; let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0]);
10896 let b_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0]);
10897 let graph = Arc::new(RwLock::new(Graph::new()));
10898 let (cond_id, a_id, b_id, where_id) = {
10899 let mut g = graph.write().unwrap();
10900 let cond = g.push(Node {
10901 op: Op::Const,
10902 inputs: vec![],
10903 shape: Shape::from_dims(&[5]),
10904 dtype: DType::Bool,
10905 });
10906 let a = g.push(Node {
10907 op: Op::Const,
10908 inputs: vec![],
10909 shape: Shape::from_dims(&[5]),
10910 dtype: DType::F32,
10911 });
10912 let b = g.push(Node {
10913 op: Op::Const,
10914 inputs: vec![],
10915 shape: Shape::from_dims(&[5]),
10916 dtype: DType::F32,
10917 });
10918 let w = g.push(Node {
10919 op: Op::Where,
10920 inputs: vec![cond, a, b],
10921 shape: Shape::from_dims(&[5]),
10922 dtype: DType::F32,
10923 });
10924 g.set_target_backend(w, BackendId::Cpu);
10925 (cond, a, b, w)
10926 };
10927 let mut inputs = StorageCache::new();
10928 inputs.insert(cond_id, Arc::new(RwLock::new(cond_storage)));
10929 inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
10930 inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
10931 let (result_arc, _) = PipelinedExecutor::realize(graph, where_id, inputs).expect("realize");
10932 let guard = result_arc.read().unwrap();
10933 assert_eq!(guard.dtype, DType::F32);
10934 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
10935 panic!("expected Cpu storage");
10936 };
10937 let out: &[f32] = c.as_slice().expect("f32 view");
10938 assert_eq!(out, &[1.0, 20.0, 3.0, 40.0, 5.0]);
10939 }
10940
10941 #[test]
10946 fn pipelined_realize_eq_then_where_full_chain() {
10947 let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
10951 let b_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0]);
10952 let pick_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
10953 let fb_storage = fuel_memory::from_slice_cpu(&[99.0_f32, 99.0, 99.0]);
10954 let graph = Arc::new(RwLock::new(Graph::new()));
10955 let (a_id, b_id, pick_id, fb_id, where_id) = {
10956 let mut g = graph.write().unwrap();
10957 let a = g.push(Node {
10958 op: Op::Const,
10959 inputs: vec![],
10960 shape: Shape::from_dims(&[3]),
10961 dtype: DType::F32,
10962 });
10963 let b = g.push(Node {
10964 op: Op::Const,
10965 inputs: vec![],
10966 shape: Shape::from_dims(&[3]),
10967 dtype: DType::F32,
10968 });
10969 let pick = g.push(Node {
10970 op: Op::Const,
10971 inputs: vec![],
10972 shape: Shape::from_dims(&[3]),
10973 dtype: DType::F32,
10974 });
10975 let fb = g.push(Node {
10976 op: Op::Const,
10977 inputs: vec![],
10978 shape: Shape::from_dims(&[3]),
10979 dtype: DType::F32,
10980 });
10981 let eq = g.push(Node {
10982 op: Op::Equal,
10983 inputs: vec![a, b],
10984 shape: Shape::from_dims(&[3]),
10985 dtype: DType::Bool,
10986 });
10987 let w = g.push(Node {
10988 op: Op::Where,
10989 inputs: vec![eq, pick, fb],
10990 shape: Shape::from_dims(&[3]),
10991 dtype: DType::F32,
10992 });
10993 g.set_target_backend(eq, BackendId::Cpu);
10994 g.set_target_backend(w, BackendId::Cpu);
10995 (a, b, pick, fb, w)
10996 };
10997 let mut inputs = StorageCache::new();
10998 inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
10999 inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
11000 inputs.insert(pick_id, Arc::new(RwLock::new(pick_storage)));
11001 inputs.insert(fb_id, Arc::new(RwLock::new(fb_storage)));
11002 let (result_arc, _) = PipelinedExecutor::realize(graph, where_id, inputs).expect("realize");
11003 let guard = result_arc.read().unwrap();
11004 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11005 panic!("expected Cpu storage");
11006 };
11007 let out: &[f32] = c.as_slice().expect("f32 view");
11008 assert_eq!(out, &[10.0, 99.0, 30.0]);
11009 }
11010
11011 #[test]
11018 fn pipelined_realize_qmatmul_q4_0_unit_weight_sums_activations() {
11019 use fuel_graph::QuantType;
11020 use half::f16;
11021 let block_size = std::mem::size_of::<fuel_quantized::BlockQ4_0>();
11022 let mut w_bytes = vec![0u8; 2 * block_size];
11023 for block_idx in 0..2 {
11024 let off = block_idx * block_size;
11025 let d_bytes = f16::from_f32(1.0).to_le_bytes();
11026 w_bytes[off..off + 2].copy_from_slice(&d_bytes);
11027 for i in 0..16 {
11028 w_bytes[off + 2 + i] = 0x99;
11029 }
11030 }
11031 let w_storage = fuel_memory::Storage::new(
11033 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_bytes(
11034 &w_bytes,
11035 )),
11036 DType::U32,
11037 );
11038
11039 let act_vec: Vec<f32> = (1..=32).map(|x| x as f32).collect();
11040 let act_storage = fuel_memory::from_slice_cpu(&act_vec);
11041
11042 let graph = Arc::new(RwLock::new(Graph::new()));
11043 let (act_id, w_id, mm_id) = {
11044 let mut g = graph.write().unwrap();
11045 let act = g.push(Node {
11046 op: Op::Const,
11047 inputs: vec![],
11048 shape: Shape::from_dims(&[1, 32]),
11049 dtype: DType::F32,
11050 });
11051 let w = g.push(Node {
11052 op: Op::Const,
11053 inputs: vec![],
11054 shape: Shape::from_dims(&[w_bytes.len() / 4]),
11055 dtype: DType::U32,
11056 });
11057 let mm = g.push(Node {
11058 op: Op::Fused(
11059 fuel_graph::registry::FusedOps::QMATMUL,
11060 fuel_graph::registry::FusedOpParams::QMatMul {
11061 quant_type: QuantType::Q4_0,
11062 k: 32,
11063 n: 2,
11064 },
11065 ),
11066 inputs: vec![act, w],
11067 shape: Shape::from_dims(&[1, 2]),
11068 dtype: DType::F32,
11069 });
11070 g.set_target_backend(mm, BackendId::Cpu);
11071 (act, w, mm)
11072 };
11073 let mut inputs = StorageCache::new();
11074 inputs.insert(act_id, Arc::new(RwLock::new(act_storage)));
11075 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
11076 let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
11077 let guard = result_arc.read().unwrap();
11078 assert_eq!(guard.dtype, DType::F32);
11079 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11080 panic!("expected Cpu storage");
11081 };
11082 let r: &[f32] = c.as_slice().unwrap();
11083 assert!((r[0] - 528.0).abs() < 0.5, "got {}, want 528", r[0]);
11086 assert!((r[1] - 528.0).abs() < 0.5, "got {}, want 528", r[1]);
11087 }
11088
11089 #[test]
11095 fn pipelined_realize_qmatmul_q5_0_against_reference() {
11096 use fuel_graph::QuantType;
11097 use fuel_quantized::{BlockQ5_0, GgmlType};
11098 let n = 2;
11099 let k = 64; let w_f32 = vec![1.0_f32; n * k];
11102 let blocks_per_row = k / BlockQ5_0::BLCK_SIZE;
11103 let mut w_blocks = vec![BlockQ5_0::zeros(); n * blocks_per_row];
11104 BlockQ5_0::from_float(&w_f32, &mut w_blocks);
11105 let bytes_per_block = std::mem::size_of::<BlockQ5_0>();
11107 let w_bytes: Vec<u8> = unsafe {
11108 std::slice::from_raw_parts(
11109 w_blocks.as_ptr() as *const u8,
11110 w_blocks.len() * bytes_per_block,
11111 )
11112 }
11113 .to_vec();
11114 let w_storage = fuel_memory::Storage::new(
11115 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_bytes(
11116 &w_bytes,
11117 )),
11118 DType::U32,
11119 );
11120 let act_vec: Vec<f32> = (1..=k).map(|x| x as f32).collect();
11121 let act_storage = fuel_memory::from_slice_cpu(&act_vec);
11122 let mut ref_out = vec![0.0_f32; n];
11124 fuel_quantized::matmul::<BlockQ5_0>((1, k, n), &act_vec, &w_blocks, &mut ref_out)
11125 .expect("ref matmul");
11126
11127 let graph = Arc::new(RwLock::new(Graph::new()));
11128 let (act_id, w_id, mm_id) = {
11129 let mut g = graph.write().unwrap();
11130 let act = g.push(Node {
11131 op: Op::Const,
11132 inputs: vec![],
11133 shape: Shape::from_dims(&[1, k]),
11134 dtype: DType::F32,
11135 });
11136 let w = g.push(Node {
11137 op: Op::Const,
11138 inputs: vec![],
11139 shape: Shape::from_dims(&[w_bytes.len() / 4]),
11140 dtype: DType::U32,
11141 });
11142 let mm = g.push(Node {
11143 op: Op::Fused(
11144 fuel_graph::registry::FusedOps::QMATMUL,
11145 fuel_graph::registry::FusedOpParams::QMatMul {
11146 quant_type: QuantType::Q5_0,
11147 k,
11148 n,
11149 },
11150 ),
11151 inputs: vec![act, w],
11152 shape: Shape::from_dims(&[1, n]),
11153 dtype: DType::F32,
11154 });
11155 g.set_target_backend(mm, BackendId::Cpu);
11156 (act, w, mm)
11157 };
11158 let mut inputs = StorageCache::new();
11159 inputs.insert(act_id, Arc::new(RwLock::new(act_storage)));
11160 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
11161 let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
11162 let guard = result_arc.read().unwrap();
11163 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11164 panic!("expected Cpu storage");
11165 };
11166 let r: &[f32] = c.as_slice().unwrap();
11167 assert_eq!(
11170 r,
11171 ref_out.as_slice(),
11172 "pipelined Q5_0 differs from reference: got {r:?}, want {ref_out:?}"
11173 );
11174 }
11175
11176 #[test]
11181 fn pipelined_realize_qmatmul_q6k_against_reference() {
11182 use fuel_graph::QuantType;
11183 use fuel_quantized::{BlockQ6K, GgmlType};
11184 let n = 2;
11185 let k = 256; let w_f32 = vec![1.0_f32; n * k];
11187 let blocks_per_row = k / BlockQ6K::BLCK_SIZE;
11188 let mut w_blocks = vec![BlockQ6K::zeros(); n * blocks_per_row];
11189 BlockQ6K::from_float(&w_f32, &mut w_blocks);
11190 let bytes_per_block = std::mem::size_of::<BlockQ6K>();
11191 let w_bytes: Vec<u8> = unsafe {
11192 std::slice::from_raw_parts(
11193 w_blocks.as_ptr() as *const u8,
11194 w_blocks.len() * bytes_per_block,
11195 )
11196 }
11197 .to_vec();
11198 let w_storage = fuel_memory::Storage::new(
11199 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_bytes(
11200 &w_bytes,
11201 )),
11202 DType::U32,
11203 );
11204 let act_vec: Vec<f32> = (1..=k).map(|x| x as f32 / 100.0).collect();
11205 let act_storage = fuel_memory::from_slice_cpu(&act_vec);
11206 let mut ref_out = vec![0.0_f32; n];
11207 fuel_quantized::matmul::<BlockQ6K>((1, k, n), &act_vec, &w_blocks, &mut ref_out)
11208 .expect("ref matmul");
11209
11210 let graph = Arc::new(RwLock::new(Graph::new()));
11211 let (act_id, w_id, mm_id) = {
11212 let mut g = graph.write().unwrap();
11213 let act = g.push(Node {
11214 op: Op::Const,
11215 inputs: vec![],
11216 shape: Shape::from_dims(&[1, k]),
11217 dtype: DType::F32,
11218 });
11219 let w = g.push(Node {
11220 op: Op::Const,
11221 inputs: vec![],
11222 shape: Shape::from_dims(&[w_bytes.len() / 4]),
11223 dtype: DType::U32,
11224 });
11225 let mm = g.push(Node {
11226 op: Op::Fused(
11227 fuel_graph::registry::FusedOps::QMATMUL,
11228 fuel_graph::registry::FusedOpParams::QMatMul {
11229 quant_type: QuantType::Q6K,
11230 k,
11231 n,
11232 },
11233 ),
11234 inputs: vec![act, w],
11235 shape: Shape::from_dims(&[1, n]),
11236 dtype: DType::F32,
11237 });
11238 g.set_target_backend(mm, BackendId::Cpu);
11239 (act, w, mm)
11240 };
11241 let mut inputs = StorageCache::new();
11242 inputs.insert(act_id, Arc::new(RwLock::new(act_storage)));
11243 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
11244 let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
11245 let guard = result_arc.read().unwrap();
11246 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11247 panic!("expected Cpu storage");
11248 };
11249 let r: &[f32] = c.as_slice().unwrap();
11250 assert_eq!(
11251 r,
11252 ref_out.as_slice(),
11253 "pipelined Q6K differs from reference: got {r:?}, want {ref_out:?}"
11254 );
11255 }
11256
11257 #[test]
11262 fn pipelined_realize_rms_norm_bf16() {
11263 let v: Vec<half::bf16> = [3.0_f32, 4.0]
11264 .iter()
11265 .map(|&x| half::bf16::from_f32(x))
11266 .collect();
11267 let storage = fuel_memory::from_slice_cpu(&v);
11268 let graph = Arc::new(RwLock::new(Graph::new()));
11269 let (in_id, op_id) = {
11270 let mut g = graph.write().unwrap();
11271 let in_id = g.push(Node {
11272 op: Op::Const,
11273 inputs: vec![],
11274 shape: Shape::from_dims(&[1, 2]),
11275 dtype: DType::BF16,
11276 });
11277 let op_id = g.push(Node {
11278 op: Op::Fused(
11279 fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM,
11280 fuel_graph::registry::FusedOpParams::RmsNormLastDim { eps: 0.0 },
11281 ),
11282 inputs: vec![in_id],
11283 shape: Shape::from_dims(&[1, 2]),
11284 dtype: DType::BF16,
11285 });
11286 g.set_target_backend(op_id, BackendId::Cpu);
11287 (in_id, op_id)
11288 };
11289 let mut inputs = StorageCache::new();
11290 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11291 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11292 let guard = result_arc.read().unwrap();
11293 assert_eq!(guard.dtype, DType::BF16);
11294 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11295 panic!("expected Cpu storage");
11296 };
11297 let r: &[half::bf16] = c.as_slice().unwrap();
11298 let rms = (12.5_f32).sqrt();
11299 assert!((r[0].to_f32() - 3.0 / rms).abs() < 0.05);
11301 assert!((r[1].to_f32() - 4.0 / rms).abs() < 0.05);
11302 }
11303
11304 #[test]
11309 fn pipelined_realize_matmul_bf16_identity() {
11310 let lhs_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0]
11311 .iter()
11312 .map(|&x| half::bf16::from_f32(x))
11313 .collect();
11314 let rhs_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0]
11315 .iter()
11316 .map(|&x| half::bf16::from_f32(x))
11317 .collect();
11318 let lhs = fuel_memory::from_slice_cpu(&lhs_v);
11319 let rhs = fuel_memory::from_slice_cpu(&rhs_v);
11320 let graph = Arc::new(RwLock::new(Graph::new()));
11321 let (l_id, r_id, mm_id) = {
11322 let mut g = graph.write().unwrap();
11323 let l = g.push(Node {
11324 op: Op::Const,
11325 inputs: vec![],
11326 shape: Shape::from_dims(&[2, 2]),
11327 dtype: DType::BF16,
11328 });
11329 let r = g.push(Node {
11330 op: Op::Const,
11331 inputs: vec![],
11332 shape: Shape::from_dims(&[2, 2]),
11333 dtype: DType::BF16,
11334 });
11335 let mm = g.push(Node {
11336 op: Op::MatMul,
11337 inputs: vec![l, r],
11338 shape: Shape::from_dims(&[2, 2]),
11339 dtype: DType::BF16,
11340 });
11341 g.set_target_backend(mm, BackendId::Cpu);
11342 (l, r, mm)
11343 };
11344 let mut inputs = StorageCache::new();
11345 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
11346 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
11347 let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
11348 let guard = result_arc.read().unwrap();
11349 assert_eq!(guard.dtype, DType::BF16);
11350 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11351 panic!("expected Cpu storage");
11352 };
11353 let r: &[half::bf16] = c.as_slice().unwrap();
11354 let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
11355 assert_eq!(result_f32, vec![1.0, 2.0, 3.0, 4.0]);
11356 }
11357
11358 #[test]
11361 fn pipelined_realize_sum_dim_f16() {
11362 let v: Vec<half::f16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
11363 .iter()
11364 .map(|&x| half::f16::from_f32(x))
11365 .collect();
11366 let storage = fuel_memory::from_slice_cpu(&v);
11367 let graph = Arc::new(RwLock::new(Graph::new()));
11368 let (in_id, sum_id) = {
11369 let mut g = graph.write().unwrap();
11370 let in_id = g.push(Node {
11371 op: Op::Const,
11372 inputs: vec![],
11373 shape: Shape::from_dims(&[2, 3]),
11374 dtype: DType::F16,
11375 });
11376 let sum_id = g.push(Node {
11377 op: Op::SumDim(1),
11378 inputs: vec![in_id],
11379 shape: Shape::from_dims(&[2]),
11380 dtype: DType::F16,
11381 });
11382 g.set_target_backend(sum_id, BackendId::Cpu);
11383 (in_id, sum_id)
11384 };
11385 let mut inputs = StorageCache::new();
11386 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11387 let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
11388 let guard = result_arc.read().unwrap();
11389 assert_eq!(guard.dtype, DType::F16);
11390 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11391 panic!("expected Cpu storage");
11392 };
11393 let r: &[half::f16] = c.as_slice().unwrap();
11394 let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
11395 assert_eq!(result_f32, vec![6.0, 15.0]);
11396 }
11397
11398 #[test]
11400 fn pipelined_realize_add_bf16() {
11401 let lhs_vec: Vec<half::bf16> = [1.0_f32, 2.0, 3.0]
11402 .iter()
11403 .map(|&x| half::bf16::from_f32(x))
11404 .collect();
11405 let rhs_vec: Vec<half::bf16> = [10.0_f32, 20.0, 30.0]
11406 .iter()
11407 .map(|&x| half::bf16::from_f32(x))
11408 .collect();
11409 let lhs = fuel_memory::from_slice_cpu(&lhs_vec);
11410 let rhs = fuel_memory::from_slice_cpu(&rhs_vec);
11411 let graph = Arc::new(RwLock::new(Graph::new()));
11412 let (l_id, r_id, op_id) = {
11413 let mut g = graph.write().unwrap();
11414 let l = g.push(Node {
11415 op: Op::Const,
11416 inputs: vec![],
11417 shape: Shape::from_dims(&[3]),
11418 dtype: DType::BF16,
11419 });
11420 let r = g.push(Node {
11421 op: Op::Const,
11422 inputs: vec![],
11423 shape: Shape::from_dims(&[3]),
11424 dtype: DType::BF16,
11425 });
11426 let op = g.push(Node {
11427 op: Op::Add,
11428 inputs: vec![l, r],
11429 shape: Shape::from_dims(&[3]),
11430 dtype: DType::BF16,
11431 });
11432 g.set_target_backend(op, BackendId::Cpu);
11433 (l, r, op)
11434 };
11435 let mut inputs = StorageCache::new();
11436 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
11437 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
11438 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11439 let guard = result_arc.read().unwrap();
11440 assert_eq!(guard.dtype, DType::BF16);
11441 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11442 panic!("expected Cpu storage");
11443 };
11444 let r: &[half::bf16] = c.as_slice().unwrap();
11445 let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
11446 assert_eq!(result_f32, vec![11.0, 22.0, 33.0]);
11447 }
11448
11449 #[test]
11453 fn pipelined_realize_sqr_then_sqrt_f16() {
11454 let v: Vec<half::f16> = [1.0_f32, 4.0, 9.0]
11455 .iter()
11456 .map(|&x| half::f16::from_f32(x))
11457 .collect();
11458 let storage = fuel_memory::from_slice_cpu(&v);
11459 let graph = Arc::new(RwLock::new(Graph::new()));
11460 let (in_id, sqrt_id) = {
11461 let mut g = graph.write().unwrap();
11462 let in_id = g.push(Node {
11463 op: Op::Const,
11464 inputs: vec![],
11465 shape: Shape::from_dims(&[3]),
11466 dtype: DType::F16,
11467 });
11468 let sqr_id = g.push(Node {
11469 op: Op::Sqr,
11470 inputs: vec![in_id],
11471 shape: Shape::from_dims(&[3]),
11472 dtype: DType::F16,
11473 });
11474 let sqrt_id = g.push(Node {
11475 op: Op::Sqrt,
11476 inputs: vec![sqr_id],
11477 shape: Shape::from_dims(&[3]),
11478 dtype: DType::F16,
11479 });
11480 g.set_target_backend(sqr_id, BackendId::Cpu);
11481 g.set_target_backend(sqrt_id, BackendId::Cpu);
11482 (in_id, sqrt_id)
11483 };
11484 let mut inputs = StorageCache::new();
11485 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11486 let (result_arc, _) = PipelinedExecutor::realize(graph, sqrt_id, inputs).expect("realize");
11487 let guard = result_arc.read().unwrap();
11488 assert_eq!(guard.dtype, DType::F16);
11489 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11490 panic!("expected Cpu storage");
11491 };
11492 let r: &[half::f16] = c.as_slice().unwrap();
11493 let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
11494 for (got, want) in result_f32.iter().zip(&[1.0_f32, 4.0, 9.0]) {
11496 assert!((got - want).abs() < 0.05, "got {got}, want {want}");
11497 }
11498 }
11499
11500 #[test]
11503 fn pipelined_realize_sum_dim_f64() {
11504 let storage = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
11505 let graph = Arc::new(RwLock::new(Graph::new()));
11506 let (in_id, sum_id) = {
11507 let mut g = graph.write().unwrap();
11508 let in_id = g.push(Node {
11509 op: Op::Const,
11510 inputs: vec![],
11511 shape: Shape::from_dims(&[2, 3]),
11512 dtype: DType::F64,
11513 });
11514 let sum_id = g.push(Node {
11515 op: Op::SumDim(1),
11516 inputs: vec![in_id],
11517 shape: Shape::from_dims(&[2]),
11518 dtype: DType::F64,
11519 });
11520 g.set_target_backend(sum_id, BackendId::Cpu);
11521 (in_id, sum_id)
11522 };
11523 let mut inputs = StorageCache::new();
11524 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11525 let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
11526 let guard = result_arc.read().unwrap();
11527 assert_eq!(guard.dtype, DType::F64);
11528 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11529 panic!("expected Cpu storage");
11530 };
11531 assert_eq!(c.as_slice::<f64>().unwrap(), &[6.0_f64, 15.0]);
11532 }
11533
11534 #[test]
11536 fn pipelined_realize_matmul_2x3_times_3x2_f64() {
11537 let lhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
11538 let rhs = fuel_memory::from_slice_cpu(&[7.0_f64, 8.0, 9.0, 10.0, 11.0, 12.0]);
11539 let graph = Arc::new(RwLock::new(Graph::new()));
11540 let (l_id, r_id, mm_id) = {
11541 let mut g = graph.write().unwrap();
11542 let l = g.push(Node {
11543 op: Op::Const,
11544 inputs: vec![],
11545 shape: Shape::from_dims(&[2, 3]),
11546 dtype: DType::F64,
11547 });
11548 let r = g.push(Node {
11549 op: Op::Const,
11550 inputs: vec![],
11551 shape: Shape::from_dims(&[3, 2]),
11552 dtype: DType::F64,
11553 });
11554 let mm = g.push(Node {
11555 op: Op::MatMul,
11556 inputs: vec![l, r],
11557 shape: Shape::from_dims(&[2, 2]),
11558 dtype: DType::F64,
11559 });
11560 g.set_target_backend(mm, BackendId::Cpu);
11561 (l, r, mm)
11562 };
11563 let mut inputs = StorageCache::new();
11564 inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
11565 inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
11566 let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
11567 let guard = result_arc.read().unwrap();
11568 assert_eq!(guard.dtype, DType::F64);
11569 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11570 panic!("expected Cpu storage");
11571 };
11572 assert_eq!(
11573 c.as_slice::<f64>().unwrap(),
11574 &[58.0_f64, 64.0, 139.0, 154.0]
11575 );
11576 }
11577
11578 #[test]
11582 fn pipelined_realize_sqr_then_sqrt_f64() {
11583 let storage = fuel_memory::from_slice_cpu(&[1.0_f64, 4.0, 9.0]);
11584 let graph = Arc::new(RwLock::new(Graph::new()));
11585 let (in_id, sqrt_id) = {
11586 let mut g = graph.write().unwrap();
11587 let in_id = g.push(Node {
11588 op: Op::Const,
11589 inputs: vec![],
11590 shape: Shape::from_dims(&[3]),
11591 dtype: DType::F64,
11592 });
11593 let sqr_id = g.push(Node {
11594 op: Op::Sqr,
11595 inputs: vec![in_id],
11596 shape: Shape::from_dims(&[3]),
11597 dtype: DType::F64,
11598 });
11599 let sqrt_id = g.push(Node {
11600 op: Op::Sqrt,
11601 inputs: vec![sqr_id],
11602 shape: Shape::from_dims(&[3]),
11603 dtype: DType::F64,
11604 });
11605 g.set_target_backend(sqr_id, BackendId::Cpu);
11606 g.set_target_backend(sqrt_id, BackendId::Cpu);
11607 (in_id, sqrt_id)
11608 };
11609 let mut inputs = StorageCache::new();
11610 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11611 let (result_arc, _) = PipelinedExecutor::realize(graph, sqrt_id, inputs).expect("realize");
11612 let guard = result_arc.read().unwrap();
11613 assert_eq!(guard.dtype, DType::F64);
11614 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11615 panic!("expected Cpu storage");
11616 };
11617 assert_eq!(c.as_slice::<f64>().unwrap(), &[1.0_f64, 4.0, 9.0]);
11618 }
11619
11620 #[test]
11622 fn pipelined_realize_argmax_dim() {
11623 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 2.0, 9.0, 0.0, 4.0]);
11626 let graph = Arc::new(RwLock::new(Graph::new()));
11627 let (in_id, op_id) = {
11628 let mut g = graph.write().unwrap();
11629 let in_id = g.push(Node {
11630 op: Op::Const,
11631 inputs: vec![],
11632 shape: Shape::from_dims(&[2, 3]),
11633 dtype: DType::F32,
11634 });
11635 let op_id = g.push(Node {
11636 op: Op::ArgMaxDim(1),
11637 inputs: vec![in_id],
11638 shape: Shape::from_dims(&[2]),
11639 dtype: DType::U32,
11640 });
11641 g.set_target_backend(op_id, BackendId::Cpu);
11642 (in_id, op_id)
11643 };
11644 let mut inputs = StorageCache::new();
11645 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11646 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11647 let guard = result_arc.read().unwrap();
11648 assert_eq!(guard.dtype, DType::U32);
11649 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11650 panic!("expected Cpu storage");
11651 };
11652 assert_eq!(c.as_slice::<u32>().unwrap(), &[1u32, 0]);
11653 }
11654
11655 #[test]
11658 fn pipelined_realize_index_add_simple() {
11659 let base = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
11660 let indices = fuel_memory::from_slice_cpu(&[0u32, 0]);
11661 let src = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
11662 let graph = Arc::new(RwLock::new(Graph::new()));
11663 let (b_id, i_id, s_id, op_id) = {
11664 let mut g = graph.write().unwrap();
11665 let b = g.push(Node {
11666 op: Op::Const,
11667 inputs: vec![],
11668 shape: Shape::from_dims(&[3]),
11669 dtype: DType::F32,
11670 });
11671 let i = g.push(Node {
11672 op: Op::Const,
11673 inputs: vec![],
11674 shape: Shape::from_dims(&[2]),
11675 dtype: DType::U32,
11676 });
11677 let s = g.push(Node {
11678 op: Op::Const,
11679 inputs: vec![],
11680 shape: Shape::from_dims(&[2]),
11681 dtype: DType::F32,
11682 });
11683 let op = g.push(Node {
11684 op: Op::IndexAdd { dim: 0 },
11685 inputs: vec![b, i, s],
11686 shape: Shape::from_dims(&[3]),
11687 dtype: DType::F32,
11688 });
11689 g.set_target_backend(op, BackendId::Cpu);
11690 (b, i, s, op)
11691 };
11692 let mut inputs = StorageCache::new();
11693 inputs.insert(b_id, Arc::new(RwLock::new(base)));
11694 inputs.insert(i_id, Arc::new(RwLock::new(indices)));
11695 inputs.insert(s_id, Arc::new(RwLock::new(src)));
11696 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11697 let guard = result_arc.read().unwrap();
11698 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11699 panic!("expected Cpu storage");
11700 };
11701 assert_eq!(c.as_slice::<f32>().unwrap(), &[13.0, 20.0, 30.0]);
11703 }
11704
11705 #[test]
11708 fn pipelined_realize_scatter_add_outer_dim() {
11709 let base = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
11713 let indices = fuel_memory::from_slice_cpu(&[0u32, 1, 2, 0]);
11714 let src = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
11715 let graph = Arc::new(RwLock::new(Graph::new()));
11716 let (b_id, i_id, s_id, op_id) = {
11717 let mut g = graph.write().unwrap();
11718 let b = g.push(Node {
11719 op: Op::Const,
11720 inputs: vec![],
11721 shape: Shape::from_dims(&[3, 2]),
11722 dtype: DType::F32,
11723 });
11724 let i = g.push(Node {
11725 op: Op::Const,
11726 inputs: vec![],
11727 shape: Shape::from_dims(&[2, 2]),
11728 dtype: DType::U32,
11729 });
11730 let s = g.push(Node {
11731 op: Op::Const,
11732 inputs: vec![],
11733 shape: Shape::from_dims(&[2, 2]),
11734 dtype: DType::F32,
11735 });
11736 let op = g.push(Node {
11737 op: Op::ScatterAdd { dim: 0 },
11738 inputs: vec![b, i, s],
11739 shape: Shape::from_dims(&[3, 2]),
11740 dtype: DType::F32,
11741 });
11742 g.set_target_backend(op, BackendId::Cpu);
11743 (b, i, s, op)
11744 };
11745 let mut inputs = StorageCache::new();
11746 inputs.insert(b_id, Arc::new(RwLock::new(base)));
11747 inputs.insert(i_id, Arc::new(RwLock::new(indices)));
11748 inputs.insert(s_id, Arc::new(RwLock::new(src)));
11749 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11750 let guard = result_arc.read().unwrap();
11751 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11752 panic!("expected Cpu storage");
11753 };
11754 assert_eq!(
11755 c.as_slice::<f32>().unwrap(),
11756 &[1.0, 4.0, 0.0, 2.0, 3.0, 0.0]
11757 );
11758 }
11759
11760 #[test]
11764 fn pipelined_realize_rope_pi_over_two() {
11765 let x = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
11768 let cos = fuel_memory::from_slice_cpu(&[0.0_f32, 0.0, 0.0, 0.0]);
11769 let sin = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
11770 let graph = Arc::new(RwLock::new(Graph::new()));
11771 let (x_id, cos_id, sin_id, r_id) = {
11772 let mut g = graph.write().unwrap();
11773 let x = g.push(Node {
11774 op: Op::Const,
11775 inputs: vec![],
11776 shape: Shape::from_dims(&[1, 1, 4]),
11777 dtype: DType::F32,
11778 });
11779 let c = g.push(Node {
11780 op: Op::Const,
11781 inputs: vec![],
11782 shape: Shape::from_dims(&[1, 4]),
11783 dtype: DType::F32,
11784 });
11785 let s = g.push(Node {
11786 op: Op::Const,
11787 inputs: vec![],
11788 shape: Shape::from_dims(&[1, 4]),
11789 dtype: DType::F32,
11790 });
11791 let r = g.push(Node {
11792 op: Op::Fused(
11793 fuel_graph::registry::FusedOps::ROPE,
11794 fuel_graph::registry::FusedOpParams::Rope,
11795 ),
11796 inputs: vec![x, c, s],
11797 shape: Shape::from_dims(&[1, 1, 4]),
11798 dtype: DType::F32,
11799 });
11800 g.set_target_backend(r, BackendId::Cpu);
11801 (x, c, s, r)
11802 };
11803 let mut inputs = StorageCache::new();
11804 inputs.insert(x_id, Arc::new(RwLock::new(x)));
11805 inputs.insert(cos_id, Arc::new(RwLock::new(cos)));
11806 inputs.insert(sin_id, Arc::new(RwLock::new(sin)));
11807 let (result_arc, _) = PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
11808 let guard = result_arc.read().unwrap();
11809 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11810 panic!("expected Cpu storage");
11811 };
11812 assert_eq!(c.as_slice::<f32>().unwrap(), &[-3.0, -4.0, 1.0, 2.0]);
11813 }
11814
11815 #[test]
11818 fn pipelined_realize_gather_inner_dim() {
11819 let source =
11820 fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]);
11821 let indices = fuel_memory::from_slice_cpu(&[0u32, 2, 1, 3, 0, 0]);
11822 let graph = Arc::new(RwLock::new(Graph::new()));
11823 let (src_id, idx_id, g_id) = {
11824 let mut g = graph.write().unwrap();
11825 let s = g.push(Node {
11826 op: Op::Const,
11827 inputs: vec![],
11828 shape: Shape::from_dims(&[2, 4]),
11829 dtype: DType::F32,
11830 });
11831 let i = g.push(Node {
11832 op: Op::Const,
11833 inputs: vec![],
11834 shape: Shape::from_dims(&[2, 3]),
11835 dtype: DType::U32,
11836 });
11837 let g_id = g.push(Node {
11838 op: Op::Gather { dim: 1 },
11839 inputs: vec![s, i],
11840 shape: Shape::from_dims(&[2, 3]),
11841 dtype: DType::F32,
11842 });
11843 g.set_target_backend(g_id, BackendId::Cpu);
11844 (s, i, g_id)
11845 };
11846 let mut inputs = StorageCache::new();
11847 inputs.insert(src_id, Arc::new(RwLock::new(source)));
11848 inputs.insert(idx_id, Arc::new(RwLock::new(indices)));
11849 let (result_arc, _) = PipelinedExecutor::realize(graph, g_id, inputs).expect("realize");
11850 let guard = result_arc.read().unwrap();
11851 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11852 panic!("expected Cpu storage");
11853 };
11854 assert_eq!(
11855 c.as_slice::<f32>().unwrap(),
11856 &[10.0, 30.0, 20.0, 80.0, 50.0, 50.0]
11857 );
11858 }
11859
11860 #[test]
11864 fn pipelined_realize_index_select_embedding_lookup() {
11865 let table = fuel_memory::from_slice_cpu(&[
11866 10.0_f32, 11.0, 12.0, 20.0, 21.0, 22.0, 30.0, 31.0, 32.0, 40.0, 41.0, 42.0, ]);
11871 let indices = fuel_memory::from_slice_cpu(&[2u32, 0, 2]);
11872 let graph = Arc::new(RwLock::new(Graph::new()));
11873 let (table_id, idx_id, sel_id) = {
11874 let mut g = graph.write().unwrap();
11875 let t = g.push(Node {
11876 op: Op::Const,
11877 inputs: vec![],
11878 shape: Shape::from_dims(&[4, 3]),
11879 dtype: DType::F32,
11880 });
11881 let i = g.push(Node {
11882 op: Op::Const,
11883 inputs: vec![],
11884 shape: Shape::from_dims(&[3]),
11885 dtype: DType::U32,
11886 });
11887 let s = g.push(Node {
11888 op: Op::IndexSelect { dim: 0 },
11889 inputs: vec![t, i],
11890 shape: Shape::from_dims(&[3, 3]),
11891 dtype: DType::F32,
11892 });
11893 g.set_target_backend(s, BackendId::Cpu);
11894 (t, i, s)
11895 };
11896 let mut inputs = StorageCache::new();
11897 inputs.insert(table_id, Arc::new(RwLock::new(table)));
11898 inputs.insert(idx_id, Arc::new(RwLock::new(indices)));
11899 let (result_arc, result_layout) =
11900 PipelinedExecutor::realize(graph, sel_id, inputs).expect("realize");
11901 assert_eq!(result_layout.shape().dims(), &[3, 3]);
11902 let guard = result_arc.read().unwrap();
11903 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11904 panic!("expected Cpu storage");
11905 };
11906 assert_eq!(
11907 c.as_slice::<f32>().unwrap(),
11908 &[30.0, 31.0, 32.0, 10.0, 11.0, 12.0, 30.0, 31.0, 32.0]
11909 );
11910 }
11911
11912 #[test]
11915 fn pipelined_realize_rms_norm_last_dim() {
11916 let storage = fuel_memory::from_slice_cpu(&[
11917 3.0_f32, 4.0, 6.0, 8.0, ]);
11920 let graph = Arc::new(RwLock::new(Graph::new()));
11921 let (in_id, op_id) = {
11922 let mut g = graph.write().unwrap();
11923 let in_id = g.push(Node {
11924 op: Op::Const,
11925 inputs: vec![],
11926 shape: Shape::from_dims(&[2, 2]),
11927 dtype: DType::F32,
11928 });
11929 let op_id = g.push(Node {
11930 op: Op::Fused(
11931 fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM,
11932 fuel_graph::registry::FusedOpParams::RmsNormLastDim { eps: 0.0 },
11933 ),
11934 inputs: vec![in_id],
11935 shape: Shape::from_dims(&[2, 2]),
11936 dtype: DType::F32,
11937 });
11938 g.set_target_backend(op_id, BackendId::Cpu);
11939 (in_id, op_id)
11940 };
11941 let mut inputs = StorageCache::new();
11942 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11943 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11944 let guard = result_arc.read().unwrap();
11945 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11946 panic!("expected Cpu storage");
11947 };
11948 let result: &[f32] = c.as_slice().unwrap();
11949 let rms0 = (12.5_f32).sqrt();
11951 assert!((result[0] - 3.0 / rms0).abs() < 1e-6);
11952 assert!((result[1] - 4.0 / rms0).abs() < 1e-6);
11953 let rms1 = (50.0_f32).sqrt();
11955 assert!((result[2] - 6.0 / rms1).abs() < 1e-6);
11956 assert!((result[3] - 8.0 / rms1).abs() < 1e-6);
11957 }
11958
11959 #[test]
11962 fn pipelined_realize_layer_norm_last_dim() {
11963 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 10.0, 20.0, 30.0]);
11964 let graph = Arc::new(RwLock::new(Graph::new()));
11965 let (in_id, op_id) = {
11966 let mut g = graph.write().unwrap();
11967 let in_id = g.push(Node {
11968 op: Op::Const,
11969 inputs: vec![],
11970 shape: Shape::from_dims(&[2, 3]),
11971 dtype: DType::F32,
11972 });
11973 let op_id = g.push(Node {
11974 op: Op::Fused(
11975 fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM,
11976 fuel_graph::registry::FusedOpParams::LayerNormLastDim { eps: 0.0 },
11977 ),
11978 inputs: vec![in_id],
11979 shape: Shape::from_dims(&[2, 3]),
11980 dtype: DType::F32,
11981 });
11982 g.set_target_backend(op_id, BackendId::Cpu);
11983 (in_id, op_id)
11984 };
11985 let mut inputs = StorageCache::new();
11986 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
11987 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
11988 let guard = result_arc.read().unwrap();
11989 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
11990 panic!("expected Cpu storage");
11991 };
11992 let result: &[f32] = c.as_slice().unwrap();
11993 for row in 0..2 {
11995 let off = row * 3;
11996 let sum: f32 = result[off..off + 3].iter().sum();
11997 let mean = sum / 3.0;
11998 assert!(mean.abs() < 1e-6, "row {row} mean should be 0, got {mean}");
11999 let var: f32 = result[off..off + 3]
12000 .iter()
12001 .map(|v| (v - mean).powi(2))
12002 .sum::<f32>()
12003 / 3.0;
12004 assert!(
12005 (var - 1.0).abs() < 1e-6,
12006 "row {row} var should be 1, got {var}"
12007 );
12008 }
12009 }
12010
12011 #[test]
12014 fn pipelined_realize_softmax_last_dim() {
12015 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 100.0]);
12018 let graph = Arc::new(RwLock::new(Graph::new()));
12019 let (in_id, op_id) = {
12020 let mut g = graph.write().unwrap();
12021 let in_id = g.push(Node {
12022 op: Op::Const,
12023 inputs: vec![],
12024 shape: Shape::from_dims(&[2, 4]),
12025 dtype: DType::F32,
12026 });
12027 let op_id = g.push(Node {
12028 op: Op::Fused(
12029 fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM,
12030 fuel_graph::registry::FusedOpParams::SoftmaxLastDim,
12031 ),
12032 inputs: vec![in_id],
12033 shape: Shape::from_dims(&[2, 4]),
12034 dtype: DType::F32,
12035 });
12036 g.set_target_backend(op_id, BackendId::Cpu);
12037 (in_id, op_id)
12038 };
12039 let mut inputs = StorageCache::new();
12040 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
12041 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12042 let guard = result_arc.read().unwrap();
12043 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12044 panic!("expected Cpu storage");
12045 };
12046 let result: &[f32] = c.as_slice().unwrap();
12047
12048 for v in &result[..4] {
12050 assert!((v - 0.25).abs() < 1e-7);
12051 }
12052 for v in &result[4..7] {
12055 assert!(
12056 *v < 1e-30,
12057 "row-1 leading positions should be near 0, got {v}"
12058 );
12059 }
12060 assert!(
12061 result[7] > 0.999,
12062 "row-1 last position should dominate, got {}",
12063 result[7]
12064 );
12065 let row0_sum: f32 = result[..4].iter().sum();
12067 let row1_sum: f32 = result[4..].iter().sum();
12068 assert!((row0_sum - 1.0).abs() < 1e-6);
12069 assert!((row1_sum - 1.0).abs() < 1e-6);
12070 }
12071
12072 #[test]
12074 fn pipelined_realize_concat_inner_dim() {
12075 let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
12076 let b = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
12077 let graph = Arc::new(RwLock::new(Graph::new()));
12078 let (a_id, b_id, c_id) = {
12079 let mut g = graph.write().unwrap();
12080 let a = g.push(Node {
12081 op: Op::Const,
12082 inputs: vec![],
12083 shape: Shape::from_dims(&[2, 3]),
12084 dtype: DType::F32,
12085 });
12086 let b = g.push(Node {
12087 op: Op::Const,
12088 inputs: vec![],
12089 shape: Shape::from_dims(&[2, 3]),
12090 dtype: DType::F32,
12091 });
12092 let c = g.push(Node {
12093 op: Op::Concat { dim: 1 },
12094 inputs: vec![a, b],
12095 shape: Shape::from_dims(&[2, 6]),
12096 dtype: DType::F32,
12097 });
12098 g.set_target_backend(c, BackendId::Cpu);
12099 (a, b, c)
12100 };
12101 let mut inputs = StorageCache::new();
12102 inputs.insert(a_id, Arc::new(RwLock::new(a)));
12103 inputs.insert(b_id, Arc::new(RwLock::new(b)));
12104 let (result_arc, result_layout) =
12105 PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
12106 assert_eq!(result_layout.shape().dims(), &[2, 6]);
12107 let guard = result_arc.read().unwrap();
12108 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12109 panic!("expected Cpu storage");
12110 };
12111 assert_eq!(
12112 c.as_slice::<f32>().unwrap(),
12113 &[
12114 1.0, 2.0, 3.0, 7.0, 8.0, 9.0, 4.0, 5.0, 6.0, 10.0, 11.0, 12.0
12115 ]
12116 );
12117 }
12118
12119 #[test]
12122 fn pipelined_realize_concat_three_inputs_outer() {
12123 let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
12124 let b = fuel_memory::from_slice_cpu(&[3.0_f32, 4.0]);
12125 let c = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0]);
12126 let graph = Arc::new(RwLock::new(Graph::new()));
12127 let (a_id, b_id, c_id, cat_id) = {
12128 let mut g = graph.write().unwrap();
12129 let a = g.push(Node {
12130 op: Op::Const,
12131 inputs: vec![],
12132 shape: Shape::from_dims(&[2]),
12133 dtype: DType::F32,
12134 });
12135 let b = g.push(Node {
12136 op: Op::Const,
12137 inputs: vec![],
12138 shape: Shape::from_dims(&[2]),
12139 dtype: DType::F32,
12140 });
12141 let c = g.push(Node {
12142 op: Op::Const,
12143 inputs: vec![],
12144 shape: Shape::from_dims(&[2]),
12145 dtype: DType::F32,
12146 });
12147 let cat = g.push(Node {
12148 op: Op::Concat { dim: 0 },
12149 inputs: vec![a, b, c],
12150 shape: Shape::from_dims(&[6]),
12151 dtype: DType::F32,
12152 });
12153 g.set_target_backend(cat, BackendId::Cpu);
12154 (a, b, c, cat)
12155 };
12156 let mut inputs = StorageCache::new();
12157 inputs.insert(a_id, Arc::new(RwLock::new(a)));
12158 inputs.insert(b_id, Arc::new(RwLock::new(b)));
12159 inputs.insert(c_id, Arc::new(RwLock::new(c)));
12160 let (result_arc, _) = PipelinedExecutor::realize(graph, cat_id, inputs).expect("realize");
12161 let guard = result_arc.read().unwrap();
12162 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12163 panic!("expected Cpu storage");
12164 };
12165 assert_eq!(
12166 c.as_slice::<f32>().unwrap(),
12167 &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
12168 );
12169 }
12170
12171 #[test]
12174 fn pipelined_realize_add_scalar() {
12175 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
12176 let graph = Arc::new(RwLock::new(Graph::new()));
12177 let (in_id, op_id) = {
12178 let mut g = graph.write().unwrap();
12179 let in_id = g.push(Node {
12180 op: Op::Const,
12181 inputs: vec![],
12182 shape: Shape::from_dims(&[3]),
12183 dtype: DType::F32,
12184 });
12185 let op_id = g.push(Node {
12186 op: Op::AddScalar(10.0),
12187 inputs: vec![in_id],
12188 shape: Shape::from_dims(&[3]),
12189 dtype: DType::F32,
12190 });
12191 g.set_target_backend(op_id, BackendId::Cpu);
12192 (in_id, op_id)
12193 };
12194 let mut inputs = StorageCache::new();
12195 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
12196 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12197 let guard = result_arc.read().unwrap();
12198 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12199 panic!("expected Cpu storage");
12200 };
12201 assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 12.0, 13.0]);
12202 }
12203
12204 #[test]
12206 fn pipelined_realize_clamp() {
12207 let storage = fuel_memory::from_slice_cpu(&[-5.0_f32, 0.5, 100.0]);
12208 let graph = Arc::new(RwLock::new(Graph::new()));
12209 let (in_id, op_id) = {
12210 let mut g = graph.write().unwrap();
12211 let in_id = g.push(Node {
12212 op: Op::Const,
12213 inputs: vec![],
12214 shape: Shape::from_dims(&[3]),
12215 dtype: DType::F32,
12216 });
12217 let op_id = g.push(Node {
12218 op: Op::Clamp {
12219 min: -2.0,
12220 max: 2.0,
12221 },
12222 inputs: vec![in_id],
12223 shape: Shape::from_dims(&[3]),
12224 dtype: DType::F32,
12225 });
12226 g.set_target_backend(op_id, BackendId::Cpu);
12227 (in_id, op_id)
12228 };
12229 let mut inputs = StorageCache::new();
12230 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
12231 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12232 let guard = result_arc.read().unwrap();
12233 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12234 panic!("expected Cpu storage");
12235 };
12236 assert_eq!(c.as_slice::<f32>().unwrap(), &[-2.0, 0.5, 2.0]);
12237 }
12238
12239 #[test]
12241 fn pipelined_realize_maximum_elementwise() {
12242 let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, -3.0]);
12243 let rhs_storage = fuel_memory::from_slice_cpu(&[2.0_f32, 1.0, -1.0]);
12244 let graph = Arc::new(RwLock::new(Graph::new()));
12245 let (lhs_id, rhs_id, op_id) = {
12246 let mut g = graph.write().unwrap();
12247 let lhs = g.push(Node {
12248 op: Op::Const,
12249 inputs: vec![],
12250 shape: Shape::from_dims(&[3]),
12251 dtype: DType::F32,
12252 });
12253 let rhs = g.push(Node {
12254 op: Op::Const,
12255 inputs: vec![],
12256 shape: Shape::from_dims(&[3]),
12257 dtype: DType::F32,
12258 });
12259 let op = g.push(Node {
12260 op: Op::Maximum,
12261 inputs: vec![lhs, rhs],
12262 shape: Shape::from_dims(&[3]),
12263 dtype: DType::F32,
12264 });
12265 g.set_target_backend(op, BackendId::Cpu);
12266 (lhs, rhs, op)
12267 };
12268 let mut inputs = StorageCache::new();
12269 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
12270 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
12271 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12272 let guard = result_arc.read().unwrap();
12273 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12274 panic!("expected Cpu storage");
12275 };
12276 assert_eq!(c.as_slice::<f32>().unwrap(), &[2.0, 5.0, -1.0]);
12277 }
12278
12279 #[test]
12282 fn pipelined_realize_conv2d_2x2_sum_kernel() {
12283 let x_storage =
12287 fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
12288 let w_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
12289
12290 let graph = Arc::new(RwLock::new(Graph::new()));
12291 let (x_id, w_id, c_id) = {
12292 let mut g = graph.write().unwrap();
12293 let x = g.push(Node {
12294 op: Op::Const,
12295 inputs: vec![],
12296 shape: Shape::from_dims(&[1, 1, 3, 3]),
12297 dtype: DType::F32,
12298 });
12299 let w = g.push(Node {
12300 op: Op::Const,
12301 inputs: vec![],
12302 shape: Shape::from_dims(&[1, 1, 2, 2]),
12303 dtype: DType::F32,
12304 });
12305 let c = g.push(Node {
12306 op: Op::Fused(
12307 fuel_graph::registry::FusedOps::CONV2D,
12308 fuel_graph::registry::FusedOpParams::Conv2D {
12309 stride: (1, 1),
12310 padding: (0, 0),
12311 groups: 1,
12312 },
12313 ),
12314 inputs: vec![x, w],
12315 shape: Shape::from_dims(&[1, 1, 2, 2]),
12316 dtype: DType::F32,
12317 });
12318 g.set_target_backend(c, BackendId::Cpu);
12319 (x, w, c)
12320 };
12321 let mut inputs = StorageCache::new();
12322 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
12323 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
12324
12325 let (result_arc, result_layout) =
12326 PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
12327 assert_eq!(result_layout.shape().dims(), &[1, 1, 2, 2]);
12328
12329 let guard = result_arc.read().unwrap();
12330 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12331 panic!("expected Cpu storage");
12332 };
12333 assert_eq!(c.as_slice::<f32>().unwrap(), &[12.0, 16.0, 24.0, 28.0]);
12334 }
12335
12336 #[test]
12338 fn pipelined_realize_conv2d_with_bias() {
12339 let x_storage =
12340 fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
12341 let w_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
12342 let bias_storage = fuel_memory::from_slice_cpu(&[100.0_f32]);
12343
12344 let graph = Arc::new(RwLock::new(Graph::new()));
12345 let (x_id, w_id, b_id, c_id) = {
12346 let mut g = graph.write().unwrap();
12347 let x = g.push(Node {
12348 op: Op::Const,
12349 inputs: vec![],
12350 shape: Shape::from_dims(&[1, 1, 3, 3]),
12351 dtype: DType::F32,
12352 });
12353 let w = g.push(Node {
12354 op: Op::Const,
12355 inputs: vec![],
12356 shape: Shape::from_dims(&[1, 1, 2, 2]),
12357 dtype: DType::F32,
12358 });
12359 let b = g.push(Node {
12360 op: Op::Const,
12361 inputs: vec![],
12362 shape: Shape::from_dims(&[1]),
12363 dtype: DType::F32,
12364 });
12365 let c = g.push(Node {
12366 op: Op::Fused(
12367 fuel_graph::registry::FusedOps::CONV2D,
12368 fuel_graph::registry::FusedOpParams::Conv2D {
12369 stride: (1, 1),
12370 padding: (0, 0),
12371 groups: 1,
12372 },
12373 ),
12374 inputs: vec![x, w, b],
12375 shape: Shape::from_dims(&[1, 1, 2, 2]),
12376 dtype: DType::F32,
12377 });
12378 g.set_target_backend(c, BackendId::Cpu);
12379 (x, w, b, c)
12380 };
12381 let mut inputs = StorageCache::new();
12382 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
12383 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
12384 inputs.insert(b_id, Arc::new(RwLock::new(bias_storage)));
12385
12386 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
12387 let guard = result_arc.read().unwrap();
12388 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12389 panic!("expected Cpu storage");
12390 };
12391 assert_eq!(c.as_slice::<f32>().unwrap(), &[112.0, 116.0, 124.0, 128.0]);
12392 }
12393
12394 #[test]
12396 fn pipelined_realize_conv2d_f64() {
12397 let x_storage =
12398 fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
12399 let w_storage = fuel_memory::from_slice_cpu(&[1.0_f64, 1.0, 1.0, 1.0]);
12400
12401 let graph = Arc::new(RwLock::new(Graph::new()));
12402 let (x_id, w_id, c_id) = {
12403 let mut g = graph.write().unwrap();
12404 let x = g.push(Node {
12405 op: Op::Const,
12406 inputs: vec![],
12407 shape: Shape::from_dims(&[1, 1, 3, 3]),
12408 dtype: DType::F64,
12409 });
12410 let w = g.push(Node {
12411 op: Op::Const,
12412 inputs: vec![],
12413 shape: Shape::from_dims(&[1, 1, 2, 2]),
12414 dtype: DType::F64,
12415 });
12416 let c = g.push(Node {
12417 op: Op::Fused(
12418 fuel_graph::registry::FusedOps::CONV2D,
12419 fuel_graph::registry::FusedOpParams::Conv2D {
12420 stride: (1, 1),
12421 padding: (0, 0),
12422 groups: 1,
12423 },
12424 ),
12425 inputs: vec![x, w],
12426 shape: Shape::from_dims(&[1, 1, 2, 2]),
12427 dtype: DType::F64,
12428 });
12429 g.set_target_backend(c, BackendId::Cpu);
12430 (x, w, c)
12431 };
12432 let mut inputs = StorageCache::new();
12433 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
12434 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
12435 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
12436 let guard = result_arc.read().unwrap();
12437 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12438 panic!("expected Cpu storage");
12439 };
12440 assert_eq!(c.as_slice::<f64>().unwrap(), &[12.0, 16.0, 24.0, 28.0]);
12441 }
12442
12443 #[test]
12445 fn pipelined_realize_conv2d_bf16() {
12446 let x_data: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
12447 .iter()
12448 .map(|v| half::bf16::from_f32(*v))
12449 .collect();
12450 let w_data: Vec<half::bf16> = [1.0_f32, 1.0, 1.0, 1.0]
12451 .iter()
12452 .map(|v| half::bf16::from_f32(*v))
12453 .collect();
12454 let x_storage = fuel_memory::from_slice_cpu(&x_data);
12455 let w_storage = fuel_memory::from_slice_cpu(&w_data);
12456
12457 let graph = Arc::new(RwLock::new(Graph::new()));
12458 let (x_id, w_id, c_id) = {
12459 let mut g = graph.write().unwrap();
12460 let x = g.push(Node {
12461 op: Op::Const,
12462 inputs: vec![],
12463 shape: Shape::from_dims(&[1, 1, 3, 3]),
12464 dtype: DType::BF16,
12465 });
12466 let w = g.push(Node {
12467 op: Op::Const,
12468 inputs: vec![],
12469 shape: Shape::from_dims(&[1, 1, 2, 2]),
12470 dtype: DType::BF16,
12471 });
12472 let c = g.push(Node {
12473 op: Op::Fused(
12474 fuel_graph::registry::FusedOps::CONV2D,
12475 fuel_graph::registry::FusedOpParams::Conv2D {
12476 stride: (1, 1),
12477 padding: (0, 0),
12478 groups: 1,
12479 },
12480 ),
12481 inputs: vec![x, w],
12482 shape: Shape::from_dims(&[1, 1, 2, 2]),
12483 dtype: DType::BF16,
12484 });
12485 g.set_target_backend(c, BackendId::Cpu);
12486 (x, w, c)
12487 };
12488 let mut inputs = StorageCache::new();
12489 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
12490 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
12491 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
12492 let guard = result_arc.read().unwrap();
12493 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12494 panic!("expected Cpu storage");
12495 };
12496 let got: Vec<f32> = c
12497 .as_slice::<half::bf16>()
12498 .unwrap()
12499 .iter()
12500 .map(|v| v.to_f32())
12501 .collect();
12502 let want = [12.0_f32, 16.0, 24.0, 28.0];
12503 for (g, w) in got.iter().zip(want.iter()) {
12504 assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
12505 }
12506 }
12507
12508 #[test]
12510 fn pipelined_realize_conv2d_f16() {
12511 let x_data: Vec<half::f16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
12512 .iter()
12513 .map(|v| half::f16::from_f32(*v))
12514 .collect();
12515 let w_data: Vec<half::f16> = [1.0_f32, 1.0, 1.0, 1.0]
12516 .iter()
12517 .map(|v| half::f16::from_f32(*v))
12518 .collect();
12519 let x_storage = fuel_memory::from_slice_cpu(&x_data);
12520 let w_storage = fuel_memory::from_slice_cpu(&w_data);
12521
12522 let graph = Arc::new(RwLock::new(Graph::new()));
12523 let (x_id, w_id, c_id) = {
12524 let mut g = graph.write().unwrap();
12525 let x = g.push(Node {
12526 op: Op::Const,
12527 inputs: vec![],
12528 shape: Shape::from_dims(&[1, 1, 3, 3]),
12529 dtype: DType::F16,
12530 });
12531 let w = g.push(Node {
12532 op: Op::Const,
12533 inputs: vec![],
12534 shape: Shape::from_dims(&[1, 1, 2, 2]),
12535 dtype: DType::F16,
12536 });
12537 let c = g.push(Node {
12538 op: Op::Fused(
12539 fuel_graph::registry::FusedOps::CONV2D,
12540 fuel_graph::registry::FusedOpParams::Conv2D {
12541 stride: (1, 1),
12542 padding: (0, 0),
12543 groups: 1,
12544 },
12545 ),
12546 inputs: vec![x, w],
12547 shape: Shape::from_dims(&[1, 1, 2, 2]),
12548 dtype: DType::F16,
12549 });
12550 g.set_target_backend(c, BackendId::Cpu);
12551 (x, w, c)
12552 };
12553 let mut inputs = StorageCache::new();
12554 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
12555 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
12556 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
12557 let guard = result_arc.read().unwrap();
12558 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12559 panic!("expected Cpu storage");
12560 };
12561 let got: Vec<f32> = c
12562 .as_slice::<half::f16>()
12563 .unwrap()
12564 .iter()
12565 .map(|v| v.to_f32())
12566 .collect();
12567 let want = [12.0_f32, 16.0, 24.0, 28.0];
12568 for (g, w) in got.iter().zip(want.iter()) {
12569 assert!((g - w).abs() < 0.05, "got {got:?} want {want:?}");
12570 }
12571 }
12572
12573 #[test]
12588 fn pipelined_realize_paged_attn_f32() {
12589 let q = fuel_memory::from_slice_cpu(&[2.0_f32, 0.0]);
12590 let k_cache = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
12591 let v_cache = fuel_memory::from_slice_cpu(&[10.0_f32, 0.0, 0.0, 10.0]);
12592 let block_table_u32 = fuel_memory::Storage::new(
12593 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
12594 0_u32,
12595 ])),
12596 DType::U32,
12597 );
12598 let context_lens_u32 = fuel_memory::Storage::new(
12599 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
12600 2_u32,
12601 ])),
12602 DType::U32,
12603 );
12604 let graph = Arc::new(RwLock::new(Graph::new()));
12605 let (q_id, kc_id, vc_id, bt_id, cl_id, op_id) = {
12606 let mut g = graph.write().unwrap();
12607 let q = g.push(Node {
12608 op: Op::Const,
12609 inputs: vec![],
12610 shape: Shape::from_dims(&[1, 1, 1, 2]),
12611 dtype: DType::F32,
12612 });
12613 let kc = g.push(Node {
12614 op: Op::Const,
12615 inputs: vec![],
12616 shape: Shape::from_dims(&[1, 2, 1, 2]),
12617 dtype: DType::F32,
12618 });
12619 let vc = g.push(Node {
12620 op: Op::Const,
12621 inputs: vec![],
12622 shape: Shape::from_dims(&[1, 2, 1, 2]),
12623 dtype: DType::F32,
12624 });
12625 let bt = g.push(Node {
12626 op: Op::Const,
12627 inputs: vec![],
12628 shape: Shape::from_dims(&[1, 1]),
12629 dtype: DType::U32,
12630 });
12631 let cl = g.push(Node {
12632 op: Op::Const,
12633 inputs: vec![],
12634 shape: Shape::from_dims(&[1]),
12635 dtype: DType::U32,
12636 });
12637 let op = g.push(Node {
12638 op: Op::Fused(
12639 fuel_graph::registry::FusedOps::PAGED_ATTN,
12640 fuel_graph::registry::FusedOpParams::PagedAttn {
12641 softmax_scale: 1.0,
12642 block_size: 2,
12643 softcap: None,
12644 },
12645 ),
12646 inputs: vec![q, kc, vc, bt, cl],
12647 shape: Shape::from_dims(&[1, 1, 1, 2]),
12648 dtype: DType::F32,
12649 });
12650 g.set_target_backend(op, BackendId::Cpu);
12651 (q, kc, vc, bt, cl, op)
12652 };
12653 let mut inputs = StorageCache::new();
12654 inputs.insert(q_id, Arc::new(RwLock::new(q)));
12655 inputs.insert(kc_id, Arc::new(RwLock::new(k_cache)));
12656 inputs.insert(vc_id, Arc::new(RwLock::new(v_cache)));
12657 inputs.insert(bt_id, Arc::new(RwLock::new(block_table_u32)));
12658 inputs.insert(cl_id, Arc::new(RwLock::new(context_lens_u32)));
12659 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12660 let guard = result_arc.read().unwrap();
12661 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12662 panic!("expected Cpu storage");
12663 };
12664 let r = c.as_slice::<f32>().unwrap();
12665 let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
12666 let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
12667 assert!(
12668 (r[0] - 10.0 * expected_p0).abs() < 1e-5,
12669 "row[0]: got {} expected {}",
12670 r[0],
12671 10.0 * expected_p0
12672 );
12673 assert!(
12674 (r[1] - 10.0 * expected_p1).abs() < 1e-5,
12675 "row[1]: got {} expected {}",
12676 r[1],
12677 10.0 * expected_p1
12678 );
12679 }
12680
12681 #[test]
12683 fn pipelined_realize_paged_attn_bf16() {
12684 let q_v: Vec<half::bf16> = [2.0_f32, 0.0]
12685 .iter()
12686 .map(|x| half::bf16::from_f32(*x))
12687 .collect();
12688 let kc_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0]
12689 .iter()
12690 .map(|x| half::bf16::from_f32(*x))
12691 .collect();
12692 let vc_v: Vec<half::bf16> = [10.0_f32, 0.0, 0.0, 10.0]
12693 .iter()
12694 .map(|x| half::bf16::from_f32(*x))
12695 .collect();
12696 let q = fuel_memory::from_slice_cpu(&q_v);
12697 let k_cache = fuel_memory::from_slice_cpu(&kc_v);
12698 let v_cache = fuel_memory::from_slice_cpu(&vc_v);
12699 let bt = fuel_memory::Storage::new(
12700 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
12701 0_u32,
12702 ])),
12703 DType::U32,
12704 );
12705 let cl = fuel_memory::Storage::new(
12706 fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
12707 2_u32,
12708 ])),
12709 DType::U32,
12710 );
12711 let graph = Arc::new(RwLock::new(Graph::new()));
12712 let (q_id, kc_id, vc_id, bt_id, cl_id, op_id) = {
12713 let mut g = graph.write().unwrap();
12714 let q = g.push(Node {
12715 op: Op::Const,
12716 inputs: vec![],
12717 shape: Shape::from_dims(&[1, 1, 1, 2]),
12718 dtype: DType::BF16,
12719 });
12720 let kc = g.push(Node {
12721 op: Op::Const,
12722 inputs: vec![],
12723 shape: Shape::from_dims(&[1, 2, 1, 2]),
12724 dtype: DType::BF16,
12725 });
12726 let vc = g.push(Node {
12727 op: Op::Const,
12728 inputs: vec![],
12729 shape: Shape::from_dims(&[1, 2, 1, 2]),
12730 dtype: DType::BF16,
12731 });
12732 let bt = g.push(Node {
12733 op: Op::Const,
12734 inputs: vec![],
12735 shape: Shape::from_dims(&[1, 1]),
12736 dtype: DType::U32,
12737 });
12738 let cl = g.push(Node {
12739 op: Op::Const,
12740 inputs: vec![],
12741 shape: Shape::from_dims(&[1]),
12742 dtype: DType::U32,
12743 });
12744 let op = g.push(Node {
12745 op: Op::Fused(
12746 fuel_graph::registry::FusedOps::PAGED_ATTN,
12747 fuel_graph::registry::FusedOpParams::PagedAttn {
12748 softmax_scale: 1.0,
12749 block_size: 2,
12750 softcap: None,
12751 },
12752 ),
12753 inputs: vec![q, kc, vc, bt, cl],
12754 shape: Shape::from_dims(&[1, 1, 1, 2]),
12755 dtype: DType::BF16,
12756 });
12757 g.set_target_backend(op, BackendId::Cpu);
12758 (q, kc, vc, bt, cl, op)
12759 };
12760 let mut inputs = StorageCache::new();
12761 inputs.insert(q_id, Arc::new(RwLock::new(q)));
12762 inputs.insert(kc_id, Arc::new(RwLock::new(k_cache)));
12763 inputs.insert(vc_id, Arc::new(RwLock::new(v_cache)));
12764 inputs.insert(bt_id, Arc::new(RwLock::new(bt)));
12765 inputs.insert(cl_id, Arc::new(RwLock::new(cl)));
12766 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12767 let guard = result_arc.read().unwrap();
12768 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12769 panic!("expected Cpu storage");
12770 };
12771 let got: Vec<f32> = c
12772 .as_slice::<half::bf16>()
12773 .unwrap()
12774 .iter()
12775 .map(|v| v.to_f32())
12776 .collect();
12777 let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
12778 let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
12779 let want = [10.0 * expected_p0, 10.0 * expected_p1];
12780 for (g, w) in got.iter().zip(want.iter()) {
12781 assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
12782 }
12783 }
12784
12785 #[test]
12792 fn pipelined_realize_flash_attn_f32() {
12793 let q = fuel_memory::from_slice_cpu(&[2.0_f32, 0.0]);
12795 let k = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
12796 let v = fuel_memory::from_slice_cpu(&[10.0_f32, 0.0, 0.0, 10.0]);
12797 let graph = Arc::new(RwLock::new(Graph::new()));
12798 let (q_id, k_id, v_id, op_id) = {
12799 let mut g = graph.write().unwrap();
12800 let q = g.push(Node {
12801 op: Op::Const,
12802 inputs: vec![],
12803 shape: Shape::from_dims(&[1, 1, 1, 2]),
12804 dtype: DType::F32,
12805 });
12806 let k = g.push(Node {
12807 op: Op::Const,
12808 inputs: vec![],
12809 shape: Shape::from_dims(&[1, 1, 2, 2]),
12810 dtype: DType::F32,
12811 });
12812 let v = g.push(Node {
12813 op: Op::Const,
12814 inputs: vec![],
12815 shape: Shape::from_dims(&[1, 1, 2, 2]),
12816 dtype: DType::F32,
12817 });
12818 let op = g.push(Node {
12819 op: Op::Fused(
12820 fuel_graph::registry::FusedOps::FLASH_ATTN,
12821 fuel_graph::registry::FusedOpParams::FlashAttn {
12822 softmax_scale: 1.0,
12823 causal: false,
12824 window_size_left: None,
12825 window_size_right: None,
12826 softcap: None,
12827 k_len: None,
12828 },
12829 ),
12830 inputs: vec![q, k, v],
12831 shape: Shape::from_dims(&[1, 1, 1, 2]),
12832 dtype: DType::F32,
12833 });
12834 g.set_target_backend(op, BackendId::Cpu);
12835 (q, k, v, op)
12836 };
12837 let mut inputs = StorageCache::new();
12838 inputs.insert(q_id, Arc::new(RwLock::new(q)));
12839 inputs.insert(k_id, Arc::new(RwLock::new(k)));
12840 inputs.insert(v_id, Arc::new(RwLock::new(v)));
12841 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12842 let guard = result_arc.read().unwrap();
12843 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12844 panic!("expected Cpu storage");
12845 };
12846 let r = c.as_slice::<f32>().unwrap();
12847 let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
12848 let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
12849 assert!(
12850 (r[0] - 10.0 * expected_p0).abs() < 1e-5,
12851 "row[0]: got {} expected {}",
12852 r[0],
12853 10.0 * expected_p0
12854 );
12855 assert!(
12856 (r[1] - 10.0 * expected_p1).abs() < 1e-5,
12857 "row[1]: got {} expected {}",
12858 r[1],
12859 10.0 * expected_p1
12860 );
12861 }
12862
12863 #[test]
12867 fn pipelined_realize_flash_attn_causal_f32() {
12868 let q = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
12877 let k = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
12878 let v = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0, 7.0, 8.0]);
12879 let graph = Arc::new(RwLock::new(Graph::new()));
12880 let (q_id, k_id, v_id, op_id) = {
12881 let mut g = graph.write().unwrap();
12882 let q = g.push(Node {
12883 op: Op::Const,
12884 inputs: vec![],
12885 shape: Shape::from_dims(&[1, 1, 2, 2]),
12886 dtype: DType::F32,
12887 });
12888 let k = g.push(Node {
12889 op: Op::Const,
12890 inputs: vec![],
12891 shape: Shape::from_dims(&[1, 1, 2, 2]),
12892 dtype: DType::F32,
12893 });
12894 let v = g.push(Node {
12895 op: Op::Const,
12896 inputs: vec![],
12897 shape: Shape::from_dims(&[1, 1, 2, 2]),
12898 dtype: DType::F32,
12899 });
12900 let op = g.push(Node {
12901 op: Op::Fused(
12902 fuel_graph::registry::FusedOps::FLASH_ATTN,
12903 fuel_graph::registry::FusedOpParams::FlashAttn {
12904 softmax_scale: 1.0,
12905 causal: true,
12906 window_size_left: None,
12907 window_size_right: None,
12908 softcap: None,
12909 k_len: None,
12910 },
12911 ),
12912 inputs: vec![q, k, v],
12913 shape: Shape::from_dims(&[1, 1, 2, 2]),
12914 dtype: DType::F32,
12915 });
12916 g.set_target_backend(op, BackendId::Cpu);
12917 (q, k, v, op)
12918 };
12919 let mut inputs = StorageCache::new();
12920 inputs.insert(q_id, Arc::new(RwLock::new(q)));
12921 inputs.insert(k_id, Arc::new(RwLock::new(k)));
12922 inputs.insert(v_id, Arc::new(RwLock::new(v)));
12923 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
12924 let guard = result_arc.read().unwrap();
12925 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
12926 panic!("expected Cpu storage");
12927 };
12928 let r = c.as_slice::<f32>().unwrap();
12929 assert!((r[0] - 5.0).abs() < 1e-5, "got {}", r[0]);
12931 assert!((r[1] - 6.0).abs() < 1e-5, "got {}", r[1]);
12932 let denom = (1.0_f32).exp() + 1.0;
12934 let expected_a = 5.0 / denom + 7.0 * (1.0_f32).exp() / denom;
12935 let expected_b = 6.0 / denom + 8.0 * (1.0_f32).exp() / denom;
12936 assert!(
12937 (r[2] - expected_a).abs() < 1e-5,
12938 "row1[0]: got {} expected {}",
12939 r[2],
12940 expected_a
12941 );
12942 assert!(
12943 (r[3] - expected_b).abs() < 1e-5,
12944 "row1[1]: got {} expected {}",
12945 r[3],
12946 expected_b
12947 );
12948 }
12949
12950 #[test]
12958 fn pipelined_realize_flash_attn_dynamic_k_len() {
12959 use fuel_ir::{DynScalar, SymId};
12960 let q = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
12963 let k = fuel_memory::from_slice_cpu(&[
12964 1.0_f32, 0.0, 0.0, 1.0, 1.0, 1.0, 100.0,
12965 100.0, ]);
12967 let v = fuel_memory::from_slice_cpu(&[
12968 5.0_f32, 6.0, 7.0, 8.0, 9.0, 10.0, 999.0,
12969 999.0, ]);
12971 let sym = SymId(0);
12972 let graph = Arc::new(RwLock::new(Graph::new()));
12973 let (q_id, k_id, v_id, op_id) = {
12974 let mut g = graph.write().unwrap();
12975 let q = g.push(Node {
12976 op: Op::Const,
12977 inputs: vec![],
12978 shape: Shape::from_dims(&[1, 1, 2, 2]),
12979 dtype: DType::F32,
12980 });
12981 let k = g.push(Node {
12982 op: Op::Const,
12983 inputs: vec![],
12984 shape: Shape::from_dims(&[1, 1, 4, 2]),
12985 dtype: DType::F32,
12986 });
12987 let v = g.push(Node {
12988 op: Op::Const,
12989 inputs: vec![],
12990 shape: Shape::from_dims(&[1, 1, 4, 2]),
12991 dtype: DType::F32,
12992 });
12993 let op = g.push(Node {
12994 op: Op::Fused(
12995 fuel_graph::registry::FusedOps::FLASH_ATTN,
12996 fuel_graph::registry::FusedOpParams::FlashAttn {
12997 softmax_scale: 1.0,
12998 causal: true,
12999 window_size_left: None,
13000 window_size_right: None,
13001 softcap: None,
13002 k_len: Some(DynScalar::Sym(sym)),
13003 },
13004 ),
13005 inputs: vec![q, k, v],
13006 shape: Shape::from_dims(&[1, 1, 2, 2]),
13007 dtype: DType::F32,
13008 });
13009 g.set_target_backend(op, BackendId::Cpu);
13010 (q, k, v, op)
13011 };
13012 let mut inputs = StorageCache::new();
13013 inputs.insert(q_id, Arc::new(RwLock::new(q)));
13014 inputs.insert(k_id, Arc::new(RwLock::new(k)));
13015 inputs.insert(v_id, Arc::new(RwLock::new(v)));
13016 let mut env = SymEnv::new();
13017 env.bind(sym, 3).unwrap();
13018 let (result_arc, _) = PipelinedExecutor::realize_with_env(graph, op_id, inputs, env)
13019 .expect("realize_with_env");
13020 let guard = result_arc.read().unwrap();
13021 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13022 panic!("expected Cpu storage");
13023 };
13024 let r = c.as_slice::<f32>().unwrap();
13025 let e = 1.0_f32.exp();
13026 let d0 = e + 1.0;
13029 let e0x = (5.0 * e + 7.0) / d0;
13030 let e0y = (6.0 * e + 8.0) / d0;
13031 let d1 = 1.0 + 2.0 * e;
13033 let e1x = (5.0 + 7.0 * e + 9.0 * e) / d1;
13034 let e1y = (6.0 + 8.0 * e + 10.0 * e) / d1;
13035 for (got, want) in r.iter().zip([e0x, e0y, e1x, e1y].iter()) {
13036 assert!(
13037 (got - want).abs() < 1e-4,
13038 "flash dyn k_len: got {got}, want {want} (full output {r:?}); \
13039 a poison-contaminated value means the capacity rows past k_len leaked in",
13040 );
13041 }
13042 }
13043
13044 #[test]
13047 fn pipelined_realize_flash_attn_dynamic_k_len_unbound_errors() {
13048 use fuel_ir::{DynScalar, SymId};
13049 let q = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
13050 let k = fuel_memory::from_slice_cpu(&[0.0_f32; 16]);
13051 let v = fuel_memory::from_slice_cpu(&[0.0_f32; 16]);
13052 let graph = Arc::new(RwLock::new(Graph::new()));
13053 let (q_id, k_id, v_id, op_id) = {
13054 let mut g = graph.write().unwrap();
13055 let q = g.push(Node {
13056 op: Op::Const,
13057 inputs: vec![],
13058 shape: Shape::from_dims(&[1, 1, 2, 2]),
13059 dtype: DType::F32,
13060 });
13061 let k = g.push(Node {
13062 op: Op::Const,
13063 inputs: vec![],
13064 shape: Shape::from_dims(&[1, 1, 4, 2]),
13065 dtype: DType::F32,
13066 });
13067 let v = g.push(Node {
13068 op: Op::Const,
13069 inputs: vec![],
13070 shape: Shape::from_dims(&[1, 1, 4, 2]),
13071 dtype: DType::F32,
13072 });
13073 let op = g.push(Node {
13074 op: Op::Fused(
13075 fuel_graph::registry::FusedOps::FLASH_ATTN,
13076 fuel_graph::registry::FusedOpParams::FlashAttn {
13077 softmax_scale: 1.0,
13078 causal: true,
13079 window_size_left: None,
13080 window_size_right: None,
13081 softcap: None,
13082 k_len: Some(DynScalar::Sym(SymId(0))),
13083 },
13084 ),
13085 inputs: vec![q, k, v],
13086 shape: Shape::from_dims(&[1, 1, 2, 2]),
13087 dtype: DType::F32,
13088 });
13089 g.set_target_backend(op, BackendId::Cpu);
13090 (q, k, v, op)
13091 };
13092 let mut inputs = StorageCache::new();
13093 inputs.insert(q_id, Arc::new(RwLock::new(q)));
13094 inputs.insert(k_id, Arc::new(RwLock::new(k)));
13095 inputs.insert(v_id, Arc::new(RwLock::new(v)));
13096 let result = PipelinedExecutor::realize_with_env(graph, op_id, inputs, SymEnv::new());
13098 assert!(
13099 result.is_err(),
13100 "unbound flash k_len symbol must surface a typed error"
13101 );
13102 }
13103
13104 #[test]
13106 fn pipelined_realize_flash_attn_bf16() {
13107 let q_v: Vec<half::bf16> = [2.0_f32, 0.0]
13108 .iter()
13109 .map(|x| half::bf16::from_f32(*x))
13110 .collect();
13111 let k_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0]
13112 .iter()
13113 .map(|x| half::bf16::from_f32(*x))
13114 .collect();
13115 let v_v: Vec<half::bf16> = [10.0_f32, 0.0, 0.0, 10.0]
13116 .iter()
13117 .map(|x| half::bf16::from_f32(*x))
13118 .collect();
13119 let q = fuel_memory::from_slice_cpu(&q_v);
13120 let k = fuel_memory::from_slice_cpu(&k_v);
13121 let v = fuel_memory::from_slice_cpu(&v_v);
13122 let graph = Arc::new(RwLock::new(Graph::new()));
13123 let (q_id, k_id, v_id, op_id) = {
13124 let mut g = graph.write().unwrap();
13125 let q = g.push(Node {
13126 op: Op::Const,
13127 inputs: vec![],
13128 shape: Shape::from_dims(&[1, 1, 1, 2]),
13129 dtype: DType::BF16,
13130 });
13131 let k = g.push(Node {
13132 op: Op::Const,
13133 inputs: vec![],
13134 shape: Shape::from_dims(&[1, 1, 2, 2]),
13135 dtype: DType::BF16,
13136 });
13137 let v = g.push(Node {
13138 op: Op::Const,
13139 inputs: vec![],
13140 shape: Shape::from_dims(&[1, 1, 2, 2]),
13141 dtype: DType::BF16,
13142 });
13143 let op = g.push(Node {
13144 op: Op::Fused(
13145 fuel_graph::registry::FusedOps::FLASH_ATTN,
13146 fuel_graph::registry::FusedOpParams::FlashAttn {
13147 softmax_scale: 1.0,
13148 causal: false,
13149 window_size_left: None,
13150 window_size_right: None,
13151 softcap: None,
13152 k_len: None,
13153 },
13154 ),
13155 inputs: vec![q, k, v],
13156 shape: Shape::from_dims(&[1, 1, 1, 2]),
13157 dtype: DType::BF16,
13158 });
13159 g.set_target_backend(op, BackendId::Cpu);
13160 (q, k, v, op)
13161 };
13162 let mut inputs = StorageCache::new();
13163 inputs.insert(q_id, Arc::new(RwLock::new(q)));
13164 inputs.insert(k_id, Arc::new(RwLock::new(k)));
13165 inputs.insert(v_id, Arc::new(RwLock::new(v)));
13166 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13167 let guard = result_arc.read().unwrap();
13168 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13169 panic!("expected Cpu storage");
13170 };
13171 let got: Vec<f32> = c
13172 .as_slice::<half::bf16>()
13173 .unwrap()
13174 .iter()
13175 .map(|v| v.to_f32())
13176 .collect();
13177 let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
13178 let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
13179 let want = [10.0 * expected_p0, 10.0 * expected_p1];
13180 for (g, w) in got.iter().zip(want.iter()) {
13181 assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
13182 }
13183 }
13184
13185 #[test]
13190 fn pipelined_realize_fused_linear_f32() {
13191 let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
13192 let b = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0, 1.0, 1.0]);
13193 let bias = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
13194 let graph = Arc::new(RwLock::new(Graph::new()));
13195 let (a_id, b_id, bias_id, op_id) = {
13196 let mut g = graph.write().unwrap();
13197 let a = g.push(Node {
13198 op: Op::Const,
13199 inputs: vec![],
13200 shape: Shape::from_dims(&[1, 2, 3]),
13201 dtype: DType::F32,
13202 });
13203 let b = g.push(Node {
13204 op: Op::Const,
13205 inputs: vec![],
13206 shape: Shape::from_dims(&[1, 3, 2]),
13207 dtype: DType::F32,
13208 });
13209 let bias = g.push(Node {
13210 op: Op::Const,
13211 inputs: vec![],
13212 shape: Shape::from_dims(&[2]),
13213 dtype: DType::F32,
13214 });
13215 let op = g.push(Node {
13216 op: Op::Fused(
13217 fuel_graph::registry::FusedOps::FUSED_LINEAR,
13218 fuel_graph::registry::FusedOpParams::FusedLinear,
13219 ),
13220 inputs: vec![a, b, bias],
13221 shape: Shape::from_dims(&[1, 2, 2]),
13222 dtype: DType::F32,
13223 });
13224 g.set_target_backend(op, BackendId::Cpu);
13225 (a, b, bias, op)
13226 };
13227 let mut inputs = StorageCache::new();
13228 inputs.insert(a_id, Arc::new(RwLock::new(a)));
13229 inputs.insert(b_id, Arc::new(RwLock::new(b)));
13230 inputs.insert(bias_id, Arc::new(RwLock::new(bias)));
13231 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13232 let guard = result_arc.read().unwrap();
13233 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13234 panic!("expected Cpu storage");
13235 };
13236 assert_eq!(c.as_slice::<f32>().unwrap(), &[14.0, 25.0, 20.0, 31.0]);
13237 }
13238
13239 #[test]
13241 fn pipelined_realize_fused_linear_f64() {
13242 let a = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
13243 let b = fuel_memory::from_slice_cpu(&[1.0_f64, 0.0, 0.0, 1.0, 1.0, 1.0]);
13244 let bias = fuel_memory::from_slice_cpu(&[10.0_f64, 20.0]);
13245 let graph = Arc::new(RwLock::new(Graph::new()));
13246 let (a_id, b_id, bias_id, op_id) = {
13247 let mut g = graph.write().unwrap();
13248 let a = g.push(Node {
13249 op: Op::Const,
13250 inputs: vec![],
13251 shape: Shape::from_dims(&[1, 2, 3]),
13252 dtype: DType::F64,
13253 });
13254 let b = g.push(Node {
13255 op: Op::Const,
13256 inputs: vec![],
13257 shape: Shape::from_dims(&[1, 3, 2]),
13258 dtype: DType::F64,
13259 });
13260 let bias = g.push(Node {
13261 op: Op::Const,
13262 inputs: vec![],
13263 shape: Shape::from_dims(&[2]),
13264 dtype: DType::F64,
13265 });
13266 let op = g.push(Node {
13267 op: Op::Fused(
13268 fuel_graph::registry::FusedOps::FUSED_LINEAR,
13269 fuel_graph::registry::FusedOpParams::FusedLinear,
13270 ),
13271 inputs: vec![a, b, bias],
13272 shape: Shape::from_dims(&[1, 2, 2]),
13273 dtype: DType::F64,
13274 });
13275 g.set_target_backend(op, BackendId::Cpu);
13276 (a, b, bias, op)
13277 };
13278 let mut inputs = StorageCache::new();
13279 inputs.insert(a_id, Arc::new(RwLock::new(a)));
13280 inputs.insert(b_id, Arc::new(RwLock::new(b)));
13281 inputs.insert(bias_id, Arc::new(RwLock::new(bias)));
13282 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13283 let guard = result_arc.read().unwrap();
13284 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13285 panic!("expected Cpu storage");
13286 };
13287 assert_eq!(c.as_slice::<f64>().unwrap(), &[14.0, 25.0, 20.0, 31.0]);
13288 }
13289
13290 #[test]
13292 fn pipelined_realize_fused_linear_bf16() {
13293 let a_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
13294 .iter()
13295 .map(|x| half::bf16::from_f32(*x))
13296 .collect();
13297 let b_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0, 1.0, 1.0]
13298 .iter()
13299 .map(|x| half::bf16::from_f32(*x))
13300 .collect();
13301 let bias_v: Vec<half::bf16> = [10.0_f32, 20.0]
13302 .iter()
13303 .map(|x| half::bf16::from_f32(*x))
13304 .collect();
13305 let a = fuel_memory::from_slice_cpu(&a_v);
13306 let b = fuel_memory::from_slice_cpu(&b_v);
13307 let bias = fuel_memory::from_slice_cpu(&bias_v);
13308 let graph = Arc::new(RwLock::new(Graph::new()));
13309 let (a_id, b_id, bias_id, op_id) = {
13310 let mut g = graph.write().unwrap();
13311 let a = g.push(Node {
13312 op: Op::Const,
13313 inputs: vec![],
13314 shape: Shape::from_dims(&[1, 2, 3]),
13315 dtype: DType::BF16,
13316 });
13317 let b = g.push(Node {
13318 op: Op::Const,
13319 inputs: vec![],
13320 shape: Shape::from_dims(&[1, 3, 2]),
13321 dtype: DType::BF16,
13322 });
13323 let bias = g.push(Node {
13324 op: Op::Const,
13325 inputs: vec![],
13326 shape: Shape::from_dims(&[2]),
13327 dtype: DType::BF16,
13328 });
13329 let op = g.push(Node {
13330 op: Op::Fused(
13331 fuel_graph::registry::FusedOps::FUSED_LINEAR,
13332 fuel_graph::registry::FusedOpParams::FusedLinear,
13333 ),
13334 inputs: vec![a, b, bias],
13335 shape: Shape::from_dims(&[1, 2, 2]),
13336 dtype: DType::BF16,
13337 });
13338 g.set_target_backend(op, BackendId::Cpu);
13339 (a, b, bias, op)
13340 };
13341 let mut inputs = StorageCache::new();
13342 inputs.insert(a_id, Arc::new(RwLock::new(a)));
13343 inputs.insert(b_id, Arc::new(RwLock::new(b)));
13344 inputs.insert(bias_id, Arc::new(RwLock::new(bias)));
13345 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13346 let guard = result_arc.read().unwrap();
13347 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13348 panic!("expected Cpu storage");
13349 };
13350 let got: Vec<f32> = c
13351 .as_slice::<half::bf16>()
13352 .unwrap()
13353 .iter()
13354 .map(|v| v.to_f32())
13355 .collect();
13356 let want = [14.0_f32, 25.0, 20.0, 31.0];
13357 for (g, w) in got.iter().zip(want.iter()) {
13358 assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
13359 }
13360 }
13361
13362 #[test]
13365 fn pipelined_realize_reduce_sum_to_f32_drops_leading_axis() {
13366 let v = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
13367 let graph = Arc::new(RwLock::new(Graph::new()));
13368 let (in_id, op_id) = {
13369 let mut g = graph.write().unwrap();
13370 let in_id = g.push(Node {
13371 op: Op::Const,
13372 inputs: vec![],
13373 shape: Shape::from_dims(&[2, 3]),
13374 dtype: DType::F32,
13375 });
13376 let op_id = g.push(Node {
13377 op: Op::ReduceSumTo(Shape::from_dims(&[3])),
13378 inputs: vec![in_id],
13379 shape: Shape::from_dims(&[3]),
13380 dtype: DType::F32,
13381 });
13382 g.set_target_backend(op_id, BackendId::Cpu);
13383 (in_id, op_id)
13384 };
13385 let mut inputs = StorageCache::new();
13386 inputs.insert(in_id, Arc::new(RwLock::new(v)));
13387 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13388 let guard = result_arc.read().unwrap();
13389 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13390 panic!("expected Cpu storage");
13391 };
13392 assert_eq!(c.as_slice::<f32>().unwrap(), &[5.0, 7.0, 9.0]);
13393 }
13394
13395 #[test]
13398 fn pipelined_realize_reduce_sum_to_f32_keepdim_middle() {
13399 let mut v: Vec<f32> = (1..=24).map(|x| x as f32).collect();
13402 let _ = &mut v;
13403 let s = fuel_memory::from_slice_cpu(&v);
13404 let graph = Arc::new(RwLock::new(Graph::new()));
13405 let (in_id, op_id) = {
13406 let mut g = graph.write().unwrap();
13407 let in_id = g.push(Node {
13408 op: Op::Const,
13409 inputs: vec![],
13410 shape: Shape::from_dims(&[2, 3, 4]),
13411 dtype: DType::F32,
13412 });
13413 let op_id = g.push(Node {
13414 op: Op::ReduceSumTo(Shape::from_dims(&[2, 1, 4])),
13415 inputs: vec![in_id],
13416 shape: Shape::from_dims(&[2, 1, 4]),
13417 dtype: DType::F32,
13418 });
13419 g.set_target_backend(op_id, BackendId::Cpu);
13420 (in_id, op_id)
13421 };
13422 let mut inputs = StorageCache::new();
13423 inputs.insert(in_id, Arc::new(RwLock::new(s)));
13424 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13425 let guard = result_arc.read().unwrap();
13426 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13427 panic!("expected Cpu storage");
13428 };
13429 assert_eq!(
13432 c.as_slice::<f32>().unwrap(),
13433 &[15.0, 18.0, 21.0, 24.0, 51.0, 54.0, 57.0, 60.0],
13434 );
13435 }
13436
13437 #[test]
13439 fn pipelined_realize_reduce_sum_to_f64() {
13440 let v = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
13441 let graph = Arc::new(RwLock::new(Graph::new()));
13442 let (in_id, op_id) = {
13443 let mut g = graph.write().unwrap();
13444 let in_id = g.push(Node {
13445 op: Op::Const,
13446 inputs: vec![],
13447 shape: Shape::from_dims(&[2, 3]),
13448 dtype: DType::F64,
13449 });
13450 let op_id = g.push(Node {
13451 op: Op::ReduceSumTo(Shape::from_dims(&[3])),
13452 inputs: vec![in_id],
13453 shape: Shape::from_dims(&[3]),
13454 dtype: DType::F64,
13455 });
13456 g.set_target_backend(op_id, BackendId::Cpu);
13457 (in_id, op_id)
13458 };
13459 let mut inputs = StorageCache::new();
13460 inputs.insert(in_id, Arc::new(RwLock::new(v)));
13461 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13462 let guard = result_arc.read().unwrap();
13463 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13464 panic!("expected Cpu storage");
13465 };
13466 assert_eq!(c.as_slice::<f64>().unwrap(), &[5.0, 7.0, 9.0]);
13467 }
13468
13469 #[test]
13471 fn pipelined_realize_reduce_max_to_f32_drops_leading_axis() {
13472 let v = fuel_memory::from_slice_cpu(&[1.0_f32, 7.0, 3.0, 4.0, 2.0, 6.0]);
13475 let graph = Arc::new(RwLock::new(Graph::new()));
13476 let (in_id, op_id) = {
13477 let mut g = graph.write().unwrap();
13478 let in_id = g.push(Node {
13479 op: Op::Const,
13480 inputs: vec![],
13481 shape: Shape::from_dims(&[2, 3]),
13482 dtype: DType::F32,
13483 });
13484 let op_id = g.push(Node {
13485 op: Op::ReduceMaxTo(Shape::from_dims(&[3])),
13486 inputs: vec![in_id],
13487 shape: Shape::from_dims(&[3]),
13488 dtype: DType::F32,
13489 });
13490 g.set_target_backend(op_id, BackendId::Cpu);
13491 (in_id, op_id)
13492 };
13493 let mut inputs = StorageCache::new();
13494 inputs.insert(in_id, Arc::new(RwLock::new(v)));
13495 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13496 let guard = result_arc.read().unwrap();
13497 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13498 panic!("expected Cpu storage");
13499 };
13500 assert_eq!(c.as_slice::<f32>().unwrap(), &[4.0, 7.0, 6.0]);
13501 }
13502
13503 #[test]
13507 fn pipelined_realize_reduce_max_to_f32_keepdim_trailing() {
13508 let v = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
13510 let graph = Arc::new(RwLock::new(Graph::new()));
13511 let (in_id, op_id) = {
13512 let mut g = graph.write().unwrap();
13513 let in_id = g.push(Node {
13514 op: Op::Const,
13515 inputs: vec![],
13516 shape: Shape::from_dims(&[2, 3]),
13517 dtype: DType::F32,
13518 });
13519 let op_id = g.push(Node {
13520 op: Op::ReduceMaxTo(Shape::from_dims(&[2, 1])),
13521 inputs: vec![in_id],
13522 shape: Shape::from_dims(&[2, 1]),
13523 dtype: DType::F32,
13524 });
13525 g.set_target_backend(op_id, BackendId::Cpu);
13526 (in_id, op_id)
13527 };
13528 let mut inputs = StorageCache::new();
13529 inputs.insert(in_id, Arc::new(RwLock::new(v)));
13530 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13531 let guard = result_arc.read().unwrap();
13532 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13533 panic!("expected Cpu storage");
13534 };
13535 assert_eq!(c.as_slice::<f32>().unwrap(), &[3.0, 6.0]);
13536 }
13537
13538 #[test]
13540 fn pipelined_realize_reduce_sum_to_bf16() {
13541 let v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
13542 .iter()
13543 .map(|x| half::bf16::from_f32(*x))
13544 .collect();
13545 let s = fuel_memory::from_slice_cpu(&v);
13546 let graph = Arc::new(RwLock::new(Graph::new()));
13547 let (in_id, op_id) = {
13548 let mut g = graph.write().unwrap();
13549 let in_id = g.push(Node {
13550 op: Op::Const,
13551 inputs: vec![],
13552 shape: Shape::from_dims(&[2, 3]),
13553 dtype: DType::BF16,
13554 });
13555 let op_id = g.push(Node {
13556 op: Op::ReduceSumTo(Shape::from_dims(&[3])),
13557 inputs: vec![in_id],
13558 shape: Shape::from_dims(&[3]),
13559 dtype: DType::BF16,
13560 });
13561 g.set_target_backend(op_id, BackendId::Cpu);
13562 (in_id, op_id)
13563 };
13564 let mut inputs = StorageCache::new();
13565 inputs.insert(in_id, Arc::new(RwLock::new(s)));
13566 let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
13567 let guard = result_arc.read().unwrap();
13568 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13569 panic!("expected Cpu storage");
13570 };
13571 let got: Vec<f32> = c
13572 .as_slice::<half::bf16>()
13573 .unwrap()
13574 .iter()
13575 .map(|v| v.to_f32())
13576 .collect();
13577 let want = [5.0_f32, 7.0, 9.0];
13578 for (g, w) in got.iter().zip(want.iter()) {
13579 assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
13580 }
13581 }
13582
13583 #[test]
13591 fn pipelined_realize_conv_transpose2d_f32() {
13592 let x_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
13593 let w_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
13594 let graph = Arc::new(RwLock::new(Graph::new()));
13595 let (x_id, w_id, c_id) = {
13596 let mut g = graph.write().unwrap();
13597 let x = g.push(Node {
13598 op: Op::Const,
13599 inputs: vec![],
13600 shape: Shape::from_dims(&[1, 1, 2, 2]),
13601 dtype: DType::F32,
13602 });
13603 let w = g.push(Node {
13604 op: Op::Const,
13605 inputs: vec![],
13606 shape: Shape::from_dims(&[1, 1, 2, 2]),
13607 dtype: DType::F32,
13608 });
13609 let c = g.push(Node {
13610 op: Op::Fused(
13611 fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D,
13612 fuel_graph::registry::FusedOpParams::ConvTranspose2D {
13613 stride: (1, 1),
13614 padding: (0, 0),
13615 output_padding: (0, 0),
13616 dilation: (1, 1),
13617 groups: 1,
13618 },
13619 ),
13620 inputs: vec![x, w],
13621 shape: Shape::from_dims(&[1, 1, 3, 3]),
13622 dtype: DType::F32,
13623 });
13624 g.set_target_backend(c, BackendId::Cpu);
13625 (x, w, c)
13626 };
13627 let mut inputs = StorageCache::new();
13628 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
13629 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
13630 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
13631 let guard = result_arc.read().unwrap();
13632 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13633 panic!("expected Cpu storage");
13634 };
13635 assert_eq!(
13636 c.as_slice::<f32>().unwrap(),
13637 &[1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0],
13638 );
13639 }
13640
13641 #[test]
13643 fn pipelined_realize_conv_transpose2d_f64() {
13644 let x_storage = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0]);
13645 let w_storage = fuel_memory::from_slice_cpu(&[1.0_f64, 1.0, 1.0, 1.0]);
13646 let graph = Arc::new(RwLock::new(Graph::new()));
13647 let (x_id, w_id, c_id) = {
13648 let mut g = graph.write().unwrap();
13649 let x = g.push(Node {
13650 op: Op::Const,
13651 inputs: vec![],
13652 shape: Shape::from_dims(&[1, 1, 2, 2]),
13653 dtype: DType::F64,
13654 });
13655 let w = g.push(Node {
13656 op: Op::Const,
13657 inputs: vec![],
13658 shape: Shape::from_dims(&[1, 1, 2, 2]),
13659 dtype: DType::F64,
13660 });
13661 let c = g.push(Node {
13662 op: Op::Fused(
13663 fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D,
13664 fuel_graph::registry::FusedOpParams::ConvTranspose2D {
13665 stride: (1, 1),
13666 padding: (0, 0),
13667 output_padding: (0, 0),
13668 dilation: (1, 1),
13669 groups: 1,
13670 },
13671 ),
13672 inputs: vec![x, w],
13673 shape: Shape::from_dims(&[1, 1, 3, 3]),
13674 dtype: DType::F64,
13675 });
13676 g.set_target_backend(c, BackendId::Cpu);
13677 (x, w, c)
13678 };
13679 let mut inputs = StorageCache::new();
13680 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
13681 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
13682 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
13683 let guard = result_arc.read().unwrap();
13684 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13685 panic!("expected Cpu storage");
13686 };
13687 assert_eq!(
13688 c.as_slice::<f64>().unwrap(),
13689 &[1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0],
13690 );
13691 }
13692
13693 #[test]
13695 fn pipelined_realize_conv_transpose2d_bf16() {
13696 let x: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0]
13697 .iter()
13698 .map(|v| half::bf16::from_f32(*v))
13699 .collect();
13700 let w: Vec<half::bf16> = [1.0_f32, 1.0, 1.0, 1.0]
13701 .iter()
13702 .map(|v| half::bf16::from_f32(*v))
13703 .collect();
13704 let x_storage = fuel_memory::from_slice_cpu(&x);
13705 let w_storage = fuel_memory::from_slice_cpu(&w);
13706 let graph = Arc::new(RwLock::new(Graph::new()));
13707 let (x_id, w_id, c_id) = {
13708 let mut g = graph.write().unwrap();
13709 let x = g.push(Node {
13710 op: Op::Const,
13711 inputs: vec![],
13712 shape: Shape::from_dims(&[1, 1, 2, 2]),
13713 dtype: DType::BF16,
13714 });
13715 let w = g.push(Node {
13716 op: Op::Const,
13717 inputs: vec![],
13718 shape: Shape::from_dims(&[1, 1, 2, 2]),
13719 dtype: DType::BF16,
13720 });
13721 let c = g.push(Node {
13722 op: Op::Fused(
13723 fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D,
13724 fuel_graph::registry::FusedOpParams::ConvTranspose2D {
13725 stride: (1, 1),
13726 padding: (0, 0),
13727 output_padding: (0, 0),
13728 dilation: (1, 1),
13729 groups: 1,
13730 },
13731 ),
13732 inputs: vec![x, w],
13733 shape: Shape::from_dims(&[1, 1, 3, 3]),
13734 dtype: DType::BF16,
13735 });
13736 g.set_target_backend(c, BackendId::Cpu);
13737 (x, w, c)
13738 };
13739 let mut inputs = StorageCache::new();
13740 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
13741 inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
13742 let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
13743 let guard = result_arc.read().unwrap();
13744 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13745 panic!("expected Cpu storage");
13746 };
13747 let got: Vec<f32> = c
13748 .as_slice::<half::bf16>()
13749 .unwrap()
13750 .iter()
13751 .map(|v| v.to_f32())
13752 .collect();
13753 let want = [1.0_f32, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0];
13754 for (g, w) in got.iter().zip(want.iter()) {
13755 assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
13756 }
13757 }
13758
13759 #[test]
13763 fn pipelined_realize_matmul_gqa() {
13764 let lhs_storage =
13768 fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
13769 let rhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
13770
13771 let graph = Arc::new(RwLock::new(Graph::new()));
13772 let (lhs_id, rhs_id, mm_id) = {
13773 let mut g = graph.write().unwrap();
13774 let lhs = g.push(Node {
13775 op: Op::Const,
13776 inputs: vec![],
13777 shape: Shape::from_dims(&[4, 1, 2]),
13778 dtype: DType::F32,
13779 });
13780 let rhs = g.push(Node {
13781 op: Op::Const,
13782 inputs: vec![],
13783 shape: Shape::from_dims(&[2, 2, 1]),
13784 dtype: DType::F32,
13785 });
13786 let mm = g.push(Node {
13787 op: Op::MatMul,
13788 inputs: vec![lhs, rhs],
13789 shape: Shape::from_dims(&[4, 1, 1]),
13790 dtype: DType::F32,
13791 });
13792 g.set_target_backend(mm, BackendId::Cpu);
13793 (lhs, rhs, mm)
13794 };
13795 let mut inputs = StorageCache::new();
13796 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
13797 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
13798
13799 let (result_arc, result_layout) =
13800 PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
13801 assert_eq!(result_layout.shape().dims(), &[4, 1, 1]);
13802
13803 let guard = result_arc.read().unwrap();
13804 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13805 panic!("expected Cpu storage");
13806 };
13807 assert_eq!(c.as_slice::<f32>().unwrap(), &[1.0, 3.0, 6.0, 8.0]);
13808 }
13809
13810 #[test]
13815 fn pipelined_realize_matmul_with_transposed_rhs() {
13816 let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
13821 let rhs_storage = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0, 7.0, 8.0]);
13822
13823 let graph = Arc::new(RwLock::new(Graph::new()));
13824 let (lhs_id, rhs_id, t_id, mm_id) = {
13825 let mut g = graph.write().unwrap();
13826 let lhs = g.push(Node {
13827 op: Op::Const,
13828 inputs: vec![],
13829 shape: Shape::from_dims(&[2, 2]),
13830 dtype: DType::F32,
13831 });
13832 let rhs = g.push(Node {
13833 op: Op::Const,
13834 inputs: vec![],
13835 shape: Shape::from_dims(&[2, 2]),
13836 dtype: DType::F32,
13837 });
13838 let t = g.push(Node {
13839 op: Op::Transpose,
13840 inputs: vec![rhs],
13841 shape: Shape::from_dims(&[2, 2]),
13842 dtype: DType::F32,
13843 });
13844 let mm = g.push(Node {
13845 op: Op::MatMul,
13846 inputs: vec![lhs, t],
13847 shape: Shape::from_dims(&[2, 2]),
13848 dtype: DType::F32,
13849 });
13850 g.set_target_backend(mm, BackendId::Cpu);
13851 (lhs, rhs, t, mm)
13852 };
13853 let _ = t_id;
13854 let mut inputs = StorageCache::new();
13855 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
13856 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
13857
13858 let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
13859 let guard = result_arc.read().unwrap();
13860 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13861 panic!("expected Cpu storage");
13862 };
13863 assert_eq!(c.as_slice::<f32>().unwrap(), &[17.0, 23.0, 39.0, 53.0]);
13864 }
13865
13866 #[test]
13870 fn pipelined_realize_reshape_zero_copy_when_contiguous() {
13871 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
13872 let graph = Arc::new(RwLock::new(Graph::new()));
13873 let (in_id, r_id) = {
13874 let mut g = graph.write().unwrap();
13875 let in_id = g.push(Node {
13876 op: Op::Const,
13877 inputs: vec![],
13878 shape: Shape::from_dims(&[2, 3]),
13879 dtype: DType::F32,
13880 });
13881 let r_id = g.push(Node {
13882 op: Op::Reshape(Shape::from_dims(&[3, 2])),
13883 inputs: vec![in_id],
13884 shape: Shape::from_dims(&[3, 2]),
13885 dtype: DType::F32,
13886 });
13887 (in_id, r_id)
13888 };
13889 let mut inputs = StorageCache::new();
13890 let in_arc = Arc::new(RwLock::new(storage));
13891 inputs.insert(in_id, Arc::clone(&in_arc));
13892
13893 let (result_arc, result_layout) =
13894 PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
13895
13896 assert!(
13898 Arc::ptr_eq(&result_arc, &in_arc),
13899 "contiguous reshape must zero-copy"
13900 );
13901 assert_eq!(result_layout.shape().dims(), &[3, 2]);
13902 assert!(result_layout.is_contiguous());
13903
13904 let guard = result_arc.read().unwrap();
13906 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13907 panic!("expected Cpu storage");
13908 };
13909 assert_eq!(
13910 c.as_slice::<f32>().unwrap(),
13911 &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
13912 );
13913 }
13914
13915 #[test]
13920 fn pipelined_realize_reshape_materializes_when_strided() {
13921 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
13925 let graph = Arc::new(RwLock::new(Graph::new()));
13926 let (in_id, t_id, r_id) = {
13927 let mut g = graph.write().unwrap();
13928 let in_id = g.push(Node {
13929 op: Op::Const,
13930 inputs: vec![],
13931 shape: Shape::from_dims(&[2, 3]),
13932 dtype: DType::F32,
13933 });
13934 let t_id = g.push(Node {
13935 op: Op::Transpose,
13936 inputs: vec![in_id],
13937 shape: Shape::from_dims(&[3, 2]),
13938 dtype: DType::F32,
13939 });
13940 let r_id = g.push(Node {
13941 op: Op::Reshape(Shape::from_dims(&[6])),
13942 inputs: vec![t_id],
13943 shape: Shape::from_dims(&[6]),
13944 dtype: DType::F32,
13945 });
13946 (in_id, t_id, r_id)
13947 };
13948 let _ = t_id;
13949 let mut inputs = StorageCache::new();
13950 let in_arc = Arc::new(RwLock::new(storage));
13951 inputs.insert(in_id, Arc::clone(&in_arc));
13952
13953 let (result_arc, result_layout) =
13954 PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
13955
13956 assert!(!Arc::ptr_eq(&result_arc, &in_arc));
13958 assert_eq!(result_layout.shape().dims(), &[6]);
13959 assert!(result_layout.is_contiguous());
13960
13961 let guard = result_arc.read().unwrap();
13963 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
13964 panic!("expected Cpu storage");
13965 };
13966 assert_eq!(
13967 c.as_slice::<f32>().unwrap(),
13968 &[1.0, 4.0, 2.0, 5.0, 3.0, 6.0]
13969 );
13970 }
13971
13972 #[test]
13980 fn pipelined_realize_transpose_then_sum_dim_e2e() {
13981 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
13985 let graph = Arc::new(RwLock::new(Graph::new()));
13986 let (in_id, t_id, sum_id) = {
13987 let mut g = graph.write().unwrap();
13988 let in_id = g.push(Node {
13989 op: Op::Const,
13990 inputs: vec![],
13991 shape: Shape::from_dims(&[2, 3]),
13992 dtype: DType::F32,
13993 });
13994 let t_id = g.push(Node {
13995 op: Op::Transpose,
13996 inputs: vec![in_id],
13997 shape: Shape::from_dims(&[3, 2]),
13998 dtype: DType::F32,
13999 });
14000 let sum_id = g.push(Node {
14001 op: Op::SumDim(1),
14002 inputs: vec![t_id],
14003 shape: Shape::from_dims(&[3]),
14004 dtype: DType::F32,
14005 });
14006 g.set_target_backend(sum_id, BackendId::Cpu);
14009 (in_id, t_id, sum_id)
14010 };
14011 let _ = t_id;
14012 let mut inputs = StorageCache::new();
14013 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14014
14015 let (result_arc, result_layout) =
14016 PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
14017
14018 assert_eq!(result_layout.shape().dims(), &[3]);
14020 assert!(result_layout.is_contiguous());
14021
14022 let guard = result_arc.read().unwrap();
14023 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
14024 panic!("expected Cpu storage");
14025 };
14026 let typed: &[f32] = c.as_slice().unwrap();
14027 assert_eq!(typed, &[5.0, 7.0, 9.0]);
14028 }
14029
14030 #[test]
14034 fn pipelined_realize_broadcast_then_add_e2e() {
14035 let bc_input = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
14040 let plus_input = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
14041 let graph = Arc::new(RwLock::new(Graph::new()));
14042 let (bc_in_id, plus_in_id, b_id, add_id) = {
14043 let mut g = graph.write().unwrap();
14044 let bc_in = g.push(Node {
14045 op: Op::Const,
14046 inputs: vec![],
14047 shape: Shape::from_dims(&[3]),
14048 dtype: DType::F32,
14049 });
14050 let plus_in = g.push(Node {
14051 op: Op::Const,
14052 inputs: vec![],
14053 shape: Shape::from_dims(&[2, 3]),
14054 dtype: DType::F32,
14055 });
14056 let b = g.push(Node {
14057 op: Op::BroadcastTo(Shape::from_dims(&[2, 3])),
14058 inputs: vec![bc_in],
14059 shape: Shape::from_dims(&[2, 3]),
14060 dtype: DType::F32,
14061 });
14062 let add = g.push(Node {
14063 op: Op::Add,
14064 inputs: vec![b, plus_in],
14065 shape: Shape::from_dims(&[2, 3]),
14066 dtype: DType::F32,
14067 });
14068 g.set_target_backend(add, BackendId::Cpu);
14069 (bc_in, plus_in, b, add)
14070 };
14071 let _ = b_id;
14072 let mut inputs = StorageCache::new();
14073 inputs.insert(bc_in_id, Arc::new(RwLock::new(bc_input)));
14074 inputs.insert(plus_in_id, Arc::new(RwLock::new(plus_input)));
14075
14076 let (result_arc, _result_layout) =
14077 PipelinedExecutor::realize(graph, add_id, inputs).expect("realize");
14078 let guard = result_arc.read().unwrap();
14079 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
14080 panic!("expected Cpu storage");
14081 };
14082 let typed: &[f32] = c.as_slice().unwrap();
14083 assert_eq!(typed, &[11.0, 22.0, 33.0, 14.0, 25.0, 36.0]);
14084 }
14085
14086 #[test]
14090 fn pipelined_realize_sum_dim() {
14091 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
14093 let graph = Arc::new(RwLock::new(Graph::new()));
14094 let (in_id, sum_id) = {
14095 let mut g = graph.write().unwrap();
14096 let in_id = g.push(Node {
14097 op: Op::Const,
14098 inputs: vec![],
14099 shape: Shape::from_dims(&[2, 3]),
14100 dtype: DType::F32,
14101 });
14102 let sum_id = g.push(Node {
14103 op: Op::SumDim(1),
14104 inputs: vec![in_id],
14105 shape: Shape::from_dims(&[2]),
14106 dtype: DType::F32,
14107 });
14108 g.set_target_backend(sum_id, BackendId::Cpu);
14109 (in_id, sum_id)
14110 };
14111 let mut inputs = StorageCache::new();
14112 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14113
14114 let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
14115 let guard = result_arc.read().unwrap();
14116 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
14117 let typed: &[f32] = c.as_slice().unwrap();
14118 assert_eq!(typed, &[6.0, 15.0]);
14119 }
14120 }
14121
14122 #[test]
14125 fn pipelined_realize_sum_all_rank3() {
14126 let data: Vec<f32> = (1..=24).map(|x| x as f32).collect();
14127 let storage = fuel_memory::from_slice_cpu(&data);
14128 let graph = Arc::new(RwLock::new(Graph::new()));
14129 let (in_id, sum_id) = {
14130 let mut g = graph.write().unwrap();
14131 let in_id = g.push(Node {
14132 op: Op::Const,
14133 inputs: vec![],
14134 shape: Shape::from_dims(&[2, 3, 4]),
14135 dtype: DType::F32,
14136 });
14137 let sum_id = g.push(Node {
14138 op: Op::SumAll,
14139 inputs: vec![in_id],
14140 shape: Shape::from_dims(&[]),
14141 dtype: DType::F32,
14142 });
14143 g.set_target_backend(sum_id, BackendId::Cpu);
14144 (in_id, sum_id)
14145 };
14146 let mut inputs = StorageCache::new();
14147 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14148
14149 let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
14150 let guard = result_arc.read().unwrap();
14151 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
14152 let typed: &[f32] = c.as_slice().unwrap();
14153 assert_eq!(typed, &[300.0]);
14155 }
14156 }
14157
14158 #[test]
14161 fn pipelined_realize_max_then_mean() {
14162 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 9.0, 3.0, 4.0, 2.0, 8.0]);
14164 let graph = Arc::new(RwLock::new(Graph::new()));
14165 let (in_id, mean_id) = {
14166 let mut g = graph.write().unwrap();
14167 let in_id = g.push(Node {
14168 op: Op::Const,
14169 inputs: vec![],
14170 shape: Shape::from_dims(&[2, 3]),
14171 dtype: DType::F32,
14172 });
14173 let max_id = g.push(Node {
14174 op: Op::MaxDim(1),
14175 inputs: vec![in_id],
14176 shape: Shape::from_dims(&[2]),
14177 dtype: DType::F32,
14178 });
14179 let mean_id = g.push(Node {
14180 op: Op::MeanDim(0),
14181 inputs: vec![max_id],
14182 shape: Shape::from_dims(&[]),
14183 dtype: DType::F32,
14184 });
14185 g.set_target_backend(max_id, BackendId::Cpu);
14186 g.set_target_backend(mean_id, BackendId::Cpu);
14187 (in_id, mean_id)
14188 };
14189 let mut inputs = StorageCache::new();
14190 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14191
14192 let (result_arc, _) = PipelinedExecutor::realize(graph, mean_id, inputs).expect("realize");
14193 let guard = result_arc.read().unwrap();
14194 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
14195 let typed: &[f32] = c.as_slice().unwrap();
14196 assert_eq!(typed, &[8.5]);
14198 }
14199 }
14200
14201 #[test]
14206 fn pipelined_realize_sigmoid_then_silu() {
14207 let storage = fuel_memory::from_slice_cpu(&[0.0_f32, 1.0, -1.0]);
14208 let graph = Arc::new(RwLock::new(Graph::new()));
14209 let (in_id, sig_id, silu_id) = {
14210 let mut g = graph.write().unwrap();
14211 let in_id = g.push(Node {
14212 op: Op::Const,
14213 inputs: vec![],
14214 shape: Shape::from_dims(&[3]),
14215 dtype: DType::F32,
14216 });
14217 let sig_id = g.push(Node {
14218 op: Op::Sigmoid,
14219 inputs: vec![in_id],
14220 shape: Shape::from_dims(&[3]),
14221 dtype: DType::F32,
14222 });
14223 let silu_id = g.push(Node {
14225 op: Op::Silu,
14226 inputs: vec![sig_id],
14227 shape: Shape::from_dims(&[3]),
14228 dtype: DType::F32,
14229 });
14230 g.set_target_backend(sig_id, BackendId::Cpu);
14231 g.set_target_backend(silu_id, BackendId::Cpu);
14232 (in_id, sig_id, silu_id)
14233 };
14234 let _ = sig_id;
14235 let mut inputs = StorageCache::new();
14236 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14237
14238 let (result_arc, _) = PipelinedExecutor::realize(graph, silu_id, inputs).expect("realize");
14239 let guard = result_arc.read().unwrap();
14240 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
14241 let typed: &[f32] = c.as_slice().unwrap();
14242 assert!((typed[0] - 0.5 * (1.0 / (1.0 + (-0.5_f32).exp()))).abs() < 1e-6);
14244 }
14245 }
14246
14247 #[test]
14250 fn pipelined_realize_chained_unary_sqr_then_sqrt() {
14251 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 4.0, 9.0]);
14252 let graph = Arc::new(RwLock::new(Graph::new()));
14253 let (in_id, sqrt_id) = {
14254 let mut g = graph.write().unwrap();
14255 let in_id = g.push(Node {
14256 op: Op::Const,
14257 inputs: vec![],
14258 shape: Shape::from_dims(&[3]),
14259 dtype: DType::F32,
14260 });
14261 let sqr_id = g.push(Node {
14262 op: Op::Sqr,
14263 inputs: vec![in_id],
14264 shape: Shape::from_dims(&[3]),
14265 dtype: DType::F32,
14266 });
14267 let sqrt_id = g.push(Node {
14268 op: Op::Sqrt,
14269 inputs: vec![sqr_id],
14270 shape: Shape::from_dims(&[3]),
14271 dtype: DType::F32,
14272 });
14273 g.set_target_backend(sqr_id, BackendId::Cpu);
14274 g.set_target_backend(sqrt_id, BackendId::Cpu);
14275 (in_id, sqrt_id)
14276 };
14277 let mut inputs = StorageCache::new();
14278 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14279
14280 let (result_arc, _) = PipelinedExecutor::realize(graph, sqrt_id, inputs).expect("realize");
14281 let guard = result_arc.read().unwrap();
14282 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
14283 let typed: &[f32] = c.as_slice().unwrap();
14284 assert_eq!(typed, &[1.0, 4.0, 9.0]);
14286 }
14287 }
14288
14289 #[test]
14293 fn pipelined_realize_chained_adds() {
14294 let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
14295 let b_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
14296 let c_storage = fuel_memory::from_slice_cpu(&[100.0_f32, 200.0]);
14297
14298 let graph = Arc::new(RwLock::new(Graph::new()));
14299 let (a_id, b_id, c_id, ab_id, abc_id) = {
14300 let mut g = graph.write().unwrap();
14301 let a = g.push(Node {
14302 op: Op::Const,
14303 inputs: vec![],
14304 shape: Shape::from_dims(&[2]),
14305 dtype: DType::F32,
14306 });
14307 let b = g.push(Node {
14308 op: Op::Const,
14309 inputs: vec![],
14310 shape: Shape::from_dims(&[2]),
14311 dtype: DType::F32,
14312 });
14313 let c = g.push(Node {
14314 op: Op::Const,
14315 inputs: vec![],
14316 shape: Shape::from_dims(&[2]),
14317 dtype: DType::F32,
14318 });
14319 let ab = g.push(Node {
14320 op: Op::Add,
14321 inputs: vec![a, b],
14322 shape: Shape::from_dims(&[2]),
14323 dtype: DType::F32,
14324 });
14325 let abc = g.push(Node {
14326 op: Op::Add,
14327 inputs: vec![ab, c],
14328 shape: Shape::from_dims(&[2]),
14329 dtype: DType::F32,
14330 });
14331 g.set_target_backend(ab, BackendId::Cpu);
14332 g.set_target_backend(abc, BackendId::Cpu);
14333 (a, b, c, ab, abc)
14334 };
14335
14336 let mut inputs = StorageCache::new();
14337 inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
14338 inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
14339 inputs.insert(c_id, Arc::new(RwLock::new(c_storage)));
14340
14341 let (result_arc, _) = PipelinedExecutor::realize(graph, abc_id, inputs).expect("realize");
14342 let _ = ab_id;
14344
14345 let guard = result_arc.read().unwrap();
14346 if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
14347 let typed: &[f32] = c.as_slice().unwrap();
14348 assert_eq!(typed, &[111.0, 222.0]);
14350 }
14351 }
14352
14353 #[test]
14357 fn pipelined_realize_triu_3x3_diag0() {
14358 let storage =
14359 fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
14360 let graph = Arc::new(RwLock::new(Graph::new()));
14361 let (in_id, out_id) = {
14362 let mut g = graph.write().unwrap();
14363 let in_id = g.push(Node {
14364 op: Op::Const,
14365 inputs: vec![],
14366 shape: Shape::from_dims(&[3, 3]),
14367 dtype: DType::F32,
14368 });
14369 let out_id = g.push(Node {
14370 op: Op::Triu { diagonal: 0 },
14371 inputs: vec![in_id],
14372 shape: Shape::from_dims(&[3, 3]),
14373 dtype: DType::F32,
14374 });
14375 g.set_target_backend(out_id, BackendId::Cpu);
14376 (in_id, out_id)
14377 };
14378 let mut inputs = StorageCache::new();
14379 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14380 let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
14381 let guard = result_arc.read().unwrap();
14382 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
14383 panic!("expected Cpu storage");
14384 };
14385 let out: &[f32] = c.as_slice().unwrap();
14386 assert_eq!(out, &[1.0, 2.0, 3.0, 0.0, 5.0, 6.0, 0.0, 0.0, 9.0,]);
14387 }
14388
14389 #[test]
14392 fn pipelined_realize_tril_3x3_diag0_causal_mask() {
14393 let storage =
14394 fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
14395 let graph = Arc::new(RwLock::new(Graph::new()));
14396 let (in_id, out_id) = {
14397 let mut g = graph.write().unwrap();
14398 let in_id = g.push(Node {
14399 op: Op::Const,
14400 inputs: vec![],
14401 shape: Shape::from_dims(&[3, 3]),
14402 dtype: DType::F32,
14403 });
14404 let out_id = g.push(Node {
14405 op: Op::Tril { diagonal: 0 },
14406 inputs: vec![in_id],
14407 shape: Shape::from_dims(&[3, 3]),
14408 dtype: DType::F32,
14409 });
14410 g.set_target_backend(out_id, BackendId::Cpu);
14411 (in_id, out_id)
14412 };
14413 let mut inputs = StorageCache::new();
14414 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14415 let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
14416 let guard = result_arc.read().unwrap();
14417 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
14418 panic!("expected Cpu storage");
14419 };
14420 let out: &[f32] = c.as_slice().unwrap();
14421 assert_eq!(out, &[1.0, 0.0, 0.0, 4.0, 5.0, 0.0, 7.0, 8.0, 9.0,]);
14422 }
14423
14424 #[test]
14431 fn pipelined_realize_log_softmax_last_dim() {
14432 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 3.0, 2.0, 1.0]);
14433 let graph = Arc::new(RwLock::new(Graph::new()));
14434 let (in_id, out_id) = {
14435 let mut g = graph.write().unwrap();
14436 let in_id = g.push(Node {
14437 op: Op::Const,
14438 inputs: vec![],
14439 shape: Shape::from_dims(&[2, 3]),
14440 dtype: DType::F32,
14441 });
14442 let out_id = g.push(Node {
14443 op: Op::LogSoftmaxLastDim,
14444 inputs: vec![in_id],
14445 shape: Shape::from_dims(&[2, 3]),
14446 dtype: DType::F32,
14447 });
14448 g.set_target_backend(out_id, BackendId::Cpu);
14449 (in_id, out_id)
14450 };
14451 let mut inputs = StorageCache::new();
14452 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
14453 let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
14454 let guard = result_arc.read().unwrap();
14455 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
14456 panic!("expected Cpu storage");
14457 };
14458 let out: &[f32] = c.as_slice().unwrap();
14459 let expected = [
14460 -2.407_606,
14461 -1.4076059,
14462 -0.40760595_f32,
14463 -0.40760595,
14464 -1.4076059,
14465 -2.407_606,
14466 ];
14467 for (a, b) in out.iter().zip(expected.iter()) {
14468 assert!((a - b).abs() < 1e-5, "log_softmax mismatch: {a} vs {b}");
14469 }
14470 for row in out.chunks(3) {
14472 let sum: f32 = row.iter().map(|x| x.exp()).sum();
14473 assert!(
14474 (sum - 1.0).abs() < 1e-5,
14475 "softmax(log_softmax) row sum != 1: {sum}"
14476 );
14477 }
14478 }
14479
14480 #[test]
14494 fn pipelined_realize_merges_op_release_side_effect_root() {
14495 let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
14496 let b = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
14497 let c = fuel_memory::from_slice_cpu(&[100.0_f32, 200.0]);
14498 let graph = Arc::new(RwLock::new(Graph::new()));
14499 let (a_id, b_id, c_id, add_id, release_id) = {
14500 let mut g = graph.write().unwrap();
14501 let a = g.push(Node {
14502 op: Op::Const,
14503 inputs: vec![],
14504 shape: Shape::from_dims(&[2]),
14505 dtype: DType::F32,
14506 });
14507 let b = g.push(Node {
14508 op: Op::Const,
14509 inputs: vec![],
14510 shape: Shape::from_dims(&[2]),
14511 dtype: DType::F32,
14512 });
14513 let c = g.push(Node {
14516 op: Op::Const,
14517 inputs: vec![],
14518 shape: Shape::from_dims(&[2]),
14519 dtype: DType::F32,
14520 });
14521 let add = g.push(Node {
14522 op: Op::Add,
14523 inputs: vec![a, b],
14524 shape: Shape::from_dims(&[2]),
14525 dtype: DType::F32,
14526 });
14527 let release = g.push(Node {
14531 op: Op::Release,
14532 inputs: vec![c],
14533 shape: Shape::from_dims(&[0]),
14534 dtype: DType::F32,
14535 });
14536 g.set_target_backend(add, BackendId::Cpu);
14537 g.set_target_backend(release, BackendId::Cpu);
14538 g.add_side_effect_root(release);
14539 (a, b, c, add, release)
14540 };
14541 let mut inputs = StorageCache::new();
14542 inputs.insert(a_id, Arc::new(RwLock::new(a)));
14543 inputs.insert(b_id, Arc::new(RwLock::new(b)));
14544 let c_arc = Arc::new(RwLock::new(c));
14545 let c_arc_external = Arc::clone(&c_arc);
14546 inputs.insert(c_id, c_arc);
14547
14548 let out = PipelinedExecutor::realize_many(graph, &[add_id], inputs).expect("realize_many");
14551 assert_eq!(out.len(), 1);
14552 let add_guard = out[0].0.read().unwrap();
14553 let fuel_memory::BackendStorage::Cpu(arr) = &add_guard.inner else {
14554 panic!()
14555 };
14556 assert_eq!(arr.as_slice::<f32>().unwrap(), &[11.0, 22.0]);
14557
14558 assert_eq!(
14563 Arc::strong_count(&c_arc_external),
14564 1,
14565 "Op::Release should evict the cache's Arc to its source; \
14566 only the test's external clone should remain",
14567 );
14568 let _ = release_id;
14569 }
14570
14571 #[test]
14575 fn compile_one_emits_release_marker_with_destructive_input() {
14576 let graph_rc = Arc::new(RwLock::new(Graph::new()));
14577 let (src_id, release_id) = {
14578 let mut g = graph_rc.write().unwrap();
14579 let src = g.push(Node {
14580 op: Op::Const,
14581 inputs: vec![],
14582 shape: Shape::from_dims(&[2]),
14583 dtype: DType::F32,
14584 });
14585 let release = g.push(Node {
14586 op: Op::Release,
14587 inputs: vec![src],
14588 shape: Shape::from_dims(&[0]),
14589 dtype: DType::F32,
14590 });
14591 g.set_target_backend(release, BackendId::Cpu);
14592 (src, release)
14593 };
14594 let g = graph_rc.read().unwrap();
14595 let bindings = crate::dispatch::global_bindings();
14596 let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
14597 let item = compile_one(
14598 &g,
14599 release_id,
14600 &mut layout_cache,
14601 &bindings,
14602 &SymEnv::default(),
14603 )
14604 .expect("compile_one Op::Release");
14605 assert!(matches!(item.kind, WorkItemKind::ReleaseMarker));
14606 assert_eq!(item.destructive_input, Some(0));
14607 assert_eq!(item.inputs, vec![src_id]);
14608 assert_eq!(
14609 item.elem_count, 0,
14610 "Release output is the zero-element marker"
14611 );
14612 }
14613
14614 #[test]
14621 fn compile_one_emits_move_with_destructive_input() {
14622 let graph_rc = Arc::new(RwLock::new(Graph::new()));
14623 let (src_id, move_id) = {
14624 let mut g = graph_rc.write().unwrap();
14625 let src = g.push(Node {
14626 op: Op::Const,
14627 inputs: vec![],
14628 shape: Shape::from_dims(&[2]),
14629 dtype: DType::F32,
14630 });
14631 let mv = g.push(Node {
14632 op: Op::Move {
14633 target: DeviceLocation::Cpu,
14634 },
14635 inputs: vec![src],
14636 shape: Shape::from_dims(&[2]),
14637 dtype: DType::F32,
14638 });
14639 g.set_target_backend(mv, BackendId::Cpu);
14640 (src, mv)
14641 };
14642 let g = graph_rc.read().unwrap();
14643 let bindings = crate::dispatch::global_bindings();
14644 let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
14645 let item = compile_one(
14646 &g,
14647 move_id,
14648 &mut layout_cache,
14649 &bindings,
14650 &SymEnv::default(),
14651 )
14652 .expect("compile_one Op::Move");
14653 assert!(matches!(
14654 item.kind,
14655 WorkItemKind::Move {
14656 target_location: DeviceLocation::Cpu
14657 },
14658 ));
14659 assert_eq!(item.destructive_input, Some(0), "Move destroys its source");
14660 assert_eq!(item.inputs, vec![src_id]);
14661 assert!(
14662 item.compiled.is_some(),
14663 "Move resolves a transfer kernel at (OpKind::Copy, [dt, dt], source backend)",
14664 );
14665 }
14666
14667 #[test]
14675 fn compile_one_resolves_runtime_fused_from_the_sidecar() {
14676 fn sentinel_kernel(
14677 _inputs: &[Arc<RwLock<fuel_memory::Storage>>],
14678 _outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
14679 _layouts: &[Layout],
14680 _params: &crate::kernel::OpParams,
14681 ) -> Result<()> {
14682 Ok(())
14683 }
14684 let region = fuel_graph::jit::PatternNode::Op {
14687 op: fuel_graph::jit::OpTag::Exp,
14688 attrs: fuel_graph::jit::OpAttrs::default(),
14689 operands: vec![fuel_graph::jit::PatternNode::Op {
14690 op: fuel_graph::jit::OpTag::Sub,
14691 attrs: fuel_graph::jit::OpAttrs::default(),
14692 operands: vec![
14693 fuel_graph::jit::PatternNode::Bind { index: 0 },
14694 fuel_graph::jit::PatternNode::Bind { index: 1 },
14695 ],
14696 }],
14697 };
14698 let rid = crate::runtime_fused_kernels::adopt_runtime_fused(
14699 "test::compile_one::exp_sub",
14700 region,
14701 sentinel_kernel as crate::kernel::KernelRef,
14702 vec![DType::F32, DType::F32, DType::F32],
14703 BackendId::Cpu,
14704 )
14705 .expect("registrable region");
14706
14707 let graph_rc = Arc::new(RwLock::new(Graph::new()));
14708 let fused_id = {
14709 let mut g = graph_rc.write().unwrap();
14710 let s = Shape::from_dims(&[4]);
14711 let a = g.push(Node {
14712 op: Op::Const,
14713 inputs: vec![],
14714 shape: s.clone(),
14715 dtype: DType::F32,
14716 });
14717 let b = g.push(Node {
14718 op: Op::Const,
14719 inputs: vec![],
14720 shape: s.clone(),
14721 dtype: DType::F32,
14722 });
14723 let fused = g.push(Node {
14724 op: Op::Fused(
14725 rid,
14726 fuel_graph::registry::FusedOpParams::Runtime {
14727 scalars: Vec::new(),
14728 },
14729 ),
14730 inputs: vec![a, b],
14731 shape: s,
14732 dtype: DType::F32,
14733 });
14734 g.set_target_backend(fused, BackendId::Cpu);
14736 fused
14737 };
14738 let g = graph_rc.read().unwrap();
14739 let bindings = crate::dispatch::global_bindings();
14740 let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
14741 let item = compile_one(
14742 &g,
14743 fused_id,
14744 &mut layout_cache,
14745 &bindings,
14746 &SymEnv::default(),
14747 )
14748 .expect("compile_one runtime fused op");
14749 assert!(matches!(item.kind, WorkItemKind::Kernel));
14750 let compiled = item.compiled.expect("resolved a kernel");
14751 assert_eq!(compiled.op, fuel_ir::dispatch::OpKind::RuntimeFused);
14752 assert_eq!(compiled.backend, BackendId::Cpu);
14753 assert!(
14754 std::ptr::fn_addr_eq(compiled.kernel, sentinel_kernel as crate::kernel::KernelRef),
14755 "the kernel pointer is exactly the adopted one",
14756 );
14757 assert!(
14758 layout_cache.contains_key(&fused_id),
14759 "output layout cached for consumers"
14760 );
14761 }
14762
14763 #[test]
14767 fn compile_one_rejects_runtime_fused_dtype_mismatch() {
14768 fn noop_kernel(
14769 _inputs: &[Arc<RwLock<fuel_memory::Storage>>],
14770 _outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
14771 _layouts: &[Layout],
14772 _params: &crate::kernel::OpParams,
14773 ) -> Result<()> {
14774 Ok(())
14775 }
14776 let region = fuel_graph::jit::PatternNode::Op {
14777 op: fuel_graph::jit::OpTag::Log,
14778 attrs: fuel_graph::jit::OpAttrs::default(),
14779 operands: vec![fuel_graph::jit::PatternNode::Op {
14780 op: fuel_graph::jit::OpTag::Maximum,
14781 attrs: fuel_graph::jit::OpAttrs::default(),
14782 operands: vec![
14783 fuel_graph::jit::PatternNode::Bind { index: 0 },
14784 fuel_graph::jit::PatternNode::Bind { index: 1 },
14785 ],
14786 }],
14787 };
14788 let rid = crate::runtime_fused_kernels::adopt_runtime_fused(
14790 "test::compile_one::log_maximum",
14791 region,
14792 noop_kernel as crate::kernel::KernelRef,
14793 vec![DType::F32, DType::F32, DType::F32],
14794 BackendId::Cpu,
14795 )
14796 .expect("registrable region");
14797
14798 let graph_rc = Arc::new(RwLock::new(Graph::new()));
14799 let fused_id = {
14800 let mut g = graph_rc.write().unwrap();
14801 let s = Shape::from_dims(&[4]);
14802 let a = g.push(Node {
14804 op: Op::Const,
14805 inputs: vec![],
14806 shape: s.clone(),
14807 dtype: DType::F64,
14808 });
14809 let b = g.push(Node {
14810 op: Op::Const,
14811 inputs: vec![],
14812 shape: s.clone(),
14813 dtype: DType::F64,
14814 });
14815 let fused = g.push(Node {
14816 op: Op::Fused(
14817 rid,
14818 fuel_graph::registry::FusedOpParams::Runtime {
14819 scalars: Vec::new(),
14820 },
14821 ),
14822 inputs: vec![a, b],
14823 shape: s,
14824 dtype: DType::F64,
14825 });
14826 g.set_target_backend(fused, BackendId::Cpu);
14827 fused
14828 };
14829 let g = graph_rc.read().unwrap();
14830 let bindings = crate::dispatch::global_bindings();
14831 let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
14832 let err = match compile_one(
14833 &g,
14834 fused_id,
14835 &mut layout_cache,
14836 &bindings,
14837 &SymEnv::default(),
14838 ) {
14839 Ok(_) => panic!("dtype mismatch must be a typed error"),
14840 Err(e) => e,
14841 };
14842 assert!(
14846 matches!(
14847 err,
14848 fuel_ir::Error::NoBackendForOp {
14849 op: fuel_ir::dispatch::OpKind::RuntimeFused,
14850 ..
14851 }
14852 ),
14853 "a dtype mismatch is an honest binding miss: {err}",
14854 );
14855 }
14856
14857 #[test]
14862 fn compile_one_maps_runtime_scalars_to_jit_scalars() {
14863 fn noop_kernel(
14864 _inputs: &[Arc<RwLock<fuel_memory::Storage>>],
14865 _outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
14866 _layouts: &[Layout],
14867 _params: &crate::kernel::OpParams,
14868 ) -> Result<()> {
14869 Ok(())
14870 }
14871 let region = fuel_graph::jit::PatternNode::Op {
14873 op: fuel_graph::jit::OpTag::Floor,
14874 attrs: fuel_graph::jit::OpAttrs::default(),
14875 operands: vec![fuel_graph::jit::PatternNode::Op {
14876 op: fuel_graph::jit::OpTag::MulScalar,
14877 attrs: fuel_graph::jit::OpAttrs::default(), operands: vec![fuel_graph::jit::PatternNode::Bind { index: 0 }],
14879 }],
14880 };
14881 let rid = crate::runtime_fused_kernels::adopt_runtime_fused(
14882 "test::compile_one::floor_mul_scalar",
14883 region,
14884 noop_kernel as crate::kernel::KernelRef,
14885 vec![DType::F32, DType::F32],
14886 BackendId::Cpu,
14887 )
14888 .expect("slot template registrable");
14889
14890 let graph_rc = Arc::new(RwLock::new(Graph::new()));
14891 let fused_id = {
14892 let mut g = graph_rc.write().unwrap();
14893 let s = Shape::from_dims(&[4]);
14894 let x = g.push(Node {
14895 op: Op::Const,
14896 inputs: vec![],
14897 shape: s.clone(),
14898 dtype: DType::F32,
14899 });
14900 let fused = g.push(Node {
14901 op: Op::Fused(
14902 rid,
14903 fuel_graph::registry::FusedOpParams::Runtime {
14904 scalars: vec![3.5, -1.25],
14905 },
14906 ),
14907 inputs: vec![x],
14908 shape: s,
14909 dtype: DType::F32,
14910 });
14911 g.set_target_backend(fused, BackendId::Cpu);
14912 fused
14913 };
14914 let g = graph_rc.read().unwrap();
14915 let bindings = crate::dispatch::global_bindings();
14916 let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
14917 let item = compile_one(
14918 &g,
14919 fused_id,
14920 &mut layout_cache,
14921 &bindings,
14922 &SymEnv::default(),
14923 )
14924 .expect("compile_one runtime fused op with scalars");
14925 let compiled = item.compiled.expect("resolved a kernel");
14926 assert!(
14927 matches!(
14928 &compiled.op_params,
14929 crate::kernel::OpParams::JitScalars { scalars } if scalars == &vec![3.5, -1.25]
14930 ),
14931 "Runtime scalars ride into OpParams::JitScalars, got {:?}",
14932 compiled.op_params,
14933 );
14934 }
14935
14936 #[test]
14943 fn pipelined_op_move_same_device_is_plain_copy() {
14944 let src = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
14945 let graph = Arc::new(RwLock::new(Graph::new()));
14946 let (src_id, move_id) = {
14947 let mut g = graph.write().unwrap();
14948 let src = g.push(Node {
14949 op: Op::Const,
14950 inputs: vec![],
14951 shape: Shape::from_dims(&[4]),
14952 dtype: DType::F32,
14953 });
14954 let mv = g.push(Node {
14955 op: Op::Move {
14956 target: DeviceLocation::Cpu,
14957 },
14958 inputs: vec![src],
14959 shape: Shape::from_dims(&[4]),
14960 dtype: DType::F32,
14961 });
14962 g.set_target_backend(mv, BackendId::Cpu);
14963 (src, mv)
14964 };
14965 let src_arc = Arc::new(RwLock::new(src));
14966 let src_arc_external = Arc::clone(&src_arc);
14967 let mut inputs = StorageCache::new();
14968 inputs.insert(src_id, src_arc);
14969
14970 let (out, layout) =
14971 PipelinedExecutor::realize(graph, move_id, inputs).expect("realize Op::Move");
14972 let guard = out.read().unwrap();
14973 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
14974 panic!()
14975 };
14976 assert_eq!(
14977 c.as_slice::<f32>().unwrap(),
14978 &[1.0, 2.0, 3.0, 4.0],
14979 "Op::Move output should carry input's data to target",
14980 );
14981 assert!(layout.is_contiguous());
14982 drop(guard);
14983 assert!(
14984 !Arc::ptr_eq(&out, &src_arc_external),
14985 "same-device Move is a plain copy — fresh storage, not an alias",
14986 );
14987 assert_eq!(
14991 Arc::strong_count(&src_arc_external),
14992 1,
14993 "Op::Move should evict the cache's Arc to its source",
14994 );
14995 }
14996
14997 #[test]
15010 fn pipelined_op_move_multi_consumer_source_not_stranded() {
15011 let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
15012 let graph = Arc::new(RwLock::new(Graph::new()));
15013 let (a_id, out_id) = {
15014 let mut g = graph.write().unwrap();
15015 let a = g.push(Node {
15016 op: Op::Const,
15017 inputs: vec![],
15018 shape: Shape::from_dims(&[4]),
15019 dtype: DType::F32,
15020 });
15021 let b = g.push(Node {
15022 op: Op::Relu,
15023 inputs: vec![a],
15024 shape: Shape::from_dims(&[4]),
15025 dtype: DType::F32,
15026 });
15027 let m = g.push(Node {
15028 op: Op::Move {
15029 target: DeviceLocation::Cpu,
15030 },
15031 inputs: vec![a],
15032 shape: Shape::from_dims(&[4]),
15033 dtype: DType::F32,
15034 });
15035 let out = g.push(Node {
15036 op: Op::Add,
15037 inputs: vec![b, m],
15038 shape: Shape::from_dims(&[4]),
15039 dtype: DType::F32,
15040 });
15041 g.set_target_backend(b, BackendId::Cpu);
15042 g.set_target_backend(m, BackendId::Cpu);
15043 g.set_target_backend(out, BackendId::Cpu);
15044 (a, out)
15045 };
15046 let mut inputs = StorageCache::new();
15047 inputs.insert(a_id, Arc::new(RwLock::new(a)));
15048
15049 let (out, _) = PipelinedExecutor::realize(graph, out_id, inputs)
15050 .expect("realize add(relu(a), move(a))");
15051 let guard = out.read().unwrap();
15052 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15053 panic!()
15054 };
15055 assert_eq!(
15056 c.as_slice::<f32>().unwrap(),
15057 &[2.0, 4.0, 6.0, 8.0],
15058 "relu(a) must run before move(a) evicts a",
15059 );
15060 }
15061
15062 #[test]
15067 fn pipelined_op_move_source_in_target_set_not_evicted() {
15068 let a = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0]);
15069 let graph = Arc::new(RwLock::new(Graph::new()));
15070 let (a_id, move_id) = {
15071 let mut g = graph.write().unwrap();
15072 let a = g.push(Node {
15073 op: Op::Const,
15074 inputs: vec![],
15075 shape: Shape::from_dims(&[2]),
15076 dtype: DType::F32,
15077 });
15078 let mv = g.push(Node {
15079 op: Op::Move {
15080 target: DeviceLocation::Cpu,
15081 },
15082 inputs: vec![a],
15083 shape: Shape::from_dims(&[2]),
15084 dtype: DType::F32,
15085 });
15086 g.set_target_backend(mv, BackendId::Cpu);
15087 (a, mv)
15088 };
15089 let a_arc = Arc::new(RwLock::new(a));
15090 let a_arc_external = Arc::clone(&a_arc);
15091 let mut inputs = StorageCache::new();
15092 inputs.insert(a_id, a_arc);
15093
15094 let out = PipelinedExecutor::realize_many(graph, &[move_id, a_id], inputs)
15095 .expect("realize_many [move, source]");
15096 assert_eq!(out.len(), 2);
15097 let mv_guard = out[0].0.read().unwrap();
15098 let fuel_memory::BackendStorage::Cpu(c) = &mv_guard.inner else {
15099 panic!()
15100 };
15101 assert_eq!(c.as_slice::<f32>().unwrap(), &[5.0, 6.0]);
15102 assert!(
15103 Arc::ptr_eq(&out[1].0, &a_arc_external),
15104 "requested source must survive the Move's destructive cleanup",
15105 );
15106 }
15107
15108 #[test]
15113 fn pipelined_realize_many_empty_targets() {
15114 let graph = Arc::new(RwLock::new(Graph::new()));
15115 let inputs = StorageCache::new();
15116 let out = PipelinedExecutor::realize_many(graph, &[], inputs).expect("realize_many");
15117 assert!(out.is_empty());
15118 }
15119
15120 #[test]
15124 fn pipelined_realize_many_two_independent_targets() {
15125 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
15126 let rhs = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
15127 let graph = Arc::new(RwLock::new(Graph::new()));
15128 let (lhs_id, rhs_id, add_id, mul_id) = {
15129 let mut g = graph.write().unwrap();
15130 let lhs = g.push(Node {
15131 op: Op::Const,
15132 inputs: vec![],
15133 shape: Shape::from_dims(&[3]),
15134 dtype: DType::F32,
15135 });
15136 let rhs = g.push(Node {
15137 op: Op::Const,
15138 inputs: vec![],
15139 shape: Shape::from_dims(&[3]),
15140 dtype: DType::F32,
15141 });
15142 let add = g.push(Node {
15143 op: Op::Add,
15144 inputs: vec![lhs, rhs],
15145 shape: Shape::from_dims(&[3]),
15146 dtype: DType::F32,
15147 });
15148 let mul = g.push(Node {
15149 op: Op::Mul,
15150 inputs: vec![lhs, rhs],
15151 shape: Shape::from_dims(&[3]),
15152 dtype: DType::F32,
15153 });
15154 g.set_target_backend(add, BackendId::Cpu);
15155 g.set_target_backend(mul, BackendId::Cpu);
15156 (lhs, rhs, add, mul)
15157 };
15158 let mut inputs = StorageCache::new();
15159 inputs.insert(lhs_id, Arc::new(RwLock::new(lhs)));
15160 inputs.insert(rhs_id, Arc::new(RwLock::new(rhs)));
15161
15162 let out = PipelinedExecutor::realize_many(graph, &[mul_id, add_id], inputs)
15165 .expect("realize_many");
15166 assert_eq!(out.len(), 2);
15167
15168 let mul_guard = out[0].0.read().unwrap();
15169 let fuel_memory::BackendStorage::Cpu(c) = &mul_guard.inner else {
15170 panic!()
15171 };
15172 assert_eq!(c.as_slice::<f32>().unwrap(), &[10.0, 40.0, 90.0]);
15173
15174 let add_guard = out[1].0.read().unwrap();
15175 let fuel_memory::BackendStorage::Cpu(c) = &add_guard.inner else {
15176 panic!()
15177 };
15178 assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 22.0, 33.0]);
15179 }
15180
15181 #[test]
15184 fn pipelined_realize_many_duplicate_target_returns_same_arc() {
15185 let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
15186 let graph = Arc::new(RwLock::new(Graph::new()));
15187 let (in_id, neg_id) = {
15188 let mut g = graph.write().unwrap();
15189 let in_id = g.push(Node {
15190 op: Op::Const,
15191 inputs: vec![],
15192 shape: Shape::from_dims(&[3]),
15193 dtype: DType::F32,
15194 });
15195 let neg = g.push(Node {
15196 op: Op::Neg,
15197 inputs: vec![in_id],
15198 shape: Shape::from_dims(&[3]),
15199 dtype: DType::F32,
15200 });
15201 g.set_target_backend(neg, BackendId::Cpu);
15202 (in_id, neg)
15203 };
15204 let mut inputs = StorageCache::new();
15205 inputs.insert(in_id, Arc::new(RwLock::new(storage)));
15206
15207 let out = PipelinedExecutor::realize_many(graph, &[neg_id, neg_id], inputs)
15208 .expect("realize_many");
15209 assert_eq!(out.len(), 2);
15210 assert!(
15211 Arc::ptr_eq(&out[0].0, &out[1].0),
15212 "duplicate target should share the same storage Arc",
15213 );
15214 }
15215
15216 #[test]
15220 fn pipelined_realize_many_shared_subgraph_evaluates_once() {
15221 let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
15222 let b_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
15223 let graph = Arc::new(RwLock::new(Graph::new()));
15224 let (a_id, b_id, add_id, mul_id) = {
15225 let mut g = graph.write().unwrap();
15226 let a = g.push(Node {
15227 op: Op::Const,
15228 inputs: vec![],
15229 shape: Shape::from_dims(&[3]),
15230 dtype: DType::F32,
15231 });
15232 let b = g.push(Node {
15233 op: Op::Const,
15234 inputs: vec![],
15235 shape: Shape::from_dims(&[3]),
15236 dtype: DType::F32,
15237 });
15238 let add = g.push(Node {
15239 op: Op::Add,
15240 inputs: vec![a, b],
15241 shape: Shape::from_dims(&[3]),
15242 dtype: DType::F32,
15243 });
15244 let mul = g.push(Node {
15245 op: Op::Mul,
15246 inputs: vec![add, a],
15247 shape: Shape::from_dims(&[3]),
15248 dtype: DType::F32,
15249 });
15250 g.set_target_backend(add, BackendId::Cpu);
15251 g.set_target_backend(mul, BackendId::Cpu);
15252 (a, b, add, mul)
15253 };
15254 let mut inputs = StorageCache::new();
15255 inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
15256 inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
15257
15258 let out = PipelinedExecutor::realize_many(graph, &[add_id, mul_id], inputs)
15259 .expect("realize_many");
15260 assert_eq!(out.len(), 2);
15261
15262 let add_guard = out[0].0.read().unwrap();
15263 let fuel_memory::BackendStorage::Cpu(c) = &add_guard.inner else {
15264 panic!()
15265 };
15266 assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 22.0, 33.0]);
15267
15268 let mul_guard = out[1].0.read().unwrap();
15269 let fuel_memory::BackendStorage::Cpu(c) = &mul_guard.inner else {
15270 panic!()
15271 };
15272 assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 44.0, 99.0]);
15274 }
15275
15276 #[test]
15280 fn pipelined_realize_masked_fill_attention_pattern() {
15281 let x_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
15282 let mut mask_storage = fuel_memory::from_slice_cpu(&[0u8, 1, 0, 1]);
15283 mask_storage.dtype = DType::Bool; let graph = Arc::new(RwLock::new(Graph::new()));
15285 let (x_id, mask_id, out_id) = {
15286 let mut g = graph.write().unwrap();
15287 let x = g.push(Node {
15288 op: Op::Const,
15289 inputs: vec![],
15290 shape: Shape::from_dims(&[4]),
15291 dtype: DType::F32,
15292 });
15293 let mask = g.push(Node {
15294 op: Op::Const,
15295 inputs: vec![],
15296 shape: Shape::from_dims(&[4]),
15297 dtype: DType::Bool,
15298 });
15299 let out = g.push(Node {
15300 op: Op::MaskedFill {
15301 value: fuel_ir::Scalar::F32(-1000.0),
15302 },
15303 inputs: vec![x, mask],
15304 shape: Shape::from_dims(&[4]),
15305 dtype: DType::F32,
15306 });
15307 g.set_target_backend(out, BackendId::Cpu);
15308 (x, mask, out)
15309 };
15310 let mut inputs = StorageCache::new();
15311 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
15312 inputs.insert(mask_id, Arc::new(RwLock::new(mask_storage)));
15313 let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
15314 let guard = result_arc.read().unwrap();
15315 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15316 panic!("expected Cpu storage");
15317 };
15318 let out: &[f32] = c.as_slice().unwrap();
15319 assert_eq!(out, &[1.0, -1000.0, 3.0, -1000.0]);
15320 }
15321
15322 #[test]
15328 fn pipelined_realize_write_slice_kv_cache_append() {
15329 let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 24]);
15330 let src_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
15331 let graph = Arc::new(RwLock::new(Graph::new()));
15332 let (dest_id, src_id, ws_id) = {
15333 let mut g = graph.write().unwrap();
15334 let dest = g.push(Node {
15335 op: Op::Const,
15336 inputs: vec![],
15337 shape: Shape::from_dims(&[4, 3, 2]),
15338 dtype: DType::F32,
15339 });
15340 let src = g.push(Node {
15341 op: Op::Const,
15342 inputs: vec![],
15343 shape: Shape::from_dims(&[1, 3, 2]),
15344 dtype: DType::F32,
15345 });
15346 let ws = g.push(Node {
15347 op: Op::WriteSlice {
15348 ranges: vec![(2, 3), (0, 3), (0, 2)],
15349 dyn_offset: None,
15350 },
15351 inputs: vec![dest, src],
15352 shape: Shape::from_dims(&[4, 3, 2]), dtype: DType::F32,
15354 });
15355 g.set_target_backend(ws, BackendId::Cpu);
15356 (dest, src, ws)
15357 };
15358 let mut inputs = StorageCache::new();
15359 inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
15360 inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
15361 let (result_arc, _) = PipelinedExecutor::realize(graph, ws_id, inputs).expect("realize");
15362 let guard = result_arc.read().unwrap();
15363 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15364 panic!("expected Cpu storage");
15365 };
15366 let out: &[f32] = c.as_slice().unwrap();
15367 let expected = [
15368 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0,
15369 6.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
15371 ];
15372 assert_eq!(out, &expected);
15373 }
15374
15375 #[test]
15380 fn pipelined_realize_write_slice_aliases_dest_storage() {
15381 let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
15382 let src_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
15383 let graph = Arc::new(RwLock::new(Graph::new()));
15384 let (dest_id, src_id, ws_id, dest_arc_external) = {
15385 let mut g = graph.write().unwrap();
15386 let dest = g.push(Node {
15387 op: Op::Const,
15388 inputs: vec![],
15389 shape: Shape::from_dims(&[3, 2]),
15390 dtype: DType::F32,
15391 });
15392 let src = g.push(Node {
15393 op: Op::Const,
15394 inputs: vec![],
15395 shape: Shape::from_dims(&[1, 2]),
15396 dtype: DType::F32,
15397 });
15398 let ws = g.push(Node {
15399 op: Op::WriteSlice {
15400 ranges: vec![(1, 2), (0, 2)],
15401 dyn_offset: None,
15402 },
15403 inputs: vec![dest, src],
15404 shape: Shape::from_dims(&[3, 2]),
15405 dtype: DType::F32,
15406 });
15407 g.set_target_backend(ws, BackendId::Cpu);
15408 let arc = Arc::new(RwLock::new(dest_storage));
15409 (dest, src, ws, arc)
15410 };
15411 let mut inputs = StorageCache::new();
15412 inputs.insert(dest_id, Arc::clone(&dest_arc_external));
15413 inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
15414 let (result_arc, _) = PipelinedExecutor::realize(graph, ws_id, inputs).expect("realize");
15415 assert!(
15419 Arc::ptr_eq(&result_arc, &dest_arc_external),
15420 "WriteSlice output Arc must be the same Arc as the destination's"
15421 );
15422 let guard = dest_arc_external.read().unwrap();
15423 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15424 panic!("expected Cpu storage");
15425 };
15426 let out: &[f32] = c.as_slice().unwrap();
15427 assert_eq!(out, &[0.0, 0.0, 10.0, 20.0, 0.0, 0.0]);
15428 }
15429
15430 #[test]
15438 fn pipelined_realize_write_slice_dynamic_offset() {
15439 use fuel_ir::{DynScalar, SymId};
15440 let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
15441 let src_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0]);
15442 let graph = Arc::new(RwLock::new(Graph::new()));
15443 let sym = SymId(0);
15444 let (dest_id, src_id, ws_id) = {
15445 let mut g = graph.write().unwrap();
15446 let dest = g.push(Node {
15447 op: Op::Const,
15448 inputs: vec![],
15449 shape: Shape::from_dims(&[6]),
15450 dtype: DType::F32,
15451 });
15452 let src = g.push(Node {
15453 op: Op::Const,
15454 inputs: vec![],
15455 shape: Shape::from_dims(&[2]),
15456 dtype: DType::F32,
15457 });
15458 let ws = g.push(Node {
15459 op: Op::WriteSlice {
15462 ranges: vec![(0, 2)],
15463 dyn_offset: Some((0, DynScalar::Sym(sym))),
15464 },
15465 inputs: vec![dest, src],
15466 shape: Shape::from_dims(&[6]),
15467 dtype: DType::F32,
15468 });
15469 g.set_target_backend(ws, BackendId::Cpu);
15470 (dest, src, ws)
15471 };
15472 let mut inputs = StorageCache::new();
15473 inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
15474 inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
15475
15476 let mut env = SymEnv::new();
15478 env.bind(sym, 3).unwrap();
15479
15480 let (result_arc, _) = PipelinedExecutor::realize_with_env(graph, ws_id, inputs, env)
15481 .expect("realize_with_env");
15482 let guard = result_arc.read().unwrap();
15483 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15484 panic!("expected Cpu storage");
15485 };
15486 let out: &[f32] = c.as_slice().unwrap();
15487 assert_eq!(
15488 out,
15489 &[0.0, 0.0, 0.0, 7.0, 8.0, 0.0],
15490 "slab must land at the SymEnv-bound offset 3, not the static placeholder 0",
15491 );
15492 }
15493
15494 #[test]
15501 fn pipelined_realize_nonzero_indices_publishes_count() {
15502 use fuel_ir::SymId;
15503 use fuel_ir::storage::{OutputViewSpec, compose_bundle};
15504 let x_storage = fuel_memory::from_slice_cpu(&[0.0_f32, 1.0, 0.0, 1.0, 1.0, 0.0]);
15506 let graph = Arc::new(RwLock::new(Graph::new()));
15507 let count_sym = SymId(0);
15508 let (x_id, nzi_id) = {
15509 let mut g = graph.write().unwrap();
15510 let x = g.push(Node {
15511 op: Op::Const,
15512 inputs: vec![],
15513 shape: Shape::from_dims(&[2, 3]),
15514 dtype: DType::F32,
15515 });
15516 let capacity = 6usize;
15517 let idx_shape = Shape::from_dims(&[capacity]);
15518 let cnt_shape = Shape::from_dims(&[1]);
15519 let specs = vec![
15520 OutputViewSpec {
15521 dtype: DType::U32,
15522 shape: idx_shape.clone(),
15523 layout: Layout::contiguous(idx_shape.clone()),
15524 name: Some("indices"),
15525 },
15526 OutputViewSpec {
15527 dtype: DType::U32,
15528 shape: cnt_shape.clone(),
15529 layout: Layout::contiguous(cnt_shape),
15530 name: Some("count"),
15531 },
15532 ];
15533 let (_bytes, views) = compose_bundle(&specs).unwrap();
15534 let nzi = g.push(Node {
15535 op: Op::NonZeroIndices { count_sym },
15536 inputs: vec![x],
15537 shape: idx_shape,
15538 dtype: DType::U32,
15539 });
15540 g.set_output_views(nzi, Arc::from(views.into_boxed_slice()))
15541 .unwrap();
15542 g.set_target_backend(nzi, BackendId::Cpu);
15543 (x, nzi)
15544 };
15545 let mut inputs = StorageCache::new();
15546 inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
15547
15548 let (_arc, _layout, produced) =
15549 PipelinedExecutor::realize_with_env_producing(graph, nzi_id, inputs, SymEnv::new())
15550 .expect("realize_with_env_producing");
15551
15552 assert_eq!(
15553 produced.get(count_sym),
15554 Some(3),
15555 "the executor must publish the data-determined nonzero count into the SymEnv",
15556 );
15557 }
15558
15559 #[test]
15562 fn pipelined_realize_write_slice_dynamic_offset_unbound_errors() {
15563 use fuel_ir::{DynScalar, SymId};
15564 let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
15565 let src_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0]);
15566 let graph = Arc::new(RwLock::new(Graph::new()));
15567 let (dest_id, src_id, ws_id) = {
15568 let mut g = graph.write().unwrap();
15569 let dest = g.push(Node {
15570 op: Op::Const,
15571 inputs: vec![],
15572 shape: Shape::from_dims(&[6]),
15573 dtype: DType::F32,
15574 });
15575 let src = g.push(Node {
15576 op: Op::Const,
15577 inputs: vec![],
15578 shape: Shape::from_dims(&[2]),
15579 dtype: DType::F32,
15580 });
15581 let ws = g.push(Node {
15582 op: Op::WriteSlice {
15583 ranges: vec![(0, 2)],
15584 dyn_offset: Some((0, DynScalar::Sym(SymId(0)))),
15585 },
15586 inputs: vec![dest, src],
15587 shape: Shape::from_dims(&[6]),
15588 dtype: DType::F32,
15589 });
15590 g.set_target_backend(ws, BackendId::Cpu);
15591 (dest, src, ws)
15592 };
15593 let mut inputs = StorageCache::new();
15594 inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
15595 inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
15596 let result = PipelinedExecutor::realize_with_env(graph, ws_id, inputs, SymEnv::new());
15598 assert!(
15599 result.is_err(),
15600 "unbound dyn_offset symbol must surface a typed error, not a panic",
15601 );
15602 }
15603
15604 #[test]
15607 fn pipelined_realize_write_slice_dynamic_offset_out_of_capacity_errors() {
15608 use fuel_ir::{DynScalar, SymId};
15609 let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
15610 let src_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0]);
15611 let graph = Arc::new(RwLock::new(Graph::new()));
15612 let sym = SymId(0);
15613 let (dest_id, src_id, ws_id) = {
15614 let mut g = graph.write().unwrap();
15615 let dest = g.push(Node {
15616 op: Op::Const,
15617 inputs: vec![],
15618 shape: Shape::from_dims(&[6]),
15619 dtype: DType::F32,
15620 });
15621 let src = g.push(Node {
15622 op: Op::Const,
15623 inputs: vec![],
15624 shape: Shape::from_dims(&[2]),
15625 dtype: DType::F32,
15626 });
15627 let ws = g.push(Node {
15628 op: Op::WriteSlice {
15629 ranges: vec![(0, 2)],
15630 dyn_offset: Some((0, DynScalar::Sym(sym))),
15631 },
15632 inputs: vec![dest, src],
15633 shape: Shape::from_dims(&[6]),
15634 dtype: DType::F32,
15635 });
15636 g.set_target_backend(ws, BackendId::Cpu);
15637 (dest, src, ws)
15638 };
15639 let mut inputs = StorageCache::new();
15640 inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
15641 inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
15642 let mut env = SymEnv::new();
15644 env.bind(sym, 5).unwrap();
15645 let result = PipelinedExecutor::realize_with_env(graph, ws_id, inputs, env);
15646 assert!(
15647 result.is_err(),
15648 "resolved slab past destination capacity must surface a typed error",
15649 );
15650 }
15651
15652 #[test]
15668 fn pipelined_inplace_with_multiple_readers_orders_correctly() {
15669 use fuel_ir::Shape;
15670 let data = [1.0_f32, -2.0, 3.0, -4.0];
15671 let src_storage = fuel_memory::from_slice_cpu(&data);
15672
15673 let graph = Arc::new(RwLock::new(Graph::new()));
15674 let (x_id, step_id, sig_id) = {
15675 let mut g = graph.write().unwrap();
15676 let x = g.push(Node {
15677 op: Op::Const,
15678 inputs: vec![],
15679 shape: Shape::from_dims(&[4]),
15680 dtype: DType::F32,
15681 });
15682 let step = g.push(Node {
15683 op: Op::Step,
15684 inputs: vec![x],
15685 shape: Shape::from_dims(&[4]),
15686 dtype: DType::F32,
15687 });
15688 let sig = g.push(Node {
15689 op: Op::SigmoidInplace,
15690 inputs: vec![x],
15691 shape: Shape::from_dims(&[4]),
15692 dtype: DType::F32,
15693 });
15694 g.set_target_backend(step, BackendId::Cpu);
15695 g.set_target_backend(sig, BackendId::Cpu);
15696 (x, step, sig)
15697 };
15698
15699 let mut cache = StorageCache::new();
15700 cache.insert(x_id, Arc::new(RwLock::new(src_storage)));
15701
15702 let results = PipelinedExecutor::realize_many(graph, &[step_id, sig_id], cache)
15703 .expect("realize_many");
15704
15705 let step_guard = results[0].0.read().unwrap();
15709 let fuel_memory::BackendStorage::Cpu(c) = &step_guard.inner else {
15710 panic!("expected Cpu storage for step output");
15711 };
15712 let step_out: &[f32] = c.as_slice().expect("f32 cast");
15713 assert_eq!(
15714 step_out,
15715 &[1.0_f32, 0.0, 1.0, 0.0],
15716 "Step must run before SigmoidInplace; got post-mutation bytes"
15717 );
15718 drop(step_guard);
15719
15720 let sig_guard = results[1].0.read().unwrap();
15723 let fuel_memory::BackendStorage::Cpu(c) = &sig_guard.inner else {
15724 panic!("expected Cpu storage for sigmoid output");
15725 };
15726 let sig_out: &[f32] = c.as_slice().expect("f32 cast");
15727 let sig_ref = [
15728 1.0_f32 / (1.0 + (-1.0_f32).exp()),
15729 1.0_f32 / (1.0 + (2.0_f32).exp()),
15730 1.0_f32 / (1.0 + (-3.0_f32).exp()),
15731 1.0_f32 / (1.0 + (4.0_f32).exp()),
15732 ];
15733 for (got, want) in sig_out.iter().zip(sig_ref.iter()) {
15734 assert!((got - want).abs() < 1e-6, "sigmoid: got {got}, want {want}");
15735 }
15736 }
15737
15738 #[test]
15753 fn pipelined_residual_connection_inserts_safety_copy() {
15754 use fuel_ir::Shape;
15755 let data = [1.0_f32, -2.0, 3.0, -4.0];
15756 let src_storage = fuel_memory::from_slice_cpu(&data);
15757
15758 let graph = Arc::new(RwLock::new(Graph::new()));
15759 let (x_id, _y_id, z_id) = {
15760 let mut g = graph.write().unwrap();
15761 let x = g.push(Node {
15762 op: Op::Const,
15763 inputs: vec![],
15764 shape: Shape::from_dims(&[4]),
15765 dtype: DType::F32,
15766 });
15767 let y = g.push(Node {
15768 op: Op::ReluInplace,
15769 inputs: vec![x],
15770 shape: Shape::from_dims(&[4]),
15771 dtype: DType::F32,
15772 });
15773 let z = g.push(Node {
15774 op: Op::Add,
15775 inputs: vec![y, x],
15776 shape: Shape::from_dims(&[4]),
15777 dtype: DType::F32,
15778 });
15779 g.set_target_backend(y, BackendId::Cpu);
15780 g.set_target_backend(z, BackendId::Cpu);
15781 (x, y, z)
15782 };
15783
15784 let mut cache = StorageCache::new();
15785 cache.insert(x_id, Arc::new(RwLock::new(src_storage)));
15786
15787 let (result_arc, _) = PipelinedExecutor::realize(graph, z_id, cache)
15788 .expect("realize z — insert_safety_copies must break the cycle");
15789
15790 let guard = result_arc.read().unwrap();
15791 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15792 panic!("expected Cpu storage");
15793 };
15794 let got: &[f32] = c.as_slice().expect("f32 cast");
15795 assert_eq!(
15796 got,
15797 &[2.0_f32, -2.0, 6.0, -4.0],
15798 "Add must read pre-mutation x via the inserted safety copy"
15799 );
15800 }
15801
15802 fn cpu_const_f32(
15810 graph: &Arc<RwLock<Graph>>,
15811 cache: &mut StorageCache,
15812 data: &[f32],
15813 dims: &[usize],
15814 ) -> NodeId {
15815 let id = {
15816 let mut g = graph.write().unwrap();
15817 g.push(Node {
15818 op: Op::Const,
15819 inputs: vec![],
15820 shape: Shape::from_dims(dims),
15821 dtype: DType::F32,
15822 })
15823 };
15824 cache.insert(id, Arc::new(RwLock::new(fuel_memory::from_slice_cpu(data))));
15825 id
15826 }
15827
15828 #[test]
15833 fn pipelined_realize_selective_scan_bundled_produces_y_and_last_state() {
15834 let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
15835 let mut cache = StorageCache::new();
15836
15837 let u_id = cpu_const_f32(&graph, &mut cache, &[3.0], &[1, 1, 1]);
15838 let delta_id = cpu_const_f32(&graph, &mut cache, &[1.0], &[1, 1, 1]);
15839 let a_id = cpu_const_f32(&graph, &mut cache, &[-1.0], &[1, 1]);
15840 let b_id = cpu_const_f32(&graph, &mut cache, &[2.0], &[1, 1, 1]);
15841 let c_id = cpu_const_f32(&graph, &mut cache, &[0.5], &[1, 1, 1]);
15842
15843 let (y_id, last_state_id) = {
15844 let u_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), u_id);
15845 let delta_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), delta_id);
15846 let a_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), a_id);
15847 let b_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), b_id);
15848 let c_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), c_id);
15849 let (y, last_state) = u_t
15850 .selective_scan_bundled(&delta_t, &a_t, &b_t, &c_t, false)
15851 .expect("selective_scan_bundled");
15852 (y.id(), last_state.id())
15853 };
15854 {
15856 let mut g = graph.write().unwrap();
15857 let producer = g.node(y_id).inputs[0];
15858 g.set_target_backend(producer, BackendId::Cpu);
15859 }
15860
15861 let outputs =
15862 PipelinedExecutor::realize_many(Arc::clone(&graph), &[y_id, last_state_id], cache)
15863 .expect("realize_many");
15864 let (y_arc, y_layout) = outputs[0].clone();
15865 let (last_state_arc, last_state_layout) = outputs[1].clone();
15866
15867 assert!(
15869 Arc::ptr_eq(&y_arc, &last_state_arc),
15870 "y and last_state Views project the SAME bundled producer Storage",
15871 );
15872
15873 let guard = y_arc.read().unwrap();
15878 let fuel_memory::BackendStorage::Cpu(c_bytes) = &guard.inner else {
15879 panic!("expected Cpu storage");
15880 };
15881 let typed: &[f32] = c_bytes.as_slice().expect("f32 cast");
15882 assert_eq!(y_layout.start_offset(), 0);
15883 let y_value = typed[y_layout.start_offset()];
15884 assert!((y_value - 3.0).abs() < 1e-5, "y expected 3.0 got {y_value}");
15885 assert_eq!(
15886 last_state_layout.start_offset(),
15887 1,
15888 "slot 1 byte_offset 4 / sizeof(f32) = 1",
15889 );
15890 let ls_value = typed[last_state_layout.start_offset()];
15891 assert!(
15892 (ls_value - 6.0).abs() < 1e-5,
15893 "last_state expected 6.0 got {ls_value}"
15894 );
15895 }
15896
15897 #[test]
15908 fn pipelined_write_slice_round2_plain_slice_reader_not_stranded() {
15909 let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
15910 let mut cache = StorageCache::new();
15911
15912 let b_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
15913 let src1_id = cpu_const_f32(&graph, &mut cache, &[1.0, 2.0], &[2]);
15914 let src2_id = cpu_const_f32(&graph, &mut cache, &[99.0, 98.0], &[2]);
15915
15916 let (k_id, y_id) = {
15917 let mut g = graph.write().unwrap();
15918 let x_id = g.push(Node {
15919 op: Op::WriteSlice {
15920 ranges: vec![(0, 2)],
15921 dyn_offset: None,
15922 },
15923 inputs: vec![b_id, src1_id],
15924 shape: Shape::from_dims(&[4]),
15925 dtype: DType::F32,
15926 });
15927 let s_id = g.push(Node {
15928 op: Op::Slice {
15929 dim: 0,
15930 start: 0,
15931 len: 2,
15932 },
15933 inputs: vec![x_id],
15934 shape: Shape::from_dims(&[2]),
15935 dtype: DType::F32,
15936 });
15937 let k_id = g.push(Node {
15938 op: Op::Relu,
15939 inputs: vec![s_id],
15940 shape: Shape::from_dims(&[2]),
15941 dtype: DType::F32,
15942 });
15943 let y_id = g.push(Node {
15946 op: Op::WriteSlice {
15947 ranges: vec![(0, 2)],
15948 dyn_offset: None,
15949 },
15950 inputs: vec![x_id, src2_id],
15951 shape: Shape::from_dims(&[4]),
15952 dtype: DType::F32,
15953 });
15954 g.set_target_backend(x_id, BackendId::Cpu);
15955 g.set_target_backend(k_id, BackendId::Cpu);
15956 g.set_target_backend(y_id, BackendId::Cpu);
15957 (k_id, y_id)
15958 };
15959
15960 let plan = fuel_graph::opt::execution_plan(&graph.read().unwrap(), &[k_id, y_id]);
15961 eprintln!("[diagnostic plain-slice] execution_plan = {:?}", plan,);
15962
15963 let outputs = PipelinedExecutor::realize_many(Arc::clone(&graph), &[k_id, y_id], cache)
15964 .expect("realize_many: plain slice reader must not error");
15965 let (k_arc, _) = &outputs[0];
15966 let guard = k_arc.read().unwrap();
15967 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
15968 panic!("expected Cpu storage");
15969 };
15970 let out: &[f32] = c.as_slice().unwrap();
15971 assert_eq!(
15972 out,
15973 &[1.0, 2.0],
15974 "K = relu(slice(X,0,2)) must read PRE-round-2 bytes; got {:?}",
15975 out,
15976 );
15977 }
15978
15979 #[test]
15996 fn pipelined_write_slice_reshape_adopted_arc_reader_reads_pre_write_bytes() {
15997 let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
15998 let mut cache = StorageCache::new();
15999
16000 let b_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
16001 let src1_id = cpu_const_f32(&graph, &mut cache, &[1.0, 2.0], &[2]);
16002 let src2_id = cpu_const_f32(&graph, &mut cache, &[99.0, 98.0], &[2]);
16003
16004 let (x_id, k_id, y_id) = {
16005 let mut g = graph.write().unwrap();
16006 let x_id = g.push(Node {
16007 op: Op::WriteSlice {
16008 ranges: vec![(0, 2)],
16009 dyn_offset: None,
16010 },
16011 inputs: vec![b_id, src1_id],
16012 shape: Shape::from_dims(&[4]),
16013 dtype: DType::F32,
16014 });
16015 let s_id = g.push(Node {
16016 op: Op::Slice {
16019 dim: 0,
16020 start: 0,
16021 len: 2,
16022 },
16023 inputs: vec![x_id],
16024 shape: Shape::from_dims(&[2]),
16025 dtype: DType::F32,
16026 });
16027 let r_id = g.push(Node {
16028 op: Op::Reshape(Shape::from_dims(&[2])),
16029 inputs: vec![s_id],
16030 shape: Shape::from_dims(&[2]),
16031 dtype: DType::F32,
16032 });
16033 let k_id = g.push(Node {
16034 op: Op::Relu,
16035 inputs: vec![r_id],
16036 shape: Shape::from_dims(&[2]),
16037 dtype: DType::F32,
16038 });
16039 let y_id = g.push(Node {
16044 op: Op::WriteSlice {
16045 ranges: vec![(0, 2)],
16046 dyn_offset: None,
16047 },
16048 inputs: vec![x_id, src2_id],
16049 shape: Shape::from_dims(&[4]),
16050 dtype: DType::F32,
16051 });
16052 g.set_target_backend(x_id, BackendId::Cpu);
16053 g.set_target_backend(k_id, BackendId::Cpu);
16054 g.set_target_backend(y_id, BackendId::Cpu);
16055 (x_id, k_id, y_id)
16056 };
16057 let _ = x_id;
16058
16059 let plan = fuel_graph::opt::execution_plan(&graph.read().unwrap(), &[k_id, y_id]);
16060 eprintln!(
16061 "[diagnostic reshape-adopted-arc] execution_plan = {:?}",
16062 plan,
16063 );
16064
16065 let outputs = PipelinedExecutor::realize_many(Arc::clone(&graph), &[k_id, y_id], cache)
16066 .expect("realize_many");
16067 let (k_arc, _) = &outputs[0];
16068 let guard = k_arc.read().unwrap();
16069 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
16070 panic!("expected Cpu storage");
16071 };
16072 let out: &[f32] = c.as_slice().unwrap();
16073 assert_eq!(
16078 out,
16079 &[1.0, 2.0],
16080 "K must read PRE-round-2 bytes through the adopted Arc; got {:?} \
16081 (post-round-2 corruption if this fails)",
16082 out,
16083 );
16084 }
16085
16086 #[test]
16097 fn pipelined_write_slice_mla_two_layer_round2_reshape_reader_not_stranded() {
16098 let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
16099 let mut cache = StorageCache::new();
16100
16101 let b_l_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
16102 let src_l_n_id = cpu_const_f32(&graph, &mut cache, &[1.0, 2.0], &[2]);
16103 let b_l1_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
16104 let src_l_n1_id = cpu_const_f32(&graph, &mut cache, &[99.0, 98.0], &[2]);
16105
16106 let (x_l_n1_id, x_l1_n_id) = {
16107 let mut g = graph.write().unwrap();
16108 let x_l_n = g.push(Node {
16110 op: Op::WriteSlice {
16111 ranges: vec![(0, 2)],
16112 dyn_offset: None,
16113 },
16114 inputs: vec![b_l_id, src_l_n_id],
16115 shape: Shape::from_dims(&[4]),
16116 dtype: DType::F32,
16117 });
16118 let s_l_n = g.push(Node {
16121 op: Op::Slice {
16122 dim: 0,
16123 start: 0,
16124 len: 2,
16125 },
16126 inputs: vec![x_l_n],
16127 shape: Shape::from_dims(&[2]),
16128 dtype: DType::F32,
16129 });
16130 let rs_l_n = g.push(Node {
16131 op: Op::Reshape(Shape::from_dims(&[2])),
16132 inputs: vec![s_l_n],
16133 shape: Shape::from_dims(&[2]),
16134 dtype: DType::F32,
16135 });
16136 let r_l_n = g.push(Node {
16137 op: Op::Relu,
16138 inputs: vec![rs_l_n],
16139 shape: Shape::from_dims(&[2]),
16140 dtype: DType::F32,
16141 });
16142 let x_l1_n = g.push(Node {
16144 op: Op::WriteSlice {
16145 ranges: vec![(0, 2)],
16146 dyn_offset: None,
16147 },
16148 inputs: vec![b_l1_id, r_l_n],
16149 shape: Shape::from_dims(&[4]),
16150 dtype: DType::F32,
16151 });
16152 let x_l_n1 = g.push(Node {
16156 op: Op::WriteSlice {
16157 ranges: vec![(0, 2)],
16158 dyn_offset: None,
16159 },
16160 inputs: vec![x_l_n, src_l_n1_id],
16161 shape: Shape::from_dims(&[4]),
16162 dtype: DType::F32,
16163 });
16164 g.set_target_backend(x_l_n, BackendId::Cpu);
16165 g.set_target_backend(r_l_n, BackendId::Cpu);
16166 g.set_target_backend(x_l1_n, BackendId::Cpu);
16167 g.set_target_backend(x_l_n1, BackendId::Cpu);
16168 (x_l_n1, x_l1_n)
16169 };
16170
16171 let plan = fuel_graph::opt::execution_plan(&graph.read().unwrap(), &[x_l1_n_id, x_l_n1_id]);
16172 eprintln!("[diagnostic mla-2-layer] execution_plan = {:?}", plan);
16173
16174 let outputs =
16175 PipelinedExecutor::realize_many(Arc::clone(&graph), &[x_l1_n_id, x_l_n1_id], cache)
16176 .expect("realize_many");
16177 let (x_l1_n_arc, _) = &outputs[0];
16180 let guard = x_l1_n_arc.read().unwrap();
16181 let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
16182 panic!("expected Cpu storage");
16183 };
16184 let out: &[f32] = c.as_slice().unwrap();
16185 assert_eq!(
16186 out,
16187 &[1.0, 2.0, 0.0, 0.0],
16188 "layer L+1's round-N append must reflect PRE-round-(N+1) attention \
16189 math output; got {:?} (post-round-(N+1) corruption if this fails)",
16190 out,
16191 );
16192 }
16193}