rudnn/convolution/routines/
simple.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::server::LaunchError;
3use ruda_kernel::dsl::Runtime;
4use ruda_kernel::dsl::client::ComputeClient;
5use ruda_kernel::dsl::ir::StorageType;
6use ruda_kernel::dsl::prelude::TensorBinding;
7use rublas::kernel_ir::components::global::read::FullLoadingStrategy;
8use rublas::kernel_ir::components::{
9    global::read::sync_full_cyclic::SyncFullCyclicLoading,
10    stage::{ColMajorTilingOrder, RowMajorTilingOrder},
11};
12use rublas::kernel_ir::{
13    components::global::read::{
14        async_full_tma::AsyncFullTmaLoading, sync_full_strided::SyncFullStridedLoading,
15        sync_full_tilewise::SyncFullTilewiseLoading,
16    },
17    routines::simple::{SimpleAlgorithm, SimpleArgs},
18};
19use rublas::kernel_ir::{
20    definition::{AvailableVectorSizes, TilingBlueprint},
21    launch::{TensorArgs, TensorMapArgs},
22};
23use ruda_kernel::tiling::tile::Strided;
24use std::marker::PhantomData;
25
26use crate::convolution::{
27    components::{
28        ConvolutionOperation,
29        global::{
30            args::RuntimeArgs,
31            read::strategy::{
32                async_full_cyclic::AsyncFullCyclicLoading,
33                async_full_strided::AsyncFullStridedLoading, sync_bias::SyncBiasLoading,
34            },
35        },
36    },
37    routines::{Routine, contiguous_pitched_layout, into_tensor_handle_tma},
38};
39
40/// Cmma convolution
41pub struct SimpleConv<LL: FullLoadingStrategy<RuntimeArgs>, LR: FullLoadingStrategy<RuntimeArgs>> {
42    _loader: PhantomData<(LL, LR)>,
43}
44
45pub type SimpleSyncCyclicConv = SimpleConv<
46    SyncFullCyclicLoading<RowMajorTilingOrder>,
47    SyncFullCyclicLoading<ColMajorTilingOrder>,
48>;
49pub type SimpleSyncStridedConv = SimpleConv<SyncFullStridedLoading, SyncFullStridedLoading>;
50pub type SimpleSyncTilewiseConv = SimpleConv<
51    SyncFullTilewiseLoading<RowMajorTilingOrder>,
52    SyncFullTilewiseLoading<ColMajorTilingOrder>,
53>;
54pub type SimpleAsyncCyclicConv = SimpleConv<
55    AsyncFullCyclicLoading<RowMajorTilingOrder>,
56    AsyncFullCyclicLoading<ColMajorTilingOrder>,
57>;
58pub type SimpleAsyncStridedConv = SimpleConv<AsyncFullStridedLoading, AsyncFullStridedLoading>;
59
60pub struct SimpleAsyncTmaConv;
61
62impl<
63    LL: FullLoadingStrategy<RuntimeArgs, TileKind = Strided>,
64    LR: FullLoadingStrategy<RuntimeArgs, TileKind = Strided, SyncStrategy = LL::SyncStrategy>,
65> Routine for SimpleConv<LL, LR>
66{
67    type Blueprint = TilingBlueprint;
68    type Strategy = SimpleArgs;
69    type MatmulRoutine = SimpleAlgorithm<LL, LR, SyncBiasLoading>;
70    type Args = TensorArgs<RuntimeArgs>;
71
72    fn correct_layout<R: Runtime>(
73        client: &ComputeClient<R>,
74        handle: TensorBinding<R>,
75        dtype: StorageType,
76        _operation: ConvolutionOperation,
77    ) -> Result<TensorBinding<R>, LaunchError> {
78        contiguous_pitched_layout(client, handle, dtype)
79    }
80}
81
82impl Routine for SimpleAsyncTmaConv {
83    type Blueprint = TilingBlueprint;
84    type Strategy = SimpleArgs;
85    type MatmulRoutine = SimpleAlgorithm<AsyncFullTmaLoading, AsyncFullTmaLoading, SyncBiasLoading>;
86    type Args = TensorMapArgs<RuntimeArgs>;
87
88    fn correct_layout<R: Runtime>(
89        client: &ComputeClient<R>,
90        handle: TensorBinding<R>,
91        dtype: StorageType,
92        operation: ConvolutionOperation,
93    ) -> Result<TensorBinding<R>, LaunchError> {
94        into_tensor_handle_tma(client, handle, dtype, operation)
95    }
96
97    fn filter_vector_sizes(vector_sizes: AvailableVectorSizes) -> AvailableVectorSizes {
98        AvailableVectorSizes {
99            lhs: vec![1],
100            rhs: vec![1],
101            out: vector_sizes.out,
102        }
103    }
104}
rudnn/convolution/routines/simple.rs

rudnn/convolution/routines/
simple.rs
