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}