rudnn/attention/kernel_ir/components/global/simple/reader/
query.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use ruda_kernel::library::Swizzle;
4use ruda_kernel::library::tensor::View;
5use ruda_kernel::library::tensor::layout::Coords2d;
6use rublas::kernel_ir::components::global::memory::GlobalMemoryConfig;
7use ruda_kernel::tiling::tile::StridedTile;
8
9use crate::attention::kernel_ir::{
10    components::stage::AttentionPartitioner,
11    definition::attention_types::{QG, QGS},
12    definition::{AttentionPrecision, AttentionTileSize},
13};
14
15#[derive(RudaType)]
16pub struct QueryReader<AP: AttentionPrecision> {
17    query: View<Vector<QG<AP>, QGS<AP>>, Coords2d>,
18    #[ruda(comptime)]
19    gmem_config: GlobalMemoryConfig,
20}
21
22#[ruda]
23impl<AP: AttentionPrecision> QueryReader<AP> {
24    pub fn new(
25        stage_q_offset: u32,
26        query: View<Vector<QG<AP>, QGS<AP>>, Coords2d>,
27        #[comptime] gmem_config: GlobalMemoryConfig,
28    ) -> Self {
29        let query = query.slice((stage_q_offset, 0), query.shape());
30
31        QueryReader::<AP> { query, gmem_config }
32    }
33
34    pub fn get_tile<P: AttentionPartitioner>(
35        &self,
36        tile: Coords2d,
37        #[comptime] tile_size: AttentionTileSize,
38        #[comptime] partition_seq_q: u32,
39        #[comptime] partition_head_dim: u32,
40    ) -> StridedTile<QG<AP>, QGS<AP>> {
41        let (row_in_partition, col) = tile;
42
43        let row = row_in_partition + P::seq_q_index() * partition_seq_q;
44
45        let vector_size = self.gmem_config.vector_size.comptime() as u32;
46
47        let slice = self
48            .query
49            .slice(
50                (row * tile_size.seq_q, col * tile_size.head_dim),
51                (tile_size.seq_q, tile_size.head_dim).runtime(),
52            )
53            .to_linear_slice();
54
55        let start = 0;
56        let vectors_per_tile = tile_size.seq_q * tile_size.head_dim / vector_size;
57        let end = start + vectors_per_tile;
58        let vectors_per_partition_row = partition_head_dim * tile_size.head_dim / vector_size;
59
60        StridedTile::<QG<AP>, QGS<AP>>::new_strided(
61            slice,
62            start,
63            end,
64            vectors_per_partition_row,
65            Swizzle::none(),
66            self.gmem_config.matrix_layout,
67        )
68    }
69
70    pub fn head_dim(&self) -> u32 {
71        self.query.shape().1
72    }
73
74    pub fn get_staged_tile<P: AttentionPartitioner>(
75        &self,
76        tile: Coords2d,
77        #[comptime] tile_size: AttentionTileSize,
78        #[comptime] partition_seq_q: u32,
79        #[comptime] plane_dim: u32,
80        #[comptime] num_planes: u32,
81    ) -> StridedTile<QG<AP>, QGS<AP>> {
82        let vector_size = comptime!(self.gmem_config.vector_size as u32);
83        let vectors_per_row = comptime!(tile_size.head_dim / vector_size);
84        let vectors_per_tile = comptime!(tile_size.seq_q * vectors_per_row);
85        let start = UNIT_POS_Y * vectors_per_tile;
86        let smem_size = comptime!((vectors_per_tile * num_planes) as usize);
87        let mut storage = SharedMemory::<Vector<QG<AP>, QGS<AP>>>::new(smem_size);
88        let mut data = storage.slice_mut(start as usize, (start + vectors_per_tile) as usize);
89        let row = (tile.0 + P::seq_q_index() * partition_seq_q) * tile_size.seq_q;
90        let col = tile.1 * tile_size.head_dim;
91        let mut index = UNIT_POS_X;
92        while index < vectors_per_tile {
93            data[index as usize] = self.query.read_checked((
94                row + index / vectors_per_row,
95                col + (index % vectors_per_row) * vector_size,
96            ));
97            index += plane_dim;
98        }
99        sync_ruda();
100        StridedTile::<QG<AP>, QGS<AP>>::new_strided(
101            data.to_slice(),
102            0,
103            vectors_per_tile,
104            vectors_per_row,
105            Swizzle::none(),
106            self.gmem_config.matrix_layout,
107        )
108    }
109}
rudnn/attention/kernel_ir/components/global/simple/reader/query.rs

rudnn/attention/kernel_ir/components/global/simple/reader/
query.rs
