FFTs modulo word-size primes in floating point

The module fft_small computes number-theoretic transforms, that is, discrete Fourier transforms over \mathbb{Z}/p\mathbb{Z} [Pol1971], for primes p < 2^{50} with 2^L \mid p - 1. Residues are integers stored in double variables and the modular arithmetic is done in floating point with fused multiply-adds, four lanes per instruction with AVX2 and two with NEON, the code always working on vectors of eight. This approach to SIMD modular arithmetic is that of van der Hoeven, Lecerf and Quintin [HLQ2016]. The implementation, by Daniel Schultz, combines radix-4 butterflies with lazy reduction in the spirit of Harvey [Har2014], van der Hoeven's truncated transforms [vdH2004] and a recursive matrix decomposition of the transform [Bai1990]. How several primes are combined is described in Multiplication with several primes and CRT, and how polynomial multiplication uses the transforms in The small-prime FFT.

Throughout, p is the prime and u = 2^{-53}. Values are integers of either sign and are never reduced to [0, p) inside a transform; a bound on a value means a bound on its absolute value.

Primes and roots of unity#

A prime is usable if it passes the check of Arithmetic and its ranges, which implies p < 2^{50}, if 2^4 \mid p - 1, and if 2^L \mid p - 1 for the largest transform length 2^L used. If 2^v \| p - 1 and a is a quadratic nonresidue, g = a^{(p-1)/2^v} has order exactly 2^v, because g^{2^{v-1}} = a^{(p-1)/2} = -1. All roots of unity are powers of this g.

The transforms evaluate at the roots of unity in bit-reversed order. For j < 2^k let \operatorname{rev}_k(j) be j with its k low bits reversed, and let \zeta_{2^k} = g^{2^{v-k}}, so that \zeta_{2^{k+1}}^2 = \zeta_{2^k}. Define

This does not depend on k, since \operatorname{rev}_{k+1}(j) = 2 \operatorname{rev}_k(j) for j < 2^k. The sequence begins 1, -1, \zeta_4, -\zeta_4, \zeta_8, \zeta_8^5, \zeta_8^3, \zeta_8^7, \ldots; sd_fft_ctx_w returns its terms.

Since w_{2j+1} = -w_{2j}, only the even-indexed roots are stored, as integers in [-p/2, p/2]. Table k \ge 2 holds w_{2j} for 2^{k-1} \le j < 2^k, and it is obtained by multiplying all earlier stored roots by \zeta_{2^{k+1}}, since w_{2j} = \zeta_{2^{k+1}} w_{2j - 2^k} in that range. Tables are added on demand under a mutex and never move, so the transforms read them without locking. The inverse transform needs -w_{2j}^{-1}; since w_{2j}^{-1} is again a root of the same order, it is found in the same table at the mirrored index (sd_fft_ctx_w2inv).

The forward transform#

By Lemma 1, the values f(w_i) for i < 2^L can be computed by a tree of remainders [Ber2008]. At level h, block j (for j < 2^{L-h}) holds the 2^h coefficients of f \bmod (x^{2^h} - w_j); at level L there is the single block f \bmod (x^{2^L} - 1) = f, and at level 0 block i holds f \bmod (x - w_i) = f(w_i).

This is correct because x^{2^{h-1}} \equiv s modulo x^{2^{h-1}} - w_{2j} and x^{2^{h-1}} \equiv -s modulo x^{2^{h-1}} - w_{2j+1}. The multiplications by s are the only modular multiplications, and they vanish for j = 0, where s = 1. Two consecutive levels are merged into radix-4 butterflies: with the quarters a_0, a_1, a_2, a_3 of block j and s_1 = w_{2j}, s_2 = w_{4j}, s_3 = w_{4j+2} = \pm \zeta_4 s_2,

which costs four multiplications for four outputs and keeps the data in registers for two levels. A transform has \lfloor L/2 \rfloor such levels, plus a radix-2 level or an 8-point kernel when L is odd. The output order is the order of the blocks, so position i holds f(w_i), except that the innermost kernel saves its final 4 \times 4 transposition: for L \ge 4 the value f(w_i) is at position sd_fft_ctx_trunc_index(L, i), obtained by exchanging the two 2-bit fields of the low four bits of i. Pointwise products do not depend on the order, and the inverse transform undoes it.

