Skip to main content

ab_riscv_interpreter/v/
vector_registers.rs

1//! Vector registers
2
3use crate::Csrs;
4use ab_riscv_primitives::prelude::*;
5use core::marker::Destruct;
6
7pub(crate) const VLENB_USIZE<const VLEN: Vlen>: usize = VLEN.bytes() as usize;
8/// Element width in bytes as `usize`
9const EEW_BYTES<const EEW: Eew>: usize = EEW.bytes_width() as usize;
10
11/// Alignment wrapper for vector registers
12#[derive(Debug, Clone, Copy)]
13// Aligned to 128 bytes, which is u32 * 32 registers, the minimum reasonable value to use in most
14// cases
15#[repr(align(128))]
16pub struct VectorRegisterFile<const VLEN: Vlen>([[u8; VLENB_USIZE::<VLEN>]; 32]);
17
18const impl<const VLEN: Vlen> Default for VectorRegisterFile<VLEN> {
19    #[inline(always)]
20    fn default() -> Self {
21        Self([[0; _]; _])
22    }
23}
24
25impl<const VLEN: Vlen> VectorRegisterFile<VLEN> {
26    /// Get reference to a vector register
27    #[inline(always)]
28    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
29    pub const fn get(&self, index: VReg) -> &[u8; VLENB_USIZE::<VLEN>] {
30        // SAFETY: Always in-range
31        unsafe { self.0.get_unchecked(usize::from(index.to_bits())) }
32    }
33
34    /// Get mutable reference to a vector register
35    #[inline(always)]
36    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
37    pub const fn get_mut(&mut self, index: VReg) -> &mut [u8; VLENB_USIZE::<VLEN>] {
38        // SAFETY: Always in-range
39        unsafe { self.0.get_unchecked_mut(usize::from(index.to_bits())) }
40    }
41
42    /// All vector registers as one contiguous array of bytes.
43    ///
44    /// Register `v` occupies bytes `[v * VLENB, (v + 1) * VLENB)` of the flattened array, so a
45    /// register group is a contiguous range and its elements are at [`Self::element_offset()`].
46    #[inline(always)]
47    pub const fn as_bytes(&self) -> &[[u8; VLENB_USIZE::<VLEN>]; 32] {
48        &self.0
49    }
50
51    /// All vector registers as one contiguous mutable array of bytes, see [`Self::as_bytes()`]
52    #[inline(always)]
53    pub const fn as_bytes_mut(&mut self) -> &mut [[u8; VLENB_USIZE::<VLEN>]; 32] {
54        &mut self.0
55    }
56
57    /// Byte offset within flattened [`Self::as_bytes()`] of element `elem_i` in the register
58    /// group starting at `base_reg`, with `eew`-wide elements.
59    ///
60    /// Element widths divide `VLENB`, so elements never straddle registers and a register group
61    /// is one contiguous array of elements.
62    #[inline(always)]
63    pub const fn element_offset<W>(base_reg: VReg, elem_i: u16, eew: W) -> usize
64    where
65        W: [const] Into<Eew>,
66    {
67        usize::from(base_reg.to_bits()) * VLENB_USIZE::<VLEN>
68            + usize::from(elem_i) * usize::from(eew.into().bytes_width())
69    }
70
71    /// Read element `elem_i` of the register group starting at `base_reg`, with `eew`-wide
72    /// elements, zero-extended.
73    ///
74    /// # Safety
75    /// The element must lie within the register file, which holds for any `elem_i < vl` of a
76    /// register group `[base_reg, base_reg + group_regs)` that ends within the register file,
77    /// since `vl <= group_regs * VLENB / eew.bytes_width()`.
78    #[inline(always)]
79    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
80    pub const unsafe fn read_element<W>(&self, base_reg: VReg, elem_i: u16, eew: W) -> u64
81    where
82        W: [const] Into<Eew> + [const] Destruct,
83    {
84        // SAFETY: Guaranteed by the caller's precondition
85        unsafe {
86            match eew.into() {
87                Eew::E8 => self.read_element_const::<{ Eew::E8 }>(base_reg, elem_i),
88                Eew::E16 => self.read_element_const::<{ Eew::E16 }>(base_reg, elem_i),
89                Eew::E32 => self.read_element_const::<{ Eew::E32 }>(base_reg, elem_i),
90                Eew::E64 => self.read_element_const::<{ Eew::E64 }>(base_reg, elem_i),
91            }
92        }
93    }
94
95    /// Write the low `eew.bytes_width()` bytes of `value` into element `elem_i` of the register
96    /// group starting at `base_reg`, with `eew`-wide elements.
97    ///
98    /// # Safety
99    /// Same as [`Self::read_element()`]
100    #[inline(always)]
101    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
102    pub const unsafe fn write_element<W>(&mut self, base_reg: VReg, elem_i: u16, eew: W, value: u64)
103    where
104        W: [const] Into<Eew> + [const] Destruct,
105    {
106        // SAFETY: Guaranteed by the caller's precondition
107        unsafe {
108            match eew.into() {
109                Eew::E8 => self.write_element_const::<{ Eew::E8 }>(base_reg, elem_i, value),
110                Eew::E16 => self.write_element_const::<{ Eew::E16 }>(base_reg, elem_i, value),
111                Eew::E32 => self.write_element_const::<{ Eew::E32 }>(base_reg, elem_i, value),
112                Eew::E64 => self.write_element_const::<{ Eew::E64 }>(base_reg, elem_i, value),
113            }
114        }
115    }
116
117    /// [`Self::read_element()`] with the element width known at compile time, which makes the
118    /// access a fixed-size load
119    ///
120    /// # Safety
121    /// Same as [`Self::read_element()`]
122    #[inline(always)]
123    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
124    pub const unsafe fn read_element_const<const EEW: Eew>(
125        &self,
126        base_reg: VReg,
127        elem_i: u16,
128    ) -> u64 {
129        let offset = Self::element_offset(base_reg, elem_i, EEW);
130        // SAFETY: `offset + EEW.bytes_width() <= 32 * VLENB` by the caller's precondition
131        let element = unsafe {
132            self.as_bytes()
133                .as_flattened()
134                .get_unchecked(offset..)
135                .first_chunk::<{ EEW_BYTES::<EEW> }>()
136                .unwrap_unchecked()
137        };
138        let mut bytes = 0u64.to_le_bytes();
139        if let Some((low, _)) = bytes.split_first_chunk_mut::<{ EEW_BYTES::<EEW> }>() {
140            *low = *element;
141        }
142        u64::from_le_bytes(bytes)
143    }
144
145    /// [`Self::write_element()`] with the element width known at compile time, which makes the
146    /// access a fixed-size store
147    ///
148    /// # Safety
149    /// Same as [`Self::read_element()`]
150    #[inline(always)]
151    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
152    pub const unsafe fn write_element_const<const EEW: Eew>(
153        &mut self,
154        base_reg: VReg,
155        elem_i: u16,
156        value: u64,
157    ) {
158        let offset = Self::element_offset(base_reg, elem_i, EEW);
159        // SAFETY: `offset + EEW.bytes_width() <= 32 * VLENB` by the caller's precondition
160        let element = unsafe {
161            self.as_bytes_mut()
162                .as_flattened_mut()
163                .get_unchecked_mut(offset..)
164                .first_chunk_mut::<{ EEW_BYTES::<EEW> }>()
165                .unwrap_unchecked()
166        };
167        if let Some((low, _)) = value
168            .to_le_bytes()
169            .split_first_chunk::<{ EEW_BYTES::<EEW> }>()
170        {
171            *element = *low;
172        }
173    }
174}
175
176/// Vector register state.
177///
178/// This trait contains only methods that implementations genuinely need to provide. Derived
179/// accessors for simpler CSRs are in [`VectorRegistersExt`].
180///
181/// Note that due to Rust type system limitations, you should use [`VectorRegistersExt`] in trait
182/// bounds instead of this trait directly or else the solver will fail.
183///
184/// Methods for `vtype` and `vl` live here (not in the ext trait) because they have non-trivial
185/// update semantics: `vtype` must maintain a cached decoded form and handle the XLEN-dependent vill
186/// bit, and `vl` is read-only via CSR instructions but writable by `vsetvl{i}` and fault-only-first
187/// loads.
188pub const trait VectorRegisters {
189    /// Maximum vector element width `ELEN` in bits
190    const ELEN: Elen;
191    /// Vector register width `VLEN` in bits
192    const VLEN: Vlen;
193
194    /// Read the vector register file
195    fn read_vregs(&self) -> &VectorRegisterFile<{ Self::VLEN }>;
196
197    /// Mutable access to the vector register file
198    fn write_vregs(&mut self) -> &mut VectorRegisterFile<{ Self::VLEN }>;
199
200    /// Check whether vector instructions are currently permitted.
201    ///
202    /// Returns `false` when `mstatus.VS == Off` (or equivalent like `sstatus`/`vstatus`). In
203    /// environments without these status registers, returns `true` always.
204    fn vector_instructions_allowed(&self) -> bool;
205
206    /// Mark the vector state as dirty.
207    ///
208    /// Must set VS to Dirty in `mstatus` (and `sstatus`/`vsstatus` shadows) when those registers
209    /// exist. No-op otherwise.
210    fn mark_vs_dirty(&mut self);
211
212    /// Compute `vl` from `AVL` and `VLMAX` per spec constraints.
213    ///
214    /// The simplest compliant implementation (which is used by default) is `min(AVL, VLMAX)`. More
215    /// sophisticated implementations may return values in `[ceil(AVL/2), VLMAX]` for
216    /// `AVL < 2*VLMAX`, but this simple strategy satisfies all three spec requirements.
217    #[inline(always)]
218    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
219    fn compute_vl(&self, avl: Vl, vlmax: Vl) -> Vl {
220        avl.min(vlmax)
221    }
222
223    /// Compute `VLMAX` for a given vtype
224    #[inline(always)]
225    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
226    fn vlmax_for_vtype(&self, vtype: Vtype<{ Self::ELEN }, { Self::VLEN }>) -> Vl
227    where
228        [(); SUPPORTED_ELEN_VLEN::<{ Self::ELEN }, { Self::VLEN }>]:,
229    {
230        vtype.vlmul().vlmax::<{ Self::VLEN }>(vtype.vsew())
231    }
232}
233
234/// Derived convenience accessors for vector CSRs that are simple read/write fields (vstart, vxrm,
235/// vxsat, vcsr).
236///
237/// Intended for types that implement both [`VectorRegisters`] and [`Csrs`].
238///
239/// NOTE: While the default methods implemented via the [`Csrs`] trait are correct, custom
240/// higher-performance implementations are often possible by overriding them and, for example,
241/// caching various CSRs as separate pre-decoded values rather than going through a generic code
242/// path with XLEN-sized raw CSR values during reads.
243pub const trait VectorRegistersExt<Reg>
244where
245    Self: [const] Csrs<Reg> + [const] VectorRegisters,
246    [(); SUPPORTED_ELEN_VLEN::<{ Self::ELEN }, { Self::VLEN }>]:,
247    Reg: [const] Register,
248{
249    /// Initialize the vector state to the recommended default configuration.
250    ///
251    /// Per spec: `vtype.vill` = 1, remaining `vtype` bits = `0`, `vl` = 0.
252    /// `vstart`, `vxrm`, `vxsat` may have arbitrary values at reset but are zeroed here for
253    /// deterministic behavior.
254    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
255    fn initialize_vector_state(&mut self) {
256        self.set_vtype(None);
257        self.set_vl(Vl::ZERO);
258        self.set_vstart(Vstart::ZERO);
259        self.set_vxrm(Vxrm::default());
260        self.set_vxsat(false);
261    }
262
263    /// Get current `vstart`
264    #[inline(always)]
265    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
266    fn vstart(&self) -> Vstart {
267        let raw = self
268            .read_csr(VectorCsr::Vstart.to_csr_index())
269            .unwrap_or_default()
270            .as_u64();
271        Vstart::from(raw as u16)
272    }
273
274    /// Set `vstart`.
275    ///
276    /// The default implementation ignores writes to uninitialized CSR in release mode and panics in
277    /// debug.
278    #[inline(always)]
279    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
280    fn set_vstart(&mut self, vstart: Vstart) {
281        let result = self.write_csr(
282            VectorCsr::Vstart.to_csr_index(),
283            Reg::Type::from(u16::from(vstart)),
284        );
285        debug_assert!(
286            result.is_ok(),
287            "Implementation must initialize `vstart` CSR"
288        );
289    }
290
291    /// Reset `vstart` to zero.
292    ///
293    /// Per spec, all vector instructions reset `vstart` to zero at the end of execution.
294    #[inline(always)]
295    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
296    fn reset_vstart(&mut self) {
297        self.set_vstart(Vstart::ZERO);
298    }
299
300    /// Get `vxsat` (single bit)
301    #[inline(always)]
302    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
303    fn vxsat(&self) -> bool {
304        let raw = self
305            .read_csr(VectorCsr::Vxsat.to_csr_index())
306            .unwrap_or_default()
307            .as_u64();
308        (raw & 1) == 1
309    }
310
311    /// Set `vxsat`.
312    ///
313    /// The default implementation ignores writes to uninitialized CSR in release mode and panics in
314    /// debug.
315    #[inline(always)]
316    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
317    fn set_vxsat(&mut self, vxsat: bool) {
318        let masked = Reg::Type::from(u8::from(vxsat));
319        let result = self.write_csr(VectorCsr::Vxsat.to_csr_index(), masked);
320        debug_assert!(result.is_ok(), "Implementation must initialize `vxsat` CSR");
321        // Mirror `vxsat` into `vcsr[0]`, preserving `vcsr[2:1]` (`vxrm`)
322        let old_vcsr = self
323            .read_csr(VectorCsr::Vcsr.to_csr_index())
324            .unwrap_or_default();
325        let new_vcsr = (old_vcsr & !Reg::Type::from(1u8)) | masked;
326        let result = self.write_csr(VectorCsr::Vcsr.to_csr_index(), new_vcsr);
327        debug_assert!(result.is_ok(), "Implementation must initialize `vcsr` CSR");
328    }
329
330    /// Get `vxrm`
331    #[inline(always)]
332    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
333    fn vxrm(&self) -> Vxrm {
334        let raw = self
335            .read_csr(VectorCsr::Vxrm.to_csr_index())
336            .unwrap_or_default()
337            .as_u64();
338        Vxrm::from_bits(raw as u8)
339    }
340
341    /// Set `vxrm`.
342    ///
343    /// The default implementation ignores writes to uninitialized CSR in release mode and panics in
344    /// debug.
345    #[inline(always)]
346    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
347    fn set_vxrm(&mut self, vxrm: Vxrm) {
348        let masked = Reg::Type::from(vxrm.to_bits());
349        let result = self.write_csr(VectorCsr::Vxrm.to_csr_index(), masked);
350        debug_assert!(result.is_ok(), "Implementation must initialize `vxrm` CSR");
351        // Mirror `vxrm` into `vcsr[2:1]`, preserving `vcsr[0]` (`vxsat`)
352        let old_vcsr = self
353            .read_csr(VectorCsr::Vcsr.to_csr_index())
354            .unwrap_or_default();
355        let new_vcsr = (old_vcsr & !Reg::Type::from(0b110u8)) | (masked << 1u8);
356        let result = self.write_csr(VectorCsr::Vcsr.to_csr_index(), new_vcsr);
357        debug_assert!(result.is_ok(), "Implementation must initialize `vcsr` CSR");
358    }
359
360    /// Get the current vl
361    #[inline(always)]
362    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
363    fn vl(&self) -> Vl {
364        let vl = self
365            .read_csr(VectorCsr::Vl.to_csr_index())
366            .unwrap_or_default()
367            .as_u64() as u32;
368        // Should always be `Some()`, but can't be guaranteed here
369        Vl::new(vl).unwrap_or_default()
370    }
371
372    /// Set vl.
373    ///
374    /// The implementation must update both its internal decoded cache and the raw CSR value (for
375    /// reads via Zicsr, writes via Zicsr are not allowed).
376    ///
377    /// The default implementation ignores writes to uninitialized CSR in release mode and panics in
378    /// debug.
379    #[inline(always)]
380    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
381    fn set_vl(&mut self, vl: Vl) {
382        let result = self.write_csr(VectorCsr::Vl.to_csr_index(), Reg::Type::from(u32::from(vl)));
383        debug_assert!(result.is_ok(), "Implementation must initialize `vl` CSR");
384    }
385
386    /// Get the current decoded vtype
387    #[inline(always)]
388    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
389    fn vtype(&self) -> Option<Vtype<{ Self::ELEN }, { Self::VLEN }>> {
390        self.read_csr(VectorCsr::Vtype.to_csr_index())
391            .ok()
392            .and_then(Vtype::from_raw::<Reg>)
393    }
394
395    /// Set the vtype register from a decoded `Vtype`.
396    ///
397    /// The implementation must update both its internal decoded cache and the raw CSR value (for
398    /// reads via Zicsr, writes via Zicsr are not allowed).
399    ///
400    /// The default implementation ignores writes to uninitialized CSR in release mode and panics in
401    /// debug.
402    #[inline(always)]
403    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
404    fn set_vtype(&mut self, vtype: Option<Vtype<{ Self::ELEN }, { Self::VLEN }>>) {
405        let vtype_raw = if let Some(vt) = vtype {
406            vt.to_raw::<Reg>()
407        } else {
408            Vtype::<{ Self::ELEN }, { Self::VLEN }>::illegal_raw::<Reg>()
409        };
410
411        let result = self.write_csr(VectorCsr::Vtype.to_csr_index(), vtype_raw);
412        debug_assert!(result.is_ok(), "Implementation must initialize `vtype` CSR");
413    }
414}
415
416// Convenience for threaded execution
417// TODO: Forward generically instead, once the compiler normalizes
418//  `<&mut T as VectorRegisters>::VLEN` to `T::VLEN`:
419//  https://github.com/rust-lang/rust/issues/161264
420#[macro_export]
421macro_rules! impl_vector_registers_for_mut_ref {
422    ($env:ty, $reg:ty) => {
423        impl VectorRegisters for &mut $env {
424            const ELEN: Elen = <$env as VectorRegisters>::ELEN;
425            const VLEN: Vlen = <$env as VectorRegisters>::VLEN;
426
427            #[inline(always)]
428            fn read_vregs(&self) -> &VectorRegisterFile<{ Self::VLEN }> {
429                <$env as VectorRegisters>::read_vregs(self)
430            }
431
432            #[inline(always)]
433            fn write_vregs(&mut self) -> &mut VectorRegisterFile<{ Self::VLEN }> {
434                <$env as VectorRegisters>::write_vregs(self)
435            }
436
437            #[inline(always)]
438            fn vector_instructions_allowed(&self) -> bool {
439                <$env as VectorRegisters>::vector_instructions_allowed(self)
440            }
441
442            #[inline(always)]
443            fn mark_vs_dirty(&mut self) {
444                <$env as VectorRegisters>::mark_vs_dirty(self);
445            }
446        }
447
448        // Every method is forwarded explicitly rather than left to `VectorRegistersExt`'s
449        // defaults: those go through `Csrs::write_csr()`, which the blanket `Csrs for &mut T`
450        // impl also forwards to `$env`, so an empty impl here would silently observe `$env`'s
451        // overrides for some accessors and the trait defaults for others
452        impl VectorRegistersExt<$reg> for &mut $env {
453            #[inline(always)]
454            fn vstart(&self) -> Vstart {
455                <$env as VectorRegistersExt<$reg>>::vstart(self)
456            }
457
458            #[inline(always)]
459            fn set_vstart(&mut self, vstart: Vstart) {
460                <$env as VectorRegistersExt<$reg>>::set_vstart(self, vstart);
461            }
462
463            #[inline(always)]
464            fn reset_vstart(&mut self) {
465                <$env as VectorRegistersExt<$reg>>::reset_vstart(self);
466            }
467
468            #[inline(always)]
469            fn vxsat(&self) -> bool {
470                <$env as VectorRegistersExt<$reg>>::vxsat(self)
471            }
472
473            #[inline(always)]
474            fn set_vxsat(&mut self, vxsat: bool) {
475                <$env as VectorRegistersExt<$reg>>::set_vxsat(self, vxsat);
476            }
477
478            #[inline(always)]
479            fn vxrm(&self) -> Vxrm {
480                <$env as VectorRegistersExt<$reg>>::vxrm(self)
481            }
482
483            #[inline(always)]
484            fn set_vxrm(&mut self, vxrm: Vxrm) {
485                <$env as VectorRegistersExt<$reg>>::set_vxrm(self, vxrm);
486            }
487
488            #[inline(always)]
489            fn vl(&self) -> Vl {
490                <$env as VectorRegistersExt<$reg>>::vl(self)
491            }
492
493            #[inline(always)]
494            fn set_vl(&mut self, vl: Vl) {
495                <$env as VectorRegistersExt<$reg>>::set_vl(self, vl);
496            }
497
498            #[inline(always)]
499            fn vtype(&self) -> Option<Vtype<{ Self::ELEN }, { Self::VLEN }>> {
500                <$env as VectorRegistersExt<$reg>>::vtype(self)
501            }
502
503            #[inline(always)]
504            fn set_vtype(&mut self, vtype: Option<Vtype<{ Self::ELEN }, { Self::VLEN }>>) {
505                <$env as VectorRegistersExt<$reg>>::set_vtype(self, vtype);
506            }
507
508            #[inline(always)]
509            fn initialize_vector_state(&mut self) {
510                <$env as VectorRegistersExt<$reg>>::initialize_vector_state(self);
511            }
512        }
513    };
514}