FFT context initialization speedup

doc/source/ulong_extras.rst

@@ -1632,10 +1632,22 @@ Primitive roots and discrete logarithms
 --------------------------------------------------------------------------------
 
 
+.. function:: ulong n_quadratic_nonresidue(ulong n)
+
+    Given an odd non-square `n`, returns an `a` with Jacobi symbol
+    `\left(\frac{a}{n}\right) = -1`. In particular, if `n` is prime, `a`
+    is a quadratic nonresidue modulo `n`.
+
+
 .. function:: ulong n_primitive_root_prime_prefactor(ulong p, n_factor_t * factors)
 
-    Returns a primitive root for the multiplicative subgroup of `\mathbb{Z}/p\mathbb{Z}`
-    where `p` is prime given the factorisation (``factors``) of `p - 1`.
+    Given a prime `p` and a list of prime divisors of `p - 1` in
+    ``factors``, returns an element `a` of `(\mathbb{Z}/p\mathbb{Z})^*`
+    such that `a^{(p-1)/q} \ne 1` for every prime `q` in ``factors``.
+    The exponents in ``factors`` are ignored.
+
+    In particular, if ``factors`` contains all prime divisors of `p - 1`,
+    this returns a primitive root modulo `p`.
 
 
 .. function:: ulong n_primitive_root_prime(ulong p)

src/fft_small.h

@@ -145,8 +145,9 @@ int fft_small_mulmod_satisfies_bounds(ulong n);
         [j_bits][j_r]  where j_bits = nbits(j), j_r = j - 2^(j_bits-1)
 
     with the special case j_bits = j_r = 0 for j = 0.
-    The first SD_FFT_CTX_W2TAB_INIT tables are stored consecutively to ease the
-    lookup of small indices, which must currently be at least max(4, LG_BLK_SZ).
+    Up to the first SD_FFT_CTX_W2TAB_INIT tables are stored consecutively to ease
+    the lookup of small indices, which must currently be at least max(4,
+    LG_BLK_SZ).
 */
 
 /* for the fft look up of powers of w */
