p3-binary-dft 0.9.0-rc.1

Additive NTT over F_2-linear subspaces of the Plonky3 binary tower fields
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

//! Paired polynomial-field butterflies without extracting scalar lanes.

use core::arch::aarch64::*;

use p3_binary_field::Poly64;

/// Apply each butterfly to two words at a time, returning the covered prefix length.
#[inline]
pub(super) fn butterfly<const INVERSE: bool>(
    lo: &mut [Poly64],
    hi: &mut [Poly64],
    t: Poly64,
) -> usize {
    let covered = lo.len() / 2 * 2;
    // SAFETY: both slices have the same length and each iteration accesses two valid elements.
    // The module is only compiled when the carryless-multiply feature is enabled.
    unsafe {
        let twiddle = vdupq_n_u64(t.to_bits());
        let bulk = covered / 8 * 8;
        for start in (0..bulk).step_by(8) {
            // Constant bounds let all four pairs stay in registers, without initializing tail slots.
            let mut low: [uint64x2_t; 4] =
                core::array::from_fn(|pair| vld1q_u64(lo.as_ptr().add(start + 2 * pair).cast()));
            let mut high: [uint64x2_t; 4] =
                core::array::from_fn(|pair| vld1q_u64(hi.as_ptr().add(start + 2 * pair).cast()));
            if INVERSE {
                high = core::array::from_fn(|pair| veorq_u64(high[pair], low[pair]));
            }
            let products: [uint64x2_t; 4] = high.map(|value| multiply(value, twiddle));
            low = core::array::from_fn(|pair| veorq_u64(low[pair], products[pair]));
            if !INVERSE {
                high = core::array::from_fn(|pair| veorq_u64(high[pair], low[pair]));
            }
            for pair in 0..4 {
                vst1q_u64(lo.as_mut_ptr().add(start + 2 * pair).cast(), low[pair]);
                vst1q_u64(hi.as_mut_ptr().add(start + 2 * pair).cast(), high[pair]);
            }
        }
        for start in (bulk..covered).step_by(2) {
            let low = vld1q_u64(lo.as_ptr().add(start).cast());
            let high = vld1q_u64(hi.as_ptr().add(start).cast());
            let high = if INVERSE { veorq_u64(high, low) } else { high };
            let low = veorq_u64(low, multiply(high, twiddle));
            let high = if INVERSE { high } else { veorq_u64(high, low) };
            vst1q_u64(lo.as_mut_ptr().add(start).cast(), low);
            vst1q_u64(hi.as_mut_ptr().add(start).cast(), high);
        }
    }
    covered
}

/// Keep eight rows in registers through all twelve butterflies of three stages.
#[inline]
pub(super) fn radix8<const INVERSE: bool>(rows: &mut [&mut [Poly64]; 8], t: &[Poly64; 7]) -> usize {
    let covered = rows[0].len() / 2 * 2;
    let bulk = covered / 4 * 4;
    radix8_groups::<INVERSE, 2>(rows, t, 0..bulk);
    radix8_groups::<INVERSE, 1>(rows, t, bulk..covered);
    covered
}

