ruprim/reduce/components/writers/
layout.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use ruda_kernel::library::tensor::layout::Coords1d;
4use ruda_kernel::library::tensor::layout::Coords2d;
5use ruda_kernel::library::tensor::layout::Layout;
6use ruda_kernel::library::tensor::layout::LayoutExpand;
7use ruda_kernel::library::tensor::r#virtual::VirtualTensor;
8
9use crate::reduce::components::args::NumericVector;
10use crate::reduce::components::layout::ReductionLayout;
11
12/// Maps a `(write_index, k_iter)` coordinate to a flat vector position in the
13/// output buffer. Strides are expressed in vector units (one step along the
14/// output's SIMD axis = one unit in `write_stride`).
15///
16/// For rank-1 outputs (or any case where `reduce_axis == out_vec_axis`), the
17/// caller should pass `write_stride = 0` and `num_writes = 1`, so the layout
18/// collapses to `position = k_iter * k_stride`.
19#[derive(RudaType, Clone)]
20pub struct ReduceOutputLayout {
21    k_stride: usize,
22    write_stride: usize,
23    num_writes: usize,
24    accumulator_length: usize,
25}
26
27#[ruda]
28impl ReduceOutputLayout {
29    pub fn new(
30        k_stride: usize,
31        write_stride: usize,
32        num_writes: usize,
33        accumulator_length: usize,
34    ) -> ReduceOutputLayout {
35        ReduceOutputLayout {
36            k_stride,
37            write_stride,
38            num_writes,
39            accumulator_length,
40        }
41    }
42}
43
44#[ruda]
45impl Layout for ReduceOutputLayout {
46    type Coordinates = Coords2d;
47    type SourceCoordinates = Coords1d;
48
49    fn to_source_pos(&self, coords: Self::Coordinates) -> Coords1d {
50        let write_index = coords.0 as usize;
51        let k_iter = coords.1 as usize;
52        k_iter * self.k_stride + write_index * self.write_stride
53    }
54
55    fn to_source_pos_checked(&self, coords: Self::Coordinates) -> (Coords1d, bool) {
56        (self.to_source_pos(coords), self.is_in_bounds(coords))
57    }
58
59    fn shape(&self) -> Self::Coordinates {
60        (self.num_writes as u32, self.accumulator_length as u32)
61    }
62
63    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
64        pos.0 < self.num_writes as u32 && pos.1 < self.accumulator_length as u32
65    }
66}
67
68/// Build the output layout from reduction ordinals and the tensor's actual strides.
69#[ruda]
70pub(crate) fn build_reduce_output_layout<Out: NumericVector>(
71    output: &VirtualTensor<Out::T, Out::N, ReadWrite>,
72    reduce_axis: usize,
73    _out_vec_axis: usize,
74    #[comptime] accumulator_length: usize,
75) -> ReductionLayout<Out::T, Out::N> {
76    ReductionLayout::new(output, reduce_axis, accumulator_length)
77}
ruprim/reduce/components/writers/layout.rs

ruprim/reduce/components/writers/
layout.rs