Blocking#

The data is divided into blocks of 256 doubles (2 KB). Up to 256 points, the transform is done by kernels that keep the data in registers: radix-4 levels on vectors of eight lanes down to 16 points, then 4 \times 4 transposes and two levels on 4-lane vectors. A longer transform of 2^k blocks treats the blocks as the entries of a 2^{k_1} \times 2^{k_2} matrix, k_1 = \lfloor k/2 \rfloor, stored by rows. The first k_1 levels are transforms of length 2^{k_1} along the columns, each acting on whole blocks at stride 2^{k_2}; the remaining levels are independent transforms of the rows, each a contiguous transform of 2^{k_2} blocks with the roots of the row's position j 2^{k_1} + b in the tree. Both are done recursively in the same way. This is Bailey's four-step method [Bai1990] in the remainder-tree formulation, which needs neither twiddle multiplications between the steps nor a transposition (compare the formulation with twiddle factors in The matrix Fourier algorithm), and the recursive halving makes it cache-oblivious in the sense of [FLPR1999]: every subproblem eventually fits in some level of the cache. In the outer levels each butterfly is a loop over the 256 lanes of a block, which gives long vectorised loops with one set of roots.

Truncation#

For a product of length t < 2^L, only t values are needed and the inputs have at most t nonzero coefficients. sd_fft_trunc takes an input truncation (coefficients from there on are zero) and an output truncation (only positions below it are needed), in whole blocks for L > 8. In the matrix decomposition, a column transform is skipped if all its inputs are zero and computes only as many outputs as there are needed rows, and of the row transforms only the needed ones are computed, the last one truncated. At the radix-4 level, butterflies with known zero inputs or unneeded outputs are replaced by special versions, one for each of the 12 combinations of 2 to 4 inputs and 1 to 4 outputs. Since a product of length t \in (2^{L-1}, 2^L] needs a transform of length 2^L, truncation avoids the factor of up to two that an untruncated transform would cost [vdH2004].

The inverse transform#

The inverse runs the tree from the leaves to the root. If block 2j holds c = A + sB and block 2j+1 holds d = A - sB (all operations coefficient-wise), with s = w_{2j}, then c + d = 2A and (c - d) s^{-1} = 2B. The code computes (c + d, (d - c) W) with W = -s^{-1}, in radix-4 form, like the butterflies of Gentleman and Sande [GS1966]. The factors 2 are not removed, so the result is 2^L times the input polynomial; the factor 2^{-L} is applied together with the pointwise product.

The truncated inverse [vdH2004] starts from the values at the first t positions of the transform of a polynomial of length at most t. In a block where some values and some coefficients are known, the unknowns are solved for from c = A + sB and d = A - sB instead of inverting the butterfly: if c and B are known, then A = c - sB and the value d = c - 2sB follows, which the next level up may need. The recursion of sd_ifft_trunc follows the matrix decomposition of the forward transform, with three parameters per subproblem (the number of known entries, the number of wanted coefficients, and whether one further value is wanted), and special radix-2 and radix-4 butterflies for the combinations that occur. Truncations are in whole blocks for L > 8, and for L \le 8 the full transform is computed. See also Truncated FFTs.

Transposed transforms#

sd_fft_trunc_tr and sd_ifft_trunc_tr compute the transposes of the linear maps d \mapsto sd_fft_trunc(d) and d \mapsto sd_ifft_trunc(d), with the same parameters, so that \langle F(a), v \rangle = \langle a, F^T(v) \rangle for every input pair. They are used for the transposed products of Transposed FFT, which compute middle products at the cost of a product of the output length. Their derivation is mechanical. The transforms are straight-line programs of butterflies, block copies and block zeroings, and the transpose of such a program performs the transposed steps in the reverse order (Theorem 1 of Middle products and the transposition principle):

