include/flint/fft_small.h

0001 /*
0002     Copyright (C) 2022 Daniel Schultz
0003
0004     This file is part of FLINT.
0005
0006     FLINT is free software: you can redistribute it and/or modify it under
0007     the terms of the GNU Lesser General Public License (LGPL) as published
0008     by the Free Software Foundation; either version 3 of the License, or
0009     (at your option) any later version.  See <https://www.gnu.org/licenses/>.
0010 */
0011
0012 #ifndef FFT_SMALL_H
0013 #define FFT_SMALL_H
0014
0015 #include "longlong.h"
0016 #include "machine_vectors.h"
0017
0018 #if FLINT_USES_PTHREAD
0019 # include <pthread.h>
0020 # include <stdatomic.h>
0021 #endif
0022
0023 #define LG_BLK_SZ 8
0024 #define BLK_SZ 256
0025
0026 #ifdef __cplusplus
0027 extern "C" {
0028 #endif
0029
0030 /* Check that a modulus n satisfies the assumptions for mulmod
0031    documented in machine_vectors.h */
0032 int fft_small_mulmod_satisfies_bounds(ulong n);
0033
0034 FLINT_FORCE_INLINE ulong n_pow2(int k)
0035 {
0036     return UWORD(1) << k;
0037 }
0038
0039 FLINT_FORCE_INLINE ulong n_min(ulong a, ulong b)
0040 {
0041     return FLINT_MIN(a, b);
0042 }
0043
0044 FLINT_FORCE_INLINE ulong n_max(ulong a, ulong b)
0045 {
0046     return FLINT_MAX(a, b);
0047 }
0048
0049 FLINT_FORCE_INLINE ulong n_cdiv(ulong a, ulong b)
0050 {
0051     /* not technically correct because the addition can overflow */
0052     return (a + b - 1)/b;
0053 }
0054
0055 FLINT_FORCE_INLINE ulong n_round_up(ulong a, ulong b)
0056 {
0057     return n_cdiv(a, b)*b;
0058 }
0059
0060 FLINT_FORCE_INLINE ulong n_round_down(ulong a, ulong b)
0061 {
0062     return a/b*b;
0063 }
0064
0065 /* 0 -> 0, 1 -> 1, [2,3] -> 3, [4,7] -> 7, [8,15] -> 15, ... */
0066 FLINT_FORCE_INLINE ulong n_next_pow2m1(ulong a)
0067 {
0068     a |= a >> 1;
0069     a |= a >> 2;
0070     a |= a >> 4;
0071     a |= a >> 8;
0072     a |= a >> 16;
0073 #if FLINT64
0074     a |= a >> 32;
0075 #endif
0076     return a;
0077 }
0078
0079 FLINT_FORCE_INLINE ulong n_leading_zeros(ulong x) {
0080     return x == 0 ? FLINT_BITS : flint_clz(x);
0081 }
0082
0083 FLINT_FORCE_INLINE ulong n_trailing_zeros(ulong x) {
0084     return x == 0 ? FLINT_BITS : flint_ctz(x);
0085 }
0086
0087 FLINT_FORCE_INLINE ulong n_nbits(ulong x) {
0088     return FLINT_BIT_COUNT(x);
0089 }
0090
0091 FLINT_FORCE_INLINE ulong n_nbits_nz(ulong x) {
0092     FLINT_ASSERT(x != 0);
0093     return (flint_clz(x)^(FLINT_BITS-1)) + 1;
0094 }
0095
0096 FLINT_FORCE_INLINE ulong n_clog2(ulong x) {
0097     return (x <= 2) ? (x == 2) : FLINT_BITS - flint_clz(x - 1);
0098 }
0099
0100 FLINT_FORCE_INLINE ulong n_flog2(ulong x) {
0101     return (x <= 2) ? (x == 2) : FLINT_BITS - flint_clz(x);
0102 }
0103
0104 FLINT_FORCE_INLINE slong z_min(slong a, slong b) {return FLINT_MIN(a, b);}
0105
0106 FLINT_FORCE_INLINE slong z_max(slong a, slong b) {return FLINT_MAX(a, b);}
0107
0108 int fft_small_mulmod_satisfies_bounds(ulong n);
0109
0110 /*
0111     Let `e(theta) = exp(2*pi*i*theta)`, let `++` denotes concatenation, and `*`
0112     denotes pointwise multiplication.
0113     Consider an infinite sequence `w` of complex numbers defined as follows.
0114
0115         w_0 = {1}
0116         w_i = w_{i-1} ++ e(2^-i) * w_{i-1}
0117
0118     For example
0119
0120         w_1 = {1, -1}
0121         w_2 = {1, -1, i, -i}
0122         w_3 = {1, -1, i, -i, e(1/8), e(5/8), e(3/8), e(7/8)}
0123
0124     Note that the `2i+1`-th element of `w` is the negation of the `2i`-th element,
0125     so we only need to store every other element of `w`. Specific elements of `w`
0126     can be directly accessed with `sd_fft_ctx_w`.
0127
0128     In actual implementation, we work modulo a prime `p` such that `p-1` has
0129     high, but finite, 2-valuation. Most of the calculation works there as well,
0130     except the sequence `w` cannot be infinite.
0131
0132     The twiddle factors are split across SD_FFT_CTX_W2TAB_SIZE tables:
0133
0134         [0] = {e(0)}                                original index 0
0135         [1] = {e(1/4)}                              original index 1
0136         [2] = {e(1/8), e(3/8)}                      original index 2,3
0137         [3] = {e(1/16), e(5/16), e(3/16), e(7/16)}  original index 4,5,6,7
0138         ...
0139
0140     The unallocated ones start out as NULL, and once they are filled in, they
0141     never have to move. This simplifies thread safe enlargement but complicates
0142     random access into the original table. If j is the index into the original
0143     table, the new indices are
0144
0145         [j_bits][j_r]  where j_bits = nbits(j), j_r = j - 2^(j_bits-1)
0146
0147     with the special case j_bits = j_r = 0 for j = 0.
0148     Up to the first SD_FFT_CTX_W2TAB_INIT tables are stored consecutively to ease
0149     the lookup of small indices, which must currently be at least max(4,
0150     LG_BLK_SZ).
0151 */
0152
0153 /* for the fft look up of powers of w */
0154 #define SET_J_BITS_AND_J_R(j_bits, j_r, j) \
0155 do { \
0156     if (j == 0) \
0157     { \
0158         j_bits = 0; \
0159         j_r = 0; \
0160     } \
0161     else \
0162     { \
0163         j_bits = n_nbits_nz(j); \
0164         j_r = j - n_pow2(j_bits - 1); \
0165     } \
0166 } while (0)
0167
0168 /* for the ifft look up of powers of w^-1: the remainder has to be flipped */
0169 #define SET_J_BITS_AND_J_MR(j_bits, j_mr, j) \
0170 do { \
0171     if (j == 0) \
0172     { \
0173         j_bits = 0; \
0174         j_mr = 0; \
0175     } \
0176     else \
0177     { \
0178         j_bits = n_nbits_nz(j); \
0179         j_mr = n_pow2(j_bits) - 1 - j; \
0180     } \
0181 } while (0)
0182
0183
0184 #define SD_FFT_CTX_W2TAB_INIT 12
0185 #define SD_FFT_CTX_W2TAB_SIZE 40
0186
0187 /*
0188     This context is the one expected to sit in a global position.
0189     One container of this context is the thread_local storage for mpn_ctx_t.
0190     See sd_fft_ctx_init_prime for the meaning of each field.
0191 */
0192 typedef struct {
0193     double p;   /* the FFT prime */
0194     double pinv;
0195     nmod_t mod;
0196     ulong primitive_2power_root; /* primitive 2^v-th root, v = valuation(p - 1, 2) */
0197 #if FLINT_USES_PTHREAD
0198     _Atomic(unsigned int) w2tab_depth;
0199 #else
0200     unsigned int w2tab_depth;
0201 #endif
0202     double* w2tab[SD_FFT_CTX_W2TAB_SIZE];
0203     /* see above for the layout of w2tab. Each entry is an integer
0204        in the range [-p/2, p/2]. */
0205 #if FLINT_USES_PTHREAD
0206     pthread_mutex_t mutex;
0207 #endif
0208 } sd_fft_ctx_struct;
0209
0210 typedef sd_fft_ctx_struct sd_fft_ctx_t[1];
0211
0212 FLINT_FORCE_INLINE ulong sd_fft_ctx_blk_offset(ulong I)
0213 {
0214     return I << LG_BLK_SZ;
0215 }
0216
0217 FLINT_FORCE_INLINE ulong sd_fft_ctx_data_size(ulong L)
0218 {
0219     return n_pow2(L);
0220 }
0221
0222 /* Return the address of block `I` in `d`, where each block has `BLK_SZ` entries. */
0223 FLINT_FORCE_INLINE double* sd_fft_ctx_blk_index(double* d, ulong I)
0224 {
0225     return d + sd_fft_ctx_blk_offset(I);
0226 }
0227
0228 /*
0229 location of the bit-reversed eval:
0230     with out_data = fft(in_data) of length 2^L, then
0231     eval_poly(in_data, sd_fft_ctx_w(Q, i)) = out_data[sd_fft_ctx_trunc_index(L, i)]
0232 */
0233 FLINT_FORCE_INLINE ulong sd_fft_ctx_trunc_index(ulong L, ulong i)
0234 {
0235     /* 4x4 transposed blocks in basecases if depth >= 4 */
0236     if (L >= 4)
0237         i = (i&(-16)) | ((i>>2)&3) | ((i&3)<<2);
0238     return i;
0239 }
0240
0241 /* sd_fft.c */
0242 void sd_fft_trunc(sd_fft_ctx_t Q, double* d, ulong L, ulong itrunc, ulong otrunc);
0243
0244 /* sd_ifft.c */
0245 void sd_ifft_trunc(sd_fft_ctx_t Q, double* d, ulong L, ulong trunc);
0246
0247 /* sd_fft_ctx.c */
0248 void sd_fft_ctx_clear(sd_fft_ctx_t Q);
0249 void sd_fft_ctx_init_prime(sd_fft_ctx_t Q, ulong pp);
0250 void sd_fft_ctx_fit_depth_with_lock(sd_fft_ctx_t Q, ulong k);
0251
0252 FLINT_FORCE_INLINE void sd_fft_ctx_fit_depth(sd_fft_ctx_t Q, ulong depth)
0253 {
0254 #if FLINT_USES_PTHREAD
0255     ulong tdepth = (ulong)atomic_load_explicit(&Q->w2tab_depth, memory_order_relaxed);
0256 #else
0257     ulong tdepth = (ulong)Q->w2tab_depth;
0258 #endif
0259     if (FLINT_UNLIKELY(tdepth < depth))
0260         sd_fft_ctx_fit_depth_with_lock(Q, depth);
0261 }
0262
0263 void sd_fft_ctx_point_mul(const sd_fft_ctx_t Q,
0264                             double* a, const double* b, ulong m_, ulong depth);
0265 void sd_fft_ctx_point_sqr(const sd_fft_ctx_t Q,
0266                             double* a, ulong m_, ulong depth);
0267
0268
0269 /*
0270     The bit reversed table is
0271         w = {e(0), e(1/2), e(1/4), e(3/4), e(1/8), e(5/8), e(3/8), e(7/8), ...}
0272     Only the terms of even index are explicitly stored, and they are split
0273     among several tables.
0274     The recursive definition of the infinite sequence `w` can be found above.
0275 */
0276
0277 /* look up w[2*j] */
0278 FLINT_FORCE_INLINE double sd_fft_ctx_w2(const sd_fft_ctx_t Q, ulong j)
0279 {
0280     ulong j_bits, j_r;
0281     SET_J_BITS_AND_J_R(j_bits, j_r, j);
0282     return Q->w2tab[j_bits][j_r];
0283 }
0284
0285 /* look up -w[2*j]^-1 */
0286 FLINT_FORCE_INLINE double sd_fft_ctx_w2inv(const sd_fft_ctx_t Q, ulong j)
0287 {
0288     ulong j_bits, j_mr;
0289     SET_J_BITS_AND_J_MR(j_bits, j_mr, j);
0290     return (j == 0) ? -1.0 : Q->w2tab[j_bits][j_mr];
0291 }
0292
0293 /* look up w[j]. Definition of w above */
0294 FLINT_FORCE_INLINE double sd_fft_ctx_w(const sd_fft_ctx_t Q, ulong j)
0295 {
0296     double r = sd_fft_ctx_w2(Q, j/2);
0297     return (j&1) ? -r : r;
0298 }
0299
0300 typedef struct {
0301     ulong prime;
0302     ulong coeff_len;
0303     ulong nprimes;
0304     ulong* data;
0305 } crt_data_struct;
0306
0307 typedef crt_data_struct crt_data_t[1];
0308
0309 void crt_data_init(crt_data_t C, ulong prime, ulong coeff_len, ulong nprimes);
0310
0311 void crt_data_clear(crt_data_t C);
0312
0313 /* return mpn of length C->coeff_len */
0314 FLINT_FORCE_INLINE ulong* crt_data_co_prime(const crt_data_t C, ulong i)
0315 {
0316     FLINT_ASSERT(i < C->nprimes);
0317     return C->data + i*C->coeff_len;
0318 }
0319
0320 FLINT_FORCE_INLINE ulong* _crt_data_co_prime(const crt_data_t C, ulong i, ulong n)
0321 {
0322     FLINT_ASSERT(i < C->nprimes);
0323     FLINT_ASSERT(n == C->coeff_len);
0324     return C->data + i*n;
0325 }
0326
0327 /* return mpn of length C->coeff_len */
0328 FLINT_FORCE_INLINE ulong* crt_data_prod_primes(const crt_data_t C)
0329 {
0330     return C->data + C->nprimes*C->coeff_len;
0331 }
0332
0333 /* the reduction of co_prime mod the i^th prime */
0334 FLINT_FORCE_INLINE ulong* crt_data_co_prime_red(const crt_data_t C, ulong i)
0335 {
0336     FLINT_ASSERT(i < C->nprimes);
0337     return C->data + C->nprimes*C->coeff_len + C->coeff_len + i;
0338 }
0339
0340
0341 typedef void (*to_ffts_func)(
0342     sd_fft_ctx_struct* Qffts, double* d, ulong dstride,
0343     const ulong* a_, ulong an_, ulong atrunc,
0344     const vec4d* two_pow,
0345     ulong start_easy, ulong stop_easy,
0346     ulong start_hard, ulong stop_hard);
0347
0348 typedef struct {
0349     ulong np;
0350     ulong bits;
0351     ulong bn_bound;
0352     to_ffts_func to_ffts;
0353 } profile_entry_struct;
0354
0355 typedef profile_entry_struct profile_entry_t[1];
0356
0357 #define MPN_CTX_NCRTS 8
0358 #define MAX_NPROFILES 20
0359 #define VEC_SZ 4
0360
0361 /*
0362     The tables for powers of two each have this fixed length. This has to go up
0363     linearly with the max number of primes MPN_CTX_NCRTS involved in chinese
0364     remaindering. This length is checked with asserts in the code.
0365 */
0366 #define MPN_CTX_TWO_POWER_TAB_SIZE 256
0367
0368 typedef struct {
0369     sd_fft_ctx_struct ffts[MPN_CTX_NCRTS];
0370     crt_data_struct crts[MPN_CTX_NCRTS];
0371
0372     /*
0373         For each table of tables of powers of two, the whole collection is held
0374         in one big buffer and the table is an array of pointer into it.
0375     */
0376     vec4d* vec_two_pow_tab[(MPN_CTX_NCRTS + VEC_SZ - 1)/VEC_SZ];
0377     vec4d* vec_two_pow_buffer;
0378     double* slow_two_pow_tab[MPN_CTX_NCRTS];
0379     double* slow_two_pow_buffer;
0380
0381     profile_entry_struct profiles[MAX_NPROFILES];
0382     ulong profiles_size;
0383     void* buffer;
0384     ulong buffer_alloc;
0385 } mpn_ctx_struct;
0386
0387 typedef mpn_ctx_struct mpn_ctx_t[1];
0388
0389 void _convert_block(ulong* Xs, sd_fft_ctx_struct* Rffts, double* d, ulong dstride, ulong np, ulong Iv);
0390 void _convert_block_1_mod(ulong* Xs, sd_fft_ctx_struct* Rffts, double* d, ulong dstride, ulong I, double p2, double p2inv);
0391 ulong flint_mpn_nbits(const ulong* a, ulong an);
0392 int flint_mpn_cmp_ui_2exp(const ulong* a, ulong an, ulong b, ulong e);
0393 unsigned char flint_mpn_add_inplace_c(ulong* z, ulong zn, ulong* a, ulong an, unsigned char cf);
0394
0395
0396 void mpn_ctx_init(mpn_ctx_t R, ulong p);
0397 void mpn_ctx_clear(mpn_ctx_t R);
0398 void* mpn_ctx_fit_buffer(mpn_ctx_t R, ulong n);
0399 void mpn_ctx_mpn_mul(mpn_ctx_t R, ulong* z, const ulong* a, ulong an, const ulong* b, ulong bn);
0400
0401 void _nmod_poly_mul_mid_mpn_ctx(
0402     ulong* z, ulong zl, ulong zh,
0403     const ulong* a, ulong an,
0404     const ulong* b, ulong bn,
0405     nmod_t mod,
0406     mpn_ctx_t R);
0407
0408 void _nmod_poly_divrem_mpn_ctx(
0409     ulong* q,
0410     ulong* r,
0411     const ulong* a, ulong an,
0412     const ulong* b, ulong bn,
0413     nmod_t mod,
0414     mpn_ctx_t R);
0415
0416 void _nmod_poly_mul_mid_classical(
0417     ulong* z, slong zl, slong zh,
0418     const ulong* a, slong an,
0419     const ulong* b, slong bn,
0420     nmod_t mod);
0421
0422 void _nmod_poly_mul_mid(
0423     ulong* z, slong zl, slong zh,
0424     const ulong* a, slong an,
0425     const ulong* b, slong bn,
0426     nmod_t mod);
0427
0428 typedef struct {
0429     ulong depth;
0430     ulong N;
0431     ulong offset;
0432     ulong np;
0433     ulong stride;
0434     ulong bn;
0435     ulong btrunc;
0436     double* bbuf;
0437 } mul_precomp_struct;
0438
0439 void _mul_precomp_init(
0440     mul_precomp_struct* M,
0441     const ulong * b, ulong bn, ulong btrunc,
0442     ulong depth,
0443     nmod_t mod,
0444     mpn_ctx_t R);
0445
0446 FLINT_FORCE_INLINE void _mul_precomp_clear(mul_precomp_struct* M)
0447 {
0448     flint_aligned_free(M->bbuf);
0449 }
0450
0451 int _nmod_poly_mul_mid_precomp(
0452     ulong* z, ulong zl, ulong zh,
0453     const ulong* a, ulong an,
0454     mul_precomp_struct* M,
0455     nmod_t mod,
0456     mpn_ctx_t R);
0457
0458 typedef struct {
0459     mul_precomp_struct quo_maker[1];
0460     mul_precomp_struct rem_maker[1];
0461 } nmod_poly_divrem_precomp_struct;
0462
0463 FLINT_FORCE_INLINE void _nmod_poly_divrem_precomp_clear(nmod_poly_divrem_precomp_struct* M)
0464 {
0465     _mul_precomp_clear(M->quo_maker);
0466     _mul_precomp_clear(M->rem_maker);
0467 }
0468
0469 void _nmod_poly_divrem_precomp_init(
0470     nmod_poly_divrem_precomp_struct* M,
0471     const ulong* b, ulong bn,
0472     ulong Bn,
0473     nmod_t mod,
0474     mpn_ctx_t R);
0475
0476
0477 int _nmod_poly_divrem_precomp(
0478     ulong* q,
0479     ulong* r,
0480     const ulong* a, ulong an,
0481     nmod_poly_divrem_precomp_struct* M,
0482     nmod_t mod,
0483     mpn_ctx_t R);
0484
0485 mpn_ctx_struct * get_default_mpn_ctx(void);
0486
0487 void mpn_mul_default_mpn_ctx(nn_ptr r1, nn_srcptr i1, slong n1, nn_srcptr i2, slong n2);
0488 void _nmod_poly_mul_mid_default_mpn_ctx(nn_ptr res, slong zl, slong zh, nn_srcptr a, slong an, nn_srcptr b, slong bn, nmod_t mod);
0489
0490
0491 int _fmpz_poly_mul_mid_mpn_ctx(
0492     fmpz * z, ulong zl, ulong zh,
0493     const fmpz * a, ulong an,
0494     const fmpz * b, ulong bn,
0495     mpn_ctx_t R);
0496
0497 int _fmpz_poly_mul_mid_default_mpn_ctx(
0498     fmpz * z, slong zl, slong zh,
0499     const fmpz * a, slong an,
0500     const fmpz * b, slong bn);
0501
0502 #ifdef __cplusplus
0503 }
0504 #endif
0505
0506 #endif