rudnn/convolution/kernels/backward_weight/
launch.rs

1use ruda_kernel::dsl as kernel_dsl;
2use crate::convolution::components::{ConvolutionProblem, Dimensionality};
3use crate::convolution::kernels::backward_weight::selector::launch_kernel_concrete;
4use crate::convolution::launch::ConvolutionArgs;
5use crate::convolution::{backward_weight::args::ConcreteArgs, components::ConvSetupError};
6use crate::convolution::{
7    components::{ConvolutionOperation, global::args::RuntimeArgs},
8    routines::Routine,
9};
10use ruda_kernel::dsl::Runtime;
11use ruda_kernel::dsl::client::ComputeClient;
12use ruda_kernel::dsl::prelude::*;
13use rublas::kernel_ir::{
14    definition::{AvailableVectorSizes, MatmulElems},
15    routines::BlueprintStrategy,
16};
17use ruda_kernel::tiling::{InputBinding, MatrixLayout};
18
19/// Backward-weight dispatch helper.
20///
21/// Called by `rudnn::convolution::launch_ref` after the routine and
22/// blueprint-strategy have been resolved.
23#[allow(clippy::result_large_err, clippy::too_many_arguments)]
24pub(crate) fn launch_internal<R: Runtime, const N_SPATIAL: usize, Rt: Routine>(
25    client: &ComputeClient<R>,
26    input: InputBinding<R>,
27    out_grad: InputBinding<R>,
28    weight_grad: TensorBinding<R>,
29    args: ConvolutionArgs<N_SPATIAL>,
30    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Rt::MatmulRoutine>,
31    dtypes: MatmulElems,
32) -> Result<(), ConvSetupError>
33where
34    Rt::Args: ConcreteArgs<Rt::MatmulRoutine>,
35{
36    let ConvolutionArgs {
37        stride,
38        padding,
39        dilation,
40    } = args;
41
42    let dimensionality = match N_SPATIAL {
43        1 => Dimensionality::Dim1,
44        2 => Dimensionality::Dim2,
45        3 => Dimensionality::Dim3,
46        other => unimplemented!("Unsupported dimensionality {other}"),
47    };
48
49    launch_with_routine::<R, Rt>(
50        client,
51        input,
52        out_grad,
53        weight_grad,
54        (&stride, &padding, &dilation),
55        dimensionality,
56        blueprint_strategy,
57        dtypes,
58    )
59}
60
61#[allow(clippy::too_many_arguments)]
62fn launch_with_routine<R: Runtime, Rt: Routine>(
63    client: &ComputeClient<R>,
64    input: InputBinding<R>,
65    out_grad: InputBinding<R>,
66    weight_grad: TensorBinding<R>,
67    (stride, padding, dilation): (&[usize], &[usize], &[usize]),
68    dimensionality: Dimensionality,
69    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Rt::MatmulRoutine>,
70    dtypes: MatmulElems,
71) -> Result<(), ConvSetupError>
72where
73    Rt::Args: ConcreteArgs<Rt::MatmulRoutine>,
74{
75    let rank = input.data().shape.len();
76    let dim_c = rank - 1;
77
78    let n = input.shape()[0];
79    let c = input.shape()[dim_c];
80
81    let out_c = out_grad.shape()[dim_c];
82
83    let in_shape = &input.shape()[1..dim_c];
84    let kernel_shape = &weight_grad.shape[1..dim_c];
85    let out_shape = &out_grad.shape()[1..dim_c];
86
87    let op = ConvolutionOperation::BackwardWeight;
88
89    let input_data = Rt::correct_layout(client, input.clone().into_data(), dtypes.lhs_global, op)?;
90    let out_grad_data =
91        Rt::correct_layout(client, out_grad.clone().into_data(), dtypes.rhs_global, op)?;
92
93    let mut input = input.clone();
94    let mut out_grad = out_grad.clone();
95
96    *input.data_mut() = input_data;
97    *out_grad.data_mut() = out_grad_data;
98
99    let address_type = input
100        .required_address_type()
101        .max(out_grad.required_address_type())
102        .max(weight_grad.required_address_type(dtypes.acc_global.size()));
103
104    let problem = ConvolutionProblem {
105        m: out_c,
106        n: c * kernel_shape.iter().product::<usize>(),
107        k: n * out_shape.iter().product::<usize>(),
108        lhs_strides: input.data().strides.clone(),
109        rhs_strides: out_grad.data().strides.clone(),
110        lhs_layout: MatrixLayout::ColMajor,
111        rhs_layout: MatrixLayout::RowMajor,
112        kernel_size: kernel_shape.iter().map(|it| *it as u32).collect(),
113        stride: stride.iter().map(|it| *it as u32).collect(),
114        padding: padding.iter().map(|it| *it as i32).collect(),
115        dilation: dilation.iter().map(|it| *it as u32).collect(),
116
117        batches: n,
118        in_shape: in_shape.into(),
119        out_shape: out_shape.into(),
120        channels: c,
121        out_channels: out_c,
122
123        padded_channels: c,
124        operation: op,
125
126        dimensionality,
127        global_dtypes: dtypes.as_global_elems(),
128        address_type,
129    };
130
131    launch_kernel::<R, Rt>(
132        client,
133        input,
134        out_grad,
135        weight_grad,
136        problem,
137        blueprint_strategy,
138        dtypes,
139    )
140}
141
142#[allow(clippy::result_large_err, clippy::too_many_arguments)]
143pub fn launch_kernel<R: Runtime, Rt: Routine>(
144    client: &ComputeClient<R>,
145    input: InputBinding<R>,
146    out_grad: InputBinding<R>,
147    weight_grad: TensorBinding<R>,
148    problem: ConvolutionProblem,
149    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Rt::MatmulRoutine>,
150    dtypes: MatmulElems,
151) -> Result<(), ConvSetupError>
152where
153    Rt::Args: ConcreteArgs<Rt::MatmulRoutine>,
154{
155    // Shape/strides are treated as k-major, with the last dim always being the contiguous one.
156    // So for the sake of selecting a vector size, the shape/strides are always row-major.
157    let vector_sizes = AvailableVectorSizes::from_type_sizes(
158        client,
159        input.data_elem_size(),
160        out_grad.data_elem_size(),
161        dtypes.acc_global.size(),
162    )
163    .filter_lhs_with_tensor(
164        &out_grad.data().strides,
165        &out_grad.data().shape,
166        MatrixLayout::RowMajor,
167    )
168    .filter_rhs_with_tensor(
169        &input.data().strides,
170        &input.data().shape,
171        MatrixLayout::RowMajor,
172    )
173    .filter_out_with_tensor(&weight_grad.strides, &weight_grad.shape);
174
175    let vector_sizes = Rt::filter_vector_sizes(vector_sizes).pick_max()?;
176
177    launch_kernel_concrete::<R, Rt::Args, Rt::MatmulRoutine>(
178        client,
179        input,
180        out_grad,
181        weight_grad,
182        problem,
183        vector_sizes,
184        blueprint_strategy,
185        &dtypes,
186    )
187}
rudnn/convolution/kernels/backward_weight/launch.rs

rudnn/convolution/kernels/backward_weight/
launch.rs