/// Process independent lane pairs together to overlap their multiplication chains.
#[inline(always)]
fn radix8_groups<const INVERSE: bool, const PAIRS: usize>(
    rows: &mut [&mut [Poly64]; 8],
    t: &[Poly64; 7],
    range: core::ops::Range<usize>,
) {
    // SAFETY: the caller checked equal row lengths. Every load/store covers one valid
    // pair in a disjoint row, and this module requires the carryless-multiply feature.
    unsafe {
        let twiddles = t.map(|t| vdupq_n_u64(t.to_bits()));
        for start in range.step_by(2 * PAIRS) {
            let mut values: [[uint64x2_t; PAIRS]; 8] = core::array::from_fn(|row| {
                core::array::from_fn(|pair| {
                    vld1q_u64(rows[row].as_ptr().add(start + 2 * pair).cast())
                })
            });
            macro_rules! step {
                ($lo:literal, $hi:literal, $t:literal) => {{
                    let (mut lo, mut hi) = (values[$lo], values[$hi]);
                    if INVERSE {
                        hi = core::array::from_fn(|pair| veorq_u64(hi[pair], lo[pair]));
                    }
                    if t[$t].to_bits() != 0 {
                        let products = hi.map(|value| multiply(value, twiddles[$t]));
                        lo = core::array::from_fn(|pair| veorq_u64(lo[pair], products[pair]));
                    }
                    if !INVERSE {
                        hi = core::array::from_fn(|pair| veorq_u64(hi[pair], lo[pair]));
                    }
                    values[$lo] = lo;
                    values[$hi] = hi;
                }};
            }
            if INVERSE {
                step!(0, 1, 3);
                step!(2, 3, 4);
                step!(4, 5, 5);
                step!(6, 7, 6);
                step!(0, 2, 1);
                step!(1, 3, 1);
                step!(4, 6, 2);
                step!(5, 7, 2);
                step!(0, 4, 0);
                step!(1, 5, 0);
                step!(2, 6, 0);
                step!(3, 7, 0);
            } else {
                step!(0, 4, 0);
                step!(1, 5, 0);
                step!(2, 6, 0);
                step!(3, 7, 0);
                step!(0, 2, 1);
                step!(1, 3, 1);
                step!(4, 6, 2);
                step!(5, 7, 2);
                step!(0, 1, 3);
                step!(2, 3, 4);
                step!(4, 5, 5);
                step!(6, 7, 6);
            }
            for (row, pairs) in rows.iter_mut().zip(values) {
                for (pair, value) in pairs.into_iter().enumerate() {
                    vst1q_u64(row.as_mut_ptr().add(start + 2 * pair).cast(), value);
                }
            }
        }
    }
}

/// Multiply and reduce two independent word products in the vector register file.
///
/// # Safety
/// Requires carryless multiplication enabled by the enclosing module's target gate.
#[inline(always)]
unsafe fn multiply(value: uint64x2_t, twiddle: uint64x2_t) -> uint64x2_t {
    // SAFETY: all operations use registers and the module requires the multiplication feature.
    unsafe {
        let a = vreinterpretq_p64_u64(value);
        let b = vreinterpretq_p64_u64(twiddle);
        let p0 = vreinterpretq_u64_p128(vmull_p64(vgetq_lane_p64::<0>(a), vgetq_lane_p64::<0>(b)));
        let p1 = vreinterpretq_u64_p128(vmull_high_p64(a, b));
        let tail = vreinterpretq_p64_u64(vdupq_n_u64(0x1b));
        let t0 = vreinterpretq_u64_p128(vmull_high_p64(vreinterpretq_p64_u64(p0), tail));
        let t1 = vreinterpretq_u64_p128(vmull_high_p64(vreinterpretq_p64_u64(p1), tail));
        let u0 = vreinterpretq_u64_p128(vmull_high_p64(vreinterpretq_p64_u64(t0), tail));
        let u1 = vreinterpretq_u64_p128(vmull_high_p64(vreinterpretq_p64_u64(t1), tail));
        let (mut e0, mut e1) = (xor3(p0, t0, u0), xor3(p1, t1, u1));
        // Keep the reductions separate until their low lanes are interleaved.
        core::arch::asm!(
            "/* {0:v} {1:v} */",
            inout(vreg) e0,
            inout(vreg) e1,
            options(pure, nomem, nostack, preserves_flags)
        );
        vzip1q_u64(e0, e1)
    }
}

/// Sum three vectors, using the three-input instruction when available.
#[inline(always)]
unsafe fn xor3(a: uint64x2_t, b: uint64x2_t, c: uint64x2_t) -> uint64x2_t {
    // SAFETY: the optional instruction is guarded by its compile-time target feature.
    #[cfg(target_feature = "sha3")]
    {
        unsafe { veor3q_u64(a, b, c) }
    }
    #[cfg(not(target_feature = "sha3"))]
    {
        unsafe { veorq_u64(a, veorq_u64(b, c)) }
    }
}