use std::collections::{HashMap, HashSet};
use std::sync::{Arc, OnceLock, RwLock};
use crate::dispatch::{global_bindings, global_registry, topology_generation};
use fuel_ir::backend::{BackendCapabilities, SubstrateClass, TransferPath};
use fuel_ir::dispatch::OpKind;
use fuel_ir::probe::BackendId;
use fuel_ir::{DType, DeviceLocation};
use fuel_hardware::probe::ProbeReport;
use fuel_hardware::transfer_cost::{TransferCalibration, TransferEstimate};
static CURRENT_TOPOLOGY: RwLock<Option<Arc<SystemTopology>>> = RwLock::new(None);
#[derive(Debug)]
pub struct SystemTopology {
generation: u64,
devices: Vec<DeviceLocation>,
backends: Vec<BackendId>,
backends_for_device: HashMap<DeviceLocation, Vec<BackendId>>,
devices_for_backend: HashMap<BackendId, Vec<DeviceLocation>>,
substrate_for: HashMap<BackendId, SubstrateClass>,
capabilities: HashMap<BackendId, BackendCapabilities>,
transfer_paths: HashMap<(DeviceLocation, DeviceLocation), TransferPath>,
binding_op_coverage: HashMap<BackendId, HashSet<(OpKind, DType)>>,
transfer_calibration: OnceLock<TransferCalibration>,
}
impl SystemTopology {
pub fn current() -> Arc<SystemTopology> {
let cur_gen = topology_generation();
if let Some(t) = CURRENT_TOPOLOGY.read().unwrap().as_ref()
&& t.generation == cur_gen
{
return Arc::clone(t);
}
let fresh = Arc::new(SystemTopology::build_at(cur_gen));
let mut guard = CURRENT_TOPOLOGY.write().unwrap();
match guard.as_ref() {
Some(existing) if existing.generation >= fresh.generation => Arc::clone(existing),
_ => {
*guard = Some(Arc::clone(&fresh));
fresh
}
}
}
pub fn refresh() {
crate::dispatch::bump_topology_generation();
}
pub fn generation(&self) -> u64 {
self.generation
}
pub fn devices(&self) -> &[DeviceLocation] {
&self.devices
}
pub fn backends(&self) -> &[BackendId] {
&self.backends
}
pub fn backends_for(&self, dev: DeviceLocation) -> &[BackendId] {
self.backends_for_device
.get(&dev)
.map(Vec::as_slice)
.unwrap_or(&[])
}
pub fn devices_for(&self, backend: BackendId) -> &[DeviceLocation] {
self.devices_for_backend
.get(&backend)
.map(Vec::as_slice)
.unwrap_or(&[])
}
pub fn substrate_class(
&self,
backend: BackendId,
_device: DeviceLocation,
) -> Option<SubstrateClass> {
self.substrate_for.get(&backend).copied()
}
pub fn shares_storage(
&self,
a: (BackendId, DeviceLocation),
b: (BackendId, DeviceLocation),
) -> bool {
let (a_backend, a_device) = a;
let (b_backend, b_device) = b;
if a_device != b_device {
return false;
}
match (
self.substrate_for.get(&a_backend),
self.substrate_for.get(&b_backend),
) {
(Some(sa), Some(sb)) => sa == sb,
_ => false,
}
}
pub fn transfer_path(&self, src: DeviceLocation, dst: DeviceLocation) -> TransferPath {
if src == dst {
return TransferPath::SameDevice;
}
self.transfer_paths
.get(&(src, dst))
.copied()
.unwrap_or(TransferPath::HostStaging)
}
pub fn transfer_calibration(&self) -> &TransferCalibration {
self.transfer_calibration
.get_or_init(|| TransferCalibration::calibrate(&self.devices))
}
pub fn transfer_estimate(&self, src: DeviceLocation, dst: DeviceLocation) -> TransferEstimate {
if src == dst {
return TransferEstimate::fallback_for(TransferPath::SameDevice);
}
let cal = self.transfer_calibration();
if let Some(e) = cal.probed(src, dst) {
return e;
}
if src != DeviceLocation::Cpu
&& dst != DeviceLocation::Cpu
&& let (Some(d2h), Some(h2d)) = (
cal.probed(src, DeviceLocation::Cpu),
cal.probed(DeviceLocation::Cpu, dst),
)
{
return TransferEstimate::compose_staged(d2h, h2d);
}
TransferEstimate::fallback_for(self.transfer_path(src, dst))
}
pub fn estimate_transfer_ns(
&self,
src: DeviceLocation,
dst: DeviceLocation,
bytes: u64,
) -> u64 {
if src == dst {
return 0;
}
self.transfer_estimate(src, dst).estimate_ns(bytes)
}
pub fn capabilities(&self, backend: BackendId) -> Option<&BackendCapabilities> {
self.capabilities.get(&backend)
}
pub fn binding_op_coverage(&self, backend: BackendId) -> &HashSet<(OpKind, DType)> {
static EMPTY: std::sync::OnceLock<HashSet<(OpKind, DType)>> = std::sync::OnceLock::new();
self.binding_op_coverage
.get(&backend)
.unwrap_or_else(|| EMPTY.get_or_init(HashSet::new))
}
pub fn capabilities_op_coverage_divergence(&self) -> Vec<(BackendId, OpKind, DType)> {
let mut missing = Vec::new();
for (backend, caps) in &self.capabilities {
let registered = self.binding_op_coverage(*backend);
for &(op, dtype) in &caps.op_dtype_support {
if !registered.contains(&(op, dtype)) {
missing.push((*backend, op, dtype));
}
}
}
missing.sort_by(|a, b| {
a.0.as_str()
.cmp(b.0.as_str())
.then_with(|| format!("{:?}", a.1).cmp(&format!("{:?}", b.1)))
});
missing
}
fn build_at(_caller_gen: u64) -> SystemTopology {
let built_gen = topology_generation();
let probe = ProbeReport::probe_all();
let mut capabilities: HashMap<BackendId, BackendCapabilities> = HashMap::new();
{
let registry = global_registry();
for caps in registry.backends() {
capabilities
.entry(caps.backend_id)
.or_insert_with(|| caps.clone());
}
}
let mut binding_op_coverage: HashMap<BackendId, HashSet<(OpKind, DType)>> = HashMap::new();
{
let bindings = global_bindings();
for (op, dtypes, backend) in bindings.iter_keys() {
if let Some(&output_dt) = dtypes.last() {
binding_op_coverage
.entry(backend)
.or_default()
.insert((op, output_dt));
}
}
}
let mut backends: Vec<BackendId> = capabilities
.keys()
.copied()
.chain(binding_op_coverage.keys().copied())
.collect::<HashSet<_>>()
.into_iter()
.collect();
backends.sort_by(|a, b| a.as_str().cmp(b.as_str()));
let mut substrate_for: HashMap<BackendId, SubstrateClass> = HashMap::new();
for &b in &backends {
let cls = capabilities
.get(&b)
.map(|c| c.storage_substrate)
.unwrap_or_else(|| default_substrate_for(b));
substrate_for.insert(b, cls);
}
let mut device_set: HashSet<DeviceLocation> = HashSet::new();
device_set.insert(DeviceLocation::Cpu);
for d in &probe.devices {
device_set.insert(d.location);
}
for caps in capabilities.values() {
device_set.insert(caps.device_location);
}
for &b in &backends {
device_set.insert(default_device_for(b));
}
let mut devices: Vec<DeviceLocation> = device_set.into_iter().collect();
devices.sort_by(|a, b| match (a, b) {
(DeviceLocation::Cpu, DeviceLocation::Cpu) => std::cmp::Ordering::Equal,
(DeviceLocation::Cpu, _) => std::cmp::Ordering::Less,
(_, DeviceLocation::Cpu) => std::cmp::Ordering::Greater,
(x, y) => format!("{x:?}").cmp(&format!("{y:?}")),
});
let mut backends_for_device: HashMap<DeviceLocation, Vec<BackendId>> = HashMap::new();
let mut devices_for_backend: HashMap<BackendId, Vec<DeviceLocation>> = HashMap::new();
let mut probe_devices_by_backend: HashMap<BackendId, Vec<DeviceLocation>> = HashMap::new();
for d in &probe.devices {
probe_devices_by_backend
.entry(d.backend)
.or_default()
.push(d.location);
}
for &b in &backends {
let probe_devs = probe_devices_by_backend.remove(&b);
let bd_devices = match probe_devs {
Some(v) if !v.is_empty() => v,
_ => {
if let Some(caps) = capabilities.get(&b) {
vec![caps.device_location]
} else {
vec![default_device_for(b)]
}
}
};
for &dev in &bd_devices {
backends_for_device.entry(dev).or_default().push(b);
}
devices_for_backend.insert(b, bd_devices);
}
for &dev in &devices {
backends_for_device.entry(dev).or_default();
}
let mut transfer_paths: HashMap<(DeviceLocation, DeviceLocation), TransferPath> =
HashMap::new();
for caps in capabilities.values() {
for (dst, path) in &caps.transfer_paths {
transfer_paths.insert((caps.device_location, *dst), *path);
}
}
SystemTopology {
generation: built_gen,
devices,
backends,
backends_for_device,
devices_for_backend,
substrate_for,
capabilities,
transfer_paths,
binding_op_coverage,
transfer_calibration: OnceLock::new(),
}
}
}
impl crate::ranker::TransferEstimator for SystemTopology {
fn estimate_transfer_ns(&self, src: DeviceLocation, dst: DeviceLocation, bytes: u64) -> u64 {
SystemTopology::estimate_transfer_ns(self, src, dst, bytes)
}
}
fn default_substrate_for(backend: BackendId) -> SubstrateClass {
match backend {
BackendId::Cpu => SubstrateClass::HostBytes,
BackendId::Cuda => SubstrateClass::CudaUntyped,
BackendId::Vulkan => SubstrateClass::VulkanBuffer,
BackendId::Metal => SubstrateClass::MetalBuffer,
_ => SubstrateClass::HostBytes,
}
}
fn default_device_for(backend: BackendId) -> DeviceLocation {
match backend {
BackendId::Cpu => DeviceLocation::Cpu,
BackendId::Cuda => DeviceLocation::Cuda { gpu_id: 0 },
BackendId::Vulkan => DeviceLocation::Vulkan { gpu_id: 0 },
BackendId::Metal => DeviceLocation::Metal { gpu_id: 0 },
_ => DeviceLocation::Cpu,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cpu_always_present() {
let topology = SystemTopology::current();
assert!(topology.devices().contains(&DeviceLocation::Cpu));
assert!(topology.backends().contains(&BackendId::Cpu));
let backends_at_cpu = topology.backends_for(DeviceLocation::Cpu);
assert!(
backends_at_cpu.contains(&BackendId::Cpu),
"Cpu should be among the backends at DeviceLocation::Cpu, got: {:?}",
backends_at_cpu,
);
}
#[test]
fn cpu_shares_with_itself() {
let topology = SystemTopology::current();
assert!(topology.shares_storage(
(BackendId::Cpu, DeviceLocation::Cpu),
(BackendId::Cpu, DeviceLocation::Cpu),
));
}
#[test]
fn transfer_path_defaults() {
let topology = SystemTopology::current();
assert_eq!(
topology.transfer_path(DeviceLocation::Cpu, DeviceLocation::Cpu),
TransferPath::SameDevice,
);
}
#[test]
fn transfer_calibration_cached_per_snapshot() {
let topology = SystemTopology::current();
let a = topology.transfer_calibration();
let b = topology.transfer_calibration();
assert!(
std::ptr::eq(a, b),
"transfer_calibration must be probed once and cached on the snapshot",
);
SystemTopology::refresh();
let next = SystemTopology::current();
assert!(
!std::ptr::eq(a, next.transfer_calibration()),
"a new generation's snapshot must carry its own calibration",
);
}
#[test]
fn transfer_estimate_same_device_zero_and_fallback_monotonic() {
let topology = SystemTopology::current();
assert_eq!(
topology.estimate_transfer_ns(DeviceLocation::Cpu, DeviceLocation::Cpu, 64 << 20),
0,
);
let phantom = DeviceLocation::Cuda { gpu_id: 7 };
let est = topology.transfer_estimate(DeviceLocation::Cpu, phantom);
assert_eq!(
est,
fuel_hardware::transfer_cost::TransferEstimate::fallback_for(TransferPath::HostStaging),
"unprobed path must price via the per-path-class fallback",
);
let small = topology.estimate_transfer_ns(DeviceLocation::Cpu, phantom, 64 * 1024);
let large = topology.estimate_transfer_ns(DeviceLocation::Cpu, phantom, 64 << 20);
assert!(
small > 0,
"unprobed cross-device transfer must never cost zero"
);
assert!(large > small, "estimate must be monotonic in bytes");
}
#[cfg(feature = "cuda")]
#[test]
#[ignore = "requires a live CUDA GPU"]
fn live_cuda_transfer_calibration_sane() {
let topology = SystemTopology::current();
let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
assert!(
topology.devices().contains(&cuda0),
"live test expects a CUDA device, got {:?}",
topology.devices(),
);
let cal = topology.transfer_calibration();
let h2d = cal
.probed(DeviceLocation::Cpu, cuda0)
.expect("H2D path must be probed on a CUDA host");
let d2h = cal
.probed(cuda0, DeviceLocation::Cpu)
.expect("D2H path must be probed on a CUDA host");
const GB: u64 = 1_000_000_000;
for (label, est) in [("H2D", h2d), ("D2H", d2h)] {
assert!(
(GB..=100 * GB).contains(&est.bandwidth_bytes_per_sec),
"{label} bandwidth out of sane bounds (1-100 GB/s): {est:?}",
);
}
let est = topology.transfer_estimate(DeviceLocation::Cpu, cuda0);
let ns_small = est.estimate_ns(64 * 1024);
let ns_mid = est.estimate_ns(4 << 20);
let ns_large = est.estimate_ns(64 << 20);
assert!(ns_small <= ns_mid && ns_mid <= ns_large);
assert!(ns_large > ns_small, "64 MB must cost more than 64 KB");
}
#[cfg(feature = "vulkan")]
#[test]
#[ignore = "requires a live Vulkan GPU"]
fn live_vulkan_transfer_calibration_sane() {
let topology = SystemTopology::current();
let vk0 = DeviceLocation::Vulkan { gpu_id: 0 };
assert!(
topology.devices().contains(&vk0),
"live test expects a Vulkan device, got {:?}",
topology.devices(),
);
let cal = topology.transfer_calibration();
let h2d = cal
.probed(DeviceLocation::Cpu, vk0)
.expect("H2D path must be probed on a Vulkan host");
let d2h = cal
.probed(vk0, DeviceLocation::Cpu)
.expect("D2H path must be probed on a Vulkan host");
const GB: u64 = 1_000_000_000;
assert!(
(GB / 5..=100 * GB).contains(&h2d.bandwidth_bytes_per_sec),
"H2D bandwidth out of sane bounds (0.2-100 GB/s): {h2d:?}",
);
assert!(
(10_000_000..=100 * GB).contains(&d2h.bandwidth_bytes_per_sec),
"D2H bandwidth out of sane bounds (10 MB/s-100 GB/s): {d2h:?}",
);
let est = topology.transfer_estimate(DeviceLocation::Cpu, vk0);
assert!(est.estimate_ns(64 << 20) > est.estimate_ns(64 * 1024));
}
#[test]
fn cpu_capabilities_present() {
let topology = SystemTopology::current();
let caps = topology
.capabilities(BackendId::Cpu)
.expect("CPU caps always registered via default_cpu_caps");
assert_eq!(caps.backend_id, BackendId::Cpu);
assert_eq!(caps.storage_substrate, SubstrateClass::HostBytes);
}
#[test]
fn cpu_op_coverage_no_divergence() {
let topology = SystemTopology::current();
let divergence = topology.capabilities_op_coverage_divergence();
let cpu_divergence: Vec<_> = divergence
.iter()
.filter(|(b, _, _)| *b == BackendId::Cpu)
.collect();
assert!(
cpu_divergence.is_empty(),
"CPU advertises ops it didn't register: {:#?}",
cpu_divergence,
);
}
#[test]
fn no_change_reuses_arc() {
use crate::dispatch::topology_generation;
for _ in 0..32 {
let gen_before = topology_generation();
let a = SystemTopology::current();
let b = SystemTopology::current();
let gen_after = topology_generation();
if gen_before == gen_after && a.generation() == b.generation() {
assert!(
Arc::ptr_eq(&a, &b),
"two current() calls with no intervening generation \
change should return the same Arc (gen={})",
gen_before,
);
return;
}
}
panic!(
"Could not observe a stable-generation window across 32 \
attempts — parallel tests may be saturating the counter; \
rerun with --test-threads=1 if this persists",
);
}
#[test]
fn bump_forces_rebuild() {
let before = SystemTopology::current();
SystemTopology::refresh();
let after = SystemTopology::current();
assert!(
after.generation() > before.generation(),
"refresh() should advance the generation: before={}, after={}",
before.generation(),
after.generation(),
);
assert!(
!Arc::ptr_eq(&before, &after),
"refresh() should produce a fresh Arc",
);
}
#[test]
fn concurrent_access_is_safe() {
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicBool, AtomicUsize};
use std::thread;
use std::time::Duration;
let stop = StdArc::new(AtomicBool::new(false));
let max_backends = StdArc::new(AtomicUsize::new(0));
let mut readers = Vec::new();
for _ in 0..8 {
let stop = StdArc::clone(&stop);
let max_backends = StdArc::clone(&max_backends);
readers.push(thread::spawn(move || {
while !stop.load(std::sync::atomic::Ordering::Acquire) {
let t = SystemTopology::current();
for &b in t.backends() {
let devs = t.devices_for(b);
assert!(
!devs.is_empty(),
"backend {:?} has no devices in snapshot gen={}",
b,
t.generation(),
);
}
let n = t.backends().len();
let mut cur = max_backends.load(std::sync::atomic::Ordering::Acquire);
while n > cur {
match max_backends.compare_exchange_weak(
cur,
n,
std::sync::atomic::Ordering::AcqRel,
std::sync::atomic::Ordering::Acquire,
) {
Ok(_) => break,
Err(prev) => cur = prev,
}
}
}
}));
}
let bumper_stop = StdArc::clone(&stop);
let bumper = thread::spawn(move || {
for _ in 0..50 {
if bumper_stop.load(std::sync::atomic::Ordering::Acquire) {
return;
}
SystemTopology::refresh();
thread::sleep(Duration::from_micros(50));
}
});
bumper.join().expect("bumper thread");
stop.store(true, std::sync::atomic::Ordering::Release);
for r in readers {
r.join().expect("reader thread");
}
}
#[cfg(feature = "cuda")]
#[test]
fn cuda_backend_discovered_when_feature_enabled() {
let topology = SystemTopology::current();
assert!(
topology.backends().contains(&BackendId::Cuda),
"cuda feature is on; topology should report Cuda. Got: {:?}",
topology.backends(),
);
let cuda_dev = DeviceLocation::Cuda { gpu_id: 0 };
let here = topology.backends_for(cuda_dev);
assert!(
here.contains(&BackendId::Cuda),
"Cuda should be among the backends targeting {:?}, got {:?}",
cuda_dev,
here,
);
}
#[cfg(feature = "vulkan")]
#[test]
fn vulkan_backend_discovered_when_feature_enabled() {
let topology = SystemTopology::current();
assert!(
topology.backends().contains(&BackendId::Vulkan),
"vulkan feature is on; topology should report Vulkan. Got: {:?}",
topology.backends(),
);
let vk_dev = DeviceLocation::Vulkan { gpu_id: 0 };
let here = topology.backends_for(vk_dev);
assert!(
here.contains(&BackendId::Vulkan),
"Vulkan should be among the backends targeting {:?}, got {:?}",
vk_dev,
here,
);
}
#[cfg(all(feature = "cuda", feature = "vulkan"))]
#[test]
fn cuda_and_vulkan_dont_share_storage() {
let topology = SystemTopology::current();
let shares = topology.shares_storage(
(BackendId::Cuda, DeviceLocation::Cuda { gpu_id: 0 }),
(BackendId::Vulkan, DeviceLocation::Vulkan { gpu_id: 0 }),
);
assert!(
!shares,
"CUDA and Vulkan on the same silicon must not share storage substrate",
);
let path = topology.transfer_path(
DeviceLocation::Cuda { gpu_id: 0 },
DeviceLocation::Vulkan { gpu_id: 0 },
);
assert_eq!(
path,
TransferPath::HostStaging,
"Cross-vendor GPU transfer should fall back to host-staging today",
);
}
}