use std::sync::{Arc, RwLock};
use fuel_backend_contract::backend::{BackendRuntime, BackendStreams};
use fuel_backend_contract::dyn_backend::DynBackendDevice;
use fuel_dispatch::dispatch::global_bindings;
use fuel_dispatch::dispatched_kernel_ident;
use fuel_dispatch::optimize::{OptimizedGraph, optimize_graph_with_runtime_fusion};
use fuel_dispatch::pipelined::{PipelinedExecutor, StorageCache};
use fuel_dispatch::plan::PlanOptions;
use fuel_dispatch::ranker::{
BackendRuntimeHandle, BackendRuntimeLookup, ChainedSelector, JudgeOracle, RuntimeSelector,
};
use fuel_graph::{Graph, Node, NodeId, Op, topo_order_multi};
use fuel_ir::backend::FitStatus;
use fuel_ir::{DeviceLocation, Error, HostBuffer, Layout, Result, SymEnv, probe::BackendId};
use fuel_memory::{BackendStorage, Storage};
use crate::Device;
use fuel_dispatch::topology::SystemTopology;
use fuel_ir::dtype::WithDType;
type RootKernelIdent = Option<(&'static str, u64)>;
static OPTIMIZE_CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
std::thread_local! {
static OPTIMIZE_CALLS_TL: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
pub fn optimize_calls() -> usize {
OPTIMIZE_CALLS.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn optimize_calls_thread_local() -> usize {
OPTIMIZE_CALLS_TL.with(|c| c.get())
}
pub fn realize_one_as<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
) -> Result<Vec<T>> {
realize_one_as_with_initial::<T>(graph, target, device, StorageCache::new())
}
pub fn realize_one_as_multi_device<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
extra_devices: &[&Device],
) -> Result<Vec<T>> {
let initial = seed_extra_device_handles(graph, device, extra_devices)?;
realize_one_as_with_initial::<T>(graph, target, device, initial)
}
fn seed_extra_device_handles(
graph: &Arc<RwLock<Graph>>,
primary: &Device,
extra_devices: &[&Device],
) -> Result<StorageCache> {
let mut cache = StorageCache::new();
if extra_devices.is_empty() {
return Ok(cache);
}
let primary_backend = device_to_backend_id(primary);
let mut seeded: Vec<BackendId> = vec![primary_backend];
for dev in extra_devices {
let backend = device_to_backend_id(dev);
if seeded.contains(&backend) {
continue;
}
let Some(seed) = device_seed_storage(dev)? else {
seeded.push(backend);
continue;
};
let anchor_id = {
let mut g = graph.write().map_err(|_| {
Error::Msg("graph lock poisoned during device-handle seed".into()).bt()
})?;
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: fuel_ir::Shape::from_dims(&[4]),
dtype: fuel_ir::DType::U8,
})
};
cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
seeded.push(backend);
}
Ok(cache)
}
pub fn realize_many_as<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
targets: &[NodeId],
device: &Device,
) -> Result<Vec<Vec<T>>> {
realize_many_as_with_initial::<T>(graph, targets, device, StorageCache::new())
}
pub fn realize_one_as_with_initial<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
) -> Result<Vec<T>> {
realize_one_as_with_initial_reporting::<T>(graph, target, device, initial, &SymEnv::default())
.map(|(bytes, _root_kernel_source)| bytes)
}
pub fn realize_one_as_with_initial_env<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
) -> Result<Vec<T>> {
realize_one_as_with_initial_reporting::<T>(graph, target, device, initial, sym_env)
.map(|(bytes, _root_kernel_source)| bytes)
}
pub fn realize_one_as_with_initial_reporting<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
) -> Result<(Vec<T>, RootKernelIdent)> {
realize_one_as_reporting_impl::<T>(graph, target, device, initial, sym_env, true)
}
pub fn realize_one_reference_as<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
) -> Result<Vec<T>> {
realize_one_as_reporting_impl::<T>(
graph,
target,
device,
StorageCache::new(),
&SymEnv::default(),
false,
)
.map(|(bytes, _root_kernel_source)| bytes)
}
fn realize_one_to_storage(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
allow_cost_placement: bool,
) -> Result<(Arc<RwLock<Storage>>, RootKernelIdent)> {
let (cache, _backend_id, mut effective_targets) = prepare(graph, &[target], device, initial)?;
let Some(cpu_target) = effective_targets.pop() else {
return Err(Error::Msg(
"pipelined_bridge: prepare returned no effective target for a \
single-target realize — internal bug"
.into(),
)
.bt());
};
dispatch_with_plan_retry(
graph,
cpu_target,
cache,
device,
target,
sym_env,
allow_cost_placement,
)
}
fn realize_one_as_reporting_impl<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
allow_cost_placement: bool,
) -> Result<(Vec<T>, RootKernelIdent)> {
let (storage, root_kernel_source) = realize_one_to_storage(
graph,
target,
device,
initial,
sym_env,
allow_cost_placement,
)?;
Ok((extract_cpu_bytes_typed::<T>(&storage)?, root_kernel_source))
}
pub fn realize_one_bytes(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
) -> Result<Vec<u8>> {
realize_one_bytes_with_initial(graph, target, device, StorageCache::new())
}
pub fn realize_one_bytes_with_initial(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
) -> Result<Vec<u8>> {
let (storage, _root_kernel_source) =
realize_one_to_storage(graph, target, device, initial, &SymEnv::default(), true)?;
extract_cpu_bytes_raw(&storage)
}
pub fn prebuild_optimized_env<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
) -> Result<(NodeId, OptimizedGraph, Vec<T>)> {
let (cache, _backend_id, mut effective_targets) = prepare(graph, &[target], device, initial)?;
let Some(cpu_target) = effective_targets.pop() else {
return Err(Error::Msg(
"pipelined_bridge: prepare returned no effective target for a \
single-target prebuild — internal bug"
.into(),
)
.bt());
};
let (storage, optimized, _full_cache) =
dispatch_with_plan_retry_capturing(graph, cpu_target, cache, device, sym_env)?;
let bytes = extract_cpu_bytes_typed::<T>(&storage)?;
Ok((cpu_target, optimized, bytes))
}
pub fn prebuild_optimized_env_capturing_cache<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
target: NodeId,
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
) -> Result<(NodeId, OptimizedGraph, StorageCache, Vec<T>)> {
let (cache, _backend_id, mut effective_targets) = prepare(graph, &[target], device, initial)?;
let Some(cpu_target) = effective_targets.pop() else {
return Err(Error::Msg(
"pipelined_bridge: prepare returned no effective target for a \
single-target prebuild — internal bug"
.into(),
)
.bt());
};
let (storage, optimized, full_cache) =
dispatch_with_plan_retry_capturing(graph, cpu_target, cache, device, sym_env)?;
let bytes = extract_cpu_bytes_typed::<T>(&storage)?;
Ok((cpu_target, optimized, full_cache, bytes))
}
pub fn realize_one_prebuilt_env<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
effective_target: NodeId,
optimized: &OptimizedGraph,
device: &Device,
cache: StorageCache,
sym_env: &SymEnv,
) -> Result<Vec<T>> {
let (selector, lookup) = match production_selector_for(device) {
Some((s, l)) => (Some(s), Some(l)),
None => (None, None),
};
let (storage, _layout) = PipelinedExecutor::realize_with_optimized_picking_env(
graph.clone(),
effective_target,
cache,
optimized,
selector,
lookup,
sym_env.clone(),
)?;
extract_cpu_bytes_typed::<T>(&storage)
}
pub fn realize_one_prebuilt_cached_env<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
effective_target: NodeId,
optimized: &OptimizedGraph,
order_cache: &std::sync::OnceLock<Vec<NodeId>>,
device: &Device,
cache: StorageCache,
sym_env: &SymEnv,
) -> Result<Vec<T>> {
let (selector, lookup) = match production_selector_for(device) {
Some((s, l)) => (Some(s), Some(l)),
None => (None, None),
};
let (storage, _layout) = PipelinedExecutor::realize_with_optimized_picking_cached_env(
graph.clone(),
effective_target,
cache,
optimized,
order_cache,
selector,
lookup,
sym_env.clone(),
)?;
extract_cpu_bytes_typed::<T>(&storage)
}
fn dispatch_with_plan_retry_capturing(
graph: &Arc<RwLock<Graph>>,
cpu_target: NodeId,
cache: StorageCache,
device: &Device,
sym_env: &SymEnv,
) -> Result<(Arc<RwLock<Storage>>, OptimizedGraph, StorageCache)> {
let pinned_loc = device.location();
let mut retry = TopologyRetryState::new();
loop {
let optimized = build_optimized_graph(graph, &[cpu_target], pinned_loc, &cache, true)?;
let (selector, lookup) = match production_selector_for(device) {
Some((s, l)) => (Some(s), Some(l)),
None => (None, None),
};
let cache_for_attempt = cache.clone();
let result = PipelinedExecutor::realize_with_optimized_picking_env(
graph.clone(),
cpu_target,
cache_for_attempt,
&optimized,
selector,
lookup,
sym_env.clone(),
);
match result {
Ok((storage, _layout)) => return Ok((storage, optimized, cache)),
Err(e) if matches!(e, Error::TopologyChanged { .. }) && retry.permit_retry() => {
continue;
}
Err(e) => return Err(e),
}
}
}
const TOPOLOGY_SETTLE_BUDGET: std::time::Duration = std::time::Duration::from_secs(2);
const MAX_PLAN_REBUILDS: usize = 256;
fn topology_retry_backoff(attempt: usize) {
let ms = 1u64 << attempt.min(4);
std::thread::sleep(std::time::Duration::from_millis(ms));
}
struct TopologyRetryState {
attempt: usize,
last_gen: u64,
last_movement: std::time::Instant,
}
impl TopologyRetryState {
fn new() -> Self {
Self {
attempt: 0,
last_gen: fuel_dispatch::dispatch::topology_generation(),
last_movement: std::time::Instant::now(),
}
}
fn permit_retry(&mut self) -> bool {
let now_gen = fuel_dispatch::dispatch::topology_generation();
if now_gen != self.last_gen {
self.last_gen = now_gen;
self.last_movement = std::time::Instant::now();
}
if self.attempt + 1 >= MAX_PLAN_REBUILDS
|| self.last_movement.elapsed() >= TOPOLOGY_SETTLE_BUDGET
{
return false;
}
topology_retry_backoff(self.attempt);
self.attempt += 1;
true
}
}
fn feasible_fallback_devices(
all_devices: &[DeviceLocation],
resident: &[DeviceLocation],
node_dev: DeviceLocation,
) -> Vec<DeviceLocation> {
all_devices
.iter()
.copied()
.filter(|&d| d != node_dev && (matches!(d, DeviceLocation::Cpu) || resident.contains(&d)))
.collect()
}
fn build_optimized_graph(
graph: &Arc<RwLock<Graph>>,
roots: &[NodeId],
pinned_device: DeviceLocation,
cache: &StorageCache,
allow_cost_placement: bool,
) -> Result<OptimizedGraph> {
OPTIMIZE_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
OPTIMIZE_CALLS_TL.with(|c| c.set(c.get() + 1));
let topology = SystemTopology::current();
let placements_for =
|dev: DeviceLocation| -> Vec<BackendId> { topology.backends_for(dev).to_vec() };
let capabilities_for = |b: BackendId| -> Option<&fuel_ir::backend::BackendCapabilities> {
topology.capabilities(b)
};
let resident_devices: Vec<DeviceLocation> = {
let mut v = Vec::new();
for arc in cache.values() {
if let Ok(guard) = arc.read()
&& let Some(loc) = cached_storage_location(&guard)
&& !v.contains(&loc)
{
v.push(loc);
}
}
v
};
let fallback_for = |dev: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
feasible_fallback_devices(topology.devices(), &resident_devices, dev)
.into_iter()
.flat_map(|d| topology.backends_for(d).iter().map(move |&b| (b, d)))
.collect()
};
let judge_oracle = crate::judge::cached_oracle();
let input_residency = |id: NodeId| -> Option<DeviceLocation> {
let slot = cache.get(&id)?;
let guard = slot.read().ok()?;
cached_storage_location(&guard)
};
#[cfg(feature = "telemetry")]
let tele_install = fuel_dispatch::telemetry::judge_sink::TelemetryInstall::new(pinned_device);
#[cfg(feature = "telemetry")]
let tele_hooks = tele_install.hooks();
let mut options = PlanOptions::new()
.with_placements_for_device(&placements_for)
.with_capabilities_for(&capabilities_for)
.with_pinned_device(pinned_device)
.with_fallback_placements_for(&fallback_for)
.with_transfer_estimator(&*topology)
.with_input_residency(&input_residency)
.with_precision_requirement(crate::reproducibility::precision_requirement());
if !allow_cost_placement {
options = options.without_cost_placement();
} else {
options = options.require_transfer_pricing();
}
if let Some(oracle) = judge_oracle.as_deref() {
options = options.with_judge(oracle);
}
#[cfg(feature = "telemetry")]
if let Some(ref hooks) = tele_hooks {
options = options.with_telemetry(hooks);
}
let bindings_guard = global_bindings();
let mut g = graph
.write()
.map_err(|_| Error::Msg("graph lock poisoned".into()).bt())?;
optimize_graph_with_runtime_fusion(&mut g, roots, &bindings_guard, &options)
}
fn dispatch_with_plan_retry(
graph: &Arc<RwLock<Graph>>,
cpu_target: NodeId,
cache: StorageCache,
device: &Device,
report_node: NodeId,
sym_env: &SymEnv,
allow_cost_placement: bool,
) -> Result<(Arc<RwLock<Storage>>, RootKernelIdent)> {
let pinned_loc = device.location();
let mut retry = TopologyRetryState::new();
loop {
let optimized = build_optimized_graph(
graph,
&[cpu_target],
pinned_loc,
&cache,
allow_cost_placement,
)?;
let (selector, lookup) = match production_selector_for(device) {
Some((s, l)) => (Some(s), Some(l)),
None => (None, None),
};
let mut cache_for_attempt = cache.clone();
seed_placed_device_handles(graph, cpu_target, pinned_loc, &mut cache_for_attempt)?;
let result = PipelinedExecutor::realize_with_optimized_picking_env(
graph.clone(),
cpu_target,
cache_for_attempt,
&optimized,
selector,
lookup,
sym_env.clone(),
);
match result {
Ok((storage, _layout)) => {
let dispatched = graph
.read()
.ok()
.and_then(|g| dispatched_kernel_ident(&g, report_node, &global_bindings()));
return Ok((storage, dispatched));
}
Err(e) if matches!(e, Error::TopologyChanged { .. }) && retry.permit_retry() => {
continue;
}
Err(e) => return Err(e),
}
}
}
pub fn realize_many_as_with_initial<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
targets: &[NodeId],
device: &Device,
initial: StorageCache,
) -> Result<Vec<Vec<T>>> {
realize_many_as_with_initial_env::<T>(graph, targets, device, initial, &SymEnv::default())
}
pub fn realize_many_as_with_initial_env<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
targets: &[NodeId],
device: &Device,
initial: StorageCache,
sym_env: &SymEnv,
) -> Result<Vec<Vec<T>>> {
if targets.is_empty() {
return Ok(Vec::new());
}
let (cache, _, effective_targets) = prepare(graph, targets, device, initial)?;
let results = dispatch_many_with_plan_retry(graph, &effective_targets, cache, device, sym_env)?;
let mut out = Vec::with_capacity(results.len());
for (storage, _layout) in results {
out.push(extract_cpu_bytes_typed::<T>(&storage)?);
}
Ok(out)
}
pub type SplitRealizeOutput<T> = (Vec<Vec<T>>, Vec<(Arc<RwLock<Storage>>, Layout)>);
pub fn realize_split_as_with_initial<T: bytemuck::Pod + WithDType>(
graph: &Arc<RwLock<Graph>>,
targets: &[NodeId],
n_host: usize,
device: &Device,
initial: StorageCache,
) -> Result<SplitRealizeOutput<T>> {
if n_host > targets.len() {
return Err(Error::Msg(format!(
"realize_split_as_with_initial: n_host ({n_host}) exceeds \
target count ({})",
targets.len(),
))
.bt());
}
if targets.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let (cache, _, effective_targets) = prepare_split(graph, targets, n_host, device, initial)?;
let results = dispatch_many_with_plan_retry(
graph,
&effective_targets,
cache,
device,
&SymEnv::default(),
)?;
let mut host_out = Vec::with_capacity(n_host);
let mut resident_out = Vec::with_capacity(results.len().saturating_sub(n_host));
for (i, (storage, layout)) in results.into_iter().enumerate() {
if i < n_host {
host_out.push(extract_cpu_bytes_typed::<T>(&storage)?);
} else {
resident_out.push((storage, layout));
}
}
Ok((host_out, resident_out))
}
fn dispatch_many_with_plan_retry(
graph: &Arc<RwLock<Graph>>,
effective_targets: &[NodeId],
cache: StorageCache,
device: &Device,
sym_env: &SymEnv,
) -> Result<Vec<(Arc<RwLock<Storage>>, Layout)>> {
let pinned_loc = device.location();
let mut retry = TopologyRetryState::new();
loop {
let optimized = build_optimized_graph(graph, effective_targets, pinned_loc, &cache, true)?;
let (selector, lookup) = match production_selector_for(device) {
Some((s, l)) => (Some(s), Some(l)),
None => (None, None),
};
let cache_for_attempt = cache.clone();
let result = PipelinedExecutor::realize_many_with_optimized_picking_env(
graph.clone(),
effective_targets,
cache_for_attempt,
&optimized,
selector,
lookup,
sym_env.clone(),
);
match result {
Ok(r) => break Ok(r),
Err(e) if matches!(e, Error::TopologyChanged { .. }) && retry.permit_retry() => {
continue;
}
Err(e) => break Err(e),
}
}
}
fn extract_cpu_bytes_typed<T: bytemuck::Pod + WithDType>(
storage: &Arc<RwLock<Storage>>,
) -> Result<Vec<T>> {
let guard = storage
.read()
.map_err(|_| Error::Msg("storage lock poisoned".into()).bt())?;
let root = guard.dtype();
if root != T::DTYPE {
return Err(Error::UnexpectedDType {
msg: "realize: root dtype does not match the requested element type \
(use realize_one_bytes* for a raw byte view)",
expected: T::DTYPE,
got: root,
}
.bt());
}
let bytes: &[u8] = match &guard.inner {
BackendStorage::Cpu(s) => s.bytes(),
#[allow(unreachable_patterns)]
other => {
return Err(Error::Msg(format!(
"pipelined_bridge: realize root produced non-CPU storage \
({other:?}) — the Op::Copy splice in `prepare()` should \
have made the root CPU-resident. This is a bug.",
))
.bt());
}
};
Ok(bytemuck::cast_slice::<u8, T>(bytes).to_vec())
}
fn extract_cpu_bytes_raw(storage: &Arc<RwLock<Storage>>) -> Result<Vec<u8>> {
let guard = storage
.read()
.map_err(|_| Error::Msg("storage lock poisoned".into()).bt())?;
match &guard.inner {
BackendStorage::Cpu(s) => Ok(s.bytes().to_vec()),
#[allow(unreachable_patterns)]
other => Err(Error::Msg(format!(
"pipelined_bridge: realize root produced non-CPU storage ({other:?}) — the \
Op::Copy splice in `prepare()` should have made the root CPU-resident. This is a bug.",
))
.bt()),
}
}
fn prepare(
graph: &Arc<RwLock<Graph>>,
targets: &[NodeId],
device: &Device,
initial: StorageCache,
) -> Result<(StorageCache, BackendId, Vec<NodeId>)> {
prepare_split(graph, targets, targets.len(), device, initial)
}
fn prepare_split(
graph: &Arc<RwLock<Graph>>,
targets: &[NodeId],
n_host: usize,
device: &Device,
initial: StorageCache,
) -> Result<(StorageCache, BackendId, Vec<NodeId>)> {
let backend_id = device_to_backend_id(device);
let effective_targets = {
let mut g = graph
.write()
.map_err(|_| Error::Msg("graph lock poisoned".into()).bt())?;
targets
.iter()
.enumerate()
.map(|(i, &src_id)| {
if i >= n_host {
return src_id;
}
let (shape, dtype) = {
let n = g.node(src_id);
(n.shape.clone(), n.dtype)
};
g.push(Node {
op: Op::Copy {
target: DeviceLocation::Cpu,
},
inputs: vec![src_id],
shape,
dtype,
})
})
.collect::<Vec<_>>()
};
let order = {
let g = graph
.read()
.map_err(|_| Error::Msg("graph lock poisoned".into()).bt())?;
topo_order_multi(&g, &effective_targets)
};
let cache = build_const_cache(graph, &order, device, initial)?;
Ok((cache, backend_id, effective_targets))
}
pub(crate) fn build_const_cache(
graph: &Arc<RwLock<Graph>>,
order: &[NodeId],
device: &Device,
initial: StorageCache,
) -> Result<StorageCache> {
let mut cache = initial;
cache.reserve(order.len() / 4);
let consts_to_upload: Vec<(NodeId, Vec<u8>, fuel_ir::DType)> = {
let g = graph
.read()
.map_err(|_| Error::Msg("graph lock poisoned".into()).bt())?;
let mut out: Vec<(NodeId, Vec<u8>, fuel_ir::DType)> = Vec::with_capacity(order.len() / 4);
for &id in order {
if cache.contains_key(&id) {
continue;
}
let node = g.node(id);
if !matches!(node.op, Op::Const) {
continue;
}
let slot_arc = g.storage_for(id).ok_or_else(|| {
Error::Msg(format!(
"pipelined_bridge: Op::Const node {id:?} has no \
storage in graph.storage_map (constructor failed \
to seed the slot)",
))
.bt()
})?;
let (host_buf, dtype) = {
let slot = slot_arc
.read()
.map_err(|_| Error::Msg("slot lock poisoned".into()).bt())?;
(slot.as_dyn().to_host_buffer_dyn()?, slot.dtype())
};
let need_bytes = node.shape.elem_count() * dtype.size_in_bytes();
let mut bytes = host_buffer_to_bytes(&host_buf);
if bytes.len() > need_bytes {
bytes.truncate(need_bytes);
}
out.push((id, bytes, dtype));
}
out
};
if consts_to_upload.is_empty() {
return Ok(cache);
}
let target_loc = device.location();
if target_loc == DeviceLocation::Cpu {
for (id, bytes, dtype) in consts_to_upload {
let storage = Storage::new(
BackendStorage::Cpu(fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(
&bytes,
)),
dtype,
);
cache.insert(id, Arc::new(RwLock::new(storage)));
}
return Ok(cache);
}
let transient = Arc::new(RwLock::new(Graph::new()));
let mut transient_cache = StorageCache::new();
if let Some(seed) = device_seed_storage(device)? {
let anchor_id = {
let mut g = transient
.write()
.map_err(|_| Error::Msg("transient graph lock poisoned".into()).bt())?;
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: fuel_ir::Shape::from_dims(&[4]),
dtype: fuel_ir::DType::U8,
})
};
transient_cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
}
let mut user_to_copy: Vec<(NodeId, NodeId)> = Vec::with_capacity(consts_to_upload.len());
{
let mut g = transient
.write()
.map_err(|_| Error::Msg("transient graph lock poisoned".into()).bt())?;
for (user_id, bytes, dtype) in consts_to_upload.into_iter() {
let n_elem = if dtype.size_in_bytes() == 0 {
0
} else {
bytes.len() / dtype.size_in_bytes()
};
let shape = fuel_ir::Shape::from_dims(&[n_elem]);
let trans_const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: shape.clone(),
dtype,
});
let cpu_storage = Storage::new(
BackendStorage::Cpu(fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(
&bytes,
)),
dtype,
);
transient_cache.insert(trans_const_id, Arc::new(RwLock::new(cpu_storage)));
let copy_id = g.push(Node {
op: Op::Copy { target: target_loc },
inputs: vec![trans_const_id],
shape,
dtype,
});
g.set_target_backend(copy_id, BackendId::Cpu);
user_to_copy.push((user_id, copy_id));
}
}
let copy_targets: Vec<NodeId> = user_to_copy.iter().map(|(_, c)| *c).collect();
let realized =
PipelinedExecutor::realize_many(Arc::clone(&transient), ©_targets, transient_cache)?;
if realized.len() != user_to_copy.len() {
return Err(Error::Msg(format!(
"build_const_cache: realize_many returned {} storages for {} \
Op::Copy targets — internal bug",
realized.len(),
user_to_copy.len(),
))
.bt());
}
for ((user_id, _), (arc, _layout)) in user_to_copy.into_iter().zip(realized) {
cache.insert(user_id, arc);
}
Ok(cache)
}
pub fn upload_host_buffer_to_device(
device: &Device,
buf: HostBuffer,
) -> Result<Arc<RwLock<Storage>>> {
let dtype = host_buffer_dtype(&buf);
let bytes = host_buffer_to_bytes(&buf);
let target_loc = device.location();
if target_loc == DeviceLocation::Cpu {
let storage = Storage::new(
BackendStorage::Cpu(fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(
&bytes,
)),
dtype,
);
return Ok(Arc::new(RwLock::new(storage)));
}
let transient = Arc::new(RwLock::new(Graph::new()));
let mut transient_cache = StorageCache::new();
if let Some(seed) = device_seed_storage(device)? {
let anchor_id = {
let mut g = transient
.write()
.map_err(|_| Error::Msg("transient graph lock poisoned".into()).bt())?;
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: fuel_ir::Shape::from_dims(&[4]),
dtype: fuel_ir::DType::U8,
})
};
transient_cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
}
let n_elem = if dtype.size_in_bytes() == 0 {
0
} else {
bytes.len() / dtype.size_in_bytes()
};
let shape = fuel_ir::Shape::from_dims(&[n_elem]);
let copy_id = {
let mut g = transient
.write()
.map_err(|_| Error::Msg("transient graph lock poisoned".into()).bt())?;
let trans_const_id = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: shape.clone(),
dtype,
});
let cpu_storage = Storage::new(
BackendStorage::Cpu(fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(
&bytes,
)),
dtype,
);
transient_cache.insert(trans_const_id, Arc::new(RwLock::new(cpu_storage)));
let copy_id = g.push(Node {
op: Op::Copy { target: target_loc },
inputs: vec![trans_const_id],
shape,
dtype,
});
g.set_target_backend(copy_id, BackendId::Cpu);
copy_id
};
let realized =
PipelinedExecutor::realize_many(Arc::clone(&transient), &[copy_id], transient_cache)?;
let (arc, _layout) = realized.into_iter().next().ok_or_else(|| {
Error::Msg(
"upload_host_buffer_to_device: realize_many returned no storage \
for the Op::Copy target — internal bug"
.into(),
)
.bt()
})?;
Ok(arc)
}
fn host_buffer_dtype(buf: &HostBuffer) -> fuel_ir::DType {
use fuel_ir::DType;
match buf {
HostBuffer::U8(_) => DType::U8,
HostBuffer::I8(_) => DType::I8,
HostBuffer::U32(_) => DType::U32,
HostBuffer::I16(_) => DType::I16,
HostBuffer::I32(_) => DType::I32,
HostBuffer::I64(_) => DType::I64,
HostBuffer::BF16(_) => DType::BF16,
HostBuffer::F16(_) => DType::F16,
HostBuffer::F32(_) => DType::F32,
HostBuffer::F64(_) => DType::F64,
HostBuffer::F8E4M3(_) => DType::F8E4M3,
HostBuffer::F6E2M3(_) => DType::F6E2M3,
HostBuffer::F6E3M2(_) => DType::F6E3M2,
HostBuffer::F4(_) => DType::F4,
HostBuffer::F8E8M0(_) => DType::F8E8M0,
HostBuffer::Bool(_) => DType::Bool,
}
}
pub(crate) fn host_buffer_to_bytes(buf: &HostBuffer) -> Vec<u8> {
match buf {
HostBuffer::U8(v) => v.clone(),
HostBuffer::Bool(v) => v.clone(),
HostBuffer::I8(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::U32(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::I16(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::I32(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::I64(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::BF16(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::F16(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::F32(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::F64(v) => bytemuck::cast_slice(v).to_vec(),
HostBuffer::F8E4M3(v) => {
let bytes: &[u8] = unsafe {
std::slice::from_raw_parts(
v.as_ptr() as *const u8,
v.len() * std::mem::size_of::<float8::F8E4M3>(),
)
};
bytes.to_vec()
}
HostBuffer::F6E2M3(v) => v.clone(),
HostBuffer::F6E3M2(v) => v.clone(),
HostBuffer::F4(v) => v.clone(),
HostBuffer::F8E8M0(v) => v.clone(),
}
}
fn device_to_backend_id(device: &Device) -> BackendId {
location_to_backend_id(device.location())
}
fn location_to_backend_id(loc: DeviceLocation) -> BackendId {
match loc {
DeviceLocation::Cpu => BackendId::Cpu,
DeviceLocation::Cuda { .. } => BackendId::Cuda,
DeviceLocation::Vulkan { .. } => BackendId::Vulkan,
DeviceLocation::Metal { .. } => BackendId::Metal,
}
}
fn runtime_selector_disabled() -> bool {
std::env::var("FUEL_DISABLE_RUNTIME_SELECTOR")
.ok()
.as_deref()
== Some("1")
}
fn production_selector_for(
device: &Device,
) -> Option<(Arc<dyn RuntimeSelector>, BackendRuntimeLookup)> {
if runtime_selector_disabled() {
return None;
}
let judge: Option<Arc<dyn JudgeOracle>> =
crate::judge::cached_oracle().map(|oracle| oracle as Arc<dyn JudgeOracle>);
let lookup = backend_runtime_lookup_for(device);
let selector: Arc<dyn RuntimeSelector> = Arc::new(ChainedSelector::with_default_estimator(
judge,
Some(lookup.clone()),
));
Some((selector, lookup))
}
struct DeviceRuntimeHandle(Arc<dyn DynBackendDevice>);
impl BackendRuntime for DeviceRuntimeHandle {
fn available_bytes(&self) -> Option<u64> {
self.0
.as_backend_runtime()
.and_then(|r| r.available_bytes())
}
fn total_bytes(&self) -> Option<u64> {
self.0.as_backend_runtime().and_then(|r| r.total_bytes())
}
fn would_fit(&self, size: u64) -> FitStatus {
match self.0.as_backend_runtime() {
Some(r) => r.would_fit(size),
None => FitStatus::Unknown,
}
}
fn as_backend_streams(&self) -> Option<&dyn BackendStreams> {
match self.0.location_dyn() {
DeviceLocation::Cpu => None,
_ => Some(self),
}
}
}
impl BackendStreams for DeviceRuntimeHandle {
fn pending_work_count(&self) -> Option<u32> {
match self.0.location_dyn() {
DeviceLocation::Cpu => None,
loc => Some(fuel_dispatch::dispatch::inflight_count(loc)),
}
}
fn slot_capacity(&self) -> u32 {
1
}
fn flush(&self) -> Result<()> {
self.0.synchronize_dyn()
}
}
fn backend_runtime_lookup_for(device: &Device) -> BackendRuntimeLookup {
let realize_loc = device.location();
let realize_backend = location_to_backend_id(realize_loc);
let inner: Arc<dyn DynBackendDevice> = device.inner.clone();
Arc::new(move |backend, loc| {
if backend == realize_backend && loc == realize_loc {
return Some(Box::new(DeviceRuntimeHandle(Arc::clone(&inner))) as BackendRuntimeHandle);
}
if backend == BackendId::Cpu && loc == DeviceLocation::Cpu {
return Some(
Box::new(fuel_cpu_backend::dyn_impl::CpuBackendDevice) as BackendRuntimeHandle
);
}
None
})
}
fn cached_storage_location(storage: &Storage) -> Option<DeviceLocation> {
match &storage.inner {
BackendStorage::Cpu(_) => Some(DeviceLocation::Cpu),
#[cfg(feature = "cuda")]
BackendStorage::Cuda(c) => Some(c.device().location()),
#[cfg(feature = "vulkan")]
BackendStorage::Vulkan(v) => v
.backend()
.map(|b| DeviceLocation::Vulkan { gpu_id: b.gpu_id }),
#[cfg(all(feature = "metal", any(target_os = "macos", target_os = "ios")))]
BackendStorage::Metal(_) => None,
}
}
fn placed_device_locations_needing_seed(
graph: &Arc<RwLock<Graph>>,
root: NodeId,
pinned_loc: DeviceLocation,
cache: &StorageCache,
) -> Result<Vec<DeviceLocation>> {
let already: Vec<DeviceLocation> = cache
.values()
.filter_map(|arc| {
let guard = arc.read().ok()?;
cached_storage_location(&guard)
})
.collect();
let g = graph
.read()
.map_err(|_| Error::Msg("graph lock poisoned during device-seed scan".into()).bt())?;
let mut needed: Vec<DeviceLocation> = Vec::new();
for id in topo_order_multi(&g, &[root]) {
let loc = match g.node(id).op {
Op::Copy { target } | Op::Move { target } | Op::Alloc { target } => target,
_ => continue,
};
if loc == DeviceLocation::Cpu || loc == pinned_loc {
continue;
}
if already.contains(&loc) || needed.contains(&loc) {
continue;
}
needed.push(loc);
}
Ok(needed)
}
fn device_for_location(loc: DeviceLocation) -> Result<Device> {
match loc {
DeviceLocation::Cpu => Ok(Device::cpu()),
DeviceLocation::Cuda { gpu_id } => crate::cuda_backend::new_device(gpu_id),
DeviceLocation::Vulkan { gpu_id } => Err(Error::Msg(format!(
"post-optimize device seed: the plan placed nodes on Vulkan \
{{ gpu_id: {gpu_id} }} but this single-device realize has no \
Vulkan handle to seed, and a Vulkan device can't be derived from \
an ordinal alone (it needs an explicit DeviceSelection). Realize \
the mixed graph via `realize_one_as_multi_device` with the Vulkan \
device in `extra_devices`.",
))
.bt()),
other => Err(Error::Msg(format!(
"post-optimize device seed: target {other:?} not wired (CPU + CUDA \
auto-seed today; Vulkan via realize_one_as_multi_device).",
))
.bt()),
}
}
fn seed_placed_device_handles(
graph: &Arc<RwLock<Graph>>,
root: NodeId,
pinned_loc: DeviceLocation,
cache: &mut StorageCache,
) -> Result<()> {
let needed = placed_device_locations_needing_seed(graph, root, pinned_loc, cache)?;
for loc in needed {
let dev = device_for_location(loc)?;
let Some(seed) = device_seed_storage(&dev)? else {
continue; };
let anchor_id = {
let mut g = graph.write().map_err(|_| {
Error::Msg("graph lock poisoned during placed-device seed".into()).bt()
})?;
g.push(Node {
op: Op::Const,
inputs: vec![],
shape: fuel_ir::Shape::from_dims(&[4]),
dtype: fuel_ir::DType::U8,
})
};
cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
}
Ok(())
}
pub fn device_seed_storage(device: &Device) -> Result<Option<Storage>> {
#[cfg(any(feature = "cuda", feature = "vulkan"))]
const SEED_BYTES: usize = 4;
match device.location() {
DeviceLocation::Cpu => Ok(None),
#[cfg(feature = "cuda")]
DeviceLocation::Cuda { .. } => {
let cuda_dev = crate::cuda_backend::as_device(device)?;
let cuda_bytes = fuel_cuda_backend::CudaStorageBytes::alloc(cuda_dev, SEED_BYTES)?;
Ok(Some(Storage::new(
BackendStorage::Cuda(cuda_bytes),
fuel_ir::DType::U8,
)))
}
#[cfg(not(feature = "cuda"))]
DeviceLocation::Cuda { .. } => Err(Error::Msg(
"device_seed_storage: CUDA device requested but fuel-core wasn't built \
with --features cuda"
.into(),
)
.bt()),
#[cfg(feature = "vulkan")]
DeviceLocation::Vulkan { .. } => {
let backend = crate::vulkan_backend::as_device(device)?;
let zeros = vec![0_u8; SEED_BYTES];
let vk_bytes = backend.upload_bytes_handle(&zeros)?;
Ok(Some(Storage::new(
BackendStorage::Vulkan(vk_bytes),
fuel_ir::DType::U8,
)))
}
#[cfg(not(feature = "vulkan"))]
DeviceLocation::Vulkan { .. } => Err(Error::Msg(
"device_seed_storage: Vulkan device requested but fuel-core wasn't built \
with --features vulkan"
.into(),
)
.bt()),
other => Err(Error::Msg(format!(
"device_seed_storage: device {other:?} not wired (CPU + CUDA + Vulkan \
today; Metal pending its byte-storage substrate)",
))
.bt()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_dispatch::plan::compile_plan;
use fuel_ir::{DType, Shape};
#[test]
fn feasible_fallback_devices_scopes_to_resident() {
let cpu = DeviceLocation::Cpu;
let cuda = DeviceLocation::Cuda { gpu_id: 0 };
let all = [cpu, cuda];
assert_eq!(
feasible_fallback_devices(&all, &[cpu], cpu),
Vec::<DeviceLocation>::new(),
"cpu-only realize must not offer CUDA as a fallback (the spill fix)",
);
assert_eq!(
feasible_fallback_devices(&all, &[cpu, cuda], cpu),
vec![cuda],
"cuda-seeded realize must still offer CUDA (positive control — not amputated)",
);
assert_eq!(
feasible_fallback_devices(&all, &[cpu, cuda], cuda),
vec![cpu]
);
assert_eq!(
feasible_fallback_devices(&all, &[cuda], cuda),
vec![cpu],
"cuda-only realize must still offer CPU fallback (conv2d-with-no-cuda-kernel case)",
);
}
fn push_node(g: &mut Graph, op: Op, inputs: Vec<NodeId>) -> NodeId {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(&[4]),
dtype: DType::F32,
})
}
fn cpu_storage_f32(vals: &[f32]) -> Arc<RwLock<Storage>> {
let bytes: &[u8] = bytemuck::cast_slice(vals);
Arc::new(RwLock::new(Storage::new(
BackendStorage::Cpu(fuel_cpu_backend::byte_storage::CpuStorageBytes::from_bytes(
bytes,
)),
DType::F32,
)))
}
fn noop_kernel_for_tests(
_i: &[Arc<RwLock<Storage>>],
_o: &mut [Arc<RwLock<Storage>>],
_l: &[fuel_ir::Layout],
_p: &fuel_dispatch::kernel::OpParams,
) -> Result<()> {
Ok(())
}
#[test]
fn option_b_detects_offloaded_device_only_when_cpu_pinned() {
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
let mut g = Graph::new();
let c = push_node(&mut g, Op::Const, vec![]);
let copy = push_node(&mut g, Op::Copy { target: cuda0 }, vec![c]);
let graph = Arc::new(RwLock::new(g));
let empty = StorageCache::new();
assert_eq!(
placed_device_locations_needing_seed(&graph, copy, DeviceLocation::Cpu, &empty)
.unwrap(),
vec![cuda0],
"CPU-pinned realize with a Copy→Cuda must flag Cuda for seeding",
);
assert!(
placed_device_locations_needing_seed(&graph, copy, cuda0, &empty)
.unwrap()
.is_empty(),
"a Copy to the pinned device needs no extra seed",
);
let mut g2 = Graph::new();
let a = push_node(&mut g2, Op::Const, vec![]);
let cpu_copy = push_node(
&mut g2,
Op::Copy {
target: DeviceLocation::Cpu,
},
vec![a],
);
let graph2 = Arc::new(RwLock::new(g2));
assert!(
placed_device_locations_needing_seed(&graph2, cpu_copy, DeviceLocation::Cpu, &empty)
.unwrap()
.is_empty(),
"an all-CPU graph never needs a device seed",
);
}
#[test]
fn cpu_runtime_handle_is_not_backend_streams() {
let handle = DeviceRuntimeHandle(
Arc::new(fuel_cpu_backend::dyn_impl::CpuBackendDevice) as Arc<dyn DynBackendDevice>
);
let as_runtime: &dyn BackendRuntime = &handle;
assert!(
as_runtime.as_backend_streams().is_none(),
"CPU has no queue concept — it must not expose BackendStreams",
);
assert_eq!(
handle.pending_work_count(),
None,
"CPU pending_work_count must be None (no queue), never Some(0)",
);
}
#[test]
fn stale_stamps_overwritten_per_realize() {
let graph = Arc::new(RwLock::new(Graph::new()));
let (c1, c2, add) = {
let mut g = graph.write().unwrap();
let c1 = push_node(&mut g, Op::Const, vec![]);
let c2 = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![c1, c2]);
g.set_target_backend(add, BackendId::Cuda);
(c1, c2, add)
};
let mut initial = StorageCache::new();
initial.insert(c1, cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]));
initial.insert(c2, cpu_storage_f32(&[10.0, 20.0, 30.0, 40.0]));
let device = crate::Device::cpu();
let out = realize_one_as_with_initial::<f32>(&graph, add, &device, initial)
.expect("CPU realize despite stale CUDA stamp");
assert_eq!(out, vec![11.0, 22.0, 33.0, 44.0]);
assert_eq!(
graph.read().unwrap().target_backend(add),
Some(BackendId::Cpu),
"stale stamp overwritten by this call's plan winner",
);
}
#[test]
fn stage2_cpu_only_estimator_leaves_plan_unchanged() {
use fuel_dispatch::kernel::{KernelBindingTable, KernelCaps, unknown_cost};
use fuel_ir::dispatch::OpKind;
let mut table = KernelBindingTable::new();
table.register_full(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
noop_kernel_for_tests,
KernelCaps::empty(),
fuel_dispatch::fused::PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
unknown_cost,
);
let mut g = Graph::new();
let c1 = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![c1, c1]);
let order = topo_order_multi(&g, &[add]);
let mut cache = StorageCache::new();
cache.insert(c1, cpu_storage_f32(&[1.0; 4]));
let topology = SystemTopology::current();
let placements_for =
|dev: DeviceLocation| -> Vec<BackendId> { topology.backends_for(dev).to_vec() };
let capabilities_for = |b: BackendId| -> Option<&fuel_ir::backend::BackendCapabilities> {
topology.capabilities(b)
};
let input_residency = |id: NodeId| -> Option<DeviceLocation> {
let slot = cache.get(&id)?;
let guard = slot.read().ok()?;
cached_storage_location(&guard)
};
let base_opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_for)
.with_capabilities_for(&capabilities_for);
let base = compile_plan(&g, &order, &table, &base_opts).expect("base plan");
let wired_opts = PlanOptions::new()
.with_pinned_device(DeviceLocation::Cpu)
.with_placements_for_device(&placements_for)
.with_capabilities_for(&capabilities_for)
.with_transfer_estimator(&*topology)
.with_input_residency(&input_residency);
let wired = compile_plan(&g, &order, &table, &wired_opts).expect("wired plan");
let a = base.alternatives(add).expect("base set");
let b = wired.alternatives(add).expect("wired set");
assert_eq!(a.len(), b.len(), "same candidate count");
for (ca, cb) in a.alternatives().iter().zip(b.alternatives()) {
assert_eq!(ca.backend, cb.backend);
assert_eq!(ca.device, cb.device);
assert_eq!(ca.kernel as usize, cb.kernel as usize, "same kernel ref");
assert_eq!(ca.kernel_source, cb.kernel_source);
assert_eq!(
cb.inbound_transfer_ns, 0,
"CPU-only host: zero transfer terms",
);
}
}
#[test]
fn optimized_cpu_realize_end_to_end() {
let graph = Arc::new(RwLock::new(Graph::new()));
let (c1, c2, add) = {
let mut g = graph.write().unwrap();
let c1 = push_node(&mut g, Op::Const, vec![]);
let c2 = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![c1, c2]);
(c1, c2, add)
};
let mut initial = StorageCache::new();
initial.insert(c1, cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]));
initial.insert(c2, cpu_storage_f32(&[10.0, 20.0, 30.0, 40.0]));
let device = crate::Device::cpu();
assert!(
production_selector_for(&device).is_some(),
"the runtime route picker defaults ON (no opt-out env set)",
);
let out = realize_one_as_with_initial::<f32>(&graph, add, &device, initial)
.expect("realize through the optimized + route-picker path");
assert_eq!(out, vec![11.0, 22.0, 33.0, 44.0]);
}
#[test]
fn d2a_prebuilt_realize_skips_optimize_and_does_not_grow_graph() {
let graph = Arc::new(RwLock::new(Graph::new()));
let (c1, c2, add) = {
let mut g = graph.write().unwrap();
let c1 = push_node(&mut g, Op::Const, vec![]);
let c2 = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![c1, c2]);
(c1, c2, add)
};
let c1_arc = cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]);
let c2_arc = cpu_storage_f32(&[10.0, 20.0, 30.0, 40.0]);
let device = crate::Device::cpu();
let mut initial = StorageCache::new();
initial.insert(c1, Arc::clone(&c1_arc));
initial.insert(c2, Arc::clone(&c2_arc));
let calls_before = optimize_calls_thread_local();
let (effective_target, optimized, first) =
prebuild_optimized_env::<f32>(&graph, add, &device, initial, &SymEnv::default())
.expect("prebuild (first) realize");
assert_eq!(first, vec![11.0, 22.0, 33.0, 44.0]);
let calls_after_first = optimize_calls_thread_local();
assert!(
calls_after_first > calls_before,
"the FIRST realize must run the optimizer (prepare + optimize_graph): \
{calls_before} -> {calls_after_first}",
);
let len_after_first = graph.read().unwrap().len();
let mut cache2 = StorageCache::new();
cache2.insert(c1, Arc::clone(&c1_arc));
cache2.insert(c2, Arc::clone(&c2_arc));
let second = realize_one_prebuilt_env::<f32>(
&graph,
effective_target,
&optimized,
&device,
cache2,
&SymEnv::default(),
)
.expect("prebuilt (second) realize");
let calls_after_second = optimize_calls_thread_local();
let len_after_second = graph.read().unwrap().len();
assert_eq!(
calls_after_second, calls_after_first,
"prebuilt realize must NOT re-run optimize_graph: \
{calls_after_first} -> {calls_after_second}",
);
assert_eq!(
second, first,
"prebuilt realize must reproduce the first result byte-for-byte",
);
assert_eq!(
len_after_second, len_after_first,
"prebuilt realize must NOT grow the graph (no re-spliced D2H Copy \
/ re-inserted residency or Contiguize)",
);
let mut initial3 = StorageCache::new();
initial3.insert(c1, Arc::clone(&c1_arc));
initial3.insert(c2, Arc::clone(&c2_arc));
let third = realize_one_as_with_initial::<f32>(&graph, add, &device, initial3)
.expect("full-path (third) realize");
assert_eq!(third, first, "full path still computes the same value");
assert!(
optimize_calls_thread_local() > calls_after_second,
"the FULL path re-optimizes (this is what the prebuilt seam skips)",
);
assert!(
graph.read().unwrap().len() > len_after_second,
"the FULL path re-splices a D2H Op::Copy at the root, growing the \
graph (this is the double-splice the prebuilt seam skips)",
);
}
#[test]
fn production_selector_for_returns_selector_and_lookup() {
let device = crate::Device::cpu();
let built = production_selector_for(&device);
assert!(built.is_some(), "picker defaults ON without the opt-out");
let (_selector, lookup) = built.unwrap();
assert!(
lookup(BackendId::Cpu, DeviceLocation::Cpu).is_some(),
"the lookup resolves the always-present CPU backend",
);
}
#[test]
fn backend_runtime_lookup_resolves_cpu_and_misses_others() {
let device = crate::Device::cpu();
let lookup = backend_runtime_lookup_for(&device);
let cpu =
lookup(BackendId::Cpu, DeviceLocation::Cpu).expect("CPU handle always resolvable");
let _ = cpu.would_fit(1);
assert!(
lookup(BackendId::Cuda, DeviceLocation::Cuda { gpu_id: 0 }).is_none(),
"no live handle for a backend that isn't the realize device",
);
}
#[test]
fn reporting_realize_returns_root_kernel_source() {
use fuel_ir::dispatch::OpKind;
let graph = Arc::new(RwLock::new(Graph::new()));
let (c1, c2, add) = {
let mut g = graph.write().unwrap();
let c1 = push_node(&mut g, Op::Const, vec![]);
let c2 = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![c1, c2]);
(c1, c2, add)
};
let mut initial = StorageCache::new();
initial.insert(c1, cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]));
initial.insert(c2, cpu_storage_f32(&[10.0, 20.0, 30.0, 40.0]));
let device = crate::Device::cpu();
let (out, root_kernel_source) = realize_one_as_with_initial_reporting::<f32>(
&graph,
add,
&device,
initial,
&SymEnv::default(),
)
.expect("reporting realize");
assert_eq!(out, vec![11.0, 22.0, 33.0, 44.0]);
let (src, rev) = root_kernel_source
.expect("plan covers the Add root → dispatched-sibling report present");
let bindings = global_bindings();
let alts = bindings.lookup_alternatives(
OpKind::AddElementwise,
&[DType::F32, DType::F32, DType::F32],
BackendId::Cpu,
);
assert!(
alts.iter()
.any(|e| e.kernel_source == src && e.kernel_revision_hash == rev),
"reported (kernel_source {src:?}, revision {rev:#x}) must be a \
registered sibling at the (AddElementwise, [F32;3], Cpu) cell — \
the tag and the revision come from ONE find() so they agree",
);
}
#[test]
fn realize_split_host_and_resident() {
let graph = Arc::new(RwLock::new(Graph::new()));
let (c1, c2, add, mul) = {
let mut g = graph.write().unwrap();
let c1 = push_node(&mut g, Op::Const, vec![]);
let c2 = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![c1, c2]);
let mul = push_node(&mut g, Op::Mul, vec![c1, c2]);
(c1, c2, add, mul)
};
let mut cache = StorageCache::new();
cache.insert(c1, cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]));
cache.insert(c2, cpu_storage_f32(&[10.0, 20.0, 30.0, 40.0]));
let (host, resident) = realize_split_as_with_initial::<f32>(
&graph,
&[add, mul],
1,
&crate::Device::cpu(),
cache,
)
.expect("realize_split");
assert_eq!(host.len(), 1, "exactly n_host host results");
assert_eq!(host[0], vec![11.0, 22.0, 33.0, 44.0]);
assert_eq!(resident.len(), 1, "remaining targets stay resident");
let (storage, layout) = &resident[0];
assert_eq!(layout.shape().dims(), &[4]);
let guard = storage.read().unwrap();
match &guard.inner {
BackendStorage::Cpu(c) => {
let vals: &[f32] = bytemuck::cast_slice(c.bytes());
assert_eq!(vals, &[10.0, 40.0, 90.0, 160.0]);
}
#[allow(unreachable_patterns)]
other => panic!("expected CPU storage, got {other:?}"),
}
}
#[test]
fn realize_split_resident_feeds_next_step() {
let g1 = Arc::new(RwLock::new(Graph::new()));
let (c1, dbl) = {
let mut g = g1.write().unwrap();
let c1 = push_node(&mut g, Op::Const, vec![]);
let dbl = push_node(&mut g, Op::Add, vec![c1, c1]);
(c1, dbl)
};
let mut cache1 = StorageCache::new();
cache1.insert(c1, cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]));
let (host1, resident1) =
realize_split_as_with_initial::<f32>(&g1, &[dbl], 0, &crate::Device::cpu(), cache1)
.expect("step 1");
assert!(host1.is_empty());
let (carried, _) = resident1.into_iter().next().expect("one resident");
let g2 = Arc::new(RwLock::new(Graph::new()));
let (ph, sq) = {
let mut g = g2.write().unwrap();
let ph = push_node(&mut g, Op::Const, vec![]);
let sq = push_node(&mut g, Op::Sqr, vec![ph]);
(ph, sq)
};
let mut cache2 = StorageCache::new();
cache2.insert(ph, carried);
let (host2, resident2) =
realize_split_as_with_initial::<f32>(&g2, &[sq], 1, &crate::Device::cpu(), cache2)
.expect("step 2");
assert_eq!(host2[0], vec![4.0, 16.0, 36.0, 64.0]);
assert!(resident2.is_empty());
}
#[test]
fn optimized_path_runs_full_correctness_sequence() {
let expected = vec![11.0_f32, 44.0, 99.0, 176.0];
let a_vals = [1.0_f32, 2.0, 3.0, 4.0];
let b_vals = [10.0_f32, 20.0, 30.0, 40.0];
let graph = Arc::new(RwLock::new(Graph::new()));
let (a, b, _add, out) = {
let mut g = graph.write().unwrap();
let a = push_node(&mut g, Op::Const, vec![]);
let b = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![a, b]);
let out = push_node(&mut g, Op::Mul, vec![add, a]);
(a, b, add, out)
};
let device = crate::Device::cpu();
let pinned = device.location();
let mut initial = StorageCache::new();
initial.insert(a, cpu_storage_f32(&a_vals));
initial.insert(b, cpu_storage_f32(&b_vals));
let (cache, _backend, mut eff) =
prepare(&graph, &[out], &device, initial).expect("prepare");
let cpu_target = eff.pop().expect("one effective target");
let optimized = build_optimized_graph(&graph, &[cpu_target], pinned, &cache, true)
.expect("optimize_graph");
let (storage, _layout) =
PipelinedExecutor::realize_with_optimized(graph.clone(), cpu_target, cache, &optimized)
.expect("optimized realize");
let out_vals = extract_cpu_bytes_typed::<f32>(&storage).expect("host bytes");
assert_eq!(out_vals, expected, "optimized path: (a+b)*a");
}
#[test]
fn default_realize_entry_point_produces_correct_values() {
let graph = Arc::new(RwLock::new(Graph::new()));
let (a, b, add, out) = {
let mut g = graph.write().unwrap();
let a = push_node(&mut g, Op::Const, vec![]);
let b = push_node(&mut g, Op::Const, vec![]);
let add = push_node(&mut g, Op::Add, vec![a, b]);
let out = push_node(&mut g, Op::Mul, vec![add, a]);
(a, b, add, out)
};
let _ = add;
let mut initial = StorageCache::new();
initial.insert(a, cpu_storage_f32(&[1.0, 2.0, 3.0, 4.0]));
initial.insert(b, cpu_storage_f32(&[10.0, 20.0, 30.0, 40.0]));
let out_vals =
realize_one_as_with_initial::<f32>(&graph, out, &crate::Device::cpu(), initial)
.expect("default realize of (a+b)*a");
assert_eq!(out_vals, vec![11.0, 44.0, 99.0, 176.0]);
}
#[test]
fn realize_split_n_host_too_large_errors() {
let graph = Arc::new(RwLock::new(Graph::new()));
let c1 = {
let mut g = graph.write().unwrap();
push_node(&mut g, Op::Const, vec![])
};
let err = realize_split_as_with_initial::<f32>(
&graph,
&[c1],
2,
&crate::Device::cpu(),
StorageCache::new(),
);
assert!(err.is_err(), "n_host > targets.len() must be a typed Err");
}
}
#[cfg(test)]
mod gap327_funnel_check {
#[test]
fn cast_slice_u8_is_a_single_guarded_site() {
let src = include_str!("pipelined_bridge.rs");
let call = concat!("cast_slice::<u8, ", "T>(");
let n = src
.lines()
.filter(|l| !l.trim_start().starts_with("//") && l.contains(call))
.count();
assert_eq!(
n, 1,
"the generic byte-reinterpret funnel must occur at exactly ONE code site; found {n}"
);
let guard = "!= T::DTYPE";
let cast_pos = src.find(call).expect("the funnel cast site");
let guard_pos = src.find(guard).expect("the dtype guard");
assert!(
guard_pos < cast_pos,
"the dtype guard must precede the byte-reinterpret funnel"
);
}
}