fgf 3.0.1

Faster Galois Fields: SIMD and portable finite-field arithmetic (binary towers and prime fields) and vector kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197

//! Byte-lane single-row AVX-512 kernels for the generic residue rows.
//!
//! Fixed-coefficient products over one-byte lanes at `P <= 127` through
//! two `vpshufb` nibble tables combined by modular sum: the 64-byte
//! vector body processes 64 raw bytes per iteration and tails run the
//! scalar table form. Entries are safe capability-token functions taking
//! an `X64V4Token` and validating their own geometry.

#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;

use super::src_tail_avx512;

// ---------------------------------------------------------------------------
// Byte fixed-coefficient (nibble-table lookups)
// ---------------------------------------------------------------------------

/// `dst[i] = c * src[i] mod P` over all raw bytes through two `vpshufb`
/// nibble tables combined by modular sum, `P <= 127`.
///
/// For moduli `P <= 127` the product `c * x` over a raw byte splits into
/// two nibble lookups: `c*(x mod 16) + c*16*(x div 16)`, each table holding
/// values below `P` and their sum below `2P`, so one unsigned-min
/// correction canonicalizes. Both 16-entry tables are generated through
/// bounded Barrett reduction; the vector body shuffles both nibbles of
/// 64 raw bytes at once, and tails use the same table values.
///
/// # Panics
///
/// Panics unless `dst` and `src` have equal lengths, or when `p` is zero
/// or exceeds 127.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_into_u8_avx512(_token: archmage::X64V4Token, dst: &mut [u8], src: &[u8], c: u8, p: u8) {
    assert!(p <= 127, "residue u8: modulus {p} exceeds 127");
    assert_eq!(
        dst.len(),
        src.len(),
        "residue::mul_into_u8_avx512: dst length differs from src"
    );
    let (lo_table, hi_table) = tables(c, p);
    let mut lo = [0u8; 16];
    let mut hi = [0u8; 16];
    _mm_storeu_si128(&mut lo, lo_table);
    _mm_storeu_si128(&mut hi, hi_table);
    let lo_v = _mm512_broadcast_i32x4(lo_table);
    let hi_v = _mm512_broadcast_i32x4(hi_table);
    let mask0f = _mm512_set1_epi8(0x0F);
    // `p <= 127` fits `i8` exactly; the cast only reinterprets the bits.
    #[allow(clippy::cast_possible_wrap)]
    let cp = _mm512_set1_epi8(p as i8);
    let (dst_lanes, dst_tail) = dst.as_chunks_mut::<64>();
    let (src_lanes, _) = src.as_chunks::<64>();
    for (d, s) in dst_lanes.iter_mut().zip(src_lanes) {
        let v = _mm512_loadu_si512(s);
        let rlo = _mm512_shuffle_epi8(lo_v, _mm512_and_si512(v, mask0f));
        let rhi = _mm512_shuffle_epi8(hi_v, _mm512_and_si512(_mm512_srli_epi16::<4>(v), mask0f));
        let sum = _mm512_add_epi8(rlo, rhi);
        let sub = _mm512_sub_epi8(sum, cp);
        _mm512_storeu_si512(d, _mm512_min_epu8(sum, sub));
    }
    for (d, s) in dst_tail.iter_mut().zip(src_tail_avx512(src)) {
        let t = lo[*s as usize & 15].wrapping_add(hi[*s as usize >> 4]);
        *d = if t >= p { t - p } else { t };
    }
}

