Skip to main content

hermes_simd_core/view/
tile.rs

1//! 2D matrix tile views for high-performance tiled and register-blocked matrix multiply kernels.
2
3use crate::align::{Alignment, Unaligned};
4use crate::scalar::NumericElement;
5use core::marker::PhantomData;
6
7/// A 2D matrix tile view, parameterized by dimensions, alignment, execution mode, and backing reference type.
8///
9/// Under the hood, this represents a 2D tile of shape `ROWS x COLS` with a specified row `stride`.
10/// It is represented as `#[repr(C)]` to guarantee layout compatibility.
11#[repr(C)]
12pub struct TileView<
13    'a,
14    T: NumericElement,
15    Backend,
16    Arch,
17    const ROWS: usize,
18    const COLS: usize,
19    Align: Alignment = Unaligned,
20    Ref: 'a = &'a [T],
21> {
22    ptr: *mut T,
23    stride: usize,
24    _marker: PhantomData<(&'a T, Backend, Arch, Align, Ref)>,
25}
26
27unsafe impl<
28        'a,
29        T: NumericElement,
30        Backend,
31        Arch,
32        const ROWS: usize,
33        const COLS: usize,
34        Align: Alignment,
35        Ref: 'a,
36    > Send for TileView<'a, T, Backend, Arch, ROWS, COLS, Align, Ref>
37where
38    Ref: Send,
39{
40}
41
42unsafe impl<
43        'a,
44        T: NumericElement,
45        Backend,
46        Arch,
47        const ROWS: usize,
48        const COLS: usize,
49        Align: Alignment,
50        Ref: 'a,
51    > Sync for TileView<'a, T, Backend, Arch, ROWS, COLS, Align, Ref>
52where
53    Ref: Sync,
54{
55}
56
57impl<
58        'a,
59        T: NumericElement,
60        Backend,
61        Arch,
62        const ROWS: usize,
63        const COLS: usize,
64        Align: Alignment,
65        Ref: 'a,
66    > Clone for TileView<'a, T, Backend, Arch, ROWS, COLS, Align, Ref>
67where
68    Ref: Clone,
69{
70    #[inline(always)]
71    fn clone(&self) -> Self {
72        Self {
73            ptr: self.ptr,
74            stride: self.stride,
75            _marker: PhantomData,
76        }
77    }
78}
79
80impl<
81        'a,
82        T: NumericElement,
83        Backend,
84        Arch,
85        const ROWS: usize,
86        const COLS: usize,
87        Align: Alignment,
88        Ref: 'a,
89    > Copy for TileView<'a, T, Backend, Arch, ROWS, COLS, Align, Ref>
90where
91    Ref: Copy,
92{
93}
94
95impl<
96        'a,
97        T: NumericElement,
98        Backend,
99        Arch: crate::arch::SimdArch,
100        const ROWS: usize,
101        const COLS: usize,
102        Align: Alignment,
103    > TileView<'a, T, Backend, Arch, ROWS, COLS, Align, &'a [T]>
104{
105    /// Create a new read-only `TileView` after verifying bounds and alignment invariants.
106    /// Returns `None` if the input slice is too small or if alignment constraints are not met.
107    #[inline]
108    pub fn new(data: &'a [T], stride: usize) -> Option<Self> {
109        if data.len() < ROWS * stride {
110            return None;
111        }
112        if Align::IS_ALIGNED {
113            let req_align = Arch::REGISTER_WIDTH_BITS as usize / 8;
114            if req_align > 0 && Align::ALIGN_BYTES < req_align {
115                return None;
116            }
117            let addr = data.as_ptr() as usize;
118            if addr % Align::ALIGN_BYTES != 0 {
119                return None;
120            }
121        }
122        Some(Self {
123            ptr: data.as_ptr() as *mut T,
124            stride,
125            _marker: PhantomData,
126        })
127    }
128}
129
130impl<
131        'a,
132        T: NumericElement,
133        Backend,
134        Arch: crate::arch::SimdArch,
135        const ROWS: usize,
136        const COLS: usize,
137        Align: Alignment,
138    > TileView<'a, T, Backend, Arch, ROWS, COLS, Align, &'a mut [T]>
139{
140    /// Create a new mutable `TileView` after verifying bounds and alignment invariants.
141    /// Returns `None` if the input slice is too small or if alignment constraints are not met.
142    #[inline]
143    pub fn new_mut(data: &'a mut [T], stride: usize) -> Option<Self> {
144        if data.len() < ROWS * stride {
145            return None;
146        }
147        if Align::IS_ALIGNED {
148            let req_align = Arch::REGISTER_WIDTH_BITS as usize / 8;
149            if req_align > 0 && Align::ALIGN_BYTES < req_align {
150                return None;
151            }
152            let addr = data.as_ptr() as usize;
153            if addr % Align::ALIGN_BYTES != 0 {
154                return None;
155            }
156        }
157        Some(Self {
158            ptr: data.as_mut_ptr(),
159            stride,
160            _marker: PhantomData,
161        })
162    }
163
164    /// Access the underlying raw mutable pointer.
165    #[inline(always)]
166    pub fn as_mut_ptr(&mut self) -> *mut T {
167        self.ptr
168    }
169}
170
171impl<
172        'a,
173        T: NumericElement,
174        Backend,
175        Arch,
176        const ROWS: usize,
177        const COLS: usize,
178        Align: Alignment,
179        Ref: 'a,
180    > TileView<'a, T, Backend, Arch, ROWS, COLS, Align, Ref>
181{
182    /// Access the underlying raw pointer.
183    #[inline(always)]
184    pub fn as_ptr(&self) -> *const T {
185        self.ptr
186    }
187
188    /// Returns the row stride of this tile view.
189    #[inline(always)]
190    pub fn stride(&self) -> usize {
191        self.stride
192    }
193
194    /// Returns the number of rows.
195    #[inline(always)]
196    pub fn rows(&self) -> usize {
197        ROWS
198    }
199
200    /// Returns the number of columns.
201    #[inline(always)]
202    pub fn cols(&self) -> usize {
203        COLS
204    }
205}
206
207/// Trait mediating zero-overhead matrix multiplication on 2D tiles.
208///
209/// Implementations are fully monomorphized to optimize layout, vectorization, and register pressure.
210pub trait TileMatrixMultiply<
211    TA,
212    TB,
213    TC,
214    Backend,
215    Arch,
216    const M: usize,
217    const N: usize,
218    const K: usize,
219>
220{
221    /// Performs tile matrix multiplication: C += A * B
222    ///
223    /// # Safety
224    /// - Pointers `a`, `b`, and `c` must be valid for reads/writes of size M*a_stride, K*b_stride, M*c_stride.
225    unsafe fn tile_matmul(
226        c: *mut TC,
227        c_stride: usize,
228        a: *const TA,
229        a_stride: usize,
230        b: *const TB,
231        b_stride: usize,
232    );
233}