Optimize trial division and half-limb probable prime test in n_is_prime

Author: fredrik-johansson

src/ulong_extras/is_prime.c

@@ -82,63 +82,136 @@ static int u32_is_base2_pseudoprime(uint32_t x)
 /* Faster arithmetic for 32-bit probable prime test on 64-bit machines. */
 #if FLINT_BITS == 64
 
-static uint32_t
-u32_addmod(uint32_t x, uint32_t y, uint32_t n)
+/* Use the division routines currently defined internally in radix.h;
+   todo - move to ulong_extras */
+#include "radix.h"
+
+static ulong
+n_rem_precomp_c0(ulong x, ulong d, const n_div_precomp_t pre)
 {
-    return (n - y > x ? x + y : x + y - n);
+    ulong q = n_div_precomp_c0(x, pre);
+    return x - d * q;
 }
 
-static uint32_t
-u32_rem_u64_shoup(uint64_t x, uint32_t n, uint64_t ninv)
+static ulong
+n_rem_precomp_c1_unsafe(ulong x, ulong d, const n_div_precomp_t pre)
 {
-    return n_mulmod_shoup(1, x, ninv, n);
+    ulong q = n_div_precomp_c1_unsafe(x, pre);
+    return x - d * q;
 }
 
-static uint32_t
-u32_2_powmod(uint32_t exp, uint32_t n, uint64_t ninv)
+static ulong
+n_2_powmod_c0(uint32_t exp, ulong n, const n_div_precomp_t pre)
 {
-    uint32_t x, bit;
-    unsigned int ebits;
+    ulong x;
+    int k, ebits;
 
     if (exp < FLINT_BITS)
-        return u32_rem_u64_shoup(UWORD(1) << exp, n, ninv);
+        return n_rem_precomp_c0(UWORD(1) << exp, n, pre);
 
     ebits = FLINT_BITS - flint_clz(exp);
-    bit = UWORD(1) << (ebits - 6);
+    x = n_rem_precomp_c0(UWORD(1) << (exp >> (ebits - 6)), n, pre);
 
-    x = u32_rem_u64_shoup(((uint64_t) 1) << (exp >> (ebits - 6)), n, ninv);
+    if (n < UWORD(1) << (FLINT_BITS / 2 - 1))
+    {
+        for (k = ebits - 7; k >= 0; k--)
+        {
+            x = n_rem_precomp_c0((ulong) x * x, n, pre);
+            x <<= ((exp >> k) & 1);
+        }
 
-    while (bit >>= 1)
+        if (x >= n)
+            x -= n;
+    }
+    else
     {
-        x = u32_rem_u64_shoup((uint64_t) x * x, n, ninv);
+        for (k = ebits - 7; k >= 0; k--)
+        {
+            x = n_rem_precomp_c0((ulong) x * x, n, pre);
+            x <<= ((exp >> k) & 1);
+            if (x >= n)
+                x -= n;
+        }
+    }
+
+    return x;
+}
+
+static ulong
+n_2_powmod_c1(uint32_t exp, ulong n, const n_div_precomp_t pre)
+{
+    ulong x;
+    int k, ebits;
+
+    if (exp < FLINT_BITS)
+        return n_rem_precomp_c1_unsafe(UWORD(1) << exp, n, pre);
 
-        if (bit & exp)
-            x = u32_addmod(x, x, n);
+    ebits = FLINT_BITS - flint_clz(exp);
+    x = n_rem_precomp_c1_unsafe(UWORD(1) << (exp >> (ebits - 6)), n, pre);
+
+    if (n < UWORD(1) << (FLINT_BITS / 2 - 1))
+    {
+        for (k = ebits - 7; k >= 0; k--)
+        {
+            x = n_rem_precomp_c1_unsafe(x * x, n, pre);
+            x <<= ((exp >> k) & 1);
+        }
+
+        if (x >= n)
+            x -= n;
+    }
+    else
+    {
+        for (k = ebits - 7; k >= 0; k--)
+        {
+            x = n_rem_precomp_c1_unsafe(x * x, n, pre);
+            x <<= ((exp >> k) & 1);
+            if (x >= n)
+                x -= n;
+        }
     }
 
     return x;
 }
 
 static int u32_is_base2_probabprime(uint32_t n)
 {
-    uint32_t d = n-1, s = 0;
-    uint64_t ninv = n_mulmod_precomp_shoup(UWORD(1), n);
+    ulong d = n-1, s = 0;
+
+    n_div_precomp_t pre;
+    n_div_precomp_init(pre, n);
 
     s = flint_ctz(d);
     d >>= s;
 
-    uint64_t t = d;
-    uint64_t y;
+    ulong t = d;
+    ulong y;
 
-    y = u32_2_powmod(t, n, ninv);
-    if (y == 1)
-        return 1;
+    if (pre->c == 0)
+    {
+        y = n_2_powmod_c0(t, n, pre);
+        if (y == 1)
+            return 1;
 
-    t <<= 1;
-    while ((t != n - 1) && (y != n - 1))
+        t <<= 1;
+        while ((t != n - 1) && (y != n - 1))
+        {
+            y = n_rem_precomp_c0(y * y, n, pre);
+            t <<= 1;
+        }
+    }
+    else
     {
-        y = u32_rem_u64_shoup(y * y, n, ninv);
+        y = n_2_powmod_c1(t, n, pre);
+        if (y == 1)
+            return 1;
+
         t <<= 1;
+        while ((t != n - 1) && (y != n - 1))
+        {
+            y = n_rem_precomp_c1_unsafe(y * y, n, pre);
+            t <<= 1;
+        }
     }
 
     return y == n - 1;
@@ -280,34 +353,129 @@ n_is_prime_odd_no_trial(ulong n)
     }
 }
 
