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
use crate::memory_management::{HardwareProperties, MemoryDeviceProperties};
use alloc::collections::BTreeSet;
use cubecl_common::profile::TimingMethod;
/// 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>,
/// The memory properties of this client.
pub memory: MemoryDeviceProperties,
/// The topology properties of this client.
pub hardware: HardwareProperties,
/// The method used for profiling on the device.
pub timing_method: TimingMethod,
}
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,
timing_method: TimingMethod,
) -> Self {
let mut set = BTreeSet::new();
for feature in features {
set.insert(*feature);
}
DeviceProperties {
set,
memory: memory_props,
hardware,
timing_method,
}
}
/// 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)
}
/// Removes a `Feature` from the compute server.
///
/// This should only be used by a [runtime](cubecl_core::Runtime) when initializing a device.
pub fn remove_feature(&mut self, feature: Feature) {
self.set.remove(&feature);
}
}