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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
30pub struct HardwareProperties {
31 pub load_width: u32,
33 pub vector_register_count: Option<u32>,
36 pub plane_size_min: u32,
38 pub plane_size_max: u32,
40 pub max_bindings: u32,
42 pub max_shared_memory_size: usize,
44 pub max_cube_count: (u32, u32, u32),
46 pub max_units_per_cube: u32,
48 pub max_cube_dim: (u32, u32, u32),
50 pub num_streaming_multiprocessors: Option<u32>,
52 pub num_cpu_cores: Option<u32>,
54 pub last_level_cache_size: Option<usize>,
60 pub num_tensor_cores: Option<u32>,
62 pub min_tensor_cores_dim: Option<u32>,
67 pub max_vector_size: VectorSize,
69 pub cube_mma_reserved_shared_memory: usize,
71}
72
73#[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 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 pub fn registers_for(&self, vector_size: VectorSize) -> usize {
97 vector_size.div_ceil(self.register_vector_size).max(1)
98 }
99
100 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 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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
114#[non_exhaustive]
115pub struct MemoryDeviceProperties {
116 pub max_page_size: u64,
118 pub alignment: u64,
120 max_memory: Option<u64>,
123}
124
125impl MemoryDeviceProperties {
126 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 pub const fn max_memory(&self) -> Option<u64> {
148 self.max_memory
149 }
150
151 pub const fn with_max_memory(mut self, max_memory: u64) -> Self {
153 self.set_max_memory(max_memory);
154 self
155 }
156
157 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#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
182pub struct DeviceIdentity {
183 pub name: String,
187 pub fingerprint: String,
191 pub physical: Option<PhysicalDevice>,
194}
195
196#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
199#[non_exhaustive]
200pub struct PhysicalDevice {
201 pub pci_address: Option<PciAddress>,
203 pub luid: Option<AdapterLuid>,
205 pub registry_entry_id: Option<RegistryEntryId>,
207 pub vendor: Option<PciVendor>,
208}
209
210impl PhysicalDevice {
211 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 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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
261pub struct AdapterLuid([u8; 8]);
262
263impl AdapterLuid {
264 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
283pub enum PciVendor {
284 Nvidia,
285 Amd,
286 Intel,
287 Apple,
288 Arm,
290 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#[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 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#[derive(Debug, Clone, PartialEq, Eq)]
390pub struct DeviceProperties {
391 pub features: Features,
393 pub memory: MemoryDeviceProperties,
395 pub hardware: HardwareProperties,
397 pub timing_method: TimingMethod,
399 pub identity: DeviceIdentity,
401 pub io_width_override: Option<u32>,
403}
404
405impl TypeHash for DeviceProperties {
406 fn write_hash(_hasher: &mut impl core::hash::Hasher) {
407 }
409}
410
411impl DeviceProperties {
412 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 pub fn io_width(&self) -> u32 {
437 self.io_width_override.unwrap_or(self.hardware.load_width)
438 }
439
440 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 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 pub fn type_usage(&self, ty: ElemType) -> EnumSet<TypeUsage> {
464 self.features.type_usage(ty)
465 }
466
467 pub fn complex_usage(&self, ty: ElemType) -> EnumSet<ComplexUsage> {
469 self.features.complex_usage(ty)
470 }
471
472 pub fn supports_complex_usage(&self, ty: ElemType, usage: ComplexUsage) -> bool {
474 self.features.supports_complex_usage(ty, usage)
475 }
476
477 pub fn atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
479 self.features.atomic_type_usage(ty)
480 }
481
482 pub fn shared_atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
484 self.features.shared_atomic_type_usage(ty)
485 }
486
487 pub fn supports_type(&self, ty: impl Into<Type>) -> bool {
489 self.features.supports_type(ty)
490 }
491
492 pub fn supports_address(&self, ty: impl Into<AddressType>) -> bool {
494 self.features.supports_address(ty)
495 }
496
497 pub fn register_address_type(&mut self, ty: impl Into<AddressType>) {
499 self.features.types.address.insert(ty.into());
500 }
501
502 pub fn register_atomic_type_usage(&mut self, ty: Type, uses: impl Into<EnumSet<AtomicUsage>>) {
504 *self.features.types.atomic.entry(ty).or_default() |= uses.into();
505 }
506
507 pub fn register_shared_atomic_type_usage(
509 &mut self,
510 ty: Type,
511 uses: impl Into<EnumSet<AtomicUsage>>,
512 ) {
513 *self.features.types.shared_atomic.entry(ty).or_default() |= uses.into();
514 }
515
516 pub fn register_type_usage(
518 &mut self,
519 ty: impl Into<ElemType>,
520 uses: impl Into<EnumSet<TypeUsage>>,
521 ) {
522 *self.features.types.elem.entry(ty.into()).or_default() |= uses.into();
523 }
524
525 pub fn register_emulated_conversion(&mut self, ty: impl Into<ElemType>) {
529 self.features.types.emulated_conversion.insert(ty.into());
530 }
531
532 pub fn register_emulated_arithmetic(&mut self, ty: impl Into<ElemType>) {
537 self.features.types.emulated_arithmetic.insert(ty.into());
538 }
539
540 pub fn register_complex_usage(
542 &mut self,
543 ty: impl Into<ElemType>,
544 uses: impl Into<EnumSet<ComplexUsage>>,
545 ) {
546 *self.features.types.complex.entry(ty.into()).or_default() |= uses.into();
547 }
548
549 pub fn register_semantic_type(&mut self, ty: SemanticType) {
551 self.features.types.semantic.insert(ty);
552 }
553
554 pub fn register_opaque_type(&mut self, ty: OpaqueType) {
556 self.features.types.opaque.insert(ty);
557 }
558
559 pub fn checksum(&self) -> u64 {
562 let state = foldhash::fast::FixedState::default();
563 let mut hasher = state.build_hasher();
564 self.features.hash(&mut hasher);
565 self.hardware.hash(&mut hasher);
566 hasher.finish()
567 }
568
569 fn widest_io_vector_size(&self, elem_size: usize) -> VectorSize {
571 self.hardware.vector_size_in(self.io_width(), elem_size)
572 }
573}
574
575#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
578#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, EnumSetType)]
579pub enum FastMath {
580 NotNaN,
582 NotInf,
585 UnsignedZero,
587 AllowReciprocal,
589 AllowContraction,
592 AllowReassociation,
594 AllowTransform,
597 ReducedPrecision,
599}
600
601impl FastMath {
602 pub const fn all() -> EnumSet<FastMath> {
603 EnumSet::all()
604 }
605}
606
607impl HardwareProperties {
608 fn vector_size_in(&self, width: u32, elem_size: usize) -> VectorSize {
610 let elems = width as usize / (elem_size * 8);
611 prev_power_of_two(elems.min(self.max_vector_size).max(1))
612 }
613}
614
615fn prev_power_of_two(value: usize) -> usize {
616 1 << value.ilog2()
617}
618
619fn vector_sizes_down_from(widest: usize) -> impl Iterator<Item = VectorSize> + Clone {
620 (0..=widest.ilog2()).rev().map(|exponent| 1 << exponent)
621}
622
623#[cfg(test)]
624mod tests {
625 use super::*;
626 use alloc::string::ToString;
627
628 #[test]
629 fn a_device_without_a_fixed_register_set_has_no_budget() {
630 assert_eq!(VectorRegisters::new(&hardware(128, None), 4), None);
631 }
632
633 #[test]
634 fn the_widest_vector_size_keeps_every_live_vector_in_registers() {
635 let f32 = registers(AVX2, 4);
636 assert_eq!(f32.widest_vector_size(6), 16);
637 assert_eq!(f32.widest_vector_size(8), 16);
638 assert_eq!(f32.widest_vector_size(9), 8);
639 assert_eq!(registers(AVX2, 8).widest_vector_size(6), 8);
640 assert_eq!(registers(AVX512, 4).widest_vector_size(6), 64);
641 }
642
643 #[test]
644 fn neon_and_avx2_budget_the_same_vector_sizes_from_equal_register_files() {
645 let (neon, avx2) = (registers(NEON, 4), registers(AVX2, 4));
646 assert_eq!(neon.widest_vector_size(3), avx2.widest_vector_size(3));
647 assert_eq!(neon.widest_vector_size(6), avx2.widest_vector_size(6));
648 }
649
650 #[test]
651 fn more_live_vectors_than_registers_still_get_one_register_each() {
652 assert_eq!(registers(AVX2, 4).widest_vector_size(40), 8);
653 assert_eq!(registers(NEON, 4).widest_vector_size(40), 4);
654 }
655
656 #[test]
657 fn vectors_fit_in_the_registers_left_unreserved() {
658 let f32 = registers(AVX2, 4);
659 assert_eq!(f32.registers_for(16), 2);
660 assert_eq!(f32.vectors_fitting(16, 0), 8);
661 assert_eq!(f32.vectors_fitting(8, 4), 12);
662 assert_eq!(f32.vectors_fitting(8, 20), 0);
663 }
664
665 #[test]
666 fn a_capped_vector_size_caps_the_register_vector_size() {
667 let mut capped = hardware(256, Some(16));
668 capped.max_vector_size = 4;
669 let f32 = VectorRegisters::new(&capped, 4).unwrap();
670 assert_eq!(f32.registers_for(8), 2);
671 assert_eq!(f32.widest_vector_size(1), 4);
672 }
673
674 #[test]
675 fn io_follows_a_load_width_raised_after_construction() {
676 let mut device = device(128);
677 device.hardware.load_width = 256;
678 assert_eq!(device.io_width(), 256);
679 }
680
681 #[test]
682 fn a_stated_io_width_outlasts_a_load_width_change() {
683 let mut device = device(256);
684 device.io_width_override = Some(512);
685 device.hardware.load_width = 128;
686 assert_eq!(device.io_width(), 512);
687 }
688
689 #[test]
690 fn a_cap_between_powers_of_two_rounds_down_alike() {
691 let mut device = device(512);
692 device.hardware.max_vector_size = 12;
693 device.hardware.vector_register_count = Some(32);
694
695 assert_eq!(
696 device
697 .io_optimized_vector_sizes(4)
698 .collect::<alloc::vec::Vec<_>>(),
699 alloc::vec![8, 4, 2, 1]
700 );
701 assert_eq!(device.vector_sizes_in_registers(4, 1).next(), Some(8));
702 }
703
704 #[test]
705 fn a_device_without_registers_sizes_by_its_io() {
706 let device = device(128);
707 assert!(
708 device
709 .vector_sizes_in_registers(4, 6)
710 .eq(device.io_optimized_vector_sizes(4))
711 );
712 }
713
714 #[test]
719 fn a_memory_states_no_capacity_until_one_is_read() {
720 let props = MemoryDeviceProperties::new(1024, 32);
721 assert_eq!(props.max_memory(), None);
722 assert_eq!(props.clone().with_max_memory(4096).max_memory(), Some(4096));
723 }
724
725 #[test]
730 fn a_capacity_of_zero_is_no_capacity() {
731 let mut props = MemoryDeviceProperties::new(1024, 32).with_max_memory(4096);
732 props.set_max_memory(0);
733 assert_eq!(props.max_memory(), None);
734 }
735
736 #[test]
737 fn a_vendor_keeps_its_id_whether_named_or_not() {
738 for id in [0x10de, 0x1002, 0x8086, 0x106b, 0x13b5, 0x5143, 0x1af4] {
739 assert_eq!(PciVendor::from(id).id(), id);
740 }
741 assert_eq!(PciVendor::from(0x10de), PciVendor::Nvidia);
742 assert_eq!(PciVendor::from(0x1af4), PciVendor::Other(0x1af4));
743 assert_eq!(PciVendor::Other(0x1af4).to_string(), "0x1af4");
744 }
745
746 #[test]
747 fn a_luid_from_parts_is_the_bytes_vulkan_reports() {
748 let bytes = [0x8a, 0x1d, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00];
749 assert_eq!(
750 AdapterLuid::from_parts(0x0001_1d8a, 0),
751 AdapterLuid::new(bytes)
752 );
753 assert_eq!(
754 AdapterLuid::from_parts(1, -1).bytes(),
755 [1, 0, 0, 0, 0xff, 0xff, 0xff, 0xff]
756 );
757 }
758
759 #[test]
760 fn a_card_is_matched_by_address_then_by_luid() {
761 let address = |bus| {
762 Some(PciAddress {
763 domain: 0,
764 bus,
765 device: 0,
766 function: 0,
767 })
768 };
769 let luid = |low| Some(AdapterLuid::from_parts(low, 0));
770 let card = |pci_address, luid| PhysicalDevice {
771 pci_address,
772 luid,
773 registry_entry_id: None,
774 vendor: None,
775 };
776
777 assert!(card(address(7), None).is_same_card(&card(address(7), luid(1))));
778 assert!(!card(address(7), luid(1)).is_same_card(&card(address(8), luid(1))));
779 assert!(card(None, luid(1)).is_same_card(&card(address(7), luid(1))));
780 assert!(!card(None, luid(1)).is_same_card(&card(None, luid(2))));
781 assert!(!card(None, None).is_same_card(&card(None, None)));
782 }
783
784 #[test]
785 fn a_card_filled_from_another_runtime_keeps_what_it_reported() {
786 let address = |bus| {
787 Some(PciAddress {
788 domain: 0,
789 bus,
790 device: 0,
791 function: 0,
792 })
793 };
794 let luid = Some(AdapterLuid::from_parts(1, 0));
795 let mut card = PhysicalDevice {
796 pci_address: address(7),
797 luid: None,
798 registry_entry_id: None,
799 vendor: None,
800 };
801
802 card.fill_from(&PhysicalDevice {
803 pci_address: address(8),
804 luid,
805 registry_entry_id: None,
806 vendor: Some(PciVendor::Nvidia),
807 });
808
809 assert_eq!(
810 card,
811 PhysicalDevice {
812 pci_address: address(7),
813 luid,
814 registry_entry_id: None,
815 vendor: Some(PciVendor::Nvidia),
816 }
817 );
818 }
819
820 #[test]
821 fn a_card_with_neither_address_nor_luid_is_matched_by_registry_entry_id() {
822 let card = |registry_entry_id| PhysicalDevice {
823 pci_address: None,
824 luid: None,
825 registry_entry_id: Some(RegistryEntryId::new(registry_entry_id)),
826 vendor: None,
827 };
828
829 assert!(card(0x1_0000_04c8).is_same_card(&card(0x1_0000_04c8)));
830 assert!(!card(0x1_0000_04c8).is_same_card(&card(0x1_0000_04c9)));
831 }
832
833 #[test]
834 fn the_address_and_the_luid_outrank_the_registry_entry_id() {
835 let address = |bus| {
836 Some(PciAddress {
837 domain: 0,
838 bus,
839 device: 0,
840 function: 0,
841 })
842 };
843 let luid = |low| Some(AdapterLuid::from_parts(low, 0));
844 let card = |pci_address, luid| PhysicalDevice {
845 pci_address,
846 luid,
847 registry_entry_id: Some(RegistryEntryId::new(0x1_0000_04c8)),
848 vendor: None,
849 };
850
851 assert!(!card(address(7), None).is_same_card(&card(address(8), None)));
852 assert!(!card(None, luid(1)).is_same_card(&card(None, luid(2))));
853 }
854
855 #[test]
856 fn a_card_filled_from_another_runtime_takes_the_registry_entry_id_it_left_out() {
857 let registry_entry_id = |id| Some(RegistryEntryId::new(id));
858 let card = |registry_entry_id| PhysicalDevice {
859 pci_address: None,
860 luid: None,
861 registry_entry_id,
862 vendor: None,
863 };
864
865 let mut without_one = card(None);
866 without_one.fill_from(&card(registry_entry_id(0x1_0000_04c8)));
867 let mut with_one = card(registry_entry_id(0x1_0000_04c8));
868 with_one.fill_from(&card(registry_entry_id(0x1_0000_04c9)));
869
870 assert_eq!(
871 without_one.registry_entry_id,
872 registry_entry_id(0x1_0000_04c8)
873 );
874 assert_eq!(with_one.registry_entry_id, registry_entry_id(0x1_0000_04c8));
875 }
876
877 #[test]
878 fn a_pci_address_round_trips_and_defaults_its_domain() {
879 let id = PciAddress {
880 domain: 0,
881 bus: 7,
882 device: 0,
883 function: 0,
884 };
885 assert_eq!(id.to_string(), "0000:07:00.0");
886 assert_eq!("0000:07:00.0".parse::<PciAddress>(), Ok(id));
887 assert_eq!("07:00.0".parse::<PciAddress>(), Ok(id));
888 assert_eq!(
889 "0001:a3:1f.7".parse::<PciAddress>(),
890 Ok(PciAddress {
891 domain: 1,
892 bus: 0xa3,
893 device: 0x1f,
894 function: 7
895 })
896 );
897 assert!("07:00".parse::<PciAddress>().is_err());
898 assert!("gpu".parse::<PciAddress>().is_err());
899 }
900
901 const AVX2: (u32, u32) = (256, 16);
902 const AVX512: (u32, u32) = (512, 32);
903 const NEON: (u32, u32) = (128, 32);
904
905 fn hardware(load_width: u32, vector_register_count: Option<u32>) -> HardwareProperties {
906 HardwareProperties {
907 load_width,
908 vector_register_count,
909 plane_size_min: 1,
910 plane_size_max: 1,
911 max_bindings: u32::MAX,
912 max_shared_memory_size: 32 * 1024,
913 max_cube_count: (u32::MAX, u32::MAX, u32::MAX),
914 max_units_per_cube: 16,
915 max_cube_dim: (16, 16, 16),
916 num_streaming_multiprocessors: None,
917 num_cpu_cores: Some(16),
918 last_level_cache_size: None,
919 num_tensor_cores: None,
920 min_tensor_cores_dim: None,
921 max_vector_size: VectorSize::MAX,
922 cube_mma_reserved_shared_memory: 0,
923 }
924 }
925
926 fn registers((width, count): (u32, u32), elem_size: usize) -> VectorRegisters {
927 VectorRegisters::new(&hardware(width, Some(count)), elem_size).unwrap()
928 }
929
930 fn device(load_width: u32) -> DeviceProperties {
931 DeviceProperties::new(
932 Features::default(),
933 MemoryDeviceProperties::new(1024, 32),
934 hardware(load_width, None),
935 TimingMethod::System,
936 DeviceIdentity::default(),
937 )
938 }
939}