use bitflags::bitflags;
bitflags! {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CapabilityFlags: u64 {
const FP64 = 1 << 0;
const FP32 = 1 << 1;
const TF32 = 1 << 2;
const FP16 = 1 << 3;
const BF16 = 1 << 4;
const FP8_E4M3 = 1 << 5;
const FP8_E5M2 = 1 << 6;
const FP4 = 1 << 7;
const INT8 = 1 << 8;
const INT4 = 1 << 9;
const TENSOR_CORES = 1 << 16; const WMMA = 1 << 17; const SPARSE_2_4 = 1 << 18; const TRANSFORMER_ENGINE = 1 << 19;
const UNIFIED_MEMORY = 1 << 24;
const HOST_PINNED = 1 << 25;
const PEER_ACCESS = 1 << 26;
const ASYNC_ALLOC = 1 << 27; const HBM = 1 << 28;
const NVLINK = 1 << 29;
const INFINITY_FABRIC = 1 << 30;
const COOPERATIVE_LAUNCH = 1 << 32;
const DYNAMIC_PARALLELISM = 1 << 33;
const GRAPHS = 1 << 34; const MULTI_STREAM = 1 << 35;
const FLASH_ATTN = 1 << 36; const RAYTRACING = 1 << 37;
const NCCL = 1 << 40;
const RCCL = 1 << 41;
const MPI_AWARE = 1 << 42;
const HVX = 1 << 48; const HMX = 1 << 49; const ANE = 1 << 50; }
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum ComputeTier {
Baseline,
Mobile,
Consumer,
Workstation,
Datacenter,
}
#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Capability {
pub flags: CapabilityFlags,
pub tier: ComputeTier,
pub fp16_tflops: Option<f32>,
pub fp8_tflops: Option<f32>,
pub mem_bandwidth_gbs: Option<f32>,
}