qfall_math/integer_mod_q/ntt_basis_polynomial_ring_zq/
ntt.rs

1// Copyright 2025 Marvin Beckmann
2//
3// This file is part of qFALL-math.
4//
5// qFALL-math is free software: you can redistribute it and/or modify it under
6// the terms of the Mozilla Public License Version 2.0 as published by the
7// Mozilla Foundation. See <https://mozilla.org/en-US/MPL/2.0/>.
8
9//! This module contains implementations to perform the NTT for [`PolyOverZq`]
10//! in the respective polynomial ring.
11//! The implementation mostly follows the description in <https://higashi.blog/2023/06/23/ntt-02/>.
12//!
13//! The explicit functions contain the documentation.
14
15use super::{NTTBasisPolynomialRingZq, from::ConvolutionType};
16use crate::{
17    integer::Z,
18    integer_mod_q::{Modulus, PolyOverZq},
19    traits::GetCoefficient,
20    utils::index::bit_reverse_permutation,
21};
22use flint3_sys::{
23    fmpz_mod_ctx, {fmpz_mod_add, fmpz_mod_mul, fmpz_mod_sub},
24};
25
26impl NTTBasisPolynomialRingZq {
27    /// For a given polynomial viewed in the polynomial ring defined by the `self`, it computes the NTT.
28    ///
29    /// It computes the iterative Cooley-Tukey transformation algorithm to compute the transform.
30    /// Polynomials of degree smaller than `n-1` are with `0` coefficients.
31    ///
32    /// Parameters:
33    /// - `poly`: The polynomial for which it is desired to compute the NTT
34    ///
35    /// Returns the NTT of a polynomial in the context of the polynomial ring.
36    ///
37    /// # Examples
38    /// ```
39    /// use qfall_math::integer::Z;
40    /// use qfall_math::integer_mod_q::{Modulus, PolyOverZq, NTTBasisPolynomialRingZq, ConvolutionType};
41    /// use std::str::FromStr;
42    ///
43    /// let g_poly = PolyOverZq::from_str("4  1 2 3 4 mod 7681").unwrap();
44    /// let modulus = Modulus::from(7681);
45    ///
46    /// let ntt_basis =
47    ///     NTTBasisPolynomialRingZq::init(4, 3383, &modulus, ConvolutionType::Cyclic);
48    ///
49    /// let ghat = ntt_basis.ntt(&g_poly);
50    ///
51    /// let cmp_ghat = vec![
52    ///     Z::from(10),
53    ///     Z::from(913),
54    ///     Z::from(7679),
55    ///     Z::from(6764),
56    /// ];
57    /// assert_eq!(cmp_ghat, ghat);
58    /// ```
59    /// ```
60    /// use qfall_math::integer::Z;
61    /// use qfall_math::integer_mod_q::{Modulus, PolyOverZq, NTTBasisPolynomialRingZq, ConvolutionType};
62    /// use std::str::FromStr;
63    ///
64    /// let g_poly = PolyOverZq::from_str("4  1 2 3 4 mod 7681").unwrap();
65    /// let modulus = Modulus::from(7681);
66    ///
67    /// let ntt_basis =
68    ///     NTTBasisPolynomialRingZq::init(4, 1925, &modulus, ConvolutionType::Negacyclic);
69    ///
70    /// let ghat = ntt_basis.ntt(&g_poly);
71    ///
72    /// let cmp_ghat = vec![
73    ///     Z::from(1467),
74    ///     Z::from(2807),
75    ///     Z::from(3471),
76    ///     Z::from(7621),
77    /// ];
78    /// assert_eq!(cmp_ghat, ghat);
79    /// ```
80    ///
81    /// # Panics if ...
82    /// - it is not reduced, i.e. has a coefficient of degree > n
83    /// - the modulus differs from the modulus over which we view the polynomial
84    ///
85    /// The algorithm is based on the Cooley-Tukey algorithm:
86    /// -\[1\] Cooley, James W., and John W. Tukey.
87    ///     "An algorithm for the machine calculation of complex Fourier series."
88    ///     Mathematics of computation 19.90 (1965): 297-301.
89    pub fn ntt(&self, poly: &PolyOverZq) -> Vec<Z> {
90        assert!(poly.get_degree() < self.n);
91        assert_eq!(poly.get_mod(), self.modulus);
92        // we use the unsafe getter, because we know that all indices are in the range
93        // and no error can occur here
94        let mut poly_coeffs: Vec<Z> = (0..self.n)
95            .map(|i| unsafe { poly.get_coeff_unchecked(i) })
96            .collect();
97        for _ in poly_coeffs.len()..(self.n as usize) {
98            poly_coeffs.push(Z::default());
99        }
100
101        // Negacyclic: perform preprocessing
102        if self.convolution_type == ConvolutionType::Negacyclic {
103            for (i, x) in poly_coeffs.iter_mut().enumerate() {
104                unsafe {
105                    fmpz_mod_mul(
106                        &mut x.value,
107                        &x.value,
108                        &self.powers_of_psi[i].value,
109                        self.modulus.get_fmpz_mod_ctx_struct(),
110                    );
111                }
112            }
113        }
114
115        iterative_ntt(poly_coeffs, &self.powers_of_omega, &self.modulus)
116    }
117}
118
119/// This function essentially computes the included butterfly computations for each provided
120/// chunk.
121/// The chunk is double the size of the stride.
122/// The computation currently performs the standard butterfly operation from Cooley-Tukey
123unsafe fn ntt_stride_steps(
124    chunk: &mut [Z],
125    stride: usize,
126    power_pointer: i64,
127    modulus_pointer: &fmpz_mod_ctx,
128    powers_of_omega_pointers: &[Z],
129) {
130    for i in 0..stride {
131        // compute power of the current level
132        let current_power = &powers_of_omega_pointers[2_usize.pow(power_pointer as u32) * (i)];
133
134        // CT butterfly
135        // by using Z, we can manage not to initialize additional modulus objects in this part
136        // and save runtime
137        let mut temp = Z::default();
138
139        unsafe {
140            fmpz_mod_mul(
141                &mut temp.value,
142                &current_power.value,
143                &chunk[i + stride].value,
144                modulus_pointer,
145            );
146            fmpz_mod_sub(
147                &mut chunk[i + stride].value,
148                &chunk[i].value,
149                &temp.value,
150                modulus_pointer,
151            );
152            fmpz_mod_add(
153                &mut chunk[i].value,
154                &chunk[i].value,
155                &temp.value,
156                modulus_pointer,
157            )
158        }
159    }
160}
161
162/// This algorithm performs an iterative version of the Cooley-Tukey algorithm.
163/// It takes in the coefficients of the polynomial and the precomputed powers of the
164/// root of unity.
165/// Here, we assume that precomputation has been done, i.e.: if the algorithm is
166/// called for negatively wrapped convolution, then this has been accounted for in the previous algorithm.
167///
168/// The algorithm possesses the option to be multi-threaded, but benchmarking has shown,
169/// that it makes the algorithm less efficient, so we turned it off.
170fn iterative_ntt(coefficients: Vec<Z>, powers_of_omega: &[Z], modulus: &Modulus) -> Vec<Z> {
171    let n = coefficients.len();
172    let nr_iterations = n.ilog2() as i64;
173
174    // compute the bit reversed order of the coefficients
175    let mut res = coefficients;
176    bit_reverse_permutation(&mut res);
177    let modulus_pointer = modulus.get_fmpz_mod_ctx_struct();
178
179    let mut power_pointer: i64 = nr_iterations - 1;
180    let mut stride = 1;
181    // iterate through all layers
182    while stride < n {
183        // split into strides and perform action for each respective stride
184        res.chunks_mut(2 * stride).for_each(|chunk| unsafe {
185            ntt_stride_steps(
186                chunk,
187                stride,
188                power_pointer,
189                modulus_pointer,
190                powers_of_omega,
191            )
192        });
193        stride *= 2;
194        power_pointer -= 1;
195    }
196    res
197}
198
199#[cfg(test)]
200mod test_ntt {
201    use crate::{
202        integer::Z,
203        integer_mod_q::{ConvolutionType, Modulus, NTTBasisPolynomialRingZq, PolyOverZq},
204    };
205    use std::str::FromStr;
206
207    /// This example is taken from: https://eprint.iacr.org/2024/585.pdf Example 3.4
208    #[test]
209    fn example_34_multiplication_with_ntt() {
210        let g_poly = PolyOverZq::from_str("4  1 2 3 4 mod 7681").unwrap();
211        let modulus = Modulus::from(7681);
212
213        let ntt_basis = NTTBasisPolynomialRingZq::init(4, 3383, &modulus, ConvolutionType::Cyclic);
214
215        let ghat = ntt_basis.ntt(&g_poly);
216        let cmp_ghat = vec![Z::from(10), Z::from(913), Z::from(7679), Z::from(6764)];
217        assert_eq!(cmp_ghat, ghat);
218    }
219
220    /// Ensure that NTT panics, if the degree of the polynomial is too high, i.e. not reduced.
221    #[test]
222    #[should_panic]
223    fn degree_too_high() {
224        let g_poly = PolyOverZq::from_str("5  1 2 3 4 5 mod 7681").unwrap();
225        let modulus = Modulus::from(7681);
226
227        let ntt_basis = NTTBasisPolynomialRingZq::init(4, 3383, &modulus, ConvolutionType::Cyclic);
228
229        let _ = ntt_basis.ntt(&g_poly);
230    }
231
232    /// Ensure that NTT panics, if the modulus of the polynomial is different
233    #[test]
234    #[should_panic]
235    fn different_modulus() {
236        let g_poly = PolyOverZq::from_str("4  1 2 3 4 mod 7681").unwrap();
237        let modulus = Modulus::from(7682);
238
239        let ntt_basis = NTTBasisPolynomialRingZq::init(4, 3383, &modulus, ConvolutionType::Cyclic);
240
241        let _ = ntt_basis.ntt(&g_poly);
242    }
243
244    /// Ensure that NTT works for smaller degree polynomials
245    #[test]
246    fn small_degree() {
247        let g_poly = PolyOverZq::from_str("2  1 2 mod 7681").unwrap();
248        let modulus = Modulus::from(7681);
249
250        let ntt_basis =
251            NTTBasisPolynomialRingZq::init(4, 1925, &modulus, ConvolutionType::Negacyclic);
252
253        let ghat = ntt_basis.ntt(&g_poly);
254        let cmp_ghat = vec![Z::from(3851), Z::from(5256), Z::from(3832), Z::from(2427)];
255        assert_eq!(cmp_ghat, ghat);
256    }
257}
qfall_math/integer_mod_q/ntt_basis_polynomial_ring_zq/ntt.rs

qfall_math/integer_mod_q/ntt_basis_polynomial_ring_zq/
ntt.rs
