cubecl_core/
pod.rs

1use cubecl_common::{e2m1, e2m1x2, e4m3, e5m2, flex32, tf32, ue8m0};
2
3use crate::{
4    ir::{ElemType, FloatKind, IntKind, UIntKind},
5    prelude::{Numeric, Scalar},
6};
7
8/// The base element trait for the jit backend.
9pub trait CubeElement: core::fmt::Debug + Send + Sync + 'static + Clone + bytemuck::Pod {
10    /// Returns the name of the type.
11    fn type_name() -> &'static str;
12    /// Convert a slice of elements to a slice of bytes.
13    fn as_bytes(slice: &[Self]) -> &[u8];
14    /// Convert a slice of bytes to a slice of elements.
15    fn from_bytes(bytes: &[u8]) -> &[Self];
16    /// Element representation for `cubecl`.
17    fn cube_type() -> ElemType;
18    /// Highest possible value
19    fn maximum_value() -> Self;
20    /// Lowest possible value
21    fn minimum_value() -> Self;
22}
23
24pub trait ScalarArgType: CubeElement + Scalar + num_traits::NumCast {}
25
26impl<E: CubeElement + Scalar + num_traits::NumCast> ScalarArgType for E {}
27
28impl CubeElement for u64 {
29    fn type_name() -> &'static str {
30        "u64"
31    }
32    fn as_bytes(slice: &[Self]) -> &[u8] {
33        bytemuck::cast_slice(slice)
34    }
35    fn from_bytes(bytes: &[u8]) -> &[Self] {
36        bytemuck::cast_slice(bytes)
37    }
38    fn cube_type() -> ElemType {
39        ElemType::UInt(UIntKind::U64)
40    }
41    fn maximum_value() -> Self {
42        u64::MAX
43    }
44    fn minimum_value() -> Self {
45        u64::MIN
46    }
47}
48
49impl CubeElement for u32 {
50    fn type_name() -> &'static str {
51        "u32"
52    }
53    fn as_bytes(slice: &[Self]) -> &[u8] {
54        bytemuck::cast_slice(slice)
55    }
56    fn from_bytes(bytes: &[u8]) -> &[Self] {
57        bytemuck::cast_slice(bytes)
58    }
59    fn cube_type() -> ElemType {
60        ElemType::UInt(UIntKind::U32)
61    }
62    fn maximum_value() -> Self {
63        u32::MAX
64    }
65    fn minimum_value() -> Self {
66        u32::MIN
67    }
68}
69
70impl CubeElement for u16 {
71    fn type_name() -> &'static str {
72        "u16"
73    }
74    fn as_bytes(slice: &[Self]) -> &[u8] {
75        bytemuck::cast_slice(slice)
76    }
77    fn from_bytes(bytes: &[u8]) -> &[Self] {
78        bytemuck::cast_slice(bytes)
79    }
80    fn cube_type() -> ElemType {
81        ElemType::UInt(UIntKind::U16)
82    }
83    fn maximum_value() -> Self {
84        u16::MAX
85    }
86    fn minimum_value() -> Self {
87        u16::MIN
88    }
89}
90
91impl CubeElement for u8 {
92    fn type_name() -> &'static str {
93        "u8"
94    }
95    fn as_bytes(slice: &[Self]) -> &[u8] {
96        bytemuck::cast_slice(slice)
97    }
98    fn from_bytes(bytes: &[u8]) -> &[Self] {
99        bytemuck::cast_slice(bytes)
100    }
101    fn cube_type() -> ElemType {
102        ElemType::UInt(UIntKind::U8)
103    }
104    fn maximum_value() -> Self {
105        u8::MAX
106    }
107    fn minimum_value() -> Self {
108        u8::MIN
109    }
110}
111
112impl CubeElement for i64 {
113    fn type_name() -> &'static str {
114        "i64"
115    }
116    fn as_bytes(slice: &[Self]) -> &[u8] {
117        bytemuck::cast_slice(slice)
118    }
119    fn from_bytes(bytes: &[u8]) -> &[Self] {
120        bytemuck::cast_slice(bytes)
121    }
122    fn cube_type() -> ElemType {
123        ElemType::Int(IntKind::I64)
124    }
125    fn maximum_value() -> Self {
126        // Seems to cause problem for some GPU
127        i64::MAX - 1
128    }
129    fn minimum_value() -> Self {
130        // Seems to cause problem for some GPU
131        i64::MIN + 1
132    }
133}
134
135impl CubeElement for i32 {
136    fn type_name() -> &'static str {
137        "i32"
138    }
139    fn as_bytes(slice: &[Self]) -> &[u8] {
140        bytemuck::cast_slice(slice)
141    }
142    fn from_bytes(bytes: &[u8]) -> &[Self] {
143        bytemuck::cast_slice(bytes)
144    }
145    fn cube_type() -> ElemType {
146        ElemType::Int(IntKind::I32)
147    }
148    fn maximum_value() -> Self {
149        // Seems to cause problem for some GPU
150        i32::MAX - 1
151    }
152    fn minimum_value() -> Self {
153        // Seems to cause problem for some GPU
154        i32::MIN + 1
155    }
156}
157
158impl CubeElement for i16 {
159    fn type_name() -> &'static str {
160        "i16"
161    }
162    fn as_bytes(slice: &[Self]) -> &[u8] {
163        bytemuck::cast_slice(slice)
164    }
165    fn from_bytes(bytes: &[u8]) -> &[Self] {
166        bytemuck::cast_slice(bytes)
167    }
168    fn cube_type() -> ElemType {
169        ElemType::Int(IntKind::I16)
170    }
171    fn maximum_value() -> Self {
172        // Seems to cause problem for some GPU
173        i16::MAX - 1
174    }
175    fn minimum_value() -> Self {
176        // Seems to cause problem for some GPU
177        i16::MIN + 1
178    }
179}
180
181impl CubeElement for i8 {
182    fn type_name() -> &'static str {
183        "i8"
184    }
185    fn as_bytes(slice: &[Self]) -> &[u8] {
186        bytemuck::cast_slice(slice)
187    }
188    fn from_bytes(bytes: &[u8]) -> &[Self] {
189        bytemuck::cast_slice(bytes)
190    }
191    fn cube_type() -> ElemType {
192        ElemType::Int(IntKind::I8)
193    }
194    fn maximum_value() -> Self {
195        // Seems to cause problem for some GPU
196        i8::MAX - 1
197    }
198    fn minimum_value() -> Self {
199        // Seems to cause problem for some GPU
200        i8::MIN + 1
201    }
202}
203
204impl CubeElement for f64 {
205    fn type_name() -> &'static str {
206        "f64"
207    }
208    fn as_bytes(slice: &[Self]) -> &[u8] {
209        bytemuck::cast_slice(slice)
210    }
211    fn from_bytes(bytes: &[u8]) -> &[Self] {
212        bytemuck::cast_slice(bytes)
213    }
214    fn cube_type() -> ElemType {
215        ElemType::Float(FloatKind::F64)
216    }
217    fn maximum_value() -> Self {
218        f64::MAX
219    }
220    fn minimum_value() -> Self {
221        f64::MIN
222    }
223}
224
225impl CubeElement for f32 {
226    fn type_name() -> &'static str {
227        "f32"
228    }
229    fn as_bytes(slice: &[Self]) -> &[u8] {
230        bytemuck::cast_slice(slice)
231    }
232    fn from_bytes(bytes: &[u8]) -> &[Self] {
233        bytemuck::cast_slice(bytes)
234    }
235    fn cube_type() -> ElemType {
236        ElemType::Float(FloatKind::F32)
237    }
238    fn maximum_value() -> Self {
239        f32::MAX
240    }
241    fn minimum_value() -> Self {
242        f32::MIN
243    }
244}
245
246impl CubeElement for half::f16 {
247    fn type_name() -> &'static str {
248        "f16"
249    }
250    fn as_bytes(slice: &[Self]) -> &[u8] {
251        bytemuck::cast_slice(slice)
252    }
253    fn from_bytes(bytes: &[u8]) -> &[Self] {
254        bytemuck::cast_slice(bytes)
255    }
256    fn cube_type() -> ElemType {
257        ElemType::Float(FloatKind::F16)
258    }
259    fn maximum_value() -> Self {
260        half::f16::MAX
261    }
262    fn minimum_value() -> Self {
263        half::f16::MIN
264    }
265}
266
267impl CubeElement for half::bf16 {
268    fn type_name() -> &'static str {
269        "bf16"
270    }
271    fn as_bytes(slice: &[Self]) -> &[u8] {
272        bytemuck::cast_slice(slice)
273    }
274    fn from_bytes(bytes: &[u8]) -> &[Self] {
275        bytemuck::cast_slice(bytes)
276    }
277    fn cube_type() -> ElemType {
278        ElemType::Float(FloatKind::BF16)
279    }
280    fn maximum_value() -> Self {
281        half::bf16::MAX
282    }
283    fn minimum_value() -> Self {
284        half::bf16::MIN
285    }
286}
287
288impl CubeElement for flex32 {
289    fn type_name() -> &'static str {
290        "flex32"
291    }
292    fn as_bytes(slice: &[Self]) -> &[u8] {
293        bytemuck::cast_slice(slice)
294    }
295    fn from_bytes(bytes: &[u8]) -> &[Self] {
296        bytemuck::cast_slice(bytes)
297    }
298    fn cube_type() -> ElemType {
299        ElemType::Float(FloatKind::Flex32)
300    }
301    fn maximum_value() -> Self {
302        <flex32 as num_traits::Float>::max_value()
303    }
304    fn minimum_value() -> Self {
305        <flex32 as num_traits::Float>::min_value()
306    }
307}
308
309impl CubeElement for tf32 {
310    fn type_name() -> &'static str {
311        "tf32"
312    }
313
314    fn as_bytes(slice: &[Self]) -> &[u8] {
315        bytemuck::cast_slice(slice)
316    }
317
318    fn from_bytes(bytes: &[u8]) -> &[Self] {
319        bytemuck::cast_slice(bytes)
320    }
321
322    fn cube_type() -> ElemType {
323        ElemType::Float(FloatKind::TF32)
324    }
325
326    fn maximum_value() -> Self {
327        tf32::max_value()
328    }
329
330    fn minimum_value() -> Self {
331        tf32::min_value()
332    }
333}
334
335impl CubeElement for e4m3 {
336    fn type_name() -> &'static str {
337        "e4m3"
338    }
339
340    fn as_bytes(slice: &[Self]) -> &[u8] {
341        bytemuck::cast_slice(slice)
342    }
343
344    fn from_bytes(bytes: &[u8]) -> &[Self] {
345        bytemuck::cast_slice(bytes)
346    }
347
348    fn cube_type() -> ElemType {
349        ElemType::Float(FloatKind::E4M3)
350    }
351
352    fn maximum_value() -> Self {
353        e4m3::from_f64(e4m3::MAX.to_f64())
354    }
355
356    fn minimum_value() -> Self {
357        e4m3::from_f64(e4m3::MIN.to_f64())
358    }
359}
360
361impl CubeElement for e5m2 {
362    fn type_name() -> &'static str {
363        "e5m2"
364    }
365
366    fn as_bytes(slice: &[Self]) -> &[u8] {
367        bytemuck::cast_slice(slice)
368    }
369
370    fn from_bytes(bytes: &[u8]) -> &[Self] {
371        bytemuck::cast_slice(bytes)
372    }
373
374    fn cube_type() -> ElemType {
375        ElemType::Float(FloatKind::E5M2)
376    }
377
378    fn maximum_value() -> Self {
379        e5m2::from_f64(e5m2::MAX.to_f64())
380    }
381
382    fn minimum_value() -> Self {
383        e5m2::from_f64(e5m2::MIN.to_f64())
384    }
385}
386
387impl CubeElement for ue8m0 {
388    fn type_name() -> &'static str {
389        "ue8m0"
390    }
391
392    fn as_bytes(slice: &[Self]) -> &[u8] {
393        bytemuck::cast_slice(slice)
394    }
395
396    fn from_bytes(bytes: &[u8]) -> &[Self] {
397        bytemuck::cast_slice(bytes)
398    }
399
400    fn cube_type() -> ElemType {
401        ElemType::Float(FloatKind::UE8M0)
402    }
403
404    fn maximum_value() -> Self {
405        ue8m0::MAX
406    }
407
408    fn minimum_value() -> Self {
409        ue8m0::MIN
410    }
411}
412
413impl CubeElement for e2m1x2 {
414    fn type_name() -> &'static str {
415        "e2m1x2"
416    }
417
418    fn as_bytes(slice: &[Self]) -> &[u8] {
419        bytemuck::cast_slice(slice)
420    }
421
422    fn from_bytes(bytes: &[u8]) -> &[Self] {
423        bytemuck::cast_slice(bytes)
424    }
425
426    fn cube_type() -> ElemType {
427        ElemType::Float(FloatKind::E2M1x2)
428    }
429
430    fn maximum_value() -> Self {
431        let max = e2m1::MAX.to_bits();
432        e2m1x2::from_bits(max << 4 | max)
433    }
434
435    fn minimum_value() -> Self {
436        let min = e2m1::MIN.to_bits();
437        e2m1x2::from_bits(min << 4 | min)
438    }
439}
cubecl_core/pod.rs

cubecl_core/
pod.rs
