diff --git a/ext/crates/fp/src/field/bitslice.rs b/ext/crates/fp/src/field/bitslice.rs new file mode 100644 index 0000000000..0f336ba122 --- /dev/null +++ b/ext/crates/fp/src/field/bitslice.rs @@ -0,0 +1,559 @@ +//! Bit-sliced arithmetic kernels for prime fields. +//! +//! In the bit-sliced layout, a group of [`BITS_PER_LIMB`] (64) field elements occupies +//! `k = ceil(log2 p)` consecutive limbs (the *planes*): plane `j` holds bit `j` of all 64 +//! elements, with element `i` living at bit `i` of each plane. Addition and scalar +//! multiplication then reduce to short branch-free boolean circuits over the planes that +//! act on 64 lanes at once, with no separate reduction step. +//! +//! Addition is a ripple-carry adder over the `k` planes (producing a `(k+1)`-bit sum in +//! `[0, 2p)`) followed by a single conditional subtraction of `p`. Scalar multiplication is +//! double-and-add with a modular reduction at each step. The number of planes `k` is +//! dispatched to a const-generic implementation so that, for each prime, the arrays are +//! exactly sized and the loops fully unrolled; a heap-scratch fallback covers the rare +//! primes with `k` beyond the dispatch range. +//! +//! The idea of bit-slicing finite-field vectors, and the optimized F3 addition circuit in +//! [`f3_add_planes`], were both contributed by Carl McTague (). + +use crate::{constants::BITS_PER_LIMB, limb::Limb}; + +/// Largest `k` that the const-generic dispatch covers directly (`p < 2^16`). Larger primes +/// fall back to the heap-scratch path. +const MAX_DISPATCH_K: usize = 16; + +/// The number of planes `k = ceil(log2 p)` needed to bit-slice an element of `F_p`. +pub(crate) fn planes(p: u32) -> usize { + debug_assert!(p >= 2); + (u32::BITS - (p - 1).leading_zeros()) as usize +} + +/// The bits of `p` as full-width lane masks: `out[j]` is all-ones iff bit `j` of `p` is set, +/// for `j` in `0..=k`. +fn p_masks(p: u32, k: usize) -> [Limb; BITS_PER_LIMB + 1] { + let mut masks = [0; BITS_PER_LIMB + 1]; + for (i, m) in masks.iter_mut().enumerate().take(k + 1) { + if (p >> i) & 1 == 1 { + *m = !0; + } + } + masks +} + +/// The full-width lane mask for bit `j` of `p`: all-ones if set, zero otherwise. +#[inline(always)] +fn pmask(p: u32, j: usize) -> Limb { + // Broadcast bit `j` across the limb via two's-complement negation: `1 -> !0`, `0 -> 0`. + Limb::from((p >> j) & 1).wrapping_neg() +} + +/// `dst += c * src` (mod p) over every group, where `dst` and `src` hold the same number of +/// whole groups of `k` planes. Assumes both are reduced; the result is reduced. +pub(crate) fn add_groups(p: u32, k: usize, dst: &mut [Limb], src: &[Limb], c: u32) { + if c == 0 { + return; + } + if p == 2 { + // One plane (k = 1); the only nonzero scalar is 1, so addition is XOR. + for (d, s) in dst.iter_mut().zip(src) { + *d ^= *s; + } + return; + } + if p == 3 { + return f3_add_groups(dst, src, c); + } + if p == 5 { + return f5_add_groups(dst, src, c); + } + macro_rules! dispatch { + ($($k:literal),*) => { + match k { + $($k => add_groups_k::<$k>(dst, src, c, p),)* + _ => add_groups_dyn(k, dst, src, c, &p_masks(p, k)), + } + }; + } + dispatch!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16); +} + +/// `dst += c * src` (mod p) for a single group of `k` planes, restricted to the lanes set in +/// `lane_mask`. Lanes outside the mask are unchanged. `dst` and `src` are each exactly `k` +/// limbs. +pub(crate) fn add_group_masked( + p: u32, + k: usize, + dst: &mut [Limb], + src: &[Limb], + c: u32, + lane_mask: Limb, +) { + if c == 0 { + return; + } + if p == 2 { + dst[0] ^= src[0] & lane_mask; + return; + } + if p == 3 { + return f3_add_group_masked(dst, src, c, lane_mask); + } + if p == 5 { + return f5_add_group_masked(dst, src, c, lane_mask); + } + macro_rules! dispatch { + ($($k:literal),*) => { + match k { + $($k => add_group_masked_k::<$k>(dst, src, c, p, lane_mask),)* + _ => { + // Masking `src` to the in-range lanes makes the circuit a no-op elsewhere. + let mut masked = vec![0; k]; + for j in 0..k { + masked[j] = src[j] & lane_mask; + } + add_groups(p, k, dst, &masked, c); + } + } + }; + } + dispatch!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16); +} + +/// `dst *= c` (mod p) over every group. +pub(crate) fn scale_groups(p: u32, k: usize, dst: &mut [Limb], c: u32) { + if c == 1 { + return; + } + if c == 0 { + dst.fill(0); + return; + } + if p == 3 { + return f3_scale_groups(dst, c); + } + if p == 5 { + return f5_scale_groups(dst, c); + } + macro_rules! dispatch { + ($($k:literal),*) => { + match k { + $($k => scale_groups_k::<$k>(dst, c, p),)* + _ => scale_groups_dyn(k, dst, c, &p_masks(p, k)), + } + }; + } + dispatch!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16); +} + +// --------------------------------------------------------------------------------------- +// Const-generic kernels: `K` planes known at compile time, so every array is exactly sized +// and every loop is fully unrolled. +// --------------------------------------------------------------------------------------- + +/// Reduce a `(K+1)`-bit unreduced sum (`s` low planes + `s_top`) in `[0, 2p)` to `s mod p`. +/// The per-plane mask is computed inline from `p` (no per-call mask array). +#[inline(always)] +fn cond_sub_k(s: &[Limb; K], s_top: Limb, p: u32) -> [Limb; K] { + let mut d = [0 as Limb; K]; + let mut borrow: Limb = 0; + for j in 0..K { + let sj = s[j]; + let pj = pmask(p, j); + let sxp = sj ^ pj; + d[j] = sxp ^ borrow; + borrow = (!sj & pj) | (borrow & !sxp); + } + // Top bit only affects the borrow-out (the result fits in K planes since result < p). + let pj = pmask(p, K); + let sxp = s_top ^ pj; + borrow = (!s_top & pj) | (borrow & !sxp); + let ge = !borrow; + let mut out = [0 as Limb; K]; + for j in 0..K { + out[j] = (d[j] & ge) | (s[j] & !ge); + } + out +} + +/// `(a + b) mod p` over `K` planes. +#[inline(always)] +fn add_mod_k(a: &[Limb; K], b: &[Limb; K], p: u32) -> [Limb; K] { + let mut s = [0 as Limb; K]; + let mut carry: Limb = 0; + for j in 0..K { + let aj = a[j]; + let bj = b[j]; + let axb = aj ^ bj; + s[j] = axb ^ carry; + carry = (aj & bj) | (carry & axb); + } + cond_sub_k::(&s, carry, p) +} + +/// `(2 * a) mod p` over `K` planes (doubling is a one-position plane shift). +#[inline(always)] +fn double_mod_k(a: &[Limb; K], p: u32) -> [Limb; K] { + let mut s = [0 as Limb; K]; + s[1..K].copy_from_slice(&a[..K - 1]); + let s_top = a[K - 1]; + cond_sub_k::(&s, s_top, p) +} + +/// `(c * b) mod p` over `K` planes, via double-and-add. +#[inline(always)] +fn scalar_mul_k(b: &[Limb; K], c: u32, p: u32) -> [Limb; K] { + let mut result = [0 as Limb; K]; + let mut temp = *b; + let mut cc = c; + loop { + if cc & 1 == 1 { + result = add_mod_k::(&result, &temp, p); + } + cc >>= 1; + if cc == 0 { + break; + } + temp = double_mod_k::(&temp, p); + } + result +} + +#[inline] +fn add_groups_k(dst: &mut [Limb], src: &[Limb], c: u32, p: u32) { + for (dg, sg) in dst + .as_chunks_mut::() + .0 + .iter_mut() + .zip(src.as_chunks::().0) + { + let addend = if c == 1 { + *sg + } else { + scalar_mul_k::(sg, c, p) + }; + *dg = add_mod_k::(dg, &addend, p); + } +} + +/// `dst += c * src` (mod p) for a single `K`-plane group, restricted to lanes in `lane_mask`. +#[inline] +fn add_group_masked_k( + dst: &mut [Limb], + src: &[Limb], + c: u32, + p: u32, + lane_mask: Limb, +) { + let mut a = [0 as Limb; K]; + let mut b = [0 as Limb; K]; + for j in 0..K { + a[j] = dst[j]; + b[j] = src[j] & lane_mask; + } + let addend = if c == 1 { + b + } else { + scalar_mul_k::(&b, c, p) + }; + let sum = add_mod_k::(&a, &addend, p); + dst[..K].copy_from_slice(&sum); +} + +#[inline] +fn scale_groups_k(dst: &mut [Limb], c: u32, p: u32) { + if c == 1 { + return; + } + if c == 0 { + dst.fill(0); + return; + } + for dg in dst.as_chunks_mut::().0 { + *dg = scalar_mul_k::(dg, c, p); + } +} + +// --------------------------------------------------------------------------------------- +// F3 specialization (k = 2). Plane 0 is the low bit, plane 1 the high bit, so an element +// `v in {0,1,2}` is stored as `(hi, lo)` with `v = 2*hi + lo`. Addition is a flat boolean +// circuit (no ripple-carry or borrow chain), and multiplication by 2 = negation just swaps +// the two planes — both avoid the sequential dependencies that make the generic circuit lose +// to the packed SWAR reduce at small primes. +// --------------------------------------------------------------------------------------- + +/// `(a + b) mod 3` as a flat 6-gate circuit on the `(lo, hi)` planes (each lane independent). +/// +/// Three parallel layers — two XORs, two XORs, two AND-NOTs — so it maps onto x86 `andn` +/// and has very short dependency chains. Verified exhaustively against the 9 valid input +/// pairs (the `(hi, lo) = (1, 1)` encoding never occurs for reduced inputs). +/// +/// This circuit was contributed by Carl McTague. +#[inline(always)] +fn f3_add_planes(a_lo: Limb, a_hi: Limb, b_lo: Limb, b_hi: Limb) -> (Limb, Limb) { + let t_hi = a_hi ^ b_hi; + let t_lo = a_lo ^ b_lo; + let u_hi = b_hi ^ t_lo; + let u_lo = b_lo ^ t_hi; + let r_hi = u_lo & !t_lo; + let r_lo = u_hi & !t_hi; + (r_lo, r_hi) +} + +/// Negation in F3 swaps 1 <-> 2 (and fixes 0), i.e. swaps the two planes. +#[inline(always)] +fn f3_addend(sg: &[Limb], c: u32) -> (Limb, Limb) { + // c is 1 or 2 here; c == 2 means add (-other), i.e. negate by swapping planes. + if c == 1 { + (sg[0], sg[1]) + } else { + (sg[1], sg[0]) + } +} + +fn f3_add_groups(dst: &mut [Limb], src: &[Limb], c: u32) { + for (dg, sg) in dst + .as_chunks_mut::<2>() + .0 + .iter_mut() + .zip(src.as_chunks::<2>().0) + { + let (b_lo, b_hi) = f3_addend(sg, c); + let (r_lo, r_hi) = f3_add_planes(dg[0], dg[1], b_lo, b_hi); + dg[0] = r_lo; + dg[1] = r_hi; + } +} + +fn f3_add_group_masked(dst: &mut [Limb], src: &[Limb], c: u32, lane_mask: Limb) { + let (b_lo, b_hi) = f3_addend(src, c); + // Masking the addend to the in-range lanes makes the circuit a no-op (adds 0) elsewhere. + let (r_lo, r_hi) = f3_add_planes(dst[0], dst[1], b_lo & lane_mask, b_hi & lane_mask); + dst[0] = r_lo; + dst[1] = r_hi; +} + +fn f3_scale_groups(dst: &mut [Limb], c: u32) { + // c == 2 is negation (plane swap); c == 1 is a no-op; c == 0 is handled by the caller. + if c == 2 { + for dg in dst.as_chunks_mut::<2>().0 { + dg.swap(0, 1); + } + } +} + +// --------------------------------------------------------------------------------------- +// F5 specialization (k = 3). Planes are bits 0,1,2 of the value `v in {0,..,4}`. Addition is a +// flat 17-gate boolean circuit; scalar multiplication is an "indicator" circuit — one-hot lane +// masks `is_v` for each operand value, recombined into the result with no carry/borrow chain. +// Both keep the wide instruction-level parallelism the sequential generic circuit loses. +// --------------------------------------------------------------------------------------- + +/// One-hot lane masks: `out[v]` has the bits of the lanes whose value is `v` (for `v in 0..5`). +#[inline(always)] +fn f5_indicators(p0: Limb, p1: Limb, p2: Limb) -> [Limb; 5] { + let n0 = !p0; + let n1 = !p1; + let n2 = !p2; + [ + n0 & n1 & n2, // 0 = 000 + p0 & n1 & n2, // 1 = 001 + n0 & p1 & n2, // 2 = 010 + p0 & p1 & n2, // 3 = 011 + n0 & n1 & p2, // 4 = 100 + ] +} + +/// Reassemble the three planes from per-value selection masks (`sel[v]` selects value `v`). +#[inline(always)] +fn f5_compose(sel: [Limb; 5]) -> (Limb, Limb, Limb) { + // bit 0 set for values {1,3}; bit 1 for {2,3}; bit 2 for {4}. + (sel[1] | sel[3], sel[2] | sel[3], sel[4]) +} + +/// `c * v mod 5` on the three planes (`c in 1..5`). +#[inline(always)] +fn f5_mul_planes(p0: Limb, p1: Limb, p2: Limb, c: u32) -> (Limb, Limb, Limb) { + let ind = f5_indicators(p0, p1, p2); + let mut sel = [0 as Limb; 5]; + for v in 0..5u32 { + sel[((c * v) % 5) as usize] |= ind[v as usize]; + } + f5_compose(sel) +} + +/// `(a + b) mod 5` on the three planes, as a flat 17-gate circuit (planes are bits 0,1,2 of the +/// value). Four gate layers with an all-`andn` output layer, so the dependency chains stay short. +/// Verified exhaustively against `(a + b) % 5` for all reduced inputs. `andnot(x, y) = x & !y` +/// maps onto x86 `andn`. +/// +/// This circuit was contributed by Carl McTague. +#[inline(always)] +fn f5_add_planes(a0: Limb, a1: Limb, a2: Limb, b0: Limb, b1: Limb, b2: Limb) -> (Limb, Limb, Limb) { + // Layer 1 + let g0 = a0 & b0; + let g1 = a0 | b0; + let g2 = a1 ^ b1; + let g3 = a1 | b1; + let g4 = a2 | b2; + let g5 = a2 ^ b2; + // Layer 2 + let g6 = g0 ^ g2; + let g7 = g3 | g5; + let g8 = g4 ^ g1; + let g9 = g1 & !g2; + let g10 = g4 & !g1; + // Layer 3 + let g11 = g2 ^ g7; + let g12 = g6 ^ g10; + let g13 = g6 ^ g7; + // Layer 4 + let g14 = g12 & !g11; + let g15 = g13 & !g9; + let g16 = g8 & !g13; + // m0 = g16, m1 = g14, m2 = g15 + (g16, g14, g15) +} + +fn f5_add_groups(dst: &mut [Limb], src: &[Limb], c: u32) { + for (dg, sg) in dst + .as_chunks_mut::<3>() + .0 + .iter_mut() + .zip(src.as_chunks::<3>().0) + { + let (b0, b1, b2) = if c == 1 { + (sg[0], sg[1], sg[2]) + } else { + f5_mul_planes(sg[0], sg[1], sg[2], c) + }; + let (r0, r1, r2) = f5_add_planes(dg[0], dg[1], dg[2], b0, b1, b2); + dg[0] = r0; + dg[1] = r1; + dg[2] = r2; + } +} + +fn f5_add_group_masked(dst: &mut [Limb], src: &[Limb], c: u32, lane_mask: Limb) { + let (mut b0, mut b1, mut b2) = if c == 1 { + (src[0], src[1], src[2]) + } else { + f5_mul_planes(src[0], src[1], src[2], c) + }; + // Zeroing the addend outside the mask leaves those lanes unchanged (adds 0). + b0 &= lane_mask; + b1 &= lane_mask; + b2 &= lane_mask; + let (r0, r1, r2) = f5_add_planes(dst[0], dst[1], dst[2], b0, b1, b2); + dst[0] = r0; + dst[1] = r1; + dst[2] = r2; +} + +fn f5_scale_groups(dst: &mut [Limb], c: u32) { + for dg in dst.as_chunks_mut::<3>().0 { + let (r0, r1, r2) = f5_mul_planes(dg[0], dg[1], dg[2], c); + dg[0] = r0; + dg[1] = r1; + dg[2] = r2; + } +} + +// --------------------------------------------------------------------------------------- +// Heap-scratch fallback for `k > MAX_DISPATCH_K` (very large primes). +// --------------------------------------------------------------------------------------- + +fn cond_sub_into(dst: &mut [Limb], s: &[Limb], masks: &[Limb], d: &mut [Limb]) { + let k = dst.len(); + let mut borrow: Limb = 0; + for j in 0..=k { + let sj = s[j]; + let pj = masks[j]; + let sxp = sj ^ pj; + d[j] = sxp ^ borrow; + borrow = (!sj & pj) | (borrow & !sxp); + } + let ge = !borrow; + for j in 0..k { + dst[j] = (d[j] & ge) | (s[j] & !ge); + } +} + +fn add_mod_into(dst: &mut [Limb], b: &[Limb], masks: &[Limb], s: &mut [Limb], d: &mut [Limb]) { + let k = dst.len(); + let mut carry: Limb = 0; + for j in 0..k { + let aj = dst[j]; + let bj = b[j]; + let axb = aj ^ bj; + s[j] = axb ^ carry; + carry = (aj & bj) | (carry & axb); + } + s[k] = carry; + cond_sub_into(dst, s, masks, d); +} + +fn double_mod_into(dst: &mut [Limb], masks: &[Limb], s: &mut [Limb], d: &mut [Limb]) { + let k = dst.len(); + s[0] = 0; + s[1..=k].copy_from_slice(&dst[..k]); + cond_sub_into(dst, s, masks, d); +} + +fn scalar_mul_into( + acc: &mut [Limb], + b: &[Limb], + c: u32, + masks: &[Limb], + temp: &mut [Limb], + s: &mut [Limb], + d: &mut [Limb], +) { + temp.copy_from_slice(b); + acc.fill(0); + let mut cc = c; + loop { + if cc & 1 == 1 { + add_mod_into(acc, temp, masks, s, d); + } + cc >>= 1; + if cc == 0 { + break; + } + double_mod_into(temp, masks, s, d); + } +} + +fn add_groups_dyn(k: usize, dst: &mut [Limb], src: &[Limb], c: u32, masks: &[Limb]) { + let mut s = vec![0; k + 1]; + let mut d = vec![0; k + 1]; + let mut acc = vec![0; k]; + let mut temp = vec![0; k]; + for (dg, sg) in dst.chunks_exact_mut(k).zip(src.chunks_exact(k)) { + if c == 1 { + add_mod_into(dg, sg, masks, &mut s, &mut d); + } else { + scalar_mul_into(&mut acc, sg, c, masks, &mut temp, &mut s, &mut d); + add_mod_into(dg, &acc, masks, &mut s, &mut d); + } + } +} + +fn scale_groups_dyn(k: usize, dst: &mut [Limb], c: u32, masks: &[Limb]) { + if c == 1 { + return; + } + if c == 0 { + dst.fill(0); + return; + } + let mut s = vec![0; k + 1]; + let mut d = vec![0; k + 1]; + let mut acc = vec![0; k]; + let mut temp = vec![0; k]; + for dg in dst.chunks_exact_mut(k) { + scalar_mul_into(&mut acc, dg, c, masks, &mut temp, &mut s, &mut d); + dg.copy_from_slice(&acc); + } +} + +const _: () = assert!(MAX_DISPATCH_K <= BITS_PER_LIMB); diff --git a/ext/crates/fp/src/field/field_internal.rs b/ext/crates/fp/src/field/field_internal.rs index 594bdb8d07..8065a5f842 100644 --- a/ext/crates/fp/src/field/field_internal.rs +++ b/ext/crates/fp/src/field/field_internal.rs @@ -133,6 +133,124 @@ pub trait FieldInternal: } } + // # Group layout (bit-sliced storage) + // + // Storage is organized into *groups*: a group holds [`entries_per_group`] = 64 consecutive + // entries and occupies [`limbs_per_group`] = `k` consecutive limbs, the *bit-planes*. Plane + // `j` of a group holds bit `j` of all 64 entries, so entry `i` lives at bit `i` of each of + // the `k` planes. Every field uses this layout, with `k = ceil(log2 q)` (the bits needed to + // store an encoded value in `0..q`). For `q = 2` this is `k = 1`, which coincides exactly + // with the old packed layout — so `F_2` (and its SIMD / matrix machinery) is byte-identical + // and unaffected. Entry access goes through [`gather`]/[`scatter`]; the sizing helpers + // [`number`]/[`range`] are expressed in terms of groups. + // + // [`entries_per_group`]: FieldInternal::entries_per_group + // [`limbs_per_group`]: FieldInternal::limbs_per_group + // [`gather`]: FieldInternal::gather + // [`scatter`]: FieldInternal::scatter + // [`number`]: FieldInternal::number + // [`range`]: FieldInternal::range + + /// The number of entries stored in a single group: one per bit of a [`Limb`]. + fn entries_per_group(self) -> usize { + BITS_PER_LIMB + } + + /// The number of bit-planes per group, `k = ceil(log2 q)`. Each field defines this; `q = 2` + /// gives `k = 1` (packed-compatible). + fn limbs_per_group(self) -> usize; + + /// The index of the group containing entry `idx`. + fn group_of(self, idx: usize) -> usize { + idx / self.entries_per_group() + } + + /// The position of entry `idx` within its group, in `0..entries_per_group()`. + fn lane_of(self, idx: usize) -> usize { + idx % self.entries_per_group() + } + + /// Read entry `lane` (in `0..entries_per_group()`) out of a single group's `k` planes (a + /// slice of length [`limbs_per_group`](FieldInternal::limbs_per_group)) by reassembling its + /// bit from each plane. + fn gather(self, group: &[Limb], lane: usize) -> FieldElement { + let mut value: Limb = 0; + for (j, plane) in group.iter().enumerate() { + value |= ((plane >> lane) & 1) << j; + } + self.decode(value) + } + + /// Write `value` into entry `lane` of a single group's `k` planes, dispersing the encoded + /// value's bits one per plane. Assumes the stored value fits in `k` bits. + fn scatter(self, group: &mut [Limb], lane: usize, value: FieldElement) { + let encoded = self.encode(value); + let lane_mask: Limb = 1 << lane; + for (j, plane) in group.iter_mut().enumerate() { + let bit = (encoded >> j) & 1; + *plane = (*plane & !lane_mask) | (bit << lane); + } + } + + /// Whether this field uses a genuinely multi-plane layout (`k > 1`). Only `F_2` has `k = 1`, + /// where the bit-sliced layout coincides with the packed one and the `F_2`-specific fast + /// paths (`offset`, `limb_masks`, SIMD, m4ri) apply. + fn is_bitsliced(self) -> bool { + self.limbs_per_group() > 1 + } + + /// `dst += coeff * src` (mod p) over a span of whole groups (`dst` and `src` have equal, + /// group-aligned length). Both are assumed reduced; the result is reduced. + /// + /// Default: element-wise over lanes via [`Self::gather`]/[`Self::scatter`] and the field's own + /// arithmetic — correct for any bit-sliced field (used by [`SmallFq`](super::SmallFq)). The + /// prime fields [`Fp`](super::Fp) override this with a branch-free plane circuit. + fn add_groups(self, dst: &mut [Limb], src: &[Limb], coeff: FieldElement) { + let lpg = self.limbs_per_group(); + let epg = self.entries_per_group(); + for (dgroup, sgroup) in dst.chunks_exact_mut(lpg).zip(src.chunks_exact(lpg)) { + for lane in 0..epg { + let a = self.gather(dgroup, lane); + let b = self.gather(sgroup, lane); + let result = self.add(a, self.mul(coeff.clone(), b)); + self.scatter(dgroup, lane, result); + } + } + } + + /// `dst *= coeff` (mod p) over a span of whole groups. Default: element-wise; overridden by + /// [`Fp`](super::Fp). + fn scale_groups(self, dst: &mut [Limb], coeff: FieldElement) { + let lpg = self.limbs_per_group(); + let epg = self.entries_per_group(); + for dgroup in dst.chunks_exact_mut(lpg) { + for lane in 0..epg { + let a = self.gather(dgroup, lane); + self.scatter(dgroup, lane, self.mul(a, coeff.clone())); + } + } + } + + /// `dst += coeff * src` (mod p) for a single group (each `limbs_per_group()` limbs), + /// restricted to the lanes set in `lane_mask`; other lanes are unchanged. Used for the + /// partial boundary groups of a slice add. Default: element-wise; overridden by + /// [`Fp`](super::Fp) with a masked plane circuit. + fn add_group_masked( + self, + dst: &mut [Limb], + src: &[Limb], + coeff: FieldElement, + lane_mask: Limb, + ) { + for lane in 0..self.entries_per_group() { + if (lane_mask >> lane) & 1 == 1 { + let a = self.gather(dst, lane); + let b = self.gather(src, lane); + self.scatter(dst, lane, self.add(a, self.mul(coeff.clone(), b))); + } + } + } + /// Check whether or not a limb is reduced. This may potentially not be faster than calling /// [`reduce`](FieldInternal::reduce) directly. fn is_reduced(self, limb: Limb) -> bool { @@ -166,17 +284,20 @@ pub trait FieldInternal: /// Return the number of limbs required to hold `dim` entries. fn number(self, dim: usize) -> usize { - if dim == 0 { - 0 - } else { - self.limb_bit_index_pair(dim - 1).limb + 1 - } + // Whole groups needed to hold `dim` entries, times the limbs in each group. + self.limbs_per_group() * dim.div_ceil(self.entries_per_group()) } - /// Return the `Range` starting at the index of the limb containing the `start`th entry, and - /// ending at the index of the limb containing the `end`th entry (including the latter). + /// Return the `Range` of limbs spanning entries `start..end`: from the first limb of + /// the group containing `start` to the last limb of the group containing `end - 1`. fn range(self, start: usize, end: usize) -> Range { - let min = self.limb_bit_index_pair(start).limb; + debug_assert!(start <= end); + let min = self.group_of(start) * self.limbs_per_group(); + if start == end { + // An empty entry range maps to an empty limb range; otherwise callers that guard + // on `limb_range.is_empty()` would touch the (unrelated) containing group. + return min..min; + } let max = self.number(end); min..max } diff --git a/ext/crates/fp/src/field/fp.rs b/ext/crates/fp/src/field/fp.rs index 18a9d0f9d5..00e6829ca4 100644 --- a/ext/crates/fp/src/field/fp.rs +++ b/ext/crates/fp/src/field/fp.rs @@ -142,6 +142,54 @@ impl FieldInternal for Fp

