ruprim/reduce/launch/
tune_key.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::AutotuneKey;
3use ruda_kernel::dsl::ir::ElemType;
4use serde::{Deserialize, Serialize};
5
6#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)]
7/// Autotune key representative of reduce versions
8pub struct ReduceAutotuneKey {
9    elem_input: ElemType,
10    elem_output: ElemType,
11    elem_acc: ElemType,
12    /// Whether the axis is contiguous.
13    pub axis_is_contiguous: bool,
14    /// The length of the vector to reduce.
15    ///
16    /// # Notes
17    ///
18    /// Max is 4^4, so 5 values are possible.
19    #[autotune(anchor(exp(min = 16, max = 1024, base = 4)))]
20    pub vector_size: usize,
21    /// The number of vectors to reduce.
22    ///
23    /// # Notes
24    ///
25    /// Max is 8^5, so 5 values are possible.
26    #[autotune(anchor(exp(max = 32768, base = 8)))]
27    pub vector_count: usize,
28}
29
30impl ReduceAutotuneKey {
31    pub fn generate(
32        elem_input: ElemType,
33        elem_output: ElemType,
34        elem_acc: ElemType,
35        input_shape: &[usize],
36        axis_is_contiguous: bool,
37        axis: usize,
38    ) -> Self {
39        let rank = input_shape.len();
40
41        if axis > rank {
42            panic!("axis {axis} is out-of-bound for a rank of {rank}");
43        }
44
45        let reduce_axis_shape = input_shape[axis];
46
47        let reduce_count = input_shape
48            .iter()
49            .enumerate()
50            .filter_map(|(i, shape)| (i != axis).then_some(shape))
51            .product();
52
53        ReduceAutotuneKey::new(
54            elem_input,
55            elem_output,
56            elem_acc,
57            axis_is_contiguous,
58            reduce_axis_shape,
59            reduce_count,
60        )
61    }
62}
ruprim/reduce/launch/tune_key.rs

ruprim/reduce/launch/
tune_key.rs
