Skip to main content

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