cubecl_runtime/
feature_set.rs

1use crate::memory_management::{HardwareProperties, MemoryDeviceProperties};
2use alloc::collections::BTreeSet;
3
4/// Properties of what the device can do, like what [features](Feature) are
5/// supported by it and what its memory properties are.
6#[derive(Debug)]
7pub struct DeviceProperties<Feature: Ord + Copy> {
8    set: alloc::collections::BTreeSet<Feature>,
9    memory: MemoryDeviceProperties,
10    hardware: HardwareProperties,
11}
12
13impl<Feature: Ord + Copy> DeviceProperties<Feature> {
14    /// Create a new feature set with the given features and memory properties.
15    pub fn new(
16        features: &[Feature],
17        memory_props: MemoryDeviceProperties,
18        hardware: HardwareProperties,
19    ) -> Self {
20        let mut set = BTreeSet::new();
21        for feature in features {
22            set.insert(*feature);
23        }
24
25        DeviceProperties {
26            set,
27            memory: memory_props,
28            hardware,
29        }
30    }
31
32    /// Check if the provided [feature](Feature) is supported by the runtime.
33    pub fn feature_enabled(&self, feature: Feature) -> bool {
34        self.set.contains(&feature)
35    }
36
37    /// Register a [feature](Feature) supported by the compute server.
38    ///
39    /// This should only be used by a [runtime](Runtime) when initializing a device.
40    pub fn register_feature(&mut self, feature: Feature) -> bool {
41        self.set.insert(feature)
42    }
43
44    /// The memory properties of this client.
45    pub fn memory_properties(&self) -> &MemoryDeviceProperties {
46        &self.memory
47    }
48
49    /// The topology properties of this client.
50    pub fn hardware_properties(&self) -> &HardwareProperties {
51        &self.hardware
52    }
53}
cubecl_runtime/feature_set.rs

cubecl_runtime/
feature_set.rs
