cubek_convolution/components/global/memory/
tma.rs

1use cubecl::prelude::*;
2
3#[derive(CubeType)]
4/// A view of a feature map tensor that starts reading data from a specified offset.
5/// Ensures safe access by preventing out-of-bounds errors.
6/// Includes pre-fetched shapes and strides for optimized performance.
7pub struct Im2colTmaReader<E: Numeric> {
8    pub tensor: TensorMap<Line<E>>,
9    pub n_offset: u32,
10    pub spatial_offsets: Sequence<u32>,
11    pub k_offset: u32,
12}
13
14#[cube]
15impl<E: Numeric> Im2colTmaReader<E> {
16    #[allow(clippy::too_many_arguments)]
17    pub fn new(
18        tensor: TensorMap<Line<E>>,
19        n_offset: u32,
20        spatial_offsets: Sequence<u32>,
21        k_offset: u32,
22    ) -> Im2colTmaReader<E> {
23        Im2colTmaReader::<E> {
24            tensor,
25            n_offset,
26            spatial_offsets,
27            k_offset,
28        }
29    }
30
31    /// Advance the view along the k dimension by a specified offset, `k_offset`.
32    pub fn update_view(&mut self, k_offset: u32) {
33        self.k_offset += k_offset;
34    }
35}
36
37unsafe impl<E: Numeric> Sync for Im2colTmaReader<E> {}
38unsafe impl<E: Numeric> Send for Im2colTmaReader<E> {}

cubek_convolution/components/global/memory/tma.rs

cubek_convolution/components/global/memory/
tma.rs