use std::collections::{HashMap, HashSet};
use std::sync::mpsc::{Sender, channel};
use std::sync::{Arc, RwLock};
use std::thread;
use fuel_graph::opt::{execution_plan, insert_safety_copies};
use fuel_graph::registry::FusedOpParams;
use fuel_graph::{Graph, Node, NodeId, Op, PickedRoute};
use fuel_ir::dispatch::OpKind;
use fuel_ir::probe::BackendId;
use fuel_ir::{DType, DeviceLocation, Error, Layout, Result, SymEnv};
use crate::compiled::{
CompiledNode, CompletionHandle, compile_node, execute_compiled_with_wait_hint,
};
use crate::dispatch::global_bindings;
use crate::kernel::{KernelBindingTable, KernelDTypes, MatmulM, OpParams};
use crate::optimize::OptimizedGraph;
use crate::ranker::{BackendRuntimeLookup, RuntimeSelector, resolve_branch};
#[cfg(test)]
use crate::ranker::pick_route;
use fuel_memory::Storage;
enum OrderSource<'a> {
Default,
Optimized {
optimized: &'a OptimizedGraph,
route: Option<&'a PickedRoute>,
},
Streaming {
optimized: &'a OptimizedGraph,
pick: &'a StreamingPick,
},
Cached {
optimized: &'a OptimizedGraph,
order_cache: &'a std::sync::OnceLock<Vec<NodeId>>,
},
}
#[derive(Clone)]
pub(crate) struct StreamingPick {
pub selector: Arc<dyn RuntimeSelector>,
#[allow(dead_code)]
pub lookup: Option<BackendRuntimeLookup>,
}
pub type StorageCache = HashMap<NodeId, Arc<RwLock<Storage>>>;
pub struct PersistentOutputs {
map: HashMap<NodeId, Arc<RwLock<Storage>>>,
mode: PersistentMode,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum PersistentMode {
Record,
Reuse,
}
impl PersistentOutputs {
pub fn recording() -> Self {
Self {
map: HashMap::new(),
mode: PersistentMode::Record,
}
}
pub fn reuse(map: HashMap<NodeId, Arc<RwLock<Storage>>>) -> Self {
Self {
map,
mode: PersistentMode::Reuse,
}
}
pub fn into_map(self) -> HashMap<NodeId, Arc<RwLock<Storage>>> {
self.map
}
}
#[cfg(feature = "cuda")]
fn op_kind_is_capture_writeinto(op: OpKind) -> bool {
matches!(
op,
OpKind::AddElementwise
| OpKind::SubElementwise
| OpKind::MulElementwise
| OpKind::DivElementwise
| OpKind::MaximumElementwise
| OpKind::MinimumElementwise
| OpKind::PowElementwise
| OpKind::RemElementwise
| OpKind::SoftmaxLastDim
| OpKind::LogSoftmaxLastDim
| OpKind::Affine
| OpKind::MatMul
| OpKind::IndexSelect
| OpKind::RmsNormLastDim
| OpKind::SiluElementwise
| OpKind::GeluTanhElementwise
| OpKind::GeluElementwise
| OpKind::NegElementwise
| OpKind::Concat
| OpKind::Rope
)
}
#[cfg(feature = "cuda")]
pub struct CapturedDecode {
pub run: fuel_cuda_backend::CapturedRun,
pub output: Arc<RwLock<Storage>>,
pub output_layout: Layout,
pub persistent: HashMap<NodeId, Arc<RwLock<Storage>>>,
pub retained_inputs: Vec<Arc<RwLock<Storage>>>,
}
#[cfg(feature = "cuda")]
pub struct CapturedDecodeSession {
device: fuel_cuda_backend::CudaDevice,
captured: CapturedDecode,
per_token_inputs: HashMap<NodeId, Arc<RwLock<Storage>>>,
_retained_inputs: Vec<Arc<RwLock<Storage>>>,
}
#[cfg(feature = "cuda")]
impl CapturedDecodeSession {
pub fn capture(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
per_token_input_ids: &[NodeId],
sym_env: SymEnv,
) -> Result<Self> {
use fuel_memory::BackendStorage;
let retained_inputs: Vec<Arc<RwLock<Storage>>> = inputs.values().cloned().collect();
let mut per_token_inputs = HashMap::with_capacity(per_token_input_ids.len());
for &id in per_token_input_ids {
let arc = inputs.get(&id).cloned().ok_or_else(|| {
Error::Msg(format!(
"CapturedDecodeSession::capture: per-token input {:?} not in inputs",
id,
))
.bt()
})?;
per_token_inputs.insert(id, arc);
}
let device = {
let any = per_token_inputs
.values()
.next()
.or_else(|| inputs.values().next())
.ok_or_else(|| {
Error::Msg("CapturedDecodeSession::capture: no inputs".into()).bt()
})?;
let guard = any.read().map_err(|_| poisoned("input storage"))?;
match &guard.inner {
BackendStorage::Cuda(c) => c.device().clone(),
other => {
return Err(Error::Msg(format!(
"CapturedDecodeSession::capture: input is BackendStorage::{:?}, not Cuda",
std::mem::discriminant(other),
))
.bt());
}
}
};
let captured = PipelinedExecutor::capture_decode(graph, target, inputs, sym_env)?;
Ok(Self {
device,
captured,
per_token_inputs,
_retained_inputs: retained_inputs,
})
}
pub fn node_outputs(&self) -> &HashMap<NodeId, Arc<RwLock<Storage>>> {
&self.captured.persistent
}
pub fn replay_token(&self, updates: &[(NodeId, &[u8])]) -> Result<Arc<RwLock<Storage>>> {
use fuel_memory::BackendStorage;
for (id, bytes) in updates {
let arc = self.per_token_inputs.get(id).ok_or_else(|| {
Error::Msg(format!(
"CapturedDecodeSession::replay_token: {:?} is not a per-token input",
id,
))
.bt()
})?;
let guard = arc.read().map_err(|_| poisoned("per-token input"))?;
match &guard.inner {
BackendStorage::Cuda(c) => c.write_from_host(bytes)?,
other => {
return Err(Error::Msg(format!(
"CapturedDecodeSession::replay_token: input {:?} is \
BackendStorage::{:?}, not Cuda",
id,
std::mem::discriminant(other),
))
.bt());
}
}
}
self.captured.run.replay(&self.device)?;
self.device.synchronize()?;
Ok(Arc::clone(&self.captured.output))
}
pub fn output(&self) -> &Arc<RwLock<Storage>> {
&self.captured.output
}
pub fn output_layout(&self) -> &Layout {
&self.captured.output_layout
}
}
enum WorkItemKind {
ConstAdopt,
ViewOf { input: NodeId },
ContiguizeOf { input: NodeId },
Kernel,
Iota { len: usize },
ReleaseMarker,
WriteSlice { dest: NodeId, source: NodeId },
WriteSliceRotating {
dest: NodeId,
source: NodeId,
position: NodeId,
},
WriteSliceDoff {
dest: NodeId,
source: NodeId,
offset: NodeId,
},
Copy { target_location: DeviceLocation },
Move { target_location: DeviceLocation },
Alloc { target_location: DeviceLocation },
ZeroFill,
InplaceKernel {
target_idx: usize,
},
SlotView { producer: NodeId },
SlotOwn { producer: NodeId, slot: u32 },
}
struct WorkItem {
node_id: NodeId,
inputs: Vec<NodeId>,
elem_count: usize,
dtype: DType,
target_backend: BackendId,
kind: WorkItemKind,
compiled: Option<CompiledNode>,
output_layout: Layout,
destructive_input: Option<usize>,
output_bundle: Option<Arc<[fuel_ir::storage::OutputView]>>,
}
fn extend_with_side_effect_roots(graph: &Graph, user_roots: &[NodeId]) -> Vec<NodeId> {
let side = graph.side_effect_roots();
if side.is_empty() {
return user_roots.to_vec();
}
let mut out = Vec::with_capacity(user_roots.len() + side.len());
out.extend_from_slice(user_roots);
for &s in side {
if !out.contains(&s) {
out.push(s);
}
}
out
}
pub struct PipelinedExecutor;
impl PipelinedExecutor {
pub fn realize(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
Self::realize_inner(
graph,
target,
inputs,
OrderSource::Default,
SymEnv::default(),
)
.map(|(s, l, _produced)| (s, l))
}
pub fn realize_with_env(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
Self::realize_inner(graph, target, inputs, OrderSource::Default, sym_env)
.map(|(s, l, _produced)| (s, l))
}
pub fn realize_with_env_producing(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout, SymEnv)> {
Self::realize_inner(graph, target, inputs, OrderSource::Default, sym_env)
}
pub fn realize_with_optimized(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
optimized: &OptimizedGraph,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
Self::realize_inner(
graph,
target,
inputs,
OrderSource::Optimized {
optimized,
route: None,
},
SymEnv::default(),
)
.map(|(s, l, _produced)| (s, l))
}
pub fn realize_with_optimized_env(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
optimized: &OptimizedGraph,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
Self::realize_inner(
graph,
target,
inputs,
OrderSource::Optimized {
optimized,
route: None,
},
sym_env,
)
.map(|(s, l, _produced)| (s, l))
}
pub fn realize_with_optimized_route(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
optimized: &OptimizedGraph,
route: &PickedRoute,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
Self::realize_inner(
graph,
target,
inputs,
OrderSource::Optimized {
optimized,
route: Some(route),
},
SymEnv::default(),
)
.map(|(s, l, _produced)| (s, l))
}
pub fn realize_with_optimized_route_env(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
optimized: &OptimizedGraph,
route: &PickedRoute,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
Self::realize_inner(
graph,
target,
inputs,
OrderSource::Optimized {
optimized,
route: Some(route),
},
sym_env,
)
.map(|(s, l, _produced)| (s, l))
}
pub fn realize_with_optimized_picking_env(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
optimized: &OptimizedGraph,
selector: Option<Arc<dyn RuntimeSelector>>,
lookup: Option<BackendRuntimeLookup>,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
match Self::streaming_pick_for(&graph, selector, lookup)? {
Some(pick) => Self::realize_inner(
graph,
target,
inputs,
OrderSource::Streaming {
optimized,
pick: &pick,
},
sym_env,
)
.map(|(s, l, _produced)| (s, l)),
None => Self::realize_with_optimized_env(graph, target, inputs, optimized, sym_env),
}
}
pub fn realize_with_optimized_picking_cached_env(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
optimized: &OptimizedGraph,
order_cache: &std::sync::OnceLock<Vec<NodeId>>,
selector: Option<Arc<dyn RuntimeSelector>>,
lookup: Option<BackendRuntimeLookup>,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout)> {
match Self::streaming_pick_for(&graph, selector, lookup)? {
Some(pick) => Self::realize_inner(
graph,
target,
inputs,
OrderSource::Streaming {
optimized,
pick: &pick,
},
sym_env,
)
.map(|(s, l, _produced)| (s, l)),
None => Self::realize_inner(
graph,
target,
inputs,
OrderSource::Cached {
optimized,
order_cache,
},
sym_env,
)
.map(|(s, l, _produced)| (s, l)),
}
}
fn streaming_pick_for(
graph: &Arc<RwLock<Graph>>,
selector: Option<Arc<dyn RuntimeSelector>>,
lookup: Option<BackendRuntimeLookup>,
) -> Result<Option<StreamingPick>> {
let Some(selector) = selector else {
return Ok(None);
};
let g = graph
.read()
.map_err(|_| Error::Msg("graph lock poisoned".into()).bt())?;
let has_branch = (0..g.len()).any(|i| matches!(g.node(NodeId(i)).op, Op::Branch { .. }));
if !has_branch {
return Ok(None);
}
Ok(Some(StreamingPick { selector, lookup }))
}
#[cfg(feature = "cuda")]
pub fn capture_decode(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
sym_env: SymEnv,
) -> Result<CapturedDecode> {
use fuel_memory::BackendStorage;
{
let mut g = graph.write().map_err(|_| poisoned("graph lock"))?;
let effective_roots = extend_with_side_effect_roots(&g, &[target]);
insert_safety_copies(&mut g, &effective_roots);
}
let (compiler_work, _wait_set, seed_layouts) = {
let g = graph.read().map_err(|_| poisoned("graph lock"))?;
let effective_roots = extend_with_side_effect_roots(&g, &[target]);
let (work, wait_set) =
compiler_work_and_wait_set_for(&g, &effective_roots, &OrderSource::Default);
let mut layouts: HashMap<NodeId, Layout> = HashMap::with_capacity(inputs.len());
for &id in inputs.keys() {
layouts.insert(id, g.layout(id));
}
(work, wait_set, layouts)
};
let (tx, rx) = channel::<Result<WorkItem>>();
let graph_for_compiler = Arc::clone(&graph);
let compiler = thread::spawn(move || {
compiler_thread_body(graph_for_compiler, compiler_work, sym_env, tx);
});
let mut items: Vec<WorkItem> = Vec::new();
for item in rx {
items.push(item?);
}
compiler
.join()
.map_err(|_| Error::Msg("compiler thread panicked".to_string()).bt())?;
{
let off: Vec<String> = items
.iter()
.filter_map(|it| match &it.kind {
WorkItemKind::Kernel if it.target_backend != BackendId::Cuda => {
let desc = it
.compiled
.as_ref()
.map(|c| format!("{:?} dtypes={:?}", c.op, c.dtypes))
.unwrap_or_else(|| "?".to_string());
Some(format!(
"node {:?}: {desc} on {:?}",
it.node_id, it.target_backend
))
}
_ => None,
})
.collect();
if !off.is_empty() {
eprintln!(
"[capture-diag] {} kernel(s) placed OFF Cuda (these block capture): {}",
off.len(),
off.join("; "),
);
}
}
for item in &items {
match &item.kind {
WorkItemKind::Move { .. } => {
return Err(Error::Msg(format!(
"capture_decode: graph has an Op::Move at node {:?}; \
Move's destructive source-release half is driven by \
per-realize cleanup, not per-replay, so capture is \
not supported for Move (out of scope — see the \
`capture_decode` validation loop's Move arm).",
item.node_id,
))
.bt());
}
WorkItemKind::Copy { target_location } => {
if !matches!(target_location, DeviceLocation::Cuda { .. }) {
return Err(Error::Msg(format!(
"capture_decode: graph has a cross-device Op::Copy at \
node {:?} (target {:?}); capture requires a \
single-device CUDA decode step — only same-device \
CUDA→CUDA copies are capture-safe.",
item.node_id, target_location,
))
.bt());
}
}
WorkItemKind::Kernel => {
if item.target_backend != BackendId::Cuda {
return Err(Error::Msg(format!(
"capture_decode: kernel node {:?} targets {:?}, not Cuda",
item.node_id, item.target_backend,
))
.bt());
}
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"capture_decode: kernel node {:?} has no compiled op",
item.node_id,
))
.bt()
})?;
let op = compiled.op;
if !op_kind_is_capture_writeinto(op) {
return Err(Error::Msg(format!(
"capture_decode: op {:?} (node {:?}) has no write-into CUDA \
wrapper yet — capture unsupported for this op. Extend \
op_kind_is_capture_writeinto once its write-into variant lands.",
op, item.node_id,
))
.bt());
}
if op == OpKind::Concat {
if let OpParams::Concat {
input_dim_sizes, ..
} = &compiled.op_params
{
if input_dim_sizes.len() != 2 {
return Err(Error::Msg(format!(
"capture_decode: Concat node {:?} has {} inputs; only \
2-input concat (the rope rotate-half shape) has a \
write-into CUDA path today — N-ary (N>2) concat would \
allocate mid-capture.",
item.node_id,
input_dim_sizes.len(),
))
.bt());
}
}
}
}
_ => {}
}
}
let mut persistent = PersistentOutputs::recording();
let mut warm_cache: StorageCache = inputs.clone();
let mut warm_layouts = seed_layouts.clone();
let mut warm_syms = SymEnv::new();
for item in &items {
let _handle = execute_work_item(
item,
&mut warm_cache,
&mut warm_layouts,
&mut warm_syms,
false,
Some(&mut persistent),
)
.map_err(|e| with_node_location(&graph, item.node_id, e))?;
}
sync_active_cuda_devices(&warm_cache)?;
let output = warm_cache.get(&target).cloned().ok_or_else(|| {
Error::Msg(format!(
"capture_decode: target {:?} not realized in warm run",
target,
))
.bt()
})?;
let output_layout = warm_layouts.get(&target).cloned().ok_or_else(|| {
Error::Msg(format!("capture_decode: target {:?} has no layout", target)).bt()
})?;
let device = {
let guard = output.read().map_err(|_| poisoned("output storage"))?;
match &guard.inner {
BackendStorage::Cuda(c) => c.device().clone(),
other => {
return Err(Error::Msg(format!(
"capture_decode: target {:?} is BackendStorage::{:?}, not Cuda",
target,
std::mem::discriminant(other),
))
.bt());
}
}
};
let mut reuse = PersistentOutputs::reuse(persistent.into_map());
let retained_inputs: Vec<Arc<RwLock<Storage>>> = inputs.values().cloned().collect();
let mut cap_cache: StorageCache = inputs;
let mut cap_layouts = seed_layouts;
let mut cap_syms = SymEnv::new();
let mut inner_err: Option<Error> = None;
let items_ref = &items;
let graph_ref = &graph;
let run = device.capture_run(|_stream| {
for item in items_ref {
match execute_work_item(
item,
&mut cap_cache,
&mut cap_layouts,
&mut cap_syms,
false,
Some(&mut reuse),
) {
Ok(_handle) => {}
Err(e) => {
inner_err = Some(with_node_location(graph_ref, item.node_id, e));
break;
}
}
}
Ok(())
})?;
if let Some(e) = inner_err {
return Err(e);
}
Ok(CapturedDecode {
run,
output,
output_layout,
retained_inputs,
persistent: reuse.into_map(),
})
}
fn realize_inner(
graph: Arc<RwLock<Graph>>,
target: NodeId,
inputs: StorageCache,
order_source: OrderSource<'_>,
sym_env: SymEnv,
) -> Result<(Arc<RwLock<Storage>>, Layout, SymEnv)> {
{
let mut g = graph.write().map_err(|_| poisoned("graph lock"))?;
let effective_roots = extend_with_side_effect_roots(&g, &[target]);
insert_safety_copies(&mut g, &effective_roots);
}
let (compiler_work, wait_set, mut layout_cache): (
CompilerWork,
WaitSet,
HashMap<NodeId, Layout>,
) = {
let g = graph.read().map_err(|_| poisoned("graph lock"))?;
let effective_roots = extend_with_side_effect_roots(&g, &[target]);
let (work, wait_set) =
compiler_work_and_wait_set_for(&g, &effective_roots, &order_source);
let mut layouts = HashMap::with_capacity(inputs.len());
for &id in inputs.keys() {
layouts.insert(id, g.layout(id));
}
(work, wait_set, layouts)
};
let (tx, rx) = channel::<Result<WorkItem>>();
let graph_for_compiler = Arc::clone(&graph);
let compiler = thread::spawn(move || {
compiler_thread_body(graph_for_compiler, compiler_work, sym_env, tx);
});
let plan_generation: Option<u64> = generation_for(&order_source);
let mut current_chunk_backend: Option<BackendId> = None;
let mut cache: StorageCache = inputs;
let mut produced_syms = SymEnv::new();
let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
let mut inflight_vulkan: Vec<InflightVulkan> = Vec::new();
let mut multi_backend = false;
for item in rx {
let item = item?;
if current_chunk_backend != Some(item.target_backend) {
if let Some(plan_gen) = plan_generation {
let live = crate::dispatch::topology_generation();
if live != plan_gen {
return Err(Error::TopologyChanged {
plan_generation: plan_gen,
current_generation: live,
}
.bt());
}
}
let leaving = current_chunk_backend;
if let Some(prev) = leaving
&& prev != item.target_backend
{
multi_backend = true;
if prev == BackendId::Vulkan {
eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
}
}
current_chunk_backend = Some(item.target_backend);
}
if multi_backend
&& item.target_backend == BackendId::Vulkan
&& !inflight_vulkan.is_empty()
{
drain_inflight_vulkan(&mut inflight_vulkan)?;
}
if matches!(
item.kind,
WorkItemKind::Copy { .. } | WorkItemKind::Move { .. }
) {
if multi_backend {
eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
}
if let Some(&producer) = item.inputs.first() {
if multi_backend && copy_source_is_vulkan(&cache, producer)? {
drain_inflight_vulkan(&mut inflight_vulkan)?;
}
wait_producer_handle(&mut handles, producer)?;
}
}
let handle = execute_work_item(
&item,
&mut cache,
&mut layout_cache,
&mut produced_syms,
will_be_waited(&wait_set, item.node_id),
None,
)
.map_err(|e| with_node_location(&graph, item.node_id, e))?;
store_handle(&mut handles, item.node_id, handle);
if let Some(d_idx) = item.destructive_input
&& let Some(&destroyed) = item.inputs.get(d_idx)
&& destroyed != target
{
if multi_backend {
defer_evicted_vulkan_buffer(&cache, destroyed, &mut inflight_vulkan)?;
} else if let Some(d_arc) = cache.get(&destroyed) {
force_flush_vulkan(d_arc)?;
}
cache.remove(&destroyed);
layout_cache.remove(&destroyed);
handles.remove(&destroyed);
}
}
compiler
.join()
.map_err(|_| Error::Msg("compiler thread panicked".to_string()).bt())?;
drain_handles(&mut handles)?;
sync_active_cuda_devices(&cache)?;
drain_inflight_vulkan(&mut inflight_vulkan)?;
drain_vulkan_pending(&cache)?;
debug_assert!(
inflight_vulkan.is_empty(),
"A4b-4: in-flight Vulkan batch list must be empty after realize-end drain",
);
let storage = cache.remove(&target).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor::realize: target slot {:?} not populated after execution",
target
))
.bt()
})?;
let layout = layout_cache.remove(&target).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor::realize: target layout {:?} not populated after execution",
target
))
.bt()
})?;
Ok((storage, layout, produced_syms))
}
pub fn realize_many(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
Self::realize_many_inner(
graph,
targets,
inputs,
OrderSource::Default,
SymEnv::default(),
)
}
pub fn realize_many_with_optimized(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
optimized: &OptimizedGraph,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
Self::realize_many_inner(
graph,
targets,
inputs,
OrderSource::Optimized {
optimized,
route: None,
},
SymEnv::default(),
)
}
pub fn realize_many_with_optimized_env(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
optimized: &OptimizedGraph,
sym_env: SymEnv,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
Self::realize_many_inner(
graph,
targets,
inputs,
OrderSource::Optimized {
optimized,
route: None,
},
sym_env,
)
}
pub fn realize_many_with_optimized_route(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
optimized: &OptimizedGraph,
route: &PickedRoute,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
Self::realize_many_inner(
graph,
targets,
inputs,
OrderSource::Optimized {
optimized,
route: Some(route),
},
SymEnv::default(),
)
}
pub fn realize_many_with_optimized_route_env(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
optimized: &OptimizedGraph,
route: &PickedRoute,
sym_env: SymEnv,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
Self::realize_many_inner(
graph,
targets,
inputs,
OrderSource::Optimized {
optimized,
route: Some(route),
},
sym_env,
)
}
pub fn realize_many_with_optimized_picking_env(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
optimized: &OptimizedGraph,
selector: Option<Arc<dyn RuntimeSelector>>,
lookup: Option<BackendRuntimeLookup>,
sym_env: SymEnv,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
match Self::streaming_pick_for(&graph, selector, lookup)? {
Some(pick) => Self::realize_many_inner(
graph,
targets,
inputs,
OrderSource::Streaming {
optimized,
pick: &pick,
},
sym_env,
),
None => {
Self::realize_many_with_optimized_env(graph, targets, inputs, optimized, sym_env)
}
}
}
fn realize_many_inner(
graph: Arc<RwLock<Graph>>,
targets: &[NodeId],
inputs: StorageCache,
order_source: OrderSource<'_>,
sym_env: SymEnv,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
if targets.is_empty() {
return Ok(Vec::new());
}
{
let mut g = graph.write().map_err(|_| poisoned("graph lock"))?;
let effective_roots = extend_with_side_effect_roots(&g, targets);
insert_safety_copies(&mut g, &effective_roots);
}
let (compiler_work, wait_set, mut layout_cache): (
CompilerWork,
WaitSet,
HashMap<NodeId, Layout>,
) = {
let g = graph.read().map_err(|_| poisoned("graph lock"))?;
let effective_roots = extend_with_side_effect_roots(&g, targets);
let (work, wait_set) =
compiler_work_and_wait_set_for(&g, &effective_roots, &order_source);
let mut layouts = HashMap::with_capacity(inputs.len());
for &id in inputs.keys() {
layouts.insert(id, g.layout(id));
}
(work, wait_set, layouts)
};
let target_set: std::collections::HashSet<NodeId> = targets.iter().copied().collect();
let (tx, rx) = channel::<Result<WorkItem>>();
let graph_for_compiler = Arc::clone(&graph);
let compiler = thread::spawn(move || {
compiler_thread_body(graph_for_compiler, compiler_work, sym_env, tx);
});
let plan_generation: Option<u64> = generation_for(&order_source);
let mut current_chunk_backend: Option<BackendId> = None;
let mut cache: StorageCache = inputs;
let mut produced_syms = SymEnv::new();
let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
let mut inflight_vulkan: Vec<InflightVulkan> = Vec::new();
let mut multi_backend = false;
for item in rx {
let item = item?;
if current_chunk_backend != Some(item.target_backend) {
if let Some(plan_gen) = plan_generation {
let live = crate::dispatch::topology_generation();
if live != plan_gen {
return Err(Error::TopologyChanged {
plan_generation: plan_gen,
current_generation: live,
}
.bt());
}
}
let leaving = current_chunk_backend;
if let Some(prev) = leaving
&& prev != item.target_backend
{
multi_backend = true;
if prev == BackendId::Vulkan {
eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
}
}
current_chunk_backend = Some(item.target_backend);
}
if multi_backend
&& item.target_backend == BackendId::Vulkan
&& !inflight_vulkan.is_empty()
{
drain_inflight_vulkan(&mut inflight_vulkan)?;
}
if matches!(
item.kind,
WorkItemKind::Copy { .. } | WorkItemKind::Move { .. }
) {
if multi_backend {
eager_submit_all_vulkan(&cache, &mut inflight_vulkan)?;
}
if let Some(&producer) = item.inputs.first() {
if multi_backend && copy_source_is_vulkan(&cache, producer)? {
drain_inflight_vulkan(&mut inflight_vulkan)?;
}
wait_producer_handle(&mut handles, producer)?;
}
}
let handle = execute_work_item(
&item,
&mut cache,
&mut layout_cache,
&mut produced_syms,
will_be_waited(&wait_set, item.node_id),
None,
)
.map_err(|e| with_node_location(&graph, item.node_id, e))?;
store_handle(&mut handles, item.node_id, handle);
if let Some(d_idx) = item.destructive_input
&& let Some(&destroyed) = item.inputs.get(d_idx)
&& !target_set.contains(&destroyed)
{
if multi_backend {
defer_evicted_vulkan_buffer(&cache, destroyed, &mut inflight_vulkan)?;
} else if let Some(d_arc) = cache.get(&destroyed) {
force_flush_vulkan(d_arc)?;
}
cache.remove(&destroyed);
layout_cache.remove(&destroyed);
handles.remove(&destroyed);
}
}
compiler
.join()
.map_err(|_| Error::Msg("compiler thread panicked".to_string()).bt())?;
drain_handles(&mut handles)?;
sync_active_cuda_devices(&cache)?;
drain_inflight_vulkan(&mut inflight_vulkan)?;
drain_vulkan_pending(&cache)?;
debug_assert!(
inflight_vulkan.is_empty(),
"A4b-4: in-flight Vulkan batch list must be empty after realize-end drain",
);
let mut out = Vec::with_capacity(targets.len());
for &target in targets {
let storage = cache.get(&target).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor::realize_many: target {:?} not populated after execution",
target
))
.bt()
})?;
let layout = layout_cache.get(&target).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor::realize_many: target layout {:?} not populated after execution",
target
))
.bt()
})?;
out.push((storage, layout));
}
Ok(out)
}
}
fn order_for(
graph: &Graph,
effective_roots: &[NodeId],
order_source: &OrderSource<'_>,
) -> Vec<NodeId> {
match order_source {
OrderSource::Default => execution_plan(graph, effective_roots),
OrderSource::Streaming {
optimized: _optimized,
pick,
} => {
let bindings = global_bindings();
fuel_graph::lower_picked_route_streaming(graph, effective_roots, |branch| {
resolve_branch(graph, branch, &bindings, pick.selector.as_ref())
})
}
OrderSource::Optimized { optimized, route } => {
match route {
Some(route) => fuel_graph::lower_picked_route(graph, effective_roots, route),
None => {
let view = OptimizedGraph {
roots: effective_roots.to_vec(),
generation: optimized.generation,
placements: None,
};
view.dispatch_order(graph)
}
}
}
OrderSource::Cached { .. } => {
unreachable!("OrderSource::Cached is resolved before order_for is called")
}
}
}
fn build_wait_set(graph: &Graph, order: &[NodeId]) -> HashSet<NodeId> {
let mut set = HashSet::new();
for &id in order {
let node = graph.node(id);
if matches!(node.op, Op::Copy { .. } | Op::Move { .. })
&& let Some(&producer) = node.inputs.first()
{
set.insert(producer);
}
}
set
}
type WaitSet = Option<HashSet<NodeId>>;
#[inline]
fn will_be_waited(wait_set: &WaitSet, node_id: NodeId) -> bool {
match wait_set {
None => true,
Some(set) => set.contains(&node_id),
}
}
fn compiler_work_and_wait_set_for(
graph: &Graph,
effective_roots: &[NodeId],
order_source: &OrderSource<'_>,
) -> (CompilerWork, WaitSet) {
match order_source {
OrderSource::Streaming { pick, .. } => (
CompilerWork::Streaming {
roots: effective_roots.to_vec(),
pick: (*pick).clone(),
},
None,
),
OrderSource::Cached {
optimized,
order_cache,
} => {
let order = order_cache.get_or_init(|| {
order_for(
graph,
effective_roots,
&OrderSource::Optimized {
optimized,
route: None,
},
)
});
let wait_set = build_wait_set(graph, order);
(CompilerWork::Order(order.clone()), Some(wait_set))
}
other => {
let order = order_for(graph, effective_roots, other);
let wait_set = build_wait_set(graph, &order);
(CompilerWork::Order(order), Some(wait_set))
}
}
}
fn generation_for(order_source: &OrderSource<'_>) -> Option<u64> {
match order_source {
OrderSource::Optimized { optimized, .. } => Some(optimized.generation),
OrderSource::Streaming { optimized, .. } => Some(optimized.generation),
OrderSource::Cached { optimized, .. } => Some(optimized.generation),
OrderSource::Default => None,
}
}
enum CompilerWork {
Order(Vec<NodeId>),
Streaming {
roots: Vec<NodeId>,
pick: StreamingPick,
},
}
fn compiler_thread_body(
graph: Arc<RwLock<Graph>>,
work: CompilerWork,
sym_env: SymEnv,
tx: Sender<Result<WorkItem>>,
) {
let bindings = global_bindings();
let g = match graph.read() {
Ok(g) => g,
Err(_) => {
let _ = tx.send(Err(poisoned("graph lock in compiler")));
return;
}
};
let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
let compile_and_send = |id: NodeId, layout_cache: &mut HashMap<NodeId, Layout>| -> bool {
let item = compile_one(&g, id, layout_cache, &bindings, &sym_env);
let stop_on_err = item.is_err();
if tx.send(item).is_err() {
return false;
}
!stop_on_err
};
match work {
CompilerWork::Order(order) => {
for id in order {
if !compile_and_send(id, &mut layout_cache) {
return;
}
}
}
CompilerWork::Streaming { roots, pick } => {
let mut stopped = false;
fuel_graph::walk_picked_route_streaming(
&g,
&roots,
|branch| resolve_branch(&g, branch, &bindings, pick.selector.as_ref()),
|members| {
if stopped {
return;
}
for &id in members {
if !compile_and_send(id, &mut layout_cache) {
stopped = true;
return;
}
}
},
);
}
}
}
fn resolve_compiled(
op_kind: OpKind,
dtypes: &[DType],
target_backend: BackendId,
op_params: OpParams,
bindings: &KernelBindingTable,
) -> Result<CompiledNode> {
compile_node(op_kind, dtypes, target_backend, op_params, bindings)
}
fn compile_one(
graph: &Graph,
id: NodeId,
layout_cache: &mut HashMap<NodeId, Layout>,
bindings: &KernelBindingTable,
sym_env: &SymEnv,
) -> Result<WorkItem> {
let node = graph.node(id);
let elem_count = node.shape.elem_count();
let inputs = node.inputs.clone();
let destructive_input = node.op.destructive_input();
if matches!(node.op, Op::Const) {
let output_layout = graph.layout(id);
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend: BackendId::Cpu,
kind: WorkItemKind::ConstAdopt,
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if node.op.is_view_op() {
if inputs.len() != 1 {
return Err(Error::Msg(format!(
"view op {:?} expects 1 input, got {}",
node.op,
inputs.len(),
))
.bt());
}
let output_layout = graph.layout(id);
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::ViewOf {
input: node.inputs[0],
},
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::View { .. } = node.op {
if inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::View expects 1 input (the producer), got {}",
inputs.len(),
))
.bt());
}
let output_layout = graph.layout(id);
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs: inputs.clone(),
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::SlotView {
producer: inputs[0],
},
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::ViewOwned { slot } = node.op {
if inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::ViewOwned expects 1 input (the producer), got {}",
inputs.len(),
))
.bt());
}
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs: inputs.clone(),
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::SlotOwn {
producer: inputs[0],
slot,
},
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::Release) {
if inputs.len() != 1 {
return Err(
Error::Msg(format!("Op::Release expects 1 input, got {}", inputs.len(),)).bt(),
);
}
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::ReleaseMarker,
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if !matches!(
node.op,
Op::WriteSlice { .. }
| Op::WriteSliceRotating { .. }
| Op::WriteSliceDoff { .. }
| Op::ZeroFill
| Op::Release
| Op::Move { .. },
) && let Some(target_idx) = node.op.destructive_input()
{
if inputs.len() <= target_idx {
return Err(Error::Msg(format!(
"in-place op {:?} declares destructive_input={target_idx} \
but has only {} input(s)",
node.op,
inputs.len(),
))
.bt());
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: in-place op node {id:?} has no target_backend set",
))
.bt()
})?;
let op_kind = op_to_op_kind(&node.op).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: in-place op {:?} has no op_to_op_kind mapping",
node.op,
))
.bt()
})?;
let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
let dtypes = build_lookup_dtypes(graph, node);
let compiled = resolve_compiled(op_kind, &dtypes, target_backend, op_params, bindings)?;
let output_layout = graph.layout(inputs[target_idx]);
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs: inputs.clone(),
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::InplaceKernel { target_idx },
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::WriteSlice { .. }) {
if inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::WriteSlice expects 2 inputs (destination, source), got {}",
inputs.len(),
))
.bt());
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSlice node {:?} has no target_backend set",
id
))
.bt()
})?;
let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
let dtypes = build_lookup_dtypes(graph, node);
let compiled = resolve_compiled(
OpKind::WriteSlice,
&dtypes,
target_backend,
op_params,
bindings,
)?;
let output_layout = graph.layout(inputs[0]);
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs: inputs.clone(),
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::WriteSlice {
dest: inputs[0],
source: inputs[1],
},
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::WriteSliceRotating { .. }) {
if inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::WriteSliceRotating expects 3 inputs (destination, source, position), got {}",
inputs.len(),
))
.bt());
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating node {:?} has no target_backend set",
id
))
.bt()
})?;
let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
let dtypes = build_lookup_dtypes(graph, node);
let compiled = resolve_compiled(
OpKind::WriteSliceRotating,
&dtypes,
target_backend,
op_params,
bindings,
)?;
let output_layout = graph.layout(inputs[0]);
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs: inputs.clone(),
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::WriteSliceRotating {
dest: inputs[0],
source: inputs[1],
position: inputs[2],
},
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::WriteSliceDoff { .. }) {
if inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::WriteSliceDoff expects 3 inputs (destination, source, offset), got {}",
inputs.len(),
))
.bt());
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff node {:?} has no target_backend set",
id
))
.bt()
})?;
let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
let dtypes = build_lookup_dtypes(graph, node);
let compiled = resolve_compiled(
OpKind::WriteSliceDoff,
&dtypes,
target_backend,
op_params,
bindings,
)?;
let output_layout = graph.layout(inputs[0]);
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs: inputs.clone(),
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::WriteSliceDoff {
dest: inputs[0],
source: inputs[1],
offset: inputs[2],
},
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::Alloc { target } = node.op {
if !inputs.is_empty() {
return Err(
Error::Msg(format!("Op::Alloc expects 0 inputs, got {}", inputs.len(),)).bt(),
);
}
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = match target {
DeviceLocation::Cpu => BackendId::Cpu,
DeviceLocation::Cuda { .. } => BackendId::Cuda,
DeviceLocation::Vulkan { .. } => BackendId::Vulkan,
DeviceLocation::Metal { .. } => BackendId::Metal,
};
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::Alloc {
target_location: target,
},
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::Iota { len } = node.op {
if !inputs.is_empty() {
return Err(
Error::Msg(format!("Op::Iota expects 0 inputs, got {}", inputs.len(),)).bt(),
);
}
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = graph.target_backend(id).unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::Iota { len },
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::ZeroFill) {
if inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::ZeroFill expects 1 input, got {}",
inputs.len(),
))
.bt());
}
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::ZeroFill,
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::Copy { target } = node.op {
if inputs.len() != 1 {
return Err(
Error::Msg(format!("Op::Copy expects 1 input, got {}", inputs.len(),)).bt(),
);
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: Op::Copy node {:?} has no target_backend \
set (= source backend, the one whose kernel runs the \
download)",
id
))
.bt()
})?;
let op_params = OpParams::None;
let dtypes = build_lookup_dtypes(graph, node);
let compiled =
resolve_compiled(OpKind::Copy, &dtypes, target_backend, op_params, bindings)?;
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::Copy {
target_location: target,
},
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::Move { target } = node.op {
if inputs.len() != 1 {
return Err(
Error::Msg(format!("Op::Move expects 1 input, got {}", inputs.len(),)).bt(),
);
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: Op::Move node {:?} has no target_backend \
set (= source backend, the one whose kernel runs the \
transfer)",
id
))
.bt()
})?;
let op_params = OpParams::None;
let dtypes = build_lookup_dtypes(graph, node);
let compiled =
resolve_compiled(OpKind::Copy, &dtypes, target_backend, op_params, bindings)?;
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::Move {
target_location: target,
},
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::Reshape(_)) {
if inputs.len() != 1 {
return Err(
Error::Msg(format!("Op::Reshape expects 1 input, got {}", inputs.len(),)).bt(),
);
}
let input_layout = graph.layout(inputs[0]);
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
let in_elem_count = input_layout.shape().elem_count();
if in_elem_count != elem_count {
return Err(Error::Msg(format!(
"Op::Reshape changes element count: input {} → output {}",
in_elem_count, elem_count,
))
.bt());
}
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::ContiguizeOf {
input: node.inputs[0],
},
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if matches!(node.op, Op::Contiguize) {
if inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::Contiguize expects 1 input, got {}",
inputs.len(),
))
.bt());
}
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let target_backend = graph
.target_backend(id)
.or_else(|| graph.target_backend(inputs[0]))
.unwrap_or(BackendId::Cpu);
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::ContiguizeOf {
input: node.inputs[0],
},
compiled: None,
output_layout,
destructive_input,
output_bundle: None,
});
}
if let Op::Fused(fid, _) = &node.op
&& fid.is_runtime()
{
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: runtime fused op {:?} (node {:?}) has no \
target_backend — the arm emitter pins it at emission",
fid, id,
))
.bt()
})?;
if graph.output_views_arc(id).is_some() {
return Err(Error::Msg(format!(
"PipelinedExecutor: runtime fused op {:?} (node {:?}) carries a \
multi-output bundle — synthesized kernels are single-output in v1",
fid, id,
))
.bt());
}
let dtypes = build_lookup_dtypes(graph, node);
let (kernel, caps) = bindings.lookup_with_caps(
crate::kernel::BindingKey::RuntimeFused(*fid),
&dtypes,
target_backend,
)?;
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let op_params = match &node.op {
Op::Fused(_, FusedOpParams::Runtime { scalars }) if !scalars.is_empty() => {
OpParams::JitScalars {
scalars: scalars.clone(),
}
}
_ => OpParams::None,
};
let compiled = CompiledNode {
op: OpKind::RuntimeFused,
dtypes: KernelDTypes::from_slice(&dtypes),
backend: target_backend,
kernel,
caps,
op_params,
};
return Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::Kernel,
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle: None,
});
}
let target_backend = graph.target_backend(id).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: node {:?} ({:?}) has no target_backend set",
id, node.op
))
.bt()
})?;
let op_kind = op_to_op_kind(&node.op).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: op {:?} not yet mapped to an OpKind \
(Phase C migrates more ops as they're registered)",
node.op,
))
.bt()
})?;
let op_params = op_to_op_params(graph, node, id, layout_cache, sym_env)?;
let dtypes = build_lookup_dtypes(graph, node);
let compiled = resolve_compiled(op_kind, &dtypes, target_backend, op_params, bindings)?;
let output_layout = Layout::contiguous(node.shape.clone());
layout_cache.insert(id, output_layout.clone());
let output_bundle = graph.output_views_arc(id);
Ok(WorkItem {
node_id: id,
inputs,
elem_count,
dtype: node.dtype,
target_backend,
kind: WorkItemKind::Kernel,
compiled: Some(compiled),
output_layout,
destructive_input,
output_bundle,
})
}
pub(crate) fn build_lookup_dtypes(graph: &Graph, node: &Node) -> Vec<DType> {
if matches!(node.op, Op::Concat { .. }) {
let in_dt = node
.inputs
.first()
.map(|&id| graph.node(id).dtype)
.unwrap_or(node.dtype);
return vec![in_dt, node.dtype];
}
if matches!(node.op, Op::WriteSlice { .. }) {
let src_dt = node
.inputs
.get(1)
.map(|&id| graph.node(id).dtype)
.unwrap_or(node.dtype);
return vec![src_dt, node.dtype];
}
if matches!(node.op, Op::WriteSliceRotating { .. }) {
let src_dt = node
.inputs
.get(1)
.map(|&id| graph.node(id).dtype)
.unwrap_or(node.dtype);
return vec![src_dt, node.dtype];
}
if matches!(node.op, Op::WriteSliceDoff { .. }) {
let src_dt = node
.inputs
.get(1)
.map(|&id| graph.node(id).dtype)
.unwrap_or(node.dtype);
return vec![src_dt, node.dtype];
}
let mut dts: Vec<DType> = node.inputs.iter().map(|&id| graph.node(id).dtype).collect();
dts.push(node.dtype);
dts
}
fn op_is_softmax_last_dim(op: &Op) -> bool {
matches!(op, Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM)
}
fn op_is_fused_linear(op: &Op) -> bool {
matches!(op, Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::FUSED_LINEAR)
}
pub(crate) fn op_to_op_kind(op: &Op) -> Option<OpKind> {
match op {
Op::Add => Some(OpKind::AddElementwise),
Op::Sub => Some(OpKind::SubElementwise),
Op::Mul => Some(OpKind::MulElementwise),
Op::Div => Some(OpKind::DivElementwise),
Op::Relu => Some(OpKind::ReluElementwise),
Op::Neg => Some(OpKind::NegElementwise),
Op::Sqr => Some(OpKind::SqrElementwise),
Op::Sqrt => Some(OpKind::SqrtElementwise),
Op::Tanh => Some(OpKind::TanhElementwise),
Op::Exp => Some(OpKind::ExpElementwise),
Op::Log => Some(OpKind::LogElementwise),
Op::Sin => Some(OpKind::SinElementwise),
Op::Cos => Some(OpKind::CosElementwise),
Op::Sigmoid => Some(OpKind::SigmoidElementwise),
Op::Silu => Some(OpKind::SiluElementwise),
Op::GeluTanh => Some(OpKind::GeluTanhElementwise),
Op::Step => Some(OpKind::StepElementwise),
Op::Recip => Some(OpKind::RecipElementwise),
Op::Abs => Some(OpKind::AbsElementwise),
Op::Equal => Some(OpKind::EqualElementwise),
Op::Ne => Some(OpKind::NotEqualElementwise),
Op::Lt => Some(OpKind::LessElementwise),
Op::Le => Some(OpKind::LessEqualElementwise),
Op::Gt => Some(OpKind::GreaterElementwise),
Op::Ge => Some(OpKind::GreaterEqualElementwise),
Op::Where => Some(OpKind::Where),
Op::Floor => Some(OpKind::FloorElementwise),
Op::Ceil => Some(OpKind::CeilElementwise),
Op::Round => Some(OpKind::RoundElementwise),
Op::Sign => Some(OpKind::SignElementwise),
Op::Erf => Some(OpKind::ErfElementwise),
Op::Gelu => Some(OpKind::GeluElementwise),
Op::Pow => Some(OpKind::PowElementwise),
Op::Rsqrt => Some(OpKind::RsqrtElementwise),
Op::Rem => Some(OpKind::RemElementwise),
Op::Flip { .. } => Some(OpKind::Flip),
Op::Roll { .. } => Some(OpKind::Roll),
Op::CumSum { .. } => Some(OpKind::CumSum),
Op::Triu { .. } => Some(OpKind::Triu),
Op::Tril { .. } => Some(OpKind::Tril),
Op::LogSoftmaxLastDim => Some(OpKind::LogSoftmaxLastDim),
Op::LogSoftmaxLastDimBackward => Some(OpKind::LogSoftmaxLastDimBackward),
Op::MaskedFill { .. } => Some(OpKind::MaskedFill),
Op::Pad { .. } => Some(OpKind::Pad),
Op::PadBackward { .. } => Some(OpKind::PadBackward),
Op::SumDim(_) => Some(OpKind::SumReduce),
Op::MaxDim(_) => Some(OpKind::MaxReduce),
Op::MinDim(_) => Some(OpKind::MinReduce),
Op::MeanDim(_) => Some(OpKind::MeanReduce),
Op::SumAll => Some(OpKind::SumReduce),
Op::MaxAll => Some(OpKind::MaxReduce),
Op::MinAll => Some(OpKind::MinReduce),
Op::MeanAll => Some(OpKind::MeanReduce),
Op::MatMul => Some(OpKind::MatMul),
Op::Cast(_) => Some(OpKind::Cast),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::CONV2D => Some(OpKind::Conv2D),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D => {
Some(OpKind::ConvTranspose2D)
}
Op::ReduceSumTo(_) => Some(OpKind::ReduceSumTo),
Op::ReduceMaxTo(_) => Some(OpKind::ReduceMaxTo),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::FLASH_ATTN => {
Some(OpKind::FlashAttn)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::PAGED_ATTN => {
Some(OpKind::PagedAttn)
}
Op::AddScalar(_) => Some(OpKind::Affine),
Op::MulScalar(_) => Some(OpKind::Affine),
Op::Clamp { .. } => Some(OpKind::ClampElementwise),
Op::PowI(_) => Some(OpKind::PowIElementwise),
Op::Maximum => Some(OpKind::MaximumElementwise),
Op::Minimum => Some(OpKind::MinimumElementwise),
Op::Concat { .. } => Some(OpKind::Concat),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM => {
Some(OpKind::SoftmaxLastDim)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM_BACKWARD => {
Some(OpKind::SoftmaxLastDimBackward)
}
Op::Fused(fid, _)
if *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM_BACKWARD =>
{
Some(OpKind::LayerNormLastDimBackward)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM_BACKWARD => {
Some(OpKind::RmsNormLastDimBackward)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::REDUCE_MAX_TO_BACKWARD => {
Some(OpKind::ReduceMaxToBackward)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::POWI_BACKWARD => {
Some(OpKind::PowIElementwiseBackward)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::FUSED_LINEAR => {
Some(OpKind::FusedLinear)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM => {
Some(OpKind::RmsNormLastDim)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM => {
Some(OpKind::LayerNormLastDim)
}
Op::IndexSelect { .. } => Some(OpKind::IndexSelect),
Op::Gather { .. } => Some(OpKind::Gather),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::ROPE => Some(OpKind::Rope),
Op::IndexAdd { .. } => Some(OpKind::IndexAdd),
Op::ScatterAdd { .. } => Some(OpKind::ScatterAdd),
Op::NonZeroIndices { .. } => Some(OpKind::NonZeroIndices),
Op::ArgMaxDim(_) => Some(OpKind::ArgMaxDim),
Op::ArgMinDim(_) => Some(OpKind::ArgMinDim),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::QMATMUL => {
Some(OpKind::QMatMul)
}
Op::WriteSlice { .. } => Some(OpKind::WriteSlice),
Op::WriteSliceRotating { .. } => Some(OpKind::WriteSliceRotating),
Op::WriteSliceDoff { .. } => Some(OpKind::WriteSliceDoff),
Op::Copy { .. } => Some(OpKind::Copy),
Op::Move { .. } => Some(OpKind::Copy),
Op::ReluInplace => Some(OpKind::ReluInplace),
Op::SiluInplace => Some(OpKind::SiluInplace),
Op::GeluTanhInplace => Some(OpKind::GeluTanhInplace),
Op::TanhInplace => Some(OpKind::TanhInplace),
Op::SigmoidInplace => Some(OpKind::SigmoidInplace),
Op::NegInplace => Some(OpKind::NegInplace),
Op::AbsInplace => Some(OpKind::AbsInplace),
Op::SqrInplace => Some(OpKind::SqrInplace),
Op::SqrtInplace => Some(OpKind::SqrtInplace),
Op::RsqrtInplace => Some(OpKind::RsqrtInplace),
Op::RecipInplace => Some(OpKind::RecipInplace),
Op::ExpInplace => Some(OpKind::ExpInplace),
Op::LogInplace => Some(OpKind::LogInplace),
Op::SinInplace => Some(OpKind::SinInplace),
Op::CosInplace => Some(OpKind::CosInplace),
Op::SignInplace => Some(OpKind::SignInplace),
Op::FloorInplace => Some(OpKind::FloorInplace),
Op::CeilInplace => Some(OpKind::CeilInplace),
Op::RoundInplace => Some(OpKind::RoundInplace),
Op::ErfInplace => Some(OpKind::ErfInplace),
Op::GeluInplace => Some(OpKind::GeluInplace),
Op::ClampInplace { .. } => Some(OpKind::ClampInplace),
Op::PowIInplace(_) => Some(OpKind::PowIInplace),
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::INPLACE_AFFINE => {
Some(OpKind::InplaceAffine)
}
Op::Fused(fid, _)
if *fid == fuel_graph::registry::FusedOps::FUSED_SOFTMAX_CROSS_ENTROPY =>
{
Some(OpKind::FusedSoftmaxCrossEntropy)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::CAUSAL_CONV1D => {
Some(OpKind::CausalConv1d)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SELECTIVE_SCAN => {
Some(OpKind::SelectiveScan)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SSD_CHUNK_SCAN => {
Some(OpKind::SsdChunkScan)
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::NF4_MATMUL => {
Some(OpKind::Nf4Matmul)
}
Op::Alloc { .. } | Op::ZeroFill => None,
_ => None,
}
}
fn encode_value_to_bytes(dtype: DType, value: f64) -> Result<Vec<u8>> {
match dtype {
DType::F32 => Ok((value as f32).to_le_bytes().to_vec()),
DType::F64 => Ok(value.to_le_bytes().to_vec()),
DType::BF16 => Ok(half::bf16::from_f32(value as f32).to_le_bytes().to_vec()),
DType::F16 => Ok(half::f16::from_f32(value as f32).to_le_bytes().to_vec()),
DType::U8 => Ok(vec![value as u8]),
DType::U32 => Ok((value as u32).to_le_bytes().to_vec()),
other => Err(Error::Msg(format!(
"encode_value_to_bytes: dtype {other:?} not yet supported for Pad fill",
))
.bt()),
}
}
fn scalar_to_bytes(s: fuel_ir::Scalar) -> Vec<u8> {
use fuel_ir::Scalar;
match s {
Scalar::U8(v) => vec![v],
Scalar::I8(v) => vec![v as u8],
Scalar::U32(v) => v.to_le_bytes().to_vec(),
Scalar::I16(v) => v.to_le_bytes().to_vec(),
Scalar::I32(v) => v.to_le_bytes().to_vec(),
Scalar::I64(v) => v.to_le_bytes().to_vec(),
Scalar::BF16(v) => v.to_le_bytes().to_vec(),
Scalar::F16(v) => v.to_le_bytes().to_vec(),
Scalar::F32(v) => v.to_le_bytes().to_vec(),
Scalar::F64(v) => v.to_le_bytes().to_vec(),
Scalar::Bool(b) => vec![b as u8],
Scalar::F8E4M3(v) => vec![v.to_bits()],
Scalar::F8E5M2(v) => vec![v.to_bits()],
Scalar::F8E8M0(v) => vec![v],
}
}
fn op_to_op_params(
graph: &Graph,
node: &Node,
node_id: NodeId,
layout_cache: &HashMap<NodeId, Layout>,
sym_env: &SymEnv,
) -> Result<OpParams> {
let input_layout = |input_id: NodeId| -> Layout {
layout_cache
.get(&input_id)
.cloned()
.unwrap_or_else(|| graph.layout(input_id))
};
Ok(match &node.op {
Op::Fused(_, fuel_graph::registry::FusedOpParams::QMatMul { quant_type, k, n }) => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::QMatMul expects 2 inputs (activations, weight_bytes), got {}",
node.inputs.len(),
))
.bt());
}
let act_layout = input_layout(node.inputs[0]);
let act_dims = act_layout.shape().dims();
if act_dims.len() < 2 {
return Err(Error::Msg(format!(
"Op::QMatMul: activations must be rank ≥ 2, got {act_dims:?}",
))
.bt());
}
let m = act_dims[act_dims.len() - 2];
let k_act = act_dims[act_dims.len() - 1];
if k_act != *k {
return Err(Error::Msg(format!(
"Op::QMatMul: activation last dim ({k_act}) must equal Op's k ({k})",
))
.bt());
}
let batch_count: usize = act_dims[..act_dims.len() - 2].iter().product();
OpParams::QMatMul {
quant_type: *quant_type,
batch_count,
m,
n: *n,
k: *k,
}
}
Op::ArgMaxDim(d) | Op::ArgMinDim(d) => {
OpParams::Reduce {
dims: vec![*d],
keepdim: false,
}
}
Op::SumDim(d) | Op::MaxDim(d) | Op::MinDim(d) | Op::MeanDim(d) => OpParams::Reduce {
dims: vec![*d],
keepdim: false,
},
Op::SumAll | Op::MaxAll | Op::MinAll | Op::MeanAll => {
let il = input_layout(node.inputs[0]);
let rank = il.shape().rank();
OpParams::Reduce {
dims: (0..rank).collect(),
keepdim: false,
}
}
Op::MatMul => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::MatMul expects 2 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let lhs = input_layout(node.inputs[0]);
let rhs = input_layout(node.inputs[1]);
let lhs_dims = lhs.shape().dims();
let rhs_dims = rhs.shape().dims();
if lhs_dims.len() < 2 || rhs_dims.len() < 2 {
return Err(Error::Msg(format!(
"Op::MatMul requires both inputs rank ≥ 2; got lhs={:?} rhs={:?}",
lhs_dims, rhs_dims,
))
.bt());
}
if lhs_dims.len() != rhs_dims.len() {
return Err(Error::Msg(format!(
"Op::MatMul: ranks must match (auto-broadcast happens at \
graph construction time); got lhs rank {} vs rhs rank {}",
lhs_dims.len(),
rhs_dims.len(),
))
.bt());
}
let rank = lhs_dims.len();
let batch_rank = rank - 2;
for i in 0..batch_rank {
let la = lhs_dims[i];
let ra = rhs_dims[i];
let ok = la == ra || (ra > 0 && la > ra && la % ra == 0);
if !ok {
return Err(Error::Msg(format!(
"Op::MatMul: batch dim {i} disallowed combination \
(lhs={la}, rhs={ra}); must be equal or \
GQA-divisible (lhs > rhs && lhs % rhs == 0)",
))
.bt());
}
}
let lhs_batch_dims: Vec<usize> = lhs_dims[..batch_rank].to_vec();
let rhs_batch_dims: Vec<usize> = rhs_dims[..batch_rank].to_vec();
let (m, k_lhs) = (lhs_dims[rank - 2], lhs_dims[rank - 1]);
let (k_rhs, n) = (rhs_dims[rank - 2], rhs_dims[rank - 1]);
if k_lhs != k_rhs {
return Err(Error::Msg(format!(
"Op::MatMul: contracting dims disagree — lhs trailing is \
[{m}, {k_lhs}], rhs trailing is [{k_rhs}, {n}]",
))
.bt());
}
let m_compute = match graph.matmul_row_count(node_id) {
None => MatmulM::All,
Some(dyn_m) => match dyn_m.resolve(sym_env) {
Some(c) => {
if c > m {
return Err(Error::Msg(format!(
"Op::MatMul: resolved row count {c} exceeds capacity {m}",
))
.bt());
}
MatmulM::Rows(c)
}
None => match dyn_m {
fuel_ir::DynScalar::Sym(s) if graph.binds_data_determined_sym(s) => {
MatmulM::Deferred(dyn_m)
}
_ => {
return Err(Error::Msg(format!(
"Op::MatMul: dynamic row count {dyn_m:?} is unbound in the \
SymEnv at compile and no producer binds it (not \
data-determined)",
))
.bt());
}
},
},
};
OpParams::Matmul {
lhs_batch_dims,
rhs_batch_dims,
m,
n,
k: k_lhs,
m_compute,
}
}
op if op_is_fused_linear(op) => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::FusedLinear expects 3 inputs (a, b, bias), got {}",
node.inputs.len(),
))
.bt());
}
let lhs = input_layout(node.inputs[0]);
let rhs = input_layout(node.inputs[1]);
let bias = input_layout(node.inputs[2]);
let lhs_dims = lhs.shape().dims();
let rhs_dims = rhs.shape().dims();
let bias_dims = bias.shape().dims();
if lhs_dims.len() < 2 || rhs_dims.len() < 2 {
return Err(Error::Msg(format!(
"Op::FusedLinear requires both a, b rank ≥ 2; got a={:?} b={:?}",
lhs_dims, rhs_dims,
))
.bt());
}
if lhs_dims.len() != rhs_dims.len() {
return Err(Error::Msg(format!(
"Op::FusedLinear: ranks must match (auto-broadcast happens at \
graph construction time); got a rank {} vs b rank {}",
lhs_dims.len(),
rhs_dims.len(),
))
.bt());
}
let rank = lhs_dims.len();
let batch_rank = rank - 2;
for i in 0..batch_rank {
let la = lhs_dims[i];
let ra = rhs_dims[i];
let ok = la == ra || (ra > 0 && la > ra && la % ra == 0);
if !ok {
return Err(Error::Msg(format!(
"Op::FusedLinear: batch dim {i} disallowed (a={la}, b={ra})",
))
.bt());
}
}
let lhs_batch_dims: Vec<usize> = lhs_dims[..batch_rank].to_vec();
let rhs_batch_dims: Vec<usize> = rhs_dims[..batch_rank].to_vec();
let (m, k_lhs) = (lhs_dims[rank - 2], lhs_dims[rank - 1]);
let (k_rhs, n) = (rhs_dims[rank - 2], rhs_dims[rank - 1]);
if k_lhs != k_rhs {
return Err(Error::Msg(format!(
"Op::FusedLinear: contracting dims disagree — a trailing is \
[{m}, {k_lhs}], b trailing is [{k_rhs}, {n}]",
))
.bt());
}
if bias_dims.len() != 1 || bias_dims[0] != n {
return Err(Error::Msg(format!(
"Op::FusedLinear: bias must be rank-1 [{n}], got {bias_dims:?}",
))
.bt());
}
OpParams::Matmul {
lhs_batch_dims,
rhs_batch_dims,
m,
n,
k: k_lhs,
m_compute: MatmulM::All,
}
}
Op::Fused(
_,
fuel_graph::registry::FusedOpParams::FusedSoftmaxCrossEntropy {
reduction,
ignore_index,
},
) => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"FusedSoftmaxCrossEntropy expects 2 inputs (logits, targets), got {}",
node.inputs.len(),
))
.bt());
}
let logits_layout = input_layout(node.inputs[0]);
let logits_dims = logits_layout.shape().dims();
if logits_dims.is_empty() {
return Err(Error::Msg(
"FusedSoftmaxCrossEntropy: logits must have rank ≥ 1".to_string(),
)
.bt());
}
let vocab = *logits_dims.last().unwrap();
let n_rows: usize = logits_dims[..logits_dims.len() - 1]
.iter()
.product::<usize>()
.max(1);
let targets_layout = input_layout(node.inputs[1]);
let target_count: usize = targets_layout
.shape()
.dims()
.iter()
.product::<usize>()
.max(1);
if target_count != n_rows {
return Err(Error::Msg(format!(
"FusedSoftmaxCrossEntropy: targets element count {target_count} must \
equal logits row count {n_rows} (logits {logits_dims:?}, targets {:?})",
targets_layout.shape().dims(),
))
.bt());
}
OpParams::FusedSoftmaxCrossEntropy {
n_rows,
vocab,
reduction: *reduction,
ignore_index: *ignore_index,
}
}
Op::Fused(_, fuel_graph::registry::FusedOpParams::CausalConv1d { use_silu }) => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"CausalConv1d expects 3 inputs (x, weight, bias), got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let w_layout = input_layout(node.inputs[1]);
let x_dims = x_layout.shape().dims();
let w_dims = w_layout.shape().dims();
if x_dims.len() != 3 {
return Err(Error::Msg(format!(
"CausalConv1d: x must be rank 3 [batch, channels, seq+pad], got {x_dims:?}",
))
.bt());
}
if w_dims.len() != 3 {
return Err(Error::Msg(format!(
"CausalConv1d: weight must be rank 3 [channels, 1, kernel], got {w_dims:?}",
))
.bt());
}
let batch = x_dims[0];
let channels = x_dims[1];
let seq_in = x_dims[2];
let kernel = w_dims[2];
if seq_in < kernel - 1 {
return Err(Error::Msg(format!(
"CausalConv1d: x time dim {seq_in} must be ≥ kernel-1 = {} \
(caller must pre-pad with kernel-1 zeros)",
kernel - 1,
))
.bt());
}
let seq_out = seq_in - (kernel - 1);
OpParams::CausalConv1d {
batch,
channels,
seq_in,
seq_out,
kernel,
use_silu: *use_silu,
}
}
Op::Fused(_, fuel_graph::registry::FusedOpParams::SelectiveScan { delta_softplus }) => {
if node.inputs.len() != 5 {
return Err(Error::Msg(format!(
"SelectiveScan expects 5 inputs (u, delta, a, b, c), got {}",
node.inputs.len(),
))
.bt());
}
let u_layout = input_layout(node.inputs[0]);
let a_layout = input_layout(node.inputs[2]);
let u_dims = u_layout.shape().dims();
let a_dims = a_layout.shape().dims();
if u_dims.len() != 3 {
return Err(Error::Msg(format!(
"SelectiveScan: u must be rank 3 [batch, seqlen, dim], got {u_dims:?}",
))
.bt());
}
if a_dims.len() != 2 {
return Err(Error::Msg(format!(
"SelectiveScan: a must be rank 2 [dim, dstate], got {a_dims:?}",
))
.bt());
}
let batch = u_dims[0];
let seqlen = u_dims[1];
let dim = u_dims[2];
let dstate = a_dims[1];
if a_dims[0] != dim {
return Err(Error::Msg(format!(
"SelectiveScan: a's first dim {} must equal dim {dim}",
a_dims[0],
))
.bt());
}
OpParams::SelectiveScan {
batch,
seqlen,
dim,
dstate,
delta_softplus: *delta_softplus,
}
}
Op::Fused(_, fuel_graph::registry::FusedOpParams::SsdChunkScan { chunk_size }) => {
if node.inputs.len() != 5 {
return Err(Error::Msg(format!(
"SsdChunkScan expects 5 inputs (x, dt, a, b, c), got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let b_layout = input_layout(node.inputs[3]);
let x_dims = x_layout.shape().dims();
let b_dims = b_layout.shape().dims();
if x_dims.len() != 4 {
return Err(Error::Msg(format!(
"SsdChunkScan: x must be rank 4 [batch, seqlen, heads, head_dim], got {x_dims:?}",
))
.bt());
}
if b_dims.len() != 4 {
return Err(Error::Msg(format!(
"SsdChunkScan: b must be rank 4 [batch, seqlen, heads, state_dim], got {b_dims:?}",
))
.bt());
}
let batch = x_dims[0];
let seqlen = x_dims[1];
let heads = x_dims[2];
let head_dim = x_dims[3];
let state_dim = b_dims[3];
OpParams::SsdChunkScan {
batch,
seqlen,
heads,
head_dim,
state_dim,
chunk_size: *chunk_size,
}
}
Op::Fused(_, fuel_graph::registry::FusedOpParams::Nf4Matmul { block_size }) => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Nf4Matmul expects 3 inputs (activations, w_packed, absmax), got {}",
node.inputs.len(),
))
.bt());
}
let a_layout = input_layout(node.inputs[0]);
let w_layout = input_layout(node.inputs[1]);
let a_dims = a_layout.shape().dims();
let w_dims = w_layout.shape().dims();
if a_dims.len() < 2 {
return Err(Error::Msg(format!(
"Nf4Matmul: activations must be rank ≥ 2, got {a_dims:?}",
))
.bt());
}
if w_dims.len() != 2 {
return Err(Error::Msg(format!(
"Nf4Matmul: w_packed must be rank 2 [n, k/2], got {w_dims:?}",
))
.bt());
}
let m = a_dims[a_dims.len() - 2];
let k = a_dims[a_dims.len() - 1];
let n = w_dims[0];
let batch: usize = a_dims[..a_dims.len() - 2].iter().product::<usize>().max(1);
if k % 2 != 0 {
return Err(Error::Msg(format!(
"Nf4Matmul: k={k} must be even (w_packed holds 2 nibbles per byte along k)",
))
.bt());
}
if k % *block_size != 0 {
return Err(Error::Msg(format!(
"Nf4Matmul: k={k} must be a multiple of block_size={block_size}",
))
.bt());
}
if w_dims[1] != k / 2 {
return Err(Error::Msg(format!(
"Nf4Matmul: w_packed second dim {} must equal k/2 = {}",
w_dims[1],
k / 2,
))
.bt());
}
OpParams::Nf4Matmul {
batch,
m,
n,
k,
block_size: *block_size,
}
}
op if op_is_softmax_last_dim(op)
|| matches!(op, Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM_BACKWARD) =>
{
let il = input_layout(node.inputs[0]);
let dims = il.shape().dims();
if dims.is_empty() {
return Err(Error::Msg("Op::SoftmaxLastDim requires rank ≥ 1".to_string()).bt());
}
let last_dim = *dims.last().unwrap();
let outer_count: usize = dims[..dims.len() - 1].iter().product();
OpParams::SoftmaxLastDim {
outer_count,
last_dim,
}
}
Op::Flip { dim } => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::Flip expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims = in_layout.shape().dims();
if *dim >= in_dims.len() {
return Err(Error::Msg(format!(
"Op::Flip: dim {dim} out of range for rank {}",
in_dims.len(),
))
.bt());
}
let outer_count: usize = in_dims[..*dim].iter().product();
let dim_size = in_dims[*dim];
let inner_count: usize = in_dims[*dim + 1..].iter().product();
OpParams::Flip {
outer_count,
dim_size,
inner_count,
axis: *dim,
}
}
Op::Roll { dim, shift } => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::Roll expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims = in_layout.shape().dims();
if *dim >= in_dims.len() {
return Err(Error::Msg(format!(
"Op::Roll: dim {dim} out of range for rank {}",
in_dims.len(),
))
.bt());
}
let outer_count: usize = in_dims[..*dim].iter().product();
let dim_size = in_dims[*dim];
let inner_count: usize = in_dims[*dim + 1..].iter().product();
OpParams::Roll {
outer_count,
dim_size,
inner_count,
shift: *shift,
axis: *dim,
}
}
Op::CumSum { dim } => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::CumSum expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims = in_layout.shape().dims();
if *dim >= in_dims.len() {
return Err(Error::Msg(format!(
"Op::CumSum: dim {dim} out of range for rank {}",
in_dims.len(),
))
.bt());
}
let outer_count: usize = in_dims[..*dim].iter().product();
let dim_size = in_dims[*dim];
let inner_count: usize = in_dims[*dim + 1..].iter().product();
OpParams::CumSum {
outer_count,
dim_size,
inner_count,
axis: *dim,
}
}
Op::Triu { diagonal } | Op::Tril { diagonal } => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::Triu/Tril expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims = in_layout.shape().dims();
if in_dims.len() < 2 {
return Err(Error::Msg(format!(
"Op::Triu/Tril requires rank >= 2, got {}",
in_dims.len(),
))
.bt());
}
let rows = in_dims[in_dims.len() - 2];
let cols = in_dims[in_dims.len() - 1];
let batch_count: usize = in_dims[..in_dims.len() - 2]
.iter()
.product::<usize>()
.max(1);
OpParams::Triangular {
batch_count,
rows,
cols,
diagonal: *diagonal,
}
}
Op::LogSoftmaxLastDim => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::LogSoftmaxLastDim expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims = in_layout.shape().dims();
if in_dims.is_empty() {
return Err(Error::Msg("Op::LogSoftmaxLastDim requires rank ≥ 1".to_string()).bt());
}
let last_dim = *in_dims.last().unwrap();
let outer_count: usize = in_dims[..in_dims.len() - 1].iter().product();
OpParams::LogSoftmaxLastDim {
outer_count,
last_dim,
}
}
Op::LogSoftmaxLastDimBackward => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::LogSoftmaxLastDimBackward expects 2 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims = in_layout.shape().dims();
if in_dims.is_empty() {
return Err(Error::Msg(
"Op::LogSoftmaxLastDimBackward requires rank ≥ 1".to_string(),
)
.bt());
}
let last_dim = *in_dims.last().unwrap();
let outer_count: usize = in_dims[..in_dims.len() - 1].iter().product();
OpParams::LogSoftmaxLastDim {
outer_count,
last_dim,
}
}
Op::MaskedFill { value } => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::MaskedFill expects 2 inputs (x, mask), got {}",
node.inputs.len(),
))
.bt());
}
let fill_bytes = scalar_to_bytes(*value);
OpParams::MaskedFill { fill_bytes }
}
Op::PadBackward {
in_shape,
padding,
mode,
} => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::PadBackward expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_dims: Vec<usize> = in_shape.dims().to_vec();
let out_dims: Vec<usize> = in_dims
.iter()
.zip(padding.iter())
.map(|(&d, &(b, a))| d + b + a)
.collect();
let mode_tag: u8 = match mode {
fuel_graph::PadMode::Constant => 0,
fuel_graph::PadMode::Reflect => 1,
fuel_graph::PadMode::Replicate => 2,
};
OpParams::PadBackward {
in_shape: in_dims,
out_shape: out_dims,
padding: padding.clone(),
mode_tag,
}
}
Op::Pad {
padding,
mode,
value,
} => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::Pad expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let in_dims: Vec<usize> = in_layout.shape().dims().to_vec();
if padding.len() != in_dims.len() {
return Err(Error::Msg(format!(
"Op::Pad: padding.len() ({}) != input rank ({})",
padding.len(),
in_dims.len(),
))
.bt());
}
let out_dims: Vec<usize> = in_dims
.iter()
.zip(padding.iter())
.map(|(&d, &(b, a))| d + b + a)
.collect();
let mode_tag: u8 = match mode {
fuel_graph::PadMode::Constant => 0,
fuel_graph::PadMode::Reflect => 1,
fuel_graph::PadMode::Replicate => 2,
};
let fill_bytes = encode_value_to_bytes(node.dtype, *value)?;
OpParams::Pad {
in_shape: in_dims,
out_shape: out_dims,
padding: padding.clone(),
mode_tag,
fill_bytes,
}
}
Op::IndexAdd { dim } => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::IndexAdd expects 3 inputs (base, indices, src), got {}",
node.inputs.len(),
))
.bt());
}
let base_layout = input_layout(node.inputs[0]);
let idx_layout = input_layout(node.inputs[1]);
let src_layout = input_layout(node.inputs[2]);
let base_dims = base_layout.shape().dims();
let idx_dims = idx_layout.shape().dims();
let src_dims = src_layout.shape().dims();
if *dim >= base_dims.len() {
return Err(Error::Msg(format!(
"Op::IndexAdd: dim {dim} out of range for base rank {}",
base_dims.len(),
))
.bt());
}
if idx_dims.len() != 1 {
return Err(Error::Msg(format!(
"Op::IndexAdd: indices must be rank 1, got {idx_dims:?}",
))
.bt());
}
if base_dims.len() != src_dims.len() {
return Err(Error::Msg(format!(
"Op::IndexAdd: base rank ({}) != src rank ({})",
base_dims.len(),
src_dims.len(),
))
.bt());
}
let outer_count: usize = base_dims[..*dim].iter().product();
let base_dim_size = base_dims[*dim];
let inner_count: usize = base_dims[*dim + 1..].iter().product();
let n_indices = idx_dims[0];
OpParams::IndexAdd {
outer_count,
base_dim_size,
n_indices,
inner_count,
}
}
Op::ScatterAdd { dim } => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::ScatterAdd expects 3 inputs (base, indices, src), got {}",
node.inputs.len(),
))
.bt());
}
let base_layout = input_layout(node.inputs[0]);
let src_layout = input_layout(node.inputs[2]);
let base_shape: Vec<usize> = base_layout.shape().dims().to_vec();
let src_shape: Vec<usize> = src_layout.shape().dims().to_vec();
if base_shape.len() != src_shape.len() {
return Err(Error::Msg(format!(
"Op::ScatterAdd: base rank ({}) != src rank ({})",
base_shape.len(),
src_shape.len(),
))
.bt());
}
if *dim >= base_shape.len() {
return Err(Error::Msg(format!(
"Op::ScatterAdd: dim {dim} out of range for rank {}",
base_shape.len(),
))
.bt());
}
OpParams::ScatterAdd {
base_shape,
src_shape,
dim: *dim,
}
}
Op::Fused(_, fuel_graph::registry::FusedOpParams::Rope) => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::Rope expects 3 inputs (x, cos, sin), got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let x_dims = x_layout.shape().dims();
if x_dims.len() < 2 {
return Err(
Error::Msg(format!("Op::Rope: x must have rank ≥ 2, got {x_dims:?}",)).bt(),
);
}
let head_dim = *x_dims.last().unwrap();
let seq = x_dims[x_dims.len() - 2];
let outer_count: usize = x_dims[..x_dims.len() - 2].iter().product();
OpParams::Rope {
outer_count,
seq,
head_dim,
}
}
Op::Gather { dim } => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::Gather expects 2 inputs (source, indices), got {}",
node.inputs.len(),
))
.bt());
}
let src_layout = input_layout(node.inputs[0]);
let idx_layout = input_layout(node.inputs[1]);
let source_shape: Vec<usize> = src_layout.shape().dims().to_vec();
let output_shape: Vec<usize> = idx_layout.shape().dims().to_vec();
if source_shape.len() != output_shape.len() {
return Err(Error::Msg(format!(
"Op::Gather: source rank ({}) != indices rank ({})",
source_shape.len(),
output_shape.len(),
))
.bt());
}
if *dim >= source_shape.len() {
return Err(Error::Msg(format!(
"Op::Gather: dim {dim} out of range for rank {}",
source_shape.len(),
))
.bt());
}
OpParams::Gather {
source_shape,
output_shape,
dim: *dim,
}
}
Op::NonZeroIndices { count_sym } => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::NonZeroIndices expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let capacity = x_layout.shape().elem_count();
OpParams::NonZeroIndices {
capacity,
count_sym: *count_sym,
}
}
Op::IndexSelect { dim } => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::IndexSelect expects 2 inputs (source, indices), got {}",
node.inputs.len(),
))
.bt());
}
let src_layout = input_layout(node.inputs[0]);
let idx_layout = input_layout(node.inputs[1]);
let src_dims = src_layout.shape().dims();
let idx_dims = idx_layout.shape().dims();
if *dim >= src_dims.len() {
return Err(Error::Msg(format!(
"Op::IndexSelect: dim {dim} out of range for source rank {}",
src_dims.len(),
))
.bt());
}
if idx_dims.len() != 1 {
return Err(Error::Msg(format!(
"Op::IndexSelect: indices must be rank 1, got shape {idx_dims:?}",
))
.bt());
}
let outer_count: usize = src_dims[..*dim].iter().product();
let source_dim_size = src_dims[*dim];
let inner_count: usize = src_dims[*dim + 1..].iter().product();
let n_indices = idx_dims[0];
OpParams::IndexSelect {
outer_count,
source_dim_size,
n_indices,
inner_count,
}
}
Op::Fused(fid, params)
if *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM
|| *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM
|| *fid == fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM_BACKWARD
|| *fid == fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM_BACKWARD =>
{
let eps = match params {
fuel_graph::registry::FusedOpParams::RmsNormLastDim { eps } => *eps,
fuel_graph::registry::FusedOpParams::LayerNormLastDim { eps } => *eps,
fuel_graph::registry::FusedOpParams::RmsNormLastDimBackward { eps } => *eps,
fuel_graph::registry::FusedOpParams::LayerNormLastDimBackward { eps } => *eps,
_ => {
return Err(Error::Msg(format!(
"Op::Fused(NORM_LAST_DIM[_BACKWARD], _) expected \
RmsNormLastDim[Backward] or LayerNormLastDim[Backward] params, got {params:?}",
))
.bt());
}
};
let il = input_layout(node.inputs[0]);
let dims = il.shape().dims();
if dims.is_empty() {
return Err(Error::Msg(
"Op::Fused(NORM_LAST_DIM[_BACKWARD], _) requires rank ≥ 1".to_string(),
)
.bt());
}
let last_dim = *dims.last().unwrap();
let outer_count: usize = dims[..dims.len() - 1].iter().product();
OpParams::NormLastDim {
outer_count,
last_dim,
eps,
}
}
Op::Fused(fid, _) if *fid == fuel_graph::registry::FusedOps::REDUCE_MAX_TO_BACKWARD => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::Fused(REDUCE_MAX_TO_BACKWARD) expects 2 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let up_layout = input_layout(node.inputs[1]);
let input_shape: Vec<usize> = x_layout.shape().dims().to_vec();
let output_shape: Vec<usize> = up_layout.shape().dims().to_vec();
OpParams::ReduceMaxToBackward {
input_shape,
output_shape,
}
}
Op::Concat { dim } => {
if node.inputs.is_empty() {
return Err(Error::Msg("Op::Concat requires at least 1 input".to_string()).bt());
}
let out_dims = node.shape.dims();
if *dim >= out_dims.len() {
return Err(Error::Msg(format!(
"Op::Concat: dim {dim} out of range for output rank {}",
out_dims.len(),
))
.bt());
}
let outer_count: usize = out_dims[..*dim].iter().product();
let inner_count: usize = out_dims[*dim + 1..].iter().product();
let mut input_dim_sizes: Vec<usize> = Vec::with_capacity(node.inputs.len());
for in_id in &node.inputs {
let il = input_layout(*in_id);
let il_dims = il.shape().dims();
if *dim >= il_dims.len() {
return Err(Error::Msg(format!(
"Op::Concat: input {in_id:?} has rank {} but concat dim is {dim}",
il_dims.len(),
))
.bt());
}
input_dim_sizes.push(il_dims[*dim]);
}
OpParams::Concat {
outer_count,
input_dim_sizes,
inner_count,
axis: *dim,
}
}
Op::AddScalar(c) => OpParams::Affine { mul: 1.0, add: *c },
Op::MulScalar(c) => OpParams::Affine { mul: *c, add: 0.0 },
Op::Clamp { min, max } => OpParams::Clamp {
min: *min,
max: *max,
},
Op::PowI(exp) => OpParams::PowI { exp: *exp },
Op::ClampInplace { min, max } => OpParams::Clamp {
min: *min,
max: *max,
},
Op::PowIInplace(exp) => OpParams::PowI { exp: *exp },
Op::Fused(_, fuel_graph::registry::FusedOpParams::PowIBackward { exp }) => {
OpParams::PowI { exp: *exp }
}
Op::Fused(_, fuel_graph::registry::FusedOpParams::InplaceAffine { mul, add }) => {
OpParams::Affine {
mul: *mul,
add: *add,
}
}
Op::Fused(
_,
fuel_graph::registry::FusedOpParams::Conv2D {
stride,
padding,
groups,
},
) => {
if node.inputs.len() != 2 && node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::Conv2D expects 2 or 3 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let w_layout = input_layout(node.inputs[1]);
let x_dims = x_layout.shape().dims();
let w_dims = w_layout.shape().dims();
if x_dims.len() != 4 || w_dims.len() != 4 {
return Err(Error::Msg(format!(
"Op::Conv2D requires rank-4 x and weight; got x={x_dims:?} w={w_dims:?}",
))
.bt());
}
let out_dims = node.shape.dims();
if out_dims.len() != 4 {
return Err(Error::Msg(format!(
"Op::Conv2D output must be rank 4, got {out_dims:?}",
))
.bt());
}
let x_shape = [x_dims[0], x_dims[1], x_dims[2], x_dims[3]];
let w_shape = [w_dims[0], w_dims[1], w_dims[2], w_dims[3]];
let out_shape = [out_dims[0], out_dims[1], out_dims[2], out_dims[3]];
OpParams::Conv2D {
x_shape,
w_shape,
out_shape,
stride: *stride,
padding: *padding,
dilation: (1, 1),
groups: *groups,
}
}
Op::Fused(
_,
fuel_graph::registry::FusedOpParams::PagedAttn {
softmax_scale,
block_size,
softcap,
},
) => {
if node.inputs.len() != 5 && node.inputs.len() != 6 {
return Err(Error::Msg(format!(
"Op::PagedAttn expects 5 or 6 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let q_layout = input_layout(node.inputs[0]);
let kc_layout = input_layout(node.inputs[1]);
let vc_layout = input_layout(node.inputs[2]);
let bt_layout = input_layout(node.inputs[3]);
let q_dims = q_layout.shape().dims();
let kc_dims = kc_layout.shape().dims();
let vc_dims = vc_layout.shape().dims();
let bt_dims = bt_layout.shape().dims();
if q_dims.len() != 4 || kc_dims.len() != 4 || vc_dims.len() != 4 {
return Err(Error::Msg(format!(
"Op::PagedAttn requires rank-4 q/k_cache/v_cache; \
got q={q_dims:?} k_cache={kc_dims:?} v_cache={vc_dims:?}",
))
.bt());
}
if kc_dims != vc_dims {
return Err(Error::Msg(format!(
"Op::PagedAttn: k_cache {kc_dims:?} and v_cache {vc_dims:?} must match",
))
.bt());
}
if bt_dims.len() != 2 {
return Err(Error::Msg(format!(
"Op::PagedAttn: block_table must be rank 2 [B, max_blocks_per_seq], got {bt_dims:?}",
))
.bt());
}
if kc_dims[1] != *block_size {
return Err(Error::Msg(format!(
"Op::PagedAttn: k_cache block_size dim ({}) must equal Op's block_size ({})",
kc_dims[1], block_size,
))
.bt());
}
OpParams::PagedAttn {
b: q_dims[0],
hq: q_dims[1],
hkv: kc_dims[2],
sq: q_dims[2],
d: q_dims[3],
block_size: *block_size,
max_blocks_per_seq: bt_dims[1],
num_blocks: kc_dims[0],
softmax_scale: *softmax_scale,
softcap: *softcap,
}
}
Op::Fused(
_,
fuel_graph::registry::FusedOpParams::FlashAttn {
softmax_scale,
causal,
window_size_left,
window_size_right,
softcap,
k_len,
},
) => {
if node.inputs.len() != 3 && node.inputs.len() != 4 {
return Err(Error::Msg(format!(
"Op::FlashAttn expects 3 or 4 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let q_layout = input_layout(node.inputs[0]);
let k_layout = input_layout(node.inputs[1]);
let v_layout = input_layout(node.inputs[2]);
let q_dims = q_layout.shape().dims();
let k_dims = k_layout.shape().dims();
let v_dims = v_layout.shape().dims();
if q_dims.len() != 4 || k_dims.len() != 4 || v_dims.len() != 4 {
return Err(Error::Msg(format!(
"Op::FlashAttn requires rank-4 q/k/v; got q={q_dims:?} k={k_dims:?} v={v_dims:?}",
))
.bt());
}
if k_dims != v_dims {
return Err(Error::Msg(format!(
"Op::FlashAttn: k {k_dims:?} and v {v_dims:?} must share shape",
))
.bt());
}
if q_dims[0] != k_dims[0] || q_dims[3] != k_dims[3] {
return Err(Error::Msg(format!(
"Op::FlashAttn: q {q_dims:?} and k {k_dims:?} must share B and D",
))
.bt());
}
let sk = k_dims[2];
let attended_k_len = match k_len {
None => sk,
Some(dyn_k) => {
let kl = dyn_k.resolve(sym_env).ok_or_else(|| {
Error::Msg(format!(
"Op::FlashAttn: dynamic k_len {dyn_k:?} is unbound in the SymEnv at realize",
))
.bt()
})?;
if kl > sk {
return Err(Error::Msg(format!(
"Op::FlashAttn: resolved k_len ({kl}) exceeds K capacity ({sk})",
))
.bt());
}
if kl < q_dims[2] {
return Err(Error::Msg(format!(
"Op::FlashAttn: resolved k_len ({kl}) < Sq ({}) — not a valid causal prefix",
q_dims[2],
))
.bt());
}
kl
}
};
OpParams::FlashAttn {
b: q_dims[0],
hq: q_dims[1],
hkv: k_dims[1],
sq: q_dims[2],
sk,
d: q_dims[3],
k_len: attended_k_len,
softmax_scale: *softmax_scale,
causal: *causal,
window_size_left: *window_size_left,
window_size_right: *window_size_right,
softcap: *softcap,
}
}
Op::ReduceSumTo(target_shape) => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::ReduceSumTo expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let input_shape: Vec<usize> = in_layout.shape().dims().to_vec();
let output_shape: Vec<usize> = target_shape.dims().to_vec();
OpParams::ReduceSumTo {
input_shape,
output_shape,
}
}
Op::ReduceMaxTo(target_shape) => {
if node.inputs.len() != 1 {
return Err(Error::Msg(format!(
"Op::ReduceMaxTo expects 1 input, got {}",
node.inputs.len(),
))
.bt());
}
let in_layout = input_layout(node.inputs[0]);
let input_shape: Vec<usize> = in_layout.shape().dims().to_vec();
let output_shape: Vec<usize> = target_shape.dims().to_vec();
OpParams::ReduceMaxTo {
input_shape,
output_shape,
}
}
Op::Fused(
_,
fuel_graph::registry::FusedOpParams::ConvTranspose2D {
stride,
padding,
output_padding,
dilation,
groups,
},
) => {
if node.inputs.len() != 2 && node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::ConvTranspose2D expects 2 or 3 inputs, got {}",
node.inputs.len(),
))
.bt());
}
let x_layout = input_layout(node.inputs[0]);
let w_layout = input_layout(node.inputs[1]);
let x_dims = x_layout.shape().dims();
let w_dims = w_layout.shape().dims();
if x_dims.len() != 4 || w_dims.len() != 4 {
return Err(Error::Msg(format!(
"Op::ConvTranspose2D requires rank-4 x and weight; got x={x_dims:?} w={w_dims:?}",
))
.bt());
}
let out_dims = node.shape.dims();
if out_dims.len() != 4 {
return Err(Error::Msg(format!(
"Op::ConvTranspose2D output must be rank 4, got {out_dims:?}",
))
.bt());
}
let x_shape = [x_dims[0], x_dims[1], x_dims[2], x_dims[3]];
let w_shape = [w_dims[0], w_dims[1], w_dims[2], w_dims[3]];
let out_shape = [out_dims[0], out_dims[1], out_dims[2], out_dims[3]];
OpParams::ConvTranspose2D {
x_shape,
w_shape,
out_shape,
stride: *stride,
padding: *padding,
output_padding: *output_padding,
dilation: *dilation,
groups: *groups,
}
}
Op::WriteSlice { ranges, dyn_offset } => {
if node.inputs.len() != 2 {
return Err(Error::Msg(format!(
"Op::WriteSlice expects 2 inputs (destination, source), got {}",
node.inputs.len(),
))
.bt());
}
let dest_dims = node.shape.dims().to_vec();
if ranges.len() != dest_dims.len() {
return Err(Error::Msg(format!(
"Op::WriteSlice: ranges.len() ({}) must equal destination rank ({})",
ranges.len(),
dest_dims.len(),
))
.bt());
}
let mut ranges = ranges.clone();
let mut deferred_dyn_offset: Option<(usize, fuel_ir::DynScalar)> = None;
if let Some((axis, off)) = dyn_offset {
if *axis >= ranges.len() {
return Err(Error::Msg(format!(
"Op::WriteSlice: dyn_offset axis ({axis}) out of bounds for rank {}",
ranges.len(),
))
.bt());
}
match off.resolve(sym_env) {
Some(base) => {
let (start, end) = ranges[*axis];
let width = end - start;
ranges[*axis] = (base, base + width);
}
None => {
let is_data_determined = matches!(off, fuel_ir::DynScalar::Sym(s)
if graph.binds_data_determined_sym(*s));
if is_data_determined {
deferred_dyn_offset = Some((*axis, *off));
} else {
return Err(Error::Msg(format!(
"Op::WriteSlice: dynamic offset {off:?} on axis {axis} is \
unbound in the SymEnv at realize (and no producer in the \
graph binds it, so it is not data-determined)",
))
.bt());
}
}
}
}
for (i, &(start, end)) in ranges.iter().enumerate() {
if deferred_dyn_offset.is_some_and(|(a, _)| a == i) {
continue;
}
if end < start || end > dest_dims[i] {
return Err(Error::Msg(format!(
"Op::WriteSlice: ranges[{i}] = ({start}, {end}) invalid \
for destination dim {i} = {}",
dest_dims[i],
))
.bt());
}
}
OpParams::WriteSlice {
dest_shape: dest_dims,
ranges,
deferred_dyn_offset,
}
}
Op::WriteSliceRotating {
axis,
modulus,
ranges,
} => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::WriteSliceRotating expects 3 inputs (destination, source, position), got {}",
node.inputs.len(),
))
.bt());
}
let dest_dims = node.shape.dims().to_vec();
if ranges.len() != dest_dims.len() {
return Err(Error::Msg(format!(
"Op::WriteSliceRotating: ranges.len() ({}) must equal destination rank ({})",
ranges.len(),
dest_dims.len(),
))
.bt());
}
if *axis >= dest_dims.len() {
return Err(Error::Msg(format!(
"Op::WriteSliceRotating: axis {axis} out of bounds for rank {}",
dest_dims.len(),
))
.bt());
}
if *modulus == 0 || *modulus > dest_dims[*axis] {
return Err(Error::Msg(format!(
"Op::WriteSliceRotating: modulus {modulus} invalid for dest_shape[{axis}] = {}",
dest_dims[*axis],
))
.bt());
}
OpParams::WriteSliceRotating {
dest_shape: dest_dims,
axis: *axis,
modulus: *modulus,
ranges: ranges.clone(),
}
}
Op::WriteSliceDoff { axis, ranges } => {
if node.inputs.len() != 3 {
return Err(Error::Msg(format!(
"Op::WriteSliceDoff expects 3 inputs (destination, source, offset), got {}",
node.inputs.len(),
))
.bt());
}
let dest_dims = node.shape.dims().to_vec();
if ranges.len() != dest_dims.len() {
return Err(Error::Msg(format!(
"Op::WriteSliceDoff: ranges.len() ({}) must equal destination rank ({})",
ranges.len(),
dest_dims.len(),
))
.bt());
}
if *axis >= dest_dims.len() {
return Err(Error::Msg(format!(
"Op::WriteSliceDoff: axis {axis} out of bounds for rank {}",
dest_dims.len(),
))
.bt());
}
OpParams::WriteSliceDoff {
dest_shape: dest_dims,
axis: *axis,
ranges: ranges.clone(),
}
}
_ => OpParams::None,
})
}
fn with_node_location(graph: &Arc<RwLock<Graph>>, node_id: NodeId, err: Error) -> Error {
match graph.read() {
Ok(g) => err.context(g.describe_node(node_id)),
Err(_) => err,
}
}
fn bind_data_determined_count(
arc: &Arc<RwLock<Storage>>,
capacity: usize,
count_sym: fuel_ir::SymId,
produced_syms: &mut SymEnv,
) -> Result<()> {
let guard = arc
.read()
.map_err(|_| poisoned("NonZeroIndices output storage"))?;
let count_byte_off = capacity * std::mem::size_of::<u32>();
let end = count_byte_off + std::mem::size_of::<u32>();
#[allow(unreachable_patterns)]
let count = match &guard.inner {
fuel_memory::BackendStorage::Cpu(cpu_bytes) => {
let bytes = cpu_bytes.bytes();
if end > bytes.len() {
return Err(Error::Msg(format!(
"NonZeroIndices: count slot byte range [{count_byte_off}..{end}) \
exceeds bundle byte length {}",
bytes.len(),
))
.bt());
}
u32::from_le_bytes(
bytes[count_byte_off..end]
.try_into()
.expect("4 bytes → u32"),
)
}
other => {
return Err(Error::Msg(format!(
"NonZeroIndices: data-determined count read is CPU-only today; got \
BackendStorage::{:?} (device backends need a D2H after the completion \
handle is waited — a documented follow-up)",
std::mem::discriminant(other),
))
.bt());
}
};
produced_syms.bind(count_sym, count as usize)?;
Ok(())
}
fn resolve_deferred_write_slice(
compiled: &CompiledNode,
produced_syms: &SymEnv,
node_id: NodeId,
) -> Result<Option<CompiledNode>> {
let OpParams::WriteSlice {
dest_shape,
ranges,
deferred_dyn_offset,
} = &compiled.op_params
else {
return Ok(None);
};
let Some((axis, off)) = deferred_dyn_offset else {
return Ok(None);
};
let base = off.resolve(produced_syms).ok_or_else(|| {
Error::Msg(format!(
"Op::WriteSlice {node_id:?}: data-determined offset {off:?} on axis {axis} is \
unbound in produced_syms at execute (its producer did not run / bind it)",
))
.bt()
})?;
let mut ranges = ranges.clone();
let (start, end) = ranges[*axis];
let width = end - start;
let resolved = (base, base + width);
if resolved.1 > dest_shape[*axis] {
return Err(Error::Msg(format!(
"Op::WriteSlice {node_id:?}: resolved data-determined range {resolved:?} on axis \
{axis} exceeds destination dim = {}",
dest_shape[*axis],
))
.bt());
}
ranges[*axis] = resolved;
let mut patched = compiled.clone();
patched.op_params = OpParams::WriteSlice {
dest_shape: dest_shape.clone(),
ranges,
deferred_dyn_offset: None,
};
Ok(Some(patched))
}
fn resolve_deferred_matmul(
compiled: &CompiledNode,
produced_syms: &SymEnv,
node_id: NodeId,
) -> Result<Option<CompiledNode>> {
let OpParams::Matmul {
lhs_batch_dims,
rhs_batch_dims,
m,
n,
k,
m_compute,
} = &compiled.op_params
else {
return Ok(None);
};
let MatmulM::Deferred(dyn_m) = m_compute else {
return Ok(None);
};
let rows = dyn_m.resolve(produced_syms).ok_or_else(|| {
Error::Msg(format!(
"Op::MatMul {node_id:?}: data-determined row count {dyn_m:?} is unbound in \
produced_syms at execute (its producer did not run / bind it)",
))
.bt()
})?;
if rows > *m {
return Err(Error::Msg(format!(
"Op::MatMul {node_id:?}: resolved row count {rows} exceeds capacity {m}",
))
.bt());
}
let mut patched = compiled.clone();
patched.op_params = OpParams::Matmul {
lhs_batch_dims: lhs_batch_dims.clone(),
rhs_batch_dims: rhs_batch_dims.clone(),
m: *m,
n: *n,
k: *k,
m_compute: MatmulM::Rows(rows),
};
Ok(Some(patched))
}
fn execute_work_item(
item: &WorkItem,
cache: &mut StorageCache,
layout_cache: &mut HashMap<NodeId, Layout>,
produced_syms: &mut SymEnv,
will_be_waited: bool,
mut persistent: Option<&mut PersistentOutputs>,
) -> Result<CompletionHandle> {
match &item.kind {
WorkItemKind::ConstAdopt => {
if !cache.contains_key(&item.node_id) {
return Err(Error::Msg(format!(
"PipelinedExecutor: Const node {:?} not in input cache",
item.node_id
))
.bt());
}
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::ViewOf { input } => {
let input_arc = cache.get(input).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: view-op input {:?} of {:?} not realized",
input, item.node_id,
))
.bt()
})?;
cache.insert(item.node_id, input_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::SlotView { producer } => {
let producer_arc = cache.get(producer).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: Op::View producer {:?} of {:?} not realized",
producer, item.node_id,
))
.bt()
})?;
cache.insert(item.node_id, producer_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::SlotOwn { producer, slot } => {
let producer_arc = cache.get(producer).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: Op::ViewOwned producer {:?} of {:?} not realized",
producer, item.node_id,
))
.bt()
})?;
let new_storage = {
let guard = producer_arc
.read()
.map_err(|_| poisoned("ViewOwned producer storage"))?;
let bundle = guard.bundle().ok_or_else(|| {
Error::Msg(format!(
"Op::ViewOwned: producer {:?} has no bundle side-table \
— the multi-output authoring contract requires \
Storage::with_bundle / new_bundled at allocation \
time. Node {:?} slot {} cannot resolve its source.",
producer, item.node_id, slot,
))
.bt()
})?;
let sv = bundle.get(*slot as usize).cloned().ok_or_else(|| {
Error::Msg(format!(
"Op::ViewOwned: slot {} out of range for producer {:?} \
(bundle has {} slots)",
slot,
producer,
bundle.len(),
))
.bt()
})?;
use fuel_memory::BackendStorage;
match &guard.inner {
BackendStorage::Cpu(cpu_bytes) => {
let src_bytes = cpu_bytes.bytes();
let end = sv.byte_offset + sv.len_bytes();
if end > src_bytes.len() {
return Err(Error::Msg(format!(
"Op::ViewOwned: slot {} byte range [{}..{}) \
exceeds producer's CPU byte length {}",
slot,
sv.byte_offset,
end,
src_bytes.len(),
))
.bt());
}
let owned = fuel_cpu_backend::CpuStorageBytes::from_bytes(
&src_bytes[sv.byte_offset..end],
);
fuel_memory::Storage::new(BackendStorage::Cpu(owned), sv.dtype)
}
#[cfg(feature = "cuda")]
BackendStorage::Cuda(cuda_bytes) => {
let dst = cuda_bytes.slot_copy_to_new(sv.byte_offset, sv.len_bytes())?;
fuel_memory::Storage::new(BackendStorage::Cuda(dst), sv.dtype)
}
#[cfg(feature = "vulkan")]
BackendStorage::Vulkan(vk_bytes) => {
let backend = vk_bytes.backend().ok_or_else(|| {
Error::Msg(format!(
"Op::ViewOwned: producer {:?} is on Vulkan but \
its storage has no VulkanBackend handle. \
Storages flowing through the pipelined executor \
must be constructed via alloc_bytes_handle / \
upload_bytes_handle so the bundle's slot extraction \
path can reach the backend.",
producer,
))
.bt()
})?;
let dst = backend.slot_copy_to_new_handle(
vk_bytes,
sv.byte_offset,
sv.len_bytes(),
)?;
fuel_memory::Storage::new(BackendStorage::Vulkan(dst), sv.dtype)
}
#[allow(unreachable_patterns)]
_ => {
return Err(Error::Msg(format!(
"Op::ViewOwned: producer {:?} is on a backend \
without a slot-copy-with-offset hook \
(CPU + CUDA + Vulkan wired). Add a per-backend \
`slot_copy_to_new`-equivalent and extend the \
match arm.",
producer,
))
.bt());
}
}
};
cache.insert(item.node_id, Arc::new(RwLock::new(new_storage)));
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::WriteSlice { dest, source } => {
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSlice work item {:?} has no compiled node",
item.node_id,
))
.bt()
})?;
let patched_ws = resolve_deferred_write_slice(compiled, produced_syms, item.node_id)?;
let compiled = patched_ws.as_ref().unwrap_or(compiled);
let dest_arc = cache.get(dest).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSlice destination {:?} of {:?} not realized",
dest, item.node_id,
))
.bt()
})?;
let source_arc = cache.get(source).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSlice source {:?} of {:?} not realized",
source, item.node_id,
))
.bt()
})?;
let dest_layout = layout_cache.get(dest).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSlice destination {:?} of {:?} has no cached layout",
dest, item.node_id,
))
.bt()
})?;
let source_layout = layout_cache.get(source).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSlice source {:?} of {:?} has no cached layout",
source, item.node_id,
))
.bt()
})?;
if !dest_layout.is_contiguous() || dest_layout.start_offset() != 0 {
return Err(Error::Msg(format!(
"Op::WriteSlice (Phase E.3.2 v1): destination {:?} must be \
contiguous + zero-offset; got Layout {:?}",
dest, dest_layout,
))
.bt());
}
let source_arc_contig = {
let s_dtype = source_arc
.read()
.map_err(|_| poisoned("WriteSlice source storage"))?
.dtype;
let s_len_bytes = source_arc
.read()
.map_err(|_| poisoned("WriteSlice source storage"))?
.inner
.len_bytes();
let layout_bytes = source_layout.shape().elem_count() * s_dtype.size_in_bytes();
let bytes_match_shape = s_len_bytes == layout_bytes;
let already_contig = source_layout.is_contiguous()
&& source_layout.start_offset() == 0
&& bytes_match_shape;
if already_contig {
source_arc
} else {
auto_contiguize(&source_arc, &source_layout)?
}
};
let source_layout_kernel = fuel_ir::Layout::contiguous(source_layout.shape().clone());
let input_arcs = vec![source_arc_contig];
let mut output_arcs = vec![dest_arc.clone()];
let kernel_layouts = vec![source_layout_kernel, dest_layout.clone()];
let handle = execute_compiled_with_wait_hint(
compiled,
&input_arcs,
&mut output_arcs,
&kernel_layouts,
will_be_waited,
)?;
cache.insert(item.node_id, dest_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(handle)
}
WorkItemKind::WriteSliceRotating {
dest,
source,
position,
} => {
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating work item {:?} has no compiled node",
item.node_id,
))
.bt()
})?;
let dest_arc = cache.get(dest).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating destination {:?} of {:?} not realized",
dest, item.node_id,
))
.bt()
})?;
let source_arc = cache.get(source).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating source {:?} of {:?} not realized",
source, item.node_id,
))
.bt()
})?;
let position_arc = cache.get(position).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating position {:?} of {:?} not realized",
position, item.node_id,
))
.bt()
})?;
let dest_layout = layout_cache.get(dest).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating destination {:?} of {:?} has no cached layout",
dest, item.node_id,
))
.bt()
})?;
let source_layout = layout_cache.get(source).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceRotating source {:?} of {:?} has no cached layout",
source, item.node_id,
))
.bt()
})?;
if !dest_layout.is_contiguous() || dest_layout.start_offset() != 0 {
return Err(Error::Msg(format!(
"Op::WriteSliceRotating (v1): destination {:?} must be \
contiguous + zero-offset; got Layout {:?}",
dest, dest_layout,
))
.bt());
}
let source_arc_contig = {
let s_dtype = source_arc
.read()
.map_err(|_| poisoned("WriteSliceRotating source storage"))?
.dtype;
let s_len_bytes = source_arc
.read()
.map_err(|_| poisoned("WriteSliceRotating source storage"))?
.inner
.len_bytes();
let layout_bytes = source_layout.shape().elem_count() * s_dtype.size_in_bytes();
let bytes_match_shape = s_len_bytes == layout_bytes;
let already_contig = source_layout.is_contiguous()
&& source_layout.start_offset() == 0
&& bytes_match_shape;
if already_contig {
source_arc
} else {
auto_contiguize(&source_arc, &source_layout)?
}
};
let source_layout_kernel = fuel_ir::Layout::contiguous(source_layout.shape().clone());
let position_layout = fuel_ir::Layout::contiguous(fuel_ir::Shape::from_dims(&[]));
let input_arcs = vec![source_arc_contig, position_arc];
let mut output_arcs = vec![dest_arc.clone()];
let kernel_layouts = vec![source_layout_kernel, position_layout, dest_layout.clone()];
let handle = execute_compiled_with_wait_hint(
compiled,
&input_arcs,
&mut output_arcs,
&kernel_layouts,
will_be_waited,
)?;
cache.insert(item.node_id, dest_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(handle)
}
WorkItemKind::WriteSliceDoff {
dest,
source,
offset,
} => {
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff work item {:?} has no compiled node",
item.node_id,
))
.bt()
})?;
let dest_arc = cache.get(dest).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff destination {:?} of {:?} not realized",
dest, item.node_id,
))
.bt()
})?;
let source_arc = cache.get(source).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff source {:?} of {:?} not realized",
source, item.node_id,
))
.bt()
})?;
let offset_arc = cache.get(offset).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff offset {:?} of {:?} not realized",
offset, item.node_id,
))
.bt()
})?;
let dest_layout = layout_cache.get(dest).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff destination {:?} of {:?} has no cached layout",
dest, item.node_id,
))
.bt()
})?;
let source_layout = layout_cache.get(source).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: WriteSliceDoff source {:?} of {:?} has no cached layout",
source, item.node_id,
))
.bt()
})?;
if !dest_layout.is_contiguous() || dest_layout.start_offset() != 0 {
return Err(Error::Msg(format!(
"Op::WriteSliceDoff (v1): destination {:?} must be \
contiguous + zero-offset; got Layout {:?}",
dest, dest_layout,
))
.bt());
}
let source_arc_contig = {
let s_dtype = source_arc
.read()
.map_err(|_| poisoned("WriteSliceDoff source storage"))?
.dtype;
let s_len_bytes = source_arc
.read()
.map_err(|_| poisoned("WriteSliceDoff source storage"))?
.inner
.len_bytes();
let layout_bytes = source_layout.shape().elem_count() * s_dtype.size_in_bytes();
let bytes_match_shape = s_len_bytes == layout_bytes;
let already_contig = source_layout.is_contiguous()
&& source_layout.start_offset() == 0
&& bytes_match_shape;
if already_contig {
source_arc
} else {
auto_contiguize(&source_arc, &source_layout)?
}
};
let source_layout_kernel = fuel_ir::Layout::contiguous(source_layout.shape().clone());
let offset_layout = fuel_ir::Layout::contiguous(fuel_ir::Shape::from_dims(&[]));
let input_arcs = vec![source_arc_contig, offset_arc];
let mut output_arcs = vec![dest_arc.clone()];
let kernel_layouts = vec![source_layout_kernel, offset_layout, dest_layout.clone()];
let handle = execute_compiled_with_wait_hint(
compiled,
&input_arcs,
&mut output_arcs,
&kernel_layouts,
will_be_waited,
)?;
cache.insert(item.node_id, dest_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(handle)
}
WorkItemKind::Alloc { target_location } => {
#[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
let n_bytes = item.elem_count * item.dtype.size_in_bytes();
let alloced: Storage = match target_location {
DeviceLocation::Cpu => fuel_memory::alloc_cpu_zeroed(item.dtype, item.elem_count)?,
#[cfg(feature = "cuda")]
DeviceLocation::Cuda { gpu_id } => {
let cuda_dev = find_cuda_device_in_cache(cache, *gpu_id).ok_or_else(|| {
Error::Msg(format!(
"Op::Alloc on Cuda {{ gpu_id: {} }}: no CUDA \
storage in input cache to derive the device \
handle from. The caller must seed the cache \
(e.g. via `fuel-core::pipelined_bridge::\
device_seed_storage`) before realizing.",
gpu_id,
))
.bt()
})?;
let cuda_bytes =
fuel_cuda_backend::CudaStorageBytes::alloc_uninit(&cuda_dev, n_bytes)?;
Storage::new(fuel_memory::BackendStorage::Cuda(cuda_bytes), item.dtype)
}
#[cfg(not(feature = "cuda"))]
DeviceLocation::Cuda { .. } => {
return Err(Error::Msg(
"Op::Alloc on Cuda but fuel-storage wasn't built \
with --features cuda"
.to_string(),
)
.bt());
}
#[cfg(feature = "vulkan")]
DeviceLocation::Vulkan { gpu_id } => {
let backend =
find_vulkan_backend_in_cache(cache, *gpu_id).ok_or_else(|| {
Error::Msg(format!(
"Op::Alloc on Vulkan {{ gpu_id: {} }}: no Vulkan \
storage in input cache to derive the backend \
handle from. The caller must seed the cache \
(e.g. via `fuel-core::pipelined_bridge::\
device_seed_storage`) before realizing.",
gpu_id,
))
.bt()
})?;
let vk_bytes = backend.alloc_bytes_handle(n_bytes)?;
Storage::new(fuel_memory::BackendStorage::Vulkan(vk_bytes), item.dtype)
}
#[cfg(not(feature = "vulkan"))]
DeviceLocation::Vulkan { .. } => {
return Err(Error::Msg(
"Op::Alloc on Vulkan but fuel-storage wasn't built \
with --features vulkan"
.to_string(),
)
.bt());
}
other => {
return Err(Error::Msg(format!(
"Op::Alloc on {:?}: target not yet wired (CPU + \
CUDA + Vulkan land in Phase 3a of bridge-retirement; \
Metal extends when its byte-storage substrate is \
ready).",
other,
))
.bt());
}
};
cache.insert(item.node_id, Arc::new(RwLock::new(alloced)));
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::Iota { len } => {
if item.target_backend != BackendId::Cpu {
return Err(Error::Msg(format!(
"WorkItemKind::Iota: only CPU is wired today, got {:?}",
item.target_backend,
))
.bt());
}
let v: Vec<f32> = (0..*len).map(|i| i as f32).collect();
let iota = fuel_memory::from_slice_cpu(&v);
cache.insert(item.node_id, Arc::new(RwLock::new(iota)));
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::ZeroFill => {
if item.inputs.len() != 1 {
return Err(Error::Msg(format!(
"PipelinedExecutor: ZeroFill work item {:?} expects \
1 input, got {}",
item.node_id,
item.inputs.len(),
))
.bt());
}
let src_id = item.inputs[0];
let src_arc = cache.get(&src_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: ZeroFill input {:?} of {:?} \
not realized",
src_id, item.node_id,
))
.bt()
})?;
{
let mut guard = src_arc
.write()
.map_err(|_| poisoned("ZeroFill destination storage"))?;
match &mut guard.inner {
fuel_memory::BackendStorage::Cpu(c) => {
let bytes = c.bytes_mut();
bytes.fill(0);
}
#[cfg(feature = "cuda")]
fuel_memory::BackendStorage::Cuda(c) => {
c.zero_async()?;
}
#[cfg(feature = "vulkan")]
fuel_memory::BackendStorage::Vulkan(v) => {
let backend = v
.backend()
.ok_or_else(|| {
Error::Msg(
"Op::ZeroFill on Vulkan: input has no \
attached backend handle. Storage must \
come from VulkanBackend::alloc_bytes_handle \
/ upload_bytes_handle."
.to_string(),
)
.bt()
})?
.clone();
backend.fill_bytes_zero(v)?;
}
#[allow(unreachable_patterns)]
other => {
return Err(Error::Msg(format!(
"Op::ZeroFill: backend not wired ({other:?}); \
CPU + CUDA + Vulkan covered, Metal extends \
when its byte-storage substrate is ready.",
))
.bt());
}
}
}
cache.insert(item.node_id, src_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::Copy { target_location } | WorkItemKind::Move { target_location } => {
let op_label = if matches!(item.kind, WorkItemKind::Move { .. }) {
"Move"
} else {
"Copy"
};
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: {op_label} work item {:?} has no compiled node",
item.node_id,
))
.bt()
})?;
if item.inputs.len() != 1 {
return Err(Error::Msg(format!(
"PipelinedExecutor: {op_label} work item {:?} expects 1 input, got {}",
item.node_id,
item.inputs.len(),
))
.bt());
}
let src_id = item.inputs[0];
let src_arc = cache.get(&src_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: {op_label} input {:?} of {:?} not realized",
src_id, item.node_id,
))
.bt()
})?;
{
let guard = src_arc
.read()
.map_err(|_| poisoned("Copy source storage"))?;
if guard.is_bundled() {
let layout_bytes_for_guard = layout_cache
.get(&src_id)
.map(|l| l.shape().elem_count() * guard.dtype.size_in_bytes())
.unwrap_or(usize::MAX);
let bundle_bytes = guard.inner.len_bytes();
let is_whole_bundle = layout_bytes_for_guard >= bundle_bytes;
if is_whole_bundle {
return Err(Error::Msg(format!(
"PipelinedExecutor: Op::{op_label} on a bundled producer \
({:?}) is not yet supported (Option C followup). \
Use Op::View → Op::Copy → Op::ViewOwned for \
per-slot cross-device moves, or wait for the \
per-backend whole-bundle-Copy hook.",
src_id,
))
.bt());
}
}
}
let src_layout = layout_cache.get(&src_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: {op_label} input {:?} of {:?} has no cached layout",
src_id, item.node_id,
))
.bt()
})?;
let src_dtype = src_arc
.read()
.map_err(|_| poisoned("Copy source storage"))?
.dtype;
let src_len_bytes = src_arc
.read()
.map_err(|_| poisoned("Copy source storage"))?
.inner
.len_bytes();
let layout_bytes = src_layout.shape().elem_count() * src_dtype.size_in_bytes();
let bytes_match_shape = src_len_bytes == layout_bytes;
let already_contig =
src_layout.is_contiguous() && src_layout.start_offset() == 0 && bytes_match_shape;
let src_input = if already_contig {
src_arc
} else {
auto_contiguize(&src_arc, &src_layout)?
};
let kernel_input_layout = fuel_ir::Layout::contiguous(src_layout.shape().clone());
let reuse_arc = match &persistent {
Some(p) if p.mode == PersistentMode::Reuse => match target_location {
DeviceLocation::Cuda { .. } => p.map.get(&item.node_id).cloned(),
_ => None,
},
_ => None,
};
#[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
let n_bytes = item.elem_count * item.dtype.size_in_bytes();
let input_arcs = vec![src_input];
let mut output_arcs = if let Some(arc) = reuse_arc {
vec![arc]
} else {
let output = match target_location {
DeviceLocation::Cpu => {
fuel_memory::alloc_cpu_zeroed(item.dtype, item.elem_count)?
}
#[cfg(feature = "cuda")]
DeviceLocation::Cuda { gpu_id } => {
let cuda_dev =
find_cuda_device_in_cache(cache, *gpu_id).ok_or_else(|| {
Error::Msg(format!(
"Op::{op_label} on Cuda {{ gpu_id: {} }}: no CUDA \
storage in input cache to derive the device \
handle from. The caller must seed the cache \
(e.g. via `fuel-core::pipelined_bridge::\
device_seed_storage`) before realizing an \
H2D Op::{op_label}.",
gpu_id,
))
.bt()
})?;
let cuda_bytes =
fuel_cuda_backend::CudaStorageBytes::alloc_uninit(&cuda_dev, n_bytes)?;
Storage::new(fuel_memory::BackendStorage::Cuda(cuda_bytes), item.dtype)
}
#[cfg(not(feature = "cuda"))]
DeviceLocation::Cuda { .. } => {
return Err(Error::Msg(format!(
"Op::{op_label} target Cuda but fuel-storage wasn't built \
with --features cuda",
))
.bt());
}
#[cfg(feature = "vulkan")]
DeviceLocation::Vulkan { gpu_id } => {
let backend =
find_vulkan_backend_in_cache(cache, *gpu_id).ok_or_else(|| {
Error::Msg(format!(
"Op::{op_label} on Vulkan {{ gpu_id: {} }}: no Vulkan \
storage in input cache to derive the backend \
handle from. The caller must seed the cache \
(e.g. via `fuel-core::pipelined_bridge::\
device_seed_storage`) before realizing an \
H2D Op::{op_label}.",
gpu_id,
))
.bt()
})?;
let vk_bytes = backend.alloc_bytes_handle(n_bytes)?;
Storage::new(fuel_memory::BackendStorage::Vulkan(vk_bytes), item.dtype)
}
#[cfg(not(feature = "vulkan"))]
DeviceLocation::Vulkan { .. } => {
return Err(Error::Msg(format!(
"Op::{op_label} target Vulkan but fuel-storage wasn't built \
with --features vulkan",
))
.bt());
}
other => {
return Err(Error::Msg(format!(
"PipelinedExecutor: Op::{op_label} target_location {:?} not yet \
wired (CPU + CUDA + Vulkan covered; Metal extends when \
its byte-storage substrate is ready).",
other
))
.bt());
}
};
vec![Arc::new(RwLock::new(output))]
};
let kernel_layouts = vec![kernel_input_layout, item.output_layout.clone()];
let handle = execute_compiled_with_wait_hint(
compiled,
&input_arcs,
&mut output_arcs,
&kernel_layouts,
will_be_waited,
)?;
let arc = output_arcs.into_iter().next().expect("one output");
if let Some(p) = persistent.as_deref_mut()
&& p.mode == PersistentMode::Record
{
p.map.insert(item.node_id, Arc::clone(&arc));
}
cache.insert(item.node_id, arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(handle)
}
WorkItemKind::ReleaseMarker => {
let marker = fuel_memory::alloc_cpu_zeroed(item.dtype, 0)?;
cache.insert(item.node_id, Arc::new(RwLock::new(marker)));
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::ContiguizeOf { input } => {
let input_arc = cache.get(input).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: reshape input {:?} of {:?} not realized",
input, item.node_id,
))
.bt()
})?;
let input_layout = layout_cache.get(input).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: reshape input {:?} of {:?} has no cached layout",
input, item.node_id,
))
.bt()
})?;
let reuse_arc = match &persistent {
Some(p) if p.mode == PersistentMode::Reuse => p.map.get(&item.node_id).cloned(),
_ => None,
};
let out_arc = if let Some(dest_arc) = reuse_arc {
contiguize_into(&dest_arc, &input_arc, &input_layout)?;
dest_arc
} else if input_layout.is_contiguous() && input_layout.start_offset() == 0 {
input_arc
} else {
let arc = auto_contiguize(&input_arc, &input_layout)?;
if let Some(p) = persistent.as_deref_mut()
&& p.mode == PersistentMode::Record
{
p.map.insert(item.node_id, Arc::clone(&arc));
}
arc
};
cache.insert(item.node_id, out_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(CompletionHandle::Ready)
}
WorkItemKind::Kernel => {
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: Kernel work item {:?} has no compiled node",
item.node_id,
))
.bt()
})?;
let patched_mm = resolve_deferred_matmul(compiled, produced_syms, item.node_id)?;
let compiled = patched_mm.as_ref().unwrap_or(compiled);
let mut input_arcs: Vec<Arc<RwLock<Storage>>> = Vec::with_capacity(item.inputs.len());
let mut kernel_layouts: Vec<fuel_ir::Layout> =
Vec::with_capacity(item.inputs.len() + 1);
let kernel_handles_strided = compiled.caps.strided_input;
for in_id in &item.inputs {
let in_arc = cache.get(in_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: input {:?} of {:?} not realized",
in_id, item.node_id,
))
.bt()
})?;
let in_layout = layout_cache.get(in_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: input {:?} of {:?} has no cached layout",
in_id, item.node_id,
))
.bt()
})?;
let in_dtype = in_arc.read().map_err(|_| poisoned("input storage"))?.dtype;
let in_len_bytes = in_arc
.read()
.map_err(|_| poisoned("input storage"))?
.inner
.len_bytes();
let layout_bytes = in_layout.shape().elem_count() * in_dtype.size_in_bytes();
let bytes_match_shape = in_len_bytes == layout_bytes;
let already_contig =
in_layout.is_contiguous() && in_layout.start_offset() == 0 && bytes_match_shape;
let strided_ok =
kernel_handles_strided && in_layout.start_offset() == 0 && bytes_match_shape;
if already_contig || strided_ok {
input_arcs.push(in_arc);
kernel_layouts.push(in_layout);
} else {
let contig_arc = auto_contiguize(&in_arc, &in_layout)?;
input_arcs.push(contig_arc);
kernel_layouts.push(fuel_ir::Layout::contiguous(in_layout.shape().clone()));
}
}
kernel_layouts.push(item.output_layout.clone());
let (alloc_elem_count, bundle_to_attach) = match &item.output_bundle {
Some(bundle) => {
let dtype_bytes = item.dtype.size_in_bytes().max(1);
let total_bytes = bundle.iter().fold(0usize, |acc, s| {
let end = s.byte_offset.saturating_add(s.len_bytes());
acc.max(end)
});
let elems = total_bytes.div_ceil(dtype_bytes);
(elems, Some(Arc::clone(bundle)))
}
None => (item.elem_count, None),
};
let reuse_arc = match &persistent {
Some(p) if p.mode == PersistentMode::Reuse => {
Some(p.map.get(&item.node_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: capture (reuse) mode has no persistent \
output buffer for kernel node {:?}; the warm (record) run \
must realize every kernel node before capture",
item.node_id,
))
.bt()
})?)
}
_ => None,
};
let mut output_arcs = if let Some(arc) = reuse_arc {
vec![arc]
} else {
let output = match item.target_backend {
BackendId::Cpu => fuel_memory::alloc_cpu_zeroed(item.dtype, alloc_elem_count)?,
#[cfg(feature = "cuda")]
BackendId::Cuda => {
let first_in = input_arcs.first().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: kernel {:?} on Cuda has no inputs; \
cannot derive device for output allocation",
item.node_id,
))
.bt()
})?;
let guard = first_in.read().map_err(|_| poisoned("input storage"))?;
let cuda_in = match &guard.inner {
fuel_memory::BackendStorage::Cuda(c) => c,
other => {
return Err(Error::Msg(format!(
"PipelinedExecutor: kernel {:?} target_backend=Cuda but \
input has BackendStorage::{:?}; mixed-backend kernels \
require an explicit Op::Copy first",
item.node_id,
std::mem::discriminant(other),
))
.bt());
}
};
let n_bytes = alloc_elem_count * item.dtype.size_in_bytes();
let cuda_bytes =
fuel_cuda_backend::CudaStorageBytes::alloc(cuda_in.device(), n_bytes)?;
fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cuda(cuda_bytes),
item.dtype,
)
}
#[cfg(feature = "vulkan")]
BackendId::Vulkan => {
let first_in = input_arcs.first().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: kernel {:?} on Vulkan has no inputs; \
cannot derive backend handle for output allocation",
item.node_id,
))
.bt()
})?;
let guard = first_in.read().map_err(|_| poisoned("input storage"))?;
let vk_in = match &guard.inner {
fuel_memory::BackendStorage::Vulkan(v) => v,
other => {
return Err(Error::Msg(format!(
"PipelinedExecutor: kernel {:?} target_backend=Vulkan but \
input has BackendStorage::{:?}; mixed-backend kernels \
require an explicit Op::Copy first",
item.node_id,
std::mem::discriminant(other),
))
.bt());
}
};
let backend = vk_in.backend().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: Vulkan kernel {:?} input has no backend \
handle. Storages flowing through the pipelined executor's \
Vulkan path must be constructed via \
VulkanBackend::alloc_bytes_handle / upload_bytes_handle \
(not the legacy alloc_bytes / upload_bytes which leave \
VulkanStorageBytes::backend = None).",
item.node_id,
))
.bt()
})?;
let n_bytes = alloc_elem_count * item.dtype.size_in_bytes();
let vk_bytes = backend.alloc_bytes_handle(n_bytes)?;
fuel_memory::Storage::new(
fuel_memory::BackendStorage::Vulkan(vk_bytes),
item.dtype,
)
}
other => {
return Err(Error::Msg(format!(
"PipelinedExecutor: target_backend {:?} output allocation \
not yet implemented (CPU + CUDA + Vulkan wired; Metal extends later)",
other
))
.bt());
}
};
let output = match bundle_to_attach {
Some(bundle) => output.with_bundle(bundle)?,
None => output,
};
vec![Arc::new(RwLock::new(output))]
};
let handle = execute_compiled_with_wait_hint(
compiled,
&input_arcs,
&mut output_arcs,
&kernel_layouts,
will_be_waited,
)?;
let arc = output_arcs.into_iter().next().expect("one output");
if let Some(p) = persistent
&& p.mode == PersistentMode::Record
{
p.map.insert(item.node_id, Arc::clone(&arc));
}
if let OpParams::NonZeroIndices {
capacity,
count_sym,
} = &compiled.op_params
{
bind_data_determined_count(&arc, *capacity, *count_sym, produced_syms)?;
}
cache.insert(item.node_id, arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(handle)
}
WorkItemKind::InplaceKernel { target_idx } => {
let compiled = item.compiled.as_ref().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: InplaceKernel work item {:?} has no compiled node",
item.node_id,
))
.bt()
})?;
let target_id = *item.inputs.get(*target_idx).ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: InplaceKernel work item {:?} target_idx={} \
out of bounds (inputs.len()={})",
item.node_id,
target_idx,
item.inputs.len(),
))
.bt()
})?;
let target_arc = cache.get(&target_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: InplaceKernel target {:?} of {:?} not realized",
target_id, item.node_id,
))
.bt()
})?;
let target_layout = layout_cache.get(&target_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: InplaceKernel target {:?} of {:?} has no cached layout",
target_id, item.node_id,
))
.bt()
})?;
if !target_layout.is_contiguous() || target_layout.start_offset() != 0 {
return Err(Error::Msg(format!(
"InplaceKernel (Phase 3 v1): target {:?} must be contiguous + \
zero-offset; got Layout {:?}",
target_id, target_layout,
))
.bt());
}
let mut input_arcs: Vec<Arc<RwLock<Storage>>> = Vec::new();
let mut kernel_layouts: Vec<Layout> = Vec::new();
for (i, &input_id) in item.inputs.iter().enumerate() {
if i == *target_idx {
continue;
}
let arc = cache.get(&input_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: InplaceKernel non-target input {:?} of {:?} \
not realized",
input_id, item.node_id,
))
.bt()
})?;
let lay = layout_cache.get(&input_id).cloned().ok_or_else(|| {
Error::Msg(format!(
"PipelinedExecutor: InplaceKernel non-target input {:?} of {:?} \
has no cached layout",
input_id, item.node_id,
))
.bt()
})?;
input_arcs.push(arc);
kernel_layouts.push(lay);
}
let mut output_arcs = vec![target_arc.clone()];
kernel_layouts.push(target_layout.clone());
let handle = execute_compiled_with_wait_hint(
compiled,
&input_arcs,
&mut output_arcs,
&kernel_layouts,
will_be_waited,
)?;
cache.insert(item.node_id, target_arc);
layout_cache.insert(item.node_id, item.output_layout.clone());
Ok(handle)
}
}
}
fn poisoned(what: &'static str) -> Error {
Error::Msg(format!("PipelinedExecutor: {} poisoned", what)).bt()
}
#[inline]
fn store_handle(
handles: &mut HashMap<NodeId, CompletionHandle>,
node_id: NodeId,
handle: CompletionHandle,
) {
if matches!(handle, CompletionHandle::Ready) {
if let Some(prev) = handles.remove(&node_id) {
let _ = prev.wait();
}
return;
}
if let Some(prev) = handles.insert(node_id, handle) {
let _ = prev.wait();
}
}
fn wait_producer_handle(
handles: &mut HashMap<NodeId, CompletionHandle>,
producer: NodeId,
) -> Result<()> {
if let Some(handle) = handles.remove(&producer) {
handle.wait()?;
}
Ok(())
}
fn drain_handles(handles: &mut HashMap<NodeId, CompletionHandle>) -> Result<()> {
for (_node, handle) in handles.drain() {
handle.wait()?;
}
debug_assert!(
handles.is_empty(),
"A4b-1: handle map must be empty after drain"
);
Ok(())
}
#[cfg(feature = "cuda")]
fn sync_active_cuda_devices(cache: &StorageCache) -> Result<()> {
let mut devices: HashMap<fuel_ir::DeviceLocation, fuel_cuda_backend::CudaDevice> =
HashMap::new();
for arc in cache.values() {
let guard = arc
.read()
.map_err(|_| poisoned("cache storage (cuda realize-end sync scan)"))?;
if let fuel_memory::BackendStorage::Cuda(c) = &guard.inner {
devices
.entry(c.device().location())
.or_insert_with(|| c.device().clone());
}
}
for dev in devices.values() {
dev.synchronize()?;
}
Ok(())
}
#[cfg(not(feature = "cuda"))]
fn sync_active_cuda_devices(_cache: &StorageCache) -> Result<()> {
Ok(())
}
#[cfg(feature = "cuda")]
fn find_cuda_device_in_cache(
cache: &StorageCache,
gpu_id: usize,
) -> Option<fuel_cuda_backend::CudaDevice> {
for arc in cache.values() {
let guard = arc.read().ok()?;
if let fuel_memory::BackendStorage::Cuda(c) = &guard.inner {
if matches!(c.device().location(), DeviceLocation::Cuda { gpu_id: g } if g == gpu_id) {
return Some(c.device().clone());
}
}
}
None
}
#[cfg(feature = "vulkan")]
struct VulkanCompletion {
backend: Arc<fuel_vulkan_backend::VulkanBackend>,
batch: Option<fuel_vulkan_backend::SubmittedBatch>,
loc: fuel_ir::DeviceLocation,
}
#[cfg(feature = "vulkan")]
type InflightVulkan = VulkanCompletion;
#[cfg(not(feature = "vulkan"))]
type InflightVulkan = ();
#[cfg(feature = "vulkan")]
fn copy_source_is_vulkan(cache: &StorageCache, producer: NodeId) -> Result<bool> {
if let Some(arc) = cache.get(&producer) {
let guard = arc
.read()
.map_err(|_| poisoned("storage lock probing copy-source backend"))?;
return Ok(matches!(
guard.inner,
fuel_memory::BackendStorage::Vulkan(_)
));
}
Ok(false)
}
#[cfg(not(feature = "vulkan"))]
#[inline]
fn copy_source_is_vulkan(_cache: &StorageCache, _producer: NodeId) -> Result<bool> {
Ok(false)
}
#[cfg(feature = "vulkan")]
impl crate::compiled::Completion for VulkanCompletion {
fn wait(self: Box<Self>) -> Result<()> {
(*self).into_wait()
}
}
#[cfg(feature = "vulkan")]
impl VulkanCompletion {
fn new(
backend: Arc<fuel_vulkan_backend::VulkanBackend>,
batch: fuel_vulkan_backend::SubmittedBatch,
loc: fuel_ir::DeviceLocation,
) -> Self {
crate::dispatch::inflight_inc(loc);
Self {
backend,
batch: Some(batch),
loc,
}
}
fn into_wait(mut self) -> Result<()> {
if let Some(batch) = self.batch.take() {
self.backend.wait_submitted(batch)?;
}
Ok(())
}
fn retain_data(&mut self, buf: std::sync::Arc<fuel_vulkan_backend::VulkanBuffer>) {
if let Some(batch) = self.batch.as_mut() {
batch.retain_buffer(buf);
}
}
}
#[cfg(feature = "vulkan")]
impl Drop for VulkanCompletion {
fn drop(&mut self) {
crate::dispatch::inflight_dec(self.loc);
}
}
#[cfg(feature = "vulkan")]
fn eager_submit_all_vulkan(
cache: &StorageCache,
inflight: &mut Vec<VulkanCompletion>,
) -> Result<()> {
let mut seen: Vec<usize> = Vec::new();
let mut backends: Vec<Arc<fuel_vulkan_backend::VulkanBackend>> = Vec::new();
for arc in cache.values() {
let guard = arc
.read()
.map_err(|_| poisoned("storage lock during eager vulkan submit"))?;
if let fuel_memory::BackendStorage::Vulkan(v) = &guard.inner
&& let Some(backend) = v.backend()
&& !seen.contains(&backend.gpu_id)
{
seen.push(backend.gpu_id);
backends.push(backend.clone());
}
}
for backend in backends {
let loc = DeviceLocation::Vulkan {
gpu_id: backend.gpu_id,
};
if let Some(batch) = backend.submit_pending()? {
inflight.push(VulkanCompletion::new(backend, batch, loc));
}
}
Ok(())
}
#[cfg(not(feature = "vulkan"))]
#[inline]
fn eager_submit_all_vulkan(
_cache: &StorageCache,
_inflight: &mut Vec<InflightVulkan>,
) -> Result<()> {
Ok(())
}
#[cfg(feature = "vulkan")]
fn drain_inflight_vulkan(inflight: &mut Vec<VulkanCompletion>) -> Result<()> {
for vc in inflight.drain(..) {
vc.into_wait()?;
}
Ok(())
}
#[cfg(not(feature = "vulkan"))]
#[inline]
fn drain_inflight_vulkan(_inflight: &mut Vec<InflightVulkan>) -> Result<()> {
Ok(())
}
#[cfg(feature = "vulkan")]
fn drain_vulkan_pending(cache: &StorageCache) -> Result<()> {
use crate::compiled::CompletionHandle;
for arc in cache.values() {
let backend = {
let guard = arc
.read()
.map_err(|_| poisoned("storage lock during vulkan submit_pending"))?;
match &guard.inner {
fuel_memory::BackendStorage::Vulkan(v) => v.backend().cloned(),
_ => None,
}
};
let Some(backend) = backend else { continue };
let loc = DeviceLocation::Vulkan {
gpu_id: backend.gpu_id,
};
if let Some(batch) = backend.submit_pending()? {
let handle: CompletionHandle =
CompletionHandle::Pending(Box::new(VulkanCompletion::new(backend, batch, loc)));
handle.wait()?;
}
}
Ok(())
}
#[cfg(not(feature = "vulkan"))]
fn drain_vulkan_pending(_cache: &StorageCache) -> Result<()> {
Ok(())
}
#[cfg(feature = "vulkan")]
fn force_flush_vulkan(arc: &Arc<RwLock<Storage>>) -> Result<()> {
let guard = arc
.read()
.map_err(|_| poisoned("storage lock during vulkan force_flush"))?;
if let fuel_memory::BackendStorage::Vulkan(v) = &guard.inner
&& let Some(backend) = v.backend()
{
backend.force_flush()?;
}
Ok(())
}
#[cfg(not(feature = "vulkan"))]
fn force_flush_vulkan(_arc: &Arc<RwLock<Storage>>) -> Result<()> {
Ok(())
}
#[cfg(feature = "vulkan")]
fn defer_evicted_vulkan_buffer(
cache: &StorageCache,
destroyed: NodeId,
inflight: &mut Vec<VulkanCompletion>,
) -> Result<()> {
let buf = {
let Some(arc) = cache.get(&destroyed) else {
return Ok(());
};
let guard = arc
.read()
.map_err(|_| poisoned("storage lock during deferred-eviction retain"))?;
match &guard.inner {
fuel_memory::BackendStorage::Vulkan(v) => v.device_buffer_arc(),
_ => None,
}
};
let Some(buf) = buf else { return Ok(()) };
eager_submit_all_vulkan(cache, inflight)?;
for vc in inflight.iter_mut() {
vc.retain_data(buf.clone());
}
if !inflight.is_empty() {
DEFERRED_EVICTION_RETAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
Ok(())
}
#[cfg(feature = "vulkan")]
pub static DEFERRED_EVICTION_RETAINS: std::sync::atomic::AtomicUsize =
std::sync::atomic::AtomicUsize::new(0);
#[cfg(feature = "vulkan")]
pub fn deferred_eviction_retains() -> usize {
DEFERRED_EVICTION_RETAINS.load(std::sync::atomic::Ordering::Relaxed)
}
#[cfg(not(feature = "vulkan"))]
#[inline]
fn defer_evicted_vulkan_buffer(
_cache: &StorageCache,
_destroyed: NodeId,
_inflight: &mut Vec<InflightVulkan>,
) -> Result<()> {
Ok(())
}
#[cfg(feature = "vulkan")]
fn find_vulkan_backend_in_cache(
cache: &StorageCache,
gpu_id: usize,
) -> Option<std::sync::Arc<fuel_vulkan_backend::VulkanBackend>> {
for arc in cache.values() {
let guard = arc.read().ok()?;
if let fuel_memory::BackendStorage::Vulkan(v) = &guard.inner
&& let Some(backend) = v.backend()
&& backend.gpu_id == gpu_id
{
return Some(backend.clone());
}
}
None
}
fn auto_contiguize(arc: &Arc<RwLock<Storage>>, layout: &Layout) -> Result<Arc<RwLock<Storage>>> {
let in_guard = arc
.read()
.map_err(|_| poisoned("input storage lock during auto_contiguize"))?;
let dtype = in_guard.dtype;
let dtype_size = dtype.size_in_bytes();
let new_storage = match &in_guard.inner {
fuel_memory::BackendStorage::Cpu(c) => {
let new_bytes = fuel_cpu_backend::byte_kernels::contiguize_cpu(c, layout, dtype_size)?;
Storage::new(fuel_memory::BackendStorage::Cpu(new_bytes), dtype)
}
#[cfg(feature = "cuda")]
fuel_memory::BackendStorage::Cuda(c) => {
let contig = fuel_cuda_backend::baracuda::contiguize::contiguize_to_fresh(
c, layout, dtype_size,
)?;
Storage::new(fuel_memory::BackendStorage::Cuda(contig), dtype)
}
#[cfg(feature = "vulkan")]
fuel_memory::BackendStorage::Vulkan(v) => {
let backend = v.backend().ok_or_else(|| {
Error::Msg(
"auto_contiguize: Vulkan input has no backend handle. \
Storages flowing through the pipelined executor's Vulkan \
path must be constructed via VulkanBackend::alloc_bytes_handle \
/ upload_bytes_handle (not the legacy alloc_bytes / \
upload_bytes which leave VulkanStorageBytes::backend = None)."
.to_string(),
)
.bt()
})?;
let host_bytes = backend.download_bytes(v)?;
let host_cpu = fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(&host_bytes);
let contig_cpu =
fuel_cpu_backend::byte_kernels::contiguize_cpu(&host_cpu, layout, dtype_size)?;
let host_contig_bytes = contig_cpu.bytes().to_vec();
let vk_bytes = backend.upload_bytes_handle(&host_contig_bytes)?;
Storage::new(fuel_memory::BackendStorage::Vulkan(vk_bytes), dtype)
}
#[allow(unreachable_patterns)]
_ => {
return Err(Error::Msg(
"auto_contiguize: backend not wired (CPU + CUDA + Vulkan are \
wired today; Metal extends this match when its byte-storage \
substrate lands)"
.to_string(),
)
.bt());
}
};
Ok(Arc::new(RwLock::new(new_storage)))
}
#[cfg_attr(not(feature = "cuda"), allow(unused_variables))]
fn contiguize_into(
dest: &Arc<RwLock<Storage>>,
arc: &Arc<RwLock<Storage>>,
layout: &Layout,
) -> Result<()> {
let in_guard = arc
.read()
.map_err(|_| poisoned("input storage lock during contiguize_into"))?;
match &in_guard.inner {
#[cfg(feature = "cuda")]
fuel_memory::BackendStorage::Cuda(src_cuda) => {
let dtype_size = in_guard.dtype.size_in_bytes();
let mut dest_guard = dest
.write()
.map_err(|_| poisoned("dest storage lock during contiguize_into"))?;
let dest_cuda = match &mut dest_guard.inner {
fuel_memory::BackendStorage::Cuda(d) => d,
other => {
return Err(Error::Msg(format!(
"contiguize_into: dest is BackendStorage::{:?}, not Cuda",
std::mem::discriminant(other),
))
.bt());
}
};
fuel_cuda_backend::baracuda::contiguize::contiguize_into(
src_cuda, layout, dest_cuda, dtype_size,
)
}
#[allow(unreachable_patterns)]
_ => Err(Error::Msg(
"contiguize_into: only CUDA is wired (PersistentOutputs::reuse / \
capture_decode is CUDA-only today — no CPU/Vulkan consumer exists \
yet; extend this match if one appears)"
.to_string(),
)
.bt()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_graph::Node;
use fuel_ir::Shape;
#[test]
fn scan_has_no_dispatch_op_kind() {
use fuel_graph::{Op, ScanEmit};
assert_eq!(
crate::pipelined::op_to_op_kind(&Op::Scan {
n_carries: 1,
n_xs: 0,
bound: 2,
emit: ScanEmit::All,
early_exit: None
}),
None,
);
}
#[test]
fn build_wait_set_collects_only_cross_device_copy_producers() {
let mut g = Graph::new();
let const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let same_device_id = g.push(Node {
op: Op::Contiguize,
inputs: vec![const_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let copy_id = g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cuda { gpu_id: 0 },
},
inputs: vec![same_device_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let order = vec![const_id, same_device_id, copy_id];
let set = build_wait_set(&g, &order);
assert!(
set.contains(&same_device_id),
"the Copy's producer must be in the wait set"
);
assert!(
!set.contains(&const_id),
"a node with no Copy/Move consumer must be absent from the wait set"
);
assert!(
!set.contains(©_id),
"the Copy node itself is not its own producer"
);
assert_eq!(
set.len(),
1,
"only the one true cross-device producer is in the set"
);
}
#[test]
fn build_wait_set_collects_move_producers_too() {
let mut g = Graph::new();
let const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let move_id = g.push(Node {
op: Op::Move {
target: DeviceLocation::Cpu,
},
inputs: vec![const_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let order = vec![const_id, move_id];
let set = build_wait_set(&g, &order);
assert!(
set.contains(&const_id),
"the Move's producer must be in the wait set"
);
assert_eq!(set.len(), 1);
}
#[test]
fn build_wait_set_is_empty_for_cpu_only_order() {
let mut g = Graph::new();
let const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let contig_id = g.push(Node {
op: Op::Contiguize,
inputs: vec![const_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let order = vec![const_id, contig_id];
let set = build_wait_set(&g, &order);
assert!(
set.is_empty(),
"a CPU-only order (no Copy/Move) yields an empty wait set"
);
}
#[test]
fn will_be_waited_none_is_conservative_true_for_any_node() {
let ws: WaitSet = None;
assert!(will_be_waited(&ws, NodeId(0)));
assert!(will_be_waited(&ws, NodeId(12345)));
}
#[test]
fn will_be_waited_some_checks_membership() {
let mut set = HashSet::new();
set.insert(NodeId(3));
let ws: WaitSet = Some(set);
assert!(will_be_waited(&ws, NodeId(3)), "member must be waited");
assert!(!will_be_waited(&ws, NodeId(4)), "non-member must be elided");
}
struct FlagCompletion(std::sync::Arc<std::sync::atomic::AtomicUsize>);
impl crate::compiled::Completion for FlagCompletion {
fn wait(self: Box<Self>) -> Result<()> {
self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(())
}
}
#[test]
fn handle_map_stores_pending_and_drains_once() {
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
store_handle(&mut handles, NodeId(1), CompletionHandle::Ready);
assert!(handles.is_empty(), "Ready handles must not be stored");
let waited = StdArc::new(AtomicUsize::new(0));
store_handle(
&mut handles,
NodeId(2),
CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
);
store_handle(
&mut handles,
NodeId(3),
CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
);
assert_eq!(handles.len(), 2);
drain_handles(&mut handles).expect("drain");
assert_eq!(
waited.load(Ordering::SeqCst),
2,
"each pending handle waited once"
);
assert!(handles.is_empty(), "map empty after drain");
}
#[test]
fn handle_map_rekey_waits_previous() {
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
let waited = StdArc::new(AtomicUsize::new(0));
store_handle(
&mut handles,
NodeId(7),
CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
);
store_handle(
&mut handles,
NodeId(7),
CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
);
assert_eq!(
waited.load(Ordering::SeqCst),
1,
"rekey waited the prior pending"
);
store_handle(&mut handles, NodeId(7), CompletionHandle::Ready);
assert_eq!(
waited.load(Ordering::SeqCst),
2,
"Ready-over-Pending waited it"
);
assert!(handles.is_empty());
drain_handles(&mut handles).expect("drain empty");
}
#[test]
fn wait_producer_handle_consumes_named_handle_once() {
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
let mut handles: HashMap<NodeId, CompletionHandle> = HashMap::new();
let waited = StdArc::new(AtomicUsize::new(0));
store_handle(
&mut handles,
NodeId(2),
CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
);
store_handle(
&mut handles,
NodeId(5),
CompletionHandle::Pending(Box::new(FlagCompletion(waited.clone()))),
);
wait_producer_handle(&mut handles, NodeId(99)).expect("absent is no-op");
assert_eq!(
waited.load(Ordering::SeqCst),
0,
"absent producer waited nothing"
);
assert_eq!(handles.len(), 2);
wait_producer_handle(&mut handles, NodeId(2)).expect("wait source");
assert_eq!(
waited.load(Ordering::SeqCst),
1,
"source producer waited once"
);
assert!(!handles.contains_key(&NodeId(2)), "source handle consumed");
assert!(
handles.contains_key(&NodeId(5)),
"unrelated handle untouched"
);
wait_producer_handle(&mut handles, NodeId(2)).expect("re-wait is no-op");
assert_eq!(waited.load(Ordering::SeqCst), 1, "no double-wait");
drain_handles(&mut handles).expect("drain");
assert_eq!(
waited.load(Ordering::SeqCst),
2,
"drain waited the remaining one"
);
assert!(handles.is_empty());
}
#[test]
fn op_contiguize_zero_copy_on_contiguous_input() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (const_id, contig_id) = {
let mut g = graph.write().unwrap();
let const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
let contig_id = g.push(Node {
op: Op::Contiguize,
inputs: vec![const_id],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
(const_id, contig_id)
};
let mut inputs = StorageCache::new();
let input_arc = Arc::new(RwLock::new(storage));
let input_arc_for_compare = Arc::clone(&input_arc);
inputs.insert(const_id, input_arc);
let (result_arc, _) =
PipelinedExecutor::realize(graph, contig_id, inputs).expect("realize");
assert!(
Arc::ptr_eq(&result_arc, &input_arc_for_compare),
"Op::Contiguize on contiguous input must adopt the input Arc (zero copy)",
);
}
#[test]
fn op_contiguize_materializes_on_strided_input() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (const_id, _transpose_id, contig_id) = {
let mut g = graph.write().unwrap();
let const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let transpose_id = g.push(Node {
op: Op::Transpose,
inputs: vec![const_id],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let contig_id = g.push(Node {
op: Op::Contiguize,
inputs: vec![transpose_id],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
(const_id, transpose_id, contig_id)
};
let mut inputs = StorageCache::new();
inputs.insert(const_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) =
PipelinedExecutor::realize(graph, contig_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let cpu = if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
c
} else {
panic!("expected Cpu storage");
};
let typed: &[f32] = cpu.as_slice().expect("f32 cast");
assert_eq!(
typed,
&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0],
"contiguized transpose must reflect logical (transposed) byte order",
);
}
#[test]
fn pipelined_realize_const_const_add() {
let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
let rhs_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, add_id) = {
let mut g = graph.write().unwrap();
let lhs_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add_id = g.push(Node {
op: Op::Add,
inputs: vec![lhs_id, rhs_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(add_id, BackendId::Cpu);
(lhs_id, rhs_id, add_id)
};
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
let (result_arc, _result_layout) =
PipelinedExecutor::realize(graph, add_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().expect("f32 cast");
assert_eq!(typed, &[11.0, 22.0, 33.0]);
} else {
panic!("expected CPU output");
}
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_binary_chain_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
);
};
let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId, NodeId) {
let graph = Arc::new(RwLock::new(Graph::new()));
let (a, b, c, z) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::Add,
inputs: vec![a, b],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let z = g.push(Node {
op: Op::Mul,
inputs: vec![y, c],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(y, BackendId::Cuda);
g.set_target_backend(z, BackendId::Cuda);
(a, b, c, z)
};
(graph, a, b, c, z)
};
let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
let bytes: &[u8] = bytemuck::cast_slice(v);
let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytes).expect("h2d");
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(cb),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
let host = c.to_cpu_bytes().expect("d2h");
bytemuck::cast_slice(&host).to_vec()
};
let a_data = [1.0_f32, 2.0, 3.0, 4.0];
let b_data = [10.0_f32, 20.0, 30.0, 40.0];
let c_data = [2.0_f32, 2.0, 2.0, 2.0];
let expect0 = [22.0_f32, 44.0, 66.0, 88.0];
let a_arc = up(&a_data);
let b_arc = up(&b_data);
let c_arc = up(&c_data);
let (graph, a, b, c, z) = build();
let mut inputs = StorageCache::new();
inputs.insert(a, a_arc.clone());
inputs.insert(b, b_arc.clone());
inputs.insert(c, c_arc.clone());
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), z, inputs, SymEnv::default())
.expect("capture_decode");
assert_eq!(
read(&captured.output),
&expect0,
"warm-run result after capture"
);
captured.run.replay(&dev).expect("replay");
dev.synchronize().expect("sync");
assert_eq!(read(&captured.output), &expect0, "replay reproduces result");
let a2 = [100.0_f32, 200.0, 300.0, 400.0];
{
let g = a_arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("a not cuda")
};
c.write_from_host(bytemuck::cast_slice(&a2))
.expect("in-place H2D");
}
captured.run.replay(&dev).expect("replay 2");
dev.synchronize().expect("sync 2");
let got = read(&captured.output);
let (graph_ref, ar, br, cr, zr) = build();
let mut inputs_ref = StorageCache::new();
inputs_ref.insert(ar, up(&a2));
inputs_ref.insert(br, up(&b_data));
inputs_ref.insert(cr, up(&c_data));
let (ref_arc, _) =
PipelinedExecutor::realize(graph_ref, zr, inputs_ref).expect("ref realize");
let expect1 = read(&ref_arc); assert_eq!(
got, expect1,
"replay after in-place input update matches uncaptured realize"
);
assert_eq!(expect1, vec![220.0, 440.0, 660.0, 880.0]);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_matmul_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else { return };
const H: usize = 4;
const F: usize = 8;
let a_data: Vec<f32> = (0..H).map(|i| 0.5 + i as f32).collect();
let w_data: Vec<f32> = (0..H * F).map(|i| 0.01 * (i as f32 + 1.0)).collect();
let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId) {
let graph = Arc::new(RwLock::new(Graph::new()));
let (a, w, y) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, H]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[H, F]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::MatMul,
inputs: vec![a, w],
shape: Shape::from_dims(&[1, F]),
dtype: DType::F32,
});
g.set_target_backend(y, BackendId::Cuda);
(a, w, y)
};
(graph, a, w, y)
};
let up = |v: &[f32]| {
let b: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
c.to_cpu_bytes()
.unwrap()
.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect()
};
let (gr, ar, wr, yr) = build();
let mut insr = StorageCache::new();
insr.insert(ar, up(&a_data));
insr.insert(wr, up(&w_data));
let (refarc, _) = PipelinedExecutor::realize(gr, yr, insr).unwrap();
let expect = read(&refarc);
let (graph, a, w, y) = build();
let mut inputs = StorageCache::new();
inputs.insert(a, up(&a_data));
inputs.insert(w, up(&w_data));
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
.expect("capture");
assert_eq!(read(&captured.output), expect, "warm result after capture");
let mut squatters = Vec::new();
for _ in 0..8 {
let g: Vec<u8> = (0..H * F).flat_map(|_| 999.0f32.to_le_bytes()).collect();
squatters.push(CudaStorageBytes::from_cpu_bytes(&dev, &g).unwrap());
}
dev.synchronize().unwrap();
let mut bad = 0usize;
for _ in 0..50 {
captured.run.replay(&dev).unwrap();
dev.synchronize().unwrap();
if read(&captured.output) != expect {
bad += 1;
}
}
assert_eq!(
bad, 0,
"cuBLAS matmul capture replays bit-exact ({bad}/50 wrong)"
);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_affine_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else { return };
const N: usize = 6;
let x_data: Vec<f32> = (0..N).map(|i| 1.0 + i as f32).collect();
let expect: Vec<f32> = x_data.iter().map(|v| v * 0.5).collect();
let graph = Arc::new(RwLock::new(Graph::new()));
let (x, y) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, N]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::MulScalar(0.5),
inputs: vec![x],
shape: Shape::from_dims(&[1, N]),
dtype: DType::F32,
});
g.set_target_backend(y, BackendId::Cuda);
(x, y)
};
let up = |v: &[f32]| {
let b: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
c.to_cpu_bytes()
.unwrap()
.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect()
};
let mut inputs = StorageCache::new();
inputs.insert(x, up(&x_data));
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
.expect("capture");
assert_eq!(read(&captured.output), expect, "warm result after capture");
let mut sq = Vec::new();
for _ in 0..8 {
let g: Vec<u8> = (0..N).flat_map(|_| 999.0f32.to_le_bytes()).collect();
sq.push(CudaStorageBytes::from_cpu_bytes(&dev, &g).unwrap());
}
dev.synchronize().unwrap();
let mut bad = 0usize;
for _ in 0..50 {
captured.run.replay(&dev).unwrap();
dev.synchronize().unwrap();
if read(&captured.output) != expect {
bad += 1;
}
}
assert_eq!(bad, 0, "affine capture replays bit-exact ({bad}/50 wrong)");
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_index_select_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else { return };
const V: usize = 3;
const H: usize = 4;
let wte_data: Vec<f32> = (0..V)
.flat_map(|v| (0..H).map(move |c| 1.0 + 10.0 * v as f32 + c as f32))
.collect();
let graph = Arc::new(RwLock::new(Graph::new()));
let (tok, wte, emb) = {
let mut g = graph.write().unwrap();
let tok = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1]),
dtype: DType::U32,
});
let wte = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[V, H]),
dtype: DType::F32,
});
let emb = g.push(Node {
op: Op::IndexSelect { dim: 0 },
inputs: vec![wte, tok],
shape: Shape::from_dims(&[1, H]),
dtype: DType::F32,
});
g.set_target_backend(emb, BackendId::Cuda);
(tok, wte, emb)
};
let up_f32 = |v: &[f32]| {
let b: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
DType::F32,
)))
};
let up_u32 = |v: &[u32]| {
let b: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(CudaStorageBytes::from_cpu_bytes(&dev, &b).unwrap()),
DType::U32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
c.to_cpu_bytes()
.unwrap()
.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect()
};
let tok_arc = up_u32(&[0]);
let mut inputs = StorageCache::new();
inputs.insert(tok, tok_arc.clone());
inputs.insert(wte, up_f32(&wte_data));
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), emb, inputs, SymEnv::default())
.expect("capture");
assert_eq!(
read(&captured.output),
vec![1.0, 2.0, 3.0, 4.0],
"warm result after capture"
);
let mut squatters = Vec::new();
for _ in 0..8 {
let g: Vec<u8> = (0..V * H).flat_map(|_| 999.0f32.to_le_bytes()).collect();
squatters.push(CudaStorageBytes::from_cpu_bytes(&dev, &g).unwrap());
}
dev.synchronize().unwrap();
let mut bad = 0usize;
for _ in 0..50 {
captured.run.replay(&dev).unwrap();
dev.synchronize().unwrap();
if read(&captured.output) != vec![1.0, 2.0, 3.0, 4.0] {
bad += 1;
}
}
assert_eq!(
bad, 0,
"index_select capture replays bit-exact ({bad}/50 wrong)"
);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_norm_silu_softmax_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_graph::registry::{FusedOpParams, FusedOps};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
);
};
const H: usize = 6;
let eps = 1e-5f64;
let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId) {
let graph = Arc::new(RwLock::new(Graph::new()));
let (x, y) = {
let mut g = graph.write().unwrap();
let mk = |g: &mut Graph, op: Op, ins: Vec<NodeId>, dims: &[usize]| {
g.push(Node {
op,
inputs: ins,
shape: Shape::from_dims(dims),
dtype: DType::F32,
})
};
let rms = || {
Op::Fused(
FusedOps::RMS_NORM_LAST_DIM,
FusedOpParams::RmsNormLastDim { eps },
)
};
let sm = || Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim);
let x = mk(&mut g, Op::Const, vec![], &[1, H]);
let n1 = mk(&mut g, rms(), vec![x], &[1, H]);
let s = mk(&mut g, Op::Silu, vec![n1], &[1, H]);
let p = mk(&mut g, sm(), vec![s], &[1, H]);
let r = mk(&mut g, Op::Add, vec![p, x], &[1, H]);
let n2 = mk(&mut g, rms(), vec![r], &[1, H]);
let y = mk(&mut g, Op::Add, vec![n2, x], &[1, H]);
for id in [n1, s, p, r, n2, y] {
g.set_target_backend(id, BackendId::Cuda);
}
(x, y)
};
(graph, x, y)
};
let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(cb),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
};
let x_data: Vec<f32> = (0..H).map(|c| 0.1 * (c as f32 + 1.0)).collect();
let x_arc = up(&x_data);
let (graph, x, y) = build();
let mut inputs = StorageCache::new();
inputs.insert(x, x_arc.clone());
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
.expect("capture_decode norm/silu/softmax");
let realize_ref = |xv: &[f32]| -> Vec<f32> {
let (g, xr, yr) = build();
let mut ins = StorageCache::new();
ins.insert(xr, up(xv));
let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
read(&arc)
};
let expect0 = realize_ref(&x_data);
assert_eq!(read(&captured.output), expect0, "warm result after capture");
captured.run.replay(&dev).expect("replay");
dev.synchronize().expect("sync");
assert_eq!(read(&captured.output), expect0, "replay reproduces result");
let x2: Vec<f32> = (0..H).map(|c| 0.3 * (c as f32 + 1.0)).collect();
{
let g = x_arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("x not cuda")
};
c.write_from_host(bytemuck::cast_slice(&x2))
.expect("in-place H2D");
}
captured.run.replay(&dev).expect("replay 2");
dev.synchronize().expect("sync 2");
assert_eq!(
read(&captured.output),
realize_ref(&x2),
"replay after x update matches realize"
);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_contiguize_then_concat_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
);
};
let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId) {
let graph = Arc::new(RwLock::new(Graph::new()));
let (x, b, y) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3, 4]),
dtype: DType::F32,
});
let p = g.push(Node {
op: Op::Permute(vec![1, 0, 2]),
inputs: vec![x],
shape: Shape::from_dims(&[3, 2, 4]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Reshape(Shape::from_dims(&[3, 8])),
inputs: vec![p],
shape: Shape::from_dims(&[3, 8]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 8]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::Concat { dim: 1 },
inputs: vec![r, b],
shape: Shape::from_dims(&[3, 16]),
dtype: DType::F32,
});
g.set_target_backend(p, BackendId::Cuda);
g.set_target_backend(r, BackendId::Cuda);
g.set_target_backend(y, BackendId::Cuda);
(x, b, y)
};
(graph, x, b, y)
};
let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(cb),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
};
let x_data: Vec<f32> = (0..24).map(|i| i as f32).collect();
let b_data: Vec<f32> = (0..24).map(|i| 100.0 + i as f32).collect();
let x_arc = up(&x_data);
let (graph, x, b, y) = build();
let mut inputs = StorageCache::new();
inputs.insert(x, x_arc.clone());
inputs.insert(b, up(&b_data));
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
.expect("capture_decode contiguize+concat");
let realize_ref = |xv: &[f32], bv: &[f32]| -> Vec<f32> {
let (g, xr, br, yr) = build();
let mut ins = StorageCache::new();
ins.insert(xr, up(xv));
ins.insert(br, up(bv));
let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
read(&arc)
};
let expect0 = realize_ref(&x_data, &b_data);
assert_eq!(read(&captured.output), expect0, "warm result after capture");
for i in 0..2 {
captured
.run
.replay(&dev)
.unwrap_or_else(|e| panic!("replay {i}: {e:?}"));
dev.synchronize().expect("sync");
assert_eq!(
read(&captured.output),
expect0,
"replay {i} reproduces result"
);
}
let x2: Vec<f32> = (0..24).map(|i| 1000.0 + i as f32).collect();
{
let g = x_arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("x not cuda")
};
c.write_from_host(bytemuck::cast_slice(&x2))
.expect("in-place H2D");
}
captured.run.replay(&dev).expect("replay 3");
dev.synchronize().expect("sync 3");
assert_eq!(
read(&captured.output),
realize_ref(&x2, &b_data),
"replay after x update matches independent uncaptured realize"
);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_same_device_copy_replays_bit_exact_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
);
};
let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId, NodeId) {
let graph = Arc::new(RwLock::new(Graph::new()));
let (x, b, y) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cuda { gpu_id: 0 },
},
inputs: vec![x],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::Add,
inputs: vec![c, b],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
g.set_target_backend(c, BackendId::Cuda);
g.set_target_backend(y, BackendId::Cuda);
(x, b, y)
};
(graph, x, b, y)
};
let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(cb),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
};
let x_data: Vec<f32> = (0..8).map(|i| i as f32).collect();
let b_data: Vec<f32> = (0..8).map(|i| 100.0 + i as f32).collect();
let x_arc = up(&x_data);
let (graph, x, b, y) = build();
let mut inputs = StorageCache::new();
inputs.insert(x, x_arc.clone());
inputs.insert(b, up(&b_data));
let captured =
PipelinedExecutor::capture_decode(Arc::clone(&graph), y, inputs, SymEnv::default())
.expect("capture_decode same-device copy");
let realize_ref = |xv: &[f32], bv: &[f32]| -> Vec<f32> {
let (g, xr, br, yr) = build();
let mut ins = StorageCache::new();
ins.insert(xr, up(xv));
ins.insert(br, up(bv));
let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
read(&arc)
};
let expect0 = realize_ref(&x_data, &b_data);
assert_eq!(read(&captured.output), expect0, "warm result after capture");
for i in 0..2 {
captured
.run
.replay(&dev)
.unwrap_or_else(|e| panic!("replay {i}: {e:?}"));
dev.synchronize().expect("sync");
assert_eq!(
read(&captured.output),
expect0,
"replay {i} reproduces result"
);
}
let x2: Vec<f32> = (0..8).map(|i| 1000.0 + i as f32).collect();
{
let g = x_arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("x not cuda")
};
c.write_from_host(bytemuck::cast_slice(&x2))
.expect("in-place H2D");
}
captured.run.replay(&dev).expect("replay 3");
dev.synchronize().expect("sync 3");
assert_eq!(
read(&captured.output),
realize_ref(&x2, &b_data),
"replay after x update matches independent uncaptured realize"
);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn capture_decode_rejects_cross_device_copy_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
);
};
let graph = Arc::new(RwLock::new(Graph::new()));
let (x, y) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cpu,
},
inputs: vec![x],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
g.set_target_backend(y, BackendId::Cuda); (x, y)
};
let x_data: Vec<f32> = (0..8).map(|i| i as f32).collect();
let cb =
CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(&x_data)).expect("h2d");
let x_arc = Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(cb),
DType::F32,
)));
let mut inputs = StorageCache::new();
inputs.insert(x, x_arc);
let result = PipelinedExecutor::capture_decode(graph, y, inputs, SymEnv::default());
let msg = match result {
Ok(_) => panic!("D2H Op::Copy must be rejected by capture_decode"),
Err(e) => e.to_string(),
};
assert!(
msg.contains("cross-device"),
"error should describe the cross-device rejection: {msg}"
);
}
#[test]
#[ignore = "requires a live CUDA device"]
#[cfg(feature = "cuda")]
fn captured_decode_session_replays_many_tokens_cuda() {
use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
use fuel_graph::registry::{FusedOpParams, FusedOps};
use fuel_ir::Shape;
use fuel_memory::BackendStorage;
let Ok(dev) = CudaDevice::new(0) else {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
);
};
const H: usize = 6;
let eps = 1e-5f64;
let build = || -> (Arc<RwLock<Graph>>, NodeId, NodeId) {
let graph = Arc::new(RwLock::new(Graph::new()));
let (x, y) = {
let mut g = graph.write().unwrap();
let mk = |g: &mut Graph, op: Op, ins: Vec<NodeId>, dims: &[usize]| {
g.push(Node {
op,
inputs: ins,
shape: Shape::from_dims(dims),
dtype: DType::F32,
})
};
let rms = || {
Op::Fused(
FusedOps::RMS_NORM_LAST_DIM,
FusedOpParams::RmsNormLastDim { eps },
)
};
let sm = || Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim);
let x = mk(&mut g, Op::Const, vec![], &[1, H]);
let n1 = mk(&mut g, rms(), vec![x], &[1, H]);
let s = mk(&mut g, Op::Silu, vec![n1], &[1, H]);
let p = mk(&mut g, sm(), vec![s], &[1, H]);
let r = mk(&mut g, Op::Add, vec![p, x], &[1, H]);
let n2 = mk(&mut g, rms(), vec![r], &[1, H]);
let y = mk(&mut g, Op::Add, vec![n2, x], &[1, H]);
for id in [n1, s, p, r, n2, y] {
g.set_target_backend(id, BackendId::Cuda);
}
(x, y)
};
(graph, x, y)
};
let up = |v: &[f32]| -> Arc<RwLock<Storage>> {
let cb = CudaStorageBytes::from_cpu_bytes(&dev, bytemuck::cast_slice(v)).expect("h2d");
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cuda(cb),
DType::F32,
)))
};
let read = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let g = arc.read().unwrap();
let BackendStorage::Cuda(c) = &g.inner else {
panic!("not cuda")
};
bytemuck::cast_slice(&c.to_cpu_bytes().expect("d2h")).to_vec()
};
let realize_ref = |xv: &[f32]| -> Vec<f32> {
let (g, xr, yr) = build();
let mut ins = StorageCache::new();
ins.insert(xr, up(xv));
let (arc, _) = PipelinedExecutor::realize(g, yr, ins).expect("ref realize");
read(&arc)
};
let tokens: Vec<Vec<f32>> = (0..5)
.map(|t| {
(0..H)
.map(|c| 0.1 * (t as f32 + 1.0) * (c as f32 + 1.0))
.collect()
})
.collect();
let x_arc = up(&tokens[0]);
let (graph, x, y) = build();
let mut inputs = StorageCache::new();
inputs.insert(x, x_arc.clone());
let session =
CapturedDecodeSession::capture(Arc::clone(&graph), y, inputs, &[x], SymEnv::default())
.expect("capture session");
assert_eq!(
read(session.output()),
realize_ref(&tokens[0]),
"warm token 0"
);
for (t, tok) in tokens.iter().enumerate() {
let out = session
.replay_token(&[(x, bytemuck::cast_slice(tok))])
.expect("replay_token");
assert_eq!(
read(&out),
realize_ref(tok),
"replayed token {t} matches uncaptured realize"
);
}
assert_ne!(
realize_ref(&tokens[0]),
realize_ref(&tokens[1]),
"tokens must differ"
);
}
#[test]
fn realize_with_optimized_route_empty_equals_arm0() {
use crate::optimize::OptimizedGraph;
let build = || {
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, add_id) = {
let mut g = graph.write().unwrap();
let lhs_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add_id = g.push(Node {
op: Op::Add,
inputs: vec![lhs_id, rhs_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(add_id, BackendId::Cpu);
(lhs_id, rhs_id, add_id)
};
let mut inputs = StorageCache::new();
inputs.insert(
lhs_id,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
1.0_f32, 2.0, 3.0,
]))),
);
inputs.insert(
rhs_id,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
10.0_f32, 20.0, 30.0,
]))),
);
(graph, add_id, inputs)
};
let read_bytes = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let guard = arc.read().unwrap();
match &guard.inner {
fuel_memory::BackendStorage::Cpu(c) => {
c.as_slice::<f32>().expect("f32 cast").to_vec()
}
_ => panic!("expected CPU output"),
}
};
drop(global_bindings());
let live_gen = crate::dispatch::topology_generation();
let (g_a, add_a, in_a) = build();
let opt_a = OptimizedGraph {
roots: vec![add_a],
generation: live_gen,
placements: None,
};
let (arm0_arc, _) = PipelinedExecutor::realize_with_optimized(g_a, add_a, in_a, &opt_a)
.expect("arm-0 realize");
let arm0 = read_bytes(&arm0_arc);
let (g_r, add_r, in_r) = build();
let opt_r = OptimizedGraph {
roots: vec![add_r],
generation: live_gen,
placements: None,
};
let empty_route = PickedRoute::new();
let (route_arc, _) =
PipelinedExecutor::realize_with_optimized_route(g_r, add_r, in_r, &opt_r, &empty_route)
.expect("empty-route realize");
let routed = read_bytes(&route_arc);
assert_eq!(arm0, vec![11.0, 22.0, 33.0]);
assert_eq!(
routed, arm0,
"an empty route realizes byte-identically to arm-0 (Phase B contract)",
);
}
#[test]
fn realize_with_optimized_picking_cached_env_matches_arm0_cold_and_warm() {
use crate::optimize::OptimizedGraph;
use std::sync::OnceLock;
let build = || {
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, add_id) = {
let mut g = graph.write().unwrap();
let lhs_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add_id = g.push(Node {
op: Op::Add,
inputs: vec![lhs_id, rhs_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(add_id, BackendId::Cpu);
(lhs_id, rhs_id, add_id)
};
let mut inputs = StorageCache::new();
inputs.insert(
lhs_id,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
1.0_f32, 2.0, 3.0,
]))),
);
inputs.insert(
rhs_id,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
10.0_f32, 20.0, 30.0,
]))),
);
(graph, add_id, inputs)
};
let read_bytes = |arc: &Arc<RwLock<Storage>>| -> Vec<f32> {
let guard = arc.read().unwrap();
match &guard.inner {
fuel_memory::BackendStorage::Cpu(c) => {
c.as_slice::<f32>().expect("f32 cast").to_vec()
}
_ => panic!("expected CPU output"),
}
};
drop(global_bindings());
let live_gen = crate::dispatch::topology_generation();
let (g, add, inputs) = build();
let optimized = OptimizedGraph {
roots: vec![add],
generation: live_gen,
placements: None,
};
let order_cache: OnceLock<Vec<NodeId>> = OnceLock::new();
assert!(order_cache.get().is_none(), "sanity: cache starts empty");
let (cold_arc, _) = PipelinedExecutor::realize_with_optimized_picking_cached_env(
g.clone(),
add,
inputs,
&optimized,
&order_cache,
None,
None,
SymEnv::default(),
)
.expect("cold cached realize");
assert_eq!(read_bytes(&cold_arc), vec![11.0, 22.0, 33.0]);
assert!(
order_cache.get().is_some(),
"the cold call must populate order_cache, not leave it empty"
);
let cached_order_after_cold = order_cache.get().unwrap().clone();
let mut warm_inputs = StorageCache::new();
warm_inputs.insert(
g.read().unwrap().node(add).inputs[0],
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
100.0_f32, 200.0, 300.0,
]))),
);
warm_inputs.insert(
g.read().unwrap().node(add).inputs[1],
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
1.0_f32, 2.0, 3.0,
]))),
);
let (warm_arc, _) = PipelinedExecutor::realize_with_optimized_picking_cached_env(
g,
add,
warm_inputs,
&optimized,
&order_cache,
None,
None,
SymEnv::default(),
)
.expect("warm cached realize");
assert_eq!(
read_bytes(&warm_arc),
vec![101.0, 202.0, 303.0],
"warm call must realize its OWN rebound inputs, not replay stale output",
);
assert_eq!(
order_cache.get().unwrap(),
&cached_order_after_cold,
"the warm call must reuse the SAME cached order, not recompute a different one",
);
}
#[test]
fn cached_generation_for_matches_optimized_generation_for() {
let optimized = OptimizedGraph {
roots: vec![NodeId(0)],
generation: 42,
placements: None,
};
let order_cache: std::sync::OnceLock<Vec<NodeId>> = std::sync::OnceLock::new();
let optimized_gen = generation_for(&OrderSource::Optimized {
optimized: &optimized,
route: None,
});
let cached_gen = generation_for(&OrderSource::Cached {
optimized: &optimized,
order_cache: &order_cache,
});
assert_eq!(
optimized_gen,
Some(42),
"sanity: Optimized reports the OptimizedGraph's own generation",
);
assert_eq!(
cached_gen, optimized_gen,
"Cached must report the SAME generation as Optimized for the same \
OptimizedGraph -- the TopologyChanged check cannot tell the two apart, \
by construction, so caching cannot silently weaken it",
);
}
#[test]
fn c3_reorder_is_a_valid_topo_permutation() {
let n = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
})
};
let mut g = Graph::new();
let ca = n(&mut g, Op::Const, vec![]);
let cb = n(&mut g, Op::Const, vec![]);
let mut cprev = ca;
for _ in 0..4 {
cprev = n(&mut g, Op::Relu, vec![cprev]);
g.set_target_backend(cprev, BackendId::Cuda);
}
let c2 = cprev; let mut vprev = cb;
for _ in 0..2 {
vprev = n(&mut g, Op::Relu, vec![vprev]);
g.set_target_backend(vprev, BackendId::Vulkan);
}
let v2 = vprev; let copy = g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cuda { gpu_id: 0 },
},
inputs: vec![v2],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(copy, BackendId::Cuda);
let out = n(&mut g, Op::Add, vec![c2, copy]);
g.set_target_backend(out, BackendId::Cuda);
let unreordered = execution_plan(&g, &[out]);
let reordered = fuel_graph::lower_runs_arm0(&g, &[out]);
assert_ne!(
reordered, unreordered,
"the C3 reorder must change the order on a two-device graph; \
reordered={reordered:?} unreordered={unreordered:?}",
);
let mut a = unreordered.clone();
let mut b = reordered.clone();
a.sort();
b.sort();
assert_eq!(a, b, "the reorder is a permutation of the same node set");
let pos: std::collections::HashMap<NodeId, usize> = reordered
.iter()
.enumerate()
.map(|(i, &node)| (node, i))
.collect();
for (i, &node) in reordered.iter().enumerate() {
for &inp in &g.node(node).inputs {
if let Some(&j) = pos.get(&inp) {
assert!(
j < i,
"input {inp:?} of {node:?} must precede it in the reordered \
order; reordered={reordered:?}",
);
}
}
}
let p = |node: NodeId| pos[&node];
assert!(p(c2) < p(copy), "CUDA chunk before the cross-device copy");
assert!(p(v2) < p(copy), "Vulkan chunk before the cross-device copy");
assert!(
p(c2) < p(v2),
"the heavier CUDA chunk is emitted before the lighter Vulkan chunk \
(critical-path heavy-first); reordered={reordered:?}",
);
}
#[test]
fn c3_reorder_is_operand_order_invariant() {
let build = |cuda_first: bool| -> (Graph, NodeId) {
let n = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
})
};
let mut g = Graph::new();
let ca = n(&mut g, Op::Const, vec![]);
let cb = n(&mut g, Op::Const, vec![]);
let mut cprev = ca;
for _ in 0..4 {
cprev = n(&mut g, Op::Relu, vec![cprev]);
g.set_target_backend(cprev, BackendId::Cuda);
}
let c2 = cprev;
let mut vprev = cb;
for _ in 0..2 {
vprev = n(&mut g, Op::Relu, vec![vprev]);
g.set_target_backend(vprev, BackendId::Vulkan);
}
let v2 = vprev;
let copy = g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cuda { gpu_id: 0 },
},
inputs: vec![v2],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(copy, BackendId::Cuda);
let inputs = if cuda_first {
vec![c2, copy]
} else {
vec![copy, c2]
};
let out = n(&mut g, Op::Add, inputs);
g.set_target_backend(out, BackendId::Cuda);
(g, out)
};
let (g_cf, out_cf) = build(true);
let (g_vf, out_vf) = build(false);
let order_cf = fuel_graph::lower_runs_arm0(&g_cf, &[out_cf]);
let order_vf = fuel_graph::lower_runs_arm0(&g_vf, &[out_vf]);
let tag = |g: &Graph, ids: &[NodeId]| -> Vec<(String, Option<BackendId>)> {
ids.iter()
.map(|&id| (format!("{:?}", g.node(id).op), g.target_backend(id)))
.collect()
};
let seq_cf = tag(&g_cf, &order_cf);
let seq_vf = tag(&g_vf, &order_vf);
assert_eq!(
seq_cf, seq_vf,
"the C3-reordered dispatch order MUST be identical for both \
reconverge operand orders (operand-order-invariance — the \
executor then sees the same WorkItem stream regardless of \
operand order, so overlap is operand-insensitive).\n\
cuda-first = {seq_cf:?}\nvulkan-first = {seq_vf:?}",
);
}
#[cfg(test)]
#[derive(Debug)]
struct PickLast;
#[cfg(test)]
impl RuntimeSelector for PickLast {
fn select<'a>(
&self,
set: &'a crate::ranker::AlternativeSet,
) -> Option<&'a crate::ranker::Candidate> {
set.alternatives().last()
}
}
#[cfg(test)]
fn picking_diamond() -> (Arc<RwLock<Graph>>, NodeId, NodeId) {
let node = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
})
};
let mut g = Graph::new();
let pre = node(&mut g, Op::Const, vec![]);
let diverge = node(&mut g, Op::Relu, vec![pre]);
let arm0 = node(&mut g, Op::Silu, vec![diverge]);
let arm1 = node(&mut g, Op::GeluTanh, vec![diverge]);
g.set_target_backend(arm0, BackendId::Cuda);
g.set_target_backend(arm1, BackendId::Cpu);
let reconverge = node(&mut g, Op::Relu, vec![arm0]);
let mut b = g.open_branch(diverge);
b.add_arm(arm0);
b.add_arm(arm1);
let branch = b
.finalize_branches(&mut g, reconverge)
.expect("well-formed 2-arm branch")
.expect("2 arms survive");
let post = node(&mut g, Op::Tanh, vec![reconverge]);
(Arc::new(RwLock::new(g)), branch, post)
}
#[test]
fn streaming_pick_for_gates() {
let (graph, _branch, post) = picking_diamond();
let _ = post;
let sel: Arc<dyn RuntimeSelector> = Arc::new(PickLast);
assert!(
PipelinedExecutor::streaming_pick_for(&graph, Some(sel.clone()), None)
.expect("ok")
.is_some(),
"selector + branch ⇒ a streaming pick",
);
assert!(
PipelinedExecutor::streaming_pick_for(&graph, None, None)
.expect("ok")
.is_none(),
"no selector ⇒ no streaming pick (arm-0 lowering)",
);
let node = |g: &mut Graph, op: Op, inputs: Vec<NodeId>| {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
})
};
let bl_graph = {
let mut g = Graph::new();
let c1 = node(&mut g, Op::Const, vec![]);
let c2 = node(&mut g, Op::Const, vec![]);
let add = node(&mut g, Op::Add, vec![c1, c2]);
g.set_target_backend(add, BackendId::Cpu);
Arc::new(RwLock::new(g))
};
assert!(
PipelinedExecutor::streaming_pick_for(&bl_graph, Some(sel), None)
.expect("ok")
.is_none(),
"branchless graph ⇒ no branches ⇒ no streaming pick (fast-path)",
);
}
#[test]
fn streamed_order_equals_one_shot() {
let (graph, _branch, post) = picking_diamond();
let g = graph.read().unwrap();
let bindings = global_bindings();
for sel in [
Arc::new(crate::ranker::WinnerSelector) as Arc<dyn RuntimeSelector>,
Arc::new(PickLast) as Arc<dyn RuntimeSelector>,
] {
let one_shot_route = pick_route(&g, &[post], &bindings, sel.as_ref(), None);
let one_shot = fuel_graph::lower_picked_route(&g, &[post], &one_shot_route);
let streamed = fuel_graph::lower_picked_route_streaming(&g, &[post], |branch| {
resolve_branch(&g, branch, &bindings, sel.as_ref())
});
assert_eq!(
streamed, one_shot,
"streamed order must equal one-shot for selector {sel:?}",
);
}
}
#[test]
fn pipelined_errors_on_unset_target_backend() {
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, add_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
(lhs, rhs, add)
};
let mut inputs = StorageCache::new();
inputs.insert(
lhs_id,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]))),
);
inputs.insert(
rhs_id,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[3.0_f32, 4.0]))),
);
let result = PipelinedExecutor::realize(graph, add_id, inputs);
assert!(result.is_err(), "missing target_backend must error");
let msg = result.unwrap_err().to_string();
assert!(
msg.contains("target_backend"),
"error names the unmet precondition: {msg}"
);
}
#[test]
fn pipelined_realize_const_only() {
let storage = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0, 7.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let const_id = {
let mut g = graph.write().unwrap();
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
})
};
let mut inputs = StorageCache::new();
inputs.insert(const_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) =
PipelinedExecutor::realize(graph, const_id, inputs).expect("realize const");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[5.0, 6.0, 7.0]);
}
}
#[test]
fn pipelined_errors_on_missing_const_input() {
let graph = Arc::new(RwLock::new(Graph::new()));
let const_id = {
let mut g = graph.write().unwrap();
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
})
};
let inputs = StorageCache::new();
let result = PipelinedExecutor::realize(graph, const_id, inputs);
assert!(result.is_err());
}
#[test]
fn pipelined_realize_const_relu() {
let storage = fuel_memory::from_slice_cpu(&[-1.0_f32, 0.0, 0.5, -3.5, 7.25]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, relu_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let relu_id = g.push(Node {
op: Op::Relu,
inputs: vec![in_id],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
g.set_target_backend(relu_id, BackendId::Cpu);
(in_id, relu_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, relu_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().expect("f32 cast");
assert_eq!(typed, &[0.0, 0.0, 0.5, 0.0, 7.25]);
}
}
#[test]
fn pipelined_realize_chained_binary_ops() {
let a_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
let b_storage = fuel_memory::from_slice_cpu(&[3.0_f32, 5.0]);
let c_storage = fuel_memory::from_slice_cpu(&[2.0_f32, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, c_id, sub_id, mul_id, div_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let sub = g.push(Node {
op: Op::Sub,
inputs: vec![a, b],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let mul = g.push(Node {
op: Op::Mul,
inputs: vec![sub, c],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let div = g.push(Node {
op: Op::Div,
inputs: vec![mul, b],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(sub, BackendId::Cpu);
g.set_target_backend(mul, BackendId::Cpu);
g.set_target_backend(div, BackendId::Cpu);
(a, b, c, sub, mul, div)
};
let _ = (sub_id, mul_id);
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
inputs.insert(c_id, Arc::new(RwLock::new(c_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, div_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert!((typed[0] - (14.0_f32 / 3.0)).abs() < 1e-6);
assert!((typed[1] - 12.0).abs() < 1e-6);
}
}
#[test]
fn pipelined_realize_transpose_is_metadata_only() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, t_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let t_id = g.push(Node {
op: Op::Transpose,
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
(in_id, t_id)
};
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, t_id, inputs).expect("realize");
assert!(
Arc::ptr_eq(&result_arc, &in_arc),
"transpose must share input bytes"
);
assert_eq!(result_layout.shape().dims(), &[3, 2]);
assert_eq!(result_layout.stride(), &[1, 3]);
assert!(!result_layout.is_contiguous());
}
#[test]
fn pipelined_realize_flip_is_metadata_only() {
let data: Vec<f32> = (1..=12).map(|x| x as f32).collect();
let storage = fuel_memory::from_slice_cpu(&data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, f_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 4]),
dtype: DType::F32,
});
let f_id = g.push(Node {
op: Op::Flip { dim: 0 },
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 4]),
dtype: DType::F32,
});
(in_id, f_id)
};
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, f_id, inputs).expect("realize");
assert!(
Arc::ptr_eq(&result_arc, &in_arc),
"flip must share input bytes"
);
assert_eq!(result_layout.shape().dims(), &[3, 4]);
assert_eq!(result_layout.stride(), &[-4_isize, 1]);
assert_eq!(result_layout.start_offset(), 8);
assert!(!result_layout.is_contiguous());
}
#[test]
fn pipelined_realize_permute_is_metadata_only() {
let data: Vec<f32> = (1..=24).map(|x| x as f32).collect();
let storage = fuel_memory::from_slice_cpu(&data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, p_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3, 4]),
dtype: DType::F32,
});
let p_id = g.push(Node {
op: Op::Permute(vec![2, 0, 1]),
inputs: vec![in_id],
shape: Shape::from_dims(&[4, 2, 3]),
dtype: DType::F32,
});
(in_id, p_id)
};
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, p_id, inputs).expect("realize");
assert!(Arc::ptr_eq(&result_arc, &in_arc));
assert_eq!(result_layout.shape().dims(), &[4, 2, 3]);
assert_eq!(result_layout.stride(), &[1, 12, 4]);
}
#[test]
fn pipelined_realize_cast_f32_to_f64() {
let storage = fuel_memory::from_slice_cpu(&[1.5_f32, -2.25, 100.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, c_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let c_id = g.push(Node {
op: Op::Cast(DType::F64),
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
g.set_target_backend(c_id, BackendId::Cpu);
(in_id, c_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[3]);
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F64);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[1.5_f64, -2.25, 100.0]);
}
#[test]
fn pipelined_realize_cast_round_trip_via_bf16() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, -3.0, 0.5]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, c1_id, c2_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let c1_id = g.push(Node {
op: Op::Cast(DType::BF16),
inputs: vec![in_id],
shape: Shape::from_dims(&[4]),
dtype: DType::BF16,
});
let c2_id = g.push(Node {
op: Op::Cast(DType::F32),
inputs: vec![c1_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(c1_id, BackendId::Cpu);
g.set_target_backend(c2_id, BackendId::Cpu);
(in_id, c1_id, c2_id)
};
let _ = c1_id;
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c2_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F32);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[1.0_f32, 2.0, -3.0, 0.5]);
}
#[test]
fn pipelined_realize_slice_is_metadata_only() {
let storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, s_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let s_id = g.push(Node {
op: Op::Slice {
dim: 0,
start: 1,
len: 3,
},
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
(in_id, s_id)
};
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, s_id, inputs).expect("realize");
assert!(Arc::ptr_eq(&result_arc, &in_arc));
assert_eq!(result_layout.shape().dims(), &[3]);
assert_eq!(result_layout.start_offset(), 1);
assert_eq!(result_layout.stride(), &[1]);
}
#[test]
fn pipelined_realize_slice_then_sum_all() {
let storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, s_id, sum_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let s_id = g.push(Node {
op: Op::Slice {
dim: 0,
start: 1,
len: 3,
},
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let sum_id = g.push(Node {
op: Op::SumAll,
inputs: vec![s_id],
shape: Shape::from_dims(&[]),
dtype: DType::F32,
});
g.set_target_backend(sum_id, BackendId::Cpu);
(in_id, s_id, sum_id)
};
let _ = s_id;
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[90.0]);
}
#[test]
fn pipelined_realize_flip_then_relu_materializes_in_reverse() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, _f_id, r_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let f_id = g.push(Node {
op: Op::Flip { dim: 0 },
inputs: vec![in_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let r_id = g.push(Node {
op: Op::Relu,
inputs: vec![f_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(r_id, BackendId::Cpu);
(in_id, f_id, r_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[4.0, 3.0, 2.0, 1.0]);
assert_eq!(result_layout.shape().dims(), &[4]);
assert!(result_layout.is_contiguous(), "relu output is contiguous");
}
#[test]
fn pipelined_realize_broadcast_is_metadata_only() {
let storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, b_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let b_id = g.push(Node {
op: Op::BroadcastTo(Shape::from_dims(&[4, 3])),
inputs: vec![in_id],
shape: Shape::from_dims(&[4, 3]),
dtype: DType::F32,
});
(in_id, b_id)
};
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, b_id, inputs).expect("realize");
assert!(Arc::ptr_eq(&result_arc, &in_arc));
assert_eq!(result_layout.shape().dims(), &[4, 3]);
assert_eq!(result_layout.stride(), &[0, 1]);
}
#[test]
fn pipelined_realize_matmul_2x3_times_3x2() {
let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let rhs_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, mm_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(lhs, rhs, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[2, 2]);
assert!(result_layout.is_contiguous());
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[58.0, 64.0, 139.0, 154.0]);
}
#[test]
fn pipelined_realize_matmul_batched_2x_2x2_times_2x2() {
let lhs_storage = fuel_memory::from_slice_cpu(&[
1.0_f32, 2.0, 3.0, 4.0, 1.0, 0.0, 0.0, 1.0, ]);
let rhs_storage = fuel_memory::from_slice_cpu(&[
5.0_f32, 6.0, 7.0, 8.0, 10.0, 20.0, 30.0, 40.0, ]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, mm_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2, 2]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2, 2]),
dtype: DType::F32,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[2, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(lhs, rhs, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[2, 2, 2]);
assert!(result_layout.is_contiguous());
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[19.0, 22.0, 43.0, 50.0, 10.0, 20.0, 30.0, 40.0]
);
}
#[test]
fn pipelined_realize_add_f64() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0]);
let rhs = fuel_memory::from_slice_cpu(&[10.0_f64, 20.0, 30.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, op_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
let op = g.push(Node {
op: Op::Add,
inputs: vec![l, r],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
g.set_target_backend(op, BackendId::Cpu);
(l, r, op)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F64);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[11.0_f64, 22.0, 33.0]);
}
#[test]
fn pipelined_realize_eq_f32_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, f32::NAN, 0.0]);
let rhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, eq_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let eq = g.push(Node {
op: Op::Equal,
inputs: vec![l, r],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
g.set_target_backend(eq, BackendId::Cpu);
(l, r, eq)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, eq_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[1, 0, 1, 0, 1]);
}
#[test]
fn pipelined_realize_ne_f32_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, f32::NAN, 0.0]);
let rhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, ne_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let ne = g.push(Node {
op: Op::Ne,
inputs: vec![l, r],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
g.set_target_backend(ne, BackendId::Cpu);
(l, r, ne)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, ne_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[0, 1, 0, 1, 0]);
}
#[test]
fn pipelined_realize_lt_f32_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -1.0]);
let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 3.0, 0.0, 0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, lt_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let lt = g.push(Node {
op: Op::Lt,
inputs: vec![l, r],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
g.set_target_backend(lt, BackendId::Cpu);
(l, r, lt)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, lt_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[1, 0, 0, 0, 1]);
}
#[test]
fn pipelined_realize_le_f32_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0, f32::NAN, -1.0]);
let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 2.0, 0.0, 0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, le_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let le = g.push(Node {
op: Op::Le,
inputs: vec![l, r],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
g.set_target_backend(le, BackendId::Cpu);
(l, r, le)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, le_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[1, 1, 0, 0, 1]);
}
#[test]
fn pipelined_realize_gt_f32_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[3.0_f32, 5.0, 2.0, f32::NAN, 1.0]);
let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 3.0, 0.0, 0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, gt_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let gt = g.push(Node {
op: Op::Gt,
inputs: vec![l, r],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
g.set_target_backend(gt, BackendId::Cpu);
(l, r, gt)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, gt_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[1, 0, 0, 0, 1]);
}
#[test]
fn pipelined_realize_ge_f32_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[3.0_f32, 5.0, 2.0, f32::NAN, 0.0]);
let rhs = fuel_memory::from_slice_cpu(&[2.0_f32, 5.0, 3.0, 0.0, 0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, ge_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let ge = g.push(Node {
op: Op::Ge,
inputs: vec![l, r],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
g.set_target_backend(ge, BackendId::Cpu);
(l, r, ge)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, ge_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[1, 1, 0, 0, 1]);
}
#[test]
fn pipelined_realize_eq_f64_to_bool_mask() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0]);
let rhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, eq_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
let eq = g.push(Node {
op: Op::Equal,
inputs: vec![l, r],
shape: Shape::from_dims(&[3]),
dtype: DType::Bool,
});
g.set_target_backend(eq, BackendId::Cpu);
(l, r, eq)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, eq_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::Bool);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let mask: &[u8] = c.as_slice().expect("u8 view");
assert_eq!(mask, &[1, 1, 0]);
}
#[test]
fn pipelined_realize_where_f32_picks_per_slot_from_bool_mask() {
let mut cond_storage = fuel_memory::from_slice_cpu(&[1u8, 0, 1, 0, 1]);
cond_storage.dtype = DType::Bool; let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0]);
let b_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (cond_id, a_id, b_id, where_id) = {
let mut g = graph.write().unwrap();
let cond = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::Bool,
});
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Where,
inputs: vec![cond, a, b],
shape: Shape::from_dims(&[5]),
dtype: DType::F32,
});
g.set_target_backend(w, BackendId::Cpu);
(cond, a, b, w)
};
let mut inputs = StorageCache::new();
inputs.insert(cond_id, Arc::new(RwLock::new(cond_storage)));
inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, where_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F32);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().expect("f32 view");
assert_eq!(out, &[1.0, 20.0, 3.0, 40.0, 5.0]);
}
#[test]
fn pipelined_realize_eq_then_where_full_chain() {
let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
let b_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 3.0]);
let pick_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let fb_storage = fuel_memory::from_slice_cpu(&[99.0_f32, 99.0, 99.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, pick_id, fb_id, where_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let pick = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let fb = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let eq = g.push(Node {
op: Op::Equal,
inputs: vec![a, b],
shape: Shape::from_dims(&[3]),
dtype: DType::Bool,
});
let w = g.push(Node {
op: Op::Where,
inputs: vec![eq, pick, fb],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(eq, BackendId::Cpu);
g.set_target_backend(w, BackendId::Cpu);
(a, b, pick, fb, w)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
inputs.insert(pick_id, Arc::new(RwLock::new(pick_storage)));
inputs.insert(fb_id, Arc::new(RwLock::new(fb_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, where_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().expect("f32 view");
assert_eq!(out, &[10.0, 99.0, 30.0]);
}
#[test]
fn pipelined_realize_qmatmul_q4_0_unit_weight_sums_activations() {
use fuel_graph::QuantType;
use half::f16;
let block_size = std::mem::size_of::<fuel_quantized::BlockQ4_0>();
let mut w_bytes = vec![0u8; 2 * block_size];
for block_idx in 0..2 {
let off = block_idx * block_size;
let d_bytes = f16::from_f32(1.0).to_le_bytes();
w_bytes[off..off + 2].copy_from_slice(&d_bytes);
for i in 0..16 {
w_bytes[off + 2 + i] = 0x99;
}
}
let w_storage = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_bytes(
&w_bytes,
)),
DType::U32,
);
let act_vec: Vec<f32> = (1..=32).map(|x| x as f32).collect();
let act_storage = fuel_memory::from_slice_cpu(&act_vec);
let graph = Arc::new(RwLock::new(Graph::new()));
let (act_id, w_id, mm_id) = {
let mut g = graph.write().unwrap();
let act = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 32]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[w_bytes.len() / 4]),
dtype: DType::U32,
});
let mm = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::QMATMUL,
fuel_graph::registry::FusedOpParams::QMatMul {
quant_type: QuantType::Q4_0,
k: 32,
n: 2,
},
),
inputs: vec![act, w],
shape: Shape::from_dims(&[1, 2]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(act, w, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(act_id, Arc::new(RwLock::new(act_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F32);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[f32] = c.as_slice().unwrap();
assert!((r[0] - 528.0).abs() < 0.5, "got {}, want 528", r[0]);
assert!((r[1] - 528.0).abs() < 0.5, "got {}, want 528", r[1]);
}
#[test]
fn pipelined_realize_qmatmul_q5_0_against_reference() {
use fuel_graph::QuantType;
use fuel_quantized::{BlockQ5_0, GgmlType};
let n = 2;
let k = 64; let w_f32 = vec![1.0_f32; n * k];
let blocks_per_row = k / BlockQ5_0::BLCK_SIZE;
let mut w_blocks = vec![BlockQ5_0::zeros(); n * blocks_per_row];
BlockQ5_0::from_float(&w_f32, &mut w_blocks);
let bytes_per_block = std::mem::size_of::<BlockQ5_0>();
let w_bytes: Vec<u8> = unsafe {
std::slice::from_raw_parts(
w_blocks.as_ptr() as *const u8,
w_blocks.len() * bytes_per_block,
)
}
.to_vec();
let w_storage = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_bytes(
&w_bytes,
)),
DType::U32,
);
let act_vec: Vec<f32> = (1..=k).map(|x| x as f32).collect();
let act_storage = fuel_memory::from_slice_cpu(&act_vec);
let mut ref_out = vec![0.0_f32; n];
fuel_quantized::matmul::<BlockQ5_0>((1, k, n), &act_vec, &w_blocks, &mut ref_out)
.expect("ref matmul");
let graph = Arc::new(RwLock::new(Graph::new()));
let (act_id, w_id, mm_id) = {
let mut g = graph.write().unwrap();
let act = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, k]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[w_bytes.len() / 4]),
dtype: DType::U32,
});
let mm = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::QMATMUL,
fuel_graph::registry::FusedOpParams::QMatMul {
quant_type: QuantType::Q5_0,
k,
n,
},
),
inputs: vec![act, w],
shape: Shape::from_dims(&[1, n]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(act, w, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(act_id, Arc::new(RwLock::new(act_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[f32] = c.as_slice().unwrap();
assert_eq!(
r,
ref_out.as_slice(),
"pipelined Q5_0 differs from reference: got {r:?}, want {ref_out:?}"
);
}
#[test]
fn pipelined_realize_qmatmul_q6k_against_reference() {
use fuel_graph::QuantType;
use fuel_quantized::{BlockQ6K, GgmlType};
let n = 2;
let k = 256; let w_f32 = vec![1.0_f32; n * k];
let blocks_per_row = k / BlockQ6K::BLCK_SIZE;
let mut w_blocks = vec![BlockQ6K::zeros(); n * blocks_per_row];
BlockQ6K::from_float(&w_f32, &mut w_blocks);
let bytes_per_block = std::mem::size_of::<BlockQ6K>();
let w_bytes: Vec<u8> = unsafe {
std::slice::from_raw_parts(
w_blocks.as_ptr() as *const u8,
w_blocks.len() * bytes_per_block,
)
}
.to_vec();
let w_storage = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_bytes(
&w_bytes,
)),
DType::U32,
);
let act_vec: Vec<f32> = (1..=k).map(|x| x as f32 / 100.0).collect();
let act_storage = fuel_memory::from_slice_cpu(&act_vec);
let mut ref_out = vec![0.0_f32; n];
fuel_quantized::matmul::<BlockQ6K>((1, k, n), &act_vec, &w_blocks, &mut ref_out)
.expect("ref matmul");
let graph = Arc::new(RwLock::new(Graph::new()));
let (act_id, w_id, mm_id) = {
let mut g = graph.write().unwrap();
let act = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, k]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[w_bytes.len() / 4]),
dtype: DType::U32,
});
let mm = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::QMATMUL,
fuel_graph::registry::FusedOpParams::QMatMul {
quant_type: QuantType::Q6K,
k,
n,
},
),
inputs: vec![act, w],
shape: Shape::from_dims(&[1, n]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(act, w, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(act_id, Arc::new(RwLock::new(act_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[f32] = c.as_slice().unwrap();
assert_eq!(
r,
ref_out.as_slice(),
"pipelined Q6K differs from reference: got {r:?}, want {ref_out:?}"
);
}
#[test]
fn pipelined_realize_rms_norm_bf16() {
let v: Vec<half::bf16> = [3.0_f32, 4.0]
.iter()
.map(|&x| half::bf16::from_f32(x))
.collect();
let storage = fuel_memory::from_slice_cpu(&v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2]),
dtype: DType::BF16,
});
let op_id = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM,
fuel_graph::registry::FusedOpParams::RmsNormLastDim { eps: 0.0 },
),
inputs: vec![in_id],
shape: Shape::from_dims(&[1, 2]),
dtype: DType::BF16,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::BF16);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[half::bf16] = c.as_slice().unwrap();
let rms = (12.5_f32).sqrt();
assert!((r[0].to_f32() - 3.0 / rms).abs() < 0.05);
assert!((r[1].to_f32() - 4.0 / rms).abs() < 0.05);
}
#[test]
fn pipelined_realize_matmul_bf16_identity() {
let lhs_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0]
.iter()
.map(|&x| half::bf16::from_f32(x))
.collect();
let rhs_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0]
.iter()
.map(|&x| half::bf16::from_f32(x))
.collect();
let lhs = fuel_memory::from_slice_cpu(&lhs_v);
let rhs = fuel_memory::from_slice_cpu(&rhs_v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, mm_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::BF16,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::BF16,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![l, r],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::BF16,
});
g.set_target_backend(mm, BackendId::Cpu);
(l, r, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::BF16);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[half::bf16] = c.as_slice().unwrap();
let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
assert_eq!(result_f32, vec![1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn pipelined_realize_sum_dim_f16() {
let v: Vec<half::f16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
.iter()
.map(|&x| half::f16::from_f32(x))
.collect();
let storage = fuel_memory::from_slice_cpu(&v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sum_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F16,
});
let sum_id = g.push(Node {
op: Op::SumDim(1),
inputs: vec![in_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F16,
});
g.set_target_backend(sum_id, BackendId::Cpu);
(in_id, sum_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F16);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[half::f16] = c.as_slice().unwrap();
let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
assert_eq!(result_f32, vec![6.0, 15.0]);
}
#[test]
fn pipelined_realize_add_bf16() {
let lhs_vec: Vec<half::bf16> = [1.0_f32, 2.0, 3.0]
.iter()
.map(|&x| half::bf16::from_f32(x))
.collect();
let rhs_vec: Vec<half::bf16> = [10.0_f32, 20.0, 30.0]
.iter()
.map(|&x| half::bf16::from_f32(x))
.collect();
let lhs = fuel_memory::from_slice_cpu(&lhs_vec);
let rhs = fuel_memory::from_slice_cpu(&rhs_vec);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, op_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::BF16,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::BF16,
});
let op = g.push(Node {
op: Op::Add,
inputs: vec![l, r],
shape: Shape::from_dims(&[3]),
dtype: DType::BF16,
});
g.set_target_backend(op, BackendId::Cpu);
(l, r, op)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::BF16);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[half::bf16] = c.as_slice().unwrap();
let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
assert_eq!(result_f32, vec![11.0, 22.0, 33.0]);
}
#[test]
fn pipelined_realize_sqr_then_sqrt_f16() {
let v: Vec<half::f16> = [1.0_f32, 4.0, 9.0]
.iter()
.map(|&x| half::f16::from_f32(x))
.collect();
let storage = fuel_memory::from_slice_cpu(&v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sqrt_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F16,
});
let sqr_id = g.push(Node {
op: Op::Sqr,
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F16,
});
let sqrt_id = g.push(Node {
op: Op::Sqrt,
inputs: vec![sqr_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F16,
});
g.set_target_backend(sqr_id, BackendId::Cpu);
g.set_target_backend(sqrt_id, BackendId::Cpu);
(in_id, sqrt_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sqrt_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F16);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r: &[half::f16] = c.as_slice().unwrap();
let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
for (got, want) in result_f32.iter().zip(&[1.0_f32, 4.0, 9.0]) {
assert!((got - want).abs() < 0.05, "got {got}, want {want}");
}
}
#[test]
fn pipelined_realize_sum_dim_f64() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sum_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F64,
});
let sum_id = g.push(Node {
op: Op::SumDim(1),
inputs: vec![in_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F64,
});
g.set_target_backend(sum_id, BackendId::Cpu);
(in_id, sum_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F64);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[6.0_f64, 15.0]);
}
#[test]
fn pipelined_realize_matmul_2x3_times_3x2_f64() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
let rhs = fuel_memory::from_slice_cpu(&[7.0_f64, 8.0, 9.0, 10.0, 11.0, 12.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (l_id, r_id, mm_id) = {
let mut g = graph.write().unwrap();
let l = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F64,
});
let r = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F64,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![l, r],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F64,
});
g.set_target_backend(mm, BackendId::Cpu);
(l, r, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(l_id, Arc::new(RwLock::new(lhs)));
inputs.insert(r_id, Arc::new(RwLock::new(rhs)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F64);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f64>().unwrap(),
&[58.0_f64, 64.0, 139.0, 154.0]
);
}
#[test]
fn pipelined_realize_sqr_then_sqrt_f64() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f64, 4.0, 9.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sqrt_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
let sqr_id = g.push(Node {
op: Op::Sqr,
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
let sqrt_id = g.push(Node {
op: Op::Sqrt,
inputs: vec![sqr_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
g.set_target_backend(sqr_id, BackendId::Cpu);
g.set_target_backend(sqrt_id, BackendId::Cpu);
(in_id, sqrt_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sqrt_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::F64);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[1.0_f64, 4.0, 9.0]);
}
#[test]
fn pipelined_realize_argmax_dim() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, 2.0, 9.0, 0.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::ArgMaxDim(1),
inputs: vec![in_id],
shape: Shape::from_dims(&[2]),
dtype: DType::U32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
assert_eq!(guard.dtype, DType::U32);
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<u32>().unwrap(), &[1u32, 0]);
}
#[test]
fn pipelined_realize_index_add_simple() {
let base = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let indices = fuel_memory::from_slice_cpu(&[0u32, 0]);
let src = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (b_id, i_id, s_id, op_id) = {
let mut g = graph.write().unwrap();
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let i = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::U32,
});
let s = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::IndexAdd { dim: 0 },
inputs: vec![b, i, s],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(b, i, s, op)
};
let mut inputs = StorageCache::new();
inputs.insert(b_id, Arc::new(RwLock::new(base)));
inputs.insert(i_id, Arc::new(RwLock::new(indices)));
inputs.insert(s_id, Arc::new(RwLock::new(src)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[13.0, 20.0, 30.0]);
}
#[test]
fn pipelined_realize_scatter_add_outer_dim() {
let base = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
let indices = fuel_memory::from_slice_cpu(&[0u32, 1, 2, 0]);
let src = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (b_id, i_id, s_id, op_id) = {
let mut g = graph.write().unwrap();
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let i = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::U32,
});
let s = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::ScatterAdd { dim: 0 },
inputs: vec![b, i, s],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(b, i, s, op)
};
let mut inputs = StorageCache::new();
inputs.insert(b_id, Arc::new(RwLock::new(base)));
inputs.insert(i_id, Arc::new(RwLock::new(indices)));
inputs.insert(s_id, Arc::new(RwLock::new(src)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[1.0, 4.0, 0.0, 2.0, 3.0, 0.0]
);
}
#[test]
fn pipelined_realize_rope_pi_over_two() {
let x = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let cos = fuel_memory::from_slice_cpu(&[0.0_f32, 0.0, 0.0, 0.0]);
let sin = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, cos_id, sin_id, r_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 4]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 4]),
dtype: DType::F32,
});
let s = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 4]),
dtype: DType::F32,
});
let r = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::ROPE,
fuel_graph::registry::FusedOpParams::Rope,
),
inputs: vec![x, c, s],
shape: Shape::from_dims(&[1, 1, 4]),
dtype: DType::F32,
});
g.set_target_backend(r, BackendId::Cpu);
(x, c, s, r)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x)));
inputs.insert(cos_id, Arc::new(RwLock::new(cos)));
inputs.insert(sin_id, Arc::new(RwLock::new(sin)));
let (result_arc, _) = PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[-3.0, -4.0, 1.0, 2.0]);
}
#[test]
fn pipelined_realize_gather_inner_dim() {
let source =
fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]);
let indices = fuel_memory::from_slice_cpu(&[0u32, 2, 1, 3, 0, 0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (src_id, idx_id, g_id) = {
let mut g = graph.write().unwrap();
let s = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
let i = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::U32,
});
let g_id = g.push(Node {
op: Op::Gather { dim: 1 },
inputs: vec![s, i],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
g.set_target_backend(g_id, BackendId::Cpu);
(s, i, g_id)
};
let mut inputs = StorageCache::new();
inputs.insert(src_id, Arc::new(RwLock::new(source)));
inputs.insert(idx_id, Arc::new(RwLock::new(indices)));
let (result_arc, _) = PipelinedExecutor::realize(graph, g_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[10.0, 30.0, 20.0, 80.0, 50.0, 50.0]
);
}
#[test]
fn pipelined_realize_index_select_embedding_lookup() {
let table = fuel_memory::from_slice_cpu(&[
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, ]);
let indices = fuel_memory::from_slice_cpu(&[2u32, 0, 2]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (table_id, idx_id, sel_id) = {
let mut g = graph.write().unwrap();
let t = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4, 3]),
dtype: DType::F32,
});
let i = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::U32,
});
let s = g.push(Node {
op: Op::IndexSelect { dim: 0 },
inputs: vec![t, i],
shape: Shape::from_dims(&[3, 3]),
dtype: DType::F32,
});
g.set_target_backend(s, BackendId::Cpu);
(t, i, s)
};
let mut inputs = StorageCache::new();
inputs.insert(table_id, Arc::new(RwLock::new(table)));
inputs.insert(idx_id, Arc::new(RwLock::new(indices)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, sel_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[3, 3]);
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[30.0, 31.0, 32.0, 10.0, 11.0, 12.0, 30.0, 31.0, 32.0]
);
}
#[test]
fn pipelined_realize_rms_norm_last_dim() {
let storage = fuel_memory::from_slice_cpu(&[
3.0_f32, 4.0, 6.0, 8.0, ]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::RMS_NORM_LAST_DIM,
fuel_graph::registry::FusedOpParams::RmsNormLastDim { eps: 0.0 },
),
inputs: vec![in_id],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let result: &[f32] = c.as_slice().unwrap();
let rms0 = (12.5_f32).sqrt();
assert!((result[0] - 3.0 / rms0).abs() < 1e-6);
assert!((result[1] - 4.0 / rms0).abs() < 1e-6);
let rms1 = (50.0_f32).sqrt();
assert!((result[2] - 6.0 / rms1).abs() < 1e-6);
assert!((result[3] - 8.0 / rms1).abs() < 1e-6);
}
#[test]
fn pipelined_realize_layer_norm_last_dim() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 10.0, 20.0, 30.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::LAYER_NORM_LAST_DIM,
fuel_graph::registry::FusedOpParams::LayerNormLastDim { eps: 0.0 },
),
inputs: vec![in_id],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let result: &[f32] = c.as_slice().unwrap();
for row in 0..2 {
let off = row * 3;
let sum: f32 = result[off..off + 3].iter().sum();
let mean = sum / 3.0;
assert!(mean.abs() < 1e-6, "row {row} mean should be 0, got {mean}");
let var: f32 = result[off..off + 3]
.iter()
.map(|v| (v - mean).powi(2))
.sum::<f32>()
/ 3.0;
assert!(
(var - 1.0).abs() < 1e-6,
"row {row} var should be 1, got {var}"
);
}
}
#[test]
fn pipelined_realize_softmax_last_dim() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 100.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::SOFTMAX_LAST_DIM,
fuel_graph::registry::FusedOpParams::SoftmaxLastDim,
),
inputs: vec![in_id],
shape: Shape::from_dims(&[2, 4]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let result: &[f32] = c.as_slice().unwrap();
for v in &result[..4] {
assert!((v - 0.25).abs() < 1e-7);
}
for v in &result[4..7] {
assert!(
*v < 1e-30,
"row-1 leading positions should be near 0, got {v}"
);
}
assert!(
result[7] > 0.999,
"row-1 last position should dominate, got {}",
result[7]
);
let row0_sum: f32 = result[..4].iter().sum();
let row1_sum: f32 = result[4..].iter().sum();
assert!((row0_sum - 1.0).abs() < 1e-6);
assert!((row1_sum - 1.0).abs() < 1e-6);
}
#[test]
fn pipelined_realize_concat_inner_dim() {
let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let b = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, c_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Concat { dim: 1 },
inputs: vec![a, b],
shape: Shape::from_dims(&[2, 6]),
dtype: DType::F32,
});
g.set_target_backend(c, BackendId::Cpu);
(a, b, c)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
inputs.insert(b_id, Arc::new(RwLock::new(b)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[2, 6]);
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[
1.0, 2.0, 3.0, 7.0, 8.0, 9.0, 4.0, 5.0, 6.0, 10.0, 11.0, 12.0
]
);
}
#[test]
fn pipelined_realize_concat_three_inputs_outer() {
let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
let b = fuel_memory::from_slice_cpu(&[3.0_f32, 4.0]);
let c = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, c_id, cat_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let cat = g.push(Node {
op: Op::Concat { dim: 0 },
inputs: vec![a, b, c],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
g.set_target_backend(cat, BackendId::Cpu);
(a, b, c, cat)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
inputs.insert(b_id, Arc::new(RwLock::new(b)));
inputs.insert(c_id, Arc::new(RwLock::new(c)));
let (result_arc, _) = PipelinedExecutor::realize(graph, cat_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
);
}
#[test]
fn pipelined_realize_add_scalar() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::AddScalar(10.0),
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 12.0, 13.0]);
}
#[test]
fn pipelined_realize_clamp() {
let storage = fuel_memory::from_slice_cpu(&[-5.0_f32, 0.5, 100.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::Clamp {
min: -2.0,
max: 2.0,
},
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[-2.0, 0.5, 2.0]);
}
#[test]
fn pipelined_realize_maximum_elementwise() {
let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 5.0, -3.0]);
let rhs_storage = fuel_memory::from_slice_cpu(&[2.0_f32, 1.0, -1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, op_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::Maximum,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(lhs, rhs, op)
};
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[2.0, 5.0, -1.0]);
}
#[test]
fn pipelined_realize_conv2d_2x2_sum_kernel() {
let x_storage =
fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
let w_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV2D,
fuel_graph::registry::FusedOpParams::Conv2D {
stride: (1, 1),
padding: (0, 0),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[1, 1, 2, 2]);
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[12.0, 16.0, 24.0, 28.0]);
}
#[test]
fn pipelined_realize_conv2d_with_bias() {
let x_storage =
fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
let w_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
let bias_storage = fuel_memory::from_slice_cpu(&[100.0_f32]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, b_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV2D,
fuel_graph::registry::FusedOpParams::Conv2D {
stride: (1, 1),
padding: (0, 0),
groups: 1,
},
),
inputs: vec![x, w, b],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, b, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
inputs.insert(b_id, Arc::new(RwLock::new(bias_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[112.0, 116.0, 124.0, 128.0]);
}
#[test]
fn pipelined_realize_conv2d_f64() {
let x_storage =
fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
let w_storage = fuel_memory::from_slice_cpu(&[1.0_f64, 1.0, 1.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::F64,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F64,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV2D,
fuel_graph::registry::FusedOpParams::Conv2D {
stride: (1, 1),
padding: (0, 0),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F64,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[12.0, 16.0, 24.0, 28.0]);
}
#[test]
fn pipelined_realize_conv2d_bf16() {
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]
.iter()
.map(|v| half::bf16::from_f32(*v))
.collect();
let w_data: Vec<half::bf16> = [1.0_f32, 1.0, 1.0, 1.0]
.iter()
.map(|v| half::bf16::from_f32(*v))
.collect();
let x_storage = fuel_memory::from_slice_cpu(&x_data);
let w_storage = fuel_memory::from_slice_cpu(&w_data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::BF16,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::BF16,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV2D,
fuel_graph::registry::FusedOpParams::Conv2D {
stride: (1, 1),
padding: (0, 0),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::BF16,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::bf16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let want = [12.0_f32, 16.0, 24.0, 28.0];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_conv2d_f16() {
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]
.iter()
.map(|v| half::f16::from_f32(*v))
.collect();
let w_data: Vec<half::f16> = [1.0_f32, 1.0, 1.0, 1.0]
.iter()
.map(|v| half::f16::from_f32(*v))
.collect();
let x_storage = fuel_memory::from_slice_cpu(&x_data);
let w_storage = fuel_memory::from_slice_cpu(&w_data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::F16,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F16,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV2D,
fuel_graph::registry::FusedOpParams::Conv2D {
stride: (1, 1),
padding: (0, 0),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F16,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::f16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let want = [12.0_f32, 16.0, 24.0, 28.0];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.05, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_paged_attn_f32() {
let q = fuel_memory::from_slice_cpu(&[2.0_f32, 0.0]);
let k_cache = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let v_cache = fuel_memory::from_slice_cpu(&[10.0_f32, 0.0, 0.0, 10.0]);
let block_table_u32 = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
0_u32,
])),
DType::U32,
);
let context_lens_u32 = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
2_u32,
])),
DType::U32,
);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, kc_id, vc_id, bt_id, cl_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::F32,
});
let kc = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 1, 2]),
dtype: DType::F32,
});
let vc = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 1, 2]),
dtype: DType::F32,
});
let bt = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1]),
dtype: DType::U32,
});
let cl = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1]),
dtype: DType::U32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::PAGED_ATTN,
fuel_graph::registry::FusedOpParams::PagedAttn {
softmax_scale: 1.0,
block_size: 2,
softcap: None,
},
),
inputs: vec![q, kc, vc, bt, cl],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(q, kc, vc, bt, cl, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(kc_id, Arc::new(RwLock::new(k_cache)));
inputs.insert(vc_id, Arc::new(RwLock::new(v_cache)));
inputs.insert(bt_id, Arc::new(RwLock::new(block_table_u32)));
inputs.insert(cl_id, Arc::new(RwLock::new(context_lens_u32)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r = c.as_slice::<f32>().unwrap();
let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
assert!(
(r[0] - 10.0 * expected_p0).abs() < 1e-5,
"row[0]: got {} expected {}",
r[0],
10.0 * expected_p0
);
assert!(
(r[1] - 10.0 * expected_p1).abs() < 1e-5,
"row[1]: got {} expected {}",
r[1],
10.0 * expected_p1
);
}
#[test]
fn pipelined_realize_paged_attn_bf16() {
let q_v: Vec<half::bf16> = [2.0_f32, 0.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let kc_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let vc_v: Vec<half::bf16> = [10.0_f32, 0.0, 0.0, 10.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let q = fuel_memory::from_slice_cpu(&q_v);
let k_cache = fuel_memory::from_slice_cpu(&kc_v);
let v_cache = fuel_memory::from_slice_cpu(&vc_v);
let bt = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
0_u32,
])),
DType::U32,
);
let cl = fuel_memory::Storage::new(
fuel_memory::BackendStorage::Cpu(fuel_cpu_backend::CpuStorageBytes::from_slice(&[
2_u32,
])),
DType::U32,
);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, kc_id, vc_id, bt_id, cl_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::BF16,
});
let kc = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 1, 2]),
dtype: DType::BF16,
});
let vc = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 1, 2]),
dtype: DType::BF16,
});
let bt = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1]),
dtype: DType::U32,
});
let cl = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1]),
dtype: DType::U32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::PAGED_ATTN,
fuel_graph::registry::FusedOpParams::PagedAttn {
softmax_scale: 1.0,
block_size: 2,
softcap: None,
},
),
inputs: vec![q, kc, vc, bt, cl],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::BF16,
});
g.set_target_backend(op, BackendId::Cpu);
(q, kc, vc, bt, cl, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(kc_id, Arc::new(RwLock::new(k_cache)));
inputs.insert(vc_id, Arc::new(RwLock::new(v_cache)));
inputs.insert(bt_id, Arc::new(RwLock::new(bt)));
inputs.insert(cl_id, Arc::new(RwLock::new(cl)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::bf16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
let want = [10.0 * expected_p0, 10.0 * expected_p1];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_flash_attn_f32() {
let q = fuel_memory::from_slice_cpu(&[2.0_f32, 0.0]);
let k = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let v = fuel_memory::from_slice_cpu(&[10.0_f32, 0.0, 0.0, 10.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, k_id, v_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::F32,
});
let k = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let v = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FLASH_ATTN,
fuel_graph::registry::FusedOpParams::FlashAttn {
softmax_scale: 1.0,
causal: false,
window_size_left: None,
window_size_right: None,
softcap: None,
k_len: None,
},
),
inputs: vec![q, k, v],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(q, k, v, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(k_id, Arc::new(RwLock::new(k)));
inputs.insert(v_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r = c.as_slice::<f32>().unwrap();
let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
assert!(
(r[0] - 10.0 * expected_p0).abs() < 1e-5,
"row[0]: got {} expected {}",
r[0],
10.0 * expected_p0
);
assert!(
(r[1] - 10.0 * expected_p1).abs() < 1e-5,
"row[1]: got {} expected {}",
r[1],
10.0 * expected_p1
);
}
#[test]
fn pipelined_realize_flash_attn_causal_f32() {
let q = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let k = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let v = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0, 7.0, 8.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, k_id, v_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let k = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let v = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FLASH_ATTN,
fuel_graph::registry::FusedOpParams::FlashAttn {
softmax_scale: 1.0,
causal: true,
window_size_left: None,
window_size_right: None,
softcap: None,
k_len: None,
},
),
inputs: vec![q, k, v],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(q, k, v, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(k_id, Arc::new(RwLock::new(k)));
inputs.insert(v_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r = c.as_slice::<f32>().unwrap();
assert!((r[0] - 5.0).abs() < 1e-5, "got {}", r[0]);
assert!((r[1] - 6.0).abs() < 1e-5, "got {}", r[1]);
let denom = (1.0_f32).exp() + 1.0;
let expected_a = 5.0 / denom + 7.0 * (1.0_f32).exp() / denom;
let expected_b = 6.0 / denom + 8.0 * (1.0_f32).exp() / denom;
assert!(
(r[2] - expected_a).abs() < 1e-5,
"row1[0]: got {} expected {}",
r[2],
expected_a
);
assert!(
(r[3] - expected_b).abs() < 1e-5,
"row1[1]: got {} expected {}",
r[3],
expected_b
);
}
#[test]
fn pipelined_realize_flash_attn_dynamic_k_len() {
use fuel_ir::{DynScalar, SymId};
let q = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let k = fuel_memory::from_slice_cpu(&[
1.0_f32, 0.0, 0.0, 1.0, 1.0, 1.0, 100.0,
100.0, ]);
let v = fuel_memory::from_slice_cpu(&[
5.0_f32, 6.0, 7.0, 8.0, 9.0, 10.0, 999.0,
999.0, ]);
let sym = SymId(0);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, k_id, v_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let k = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 4, 2]),
dtype: DType::F32,
});
let v = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 4, 2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FLASH_ATTN,
fuel_graph::registry::FusedOpParams::FlashAttn {
softmax_scale: 1.0,
causal: true,
window_size_left: None,
window_size_right: None,
softcap: None,
k_len: Some(DynScalar::Sym(sym)),
},
),
inputs: vec![q, k, v],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(q, k, v, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(k_id, Arc::new(RwLock::new(k)));
inputs.insert(v_id, Arc::new(RwLock::new(v)));
let mut env = SymEnv::new();
env.bind(sym, 3).unwrap();
let (result_arc, _) = PipelinedExecutor::realize_with_env(graph, op_id, inputs, env)
.expect("realize_with_env");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let r = c.as_slice::<f32>().unwrap();
let e = 1.0_f32.exp();
let d0 = e + 1.0;
let e0x = (5.0 * e + 7.0) / d0;
let e0y = (6.0 * e + 8.0) / d0;
let d1 = 1.0 + 2.0 * e;
let e1x = (5.0 + 7.0 * e + 9.0 * e) / d1;
let e1y = (6.0 + 8.0 * e + 10.0 * e) / d1;
for (got, want) in r.iter().zip([e0x, e0y, e1x, e1y].iter()) {
assert!(
(got - want).abs() < 1e-4,
"flash dyn k_len: got {got}, want {want} (full output {r:?}); \
a poison-contaminated value means the capacity rows past k_len leaked in",
);
}
}
#[test]
fn pipelined_realize_flash_attn_dynamic_k_len_unbound_errors() {
use fuel_ir::{DynScalar, SymId};
let q = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let k = fuel_memory::from_slice_cpu(&[0.0_f32; 16]);
let v = fuel_memory::from_slice_cpu(&[0.0_f32; 16]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, k_id, v_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let k = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 4, 2]),
dtype: DType::F32,
});
let v = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 4, 2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FLASH_ATTN,
fuel_graph::registry::FusedOpParams::FlashAttn {
softmax_scale: 1.0,
causal: true,
window_size_left: None,
window_size_right: None,
softcap: None,
k_len: Some(DynScalar::Sym(SymId(0))),
},
),
inputs: vec![q, k, v],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(q, k, v, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(k_id, Arc::new(RwLock::new(k)));
inputs.insert(v_id, Arc::new(RwLock::new(v)));
let result = PipelinedExecutor::realize_with_env(graph, op_id, inputs, SymEnv::new());
assert!(
result.is_err(),
"unbound flash k_len symbol must surface a typed error"
);
}
#[test]
fn pipelined_realize_flash_attn_bf16() {
let q_v: Vec<half::bf16> = [2.0_f32, 0.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let k_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let v_v: Vec<half::bf16> = [10.0_f32, 0.0, 0.0, 10.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let q = fuel_memory::from_slice_cpu(&q_v);
let k = fuel_memory::from_slice_cpu(&k_v);
let v = fuel_memory::from_slice_cpu(&v_v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (q_id, k_id, v_id, op_id) = {
let mut g = graph.write().unwrap();
let q = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::BF16,
});
let k = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::BF16,
});
let v = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::BF16,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FLASH_ATTN,
fuel_graph::registry::FusedOpParams::FlashAttn {
softmax_scale: 1.0,
causal: false,
window_size_left: None,
window_size_right: None,
softcap: None,
k_len: None,
},
),
inputs: vec![q, k, v],
shape: Shape::from_dims(&[1, 1, 1, 2]),
dtype: DType::BF16,
});
g.set_target_backend(op, BackendId::Cpu);
(q, k, v, op)
};
let mut inputs = StorageCache::new();
inputs.insert(q_id, Arc::new(RwLock::new(q)));
inputs.insert(k_id, Arc::new(RwLock::new(k)));
inputs.insert(v_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::bf16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let expected_p0 = 2.0_f32.exp() / (2.0_f32.exp() + 1.0);
let expected_p1 = 1.0_f32 / (2.0_f32.exp() + 1.0);
let want = [10.0 * expected_p0, 10.0 * expected_p1];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_fused_linear_f32() {
let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let b = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0, 1.0, 1.0]);
let bias = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, bias_id, op_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 3]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 3, 2]),
dtype: DType::F32,
});
let bias = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FUSED_LINEAR,
fuel_graph::registry::FusedOpParams::FusedLinear,
),
inputs: vec![a, b, bias],
shape: Shape::from_dims(&[1, 2, 2]),
dtype: DType::F32,
});
g.set_target_backend(op, BackendId::Cpu);
(a, b, bias, op)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
inputs.insert(b_id, Arc::new(RwLock::new(b)));
inputs.insert(bias_id, Arc::new(RwLock::new(bias)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[14.0, 25.0, 20.0, 31.0]);
}
#[test]
fn pipelined_realize_fused_linear_f64() {
let a = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
let b = fuel_memory::from_slice_cpu(&[1.0_f64, 0.0, 0.0, 1.0, 1.0, 1.0]);
let bias = fuel_memory::from_slice_cpu(&[10.0_f64, 20.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, bias_id, op_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 3]),
dtype: DType::F64,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 3, 2]),
dtype: DType::F64,
});
let bias = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F64,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FUSED_LINEAR,
fuel_graph::registry::FusedOpParams::FusedLinear,
),
inputs: vec![a, b, bias],
shape: Shape::from_dims(&[1, 2, 2]),
dtype: DType::F64,
});
g.set_target_backend(op, BackendId::Cpu);
(a, b, bias, op)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
inputs.insert(b_id, Arc::new(RwLock::new(b)));
inputs.insert(bias_id, Arc::new(RwLock::new(bias)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[14.0, 25.0, 20.0, 31.0]);
}
#[test]
fn pipelined_realize_fused_linear_bf16() {
let a_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let b_v: Vec<half::bf16> = [1.0_f32, 0.0, 0.0, 1.0, 1.0, 1.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let bias_v: Vec<half::bf16> = [10.0_f32, 20.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let a = fuel_memory::from_slice_cpu(&a_v);
let b = fuel_memory::from_slice_cpu(&b_v);
let bias = fuel_memory::from_slice_cpu(&bias_v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, bias_id, op_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2, 3]),
dtype: DType::BF16,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 3, 2]),
dtype: DType::BF16,
});
let bias = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::BF16,
});
let op = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::FUSED_LINEAR,
fuel_graph::registry::FusedOpParams::FusedLinear,
),
inputs: vec![a, b, bias],
shape: Shape::from_dims(&[1, 2, 2]),
dtype: DType::BF16,
});
g.set_target_backend(op, BackendId::Cpu);
(a, b, bias, op)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
inputs.insert(b_id, Arc::new(RwLock::new(b)));
inputs.insert(bias_id, Arc::new(RwLock::new(bias)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::bf16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let want = [14.0_f32, 25.0, 20.0, 31.0];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_reduce_sum_to_f32_drops_leading_axis() {
let v = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::ReduceSumTo(Shape::from_dims(&[3])),
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[5.0, 7.0, 9.0]);
}
#[test]
fn pipelined_realize_reduce_sum_to_f32_keepdim_middle() {
let mut v: Vec<f32> = (1..=24).map(|x| x as f32).collect();
let _ = &mut v;
let s = fuel_memory::from_slice_cpu(&v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3, 4]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::ReduceSumTo(Shape::from_dims(&[2, 1, 4])),
inputs: vec![in_id],
shape: Shape::from_dims(&[2, 1, 4]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(s)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[15.0, 18.0, 21.0, 24.0, 51.0, 54.0, 57.0, 60.0],
);
}
#[test]
fn pipelined_realize_reduce_sum_to_f64() {
let v = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F64,
});
let op_id = g.push(Node {
op: Op::ReduceSumTo(Shape::from_dims(&[3])),
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F64,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f64>().unwrap(), &[5.0, 7.0, 9.0]);
}
#[test]
fn pipelined_realize_reduce_max_to_f32_drops_leading_axis() {
let v = fuel_memory::from_slice_cpu(&[1.0_f32, 7.0, 3.0, 4.0, 2.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::ReduceMaxTo(Shape::from_dims(&[3])),
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[4.0, 7.0, 6.0]);
}
#[test]
fn pipelined_realize_reduce_max_to_f32_keepdim_trailing() {
let v = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let op_id = g.push(Node {
op: Op::ReduceMaxTo(Shape::from_dims(&[2, 1])),
inputs: vec![in_id],
shape: Shape::from_dims(&[2, 1]),
dtype: DType::F32,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(v)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[3.0, 6.0]);
}
#[test]
fn pipelined_realize_reduce_sum_to_bf16() {
let v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
.iter()
.map(|x| half::bf16::from_f32(*x))
.collect();
let s = fuel_memory::from_slice_cpu(&v);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, op_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::BF16,
});
let op_id = g.push(Node {
op: Op::ReduceSumTo(Shape::from_dims(&[3])),
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::BF16,
});
g.set_target_backend(op_id, BackendId::Cpu);
(in_id, op_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(s)));
let (result_arc, _) = PipelinedExecutor::realize(graph, op_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::bf16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let want = [5.0_f32, 7.0, 9.0];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_conv_transpose2d_f32() {
let x_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let w_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 1.0, 1.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D,
fuel_graph::registry::FusedOpParams::ConvTranspose2D {
stride: (1, 1),
padding: (0, 0),
output_padding: (0, 0),
dilation: (1, 1),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::F32,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0],
);
}
#[test]
fn pipelined_realize_conv_transpose2d_f64() {
let x_storage = fuel_memory::from_slice_cpu(&[1.0_f64, 2.0, 3.0, 4.0]);
let w_storage = fuel_memory::from_slice_cpu(&[1.0_f64, 1.0, 1.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F64,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::F64,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D,
fuel_graph::registry::FusedOpParams::ConvTranspose2D {
stride: (1, 1),
padding: (0, 0),
output_padding: (0, 0),
dilation: (1, 1),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::F64,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f64>().unwrap(),
&[1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0],
);
}
#[test]
fn pipelined_realize_conv_transpose2d_bf16() {
let x: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0]
.iter()
.map(|v| half::bf16::from_f32(*v))
.collect();
let w: Vec<half::bf16> = [1.0_f32, 1.0, 1.0, 1.0]
.iter()
.map(|v| half::bf16::from_f32(*v))
.collect();
let x_storage = fuel_memory::from_slice_cpu(&x);
let w_storage = fuel_memory::from_slice_cpu(&w);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, w_id, c_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::BF16,
});
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 1, 2, 2]),
dtype: DType::BF16,
});
let c = g.push(Node {
op: Op::Fused(
fuel_graph::registry::FusedOps::CONV_TRANSPOSE2D,
fuel_graph::registry::FusedOpParams::ConvTranspose2D {
stride: (1, 1),
padding: (0, 0),
output_padding: (0, 0),
dilation: (1, 1),
groups: 1,
},
),
inputs: vec![x, w],
shape: Shape::from_dims(&[1, 1, 3, 3]),
dtype: DType::BF16,
});
g.set_target_backend(c, BackendId::Cpu);
(x, w, c)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(w_id, Arc::new(RwLock::new(w_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, c_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: Vec<f32> = c
.as_slice::<half::bf16>()
.unwrap()
.iter()
.map(|v| v.to_f32())
.collect();
let want = [1.0_f32, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0];
for (g, w) in got.iter().zip(want.iter()) {
assert!((g - w).abs() < 0.5, "got {got:?} want {want:?}");
}
}
#[test]
fn pipelined_realize_matmul_gqa() {
let lhs_storage =
fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
let rhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 0.0, 0.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, mm_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4, 1, 2]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2, 1]),
dtype: DType::F32,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[4, 1, 1]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(lhs, rhs, mm)
};
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[4, 1, 1]);
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[1.0, 3.0, 6.0, 8.0]);
}
#[test]
fn pipelined_realize_matmul_with_transposed_rhs() {
let lhs_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let rhs_storage = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0, 7.0, 8.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, t_id, mm_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
let t = g.push(Node {
op: Op::Transpose,
inputs: vec![rhs],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
let mm = g.push(Node {
op: Op::MatMul,
inputs: vec![lhs, t],
shape: Shape::from_dims(&[2, 2]),
dtype: DType::F32,
});
g.set_target_backend(mm, BackendId::Cpu);
(lhs, rhs, t, mm)
};
let _ = t_id;
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs_storage)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mm_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[17.0, 23.0, 39.0, 53.0]);
}
#[test]
fn pipelined_realize_reshape_zero_copy_when_contiguous() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, r_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let r_id = g.push(Node {
op: Op::Reshape(Shape::from_dims(&[3, 2])),
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
(in_id, r_id)
};
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
assert!(
Arc::ptr_eq(&result_arc, &in_arc),
"contiguous reshape must zero-copy"
);
assert_eq!(result_layout.shape().dims(), &[3, 2]);
assert!(result_layout.is_contiguous());
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
);
}
#[test]
fn pipelined_realize_reshape_materializes_when_strided() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, t_id, r_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let t_id = g.push(Node {
op: Op::Transpose,
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let r_id = g.push(Node {
op: Op::Reshape(Shape::from_dims(&[6])),
inputs: vec![t_id],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
(in_id, t_id, r_id)
};
let _ = t_id;
let mut inputs = StorageCache::new();
let in_arc = Arc::new(RwLock::new(storage));
inputs.insert(in_id, Arc::clone(&in_arc));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, r_id, inputs).expect("realize");
assert!(!Arc::ptr_eq(&result_arc, &in_arc));
assert_eq!(result_layout.shape().dims(), &[6]);
assert!(result_layout.is_contiguous());
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0]
);
}
#[test]
fn pipelined_realize_transpose_then_sum_dim_e2e() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, t_id, sum_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let t_id = g.push(Node {
op: Op::Transpose,
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let sum_id = g.push(Node {
op: Op::SumDim(1),
inputs: vec![t_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(sum_id, BackendId::Cpu);
(in_id, t_id, sum_id)
};
let _ = t_id;
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, result_layout) =
PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
assert_eq!(result_layout.shape().dims(), &[3]);
assert!(result_layout.is_contiguous());
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[5.0, 7.0, 9.0]);
}
#[test]
fn pipelined_realize_broadcast_then_add_e2e() {
let bc_input = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let plus_input = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (bc_in_id, plus_in_id, b_id, add_id) = {
let mut g = graph.write().unwrap();
let bc_in = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let plus_in = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::BroadcastTo(Shape::from_dims(&[2, 3])),
inputs: vec![bc_in],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![b, plus_in],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
(bc_in, plus_in, b, add)
};
let _ = b_id;
let mut inputs = StorageCache::new();
inputs.insert(bc_in_id, Arc::new(RwLock::new(bc_input)));
inputs.insert(plus_in_id, Arc::new(RwLock::new(plus_input)));
let (result_arc, _result_layout) =
PipelinedExecutor::realize(graph, add_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[11.0, 22.0, 33.0, 14.0, 25.0, 36.0]);
}
#[test]
fn pipelined_realize_sum_dim() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sum_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let sum_id = g.push(Node {
op: Op::SumDim(1),
inputs: vec![in_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(sum_id, BackendId::Cpu);
(in_id, sum_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[6.0, 15.0]);
}
}
#[test]
fn pipelined_realize_sum_all_rank3() {
let data: Vec<f32> = (1..=24).map(|x| x as f32).collect();
let storage = fuel_memory::from_slice_cpu(&data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sum_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3, 4]),
dtype: DType::F32,
});
let sum_id = g.push(Node {
op: Op::SumAll,
inputs: vec![in_id],
shape: Shape::from_dims(&[]),
dtype: DType::F32,
});
g.set_target_backend(sum_id, BackendId::Cpu);
(in_id, sum_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sum_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[300.0]);
}
}
#[test]
fn pipelined_realize_max_then_mean() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 9.0, 3.0, 4.0, 2.0, 8.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, mean_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let max_id = g.push(Node {
op: Op::MaxDim(1),
inputs: vec![in_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let mean_id = g.push(Node {
op: Op::MeanDim(0),
inputs: vec![max_id],
shape: Shape::from_dims(&[]),
dtype: DType::F32,
});
g.set_target_backend(max_id, BackendId::Cpu);
g.set_target_backend(mean_id, BackendId::Cpu);
(in_id, mean_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, mean_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[8.5]);
}
}
#[test]
fn pipelined_realize_sigmoid_then_silu() {
let storage = fuel_memory::from_slice_cpu(&[0.0_f32, 1.0, -1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sig_id, silu_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let sig_id = g.push(Node {
op: Op::Sigmoid,
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let silu_id = g.push(Node {
op: Op::Silu,
inputs: vec![sig_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(sig_id, BackendId::Cpu);
g.set_target_backend(silu_id, BackendId::Cpu);
(in_id, sig_id, silu_id)
};
let _ = sig_id;
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, silu_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert!((typed[0] - 0.5 * (1.0 / (1.0 + (-0.5_f32).exp()))).abs() < 1e-6);
}
}
#[test]
fn pipelined_realize_chained_unary_sqr_then_sqrt() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 4.0, 9.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, sqrt_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let sqr_id = g.push(Node {
op: Op::Sqr,
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let sqrt_id = g.push(Node {
op: Op::Sqrt,
inputs: vec![sqr_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(sqr_id, BackendId::Cpu);
g.set_target_backend(sqrt_id, BackendId::Cpu);
(in_id, sqrt_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, sqrt_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[1.0, 4.0, 9.0]);
}
}
#[test]
fn pipelined_realize_chained_adds() {
let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
let b_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
let c_storage = fuel_memory::from_slice_cpu(&[100.0_f32, 200.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, c_id, ab_id, abc_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let ab = g.push(Node {
op: Op::Add,
inputs: vec![a, b],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let abc = g.push(Node {
op: Op::Add,
inputs: vec![ab, c],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(ab, BackendId::Cpu);
g.set_target_backend(abc, BackendId::Cpu);
(a, b, c, ab, abc)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
inputs.insert(c_id, Arc::new(RwLock::new(c_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, abc_id, inputs).expect("realize");
let _ = ab_id;
let guard = result_arc.read().unwrap();
if let fuel_memory::BackendStorage::Cpu(c) = &guard.inner {
let typed: &[f32] = c.as_slice().unwrap();
assert_eq!(typed, &[111.0, 222.0]);
}
}
#[test]
fn pipelined_realize_triu_3x3_diag0() {
let storage =
fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, out_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 3]),
dtype: DType::F32,
});
let out_id = g.push(Node {
op: Op::Triu { diagonal: 0 },
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 3]),
dtype: DType::F32,
});
g.set_target_backend(out_id, BackendId::Cpu);
(in_id, out_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(out, &[1.0, 2.0, 3.0, 0.0, 5.0, 6.0, 0.0, 0.0, 9.0,]);
}
#[test]
fn pipelined_realize_tril_3x3_diag0_causal_mask() {
let storage =
fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, out_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 3]),
dtype: DType::F32,
});
let out_id = g.push(Node {
op: Op::Tril { diagonal: 0 },
inputs: vec![in_id],
shape: Shape::from_dims(&[3, 3]),
dtype: DType::F32,
});
g.set_target_backend(out_id, BackendId::Cpu);
(in_id, out_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(out, &[1.0, 0.0, 0.0, 4.0, 5.0, 0.0, 7.0, 8.0, 9.0,]);
}
#[test]
fn pipelined_realize_log_softmax_last_dim() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 3.0, 2.0, 1.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, out_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let out_id = g.push(Node {
op: Op::LogSoftmaxLastDim,
inputs: vec![in_id],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
g.set_target_backend(out_id, BackendId::Cpu);
(in_id, out_id)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
let expected = [
-2.407_606,
-1.4076059,
-0.40760595_f32,
-0.40760595,
-1.4076059,
-2.407_606,
];
for (a, b) in out.iter().zip(expected.iter()) {
assert!((a - b).abs() < 1e-5, "log_softmax mismatch: {a} vs {b}");
}
for row in out.chunks(3) {
let sum: f32 = row.iter().map(|x| x.exp()).sum();
assert!(
(sum - 1.0).abs() < 1e-5,
"softmax(log_softmax) row sum != 1: {sum}"
);
}
}
#[test]
fn pipelined_realize_merges_op_release_side_effect_root() {
let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
let b = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
let c = fuel_memory::from_slice_cpu(&[100.0_f32, 200.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, c_id, add_id, release_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let c = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![a, b],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let release = g.push(Node {
op: Op::Release,
inputs: vec![c],
shape: Shape::from_dims(&[0]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
g.set_target_backend(release, BackendId::Cpu);
g.add_side_effect_root(release);
(a, b, c, add, release)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
inputs.insert(b_id, Arc::new(RwLock::new(b)));
let c_arc = Arc::new(RwLock::new(c));
let c_arc_external = Arc::clone(&c_arc);
inputs.insert(c_id, c_arc);
let out = PipelinedExecutor::realize_many(graph, &[add_id], inputs).expect("realize_many");
assert_eq!(out.len(), 1);
let add_guard = out[0].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(arr) = &add_guard.inner else {
panic!()
};
assert_eq!(arr.as_slice::<f32>().unwrap(), &[11.0, 22.0]);
assert_eq!(
Arc::strong_count(&c_arc_external),
1,
"Op::Release should evict the cache's Arc to its source; \
only the test's external clone should remain",
);
let _ = release_id;
}
#[test]
fn compile_one_emits_release_marker_with_destructive_input() {
let graph_rc = Arc::new(RwLock::new(Graph::new()));
let (src_id, release_id) = {
let mut g = graph_rc.write().unwrap();
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let release = g.push(Node {
op: Op::Release,
inputs: vec![src],
shape: Shape::from_dims(&[0]),
dtype: DType::F32,
});
g.set_target_backend(release, BackendId::Cpu);
(src, release)
};
let g = graph_rc.read().unwrap();
let bindings = crate::dispatch::global_bindings();
let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
let item = compile_one(
&g,
release_id,
&mut layout_cache,
&bindings,
&SymEnv::default(),
)
.expect("compile_one Op::Release");
assert!(matches!(item.kind, WorkItemKind::ReleaseMarker));
assert_eq!(item.destructive_input, Some(0));
assert_eq!(item.inputs, vec![src_id]);
assert_eq!(
item.elem_count, 0,
"Release output is the zero-element marker"
);
}
#[test]
fn compile_one_emits_move_with_destructive_input() {
let graph_rc = Arc::new(RwLock::new(Graph::new()));
let (src_id, move_id) = {
let mut g = graph_rc.write().unwrap();
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let mv = g.push(Node {
op: Op::Move {
target: DeviceLocation::Cpu,
},
inputs: vec![src],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(mv, BackendId::Cpu);
(src, mv)
};
let g = graph_rc.read().unwrap();
let bindings = crate::dispatch::global_bindings();
let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
let item = compile_one(
&g,
move_id,
&mut layout_cache,
&bindings,
&SymEnv::default(),
)
.expect("compile_one Op::Move");
assert!(matches!(
item.kind,
WorkItemKind::Move {
target_location: DeviceLocation::Cpu
},
));
assert_eq!(item.destructive_input, Some(0), "Move destroys its source");
assert_eq!(item.inputs, vec![src_id]);
assert!(
item.compiled.is_some(),
"Move resolves a transfer kernel at (OpKind::Copy, [dt, dt], source backend)",
);
}
#[test]
fn compile_one_resolves_runtime_fused_from_the_sidecar() {
fn sentinel_kernel(
_inputs: &[Arc<RwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
_layouts: &[Layout],
_params: &crate::kernel::OpParams,
) -> Result<()> {
Ok(())
}
let region = fuel_graph::jit::PatternNode::Op {
op: fuel_graph::jit::OpTag::Exp,
attrs: fuel_graph::jit::OpAttrs::default(),
operands: vec![fuel_graph::jit::PatternNode::Op {
op: fuel_graph::jit::OpTag::Sub,
attrs: fuel_graph::jit::OpAttrs::default(),
operands: vec![
fuel_graph::jit::PatternNode::Bind { index: 0 },
fuel_graph::jit::PatternNode::Bind { index: 1 },
],
}],
};
let rid = crate::runtime_fused_kernels::adopt_runtime_fused(
"test::compile_one::exp_sub",
region,
sentinel_kernel as crate::kernel::KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("registrable region");
let graph_rc = Arc::new(RwLock::new(Graph::new()));
let fused_id = {
let mut g = graph_rc.write().unwrap();
let s = Shape::from_dims(&[4]);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
rid,
fuel_graph::registry::FusedOpParams::Runtime {
scalars: Vec::new(),
},
),
inputs: vec![a, b],
shape: s,
dtype: DType::F32,
});
g.set_target_backend(fused, BackendId::Cpu);
fused
};
let g = graph_rc.read().unwrap();
let bindings = crate::dispatch::global_bindings();
let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
let item = compile_one(
&g,
fused_id,
&mut layout_cache,
&bindings,
&SymEnv::default(),
)
.expect("compile_one runtime fused op");
assert!(matches!(item.kind, WorkItemKind::Kernel));
let compiled = item.compiled.expect("resolved a kernel");
assert_eq!(compiled.op, fuel_ir::dispatch::OpKind::RuntimeFused);
assert_eq!(compiled.backend, BackendId::Cpu);
assert!(
std::ptr::fn_addr_eq(compiled.kernel, sentinel_kernel as crate::kernel::KernelRef),
"the kernel pointer is exactly the adopted one",
);
assert!(
layout_cache.contains_key(&fused_id),
"output layout cached for consumers"
);
}
#[test]
fn compile_one_rejects_runtime_fused_dtype_mismatch() {
fn noop_kernel(
_inputs: &[Arc<RwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
_layouts: &[Layout],
_params: &crate::kernel::OpParams,
) -> Result<()> {
Ok(())
}
let region = fuel_graph::jit::PatternNode::Op {
op: fuel_graph::jit::OpTag::Log,
attrs: fuel_graph::jit::OpAttrs::default(),
operands: vec![fuel_graph::jit::PatternNode::Op {
op: fuel_graph::jit::OpTag::Maximum,
attrs: fuel_graph::jit::OpAttrs::default(),
operands: vec![
fuel_graph::jit::PatternNode::Bind { index: 0 },
fuel_graph::jit::PatternNode::Bind { index: 1 },
],
}],
};
let rid = crate::runtime_fused_kernels::adopt_runtime_fused(
"test::compile_one::log_maximum",
region,
noop_kernel as crate::kernel::KernelRef,
vec![DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("registrable region");
let graph_rc = Arc::new(RwLock::new(Graph::new()));
let fused_id = {
let mut g = graph_rc.write().unwrap();
let s = Shape::from_dims(&[4]);
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F64,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F64,
});
let fused = g.push(Node {
op: Op::Fused(
rid,
fuel_graph::registry::FusedOpParams::Runtime {
scalars: Vec::new(),
},
),
inputs: vec![a, b],
shape: s,
dtype: DType::F64,
});
g.set_target_backend(fused, BackendId::Cpu);
fused
};
let g = graph_rc.read().unwrap();
let bindings = crate::dispatch::global_bindings();
let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
let err = match compile_one(
&g,
fused_id,
&mut layout_cache,
&bindings,
&SymEnv::default(),
) {
Ok(_) => panic!("dtype mismatch must be a typed error"),
Err(e) => e,
};
assert!(
matches!(
err,
fuel_ir::Error::NoBackendForOp {
op: fuel_ir::dispatch::OpKind::RuntimeFused,
..
}
),
"a dtype mismatch is an honest binding miss: {err}",
);
}
#[test]
fn compile_one_maps_runtime_scalars_to_jit_scalars() {
fn noop_kernel(
_inputs: &[Arc<RwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
_layouts: &[Layout],
_params: &crate::kernel::OpParams,
) -> Result<()> {
Ok(())
}
let region = fuel_graph::jit::PatternNode::Op {
op: fuel_graph::jit::OpTag::Floor,
attrs: fuel_graph::jit::OpAttrs::default(),
operands: vec![fuel_graph::jit::PatternNode::Op {
op: fuel_graph::jit::OpTag::MulScalar,
attrs: fuel_graph::jit::OpAttrs::default(), operands: vec![fuel_graph::jit::PatternNode::Bind { index: 0 }],
}],
};
let rid = crate::runtime_fused_kernels::adopt_runtime_fused(
"test::compile_one::floor_mul_scalar",
region,
noop_kernel as crate::kernel::KernelRef,
vec![DType::F32, DType::F32],
BackendId::Cpu,
)
.expect("slot template registrable");
let graph_rc = Arc::new(RwLock::new(Graph::new()));
let fused_id = {
let mut g = graph_rc.write().unwrap();
let s = Shape::from_dims(&[4]);
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let fused = g.push(Node {
op: Op::Fused(
rid,
fuel_graph::registry::FusedOpParams::Runtime {
scalars: vec![3.5, -1.25],
},
),
inputs: vec![x],
shape: s,
dtype: DType::F32,
});
g.set_target_backend(fused, BackendId::Cpu);
fused
};
let g = graph_rc.read().unwrap();
let bindings = crate::dispatch::global_bindings();
let mut layout_cache: HashMap<NodeId, Layout> = HashMap::new();
let item = compile_one(
&g,
fused_id,
&mut layout_cache,
&bindings,
&SymEnv::default(),
)
.expect("compile_one runtime fused op with scalars");
let compiled = item.compiled.expect("resolved a kernel");
assert!(
matches!(
&compiled.op_params,
crate::kernel::OpParams::JitScalars { scalars } if scalars == &vec![3.5, -1.25]
),
"Runtime scalars ride into OpParams::JitScalars, got {:?}",
compiled.op_params,
);
}
#[test]
fn pipelined_op_move_same_device_is_plain_copy() {
let src = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (src_id, move_id) = {
let mut g = graph.write().unwrap();
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let mv = g.push(Node {
op: Op::Move {
target: DeviceLocation::Cpu,
},
inputs: vec![src],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(mv, BackendId::Cpu);
(src, mv)
};
let src_arc = Arc::new(RwLock::new(src));
let src_arc_external = Arc::clone(&src_arc);
let mut inputs = StorageCache::new();
inputs.insert(src_id, src_arc);
let (out, layout) =
PipelinedExecutor::realize(graph, move_id, inputs).expect("realize Op::Move");
let guard = out.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!()
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[1.0, 2.0, 3.0, 4.0],
"Op::Move output should carry input's data to target",
);
assert!(layout.is_contiguous());
drop(guard);
assert!(
!Arc::ptr_eq(&out, &src_arc_external),
"same-device Move is a plain copy — fresh storage, not an alias",
);
assert_eq!(
Arc::strong_count(&src_arc_external),
1,
"Op::Move should evict the cache's Arc to its source",
);
}
#[test]
fn pipelined_op_move_multi_consumer_source_not_stranded() {
let a = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, out_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Relu,
inputs: vec![a],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let m = g.push(Node {
op: Op::Move {
target: DeviceLocation::Cpu,
},
inputs: vec![a],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let out = g.push(Node {
op: Op::Add,
inputs: vec![b, m],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(b, BackendId::Cpu);
g.set_target_backend(m, BackendId::Cpu);
g.set_target_backend(out, BackendId::Cpu);
(a, out)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a)));
let (out, _) = PipelinedExecutor::realize(graph, out_id, inputs)
.expect("realize add(relu(a), move(a))");
let guard = out.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!()
};
assert_eq!(
c.as_slice::<f32>().unwrap(),
&[2.0, 4.0, 6.0, 8.0],
"relu(a) must run before move(a) evicts a",
);
}
#[test]
fn pipelined_op_move_source_in_target_set_not_evicted() {
let a = fuel_memory::from_slice_cpu(&[5.0_f32, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, move_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let mv = g.push(Node {
op: Op::Move {
target: DeviceLocation::Cpu,
},
inputs: vec![a],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(mv, BackendId::Cpu);
(a, mv)
};
let a_arc = Arc::new(RwLock::new(a));
let a_arc_external = Arc::clone(&a_arc);
let mut inputs = StorageCache::new();
inputs.insert(a_id, a_arc);
let out = PipelinedExecutor::realize_many(graph, &[move_id, a_id], inputs)
.expect("realize_many [move, source]");
assert_eq!(out.len(), 2);
let mv_guard = out[0].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &mv_guard.inner else {
panic!()
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[5.0, 6.0]);
assert!(
Arc::ptr_eq(&out[1].0, &a_arc_external),
"requested source must survive the Move's destructive cleanup",
);
}
#[test]
fn pipelined_realize_many_empty_targets() {
let graph = Arc::new(RwLock::new(Graph::new()));
let inputs = StorageCache::new();
let out = PipelinedExecutor::realize_many(graph, &[], inputs).expect("realize_many");
assert!(out.is_empty());
}
#[test]
fn pipelined_realize_many_two_independent_targets() {
let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
let rhs = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (lhs_id, rhs_id, add_id, mul_id) = {
let mut g = graph.write().unwrap();
let lhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let rhs = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let mul = g.push(Node {
op: Op::Mul,
inputs: vec![lhs, rhs],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
g.set_target_backend(mul, BackendId::Cpu);
(lhs, rhs, add, mul)
};
let mut inputs = StorageCache::new();
inputs.insert(lhs_id, Arc::new(RwLock::new(lhs)));
inputs.insert(rhs_id, Arc::new(RwLock::new(rhs)));
let out = PipelinedExecutor::realize_many(graph, &[mul_id, add_id], inputs)
.expect("realize_many");
assert_eq!(out.len(), 2);
let mul_guard = out[0].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &mul_guard.inner else {
panic!()
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[10.0, 40.0, 90.0]);
let add_guard = out[1].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &add_guard.inner else {
panic!()
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 22.0, 33.0]);
}
#[test]
fn pipelined_realize_many_duplicate_target_returns_same_arc() {
let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (in_id, neg_id) = {
let mut g = graph.write().unwrap();
let in_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let neg = g.push(Node {
op: Op::Neg,
inputs: vec![in_id],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(neg, BackendId::Cpu);
(in_id, neg)
};
let mut inputs = StorageCache::new();
inputs.insert(in_id, Arc::new(RwLock::new(storage)));
let out = PipelinedExecutor::realize_many(graph, &[neg_id, neg_id], inputs)
.expect("realize_many");
assert_eq!(out.len(), 2);
assert!(
Arc::ptr_eq(&out[0].0, &out[1].0),
"duplicate target should share the same storage Arc",
);
}
#[test]
fn pipelined_realize_many_shared_subgraph_evaluates_once() {
let a_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
let b_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (a_id, b_id, add_id, mul_id) = {
let mut g = graph.write().unwrap();
let a = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let b = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let add = g.push(Node {
op: Op::Add,
inputs: vec![a, b],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
let mul = g.push(Node {
op: Op::Mul,
inputs: vec![add, a],
shape: Shape::from_dims(&[3]),
dtype: DType::F32,
});
g.set_target_backend(add, BackendId::Cpu);
g.set_target_backend(mul, BackendId::Cpu);
(a, b, add, mul)
};
let mut inputs = StorageCache::new();
inputs.insert(a_id, Arc::new(RwLock::new(a_storage)));
inputs.insert(b_id, Arc::new(RwLock::new(b_storage)));
let out = PipelinedExecutor::realize_many(graph, &[add_id, mul_id], inputs)
.expect("realize_many");
assert_eq!(out.len(), 2);
let add_guard = out[0].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &add_guard.inner else {
panic!()
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 22.0, 33.0]);
let mul_guard = out[1].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &mul_guard.inner else {
panic!()
};
assert_eq!(c.as_slice::<f32>().unwrap(), &[11.0, 44.0, 99.0]);
}
#[test]
fn pipelined_realize_masked_fill_attention_pattern() {
let x_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
let mut mask_storage = fuel_memory::from_slice_cpu(&[0u8, 1, 0, 1]);
mask_storage.dtype = DType::Bool; let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, mask_id, out_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let mask = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::Bool,
});
let out = g.push(Node {
op: Op::MaskedFill {
value: fuel_ir::Scalar::F32(-1000.0),
},
inputs: vec![x, mask],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(out, BackendId::Cpu);
(x, mask, out)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
inputs.insert(mask_id, Arc::new(RwLock::new(mask_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, out_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(out, &[1.0, -1000.0, 3.0, -1000.0]);
}
#[test]
fn pipelined_realize_write_slice_kv_cache_append() {
let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 24]);
let src_storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (dest_id, src_id, ws_id) = {
let mut g = graph.write().unwrap();
let dest = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4, 3, 2]),
dtype: DType::F32,
});
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 3, 2]),
dtype: DType::F32,
});
let ws = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(2, 3), (0, 3), (0, 2)],
dyn_offset: None,
},
inputs: vec![dest, src],
shape: Shape::from_dims(&[4, 3, 2]), dtype: DType::F32,
});
g.set_target_backend(ws, BackendId::Cpu);
(dest, src, ws)
};
let mut inputs = StorageCache::new();
inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, ws_id, inputs).expect("realize");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
let expected = [
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,
6.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
];
assert_eq!(out, &expected);
}
#[test]
fn pipelined_realize_write_slice_aliases_dest_storage() {
let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
let src_storage = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (dest_id, src_id, ws_id, dest_arc_external) = {
let mut g = graph.write().unwrap();
let dest = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[1, 2]),
dtype: DType::F32,
});
let ws = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(1, 2), (0, 2)],
dyn_offset: None,
},
inputs: vec![dest, src],
shape: Shape::from_dims(&[3, 2]),
dtype: DType::F32,
});
g.set_target_backend(ws, BackendId::Cpu);
let arc = Arc::new(RwLock::new(dest_storage));
(dest, src, ws, arc)
};
let mut inputs = StorageCache::new();
inputs.insert(dest_id, Arc::clone(&dest_arc_external));
inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, ws_id, inputs).expect("realize");
assert!(
Arc::ptr_eq(&result_arc, &dest_arc_external),
"WriteSlice output Arc must be the same Arc as the destination's"
);
let guard = dest_arc_external.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(out, &[0.0, 0.0, 10.0, 20.0, 0.0, 0.0]);
}
#[test]
fn pipelined_realize_write_slice_dynamic_offset() {
use fuel_ir::{DynScalar, SymId};
let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
let src_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let sym = SymId(0);
let (dest_id, src_id, ws_id) = {
let mut g = graph.write().unwrap();
let dest = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let ws = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: Some((0, DynScalar::Sym(sym))),
},
inputs: vec![dest, src],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
g.set_target_backend(ws, BackendId::Cpu);
(dest, src, ws)
};
let mut inputs = StorageCache::new();
inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
let mut env = SymEnv::new();
env.bind(sym, 3).unwrap();
let (result_arc, _) = PipelinedExecutor::realize_with_env(graph, ws_id, inputs, env)
.expect("realize_with_env");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(
out,
&[0.0, 0.0, 0.0, 7.0, 8.0, 0.0],
"slab must land at the SymEnv-bound offset 3, not the static placeholder 0",
);
}
#[test]
fn pipelined_realize_nonzero_indices_publishes_count() {
use fuel_ir::SymId;
use fuel_ir::storage::{OutputViewSpec, compose_bundle};
let x_storage = fuel_memory::from_slice_cpu(&[0.0_f32, 1.0, 0.0, 1.0, 1.0, 0.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let count_sym = SymId(0);
let (x_id, nzi_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, 3]),
dtype: DType::F32,
});
let capacity = 6usize;
let idx_shape = Shape::from_dims(&[capacity]);
let cnt_shape = Shape::from_dims(&[1]);
let specs = vec![
OutputViewSpec {
dtype: DType::U32,
shape: idx_shape.clone(),
layout: Layout::contiguous(idx_shape.clone()),
name: Some("indices"),
},
OutputViewSpec {
dtype: DType::U32,
shape: cnt_shape.clone(),
layout: Layout::contiguous(cnt_shape),
name: Some("count"),
},
];
let (_bytes, views) = compose_bundle(&specs).unwrap();
let nzi = g.push(Node {
op: Op::NonZeroIndices { count_sym },
inputs: vec![x],
shape: idx_shape,
dtype: DType::U32,
});
g.set_output_views(nzi, Arc::from(views.into_boxed_slice()))
.unwrap();
g.set_target_backend(nzi, BackendId::Cpu);
(x, nzi)
};
let mut inputs = StorageCache::new();
inputs.insert(x_id, Arc::new(RwLock::new(x_storage)));
let (_arc, _layout, produced) =
PipelinedExecutor::realize_with_env_producing(graph, nzi_id, inputs, SymEnv::new())
.expect("realize_with_env_producing");
assert_eq!(
produced.get(count_sym),
Some(3),
"the executor must publish the data-determined nonzero count into the SymEnv",
);
}
#[test]
fn pipelined_realize_write_slice_dynamic_offset_unbound_errors() {
use fuel_ir::{DynScalar, SymId};
let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
let src_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let (dest_id, src_id, ws_id) = {
let mut g = graph.write().unwrap();
let dest = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let ws = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: Some((0, DynScalar::Sym(SymId(0)))),
},
inputs: vec![dest, src],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
g.set_target_backend(ws, BackendId::Cpu);
(dest, src, ws)
};
let mut inputs = StorageCache::new();
inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
let result = PipelinedExecutor::realize_with_env(graph, ws_id, inputs, SymEnv::new());
assert!(
result.is_err(),
"unbound dyn_offset symbol must surface a typed error, not a panic",
);
}
#[test]
fn pipelined_realize_write_slice_dynamic_offset_out_of_capacity_errors() {
use fuel_ir::{DynScalar, SymId};
let dest_storage = fuel_memory::from_slice_cpu(&[0.0_f32; 6]);
let src_storage = fuel_memory::from_slice_cpu(&[7.0_f32, 8.0]);
let graph = Arc::new(RwLock::new(Graph::new()));
let sym = SymId(0);
let (dest_id, src_id, ws_id) = {
let mut g = graph.write().unwrap();
let dest = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
let src = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let ws = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: Some((0, DynScalar::Sym(sym))),
},
inputs: vec![dest, src],
shape: Shape::from_dims(&[6]),
dtype: DType::F32,
});
g.set_target_backend(ws, BackendId::Cpu);
(dest, src, ws)
};
let mut inputs = StorageCache::new();
inputs.insert(dest_id, Arc::new(RwLock::new(dest_storage)));
inputs.insert(src_id, Arc::new(RwLock::new(src_storage)));
let mut env = SymEnv::new();
env.bind(sym, 5).unwrap();
let result = PipelinedExecutor::realize_with_env(graph, ws_id, inputs, env);
assert!(
result.is_err(),
"resolved slab past destination capacity must surface a typed error",
);
}
#[test]
fn pipelined_inplace_with_multiple_readers_orders_correctly() {
use fuel_ir::Shape;
let data = [1.0_f32, -2.0, 3.0, -4.0];
let src_storage = fuel_memory::from_slice_cpu(&data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, step_id, sig_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let step = g.push(Node {
op: Op::Step,
inputs: vec![x],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let sig = g.push(Node {
op: Op::SigmoidInplace,
inputs: vec![x],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(step, BackendId::Cpu);
g.set_target_backend(sig, BackendId::Cpu);
(x, step, sig)
};
let mut cache = StorageCache::new();
cache.insert(x_id, Arc::new(RwLock::new(src_storage)));
let results = PipelinedExecutor::realize_many(graph, &[step_id, sig_id], cache)
.expect("realize_many");
let step_guard = results[0].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &step_guard.inner else {
panic!("expected Cpu storage for step output");
};
let step_out: &[f32] = c.as_slice().expect("f32 cast");
assert_eq!(
step_out,
&[1.0_f32, 0.0, 1.0, 0.0],
"Step must run before SigmoidInplace; got post-mutation bytes"
);
drop(step_guard);
let sig_guard = results[1].0.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &sig_guard.inner else {
panic!("expected Cpu storage for sigmoid output");
};
let sig_out: &[f32] = c.as_slice().expect("f32 cast");
let sig_ref = [
1.0_f32 / (1.0 + (-1.0_f32).exp()),
1.0_f32 / (1.0 + (2.0_f32).exp()),
1.0_f32 / (1.0 + (-3.0_f32).exp()),
1.0_f32 / (1.0 + (4.0_f32).exp()),
];
for (got, want) in sig_out.iter().zip(sig_ref.iter()) {
assert!((got - want).abs() < 1e-6, "sigmoid: got {got}, want {want}");
}
}
#[test]
fn pipelined_residual_connection_inserts_safety_copy() {
use fuel_ir::Shape;
let data = [1.0_f32, -2.0, 3.0, -4.0];
let src_storage = fuel_memory::from_slice_cpu(&data);
let graph = Arc::new(RwLock::new(Graph::new()));
let (x_id, _y_id, z_id) = {
let mut g = graph.write().unwrap();
let x = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let y = g.push(Node {
op: Op::ReluInplace,
inputs: vec![x],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let z = g.push(Node {
op: Op::Add,
inputs: vec![y, x],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(y, BackendId::Cpu);
g.set_target_backend(z, BackendId::Cpu);
(x, y, z)
};
let mut cache = StorageCache::new();
cache.insert(x_id, Arc::new(RwLock::new(src_storage)));
let (result_arc, _) = PipelinedExecutor::realize(graph, z_id, cache)
.expect("realize z — insert_safety_copies must break the cycle");
let guard = result_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let got: &[f32] = c.as_slice().expect("f32 cast");
assert_eq!(
got,
&[2.0_f32, -2.0, 6.0, -4.0],
"Add must read pre-mutation x via the inserted safety copy"
);
}
fn cpu_const_f32(
graph: &Arc<RwLock<Graph>>,
cache: &mut StorageCache,
data: &[f32],
dims: &[usize],
) -> NodeId {
let id = {
let mut g = graph.write().unwrap();
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(dims),
dtype: DType::F32,
})
};
cache.insert(id, Arc::new(RwLock::new(fuel_memory::from_slice_cpu(data))));
id
}
#[test]
fn pipelined_realize_selective_scan_bundled_produces_y_and_last_state() {
let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
let mut cache = StorageCache::new();
let u_id = cpu_const_f32(&graph, &mut cache, &[3.0], &[1, 1, 1]);
let delta_id = cpu_const_f32(&graph, &mut cache, &[1.0], &[1, 1, 1]);
let a_id = cpu_const_f32(&graph, &mut cache, &[-1.0], &[1, 1]);
let b_id = cpu_const_f32(&graph, &mut cache, &[2.0], &[1, 1, 1]);
let c_id = cpu_const_f32(&graph, &mut cache, &[0.5], &[1, 1, 1]);
let (y_id, last_state_id) = {
let u_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), u_id);
let delta_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), delta_id);
let a_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), a_id);
let b_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), b_id);
let c_t = fuel_graph::NodeHandle::from_existing(Arc::clone(&graph), c_id);
let (y, last_state) = u_t
.selective_scan_bundled(&delta_t, &a_t, &b_t, &c_t, false)
.expect("selective_scan_bundled");
(y.id(), last_state.id())
};
{
let mut g = graph.write().unwrap();
let producer = g.node(y_id).inputs[0];
g.set_target_backend(producer, BackendId::Cpu);
}
let outputs =
PipelinedExecutor::realize_many(Arc::clone(&graph), &[y_id, last_state_id], cache)
.expect("realize_many");
let (y_arc, y_layout) = outputs[0].clone();
let (last_state_arc, last_state_layout) = outputs[1].clone();
assert!(
Arc::ptr_eq(&y_arc, &last_state_arc),
"y and last_state Views project the SAME bundled producer Storage",
);
let guard = y_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c_bytes) = &guard.inner else {
panic!("expected Cpu storage");
};
let typed: &[f32] = c_bytes.as_slice().expect("f32 cast");
assert_eq!(y_layout.start_offset(), 0);
let y_value = typed[y_layout.start_offset()];
assert!((y_value - 3.0).abs() < 1e-5, "y expected 3.0 got {y_value}");
assert_eq!(
last_state_layout.start_offset(),
1,
"slot 1 byte_offset 4 / sizeof(f32) = 1",
);
let ls_value = typed[last_state_layout.start_offset()];
assert!(
(ls_value - 6.0).abs() < 1e-5,
"last_state expected 6.0 got {ls_value}"
);
}
#[test]
fn pipelined_write_slice_round2_plain_slice_reader_not_stranded() {
let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
let mut cache = StorageCache::new();
let b_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
let src1_id = cpu_const_f32(&graph, &mut cache, &[1.0, 2.0], &[2]);
let src2_id = cpu_const_f32(&graph, &mut cache, &[99.0, 98.0], &[2]);
let (k_id, y_id) = {
let mut g = graph.write().unwrap();
let x_id = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![b_id, src1_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let s_id = g.push(Node {
op: Op::Slice {
dim: 0,
start: 0,
len: 2,
},
inputs: vec![x_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let k_id = g.push(Node {
op: Op::Relu,
inputs: vec![s_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let y_id = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![x_id, src2_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(x_id, BackendId::Cpu);
g.set_target_backend(k_id, BackendId::Cpu);
g.set_target_backend(y_id, BackendId::Cpu);
(k_id, y_id)
};
let plan = fuel_graph::opt::execution_plan(&graph.read().unwrap(), &[k_id, y_id]);
eprintln!("[diagnostic plain-slice] execution_plan = {:?}", plan,);
let outputs = PipelinedExecutor::realize_many(Arc::clone(&graph), &[k_id, y_id], cache)
.expect("realize_many: plain slice reader must not error");
let (k_arc, _) = &outputs[0];
let guard = k_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(
out,
&[1.0, 2.0],
"K = relu(slice(X,0,2)) must read PRE-round-2 bytes; got {:?}",
out,
);
}
#[test]
fn pipelined_write_slice_reshape_adopted_arc_reader_reads_pre_write_bytes() {
let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
let mut cache = StorageCache::new();
let b_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
let src1_id = cpu_const_f32(&graph, &mut cache, &[1.0, 2.0], &[2]);
let src2_id = cpu_const_f32(&graph, &mut cache, &[99.0, 98.0], &[2]);
let (x_id, k_id, y_id) = {
let mut g = graph.write().unwrap();
let x_id = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![b_id, src1_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let s_id = g.push(Node {
op: Op::Slice {
dim: 0,
start: 0,
len: 2,
},
inputs: vec![x_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let r_id = g.push(Node {
op: Op::Reshape(Shape::from_dims(&[2])),
inputs: vec![s_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let k_id = g.push(Node {
op: Op::Relu,
inputs: vec![r_id],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let y_id = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![x_id, src2_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(x_id, BackendId::Cpu);
g.set_target_backend(k_id, BackendId::Cpu);
g.set_target_backend(y_id, BackendId::Cpu);
(x_id, k_id, y_id)
};
let _ = x_id;
let plan = fuel_graph::opt::execution_plan(&graph.read().unwrap(), &[k_id, y_id]);
eprintln!(
"[diagnostic reshape-adopted-arc] execution_plan = {:?}",
plan,
);
let outputs = PipelinedExecutor::realize_many(Arc::clone(&graph), &[k_id, y_id], cache)
.expect("realize_many");
let (k_arc, _) = &outputs[0];
let guard = k_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(
out,
&[1.0, 2.0],
"K must read PRE-round-2 bytes through the adopted Arc; got {:?} \
(post-round-2 corruption if this fails)",
out,
);
}
#[test]
fn pipelined_write_slice_mla_two_layer_round2_reshape_reader_not_stranded() {
let graph: Arc<RwLock<Graph>> = Arc::new(RwLock::new(Graph::new()));
let mut cache = StorageCache::new();
let b_l_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
let src_l_n_id = cpu_const_f32(&graph, &mut cache, &[1.0, 2.0], &[2]);
let b_l1_id = cpu_const_f32(&graph, &mut cache, &[0.0, 0.0, 0.0, 0.0], &[4]);
let src_l_n1_id = cpu_const_f32(&graph, &mut cache, &[99.0, 98.0], &[2]);
let (x_l_n1_id, x_l1_n_id) = {
let mut g = graph.write().unwrap();
let x_l_n = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![b_l_id, src_l_n_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let s_l_n = g.push(Node {
op: Op::Slice {
dim: 0,
start: 0,
len: 2,
},
inputs: vec![x_l_n],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let rs_l_n = g.push(Node {
op: Op::Reshape(Shape::from_dims(&[2])),
inputs: vec![s_l_n],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let r_l_n = g.push(Node {
op: Op::Relu,
inputs: vec![rs_l_n],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
let x_l1_n = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![b_l1_id, r_l_n],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
let x_l_n1 = g.push(Node {
op: Op::WriteSlice {
ranges: vec![(0, 2)],
dyn_offset: None,
},
inputs: vec![x_l_n, src_l_n1_id],
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
});
g.set_target_backend(x_l_n, BackendId::Cpu);
g.set_target_backend(r_l_n, BackendId::Cpu);
g.set_target_backend(x_l1_n, BackendId::Cpu);
g.set_target_backend(x_l_n1, BackendId::Cpu);
(x_l_n1, x_l1_n)
};
let plan = fuel_graph::opt::execution_plan(&graph.read().unwrap(), &[x_l1_n_id, x_l_n1_id]);
eprintln!("[diagnostic mla-2-layer] execution_plan = {:?}", plan);
let outputs =
PipelinedExecutor::realize_many(Arc::clone(&graph), &[x_l1_n_id, x_l_n1_id], cache)
.expect("realize_many");
let (x_l1_n_arc, _) = &outputs[0];
let guard = x_l1_n_arc.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected Cpu storage");
};
let out: &[f32] = c.as_slice().unwrap();
assert_eq!(
out,
&[1.0, 2.0, 0.0, 0.0],
"layer L+1's round-N append must reflect PRE-round-(N+1) attention \
math output; got {:?} (post-round-(N+1) corruption if this fails)",
out,
);
}
}