• The transpose of the forward butterfly (a, b) \mapsto (a + wb, a - wb) is (a, b) \mapsto (a + b, w(a - b)): a butterfly of the inverse (Gentleman–Sande) shape with the forward twiddle w. So sd_fft_trunc_tr uses the twiddle tables of sd_fft_trunc and the reduction pattern of sd_ifft_trunc, and conversely sd_ifft_trunc_tr uses forward butterflies with the inverse twiddles. In the source the four transforms share one set of butterfly macros (src/fft_small/butterflies.h), which take the modular multiplication as an argument: the inverse-shape butterfly computes (a, b) \mapsto (a + b, W(b - a)), and the transposed transforms pass a negated multiplication, so that W = -w without negating the twiddle tables. Every intermediate value then has the magnitude it has in the transform whose butterflies are shared, and the bound analyses below apply unchanged.
• The recursion is traversed in reverse: in the matrix decomposition the row transforms come before the column transforms, and within the basecases the levels run from last to first. The 4 \times 4 transposition that the forward kernels omit at the end (see sd_fft_ctx_trunc_index) becomes an omitted transposition at the start, so the transposed forward transform reads its input in the order in which the forward transform writes its output.
• Truncation transposes into truncation: not reading input i deletes a column of the matrix and not writing output o deletes a row, so the transposed truncated forward butterflies are the same macros with loads guarded by otrunc and stores by itrunc. The truncated inverse butterflies of The inverse transform are separately derived maps (thirteen radix-4 and six radix-2 variants), and each is transposed separately.

Each transposed transform therefore performs exactly the operations of the transform it transposes, and costs the same. sd_fft_trunc_tr reads the evaluation side and accepts the range (-2p, 2p) of sd_ifft_trunc inputs, which pointwise products produce, because its butterflies have the inverse shape; it also returns values in (-2p, 2p), which was checked by interval evaluation of its truncated butterflies and by propagating bounds through whole transforms up to depth 16. sd_ifft_trunc_tr has the forward shape and maps (-3p, 3p) into itself, as sd_fft_trunc does (Theorem 3). The tests check the adjoint identity on random vectors over a range of depths and truncations.

Pointwise products and normalisation#

After the forward transforms of a and b, the pointwise product and the inverse transform give 2^L times the cyclic convolution of a and b of length 2^L, which is their product if its length is at most the truncation. sd_fft_ctx_point_mul computes m x_i y_i with a constant m in the same pass, and the multiplication drivers take m = (c 2^L)^{-1} \bmod p, where c is the reduced Chinese remainder cofactor of p, so that neither the normalisation nor the inverse in the Chinese remainder formula costs anything extra.

Arithmetic and its ranges#

Each step is one of the following operations of Modular arithmetic in floating point, applied to integers stored exactly in doubles:

• reduce_to_pm1n maps an integer a with |a| < 2^{53} to an integer r \equiv a \pmod p with |r| < p/2 + 2 (by the reduction theorem);
• mulmod_fast(a, w) and mulmod(a, w) return r \equiv aw \pmod p;
• additions and subtractions, which are exact for results below 2^{53}, in particular for all sums below 8p.

For the products FLINT checks each prime with fft_small_mulmod_satisfies_bounds, from sd_fft_ctx_init_prime, for the two properties

(P1) for both functions and (P2) for mulmod. The check evaluates the error bound of the multiplication theorem with the actual error of \mathrm{fl}(1/p) instead of its maximum; every modulus below 2^{49} passes, and about one 50-bit prime in 900 fails. The roots of unity satisfy |w| \le (p - 1)/2.

The conversions of the multiplication drivers produce inputs below p + 2^{32}, so the outputs of their forward transforms are below \tfrac52 p + 2.

So an untruncated inverse transform maps (-2p, 2p) into itself, and by Proposition 4 the pointwise products are valid inputs.

The guaranteed ranges alone would allow |xy| up to 9p^2 and give no bound below 2p, so this step depends on the two specific primes. All other drivers multiply by the normalisation constant, or reduce, before forming the product, as in Proposition 4.

Cost#

A radix-4 butterfly with j \ne 0 performs four multiplications by roots, one reduction and eight additions for four points and two levels, so a full transform of length 2^L takes \tfrac12 L 2^L multiplications, \tfrac18 L 2^L reductions and L 2^L additions, on 8-lane vectors. A multiplication costs two floating-point multiplications, two fused multiply-adds, a rounding and an addition. Truncated transforms cost roughly in proportion to the number of points computed. The inverse transform has the same cost. The data is transformed in place, and the twiddle tables take half a double per point of the longest transform. The lengths from which these transforms beat the other polynomial multiplication algorithms are given in Cutoffs.