rudnn/convolution/launch/
inputs.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::Runtime;
3use ruda_kernel::dsl::prelude::TensorBinding;
4use ruda_kernel::tiling::InputBinding;
5
6use crate::convolution::components::ConvolutionOperation;
7
8/// Spatial convolution arguments (stride / padding / dilation per spatial dim).
9#[derive(Clone, Debug)]
10pub struct ConvolutionArgs<const N_SPATIAL: usize> {
11    pub stride: [usize; N_SPATIAL],
12    pub padding: [usize; N_SPATIAL],
13    pub dilation: [usize; N_SPATIAL],
14}
15
16#[allow(clippy::large_enum_variant)]
17/// Per-operation tensor bindings supplied to `launch_ref`.
18///
19/// Each variant carries exactly the bindings the corresponding operation needs.
20/// The discriminant maps 1:1 to `ConvolutionOperation`.
21pub enum ConvolutionInputs<R: Runtime> {
22    Forward {
23        input: InputBinding<R>,
24        weight: InputBinding<R>,
25        bias: Option<InputBinding<R>>,
26        out: TensorBinding<R>,
27    },
28    BackwardData {
29        out_grad: InputBinding<R>,
30        weights: InputBinding<R>,
31        in_grad: TensorBinding<R>,
32    },
33    BackwardWeight {
34        input: InputBinding<R>,
35        out_grad: InputBinding<R>,
36        weight_grad: TensorBinding<R>,
37    },
38}
39
40impl<R: Runtime> ConvolutionInputs<R> {
41    pub fn operation(&self) -> ConvolutionOperation {
42        match self {
43            ConvolutionInputs::Forward { .. } => ConvolutionOperation::Forward,
44            ConvolutionInputs::BackwardData { .. } => ConvolutionOperation::BackwardData,
45            ConvolutionInputs::BackwardWeight { .. } => ConvolutionOperation::BackwardWeight,
46        }
47    }
48}
rudnn/convolution/launch/inputs.rs

rudnn/convolution/launch/
inputs.rs
