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#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub struct HardwareProperties {
32 pub load_width: u32,
34 pub vector_register_count: Option<u32>,
37 pub plane_size_min: u32,
39 pub plane_size_max: u32,
41 pub max_bindings: u32,
43 pub max_shared_memory_size: usize,
45 pub max_cube_count: (u32, u32, u32),
47 pub max_units_per_cube: u32,
49 pub max_cube_dim: (u32, u32, u32),
51 pub num_streaming_multiprocessors: Option<u32>,
53 pub num_cpu_cores: Option<u32>,
55 pub last_level_cache_size: Option<usize>,
61 pub num_tensor_cores: Option<u32>,
63 pub min_tensor_cores_dim: Option<u32>,
68 pub max_vector_size: VectorSize,
70 pub cube_mma_reserved_shared_memory: usize,
72}
73
74#[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 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 pub fn registers_for(&self, vector_size: VectorSize) -> usize {
98 vector_size.div_ceil(self.register_vector_size).max(1)
99 }
100
101 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 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#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
115#[derive(Debug, Clone, PartialEq, Eq, Hash)]
116#[non_exhaustive]
117pub struct MemoryDeviceProperties {
118 pub max_page_size: u64,
120 pub alignment: u64,
122 #[cfg_attr(feature = "serde", serde(deserialize_with = "deserialize_max_memory"))]
126 max_memory: Option<u64>,
127}
128
129#[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 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 pub const fn max_memory(&self) -> Option<u64> {
161 self.max_memory
162 }
163
164 pub const fn with_max_memory(mut self, max_memory: u64) -> Self {
166 self.set_max_memory(max_memory);
167 self
168 }
169
170 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#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
195#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
196pub struct DeviceIdentity {
197 pub name: String,
201 pub fingerprint: String,
205 #[cfg_attr(feature = "serde", serde(skip))]
211 pub physical: Option<PhysicalDevice>,
212}
213
214#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
217#[non_exhaustive]
218pub struct PhysicalDevice {
219 pub pci_address: Option<PciAddress>,
221 pub luid: Option<AdapterLuid>,
223 pub registry_entry_id: Option<RegistryEntryId>,
225 pub vendor: Option<PciVendor>,
226}
227
228impl PhysicalDevice {
229 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 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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
279pub struct AdapterLuid([u8; 8]);
280
281impl AdapterLuid {
282 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
301pub enum PciVendor {
302 Nvidia,
303 Amd,
304 Intel,
305 Apple,
306 Arm,
308 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#[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 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#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
408#[derive(Debug, Clone, PartialEq, Eq)]
409pub struct DeviceProperties {
410 pub features: Features,
412 pub memory: MemoryDeviceProperties,
414 pub hardware: HardwareProperties,
416 pub timing_method: TimingMethod,
418 pub identity: DeviceIdentity,
420 pub io_width_override: Option<u32>,
422}
423
424impl TypeHash for DeviceProperties {
425 fn write_hash(_hasher: &mut impl core::hash::Hasher) {
426 }
428}
429
430impl DeviceProperties {
431 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 pub fn io_width(&self) -> u32 {
456 self.io_width_override.unwrap_or(self.hardware.load_width)
457 }
458
459 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 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 pub fn type_usage(&self, ty: ElemType) -> EnumSet<TypeUsage> {
483 self.features.type_usage(ty)
484 }
485
486 pub fn complex_usage(&self, ty: ElemType) -> EnumSet<ComplexUsage> {
488 self.features.complex_usage(ty)
489 }
490
491 pub fn supports_complex_usage(&self, ty: ElemType, usage: ComplexUsage) -> bool {
493 self.features.supports_complex_usage(ty, usage)
494 }
495
496 pub fn atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
498 self.features.atomic_type_usage(ty)
499 }
500
501 pub fn shared_atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
503 self.features.shared_atomic_type_usage(ty)
504 }
505
506 pub fn supports_type(&self, ty: impl Into<Type>) -> bool {
508 self.features.supports_type(ty)
509 }
510
511 pub fn supports_address(&self, ty: impl Into<AddressType>) -> bool {
513 self.features.supports_address(ty)
514 }
515
516 pub fn register_address_type(&mut self, ty: impl Into<AddressType>) {
518 self.features.types.address.insert(ty.into());
519 }
520
521 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 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 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 pub fn register_emulated_conversion(&mut self, ty: impl Into<ElemType>) {
548 self.features.types.emulated_conversion.insert(ty.into());
549 }
550
551 pub fn register_emulated_arithmetic(&mut self, ty: impl Into<ElemType>) {
556 self.features.types.emulated_arithmetic.insert(ty.into());
557 }
558
559 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 pub fn register_semantic_type(&mut self, ty: SemanticType) {
570 self.features.types.semantic.insert(ty);
571 }
572
573 pub fn register_opaque_type(&mut self, ty: OpaqueType) {
575 self.features.types.opaque.insert(ty);
576 }
577
578 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 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#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
597#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, EnumSetType)]
598pub enum FastMath {
599 NotNaN,
601 NotInf,
604 UnsignedZero,
606 AllowReciprocal,
608 AllowContraction,
611 AllowReassociation,
613 AllowTransform,
616 ReducedPrecision,
618}
619
620impl FastMath {
621 pub const fn all() -> EnumSet<FastMath> {
622 EnumSet::all()
623 }
624}
625
626impl HardwareProperties {
627 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 #[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 #[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 #[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 #[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 #[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}