Skip to main content

rudnn/attention/kernel_ir/components/tile/
operand.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3
4use ruda_kernel::tiling::tile::{Plane, Tile};
5
6#[derive(RudaType)]
7/// Query input to the Tile Attention
8pub struct Query<L: Numeric> {
9    pub tile: Tile<L, Plane, ReadWrite>,
10}
11
12#[ruda]
13impl<L: Numeric> Query<L> {
14    pub fn new(tile: Tile<L, Plane, ReadWrite>) -> Query<L> {
15        Query::<L> { tile }
16    }
17}
18
19#[derive(RudaType)]
20pub struct Key<R: Numeric> {
21    pub tile: Tile<R, Plane, ReadWrite>,
22}
23
24#[ruda]
25impl<R: Numeric> Key<R> {
26    pub fn new(tile: Tile<R, Plane, ReadWrite>) -> Key<R> {
27        Key::<R> { tile }
28    }
29}
30
31#[derive(RudaType)]
32pub struct Value<R: Numeric> {
33    pub tile: Tile<R, Plane, ReadWrite>,
34}
35
36#[ruda]
37impl<R: Numeric> Value<R> {
38    pub fn new(tile: Tile<R, Plane, ReadWrite>) -> Value<R> {
39        Value::<R> { tile }
40    }
41}