ruda_kernel/dsl/frontend/
unaligned_vector.rs

1use crate::dsl::intrinsic;
2use crate::dsl::ir::{IndexAssignOperator, IndexOperator, Instruction, Operator};
3use crate::dsl::{prelude::*};
4
5/// An extension trait for expanding the ruda frontend with the ability to
6/// request unaligned vector reads and writes
7///
8/// Typically in ruda, a buffer is declared as having a certain vector size
9/// at kernel compilation time. The buffer can then be indexed to produce
10/// vectors that are aligned to the `vector_size`.
11///
12/// This trait allows the user to request a `vector_read` from a buffer where the
13/// start of the read is not aligned to the `vector_read` requested.
14///
15/// As an example, imagine a buffer of scalar length 4. With `vector_size` = 1,
16/// this could be illustrated like so
17/// [1, 2, 3, 4]
18///
19/// Imagine the same buffer, now with `vector_size` = 2.
20/// [[1, 2], [3, 4]]
21///
22/// Vectors can now be accessed from this buffer, but only those that that are aligned
23/// with the `vector_size`. I.e. we can get the vectors [1, 2] or [3, 4], but not [2, 3]
24///
25/// This trait allows you to treat the buffer as having no `vector_size` = 1, yet asking
26/// for a vector of some kernel-compile-time known length at some offset in the buffer.
27/// I.e. if for the buffer `buf = [1, 2, 3, 4]`, `buf.unaligned_vector_read(1, 2)`
28/// will produce the vector `[2, 3]`.
29#[ruda]
30pub trait UnalignedVector<E: Scalar, N: Size>: RudaType + Sized {
31    /// Perform an unchecked read of a vector of the given length at the given index
32    ///
33    /// # Safety
34    /// Out of bounds indexing causes undefined behaviour and may segfault. Ensure `index..index+vector_size` is
35    /// always in bounds
36    fn unaligned_vector_read(&self, index: usize) -> Vector<E, N>;
37
38    /// Perform an unchecked write of a vector of the given length at the given index
39    ///
40    /// # Safety
41    /// Out of bounds indexing causes undefined behaviour and may segfault. Ensure `index..index+vector_size` is
42    /// always in bounds
43    fn unaligned_vector_write(&mut self, index: usize, value: Vector<E, N>);
44}
45
46macro_rules! impl_unaligned_vector {
47    ($type:ident) => {
48        paste::paste! {
49            type [<$type Expand>]<E> = NativeExpand<$type<E>>;
50        }
51        #[ruda]
52        impl<E: Scalar, N: Size> UnalignedVector<E, N> for $type<E> {
53            fn unaligned_vector_read(&self, index: usize) -> Vector<E, N> {
54                unaligned_vector_read::<$type<E>, E, N>(self, index)
55            }
56
57            fn unaligned_vector_write(&mut self, index: usize, value: Vector<E, N>) {
58                unaligned_vector_write::<$type<E>, E, N>(self, index, value)
59            }
60        }
61    };
62}
63
64impl_unaligned_vector!(Array);
65impl_unaligned_vector!(Tensor);
66impl_unaligned_vector!(SharedMemory);
67
68// TODO: Maybe impl unaligned IO on slices?
69// The last dimension will have to be contiguous for this to make sense,
70// as the unaligned IO isn't gather / scatter from arbitrary memory locations
71// and still needs the loaded elements to be contiguous
72
73#[ruda]
74#[allow(unused_variables)]
75fn unaligned_vector_read<T: RudaType<ExpandType = NativeExpand<T>>, E: Scalar, N: Size>(
76    this: &T,
77    index: usize,
78) -> Vector<E, N> {
79    intrinsic!(|scope| {
80        if !matches!(this.expand.ty, crate::dsl::ir::Type::Scalar(_)) {
81            todo!("Unaligned reads are only allowed on scalar arrays for now");
82        }
83        let vector_size = N::__expand_value(scope);
84        let out = scope.create_local(this.expand.ty.with_vector_size(vector_size));
85        scope.register(Instruction::new(
86            Operator::UncheckedIndex(IndexOperator {
87                list: *this.expand,
88                index: index.expand.consume(),
89                vector_size: 0,
90                unroll_factor: 1,
91            }),
92            *out,
93        ));
94        out.into()
95    })
96}
97
98#[ruda]
99#[allow(unused_variables)]
100fn unaligned_vector_write<T: RudaType<ExpandType = NativeExpand<T>>, E: Scalar, N: Size>(
101    this: &mut T,
102    index: usize,
103    value: Vector<E, N>,
104) {
105    intrinsic!(|scope| {
106        if !matches!(this.expand.ty, crate::dsl::ir::Type::Scalar(_)) {
107            todo!("Unaligned reads are only allowed on scalar arrays for now");
108        }
109        let vector_size = N::__expand_value(scope);
110        scope.register(Instruction::new(
111            Operator::UncheckedIndexAssign(IndexAssignOperator {
112                index: index.expand.consume(),
113                value: value.expand.consume(),
114                vector_size: 0,
115                unroll_factor: 1,
116            }),
117            *this.expand,
118        ));
119    })
120}
ruda_kernel/dsl/frontend/unaligned_vector.rs

ruda_kernel/dsl/frontend/
unaligned_vector.rs