+/* 3, 5 (checked separately), and 32 more primes which may be vectorized */
+#define NUM_ODD_TRIAL_PRIMES 34
+
+static const uint32_t trial_inv_32[2 * NUM_ODD_TRIAL_PRIMES] = {
+    UWORD(0xaaaaaaab), UWORD(0x55555555), /* 3 */
+    UWORD(0xcccccccd), UWORD(0x33333333), /* 5 */
+    UWORD(0xb6db6db7), UWORD(0x24924924), /* 7 */
+    UWORD(0xba2e8ba3), UWORD(0x1745d174), /* 11 */
+    UWORD(0xc4ec4ec5), UWORD(0x13b13b13), /* 13 */
+    UWORD(0xf0f0f0f1), UWORD(0xf0f0f0f), /* 17 */
+    UWORD(0x286bca1b), UWORD(0xd79435e), /* 19 */
+    UWORD(0xe9bd37a7), UWORD(0xb21642c), /* 23 */
+    UWORD(0x4f72c235), UWORD(0x8d3dcb0), /* 29 */
+    UWORD(0xbdef7bdf), UWORD(0x8421084), /* 31 */
+    UWORD(0x914c1bad), UWORD(0x6eb3e45), /* 37 */
+    UWORD(0xc18f9c19), UWORD(0x63e7063), /* 41 */
+    UWORD(0x2fa0be83), UWORD(0x5f417d0), /* 43 */
+    UWORD(0x677d46cf), UWORD(0x572620a), /* 47 */
+    UWORD(0x8c13521d), UWORD(0x4d4873e), /* 53 */
+    UWORD(0xa08ad8f3), UWORD(0x456c797), /* 59 */
+    UWORD(0xc10c9715), UWORD(0x4325c53), /* 61 */
+    UWORD(0x7a44c6b), UWORD(0x3d22635), /* 67 */
+    UWORD(0xe327a977), UWORD(0x39b0ad1), /* 71 */
+    UWORD(0xc7e3f1f9), UWORD(0x381c0e0), /* 73 */
+    UWORD(0x613716af), UWORD(0x33d91d2), /* 79 */
+    UWORD(0x2b2e43db), UWORD(0x3159721), /* 83 */
+    UWORD(0xfa3f47e9), UWORD(0x2e05c0b), /* 89 */
+    UWORD(0x5f02a3a1), UWORD(0x2a3a0fd), /* 97 */
+    UWORD(0x7c32b16d), UWORD(0x288df0c), /* 101 */
+    UWORD(0xd3431b57), UWORD(0x27c4597), /* 103 */
+    UWORD(0x8d28ac43), UWORD(0x2647c69), /* 107 */
+    UWORD(0xda6c0965), UWORD(0x2593f69), /* 109 */
+    UWORD(0xfdbc091), UWORD(0x243f6f0), /* 113 */
+    UWORD(0xefdfbf7f), UWORD(0x2040810), /* 127 */
+    UWORD(0xc9484e2b), UWORD(0x1f44659), /* 131 */
+    UWORD(0x77975b9), UWORD(0x1de5d6e), /* 137 */
+    UWORD(0x70586723), UWORD(0x1d77b65), /* 139 */
+    UWORD(0x8ce2cabd), UWORD(0x1b7d6c3), /* 149 */
+};
+
+static const uint64_t trial_inv_64[2 * NUM_ODD_TRIAL_PRIMES] = {
+    UWORD(0xaaaaaaaaaaaaaaab), UWORD(0x5555555555555555), /* 3 */
+    UWORD(0xcccccccccccccccd), UWORD(0x3333333333333333), /* 5 */
+    UWORD(0x6db6db6db6db6db7), UWORD(0x2492492492492492), /* 7 */
+    UWORD(0x2e8ba2e8ba2e8ba3), UWORD(0x1745d1745d1745d1), /* 11 */
+    UWORD(0x4ec4ec4ec4ec4ec5), UWORD(0x13b13b13b13b13b1), /* 13 */
+    UWORD(0xf0f0f0f0f0f0f0f1), UWORD(0xf0f0f0f0f0f0f0f), /* 17 */
+    UWORD(0x86bca1af286bca1b), UWORD(0xd79435e50d79435), /* 19 */
+    UWORD(0xd37a6f4de9bd37a7), UWORD(0xb21642c8590b216), /* 23 */
+    UWORD(0x34f72c234f72c235), UWORD(0x8d3dcb08d3dcb08), /* 29 */
+    UWORD(0xef7bdef7bdef7bdf), UWORD(0x842108421084210), /* 31 */
+    UWORD(0x14c1bacf914c1bad), UWORD(0x6eb3e45306eb3e4), /* 37 */
+    UWORD(0x8f9c18f9c18f9c19), UWORD(0x63e7063e7063e70), /* 41 */
+    UWORD(0x82fa0be82fa0be83), UWORD(0x5f417d05f417d05), /* 43 */
+    UWORD(0x51b3bea3677d46cf), UWORD(0x572620ae4c415c9), /* 47 */
+    UWORD(0x21cfb2b78c13521d), UWORD(0x4d4873ecade304d), /* 53 */
+    UWORD(0xcbeea4e1a08ad8f3), UWORD(0x456c797dd49c341), /* 59 */
+    UWORD(0x4fbcda3ac10c9715), UWORD(0x4325c53ef368eb0), /* 61 */
+    UWORD(0xf0b7672a07a44c6b), UWORD(0x3d226357e16ece5), /* 67 */
+    UWORD(0x193d4bb7e327a977), UWORD(0x39b0ad12073615a), /* 71 */
+    UWORD(0x7e3f1f8fc7e3f1f9), UWORD(0x381c0e070381c0e), /* 73 */
+    UWORD(0x9b8b577e613716af), UWORD(0x33d91d2a2067b23), /* 79 */
+    UWORD(0xa3784a062b2e43db), UWORD(0x3159721ed7e7534), /* 83 */
+    UWORD(0xf47e8fd1fa3f47e9), UWORD(0x2e05c0b81702e05), /* 89 */
+    UWORD(0xa3a0fd5c5f02a3a1), UWORD(0x2a3a0fd5c5f02a3), /* 97 */
+    UWORD(0x3a4c0a237c32b16d), UWORD(0x288df0cac5b3f5d), /* 101 */
+    UWORD(0xdab7ec1dd3431b57), UWORD(0x27c45979c95204f), /* 103 */
+    UWORD(0x77a04c8f8d28ac43), UWORD(0x2647c69456217ec), /* 107 */
+    UWORD(0xa6c0964fda6c0965), UWORD(0x2593f69b02593f6), /* 109 */
+    UWORD(0x90fdbc090fdbc091), UWORD(0x243f6f0243f6f02), /* 113 */
+    UWORD(0x7efdfbf7efdfbf7f), UWORD(0x204081020408102), /* 127 */
+    UWORD(0x3e88cb3c9484e2b), UWORD(0x1f44659e4a42715), /* 131 */
+    UWORD(0xe21a291c077975b9), UWORD(0x1de5d6e3f8868a4), /* 137 */
+    UWORD(0x3aef6ca970586723), UWORD(0x1d77b654b82c339), /* 139 */
+    UWORD(0xdf5b0f768ce2cabd), UWORD(0x1b7d6c3dda338b2), /* 149 */
+};
+
+#define TRIAL32(n, ii) (((uint32_t) (n) * trial_inv_32[2 * ii]) <= trial_inv_32[2 * ii + 1])
+#define TRIAL64(n, ii) (((uint64_t) (n) * trial_inv_64[2 * ii]) <= trial_inv_64[2 * ii + 1])
+
 int
 n_is_prime(ulong n)
 {
     if (n < SMALL_ODDPRIME_LIMIT)
         return (n % 2 ? n_is_oddprime_small2(n) : n == 2);
 
-    if (!(n % 2) || !(n % 3) || !(n % 5))
+    /* Trial division by 2, 3 and 5 */
+#if FLINT_BITS == 64
+    if (!(n % 2) || TRIAL64(n, 0) || TRIAL64(n, 1))
+#else
+    if (!(n % 2) || TRIAL32(n, 0) || TRIAL32(n, 1))
+#endif
         return 0;
 
     if (n < (1 << 20))
         return n_is_mod30prime_small(n);
 
-    if (n > (1 << 20) &&
-        (!(n % 7) || !(n % 11) || !(n % 13) || !(n % 17) || !(n % 19) ||
-        !(n % 23) || !(n % 29) || !(n % 31) || !(n % 37) ||
-        !(n % 41) || !(n % 43) || !(n % 47) || !(n % 53)))
-        return 0;
-
-    if (n > UINT32_MAX &&
-        (!(n % 59) || !(n % 61) || !(n % 67) || !(n % 71) || !(n % 73) ||
-         !(n % 79) || !(n % 83) || !(n % 89) || !(n % 97) || !(n %101) ||
-         !(n %103) || !(n % 107) || !(n %109) || !(n %113) || !(n % 127) ||
-         !(n %131) || !(n % 137) || !(n %139) || !(n %149)))
-        return 0;
+    if (n <= UINT32_MAX)
+    {
+        int i, c;
+        /* The compiler should unroll and convert this to a few vector
+           instructions when e.g. AVX2 is available; branchy return isn't
+           worth it. */
+        c = 0;
+        for (i = 2; i < NUM_ODD_TRIAL_PRIMES; i++)
+            c |= TRIAL32(n, i);
+
+        if (c != 0)
+            return 0;
+    }
+    else
+    {
+        int i;
+        for (i = 2; i < NUM_ODD_TRIAL_PRIMES; i++)
+            if (TRIAL64(n, i))
+                return 0;
+    }
 
     return n_is_prime_odd_no_trial(n);
 }
 
+
+
 /* first 64 primes used for modular arithmetic */
 #define N_MODULUS (UWORD(1) << (FLINT_BITS - 1))
 #define N_MOD_TAB 64