Skip to main content

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