ruPRIM 0.1.8

Parallel primitives, reductions, scans, and indexing for Ruda.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use ruda_kernel::dsl as kernel_dsl;
use crate::reduce::{
    BoundChecks, ReduceInstruction, ReducePrecision, VectorizationMode,
    components::{
        args::NumericVector,
        instructions::{Item, ReduceRequirements},
        readers::{bound_checks::ReaderBoundChecks, new_coordinates},
    },
};
use ruda_kernel::dsl::prelude::*;
use ruda_kernel::library::tensor::View;
use ruda_kernel::library::tensor::layout::Coords1d;
use ruda_kernel::library::tensor::layout::plain::PlainLayout;
use ruda_kernel::library::tensor::r#virtual::VirtualTensor;
use crate::reduce::components::layout::reduction_input_offset;

#[derive(RudaType)]
pub struct PerpendicularReader<P: ReducePrecision> {
    view: View<Vector<P::EI, P::SI>, Coords1d>,
    /// The global offset that points where the vector to reduce is located in global memory.
    batch_offset: usize,
    vector_offset_stride: usize,
    requirements: ReduceRequirements,
    bound_checks: ReaderBoundChecks<P>,
    shape: usize,
    effective_plane_dim: u32,
}

#[ruda]
impl<P: ReducePrecision> PerpendicularReader<P> {
    #[allow(clippy::too_many_arguments)]
    pub fn new<I: ReduceInstruction<P>, Out: NumericVector>(
        input: &VirtualTensor<P::EI, P::SI>,
        output: &mut VirtualTensor<Out::T, Out::N, ReadWrite>,
        inst: &I,
        reduce_axis: usize,
        reduce_index: usize,
        idle: ComptimeOption<bool>,
        effective_plane_dim: u32,
        #[comptime] bound_checks: BoundChecks,
    ) -> PerpendicularReader<P> {
        let vector_size = input.vector_size();
        let output_index = reduce_index * vector_size;

        let batch_offset = reduction_input_offset(input, output, reduce_axis, output_index);

        let requirements = I::requirements(inst);
        let vector_offset_stride = input.stride(reduce_axis) / vector_size;
        let shape = input.shape(reduce_axis);

        let bound_checks = ReaderBoundChecks::new::<I>(inst, shape, idle, bound_checks);

        PerpendicularReader::<P> {
            view: input.view(PlainLayout::new(input.len())),
            batch_offset,
            vector_offset_stride,
            requirements,
            bound_checks,
            shape,
            effective_plane_dim,
        }
    }

    pub fn length_unit(&self) -> usize {
        self.shape
    }

    pub fn length_plane(&self) -> usize {
        self.shape.div_ceil(self.effective_plane_dim as usize)
    }

    pub fn length_ruda(&self) -> usize {
        self.shape.div_ceil(RUDA_DIM as usize)
    }

    pub fn read_ruda(&self, vector_index: usize) -> Item<P> {
        let plane_pos = vector_index * RUDA_DIM as usize;
        let unit_pos = UNIT_POS as usize;
        let pos = plane_pos + unit_pos;
        let offset = plane_pos * self.vector_offset_stride
            + unit_pos * self.vector_offset_stride
            + self.batch_offset;

        let elements = self.bound_checks.read(pos, offset, &self.view);

        let args = new_coordinates(
            plane_pos + unit_pos,
            self.requirements,
            VectorizationMode::Perpendicular,
        );

        Item::<P> { elements, args }
    }

    pub fn read_plane(&self, vector_index: usize) -> Item<P> {
        self.read_plane_at(vector_index, self.effective_plane_dim, UNIT_POS_PLANE)
    }

    pub(crate) fn read_plane_at(&self, vector_index: usize, plane_dim: u32, unit_pos: u32) -> Item<P> {
        let plane_pos = vector_index * plane_dim as usize;
        let unit_pos = unit_pos as usize;
        let pos = plane_pos + unit_pos;
        let offset = plane_pos * self.vector_offset_stride
            + unit_pos * self.vector_offset_stride
            + self.batch_offset;

        let elements = self.bound_checks.read(pos, offset, &self.view);

        let args = new_coordinates(
            plane_pos + unit_pos,
            self.requirements,
            VectorizationMode::Perpendicular,
        );

        Item::<P> { elements, args }
    }

    pub fn read_unit(&self, vector_index: usize) -> Item<P> {
        let offset = self.batch_offset + vector_index * self.vector_offset_stride;
        let elements = self.bound_checks.read(vector_index, offset, &self.view);

        let args = new_coordinates(
            vector_index,
            self.requirements,
            VectorizationMode::Perpendicular,
        );

        Item::<P> { elements, args }
    }
}