rudnn/convolution/routines/
base.rs

1use ruda_kernel::dsl as kernel_dsl;
2use crate::convolution::components::{ConvolutionOperation, global::args::RuntimeArgs};
3use ruda_kernel::dsl::prelude::*;
4use ruda_kernel::dsl::{CompilationError, backtrace::BackTrace, ir::features::Tma};
5use ruda_kernel::library::tensor::into_contiguous_pitched;
6use ruda_kernel::library::tensor::is_contiguous_pitched;
7use rublas::kernel_ir::{
8    definition::{AvailableVectorSizes, Blueprint},
9    launch::MatmulArgs,
10    routines::Routine as MatmulRoutine,
11};
12use std::fmt::Display;
13
14/// Specifications for a convolution routine.
15///
16/// A `Routine` is the convolution-side counterpart of `rublas::kernel_ir::routines::Routine`:
17/// it pairs a per-operation matmul routine with the metadata needed to wire the
18/// kernel up (input args, optional layout fixups, vector-size filtering).
19///
20/// `Blueprint` and `Strategy` are surfaced as direct associated types so callers
21/// don't have to reach through `MatmulRoutine` to bound them.
22pub trait Routine {
23    type Blueprint: Blueprint;
24    type Strategy: Default + Display + Clone;
25
26    type MatmulRoutine: MatmulRoutine<RuntimeArgs, Blueprint = Self::Blueprint, Strategy = Self::Strategy>;
27    type Args: MatmulArgs<Config = RuntimeArgs>;
28
29    /// Whether to select specialized load flow in tests. Should replace with something cleaner
30    /// eventually, but this is nice and simple.
31    const IS_SPECIALIZED: bool = false;
32
33    fn correct_layout<R: Runtime>(
34        client: &ComputeClient<R>,
35        handle: TensorBinding<R>,
36        dtype: StorageType,
37        operation: ConvolutionOperation,
38    ) -> Result<TensorBinding<R>, LaunchError>;
39
40    fn filter_vector_sizes(vector_sizes: AvailableVectorSizes) -> AvailableVectorSizes {
41        vector_sizes
42    }
43}
44
45pub(crate) fn contiguous_pitched_layout<R: Runtime>(
46    client: &ComputeClient<R>,
47    binding: TensorBinding<R>,
48    dtype: StorageType,
49) -> Result<TensorBinding<R>, LaunchError> {
50    let binding = if has_valid_layout(&binding) {
51        binding
52    } else {
53        into_contiguous_pitched(client, binding, dtype).binding()
54    };
55    Ok(binding)
56}
57
58fn has_valid_layout<R: Runtime>(binding: &TensorBinding<R>) -> bool {
59    let rank = binding.shape.len();
60    let dim_c = rank - 1;
61    binding.strides[dim_c] == 1
62}
63
64const TMA_STRIDE_ALIGN: usize = 16;
65
66pub(crate) fn into_tensor_handle_tma<R: Runtime>(
67    client: &ComputeClient<R>,
68    handle: TensorBinding<R>,
69    dtype: StorageType,
70    operation: ConvolutionOperation,
71) -> Result<TensorBinding<R>, LaunchError> {
72    if !client.properties().features.tma.contains(Tma::Base) {
73        return Err(CompilationError::UnsupportedInstruction {
74            reason: "TMA convolution requires tensor-map support on this device".into(),
75            backtrace: BackTrace::capture(),
76        }.into());
77    }
78    let binding = if has_valid_layout_tma(&handle, dtype, operation) {
79        handle
80    } else {
81        into_contiguous_pitched(client, handle, dtype).binding()
82    };
83    Ok(binding)
84}
85
86pub(crate) fn has_valid_layout_tma<R: Runtime>(
87    binding: &TensorBinding<R>,
88    dtype: StorageType,
89    operation: ConvolutionOperation,
90) -> bool {
91    let stride_align = TMA_STRIDE_ALIGN / dtype.size();
92    let rank = binding.shape.len();
93    let dim_c = rank - 1;
94
95    let aligned = binding.strides[..dim_c]
96        .iter()
97        .all(|stride| stride % stride_align == 0);
98
99    let valid_layout = binding.strides[dim_c] == 1;
100
101    let is_valid_wgrad = if operation == ConvolutionOperation::BackwardWeight {
102        is_contiguous_pitched(&binding.shape, &binding.strides)
103    } else {
104        true
105    };
106
107    valid_layout && aligned && is_valid_wgrad
108}
rudnn/convolution/routines/base.rs

rudnn/convolution/routines/
base.rs
