Skip to main content

cubecl_ir/
properties.rs

1use alloc::string::String;
2use core::{
3    fmt,
4    hash::{BuildHasher, Hash, Hasher},
5    str::FromStr,
6};
7
8use crate::{
9    AddressType, ElemType, EnumSet, EnumSetType, OpaqueType, SemanticType, Type, TypeHash,
10    VectorSize,
11    features::{AtomicUsage, ComplexUsage, Features, TypeUsage},
12};
13use cubecl_common::profile::TimingMethod;
14
15/// Properties of the device related to the accelerator hardware.
16///
17/// # Plane size min/max
18///
19/// This is a range of possible values for the plane size.
20///
21/// For Nvidia GPUs and HIP, this is a single fixed value.
22///
23/// For wgpu with AMD GPUs this is a range of possible values, but the actual configured value
24/// is undefined and can only be queried at runtime. Should usually be 32, but not guaranteed.
25///
26/// For Intel GPUs, this is variable based on the number of registers used in the kernel. No way to
27/// query this at compile time is currently available. As a result, the minimum value should usually
28/// be assumed.
29#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub struct HardwareProperties {
32    /// The widest single load instruction, in bits.
33    pub load_width: u32,
34    /// How many `load_width`-bit vector registers a kernel may keep live, or `None` where the
35    /// runtime states no budget, as a GPU does.
36    pub vector_register_count: Option<u32>,
37    /// The minimum size of a plane on this device
38    pub plane_size_min: u32,
39    /// The maximum size of a plane on this device
40    pub plane_size_max: u32,
41    /// minimum number of bindings for a kernel that can be used at once.
42    pub max_bindings: u32,
43    /// Maximum amount of shared memory, in bytes
44    pub max_shared_memory_size: usize,
45    /// Maximum `CubeCount` in x, y and z dimensions
46    pub max_cube_count: (u32, u32, u32),
47    /// Maximum number of total units in a cube
48    pub max_units_per_cube: u32,
49    /// Maximum `CubeDim` in x, y, and z dimensions
50    pub max_cube_dim: (u32, u32, u32),
51    /// Number of streaming multiprocessors (SM), if available
52    pub num_streaming_multiprocessors: Option<u32>,
53    /// Number of available parallel cpu units, if the runtime is CPU.
54    pub num_cpu_cores: Option<u32>,
55    /// Bytes of the device's last level cache, and `None`, never `Some(0)`,
56    /// where the runtime cannot read one.
57    ///
58    /// The size a working set has to outgrow before what it reaches is set by
59    /// memory rather than by the chip.
60    pub last_level_cache_size: Option<usize>,
61    /// Number of tensor cores per SM, if any
62    pub num_tensor_cores: Option<u32>,
63    /// The minimum tiling dimension for a single axis in tensor cores.
64    ///
65    /// For a backend that only supports 16x16x16, the value would be 16.
66    /// For a backend that also supports 32x8x16, the value would be 8.
67    pub min_tensor_cores_dim: Option<u32>,
68    /// Maximum vector size supported by the device
69    pub max_vector_size: VectorSize,
70    /// Memory reserved for the driver when using cube-scoped matrices
71    pub cube_mma_reserved_shared_memory: usize,
72}
73
74/// A device's vector registers, counted in elements of one type.
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub struct VectorRegisters {
77    count: usize,
78    register_vector_size: VectorSize,
79    max_vector_size: VectorSize,
80}
81
82impl VectorRegisters {
83    /// The registers as `elem_size`-byte elements see them, or `None` without a register count.
84    pub fn new(hardware: &HardwareProperties, elem_size: usize) -> Option<Self> {
85        Some(VectorRegisters {
86            count: hardware.vector_register_count? as usize,
87            register_vector_size: hardware.vector_size_in(hardware.load_width, elem_size),
88            max_vector_size: prev_power_of_two(hardware.max_vector_size.max(1)),
89        })
90    }
91
92    pub fn count(&self) -> usize {
93        self.count
94    }
95
96    /// Registers a vector of `vector_size` elements occupies, at least one.
97    pub fn registers_for(&self, vector_size: VectorSize) -> usize {
98        vector_size.div_ceil(self.register_vector_size).max(1)
99    }
100
101    /// The widest power-of-two size at which `live` vectors stay in registers, at least one.
102    pub fn widest_vector_size(&self, live: usize) -> VectorSize {
103        let registers = prev_power_of_two((self.count / live.max(1)).max(1));
104        (registers * self.register_vector_size).min(self.max_vector_size)
105    }
106
107    /// How many vectors of `vector_size` elements fit beside `reserved` registers.
108    pub fn vectors_fitting(&self, vector_size: VectorSize, reserved: usize) -> usize {
109        self.count.saturating_sub(reserved) / self.registers_for(vector_size)
110    }
111}
112
113/// Properties of the device related to allocation.
114#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
115#[derive(Debug, Clone, PartialEq, Eq, Hash)]
116#[non_exhaustive]
117pub struct MemoryDeviceProperties {
118    /// The maximum nr. of bytes that can be allocated in one go.
119    pub max_page_size: u64,
120    /// The required memory offset alignment in bytes.
121    pub alignment: u64,
122    /// Private because [`set_max_memory`](Self::set_max_memory) is its only
123    /// writer, and that is where a zero becomes `None`; a deserialized one
124    /// goes through the same rule.
125    #[cfg_attr(feature = "serde", serde(deserialize_with = "deserialize_max_memory"))]
126    max_memory: Option<u64>,
127}
128
129/// A stated capacity as [`MemoryDeviceProperties::set_max_memory`] takes it: a zero is no figure.
130#[cfg(feature = "serde")]
131fn deserialize_max_memory<'de, D: serde::Deserializer<'de>>(
132    deserializer: D,
133) -> Result<Option<u64>, D::Error> {
134    let stated = <Option<u64> as serde::Deserialize>::deserialize(deserializer)?;
135    Ok(stated.filter(|bytes| *bytes != 0))
136}
137
138impl MemoryDeviceProperties {
139    /// Properties that state no capacity. A runtime that can read one adds it
140    /// with [`with_max_memory`](Self::with_max_memory).
141    pub const fn new(max_page_size: u64, alignment: u64) -> Self {
142        Self {
143            max_page_size,
144            alignment,
145            max_memory: None,
146        }
147    }
148
149    /// How many bytes this memory may be asked to hold at once, or `None`,
150    /// never `Some(0)`, where the runtime has no figure.
151    ///
152    /// This sizes a whole workload, while
153    /// [`max_page_size`](Self::max_page_size) bounds a single allocation and
154    /// is often derived from it. It is a budget, not a hardware census: each
155    /// runtime reports the largest figure its own API stands behind, and
156    /// staying under it is what keeps the device off its paging path.
157    ///
158    /// A runtime with no figure leaves it `None` rather than guess, because a
159    /// guess reads as a measurement to every caller downstream.
160    pub const fn max_memory(&self) -> Option<u64> {
161        self.max_memory
162    }
163
164    /// States the capacity, as [`set_max_memory`](Self::set_max_memory) does.
165    pub const fn with_max_memory(mut self, max_memory: u64) -> Self {
166        self.set_max_memory(max_memory);
167        self
168    }
169
170    /// States the capacity, dropping a zero: an API with nothing to report
171    /// reports `0`, and [`max_memory`](Self::max_memory) says that as `None`.
172    pub const fn set_max_memory(&mut self, max_memory: u64) {
173        self.max_memory = match max_memory {
174            0 => None,
175            size => Some(size),
176        };
177    }
178}
179
180/// Who a device is, and what its compiled code is keyed to.
181///
182/// `name` and `fingerprint` answer different questions and must not be confused. `name`
183/// is for people: it names the physical part, and two machines holding the same
184/// part report the same name. `fingerprint` is for correctness: it is verbatim
185/// the string this runtime passes to
186/// [`compilation_store`](../../cubecl_runtime/compiler/fn.compilation_store.html),
187/// which is what puts a compiled artifact out of reach of a machine that cannot
188/// run it.
189///
190/// Reporting the fingerprint here rather than recomputing it is the whole
191/// point: a backend derives it once and hands it to both consumers, so the
192/// identity a bundle is stamped with and the namespace its kernels live under
193/// cannot drift apart.
194#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
195#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
196pub struct DeviceIdentity {
197    /// The device as it names itself — `AMD Radeon 8060S Graphics`,
198    /// `NVIDIA H100 PCIe`. Distinct parts may share a name, so no capability
199    /// and no compiled artifact may be keyed to it.
200    pub name: String,
201    /// What this runtime compiles *for* — `hip-kernel_gfx1151`, `ptx_sm90`.
202    /// Verbatim the `compilation_store` fingerprint, so a namespace read back
203    /// out of a bundle compares against it directly.
204    pub fingerprint: String,
205    /// The card behind the device, `None` for a device that is no card: a CPU, or a software
206    /// rasterizer.
207    ///
208    /// Never serialized: every id on a card is local to the host that reported it, so a card read
209    /// on another host would match that host's own card at the same PCI address.
210    #[cfg_attr(feature = "serde", serde(skip))]
211    pub physical: Option<PhysicalDevice>,
212}
213
214/// The card a device runs on. Two runtimes report different fields for one card, so compare with
215/// [`is_same_card`](Self::is_same_card), not `==`.
216#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
217#[non_exhaustive]
218pub struct PhysicalDevice {
219    /// The key every runtime reports alike on Linux.
220    pub pci_address: Option<PciAddress>,
221    /// The key every runtime reports alike on Windows.
222    pub luid: Option<AdapterLuid>,
223    /// The key every runtime reports alike on macOS.
224    pub registry_entry_id: Option<RegistryEntryId>,
225    pub vendor: Option<PciVendor>,
226}
227
228impl PhysicalDevice {
229    /// Whether `other` is known to be this card through another runtime: by PCI address, otherwise
230    /// by LUID, otherwise by registry entry id. A card with none of them matches nothing, itself
231    /// included, since two such cards of one make compare equal.
232    pub fn is_same_card(&self, other: &Self) -> bool {
233        if let (Some(mine), Some(theirs)) = (self.pci_address, other.pci_address) {
234            return mine == theirs;
235        }
236        if let (Some(mine), Some(theirs)) = (self.luid, other.luid) {
237            return mine == theirs;
238        }
239        matches!(
240            (self.registry_entry_id, other.registry_entry_id),
241            (Some(mine), Some(theirs)) if mine == theirs
242        )
243    }
244
245    /// Takes what this runtime left out from `other`, the same card seen through another runtime.
246    pub fn fill_from(&mut self, other: &Self) {
247        let Self {
248            pci_address,
249            luid,
250            registry_entry_id,
251            vendor,
252        } = *other;
253        self.pci_address = self.pci_address.or(pci_address);
254        self.luid = self.luid.or(luid);
255        self.registry_entry_id = self.registry_entry_id.or(registry_entry_id);
256        self.vendor = self.vendor.or(vendor);
257    }
258}
259
260/// The id of the GPU's entry in the macOS I/O Registry, which Metal reports as `registryID`. Like
261/// [`AdapterLuid`] it is not promised to survive a restart, so it has no serialization or text
262/// form.
263#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
264pub struct RegistryEntryId(u64);
265
266impl RegistryEntryId {
267    pub fn new(id: u64) -> Self {
268        Self(id)
269    }
270
271    pub fn get(self) -> u64 {
272        self.0
273    }
274}
275
276/// The id Windows gives a graphics adapter. It changes on restart, so it has no serialization or
277/// text form: a stored key wants [`PhysicalDevice::pci_address`].
278#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
279pub struct AdapterLuid([u8; 8]);
280
281impl AdapterLuid {
282    /// From the eight bytes of a Windows `LUID`, low part first.
283    pub fn new(bytes: [u8; 8]) -> Self {
284        Self(bytes)
285    }
286
287    pub fn from_parts(low_part: u32, high_part: i32) -> Self {
288        let mut bytes = [0; 8];
289        bytes[..4].copy_from_slice(&low_part.to_le_bytes());
290        bytes[4..].copy_from_slice(&high_part.to_le_bytes());
291        Self(bytes)
292    }
293
294    pub fn bytes(self) -> [u8; 8] {
295        self.0
296    }
297}
298
299/// The maker of a card, by PCI vendor id.
300#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
301pub enum PciVendor {
302    Nvidia,
303    Amd,
304    Intel,
305    Apple,
306    /// Mali GPUs.
307    Arm,
308    /// Adreno GPUs.
309    Qualcomm,
310    Other(u32),
311}
312
313impl PciVendor {
314    pub fn id(self) -> u32 {
315        match self {
316            Self::Nvidia => 0x10de,
317            Self::Amd => 0x1002,
318            Self::Intel => 0x8086,
319            Self::Apple => 0x106b,
320            Self::Arm => 0x13b5,
321            Self::Qualcomm => 0x5143,
322            Self::Other(id) => id,
323        }
324    }
325}
326
327impl From<u32> for PciVendor {
328    fn from(id: u32) -> Self {
329        match id {
330            0x10de => Self::Nvidia,
331            0x1002 => Self::Amd,
332            0x8086 => Self::Intel,
333            0x106b => Self::Apple,
334            0x13b5 => Self::Arm,
335            0x5143 => Self::Qualcomm,
336            other => Self::Other(other),
337        }
338    }
339}
340
341impl fmt::Display for PciVendor {
342    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
343        match self {
344            Self::Nvidia => f.write_str("NVIDIA"),
345            Self::Amd => f.write_str("AMD"),
346            Self::Intel => f.write_str("Intel"),
347            Self::Apple => f.write_str("Apple"),
348            Self::Arm => f.write_str("Arm"),
349            Self::Qualcomm => f.write_str("Qualcomm"),
350            Self::Other(id) => write!(f, "{id:#06x}"),
351        }
352    }
353}
354
355/// `domain:bus:device.function`, written `0000:07:00.0`. CUDA and NVML call it the bus id.
356#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
357pub struct PciAddress {
358    pub domain: u32,
359    pub bus: u8,
360    pub device: u8,
361    pub function: u8,
362}
363
364impl fmt::Display for PciAddress {
365    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
366        write!(
367            f,
368            "{:04x}:{:02x}:{:02x}.{:x}",
369            self.domain, self.bus, self.device, self.function
370        )
371    }
372}
373
374#[derive(Debug, Clone, PartialEq, Eq)]
375pub struct PciAddressError(pub String);
376
377impl fmt::Display for PciAddressError {
378    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
379        write!(f, "not a PCI address: {}", self.0)
380    }
381}
382
383impl core::error::Error for PciAddressError {}
384
385impl FromStr for PciAddress {
386    type Err = PciAddressError;
387
388    /// Also accepts the domainless `07:00.0` CUDA emits.
389    fn from_str(text: &str) -> Result<Self, Self::Err> {
390        let err = || PciAddressError(String::from(text));
391        let (rest, function) = text.rsplit_once('.').ok_or_else(err)?;
392        let mut parts = rest.rsplitn(3, ':');
393        let device = parts.next().ok_or_else(err)?;
394        let bus = parts.next().ok_or_else(err)?;
395        let domain = parts.next().unwrap_or("0");
396        Ok(Self {
397            domain: u32::from_str_radix(domain, 16).map_err(|_| err())?,
398            bus: u8::from_str_radix(bus, 16).map_err(|_| err())?,
399            device: u8::from_str_radix(device, 16).map_err(|_| err())?,
400            function: u8::from_str_radix(function, 16).map_err(|_| err())?,
401        })
402    }
403}
404
405/// Properties of what the device can do, like what `Feature` are
406/// supported by it and what its memory properties are.
407#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
408#[derive(Debug, Clone, PartialEq, Eq)]
409pub struct DeviceProperties {
410    /// The features supported by the runtime.
411    pub features: Features,
412    /// The memory properties of this client.
413    pub memory: MemoryDeviceProperties,
414    /// The topology properties of this client.
415    pub hardware: HardwareProperties,
416    /// The method used for profiling on the device.
417    pub timing_method: TimingMethod,
418    /// Who the device is, and what its kernels are keyed to.
419    pub identity: DeviceIdentity,
420    /// Bits IO is sized to in place of the load width, where a backend measured wider as faster.
421    pub io_width_override: Option<u32>,
422}
423
424impl TypeHash for DeviceProperties {
425    fn write_hash(_hasher: &mut impl core::hash::Hasher) {
426        // ignored.
427    }
428}
429
430impl DeviceProperties {
431    /// Create a new feature set with the given features and memory properties.
432    ///
433    /// `identity` is a required argument rather than something a backend may
434    /// fill in afterwards, so a runtime cannot ship reporting an anonymous
435    /// device — the failure mode that leaves a bundle unable to say what it was
436    /// built for.
437    pub fn new(
438        features: Features,
439        memory_props: MemoryDeviceProperties,
440        hardware: HardwareProperties,
441        timing_method: TimingMethod,
442        identity: DeviceIdentity,
443    ) -> Self {
444        DeviceProperties {
445            features,
446            memory: memory_props,
447            hardware,
448            timing_method,
449            identity,
450            io_width_override: None,
451        }
452    }
453
454    /// The widest vector, in bits, that reads and writes are sized to.
455    pub fn io_width(&self) -> u32 {
456        self.io_width_override.unwrap_or(self.hardware.load_width)
457    }
458
459    /// Vector sizes, widest first, that reads and writes of `elem_size`-byte elements are sized to.
460    pub fn io_optimized_vector_sizes(
461        &self,
462        elem_size: usize,
463    ) -> impl Iterator<Item = VectorSize> + Clone {
464        vector_sizes_down_from(self.widest_io_vector_size(elem_size))
465    }
466
467    /// Vector sizes, widest first, that keep `live` vectors in registers, else the IO sizes.
468    pub fn vector_sizes_in_registers(
469        &self,
470        elem_size: usize,
471        live: usize,
472    ) -> impl Iterator<Item = VectorSize> + Clone {
473        let widest = match VectorRegisters::new(&self.hardware, elem_size) {
474            Some(registers) => registers.widest_vector_size(live),
475            None => self.widest_io_vector_size(elem_size),
476        };
477
478        vector_sizes_down_from(widest)
479    }
480
481    /// Get the usages for a type
482    pub fn type_usage(&self, ty: ElemType) -> EnumSet<TypeUsage> {
483        self.features.type_usage(ty)
484    }
485
486    /// Get the complex capability families for a type.
487    pub fn complex_usage(&self, ty: ElemType) -> EnumSet<ComplexUsage> {
488        self.features.complex_usage(ty)
489    }
490
491    /// Whether a complex type supports the requested capability family.
492    pub fn supports_complex_usage(&self, ty: ElemType, usage: ComplexUsage) -> bool {
493        self.features.supports_complex_usage(ty, usage)
494    }
495
496    /// Get the usages for an atomic type on global memory
497    pub fn atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
498        self.features.atomic_type_usage(ty)
499    }
500
501    /// Get the usages for an atomic type on shared memory
502    pub fn shared_atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
503        self.features.shared_atomic_type_usage(ty)
504    }
505
506    /// Whether the type is supported in any way
507    pub fn supports_type(&self, ty: impl Into<Type>) -> bool {
508        self.features.supports_type(ty)
509    }
510
511    /// Whether the address type is supported in any way
512    pub fn supports_address(&self, ty: impl Into<AddressType>) -> bool {
513        self.features.supports_address(ty)
514    }
515
516    /// Register an address type to the features
517    pub fn register_address_type(&mut self, ty: impl Into<AddressType>) {
518        self.features.types.address.insert(ty.into());
519    }
520
521    /// Register the usages of an atomic type on global memory
522    pub fn register_atomic_type_usage(&mut self, ty: Type, uses: impl Into<EnumSet<AtomicUsage>>) {
523        *self.features.types.atomic.entry(ty).or_default() |= uses.into();
524    }
525
526    /// Register the usages of an atomic type on shared memory
527    pub fn register_shared_atomic_type_usage(
528        &mut self,
529        ty: Type,
530        uses: impl Into<EnumSet<AtomicUsage>>,
531    ) {
532        *self.features.types.shared_atomic.entry(ty).or_default() |= uses.into();
533    }
534
535    /// Register a storage type to the features
536    pub fn register_type_usage(
537        &mut self,
538        ty: impl Into<ElemType>,
539        uses: impl Into<EnumSet<TypeUsage>>,
540    ) {
541        *self.features.types.elem.entry(ty.into()).or_default() |= uses.into();
542    }
543
544    /// Register `ty`'s conversion as emulated in software ([`Types::emulated_conversion`]).
545    ///
546    /// [`Types::emulated_conversion`]: crate::Types::emulated_conversion
547    pub fn register_emulated_conversion(&mut self, ty: impl Into<ElemType>) {
548        self.features.types.emulated_conversion.insert(ty.into());
549    }
550
551    /// Register `ty`'s arithmetic as emulated through a wider type
552    /// ([`Types::emulated_arithmetic`]).
553    ///
554    /// [`Types::emulated_arithmetic`]: crate::Types::emulated_arithmetic
555    pub fn register_emulated_arithmetic(&mut self, ty: impl Into<ElemType>) {
556        self.features.types.emulated_arithmetic.insert(ty.into());
557    }
558
559    /// Register complex capability families for an element type.
560    pub fn register_complex_usage(
561        &mut self,
562        ty: impl Into<ElemType>,
563        uses: impl Into<EnumSet<ComplexUsage>>,
564    ) {
565        *self.features.types.complex.entry(ty.into()).or_default() |= uses.into();
566    }
567
568    /// Register a semantic type to the features
569    pub fn register_semantic_type(&mut self, ty: SemanticType) {
570        self.features.types.semantic.insert(ty);
571    }
572
573    /// Register an opaque type to the features
574    pub fn register_opaque_type(&mut self, ty: OpaqueType) {
575        self.features.types.opaque.insert(ty);
576    }
577
578    /// Create a stable hash of all device properties relevant to kernel compilation. Can be used
579    /// as a stable checksum for a compilation cache.
580    pub fn checksum(&self) -> u64 {
581        let state = foldhash::fast::FixedState::default();
582        let mut hasher = state.build_hasher();
583        self.features.hash(&mut hasher);
584        self.hardware.hash(&mut hasher);
585        hasher.finish()
586    }
587
588    /// A power of two, and at least one.
589    fn widest_io_vector_size(&self, elem_size: usize) -> VectorSize {
590        self.hardware.vector_size_in(self.io_width(), elem_size)
591    }
592}
593
594/// Unchecked optimizations for float operations. May cause precision differences, or undefined
595/// behaviour if the relevant conditions are not followed.
596#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
597#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, EnumSetType)]
598pub enum FastMath {
599    /// Assume values are never `NaN`. If they are, the result is considered undefined behaviour.
600    NotNaN,
601    /// Assume values are never `Inf`/`-Inf`. If they are, the result is considered undefined
602    /// behaviour.
603    NotInf,
604    /// Ignore sign on zero values.
605    UnsignedZero,
606    /// Allow swapping float division with a reciprocal, even if that swap would change precision.
607    AllowReciprocal,
608    /// Allow contracting float operations into fewer operations, even if the precision could
609    /// change.
610    AllowContraction,
611    /// Allow reassociation for float operations, even if the precision could change.
612    AllowReassociation,
613    /// Allow all mathematical transformations for float operations, including contraction and
614    /// reassociation, even if the precision could change.
615    AllowTransform,
616    /// Allow using lower precision intrinsics
617    ReducedPrecision,
618}
619
620impl FastMath {
621    pub const fn all() -> EnumSet<FastMath> {
622        EnumSet::all()
623    }
624}
625
626impl HardwareProperties {
627    /// `width` bits of `elem_size`-byte elements, as a power of two the device accepts.
628    fn vector_size_in(&self, width: u32, elem_size: usize) -> VectorSize {
629        let elems = width as usize / (elem_size * 8);
630        prev_power_of_two(elems.min(self.max_vector_size).max(1))
631    }
632}
633
634fn prev_power_of_two(value: usize) -> usize {
635    1 << value.ilog2()
636}
637
638fn vector_sizes_down_from(widest: usize) -> impl Iterator<Item = VectorSize> + Clone {
639    (0..=widest.ilog2()).rev().map(|exponent| 1 << exponent)
640}
641
642#[cfg(test)]
643mod tests {
644    use super::*;
645    use crate::FloatKind;
646    use alloc::string::ToString;
647
648    /// Properties sent to another host read back as they were sent, the element-type tables
649    /// keyed by a non-string type included, and a stated capacity of zero reads back as none.
650    #[cfg(feature = "serde")]
651    #[test]
652    fn device_properties_round_trip_through_a_binary_format() {
653        let mut sent = DeviceProperties::new(
654            Features::default(),
655            MemoryDeviceProperties::new(1 << 30, 256).with_max_memory(8 << 30),
656            hardware(128, Some(256)),
657            TimingMethod::Device,
658            DeviceIdentity {
659                name: "a card".to_string(),
660                fingerprint: "hip-kernel_gfx1102".to_string(),
661                physical: None,
662            },
663        );
664        sent.register_type_usage(ElemType::Float(FloatKind::F16), TypeUsage::Arithmetic);
665        sent.register_emulated_conversion(ElemType::Float(FloatKind::BF16));
666
667        let mut bytes = alloc::vec::Vec::new();
668        ciborium::into_writer(&sent, &mut bytes).unwrap();
669        let received: DeviceProperties = ciborium::from_reader(bytes.as_slice()).unwrap();
670        assert_eq!(received, sent);
671
672        /// What another host could send: a capacity of zero, which no writer here produces.
673        #[derive(serde::Serialize)]
674        struct Stated {
675            max_page_size: u64,
676            alignment: u64,
677            max_memory: Option<u64>,
678        }
679        let mut zero = alloc::vec::Vec::new();
680        let stated = Stated {
681            max_page_size: 1,
682            alignment: 1,
683            max_memory: Some(0),
684        };
685        ciborium::into_writer(&stated, &mut zero).unwrap();
686        let memory: MemoryDeviceProperties = ciborium::from_reader(zero.as_slice()).unwrap();
687        assert_eq!(memory.max_memory(), None);
688    }
689
690    /// A card read on another host is no card here: its PCI address would name whichever local
691    /// card sits at the same address, and every check keyed to the card would read that one.
692    #[cfg(feature = "serde")]
693    #[test]
694    fn a_card_never_leaves_the_host_that_reported_it() {
695        let sent = DeviceIdentity {
696            name: "a card".to_string(),
697            fingerprint: "hip-kernel_gfx1102".to_string(),
698            physical: Some(PhysicalDevice {
699                pci_address: Some("0000:03:00.0".parse().unwrap()),
700                vendor: Some(PciVendor::Amd),
701                ..PhysicalDevice::default()
702            }),
703        };
704
705        let mut bytes = alloc::vec::Vec::new();
706        ciborium::into_writer(&sent, &mut bytes).unwrap();
707        let received: DeviceIdentity = ciborium::from_reader(bytes.as_slice()).unwrap();
708        assert_eq!(received.physical, None);
709        assert_eq!(received.fingerprint, sent.fingerprint);
710    }
711
712    #[test]
713    fn a_device_without_a_fixed_register_set_has_no_budget() {
714        assert_eq!(VectorRegisters::new(&hardware(128, None), 4), None);
715    }
716
717    #[test]
718    fn the_widest_vector_size_keeps_every_live_vector_in_registers() {
719        let f32 = registers(AVX2, 4);
720        assert_eq!(f32.widest_vector_size(6), 16);
721        assert_eq!(f32.widest_vector_size(8), 16);
722        assert_eq!(f32.widest_vector_size(9), 8);
723        assert_eq!(registers(AVX2, 8).widest_vector_size(6), 8);
724        assert_eq!(registers(AVX512, 4).widest_vector_size(6), 64);
725    }
726
727    #[test]
728    fn neon_and_avx2_budget_the_same_vector_sizes_from_equal_register_files() {
729        let (neon, avx2) = (registers(NEON, 4), registers(AVX2, 4));
730        assert_eq!(neon.widest_vector_size(3), avx2.widest_vector_size(3));
731        assert_eq!(neon.widest_vector_size(6), avx2.widest_vector_size(6));
732    }
733
734    #[test]
735    fn more_live_vectors_than_registers_still_get_one_register_each() {
736        assert_eq!(registers(AVX2, 4).widest_vector_size(40), 8);
737        assert_eq!(registers(NEON, 4).widest_vector_size(40), 4);
738    }
739
740    #[test]
741    fn vectors_fit_in_the_registers_left_unreserved() {
742        let f32 = registers(AVX2, 4);
743        assert_eq!(f32.registers_for(16), 2);
744        assert_eq!(f32.vectors_fitting(16, 0), 8);
745        assert_eq!(f32.vectors_fitting(8, 4), 12);
746        assert_eq!(f32.vectors_fitting(8, 20), 0);
747    }
748
749    #[test]
750    fn a_capped_vector_size_caps_the_register_vector_size() {
751        let mut capped = hardware(256, Some(16));
752        capped.max_vector_size = 4;
753        let f32 = VectorRegisters::new(&capped, 4).unwrap();
754        assert_eq!(f32.registers_for(8), 2);
755        assert_eq!(f32.widest_vector_size(1), 4);
756    }
757
758    #[test]
759    fn io_follows_a_load_width_raised_after_construction() {
760        let mut device = device(128);
761        device.hardware.load_width = 256;
762        assert_eq!(device.io_width(), 256);
763    }
764
765    #[test]
766    fn a_stated_io_width_outlasts_a_load_width_change() {
767        let mut device = device(256);
768        device.io_width_override = Some(512);
769        device.hardware.load_width = 128;
770        assert_eq!(device.io_width(), 512);
771    }
772
773    #[test]
774    fn a_cap_between_powers_of_two_rounds_down_alike() {
775        let mut device = device(512);
776        device.hardware.max_vector_size = 12;
777        device.hardware.vector_register_count = Some(32);
778
779        assert_eq!(
780            device
781                .io_optimized_vector_sizes(4)
782                .collect::<alloc::vec::Vec<_>>(),
783            alloc::vec![8, 4, 2, 1]
784        );
785        assert_eq!(device.vector_sizes_in_registers(4, 1).next(), Some(8));
786    }
787
788    #[test]
789    fn a_device_without_registers_sizes_by_its_io() {
790        let device = device(128);
791        assert!(
792            device
793                .vector_sizes_in_registers(4, 6)
794                .eq(device.io_optimized_vector_sizes(4))
795        );
796    }
797
798    /// A capacity is stated only by the runtime that read one.
799    ///
800    /// Properties built without one must answer `None`, or a caller would
801    /// size a workload against a figure nobody measured.
802    #[test]
803    fn a_memory_states_no_capacity_until_one_is_read() {
804        let props = MemoryDeviceProperties::new(1024, 32);
805        assert_eq!(props.max_memory(), None);
806        assert_eq!(props.clone().with_max_memory(4096).max_memory(), Some(4096));
807    }
808
809    /// A zero from the runtime's API reads as no capacity.
810    ///
811    /// An API with nothing to report reports `0`, which must not reach a
812    /// caller as a device that holds nothing.
813    #[test]
814    fn a_capacity_of_zero_is_no_capacity() {
815        let mut props = MemoryDeviceProperties::new(1024, 32).with_max_memory(4096);
816        props.set_max_memory(0);
817        assert_eq!(props.max_memory(), None);
818    }
819
820    #[test]
821    fn a_vendor_keeps_its_id_whether_named_or_not() {
822        for id in [0x10de, 0x1002, 0x8086, 0x106b, 0x13b5, 0x5143, 0x1af4] {
823            assert_eq!(PciVendor::from(id).id(), id);
824        }
825        assert_eq!(PciVendor::from(0x10de), PciVendor::Nvidia);
826        assert_eq!(PciVendor::from(0x1af4), PciVendor::Other(0x1af4));
827        assert_eq!(PciVendor::Other(0x1af4).to_string(), "0x1af4");
828    }
829
830    #[test]
831    fn a_luid_from_parts_is_the_bytes_vulkan_reports() {
832        let bytes = [0x8a, 0x1d, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00];
833        assert_eq!(
834            AdapterLuid::from_parts(0x0001_1d8a, 0),
835            AdapterLuid::new(bytes)
836        );
837        assert_eq!(
838            AdapterLuid::from_parts(1, -1).bytes(),
839            [1, 0, 0, 0, 0xff, 0xff, 0xff, 0xff]
840        );
841    }
842
843    #[test]
844    fn a_card_is_matched_by_address_then_by_luid() {
845        let address = |bus| {
846            Some(PciAddress {
847                domain: 0,
848                bus,
849                device: 0,
850                function: 0,
851            })
852        };
853        let luid = |low| Some(AdapterLuid::from_parts(low, 0));
854        let card = |pci_address, luid| PhysicalDevice {
855            pci_address,
856            luid,
857            registry_entry_id: None,
858            vendor: None,
859        };
860
861        assert!(card(address(7), None).is_same_card(&card(address(7), luid(1))));
862        assert!(!card(address(7), luid(1)).is_same_card(&card(address(8), luid(1))));
863        assert!(card(None, luid(1)).is_same_card(&card(address(7), luid(1))));
864        assert!(!card(None, luid(1)).is_same_card(&card(None, luid(2))));
865        assert!(!card(None, None).is_same_card(&card(None, None)));
866    }
867
868    #[test]
869    fn a_card_filled_from_another_runtime_keeps_what_it_reported() {
870        let address = |bus| {
871            Some(PciAddress {
872                domain: 0,
873                bus,
874                device: 0,
875                function: 0,
876            })
877        };
878        let luid = Some(AdapterLuid::from_parts(1, 0));
879        let mut card = PhysicalDevice {
880            pci_address: address(7),
881            luid: None,
882            registry_entry_id: None,
883            vendor: None,
884        };
885
886        card.fill_from(&PhysicalDevice {
887            pci_address: address(8),
888            luid,
889            registry_entry_id: None,
890            vendor: Some(PciVendor::Nvidia),
891        });
892
893        assert_eq!(
894            card,
895            PhysicalDevice {
896                pci_address: address(7),
897                luid,
898                registry_entry_id: None,
899                vendor: Some(PciVendor::Nvidia),
900            }
901        );
902    }
903
904    #[test]
905    fn a_card_with_neither_address_nor_luid_is_matched_by_registry_entry_id() {
906        let card = |registry_entry_id| PhysicalDevice {
907            pci_address: None,
908            luid: None,
909            registry_entry_id: Some(RegistryEntryId::new(registry_entry_id)),
910            vendor: None,
911        };
912
913        assert!(card(0x1_0000_04c8).is_same_card(&card(0x1_0000_04c8)));
914        assert!(!card(0x1_0000_04c8).is_same_card(&card(0x1_0000_04c9)));
915    }
916
917    #[test]
918    fn the_address_and_the_luid_outrank_the_registry_entry_id() {
919        let address = |bus| {
920            Some(PciAddress {
921                domain: 0,
922                bus,
923                device: 0,
924                function: 0,
925            })
926        };
927        let luid = |low| Some(AdapterLuid::from_parts(low, 0));
928        let card = |pci_address, luid| PhysicalDevice {
929            pci_address,
930            luid,
931            registry_entry_id: Some(RegistryEntryId::new(0x1_0000_04c8)),
932            vendor: None,
933        };
934
935        assert!(!card(address(7), None).is_same_card(&card(address(8), None)));
936        assert!(!card(None, luid(1)).is_same_card(&card(None, luid(2))));
937    }
938
939    #[test]
940    fn a_card_filled_from_another_runtime_takes_the_registry_entry_id_it_left_out() {
941        let registry_entry_id = |id| Some(RegistryEntryId::new(id));
942        let card = |registry_entry_id| PhysicalDevice {
943            pci_address: None,
944            luid: None,
945            registry_entry_id,
946            vendor: None,
947        };
948
949        let mut without_one = card(None);
950        without_one.fill_from(&card(registry_entry_id(0x1_0000_04c8)));
951        let mut with_one = card(registry_entry_id(0x1_0000_04c8));
952        with_one.fill_from(&card(registry_entry_id(0x1_0000_04c9)));
953
954        assert_eq!(
955            without_one.registry_entry_id,
956            registry_entry_id(0x1_0000_04c8)
957        );
958        assert_eq!(with_one.registry_entry_id, registry_entry_id(0x1_0000_04c8));
959    }
960
961    #[test]
962    fn a_pci_address_round_trips_and_defaults_its_domain() {
963        let id = PciAddress {
964            domain: 0,
965            bus: 7,
966            device: 0,
967            function: 0,
968        };
969        assert_eq!(id.to_string(), "0000:07:00.0");
970        assert_eq!("0000:07:00.0".parse::<PciAddress>(), Ok(id));
971        assert_eq!("07:00.0".parse::<PciAddress>(), Ok(id));
972        assert_eq!(
973            "0001:a3:1f.7".parse::<PciAddress>(),
974            Ok(PciAddress {
975                domain: 1,
976                bus: 0xa3,
977                device: 0x1f,
978                function: 7
979            })
980        );
981        assert!("07:00".parse::<PciAddress>().is_err());
982        assert!("gpu".parse::<PciAddress>().is_err());
983    }
984
985    const AVX2: (u32, u32) = (256, 16);
986    const AVX512: (u32, u32) = (512, 32);
987    const NEON: (u32, u32) = (128, 32);
988
989    fn hardware(load_width: u32, vector_register_count: Option<u32>) -> HardwareProperties {
990        HardwareProperties {
991            load_width,
992            vector_register_count,
993            plane_size_min: 1,
994            plane_size_max: 1,
995            max_bindings: u32::MAX,
996            max_shared_memory_size: 32 * 1024,
997            max_cube_count: (u32::MAX, u32::MAX, u32::MAX),
998            max_units_per_cube: 16,
999            max_cube_dim: (16, 16, 16),
1000            num_streaming_multiprocessors: None,
1001            num_cpu_cores: Some(16),
1002            last_level_cache_size: None,
1003            num_tensor_cores: None,
1004            min_tensor_cores_dim: None,
1005            max_vector_size: VectorSize::MAX,
1006            cube_mma_reserved_shared_memory: 0,
1007        }
1008    }
1009
1010    fn registers((width, count): (u32, u32), elem_size: usize) -> VectorRegisters {
1011        VectorRegisters::new(&hardware(width, Some(count)), elem_size).unwrap()
1012    }
1013
1014    fn device(load_width: u32) -> DeviceProperties {
1015        DeviceProperties::new(
1016            Features::default(),
1017            MemoryDeviceProperties::new(1024, 32),
1018            hardware(load_width, None),
1019            TimingMethod::System,
1020            DeviceIdentity::default(),
1021        )
1022    }
1023}