{ _ => self.pack(self.unpack(limb)), } } + + // # Bit-sliced layout + // + // Prime-field vectors use the bit-sliced group layout (see [`FieldInternal`]). The packed + // limb helpers above are not used to lay out `FqVector>` storage, but are retained + // because `decode`/`encode`/`reduce` are still called when constructing elements. The + // uniform `gather`/`scatter` defaults handle entry access; `Fp` only overrides the bulk + // kernels with a branch-free plane circuit. + + fn limbs_per_group(self) -> usize { + crate::field::bitslice::planes(self.characteristic().as_u32()) + } + + fn add_groups(self, dst: &mut [Limb], src: &[Limb], coeff: FieldElement) { + crate::field::bitslice::add_groups( + self.characteristic().as_u32(), + self.limbs_per_group(), + dst, + src, + self.encode(coeff) as u32, + ); + } + + fn scale_groups(self, dst: &mut [Limb], coeff: FieldElement) { + crate::field::bitslice::scale_groups( + self.characteristic().as_u32(), + self.limbs_per_group(), + dst, + self.encode(coeff) as u32, + ); + } + + fn add_group_masked( + self, + dst: &mut [Limb], + src: &[Limb], + coeff: FieldElement, + lane_mask: Limb, + ) { + crate::field::bitslice::add_group_masked( + self.characteristic().as_u32(), + self.limbs_per_group(), + dst, + src, + self.encode(coeff) as u32, + lane_mask, + ); + } } #[cfg(feature = "proptest")] diff --git a/ext/crates/fp/src/field/mod.rs b/ext/crates/fp/src/field/mod.rs index 714a9e735d..7348ef0475 100644 --- a/ext/crates/fp/src/field/mod.rs +++ b/ext/crates/fp/src/field/mod.rs @@ -4,6 +4,7 @@ use crate::prime::Prime; pub mod element; pub(crate) mod field_internal; +pub(crate) mod bitslice; pub mod fp; pub mod smallfq; diff --git a/ext/crates/fp/src/field/smallfq.rs b/ext/crates/fp/src/field/smallfq.rs index b1aa80115c..96f91067a6 100644 --- a/ext/crates/fp/src/field/smallfq.rs +++ b/ext/crates/fp/src/field/smallfq.rs @@ -282,6 +282,13 @@ impl FieldInternal for SmallFq

