ironaccelerator_core/capability.rs
1//! Coarse capability flags. Backends translate vendor-specific capability
2//! tables (CUDA compute capability, ROCm gfx, Metal feature sets, QNN HTP
3//! version) into this common space so callers can reason about hardware
4//! uniformly without parsing vendor version strings.
5
6use bitflags::bitflags;
7
8bitflags! {
9 /// Hardware capability flags. Stable across backends; new bits append.
10 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12 pub struct CapabilityFlags: u64 {
13 // Numeric formats --------------------------------------------------
14 const FP64 = 1 << 0;
15 const FP32 = 1 << 1;
16 const TF32 = 1 << 2;
17 const FP16 = 1 << 3;
18 const BF16 = 1 << 4;
19 const FP8_E4M3 = 1 << 5;
20 const FP8_E5M2 = 1 << 6;
21 const FP4 = 1 << 7;
22 const INT8 = 1 << 8;
23 const INT4 = 1 << 9;
24
25 // Tensor / matrix engines ------------------------------------------
26 const TENSOR_CORES = 1 << 16; // Volta+ / CDNA / AMX-style MMA
27 const WMMA = 1 << 17; // wavefront matrix-multiply ops
28 const SPARSE_2_4 = 1 << 18; // Ampere-style 2:4 sparsity
29 const TRANSFORMER_ENGINE = 1 << 19; // Hopper+ TE / equivalent
30
31 // Memory ------------------------------------------------------------
32 const UNIFIED_MEMORY = 1 << 24;
33 const HOST_PINNED = 1 << 25;
34 const PEER_ACCESS = 1 << 26;
35 const ASYNC_ALLOC = 1 << 27; // cudaMallocAsync etc.
36 const HBM = 1 << 28;
37 const NVLINK = 1 << 29;
38 const INFINITY_FABRIC = 1 << 30;
39
40 // Compute features -------------------------------------------------
41 const COOPERATIVE_LAUNCH = 1 << 32;
42 const DYNAMIC_PARALLELISM = 1 << 33;
43 const GRAPHS = 1 << 34; // CUDA graphs / equivalent
44 const MULTI_STREAM = 1 << 35;
45 const FLASH_ATTN = 1 << 36; // hardware-friendly fused attn
46 const RAYTRACING = 1 << 37;
47
48 // Distribution -----------------------------------------------------
49 const NCCL = 1 << 40;
50 const RCCL = 1 << 41;
51 const MPI_AWARE = 1 << 42;
52
53 // NPU / mobile -----------------------------------------------------
54 const HVX = 1 << 48; // Hexagon Vector eXtensions
55 const HMX = 1 << 49; // Hexagon Matrix eXtensions
56 const ANE = 1 << 50; // Apple Neural Engine bridge
57 }
58}
59
60/// A coarse "compute tier" used by strategy heuristics when capability bits
61/// are insufficient. Backends map their micro-arch to one of these.
62#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
63#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
64pub enum ComputeTier {
65 /// CPU SIMD baseline.
66 Baseline,
67 /// Mobile NPU / integrated GPU.
68 Mobile,
69 /// Consumer dGPU (e.g. RTX 4060, RX 7600, M-series Pro).
70 Consumer,
71 /// Workstation GPU (RTX 6000 Ada, Radeon Pro, M-series Max).
72 Workstation,
73 /// Datacenter GPU (H100, MI300, GB200).
74 Datacenter,
75}
76
77/// Aggregated capability description that wraps flags + tier + numeric extras.
78#[derive(Debug, Clone, Copy)]
79#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
80pub struct Capability {
81 pub flags: CapabilityFlags,
82 pub tier: ComputeTier,
83 /// Peak FP16 TFLOPS, vendor-published. `None` if unknown.
84 pub fp16_tflops: Option<f32>,
85 /// Peak FP8 TFLOPS, vendor-published.
86 pub fp8_tflops: Option<f32>,
87 /// HBM/GDDR bandwidth in GB/s.
88 pub mem_bandwidth_gbs: Option<f32>,
89}