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

ab_riscv_interpreter/rv64/zk/zbkx/
rv64_zbkx_helpers.rs
