binius_math/ntt/mod.rs
1// Copyright 2024-2025 Irreducible Inc.
2
3//! Efficient implementations of the binary field additive NTT.
4//!
5//! See [LCH14] and [DP24] Section 2.3 for mathematical background.
6//!
7//! [LCH14]: <https://arxiv.org/abs/1404.3458>
8//! [DP24]: <https://eprint.iacr.org/2024/504>
9
10pub mod domain_context;
11mod neighbors_last;
12#[cfg(test)]
13mod tests_evaluation;
14#[cfg(test)]
15pub mod tests_reference;
16
17use binius_field::{BinaryField, PackedField};
18pub use neighbors_last::{
19 NeighborsLastMultiThread, NeighborsLastReference, NeighborsLastSingleThread,
20};
21
22use super::BinarySubspace;
23use crate::FieldSliceMut;
24
25/// The binary field additive NTT.
26///
27/// A number-theoretic transform (NTT) is a linear transformation on a finite field analogous to
28/// the discrete fourier transform. The version of the additive NTT we use is originally described
29/// in [LCH14]. In [DP24] Section 4.1, the authors present the LCH additive NTT algorithm in a way
30/// that makes apparent its compatibility with the FRI proximity test. Throughout the
31/// documentation, we will refer to the notation used in [DP24].
32///
33/// The additive NTT is parameterized by a binary field $K$ and $\mathbb{F}\_2$-linear subspace. We
34/// write $\beta_0, \ldots, \beta_{\ell-1}$ for the ordered basis elements of the subspace. The
35/// basis determines a novel polynomial basis and an evaluation domain. In the forward direction,
36/// the additive NTT transforms a vector of polynomial coefficients, with respect to the novel
37/// polynomial basis, into a vector of their evaluations over the evaluation domain. The inverse
38/// transformation interpolates polynomial values over the domain into novel polynomial basis
39/// coefficients.
40///
41/// An [`AdditiveNTT`] implementation with a maximum domain dimension of $\ell$ can be applied on
42/// a sequence of $\ell + 1$ evaluation domains of sizes $2^0, \ldots, 2^\ell$. These are the
43/// domains $S^{(\ell)}, S^{(\ell - 1)}, \ldots, S^{(0)}$ defined in [DP24] Section 4. The methods
44/// [`Self::forward_transform`] and [`Self::inverse_transform`] require a parameter
45/// `log_domain_size` that indicates which of the $S^(i)$ domains to use for the transformation's
46/// evaluation domain and novel polynomial basis. (Remember, the novel polynomial basis is itself
47/// parameterized by basis). **Counterintuitively, the space $S^(i+1)$ is not necessarily
48/// a subset of $S^i$**. We choose this behavior for the [`AdditiveNTT`] trait because it
49/// facilitates compatibility with FRI when batching proximity tests for codewords of different
50/// dimensions.
51///
52/// [LCH14]: <https://arxiv.org/abs/1404.3458>
53/// [DP24]: <https://eprint.iacr.org/2024/504>
54pub trait AdditiveNTT {
55 type Field: BinaryField;
56
57 /// Forward transformation as defined in [DP24], Section 2.3.
58 ///
59 /// Arguments:
60 /// - `data` is the data on which the NTT is performed.
61 /// - `skip_early` is the number of early layers that should be skipped
62 /// - `skip_late` is the number of late layers that should be skipped
63 ///
64 /// ## Preconditons
65 ///
66 /// - `data.len()` is a power of 2
67 /// - `skip_early + skip_late <= log2(data.len()) + P::LOG_WIDTH`
68 /// - `log2(data.len()) + P::LOG_WIDTH <= self.log_domain_size() + skip_late`
69 ///
70 /// [DP24]: <https://eprint.iacr.org/2024/504>
71 fn forward_transform<P: PackedField<Scalar = Self::Field>>(
72 &self,
73 data: FieldSliceMut<P>,
74 skip_early: usize,
75 skip_late: usize,
76 );
77
78 /// Inverse transformation of [`Self::forward_transform`].
79 ///
80 /// Note that "early" layers here refer to "early" time in the forward transform, i.e. layers
81 /// with low index in the forward transform.
82 ///
83 /// ## Preconditions
84 ///
85 /// - same as [`Self::forward_transform`]
86 fn inverse_transform<P: PackedField<Scalar = Self::Field>>(
87 &self,
88 data: FieldSliceMut<P>,
89 skip_early: usize,
90 skip_late: usize,
91 );
92
93 /// The associated [`DomainContext`].
94 fn domain_context(&self) -> &impl DomainContext<Field = Self::Field>;
95
96 /// See [`DomainContext::log_domain_size`].
97 fn log_domain_size(&self) -> usize {
98 self.domain_context().log_domain_size()
99 }
100
101 /// See [`DomainContext::subspace`].
102 fn subspace(&self, i: usize) -> BinarySubspace<Self::Field> {
103 self.domain_context().subspace(i)
104 }
105
106 /// See [`DomainContext::twiddle`].
107 fn twiddle(&self, i: usize, j: usize) -> Self::Field {
108 self.domain_context().twiddle(i, j)
109 }
110}
111
112/// Provides information about the domains $S^{(i)}$ and the associated twiddle factors.
113///
114/// Needed by the NTT and by FRI.
115pub trait DomainContext {
116 type Field: BinaryField;
117
118 /// Base 2 logarithm of the size of $S^{(0)}$, i.e., $\ell$.
119 ///
120 /// In other words: Index of the first layer that can _not_ be computed anymore.
121 /// I.e., number of the latest layer that _can_ be computed, plus one.
122 /// Layers are indexed starting from 0.
123 ///
124 /// If you intend to call the NTT with `skip_late = 0`, then this should be equal to the base 2
125 /// logarithm of the number of scalars in the input.
126 fn log_domain_size(&self) -> usize;
127
128 /// Returns the binary subspace with dimension $i$.
129 ///
130 /// In [DP24], this subspace is referred to as $S^{(\ell - i)}$, where $\ell$ is the maximum
131 /// domain size of the NTT, i.e., `self.log_domain_size()`. We choose to reverse the indexing
132 /// order with respect to the paper because it is more natural in code that the $i$th subspace
133 /// has dimension $i$.
134 ///
135 /// ## Preconditions
136 ///
137 /// - `i` must be less than or equal to `self.log_domain_size()`
138 ///
139 /// [DP24]: <https://eprint.iacr.org/2024/504>
140 fn subspace(&self, i: usize) -> BinarySubspace<Self::Field>;
141
142 /// Returns the twiddle of a certain block in a certain layer.
143 ///
144 /// The layer numbers start from 0, i.e., the earliest layer is layer 0.
145 ///
146 /// ## Preconditions
147 ///
148 /// - `layer < self.log_domain_size()`
149 /// - `block < 2^layer`
150 fn twiddle(&self, layer: usize, block: usize) -> Self::Field;
151}
152
153/// Make it so that references to a [`DomainContext` implement [`DomainContext`] themselves.
154///
155/// This is useful, for example, if you need two objects that each want to _own_ a
156/// [`DomainContext`], but you don't want to clone the [`DomainContext`].
157impl<T: DomainContext> DomainContext for &T {
158 type Field = T::Field;
159
160 fn log_domain_size(&self) -> usize {
161 (*self).log_domain_size()
162 }
163
164 fn subspace(&self, i: usize) -> BinarySubspace<Self::Field> {
165 (*self).subspace(i)
166 }
167
168 fn twiddle(&self, layer: usize, block: usize) -> Self::Field {
169 (*self).twiddle(layer, block)
170 }
171}