@@ -192,7 +193,7 @@ typedef struct {
     double p;   /* the FFT prime */
     double pinv;
     nmod_t mod;
-    ulong primitive_root;
+    ulong primitive_2power_root; /* primitive 2^v-th root, v = valuation(p - 1, 2) */
 #if FLINT_USES_PTHREAD
     _Atomic(unsigned int) w2tab_depth;
 #else

src/fft_small/nmod_poly_mul.c

@@ -836,7 +836,7 @@ _nmod_poly_should_directly_fft(ulong bn, ulong depth, nmod_t mod)
         /* unlikely, the convolution length would have to be massive */
         return 0;
 
-    if (n_trailing_zeros(mod.n - 1) < n_max(depth, SD_FFT_CTX_W2TAB_INIT))
+    if (n_trailing_zeros(mod.n - 1) < depth)
         return 0;
 
     return n_is_prime(mod.n); /* check the most expensive condition last */

src/fft_small/sd_fft.c

@@ -1028,12 +1028,12 @@ void sd_fft_trunc(
     FLINT_ASSERT(itrunc <= n_pow2(L));
     FLINT_ASSERT(otrunc <= n_pow2(L));
 
+    sd_fft_ctx_fit_depth(Q, L);
+
     if (L > LG_BLK_SZ)
     {
         ulong new_itrunc, new_otrunc;
 
-        sd_fft_ctx_fit_depth(Q, L);
-
         new_itrunc = n_cdiv(itrunc, BLK_SZ);
         new_otrunc = n_cdiv(otrunc, BLK_SZ);
         /* this isn't very clever */

src/fft_small/sd_fft_ctx.c

@@ -25,6 +25,28 @@ void sd_fft_ctx_clear(sd_fft_ctx_t Q)
 #endif
 }
 
+/*
+    Return a primitive 2^depth-th root modulo the prime pp.
+    Requires depth == valuation(pp - 1, 2).
+*/
+static ulong
+sd_fft_ctx_primitive_2power_root(ulong pp, ulong depth, nmod_t mod)
+{
+    ulong a = n_quadratic_nonresidue(pp);
+    return nmod_pow_ui(a, (pp - 1) >> depth, mod);
+}
+
+/*
+    Return the primitive 2^(k+1)-th root used to generate w2tab[k].
+    Requires depth == valuation(Q->mod.n - 1, 2).
+*/
+static ulong
+sd_fft_ctx_w2tab_root(const sd_fft_ctx_t Q, ulong depth, ulong k)
+{
+    FLINT_ASSERT(k + 1 <= depth);
+    return nmod_pow_ui(Q->primitive_2power_root, UWORD(1) << (depth - k - 1), Q->mod);
+}
+
 /*
     Initialize FFT context.
     pp is a prime with at most ~ 50 bits (exactly representable with a `double`)
@@ -33,24 +55,30 @@ void sd_fft_ctx_clear(sd_fft_ctx_t Q)
 */
 void sd_fft_ctx_init_prime(sd_fft_ctx_t Q, ulong pp)
 {
-    ulong N, i, k, l;
+    ulong N, i, k, l, init_depth, two_power_depth;
     double * t;
-    double n, ninv;
+    double n, ninv, w;
 
     if (!fft_small_mulmod_satisfies_bounds(pp))
         flint_throw(FLINT_ERROR, "FFT prime %wu does not satisfy bounds for arithmetic", pp);
 
     Q->p = pp;
     Q->pinv = 1.0/Q->p;
     nmod_init(&Q->mod, pp);
-    Q->primitive_root = n_primitive_root_prime(pp);
+    two_power_depth = n_trailing_zeros(pp - 1);
+    if (two_power_depth == 0)
+        flint_throw(FLINT_ERROR, "Input %wu is either 2 or not a prime", pp);
+    Q->primitive_2power_root = sd_fft_ctx_primitive_2power_root(pp, two_power_depth, Q->mod);
+    init_depth = n_min(two_power_depth, SD_FFT_CTX_W2TAB_INIT);
+    if (init_depth < 4)
+        flint_throw(FLINT_ERROR, "Input %wu does not support FFT context initialization", pp);
 
     n = Q->p;
     ninv = Q->pinv;
 
     /*
-        fill wtab to a depth of SD_FFT_CTX_W2TAB_INIT:
-        2^(SD_FFT_CTX_W2TAB_INIT-1) entries: 1, e(1/4), e(1/8), e(3/8), ...
+        fill wtab to a depth of init_depth:
+        2^(init_depth-1) entries: 1, e(1/4), e(1/8), e(3/8), ...
 
         Q->w2tab[j] is itself a table of length 2^(j-1) containing 2^(j+1) st
         roots of unity. More documentation on the layout of w2tab can be found
@@ -60,23 +88,42 @@ void sd_fft_ctx_init_prime(sd_fft_ctx_t Q, ulong pp)
         they're stored as `double` to make use of the vectorized functions in
         machine_vectors.h.
     */
-    N = n_pow2(SD_FFT_CTX_W2TAB_INIT - 1);
+    N = n_pow2(init_depth - 1);
     t = (double*) flint_aligned_alloc(4096, n_round_up(N*sizeof(double), 4096));
 
     Q->w2tab[0] = t;
     t[0] = 1;
-    for (k = 1, l = 1; k < SD_FFT_CTX_W2TAB_INIT; k++, l *= 2)
+
+    {
+        ulong ww = sd_fft_ctx_w2tab_root(Q, two_power_depth, 3);
+        w = vec1d_reduce_0n_to_pmhn(ww, n);
+        double w2 = vec1d_reduce_pm1n_to_pmhn(vec1d_mulmod(w, w, n, ninv), n);
+
+        Q->w2tab[1] = t + 1;
+        t[1] = vec1d_reduce_pm1n_to_pmhn(vec1d_mulmod(w2, w2, n, ninv), n);
+
+        Q->w2tab[2] = t + 2;
+        t[2] = w2;
+        t[3] = vec1d_reduce_pm1n_to_pmhn(vec1d_mulmod(t[1], w2, n, ninv), n);
+    }
+
+    vec4d n4 = vec4d_set_d(n);
+    vec4d ninv4 = vec4d_set_d(ninv);
+
+    for (k = 3, l = 4; k < init_depth; k++, l *= 2)
     {
-        ulong ww = nmod_pow_ui(Q->primitive_root, (Q->mod.n - 1)>>(k + 1), Q->mod);
-        double w = vec1d_set_d(vec1d_reduce_0n_to_pmhn(ww, n));
         double* curr = t + l;
+        vec4d w4 = vec4d_set_d(w);
         Q->w2tab[k] = curr;
         i = 0; do {
-            vec1d x = vec1d_load(t + i);
-            x = vec1d_mulmod(x, w, n, ninv);
-            x = vec1d_reduce_pm1n_to_pmhn(x, n);
-            vec1d_store(curr + i, x);
-        } while (i += 1, i < l);
+            vec4d x = vec4d_load_aligned(t + i);
+            x = vec4d_mulmod(x, w4, n4, ninv4);
+            x = vec4d_reduce_pm1n_to_pmhn(x, n4);
+            vec4d_store_aligned(curr + i, x);
+        } while (i += 4, i < l);
+
+        if (k + 1 < init_depth)
+            w = vec1d_reduce_0n_to_pmhn(sd_fft_ctx_w2tab_root(Q, two_power_depth, k + 1), n);
     }
 
 #if FLINT_USES_PTHREAD
@@ -94,9 +141,9 @@ void sd_fft_ctx_init_prime(sd_fft_ctx_t Q, ulong pp)
 #endif
 
 #if FLINT_WANT_ASSERT
-    for (k = 1; k < SD_FFT_CTX_W2TAB_INIT; k++)
+    for (k = 1; k < init_depth; k++)
     {
-        ulong ww = nmod_pow_ui(Q->primitive_root, (Q->mod.n - 1)>>(k + 1), Q->mod);
+        ulong ww = sd_fft_ctx_w2tab_root(Q, two_power_depth, k);
         for (i = 0; i < n_pow2(k-1); i++)
         {
             ulong www = nmod_pow_ui(ww, n_revbin(i+n_pow2(k-1), k), Q->mod);
@@ -108,6 +155,11 @@ void sd_fft_ctx_init_prime(sd_fft_ctx_t Q, ulong pp)
 
 void sd_fft_ctx_fit_depth_with_lock(sd_fft_ctx_t Q, ulong depth)
 {
+    ulong two_power_depth = n_trailing_zeros(Q->mod.n - 1);
+
+    if (depth > two_power_depth)
+        flint_throw(FLINT_ERROR, "FFT prime %wu does not support depth %wu", Q->mod.n, depth);
+
 #if FLINT_USES_PTHREAD
     pthread_mutex_lock(&Q->mutex);
 #endif
@@ -121,7 +173,7 @@ void sd_fft_ctx_fit_depth_with_lock(sd_fft_ctx_t Q, ulong depth)
     while (k < depth)
     {
         ulong i, j, l, off;
-        ulong ww = nmod_pow_ui(Q->primitive_root, (Q->mod.n - 1)>>(k + 1), Q->mod);
+        ulong ww = sd_fft_ctx_w2tab_root(Q, two_power_depth, k);
         vec8d w    = vec8d_set_d(vec1d_reduce_0n_to_pmhn(ww, Q->p));
         vec8d n    = vec8d_set_d(Q->p);
         vec8d ninv = vec8d_set_d(Q->pinv);
@@ -153,7 +205,7 @@ void sd_fft_ctx_fit_depth_with_lock(sd_fft_ctx_t Q, ulong depth)
 
 #if FLINT_WANT_ASSERT
         {
-            ulong ww = nmod_pow_ui(Q->primitive_root, (Q->mod.n - 1)>>(k + 1), Q->mod);
+            ulong ww = sd_fft_ctx_w2tab_root(Q, two_power_depth, k);
             for (i = 0; i < n_pow2(k-1); i++)
             {
                 ulong www = nmod_pow_ui(ww, n_revbin(i+n_pow2(k-1), k), Q->mod);

src/fft_small/sd_ifft.c

@@ -1510,12 +1510,12 @@ void sd_ifft_trunc(
 {
     FLINT_ASSERT(trunc <= n_pow2(L));
 
+    sd_fft_ctx_fit_depth(Q, L);
+
     if (L > LG_BLK_SZ)
     {
         ulong new_trunc = n_cdiv(trunc, BLK_SZ);
 
-        sd_fft_ctx_fit_depth(Q, L);
-
         sd_ifft_trunc_internal(Q, d, 1, L - LG_BLK_SZ, 0, new_trunc, new_trunc, 0);
         return;
     }

src/fft_small/test/t-sd_fft.c

@@ -102,5 +102,12 @@ TEST_FUNCTION_START(sd_fft, state)
         sd_fft_ctx_clear(Q);
     }
 
+    {
+        sd_fft_ctx_t Q;
+        sd_fft_ctx_init_prime(Q, UWORD(257));
+        test_sd_fft_trunc(Q, 0, 8, 5, state);
+        sd_fft_ctx_clear(Q);
+    }
+
     TEST_FUNCTION_END(state);
 }

src/ulong_extras.h

@@ -479,6 +479,7 @@ void n_mulmod_and_precomp_shoup(ulong * ab, ulong * ab_precomp,
 
 /* Primitive roots and discrete logarithms ***********************************/
 
+ulong n_quadratic_nonresidue(ulong n);
 ulong n_primitive_root_prime_prefactor(ulong p, n_factor_t * factors);
 ulong n_primitive_root_prime(ulong p);
 

src/ulong_extras/primitive_root_prime.c

@@ -12,6 +12,18 @@
 
 #include "ulong_extras.h"
 
+ulong n_quadratic_nonresidue(ulong n)
+{
+    if ((n & UWORD(7)) == 3 || (n & UWORD(7)) == 5)
+        return 2;
+
+    for (ulong a = 3; ; a += 2)
+    {
+        if (n_jacobi_unsigned(a, n) == -1)
+            return a;
+    }
+}
+
 ulong n_primitive_root_prime_prefactor(ulong p, n_factor_t * factors)
 {
     if (p == 2)

src/ulong_extras/sqrtmod.c

@@ -83,11 +83,7 @@ ulong n_sqrtmod(ulong a, ulong p)
 
     b = n_powmod2_ui_preinv(a, p1, p, pinv);
 
-    for (k = 3; ; k+=2) /* 2 is a quadratic residue mod p = 8k + 1 */
-    {
-        if (n_jacobi_unsigned(k, p) == -1) break;
-    }
-
+    k = n_quadratic_nonresidue(p);
     g = n_powmod2_ui_preinv(k, p1, p, pinv);
     res = n_powmod2_ui_preinv(a, (p1 + 1) / 2, p, pinv);
 

src/ulong_extras/test/t-primitive_root_prime.c

@@ -17,6 +17,17 @@ TEST_FUNCTION_START(n_primitive_root_prime, state)
 {
     int i, j;
 
+    {
+        ulong n[] = { UWORD(15), UWORD(21), UWORD(33), UWORD(45) };
+
+        for (i = 0; i < 4; i++)
+        {
+            ulong nonresidue = n_quadratic_nonresidue(n[i]);
+            if (n_jacobi_unsigned(nonresidue, n[i]) != -1)
+                TEST_FUNCTION_FAIL("%wu is not a Jacobi nonresidue mod %wu\n", nonresidue, n[i]);
+        }
+    }
+
     for (i = 0; i < 100; i++)
     {
         ulong p = n_randtest_prime(state, 1);
@@ -27,6 +38,13 @@ TEST_FUNCTION_START(n_primitive_root_prime, state)
         n_factor(&factors, p - 1, 1);
 
         ulong root;
+        if (p != 2)
+        {
+            ulong nonresidue = n_quadratic_nonresidue(p);
+            if (n_jacobi_unsigned(nonresidue, p) != -1)
+                TEST_FUNCTION_FAIL("%wu is not a quadratic nonresidue mod %wu\n", nonresidue, p);
+        }
+
         root = n_primitive_root_prime(p);
 
         for (j = 0; j < factors.num; j++)