Skip to main content

ruprim/reduce/components/
precision.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::define_size;
3use ruda_kernel::dsl::flex32;
4use ruda_kernel::dsl::prelude::Numeric;
5use ruda_kernel::dsl::prelude::Size;
6
7define_size!(pub S);
8
9/// Precision used for the reduction.
10pub trait ReducePrecision: 'static {
11    /// Precision used for the input tensor.
12    type EI: Numeric;
13    /// Size used for the input tensor.
14    type SI: Size;
15    /// Precision used for the accumulation.
16    type EA: Numeric;
17}
18
19impl<EI: Numeric, SI: Size, EA: Numeric> ReducePrecision for (EI, SI, EA) {
20    type EI = EI;
21    type SI = SI;
22    type EA = EA;
23}
24
25// The below implementations are suggestion for reduction that can accumulate precision errors like
26// summations.
27
28impl ReducePrecision for f64 {
29    type EI = f64;
30    type EA = f64;
31    type SI = S;
32}
33
34impl ReducePrecision for f32 {
35    type EI = f32;
36    type EA = f32;
37    type SI = S;
38}
39
40impl ReducePrecision for flex32 {
41    type EI = f32;
42    type EA = f32;
43    type SI = S;
44}
45
46impl ReducePrecision for half::f16 {
47    type EI = half::f16;
48    type EA = f32;
49    type SI = S;
50}
51
52impl ReducePrecision for half::bf16 {
53    type EI = half::bf16;
54    type EA = f32;
55    type SI = S;
56}
57
58impl ReducePrecision for i64 {
59    type EI = i64;
60    type EA = i64;
61    type SI = S;
62}
63
64impl ReducePrecision for i32 {
65    type EI = i32;
66    type EA = i32;
67    type SI = S;
68}
69
70impl ReducePrecision for i16 {
71    type EI = i16;
72    type EA = i32;
73    type SI = S;
74}
75
76impl ReducePrecision for i8 {
77    type EI = i8;
78    type EA = i32;
79    type SI = S;
80}
81
82impl ReducePrecision for u64 {
83    type EI = u64;
84    type EA = u64;
85    type SI = S;
86}
87
88impl ReducePrecision for u32 {
89    type EI = u32;
90    type EA = u32;
91    type SI = S;
92}
93
94impl ReducePrecision for u16 {
95    type EI = u16;
96    type EA = u32;
97    type SI = S;
98}
99
100impl ReducePrecision for u8 {
101    type EI = u8;
102    type EA = u32;
103    type SI = S;
104}