1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
use crate::memory_management::{HardwareProperties, MemoryDeviceProperties};
use alloc::collections::BTreeSet;
/// Properties of what the device can do, like what `Feature` are
/// supported by it and what its memory properties are.
#[derive(Debug)]
pub struct DeviceProperties<Feature: Ord + Copy> {
set: alloc::collections::BTreeSet<Feature>,
memory: MemoryDeviceProperties,
hardware: HardwareProperties,
time_measurement: TimeMeasurement,
}
#[derive(Debug, Clone, Copy)]
/// How times are measured on a device.
pub enum TimeMeasurement {
/// Using the device own measuting capability.
///
/// Normally compatible with async.
Device,
/// Using [std::time::Instant] to measure kernel execution.
///
/// When this version is activated, we must await on async tasks.
System,
}
impl<Feature: Ord + Copy> DeviceProperties<Feature> {
/// Create a new feature set with the given features and memory properties.
pub fn new(
features: &[Feature],
memory_props: MemoryDeviceProperties,
hardware: HardwareProperties,
time_measurement: TimeMeasurement,
) -> Self {
let mut set = BTreeSet::new();
for feature in features {
set.insert(*feature);
}
DeviceProperties {
set,
memory: memory_props,
hardware,
time_measurement,
}
}
/// Get the [time measurement](TimeMeasurement) of the current device.
pub fn time_measurement(&self) -> TimeMeasurement {
self.time_measurement
}
/// Check if the provided `Feature` is supported by the runtime.
pub fn feature_enabled(&self, feature: Feature) -> bool {
self.set.contains(&feature)
}
/// Register a `Feature` supported by the compute server.
///
/// This should only be used by a [runtime](cubecl_core::Runtime) when initializing a device.
pub fn register_feature(&mut self, feature: Feature) -> bool {
self.set.insert(feature)
}
/// The memory properties of this client.
pub fn memory_properties(&self) -> &MemoryDeviceProperties {
&self.memory
}
/// The topology properties of this client.
pub fn hardware_properties(&self) -> &HardwareProperties {
&self.hardware
}
}