{ BITS_PER_LIMB - (self.q() - 1).leading_zeros() as usize + 1 } + fn limbs_per_group(self) -> usize { + // Bit-sliced layout: one plane per bit of the encoded value. `encode` maps a^n to the + // odd number `2n + 1` (and zero to `0`), which occupies exactly `bit_length()` bits. + // The default element-wise `add_groups`/`scale_groups` use Zech-log field arithmetic. + self.bit_length() + } + fn fma_limb(self, limb_a: Limb, limb_b: Limb, coeff: FieldElement) -> Limb { let bit_length = self.bit_length(); let mut result: Limb = 0; diff --git a/ext/crates/fp/src/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index 8c6dd93b38..d80efa5755 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -322,8 +322,26 @@ impl Matrix { return Self::new(p, 0, 0); } let columns = input[0].len(); + assert!( + input.iter().all(|row| row.len() == columns), + "all rows must have the same length" + ); let stride = fp.number(columns); let physical_rows = get_physical_rows(p, rows); + + if fp.is_bitsliced() { + // The bit-sliced layout interleaves an entry's bits across planes, so build each + // row through the (dispatching) row slice rather than by packing contiguous chunks. + let mut matrix = Self::new(p, rows, columns); + for (i, row) in input.iter().enumerate() { + let mut target = matrix.row_mut(i); + for (j, &x) in row.iter().enumerate() { + target.set_entry(j, x); + } + } + return matrix; + } + let mut data = AVec::with_capacity(0, physical_rows * stride); for row in input { for chunk in row.chunks(fp.entries_per_limb()) { @@ -352,6 +370,14 @@ impl Matrix { /// assert_eq!(Matrix::from_vec(TWO, &matrix_vec).to_vec(), matrix_vec); /// ``` pub fn to_vec(&self) -> Vec> { + if self.fp.is_bitsliced() { + return (0..self.rows()) + .map(|i| { + let row = self.row(i); + (0..self.columns()).map(|j| row.entry(j)).collect() + }) + .collect(); + } self.data .iter() .chunks(self.stride) diff --git a/ext/crates/fp/src/vector/fp_wrapper/mod.rs b/ext/crates/fp/src/vector/fp_wrapper/mod.rs index 112760002a..1b7810fe55 100644 --- a/ext/crates/fp/src/vector/fp_wrapper/mod.rs +++ b/ext/crates/fp/src/vector/fp_wrapper/mod.rs @@ -112,14 +112,13 @@ impl FpVector { v } - // Convenient for some matrix methods - pub(crate) fn num_limbs(p: ValidPrime, len: usize) -> usize { - Fp::new(p).number(len) - } - - // Convenient for some matrix methods + // Round `len` up to a whole number of groups, so that an augmented-matrix segment of this + // length ends on a group boundary and the next segment starts on one. A group spans 64 + // entries (across several limbs in the bit-sliced layout), and segments must align to those + // 64-entry boundaries or a single group would straddle two segments. pub(crate) fn padded_len(p: ValidPrime, len: usize) -> usize { - Self::num_limbs(p, len) * Fp::new(p).entries_per_limb() + let entries_per_group = Fp::new(p).entries_per_group(); + len.div_ceil(entries_per_group) * entries_per_group } } diff --git a/ext/crates/fp/src/vector/impl_fqslice.rs b/ext/crates/fp/src/vector/impl_fqslice.rs index e2f243f0db..19a547ea08 100644 --- a/ext/crates/fp/src/vector/impl_fqslice.rs +++ b/ext/crates/fp/src/vector/impl_fqslice.rs @@ -33,12 +33,11 @@ impl<'a, F: Field> FqSlice<'a, F> { index, self.len() ); - let bit_mask = self.fq().bitmask(); - let limb_index = self.fq().limb_bit_index_pair(index + self.start()); - let mut result = self.limbs()[limb_index.limb]; - result >>= limb_index.bit_index; - result &= bit_mask; - self.fq().decode(result) + let fq = self.fq(); + let idx = index + self.start(); + let lpg = fq.limbs_per_group(); + let base = fq.group_of(idx) * lpg; + fq.gather(&self.limbs()[base..base + lpg], fq.lane_of(idx)) } /// TODO: implement prime 2 version @@ -55,6 +54,9 @@ impl<'a, F: Field> FqSlice<'a, F> { } pub fn is_zero(&self) -> bool { + if self.fq().is_bitsliced() { + return self.first_nonzero().is_none(); + } let limb_range = self.limb_range(); if limb_range.is_empty() { return true; @@ -90,7 +92,7 @@ impl<'a, F: Field> FqSlice<'a, F> { #[must_use] pub fn to_owned(self) -> FqVector { let mut new = FqVector::new(self.fq(), self.len()); - if self.start().is_multiple_of(self.fq().entries_per_limb()) { + if !self.fq().is_bitsliced() && self.start().is_multiple_of(self.fq().entries_per_limb()) { let limb_range = self.limb_range(); new.limbs_mut()[0..limb_range.len()].copy_from_slice(&self.limbs()[limb_range]); if !new.limbs().is_empty() { diff --git a/ext/crates/fp/src/vector/impl_fqslicemut.rs b/ext/crates/fp/src/vector/impl_fqslicemut.rs index d65d302471..952f7284ed 100644 --- a/ext/crates/fp/src/vector/impl_fqslicemut.rs +++ b/ext/crates/fp/src/vector/impl_fqslicemut.rs @@ -1,7 +1,3 @@ -use std::cmp::Ordering; - -use itertools::Itertools; - use super::inner::{FqSlice, FqSliceMut, FqVector}; use crate::{ constants, @@ -30,12 +26,12 @@ impl<'a, F: Field> FqSliceMut<'a, F> { pub fn set_entry(&mut self, index: usize, value: FieldElement) { assert_eq!(self.fq(), value.field()); assert!(index < self.as_slice().len()); - let bit_mask = self.fq().bitmask(); - let limb_index = self.fq().limb_bit_index_pair(index + self.start()); - let mut result = self.limbs()[limb_index.limb]; - result &= !(bit_mask << limb_index.bit_index); - result |= self.fq().encode(value) << limb_index.bit_index; - self.limbs_mut()[limb_index.limb] = result; + let fq = self.fq(); + let idx = index + self.start(); + let lpg = fq.limbs_per_group(); + let base = fq.group_of(idx) * lpg; + let lane = fq.lane_of(idx); + fq.scatter(&mut self.limbs_mut()[base..base + lpg], lane, value); } fn reduce_limbs(&mut self) { @@ -60,6 +56,20 @@ impl<'a, F: Field> FqSliceMut<'a, F> { return; } + if fq.is_bitsliced() { + // The packed bit-offset masking does not apply to the bit-sliced layout; scale + // each in-range entry through the layout-aware gather/scatter. + if c == fq.zero() { + self.set_to_zero(); + return; + } + for i in 0..self.as_slice().len() { + let x = self.as_slice().entry(i) * c.clone(); + self.set_entry(i, x); + } + return; + } + let limb_range = self.as_slice().limb_range(); if limb_range.is_empty() { return; @@ -85,6 +95,13 @@ impl<'a, F: Field> FqSliceMut<'a, F> { } pub fn set_to_zero(&mut self) { + if self.fq().is_bitsliced() { + let zero = self.fq().zero(); + for i in 0..self.as_slice().len() { + self.set_entry(i, zero.clone()); + } + return; + } let limb_range = self.as_slice().limb_range(); if limb_range.is_empty() { return; @@ -102,25 +119,162 @@ impl<'a, F: Field> FqSliceMut<'a, F> { pub fn add(&mut self, other: FqSlice<'_, F>, c: FieldElement) { assert_eq!(self.fq(), c.field()); assert_eq!(self.fq(), other.fq()); + assert_eq!(self.as_slice().len(), other.len()); if self.as_slice().is_empty() { return; } - if self.fq().q() == 2 { - if c != self.fq().zero() { - match self.as_slice().offset().cmp(&other.offset()) { - Ordering::Equal => self.add_shift_none(other, self.fq().one()), - Ordering::Less => self.add_shift_left(other, self.fq().one()), - Ordering::Greater => self.add_shift_right(other, self.fq().one()), - }; + // Every field uses the bit-sliced layout (`F_2` is just the `k = 1` case), so a single + // code path handles them all. + self.add_bitsliced(other, c); + } + + /// Add `c * other` to `self` in the bit-sliced layout. Interior full groups are added with + /// the fast plane kernel ([`add_groups`](crate::field::field_internal)); the leading/trailing + /// partial groups go through the masked plane circuit + /// ([`add_group_masked`](crate::field::field_internal::FieldInternal::add_group_masked)). When + /// the two slices have different lane offsets within their groups, the planes are realigned + /// first via [`add_bitsliced_shifted`](Self::add_bitsliced_shifted). + /// + /// [`add_groups`]: crate::field::field_internal::FieldInternal::add_groups + fn add_bitsliced(&mut self, other: FqSlice<'_, F>, c: FieldElement) { + let fq = self.fq(); + if c == fq.zero() { + return; + } + let epg = fq.entries_per_group(); + let s_start = self.start(); + let o_start = other.start(); + let len = self.as_slice().len(); + if len == 0 { + return; + } + + // The fast plane kernel needs the two slices to share a lane offset within their groups + // (so group `g` of one lines up with group `g` of the other). This holds for whole + // vectors and for matrix-row adds that start at the same pivot column. When the offsets + // differ, the planes must be realigned first — see `add_bitsliced_shifted`. + if s_start % epg != o_start % epg { + self.add_bitsliced_shifted(other, c); + return; + } + + let k = fq.limbs_per_group(); + let s_end = s_start + len; + let first_g = s_start / epg; + let last_g = (s_end - 1) / epg; + // Group `g` of `self` lines up with group `g - first_g + o_first_g` of `other`. + let o_first_g = o_start / epg; + let group_limbs = |self_g: usize| { + let s = self_g * k; + let o = (o_first_g + (self_g - first_g)) * k; + (s, o) + }; + // Lane mask selecting bits `[lo, hi)`. + let lane_mask = |lo: usize, hi: usize| -> Limb { + let high: Limb = if hi >= epg { !0 } else { (1 << hi) - 1 }; + let low: Limb = (1 << lo) - 1; + high & !low + }; + + if first_g == last_g { + let (s, o) = group_limbs(first_g); + let mask = lane_mask(s_start - first_g * epg, s_end - first_g * epg); + let src = &other.limbs()[o..o + k]; + fq.add_group_masked(&mut self.limbs_mut()[s..s + k], src, c, mask); + return; + } + + // Leading partial group: lanes [s_start mod 64, 64). + { + let (s, o) = group_limbs(first_g); + let mask = lane_mask(s_start - first_g * epg, epg); + let src = &other.limbs()[o..o + k]; + fq.add_group_masked(&mut self.limbs_mut()[s..s + k], src, c.clone(), mask); + } + // Interior full groups, via the plane kernel in one contiguous call. + if last_g > first_g + 1 { + let (s, o) = group_limbs(first_g + 1); + let nlimbs = (last_g - first_g - 1) * k; + let src = &other.limbs()[o..o + nlimbs]; + fq.add_groups(&mut self.limbs_mut()[s..s + nlimbs], src, c.clone()); + } + // Trailing partial group: lanes [0, s_end mod 64). + { + let (s, o) = group_limbs(last_g); + let mask = lane_mask(0, s_end - last_g * epg); + let src = &other.limbs()[o..o + k]; + fq.add_group_masked(&mut self.limbs_mut()[s..s + k], src, c, mask); + } + } + + /// Add `c * other` to `self` when the two slices have different lane offsets within their + /// groups, so the planes don't line up. A bit-sliced vector is `k` independent single-bit + /// planes; realigning it is a per-plane bit shift (this is exactly what the old F_2-only + /// `add_shift_*` did, generalized from one plane to `k`). For each target group we build the + /// `k` source planes shifted into alignment, then add them in with the masked group circuit. + fn add_bitsliced_shifted(&mut self, other: FqSlice<'_, F>, c: FieldElement) { + let fq = self.fq(); + let k = fq.limbs_per_group(); + let epg = fq.entries_per_group(); + let ts = self.start(); + let len = self.as_slice().len(); + // Source lane = target lane + shift (`other`'s entry i sits `shift` lanes from `self`'s). + let shift = other.start() as isize - ts as isize; + let src_limbs = other.limbs(); + + // Plane `j` of (absolute) group `g` of the source, or 0 if out of range. The whole-limb + // reads can stray outside the valid lane range, but those bits are masked off below. + let plane_limb = |g: isize, j: usize| -> Limb { + if g < 0 { + return 0; } - } else { - match self.as_slice().offset().cmp(&other.offset()) { - Ordering::Equal => self.add_shift_none(other, c), - Ordering::Less => self.add_shift_left(other, c), - Ordering::Greater => self.add_shift_right(other, c), + let idx = g as usize * k + j; + if idx < src_limbs.len() { + src_limbs[idx] + } else { + 0 + } + }; + let lane_mask = |lo: usize, hi: usize| -> Limb { + let high: Limb = if hi >= epg { !0 } else { (1 << hi) - 1 }; + let low: Limb = (1 << lo) - 1; + high & !low + }; + + debug_assert!(k <= epg); + let mut shifted = [0 as Limb; constants::BITS_PER_LIMB]; + + let first_g = ts / epg; + let last_g = (ts + len - 1) / epg; + let epg_i = epg as isize; + for g in first_g..=last_g { + let lo = if g == first_g { ts - g * epg } else { 0 }; + let hi = if g == last_g { + (ts + len) - g * epg + } else { + epg }; + let mask = lane_mask(lo, hi); + + // Source bit `b` of target group `g` lives at absolute source lane `g*epg + b + shift`. + let src_base = g as isize * epg_i + shift; + let sg = src_base.div_euclid(epg_i); + let bs = src_base.rem_euclid(epg_i) as u32; + for (j, s) in shifted[..k].iter_mut().enumerate() { + *s = if bs == 0 { + plane_limb(sg, j) + } else { + (plane_limb(sg, j) >> bs) | (plane_limb(sg + 1, j) << (epg as u32 - bs)) + }; + } + fq.add_group_masked( + &mut self.limbs_mut()[g * k..g * k + k], + &shifted[..k], + c.clone(), + mask, + ); } } @@ -153,7 +307,7 @@ impl<'a, F: Field> FqSliceMut<'a, F> { /// TODO: improve efficiency pub fn assign(&mut self, other: FqSlice<'_, F>) { assert_eq!(self.fq(), other.fq()); - if self.as_slice().offset() != other.offset() { + if self.fq().is_bitsliced() || self.as_slice().offset() != other.offset() { self.set_to_zero(); self.add(other, self.fq().one()); return; @@ -188,6 +342,23 @@ impl<'a, F: Field> FqSliceMut<'a, F> { if shift == 0 { return; } + if self.fq().is_bitsliced() { + // The packed limb-move trick assumes an entry is a contiguous bitfield, which the + // bit-sliced layout breaks. Move entries down one at a time via gather/scatter: + // reading `i + shift` strictly ahead of writing `i` keeps it correct in place. + let len = self.as_slice().len(); + if shift >= len { + *self.end_mut() = self.start(); + return; + } + let new_len = len - shift; + for i in 0..new_len { + let v = self.as_slice().entry(i + shift); + self.set_entry(i, v); + } + *self.end_mut() -= shift; + return; + } if self.start() == 0 && shift.is_multiple_of(self.fq().entries_per_limb()) { let limb_shift = shift / self.fq().entries_per_limb(); *self.end_mut() -= shift; @@ -200,299 +371,6 @@ impl<'a, F: Field> FqSliceMut<'a, F> { } } - /// Adds `c` * `other` to `self`. `other` must have the same length, offset, and prime as self. - pub fn add_shift_none(&mut self, other: FqSlice<'_, F>, c: FieldElement) { - assert_eq!(self.fq(), c.field()); - assert_eq!(self.fq(), other.fq()); - let fq = self.fq(); - - let target_range = self.as_slice().limb_range(); - let source_range = other.limb_range(); - - let (min_mask, max_mask) = other.limb_masks(); - - self.limbs_mut()[target_range.start] = fq.fma_limb( - self.limbs()[target_range.start], - other.limbs()[source_range.start] & min_mask, - c.clone(), - ); - self.limbs_mut()[target_range.start] = fq.reduce(self.limbs()[target_range.start]); - - let target_inner_range = self.as_slice().limb_range_inner(); - let source_inner_range = other.limb_range_inner(); - if !source_inner_range.is_empty() { - for (left, right) in self.limbs_mut()[target_inner_range] - .iter_mut() - .zip_eq(&other.limbs()[source_inner_range]) - { - *left = fq.fma_limb(*left, *right, c.clone()); - *left = fq.reduce(*left); - } - } - if source_range.len() > 1 { - // The first and last limbs are distinct, so we process the last. - self.limbs_mut()[target_range.end - 1] = fq.fma_limb( - self.limbs()[target_range.end - 1], - other.limbs()[source_range.end - 1] & max_mask, - c, - ); - self.limbs_mut()[target_range.end - 1] = fq.reduce(self.limbs()[target_range.end - 1]); - } - } - - fn add_shift_left(&mut self, other: FqSlice<'_, F>, c: FieldElement) { - struct AddShiftLeftData { - offset_shift: usize, - tail_shift: usize, - zero_bits: usize, - min_source_limb: usize, - min_target_limb: usize, - number_of_source_limbs: usize, - number_of_target_limbs: usize, - min_mask: Limb, - max_mask: Limb, - } - - impl AddShiftLeftData { - fn new(fq: F, target: FqSlice<'_, F>, source: FqSlice<'_, F>) -> Self { - debug_assert!(target.prime() == source.prime()); - debug_assert!(target.offset() <= source.offset()); - debug_assert!( - target.len() == source.len(), - "self.dim {} not equal to other.dim {}", - target.len(), - source.len() - ); - let offset_shift = source.offset() - target.offset(); - let bit_length = fq.bit_length(); - let entries_per_limb = fq.entries_per_limb(); - let usable_bits_per_limb = bit_length * entries_per_limb; - let tail_shift = usable_bits_per_limb - offset_shift; - let zero_bits = constants::BITS_PER_LIMB - usable_bits_per_limb; - let source_range = source.limb_range(); - let target_range = target.limb_range(); - let min_source_limb = source_range.start; - let min_target_limb = target_range.start; - let number_of_source_limbs = source_range.len(); - let number_of_target_limbs = target_range.len(); - let (min_mask, max_mask) = source.limb_masks(); - - Self { - offset_shift, - tail_shift, - zero_bits, - min_source_limb, - min_target_limb, - number_of_source_limbs, - number_of_target_limbs, - min_mask, - max_mask, - } - } - - fn mask_first_limb(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - (other.limbs()[i] & self.min_mask) >> self.offset_shift - } - - fn mask_middle_limb_a(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - other.limbs()[i] >> self.offset_shift - } - - fn mask_middle_limb_b(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - (other.limbs()[i] << (self.tail_shift + self.zero_bits)) >> self.zero_bits - } - - fn mask_last_limb_a(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - let source_limb_masked = other.limbs()[i] & self.max_mask; - source_limb_masked << self.tail_shift - } - - fn mask_last_limb_b(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - let source_limb_masked = other.limbs()[i] & self.max_mask; - source_limb_masked >> self.offset_shift - } - } - - let dat = AddShiftLeftData::new(self.fq(), self.as_slice(), other); - let mut i = 0; - { - self.limbs_mut()[i + dat.min_target_limb] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb], - dat.mask_first_limb(other, i + dat.min_source_limb), - c.clone(), - ); - } - for i in 1..dat.number_of_source_limbs - 1 { - self.limbs_mut()[i + dat.min_target_limb] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb], - dat.mask_middle_limb_a(other, i + dat.min_source_limb), - c.clone(), - ); - self.limbs_mut()[i + dat.min_target_limb - 1] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb - 1], - dat.mask_middle_limb_b(other, i + dat.min_source_limb), - c.clone(), - ); - self.limbs_mut()[i + dat.min_target_limb - 1] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb - 1]); - } - i = dat.number_of_source_limbs - 1; - if i > 0 { - self.limbs_mut()[i + dat.min_target_limb - 1] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb - 1], - dat.mask_last_limb_a(other, i + dat.min_source_limb), - c.clone(), - ); - self.limbs_mut()[i + dat.min_target_limb - 1] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb - 1]); - if dat.number_of_source_limbs == dat.number_of_target_limbs { - self.limbs_mut()[i + dat.min_target_limb] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb], - dat.mask_last_limb_b(other, i + dat.min_source_limb), - c, - ); - self.limbs_mut()[i + dat.min_target_limb] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb]); - } - } else { - self.limbs_mut()[i + dat.min_target_limb] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb]); - } - } - - fn add_shift_right(&mut self, other: FqSlice<'_, F>, c: FieldElement) { - struct AddShiftRightData { - offset_shift: usize, - tail_shift: usize, - zero_bits: usize, - min_source_limb: usize, - min_target_limb: usize, - number_of_source_limbs: usize, - number_of_target_limbs: usize, - min_mask: Limb, - max_mask: Limb, - } - - impl AddShiftRightData { - fn new(fq: F, target: FqSlice<'_, F>, source: FqSlice<'_, F>) -> Self { - debug_assert!(target.prime() == source.prime()); - debug_assert!(target.offset() >= source.offset()); - debug_assert!( - target.len() == source.len(), - "self.dim {} not equal to other.dim {}", - target.len(), - source.len() - ); - let offset_shift = target.offset() - source.offset(); - let bit_length = fq.bit_length(); - let entries_per_limb = fq.entries_per_limb(); - let usable_bits_per_limb = bit_length * entries_per_limb; - let tail_shift = usable_bits_per_limb - offset_shift; - let zero_bits = constants::BITS_PER_LIMB - usable_bits_per_limb; - let source_range = source.limb_range(); - let target_range = target.limb_range(); - let min_source_limb = source_range.start; - let min_target_limb = target_range.start; - let number_of_source_limbs = source_range.len(); - let number_of_target_limbs = target_range.len(); - let (min_mask, max_mask) = source.limb_masks(); - Self { - offset_shift, - tail_shift, - zero_bits, - min_source_limb, - min_target_limb, - number_of_source_limbs, - number_of_target_limbs, - min_mask, - max_mask, - } - } - - fn mask_first_limb_a(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - let source_limb_masked = other.limbs()[i] & self.min_mask; - (source_limb_masked << (self.offset_shift + self.zero_bits)) >> self.zero_bits - } - - fn mask_first_limb_b(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - let source_limb_masked = other.limbs()[i] & self.min_mask; - source_limb_masked >> self.tail_shift - } - - fn mask_middle_limb_a(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - (other.limbs()[i] << (self.offset_shift + self.zero_bits)) >> self.zero_bits - } - - fn mask_middle_limb_b(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - other.limbs()[i] >> self.tail_shift - } - - fn mask_last_limb_a(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - let source_limb_masked = other.limbs()[i] & self.max_mask; - source_limb_masked << self.offset_shift - } - - fn mask_last_limb_b(&self, other: FqSlice<'_, F>, i: usize) -> Limb { - let source_limb_masked = other.limbs()[i] & self.max_mask; - source_limb_masked >> self.tail_shift - } - } - - let dat = AddShiftRightData::new(self.fq(), self.as_slice(), other); - let mut i = 0; - { - self.limbs_mut()[i + dat.min_target_limb] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb], - dat.mask_first_limb_a(other, i + dat.min_source_limb), - c.clone(), - ); - self.limbs_mut()[i + dat.min_target_limb] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb]); - if dat.number_of_target_limbs > 1 { - self.limbs_mut()[i + dat.min_target_limb + 1] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb + 1], - dat.mask_first_limb_b(other, i + dat.min_source_limb), - c.clone(), - ); - } - } - for i in 1..dat.number_of_source_limbs - 1 { - self.limbs_mut()[i + dat.min_target_limb] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb], - dat.mask_middle_limb_a(other, i + dat.min_source_limb), - c.clone(), - ); - self.limbs_mut()[i + dat.min_target_limb] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb]); - self.limbs_mut()[i + dat.min_target_limb + 1] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb + 1], - dat.mask_middle_limb_b(other, i + dat.min_source_limb), - c.clone(), - ); - } - i = dat.number_of_source_limbs - 1; - if i > 0 { - self.limbs_mut()[i + dat.min_target_limb] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb], - dat.mask_last_limb_a(other, i + dat.min_source_limb), - c.clone(), - ); - self.limbs_mut()[i + dat.min_target_limb] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb]); - if dat.number_of_target_limbs > dat.number_of_source_limbs { - self.limbs_mut()[i + dat.min_target_limb + 1] = self.fq().fma_limb( - self.limbs()[i + dat.min_target_limb + 1], - dat.mask_last_limb_b(other, i + dat.min_source_limb), - c.clone(), - ); - } - } - if dat.number_of_target_limbs > dat.number_of_source_limbs { - self.limbs_mut()[i + dat.min_target_limb + 1] = - self.fq().reduce(self.limbs()[i + dat.min_target_limb + 1]); - } - } - /// Given a mask v, add the `v[i]`th entry of `other` to the `i`th entry of `self`. pub fn add_masked(&mut self, other: FqSlice<'_, F>, c: FieldElement, mask: &[usize]) { // TODO: If this ends up being a bottleneck, try to use PDEP/PEXT diff --git a/ext/crates/fp/src/vector/impl_fqvector.rs b/ext/crates/fp/src/vector/impl_fqvector.rs index 2a472a8b5e..7e2a6a1f4e 100644 --- a/ext/crates/fp/src/vector/impl_fqvector.rs +++ b/ext/crates/fp/src/vector/impl_fqvector.rs @@ -118,12 +118,12 @@ impl FqVector { if c == fq.zero() { self.set_to_zero(); + return; } - if fq.q() != 2 { - for limb in self.limbs_mut() { - *limb = fq.reduce(fq.fma_limb(0, *limb, c.clone())); - } + if fq.q() == 2 { + return; } + fq.scale_groups(self.limbs_mut(), c); } /// Add `other` to `self` on the assumption that the first `offset` entries of `other` are @@ -134,24 +134,21 @@ impl FqVector { assert_eq!(self.len(), other.len()); let fq = self.fq(); - let min_limb = offset / fq.entries_per_limb(); + // The first limb of the group containing `offset`. Since `other`'s entries below + // `offset` are zero, starting at the group boundary and adding whole groups is safe. + let min_limb = fq.group_of(offset) * fq.limbs_per_group(); if fq.q() == 2 { if c != fq.zero() { crate::simd::add_simd(self.limbs_mut(), other.limbs(), min_limb); } } else { - for (left, right) in self - .limbs_mut() - .iter_mut() - .zip_eq(other.limbs()) - .skip(min_limb) - { - *left = fq.fma_limb(*left, *right, c.clone()); - } - for limb in self.limbs_mut()[min_limb..].iter_mut() { - *limb = fq.reduce(*limb); - } + let end = self.limbs().len(); + fq.add_groups( + &mut self.limbs_mut()[min_limb..end], + &other.limbs()[min_limb..end], + c, + ); } } @@ -207,6 +204,20 @@ impl FqVector { assert_eq!(self.len(), slice.len()); let fq = self.fq(); + if fq.is_bitsliced() { + // The bit-sliced layout interleaves an entry's bits across planes, so we cannot + // `pack` contiguous chunks; scatter each entry into its group. + let num_limbs = fq.number(self.len()); + { + let v = self.vec_mut(); + v.clear(); + v.resize(num_limbs, 0); + } + for (i, x) in slice.iter().enumerate() { + self.set_entry(i, x.clone()); + } + return; + } self.vec_mut().clear(); self.vec_mut().extend( slice @@ -237,6 +248,14 @@ impl FqVector { pub fn add_truncate(&mut self, other: &Self, c: FieldElement) -> Option<()> { assert_eq!(self.fq(), other.fq()); let fq = self.fq(); + if fq.is_bitsliced() { + // `truncate` guards against an entry's packed sum carrying into the next entry's + // bits. The bit-sliced layout reduces each lane independently (every plane is a + // separate limb), so no such carry can occur and the addition never fails. The + // packed `fma_limb`/`truncate` loop below would instead corrupt the bit-planes. + self.add(other, c); + return Some(()); + } for (left, right) in self.limbs_mut().iter_mut().zip_eq(other.limbs()) { *left = fq.fma_limb(*left, *right, c.clone()); *left = fq.truncate(*left)?; @@ -297,6 +316,9 @@ impl FqVector { /// Find the index and value of the first non-zero entry of the vector. `None` if the vector is zero. pub fn first_nonzero(&self) -> Option<(usize, FieldElement)> { + if self.fq().is_bitsliced() { + return self.as_slice().first_nonzero(); + } let entries_per_limb = self.fq().entries_per_limb(); let bit_length = self.fq().bit_length(); let bitmask = self.fq().bitmask(); diff --git a/ext/crates/fp/src/vector/iter.rs b/ext/crates/fp/src/vector/iter.rs index 3dcc00a0a4..aa959c48fc 100644 --- a/ext/crates/fp/src/vector/iter.rs +++ b/ext/crates/fp/src/vector/iter.rs @@ -4,9 +4,22 @@ use crate::{ limb::Limb, }; +/// Read entry `idx` (an absolute index into `limbs`) under the bit-sliced layout, by +/// gathering its bit from each plane of its group. +#[inline] +fn gather_at(fq: F, limbs: &[Limb], idx: usize) -> FieldElement { + let lpg = fq.limbs_per_group(); + let base = fq.group_of(idx) * lpg; + fq.gather(&limbs[base..base + lpg], fq.lane_of(idx)) +} + pub struct FqVectorIterator<'a, F> { fq: F, limbs: &'a [Limb], + // Bit-sliced path: `pos` is the absolute index of the next entry to emit. + bitsliced: bool, + pos: usize, + // Packed path state. bit_length: usize, bit_mask: Limb, entries_per_limb_m_1: usize, @@ -19,12 +32,16 @@ pub struct FqVectorIterator<'a, F> { impl<'a, F: Field> FqVectorIterator<'a, F> { pub(super) fn new(vec: FqSlice<'a, F>) -> Self { let counter = vec.len(); + let fq = vec.fq(); + let start = vec.start(); let limbs = vec.into_limbs(); if counter == 0 { return Self { - fq: vec.fq(), + fq, limbs, + bitsliced: fq.is_bitsliced(), + pos: start, bit_length: 0, entries_per_limb_m_1: 0, bit_mask: 0, @@ -34,20 +51,37 @@ impl<'a, F: Field> FqVectorIterator<'a, F> { counter, }; } - let pair = vec.fq().limb_bit_index_pair(vec.start()); - let bit_length = vec.fq().bit_length(); - let cur_limb = limbs[pair.limb] >> pair.bit_index; + if fq.is_bitsliced() { + return Self { + fq, + limbs, + bitsliced: true, + pos: start, + bit_length: 0, + entries_per_limb_m_1: 0, + bit_mask: 0, + limb_index: 0, + entries_left: 0, + cur_limb: 0, + counter, + }; + } - let entries_per_limb = vec.fq().entries_per_limb(); + let pair = fq.limb_bit_index_pair(start); + let bit_length = fq.bit_length(); + let cur_limb = limbs[pair.limb] >> pair.bit_index; + let entries_per_limb = fq.entries_per_limb(); Self { - fq: vec.fq(), + fq, limbs, + bitsliced: false, + pos: start, bit_length, entries_per_limb_m_1: entries_per_limb - 1, - bit_mask: vec.fq().bitmask(), + bit_mask: fq.bitmask(), limb_index: pair.limb, - entries_left: entries_per_limb - (vec.start() % entries_per_limb), + entries_left: entries_per_limb - (start % entries_per_limb), cur_limb, counter, } @@ -55,9 +89,15 @@ impl<'a, F: Field> FqVectorIterator<'a, F> { pub fn skip_n(&mut self, mut n: usize) { if n >= self.counter { + self.pos += self.counter; self.counter = 0; return; } + if self.bitsliced { + self.pos += n; + self.counter -= n; + return; + } let entries_per_limb = self.entries_per_limb_m_1 + 1; if n < self.entries_left { self.entries_left -= n; @@ -90,7 +130,16 @@ impl Iterator for FqVectorIterator<'_, F> { fn next(&mut self) -> Option { if self.counter == 0 { return None; - } else if self.entries_left == 0 { + } + + if self.bitsliced { + let result = gather_at(self.fq, self.limbs, self.pos); + self.pos += 1; + self.counter -= 1; + return Some(result); + } + + if self.entries_left == 0 { self.limb_index += 1; self.cur_limb = self.limbs[self.limb_index]; self.entries_left = self.entries_per_limb_m_1; @@ -117,6 +166,10 @@ impl ExactSizeIterator for FqVectorIterator<'_, F> { pub struct FqVectorNonZeroIterator<'a, F> { fq: F, limbs: &'a [Limb], + // Bit-sliced path: `start` is the slice's absolute start; `idx` is the relative cursor. + bitsliced: bool, + start: usize, + // Shared/packed path state. limb_index: usize, cur_limb_entries_left: usize, cur_limb: Limb, @@ -126,29 +179,34 @@ pub struct FqVectorNonZeroIterator<'a, F> { impl<'a, F: Field> FqVectorNonZeroIterator<'a, F> { pub(super) fn new(vec: FqSlice<'a, F>) -> Self { - let entries_per_limb = vec.fq().entries_per_limb(); - + let fq = vec.fq(); let dim = vec.len(); + let start = vec.start(); let limbs = vec.into_limbs(); - if dim == 0 { + if dim == 0 || fq.is_bitsliced() { return Self { - fq: vec.fq(), + fq, limbs, + bitsliced: fq.is_bitsliced(), + start, limb_index: 0, cur_limb_entries_left: 0, cur_limb: 0, idx: 0, - dim: 0, + dim, }; } - let min_index = vec.start(); - let pair = vec.fq().limb_bit_index_pair(min_index); + + let entries_per_limb = fq.entries_per_limb(); + let pair = fq.limb_bit_index_pair(start); let cur_limb = limbs[pair.limb] >> pair.bit_index; - let cur_limb_entries_left = entries_per_limb - (min_index % entries_per_limb); + let cur_limb_entries_left = entries_per_limb - (start % entries_per_limb); Self { - fq: vec.fq(), + fq, limbs, + bitsliced: false, + start, limb_index: pair.limb, cur_limb_entries_left, cur_limb, @@ -162,6 +220,19 @@ impl Iterator for FqVectorNonZeroIterator<'_, F> { type Item = (usize, FieldElement); fn next(&mut self) -> Option { + if self.bitsliced { + let zero = self.fq.zero(); + while self.idx < self.dim { + let value = gather_at(self.fq, self.limbs, self.start + self.idx); + let cur = self.idx; + self.idx += 1; + if value != zero { + return Some((cur, value)); + } + } + return None; + } + let bit_length: usize = self.fq.bit_length(); let bitmask: Limb = self.fq.bitmask(); let entries_per_limb: usize = self.fq.entries_per_limb(); diff --git a/ext/crates/fp/tests/serde_format.rs b/ext/crates/fp/tests/serde_format.rs index 04fa2db115..c8627a7d6a 100644 --- a/ext/crates/fp/tests/serde_format.rs +++ b/ext/crates/fp/tests/serde_format.rs @@ -249,7 +249,8 @@ fn fpvector_p3_json_format() { }, "len": 5, "limbs": [ - 5137 + 17, + 10 ] }"#]] .assert_eq(&s); @@ -266,7 +267,9 @@ fn fpvector_p5_json_format() { }, "len": 4, "limbs": [ - 98372 + 8, + 10, + 1 ] }"#]] .assert_eq(&s);