use crate::{DType, DeviceLocation, dispatch::OpKind, probe::BackendId};
use std::collections::HashSet;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum SubstrateClass {
HostBytes,
CudaUntyped,
VulkanBuffer,
MetalBuffer,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TransferPath {
SameDevice,
Peer,
SharedMemory,
DeviceCopy,
HostStaging,
}
#[derive(Debug, Clone)]
pub struct BackendCapabilities {
pub backend_id: BackendId,
pub device_location: DeviceLocation,
pub op_dtype_support: HashSet<(OpKind, DType)>,
pub required_alignment: usize,
pub access_granularity_bits: u32,
pub transfer_paths: Vec<(DeviceLocation, TransferPath)>,
pub storage_substrate: SubstrateClass,
pub compute_throughput_flops_per_ns: f64,
pub mem_bandwidth_bytes_per_ns: f64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FitStatus {
Comfortable,
Tight,
WontFit,
Unknown,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::DeviceLocation;
#[test]
fn capability_construction_round_trip() {
let caps = BackendCapabilities {
backend_id: BackendId::Cpu,
device_location: DeviceLocation::Cpu,
op_dtype_support: [
(OpKind::MatMul, DType::F32),
(OpKind::MatMul, DType::F64),
(OpKind::AddElementwise, DType::F32),
]
.into_iter()
.collect(),
required_alignment: 64,
access_granularity_bits: 8,
transfer_paths: vec![(DeviceLocation::Cpu, TransferPath::SameDevice)],
storage_substrate: SubstrateClass::HostBytes,
compute_throughput_flops_per_ns: 1.0,
mem_bandwidth_bytes_per_ns: 4.0,
};
assert_eq!(caps.backend_id, BackendId::Cpu);
assert_eq!(caps.required_alignment, 64);
assert_eq!(caps.access_granularity_bits, 8);
assert!(
caps.op_dtype_support
.contains(&(OpKind::MatMul, DType::F32))
);
assert!(
!caps
.op_dtype_support
.contains(&(OpKind::MatMul, DType::BF16))
);
}
#[test]
fn transfer_path_traits() {
let a = TransferPath::SameDevice;
let b = TransferPath::SameDevice;
let c = TransferPath::HostStaging;
assert_eq!(a, b);
assert_ne!(a, c);
let mut set = HashSet::new();
set.insert(TransferPath::Peer);
set.insert(TransferPath::Peer); set.insert(TransferPath::DeviceCopy);
assert_eq!(set.len(), 2);
}
}