ab_riscv_interpreter/v/zvexx/reduction/
zvexx_reduction_helpers.rs

1//! Opaque helpers for ZveXx extension
2use crate::v::vector_registers::VectorRegistersExt;
3use crate::v::zvexx::arith::zvexx_arith_helpers::{
4    read_element_u64, sign_extend, write_element_u64,
5};
6use crate::v::zvexx::load::zvexx_load_helpers::{mask_bit, snapshot_mask};
7use ab_riscv_primitives::prelude::*;
8use core::hint::cold_path;
9
10/// Execute a single-width integer reduction.
11///
12/// # Safety
13/// - `vs2.to_bits() % group_regs == 0` and `vs2.to_bits() + group_regs <= 32` (verified by caller)
14/// - `vstart == 0` (verified by caller; reductions with non-zero vstart are illegal)
15/// - `vl <= group_regs * VLEN.bytes() / sew_bytes`
16/// - `vl <= VLEN`
17#[inline(always)]
18#[expect(clippy::too_many_arguments, reason = "Internal API")]
19#[doc(hidden)]
20#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
21pub unsafe fn execute_reduce_op<Reg, Env, F>(
22    env: &mut Env,
23    vd: VReg,
24    vs2: VReg,
25    vs1: VReg,
26    vm: bool,
27    vl: Vl,
28    sew: Vsew,
29    op: F,
30) where
31    Reg: Register,
32    Env: VectorRegistersExt<Reg>,
33    [(); SUPPORTED_ELEN_VLEN::<{ Env::ELEN }, { Env::VLEN }>]:,
34    F: Fn(u64, u64, Vsew) -> u64,
35{
36    // Spec §5.4: when vstart >= vl, no element of vd is updated. For reductions this means
37    // vl == 0 (since caller has verified vstart == 0). In that case we must not write vd and
38    // must not mark vs dirty.
39    if vl == Vl::ZERO {
40        cold_path();
41        env.reset_vstart();
42        return;
43    }
44    // SAFETY: element 0 always fits within register vs1
45    let init = unsafe { read_element_u64(env.read_vregs(), vs1, 0, sew) };
46    // SAFETY: `vl <= VLEN`
47    let mask_buf = unsafe { snapshot_mask(env.read_vregs(), vm, vl) };
48    let mut acc = init;
49    for i in Vstart::ZERO.range_to(vl) {
50        if !mask_bit(&mask_buf, i) {
51            continue;
52        }
53        // SAFETY: `vs2 % group_regs == 0` and `i < vl <= group_regs * elems_per_reg`
54        let elem = unsafe { read_element_u64(env.read_vregs(), vs2, i, sew) };
55        acc = op(acc, elem, sew);
56    }
57    // SAFETY: element 0 always fits within register vd
58    unsafe {
59        write_element_u64(env.write_vregs(), vd, 0, sew, acc);
60    }
61    env.mark_vs_dirty();
62    env.reset_vstart();
63}
64
65/// Execute a widening integer sum reduction.
66///
67/// # Safety
68/// - `vs2.to_bits() % group_regs == 0` and `vs2.to_bits() + group_regs <= 32` (verified by caller)
69/// - `sew.double_width().is_some()` (verified by caller)
70/// - `vstart == 0` (verified by caller)
71/// - `vl <= group_regs * VLEN.bytes() / sew_bytes`
72/// - `vl <= VLEN`
73#[inline(always)]
74#[expect(clippy::too_many_arguments, reason = "Internal API")]
75#[doc(hidden)]
76#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
77pub unsafe fn execute_widening_reduce_op<const SIGN_EXTEND_SRC: bool, Reg, Env, F>(
78    env: &mut Env,
79    vd: VReg,
80    vs2: VReg,
81    vs1: VReg,
82    vm: bool,
83    vl: Vl,
84    sew: Vsew,
85    op: F,
86) where
87    Reg: Register,
88    Env: VectorRegistersExt<Reg>,
89    [(); SUPPORTED_ELEN_VLEN::<{ Env::ELEN }, { Env::VLEN }>]:,
90    F: Fn(u64, u64, Vsew) -> u64,
91{
92    let Some(wide_sew) = sew.double_width() else {
93        // SAFETY: caller verified `2*SEW <= ELEN`; E64 widening is unreachable here
94        unsafe { core::hint::unreachable_unchecked() }
95    };
96    if vl == Vl::ZERO {
97        cold_path();
98        env.reset_vstart();
99        return;
100    }
101    // SAFETY: element 0 always fits within register vs1
102    let init = unsafe { read_element_u64(env.read_vregs(), vs1, 0, wide_sew) };
103    // SAFETY: `vl <= VLEN`
104    let mask_buf = unsafe { snapshot_mask(env.read_vregs(), vm, vl) };
105    let mut acc = init;
106    for i in Vstart::ZERO.range_to(vl) {
107        if !mask_bit(&mask_buf, i) {
108            continue;
109        }
110        // SAFETY: same bounds argument as `execute_reduce_op`
111        let raw = unsafe { read_element_u64(env.read_vregs(), vs2, i, sew) };
112        let elem = if SIGN_EXTEND_SRC {
113            sign_extend(raw, sew).cast_unsigned()
114        } else {
115            raw
116        };
117        acc = op(acc, elem, wide_sew);
118    }
119    // SAFETY: element 0 always fits within register vd
120    unsafe {
121        write_element_u64(env.write_vregs(), vd, 0, wide_sew, acc);
122    }
123    env.mark_vs_dirty();
124    env.reset_vstart();
125}
ab_riscv_interpreter/v/zvexx/reduction/zvexx_reduction_helpers.rs

ab_riscv_interpreter/v/zvexx/reduction/
zvexx_reduction_helpers.rs
