burn_std/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! # Burn Standard Library
6//!
7//! This library contains core types and utilities shared across Burn, including shapes, indexing,
8//! and data types.
9
10extern crate alloc;
11
12/// Id module contains types for unique identifiers.
13pub mod id;
14
15/// Tensor utilities.
16pub mod tensor;
17pub use tensor::*;
18
19/// Burn runtime configurations.
20pub mod config;
21
22/// Common Errors.
23pub use cubecl_zspace::errors::{self, *};
24
25/// Network utilities.
26#[cfg(feature = "network")]
27pub mod network;
28
29// Re-exported types
30#[cfg(feature = "cubecl")]
31pub use cubecl::server::CommunicationId;
32pub use cubecl_common::bytes::*;
33pub use cubecl_common::device_handle::DeviceHandle;
34pub use cubecl_common::*;
35pub use half::{bf16, f16};
36
37#[cfg(feature = "cubecl")]
38pub use cubecl::flex32;
39
40#[cfg(feature = "cubecl")]
41mod cube {
42    use cubecl::ir::{ElemType, FloatKind, IntKind, StorageType, UIntKind};
43    use cubecl_common::quant::scheme::QuantScheme;
44
45    use crate::tensor::DType;
46    use crate::tensor::quantization::{QuantStore, QuantValue};
47
48    impl From<DType> for cubecl::ir::ElemType {
49        fn from(dtype: DType) -> Self {
50            match dtype {
51                DType::F64 => ElemType::Float(FloatKind::F64),
52                DType::F32 => ElemType::Float(FloatKind::F32),
53                DType::Flex32 => ElemType::Float(FloatKind::Flex32),
54                DType::F16 => ElemType::Float(FloatKind::F16),
55                DType::BF16 => ElemType::Float(FloatKind::BF16),
56                DType::I64 => ElemType::Int(IntKind::I64),
57                DType::I32 => ElemType::Int(IntKind::I32),
58                DType::I16 => ElemType::Int(IntKind::I16),
59                DType::I8 => ElemType::Int(IntKind::I8),
60                DType::U64 => ElemType::UInt(UIntKind::U64),
61                DType::U32 => ElemType::UInt(UIntKind::U32),
62                DType::U16 => ElemType::UInt(UIntKind::U16),
63                DType::U8 => ElemType::UInt(UIntKind::U8),
64                DType::Bool(store) => match store {
65                    crate::BoolStore::Native => ElemType::Bool,
66                    crate::BoolStore::U8 => ElemType::UInt(UIntKind::U8),
67                    crate::BoolStore::U32 => ElemType::UInt(UIntKind::U32),
68                },
69                DType::QFloat(scheme) => match scheme.store {
70                    QuantStore::Native => match scheme.value {
71                        QuantValue::Q8F | QuantValue::Q8S => Self::Int(IntKind::I8),
72                        QuantValue::E4M3 => Self::Float(FloatKind::E4M3),
73                        QuantValue::E5M2 => Self::Float(FloatKind::E5M2),
74                        QuantValue::Q4F
75                        | QuantValue::Q4S
76                        | QuantValue::Q2F
77                        | QuantValue::Q2S
78                        | QuantValue::E2M1 => {
79                            panic!("Can't store native sub-byte values")
80                        }
81                    },
82                    QuantStore::PackedU32(_) => Self::UInt(UIntKind::U32),
83                    QuantStore::PackedNative(_) => match scheme.value {
84                        QuantValue::E2M1 => panic!("Can't store native sub-byte values"),
85                        other => panic!("{other:?} doesn't support native packing"),
86                    },
87                },
88            }
89        }
90    }
91
92    impl From<DType> for cubecl::ir::StorageType {
93        fn from(dtype: DType) -> cubecl::ir::StorageType {
94            match dtype {
95                DType::QFloat(QuantScheme {
96                    store: QuantStore::PackedNative(_),
97                    value: QuantValue::E2M1,
98                    ..
99                }) => StorageType::Packed(ElemType::Float(FloatKind::E2M1), 2),
100                _ => {
101                    let elem: ElemType = dtype.into();
102                    elem.into()
103                }
104            }
105        }
106    }
107}
burn_std/lib.rs

burn_std/
lib.rs
