ab_riscv_interpreter/rv32/zk/zbkx/
rv32_zbkx_helpers.rs

1//! Opaque helpers for Zbkx extension
2
3use crate::const_utils::ConstRange;
4use const_fn_specialization::const_fn_specialization;
5
6#[const_fn_specialization]
7#[inline]
8#[doc(hidden)]
9#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
10pub fn xperm4(rs1: u32, rs2: u32) -> u32 {
11    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
12    cfg_select! {
13        all(not(miri), target_arch = "riscv32", target_feature = "zbkx") => unsafe {
14            core::arch::riscv32::xperm4(rs1 as usize, rs2 as usize) as u32
15        },
16        _ => {
17            use core::simd::num::SimdUint;
18            use core::simd::{simd_swizzle, u32x8};
19
20            const SHIFT: u32x8 = u32x8::from_array([0, 4, 8, 12, 16, 20, 24, 28]);
21            const MASK: u32x8 = u32x8::splat(0xf);
22
23            // Unpack nibbles of rs1 into bytes via SIMD: broadcast, shift per-lane, mask
24            let lut = (u32x8::splat(rs1) >> SHIFT) & MASK;
25            // Unpack nibbles of rs2 into byte indices via SIMD
26            let idx = (u32x8::splat(rs2) >> SHIFT) & MASK;
27            // For each nibble of rs2, look up from lut (out-of-bounds -> 0 via swizzle_dyn)
28            let nibbles = lut.cast().swizzle_dyn(idx.cast());
29            // Pack nibbles back: interleave even/odd lanes and fold into bytes
30            let lo = simd_swizzle!(nibbles, [0, 2, 4, 6]);
31            let hi = simd_swizzle!(nibbles, [1, 3, 5, 7]);
32            u32::from_le_bytes((lo | (hi << 4)).to_array())
33        }
34    }
35}
36
37#[const_fn_specialization]
38#[inline]
39#[doc(hidden)]
40#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
41pub const fn xperm4(rs1: u32, rs2: u32) -> u32 {
42    let mut result = 0;
43
44    // Each nibble of `rs2` selects a nibble of `rs1`, out of bounds selection yields zero
45    for nibble in ConstRange::new(0, u32::BITS / 4) {
46        let index = (rs2 >> (nibble * 4)) & 0xf;
47
48        if index < 8 {
49            result |= ((rs1 >> (index * 4)) & 0xf) << (nibble * 4);
50        }
51    }
52
53    result
54}
55
56#[const_fn_specialization]
57#[inline]
58#[doc(hidden)]
59#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
60pub fn xperm8(rs1: u32, rs2: u32) -> u32 {
61    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
62    cfg_select! {
63        all(not(miri), target_arch = "riscv32", target_feature = "zbkx") => unsafe {
64            core::arch::riscv32::xperm8(rs1 as usize, rs2 as usize) as u32
65        },
66        _ => {
67            use core::simd::u8x4;
68
69            let lut = u8x4::from_array(rs1.to_le_bytes());
70            let idx = u8x4::from_array(rs2.to_le_bytes());
71
72            let result = lut.swizzle_dyn(idx);
73
74            u32::from_le_bytes(result.to_array())
75        }
76    }
77}
78
79#[const_fn_specialization]
80#[inline]
81#[doc(hidden)]
82#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
83pub const fn xperm8(rs1: u32, rs2: u32) -> u32 {
84    let mut result = 0;
85
86    // Each byte of `rs2` selects a byte of `rs1`, out of bounds selection yields zero
87    for byte in ConstRange::new(0, u32::BITS / 8) {
88        let index = (rs2 >> (byte * 8)) & 0xff;
89
90        if index < 4 {
91            result |= ((rs1 >> (index * 8)) & 0xff) << (byte * 8);
92        }
93    }
94
95    result
96}
ab_riscv_interpreter/rv32/zk/zbkx/rv32_zbkx_helpers.rs

ab_riscv_interpreter/rv32/zk/zbkx/
rv32_zbkx_helpers.rs