/// `dst[i] += c * src[i] mod P` through the nibble tables, `P <= 127`.
///
/// The raw destination byte reduces through an 8-bit Barrett step
/// (`q = (x * floor(256/P)) >> 8` undershoots `floor(x / P)` by at most
/// one, so `r = x - q*P` sits below `2P` and one unsigned-min correction
/// canonicalizes it) and joins the canonical product below `2P`; the
/// output is canonical.
///
/// # Panics
///
/// Panics unless `dst` and `src` have equal lengths, or when `p` is zero
/// or exceeds 127.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_u8_avx512(_token: archmage::X64V4Token, dst: &mut [u8], src: &[u8], c: u8, p: u8) {
    assert!(p <= 127, "residue u8: modulus {p} exceeds 127");
    assert_eq!(
        dst.len(),
        src.len(),
        "residue::mul_add_u8_avx512: dst length differs from src"
    );
    let (lo_table, hi_table) = tables(c, p);
    let mut lo = [0u8; 16];
    let mut hi = [0u8; 16];
    _mm_storeu_si128(&mut lo, lo_table);
    _mm_storeu_si128(&mut hi, hi_table);
    let lo_v = _mm512_broadcast_i32x4(lo_table);
    let hi_v = _mm512_broadcast_i32x4(hi_table);
    let mask0f = _mm512_set1_epi8(0x0F);
    // `p <= 127` fits `i8` exactly; the cast only reinterprets the bits.
    #[allow(clippy::cast_possible_wrap)]
    let cp = _mm512_set1_epi8(p as i8);
    // `256 / P <= 85` for `P >= 3`: both casts are exact.
    #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
    let mu = _mm512_set1_epi16(((256u32 / u32::from(p)) as u16) as i16);
    let cp16 = _mm512_set1_epi16(i16::from(p));
    let (dst_lanes, dst_tail) = dst.as_chunks_mut::<64>();
    let (src_lanes, _) = src.as_chunks::<64>();
    for (d, s) in dst_lanes.iter_mut().zip(src_lanes) {
        let acc = _mm512_loadu_si512(&*d);
        let acc_lo = _mm512_and_si512(acc, _mm512_set1_epi16(255));
        let acc_hi = _mm512_srli_epi16::<8>(acc);
        let normed = {
            let q_lo = _mm512_srli_epi16::<8>(_mm512_mullo_epi16(acc_lo, mu));
            let q_hi = _mm512_srli_epi16::<8>(_mm512_mullo_epi16(acc_hi, mu));
            let r_lo = _mm512_sub_epi16(acc_lo, _mm512_mullo_epi16(q_lo, cp16));
            let r_hi = _mm512_sub_epi16(acc_hi, _mm512_mullo_epi16(q_hi, cp16));
            let c_lo = _mm512_min_epu16(r_lo, _mm512_sub_epi16(r_lo, cp16));
            let c_hi = _mm512_min_epu16(r_hi, _mm512_sub_epi16(r_hi, cp16));
            _mm512_or_si512(c_lo, _mm512_slli_epi16::<8>(c_hi))
        };
        let v = _mm512_loadu_si512(s);
        let rlo = _mm512_shuffle_epi8(lo_v, _mm512_and_si512(v, mask0f));
        let rhi = _mm512_shuffle_epi8(hi_v, _mm512_and_si512(_mm512_srli_epi16::<4>(v), mask0f));
        let prod = _mm512_min_epu8(
            _mm512_add_epi8(rlo, rhi),
            _mm512_sub_epi8(_mm512_add_epi8(rlo, rhi), cp),
        );
        let sum = _mm512_add_epi8(normed, prod);
        _mm512_storeu_si512(d, _mm512_min_epu8(sum, _mm512_sub_epi8(sum, cp)));
    }
    for (d, s) in dst_tail.iter_mut().zip(src_tail_avx512(src)) {
        let acc = *d % p;
        let t = lo[*s as usize & 15].wrapping_add(hi[*s as usize >> 4]);
        let prod = if t >= p { t - p } else { t };
        let sum = acc + prod;
        *d = if sum >= p { sum - p } else { sum };
    }
}

/// `dst[i] = dst[i] * c mod P` through the nibble tables, `P <= 127`.
///
/// Arbitrary raw bytes reduce in place; the output is canonical.
///
/// # Panics
///
/// Panics when `p` is zero or exceeds 127.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_assign_u8_avx512(_token: archmage::X64V4Token, dst: &mut [u8], c: u8, p: u8) {
    assert!(p <= 127, "residue u8: modulus {p} exceeds 127");
    let (lo_table, hi_table) = tables(c, p);
    let mut lo = [0u8; 16];
    let mut hi = [0u8; 16];
    _mm_storeu_si128(&mut lo, lo_table);
    _mm_storeu_si128(&mut hi, hi_table);
    let lo_v = _mm512_broadcast_i32x4(lo_table);
    let hi_v = _mm512_broadcast_i32x4(hi_table);
    let mask0f = _mm512_set1_epi8(0x0F);
    // `p <= 127` fits `i8` exactly; the cast only reinterprets the bits.
    #[allow(clippy::cast_possible_wrap)]
    let cp = _mm512_set1_epi8(p as i8);
    let (dst_lanes, dst_tail) = dst.as_chunks_mut::<64>();
    for d in dst_lanes.iter_mut() {
        let v = _mm512_loadu_si512(&*d);
        let rlo = _mm512_shuffle_epi8(lo_v, _mm512_and_si512(v, mask0f));
        let rhi = _mm512_shuffle_epi8(hi_v, _mm512_and_si512(_mm512_srli_epi16::<4>(v), mask0f));
        let sum = _mm512_add_epi8(rlo, rhi);
        _mm512_storeu_si512(d, _mm512_min_epu8(sum, _mm512_sub_epi8(sum, cp)));
    }
    for d in dst_tail.iter_mut() {
        let t = lo[*d as usize & 15].wrapping_add(hi[*d as usize >> 4]);
        *d = if t >= p { t - p } else { t };
    }
}

/// Nibble products `lo[i] = c*i mod P` and `hi[i] = c*16*i mod P`.
#[archmage::rite(v4)]
fn tables(c: u8, p: u8) -> (__m128i, __m128i) {
    let indices = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
    let modulus = _mm512_set1_epi32(i32::from(p));
    let reciprocal = _mm512_set1_epi32(65536 / i32::from(p));
    let low = _mm512_mullo_epi32(indices, _mm512_set1_epi32(i32::from(c)));
    let high = _mm512_slli_epi32::<4>(low);
    (
        table_residues(low, reciprocal, modulus),
        table_residues(high, reciprocal, modulus),
    )
}

/// Reduce products below `2^16` through a reciprocal with one correction.
#[archmage::rite(v4)]
fn table_residues(products: __m512i, reciprocal: __m512i, modulus: __m512i) -> __m128i {
    // The largest product is `255 * 16 * 15 < 2^16`. Its reciprocal
    // product fits u32, and the quotient undershoots by at most one.
    let quotient = _mm512_srli_epi32::<16>(_mm512_mullo_epi32(products, reciprocal));
    let residue = _mm512_sub_epi32(products, _mm512_mullo_epi32(quotient, modulus));
    let canonical = _mm512_min_epu32(residue, _mm512_sub_epi32(residue, modulus));
    _mm512_cvtepi32_epi8(canonical)
}