Skip to main content

burn_std/
device_settings.rs

1//! Device data types: settings and errors.
2
3use alloc::format;
4use alloc::string::String;
5use cubecl_common::backtrace::BackTrace;
6use serde::{Deserialize, Serialize};
7use thiserror::Error;
8
9use crate::{BoolDType, DType, FloatDType, IntDType, QuantConfig};
10
11/// Settings controlling the default data types for a specific device.
12///
13/// These settings are managed in a global registry that enforces strict initialization semantics:
14///
15/// 1. Manual Initialization: You can set these once at the start of your program using `set_default_dtypes`.
16/// 2. Default Initialization: If an operation (like creating a tensor) occurs before manual initialization,
17///    the settings are permanently locked to their default values.
18/// 3. Immutability: Once initialized, settings cannot be changed. This ensures consistent behavior across
19///    all threads and operations.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
21pub struct DeviceSettings {
22    /// Default floating-point data type.
23    pub float_dtype: FloatDType,
24    /// Default integer data type.
25    pub int_dtype: IntDType,
26    /// Default bool data type.
27    pub bool_dtype: BoolDType,
28    /// Quantization configuration.
29    pub quantization: QuantConfig,
30}
31
32impl DeviceSettings {
33    /// Creates a new [`DeviceSettings`] from any types convertible into the dtype kinds and the [quantization config](QuantConfig).
34    pub fn new(
35        float_dtype: impl Into<FloatDType>,
36        int_dtype: impl Into<IntDType>,
37        bool_dtype: impl Into<BoolDType>,
38        quantization: QuantConfig,
39    ) -> Self {
40        Self {
41            float_dtype: float_dtype.into(),
42            int_dtype: int_dtype.into(),
43            bool_dtype: bool_dtype.into(),
44            quantization,
45        }
46    }
47
48    /// Creates a new [`DeviceSettings`] from any types convertible into the dtype kinds.
49    pub fn with_dtypes(
50        float_dtype: impl Into<FloatDType>,
51        int_dtype: impl Into<IntDType>,
52        bool_dtype: impl Into<BoolDType>,
53    ) -> Self {
54        Self::new(float_dtype, int_dtype, bool_dtype, Default::default())
55    }
56}
57
58/// Errors returned by device-related operations.
59///
60/// Examples include attempting to use an unsupported data type on a device or initialization
61/// errors like attempting to change a settings in an invalid context.
62#[derive(Debug, Error)]
63pub enum DeviceError {
64    /// Unsupported data type by the device.
65    #[error("Device {device} does not support the requested data type {dtype:?}")]
66    UnsupportedDType {
67        /// The string representation of the device.
68        device: String,
69        /// The data type that caused the error.
70        dtype: DType,
71    },
72    /// Device settings have already been initialized.
73    #[error(
74        "Device {device} settings have already been initialized and cannot be changed.\n\
75         Default data types lock on the first tensor operation for a device (or on a previous \
76         `set_default_dtypes`/`configure` call). Configure the device before creating any tensor \
77         on it, or create tensors with an explicit dtype instead."
78    )]
79    AlreadyInitialized {
80        /// The string representation of the device.
81        device: String,
82    },
83}
84
85impl DeviceError {
86    /// Helper to create a [`DeviceError::UnsupportedDType`] from any device.
87    pub fn unsupported_dtype<D: core::fmt::Debug + ?Sized>(device: &D, dtype: DType) -> Self {
88        Self::UnsupportedDType {
89            device: format!("{device:?}"),
90            dtype,
91        }
92    }
93
94    /// Helper to create a [`DeviceError::AlreadyInitialized`] from any device.
95    pub fn already_initialized<D: core::fmt::Debug + ?Sized>(device: &D) -> Self {
96        Self::AlreadyInitialized {
97            device: format!("{device:?}"),
98        }
99    }
100}
101
102/// An error that can happen when syncing a device.
103#[derive(Error, Serialize, Deserialize)]
104pub enum ExecutionError {
105    /// A generic error happened during execution.
106    ///
107    /// The backtrace and context information should be included in the reason string.
108    #[error("An error happened during execution\nCaused by:\n  {reason}")]
109    WithContext {
110        /// The reason of the error.
111        reason: String,
112    },
113    /// A generic error happened during execution thrown in the Burn project.
114    ///
115    /// The full context isn't captured by the string alone.
116    #[error("An error happened during execution\nCaused by:\n  {reason}")]
117    Generic {
118        /// The reason of the error.
119        reason: String,
120        /// The backtrace.
121        #[serde(skip)]
122        backtrace: BackTrace,
123    },
124}
125
126impl core::fmt::Debug for ExecutionError {
127    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
128        f.write_fmt(format_args!("{self}"))
129    }
130}