rudnn/convolution/components/
problem.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::ir::AddressType;
3use ruda_kernel::dsl::zspace::Shape;
4use ruda_kernel::dsl::zspace::Strides;
5use ruda_kernel::dsl::zspace::shape;
6use rublas::kernel_ir::definition::{MatmulGlobalElems, MatmulProblem};
7use ruda_kernel::tiling::MatrixLayout;
8
9#[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
10pub enum ConvolutionOperation {
11    Forward,
12    BackwardData,
13    BackwardWeight,
14    ForwardTransposed,
15}
16
17#[derive(Clone, Debug)]
18/// Description of a matmul problem to solve, regardless of actual data
19pub struct ConvolutionProblem {
20    pub m: usize,
21    pub n: usize,
22    pub k: usize,
23
24    pub lhs_strides: Strides,
25    pub rhs_strides: Strides,
26
27    pub lhs_layout: MatrixLayout,
28    pub rhs_layout: MatrixLayout,
29
30    pub kernel_size: Vec<u32>,
31    pub stride: Vec<u32>,
32    pub padding: Vec<i32>,
33    pub dilation: Vec<u32>,
34
35    pub batches: usize,
36    pub channels: usize,
37    pub out_channels: usize,
38    pub in_shape: Shape,
39    pub out_shape: Shape,
40
41    /// Channels after applying loader-specific padding
42    pub padded_channels: usize,
43    pub operation: ConvolutionOperation,
44
45    pub dimensionality: Dimensionality,
46
47    pub global_dtypes: MatmulGlobalElems,
48    /// Address type, defined as the max of each handle's `required_address_type`
49    pub address_type: AddressType,
50}
51
52impl ConvolutionProblem {
53    pub fn as_matmul_problem(&self) -> MatmulProblem {
54        let rank = self.lhs_strides.len();
55
56        // Strides are expected to be in row major (m, n) format so for matmul checks we need to
57        // convert them to that format, with all other dims treated as batch dims so they're still
58        // checked.
59        // lhs already has the right format, but rhs needs special handling.
60        // (h, w, c, n)
61        let lhs_strides = match self.lhs_layout {
62            MatrixLayout::RowMajor => self.lhs_strides.clone(),
63            MatrixLayout::ColMajor => {
64                let mut lhs_strides: Strides = self.lhs_strides[1..rank].into();
65                lhs_strides.push(self.lhs_strides[0]);
66                lhs_strides
67            }
68        };
69        let rhs_strides = match self.rhs_layout {
70            MatrixLayout::RowMajor => self.rhs_strides.clone(),
71            MatrixLayout::ColMajor => {
72                let mut rhs_strides: Strides = self.rhs_strides[1..rank].into();
73                rhs_strides.push(self.rhs_strides[0]);
74                rhs_strides
75            }
76        };
77
78        MatmulProblem {
79            m: self.m,
80            n: self.n,
81            k: self.k,
82            lhs_batches: shape![],
83            rhs_batches: shape![],
84            out_batches: shape![],
85            lhs_strides,
86            rhs_strides,
87            lhs_layout: self.lhs_layout,
88            rhs_layout: self.rhs_layout,
89            lhs_shape: shape![self.m, self.k],
90            rhs_shape: shape![self.k, self.n],
91            out_shape: shape![self.m, self.n],
92            out_strides: MatrixLayout::RowMajor.to_strides(&[self.m, self.n]),
93            out_layout: MatrixLayout::RowMajor,
94            lhs_scheme: None,
95            rhs_scheme: None,
96            global_dtypes: self.global_dtypes.clone(),
97            address_type: self.address_type,
98        }
99    }
100
101    pub fn should_check_channel(&self) -> bool {
102        self.channels != self.padded_channels
103    }
104
105    pub fn should_check_spatial_bounds(&self) -> bool {
106        self.padding.iter().any(|&pad| pad != 0)
107    }
108}
109
110/// Spatial dimensionality of an operation
111#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
112pub enum Dimensionality {
113    Dim1,
114    Dim2,
115    Dim3,
116}
117
118impl Dimensionality {
119    pub fn num_dims(&self) -> usize {
120        match self {
121            Dimensionality::Dim1 => 1,
122            Dimensionality::Dim2 => 2,
123            Dimensionality::Dim3 => 3,
124        }
125    }
126}
rudnn/convolution/components/problem.rs

rudnn/convolution/components/
problem.rs
