cubek_interpolate/definition/
precision.rs

1use cubecl::{
2    flex32,
3    ir::{ElemType, FloatKind, StorageType},
4    prelude::Float,
5};
6
7/// Precision used for interpolation.
8pub trait InterpolatePrecision: 'static {
9    /// Precision used for the input and output tensors.
10    type EI: Float;
11    /// Precision used for accumulation and weights.
12    type EA: Float;
13}
14
15impl<EI: Float, EA: Float> InterpolatePrecision for (EI, EA) {
16    type EI = EI;
17    type EA = EA;
18}
19
20impl InterpolatePrecision for f64 {
21    type EI = f64;
22    type EA = f64;
23}
24
25impl InterpolatePrecision for f32 {
26    type EI = f32;
27    type EA = f32;
28}
29
30impl InterpolatePrecision for flex32 {
31    type EI = f32;
32    type EA = f32;
33}
34
35impl InterpolatePrecision for half::f16 {
36    type EI = half::f16;
37    type EA = f32;
38}
39
40impl InterpolatePrecision for half::bf16 {
41    type EI = half::bf16;
42    type EA = f32;
43}
44
45pub fn accumulator_dtype(input: StorageType) -> StorageType {
46    match input {
47        StorageType::Scalar(ElemType::Float(FloatKind::F16))
48        | StorageType::Scalar(ElemType::Float(FloatKind::BF16))
49        | StorageType::Scalar(ElemType::Float(FloatKind::Flex32)) => {
50            StorageType::Scalar(ElemType::Float(FloatKind::F32))
51        }
52        _ => input,
53    }
54}
cubek_interpolate/definition/precision.rs

cubek_interpolate/definition/
precision.rs
