ruprim/reduce/components/
config.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::zspace::Strides;
3
4#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
5pub enum VectorizationMode {
6    Parallel,
7    Perpendicular,
8}
9
10pub fn output_vectorization_axis(
11    input_strides: &Strides,
12    reduce_axis: usize,
13    _vectorization_mode: VectorizationMode,
14) -> usize {
15    if input_strides.len() < 2 {
16        // The axis of vectorization for input and output are both 0
17        return 0;
18    }
19
20    // Find the two smallest strides overall (tracking axis indices).
21    let mut min1 = (usize::MAX, 0); // (stride, axis)
22    let mut min2 = (usize::MAX, 0);
23
24    for (i, &s) in input_strides.iter().enumerate() {
25        if s < min1.0 {
26            min2 = min1;
27            min1 = (s, i);
28        } else if s < min2.0 {
29            min2 = (s, i);
30        }
31    }
32
33    // The vectorization axis is the smallest-stride *non-reduce* axis. For
34    // parallel reductions the reduce axis is itself the contiguous (stride 1)
35    // axis, so this falls through to the next-smallest; for perpendicular it's
36    // usually the smallest, except when the reduce axis happens to share the
37    // overall minimum (e.g. a broadcast stride of 0), which forces the fallback.
38    if min1.1 == reduce_axis {
39        min2.1
40    } else {
41        min1.1
42    }
43}
44
45#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
46/// How bound checks is handled for inner reductions.
47pub enum BoundChecks {
48    /// No bound check is necessary.
49    None,
50    /// Using a mask is enough for bound checks.
51    /// This will still read the memory in an out-of-bound location,
52    /// but will replace the value by the null value.
53    Mask,
54    /// Branching is necessary for bound checks.
55    ///
56    /// Probably the right setting when performing fuse on read.
57    Branch,
58}
59
60impl BoundChecks {
61    pub fn idle(self) -> Self {
62        Self::Mask
63    }
64}
65
66#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
67pub enum IdleMode {
68    None,
69    Mask,
70    Terminate,
71}
72
73impl IdleMode {
74    /// Whether idle is activated.
75    pub fn is_enabled(&self) -> bool {
76        !matches!(self, Self::None)
77    }
78}
ruprim/reduce/components/config.rs

ruprim/reduce/components/
config.rs
