improve dft init, recursive fft for radix-4

Author: adamant-pwn

cp-algo/math/cvector.hpp

@@ -112,10 +112,23 @@ namespace cp_algo::math::fft {
 
         void ifft() {
             size_t n = size();
-            for(size_t i = flen; i <= n / 2; i *= 2) {
-                if (4 * i <= n) { // radix-4
-                    exec_on_evals<4>(n / (4 * i), [&](size_t k, point rt) {
-                        k *= 4 * i;
+            bool parity = std::countr_zero(n) % 2;
+            if(parity) {
+                exec_on_evals<2>(n / (2 * flen), [&](size_t k, point rt) {
+                    k *= 2 * flen;
+                    vpoint cvrt = {vz + real(rt), vz - imag(rt)};
+                    auto B = at(k) - at(k + flen);
+                    at(k) += at(k + flen);
+                    at(k + flen) = B * cvrt;
+                });
+            }
+
+            for(size_t leaf = 3 * flen; leaf < n; leaf += 4 * flen) {
+                size_t level = std::countr_one(leaf + 3);
+                for(size_t lvl = 4 + parity; lvl <= level; lvl += 2) {
+                    size_t i = (1 << lvl) / 4;
+                    exec_on_eval<4>(n >> lvl, leaf >> lvl, [&](size_t k, point rt) {
+                        k <<= lvl;
                         vpoint v1 = {vz + real(rt), vz - imag(rt)};
                         vpoint v2 = v1 * v1;
                         vpoint v3 = v1 * v2;
@@ -124,21 +137,10 @@ namespace cp_algo::math::fft {
                             auto B = at(j + i);
                             auto C = at(j + 2 * i);
                             auto D = at(j + 3 * i);
-                            at(j) = (A + B + C + D);
-                            at(j + 2 * i) = (A + B - C - D) * v2;
-                            at(j +     i) = (A - B - vi(C - D)) * v1;
-                            at(j + 3 * i) = (A - B + vi(C - D)) * v3;
-                        }
-                    });
-                    i *= 2;
-                } else { // radix-2 fallback
-                    exec_on_evals<2>(n / (2 * i), [&](size_t k, point rt) {
-                        k *= 2 * i;
-                        vpoint cvrt = {vz + real(rt), vz - imag(rt)};
-                        for(size_t j = k; j < k + i; j += flen) {
-                            auto B = at(j) - at(j + i);
-                            at(j) += at(j + i);
-                            at(j + i) = B * cvrt;
+                            at(j) = ((A + B) + (C + D));
+                            at(j + 2 * i) = ((A + B) - (C + D)) * v2;
+                            at(j +     i) = ((A - B) - vi(C - D)) * v1;
+                            at(j + 3 * i) = ((A - B) + vi(C - D)) * v3;
                         }
                     });
                 }
@@ -150,11 +152,14 @@ namespace cp_algo::math::fft {
         }
         void fft() {
             size_t n = size();
-            for(size_t i = n / 2; i >= flen; i /= 2) {
-                if (i / 2 >= flen) { // radix-4
-                    i /= 2;
-                    exec_on_evals<4>(n / (4 * i), [&](size_t k, point rt) {
-                        k *= 4 * i;
+            bool parity = std::countr_zero(n) % 2;
+            for(size_t leaf = 0; leaf < n; leaf += 4 * flen) {
+                size_t level = std::countr_zero(n + leaf);
+                level -= level % 2 != parity;
+                for(size_t lvl = level; lvl >= 4; lvl -= 2) {
+                    size_t i = (1 << lvl) / 4;
+                    exec_on_eval<4>(n >> lvl, leaf >> lvl, [&](size_t k, point rt) {
+                        k <<= lvl;
                         vpoint v1 = {vz + real(rt), vz + imag(rt)};
                         vpoint v2 = v1 * v1;
                         vpoint v3 = v1 * v2;
@@ -169,18 +174,17 @@ namespace cp_algo::math::fft {
                             at(j + 3 * i) = (A - C) - vi(B - D);
                         }
                     });
-                } else { // radix-2 fallback
-                    exec_on_evals<2>(n / (2 * i), [&](size_t k, point rt) {
-                        k *= 2 * i;
-                        vpoint vrt = {vz + real(rt), vz + imag(rt)};
-                        for(size_t j = k; j < k + i; j += flen) {
-                            auto t = at(j + i) * vrt;
-                            at(j + i) = at(j) - t;
-                            at(j) += t;
-                        }
-                    });
                 }
             }
+            if(parity) {
+                exec_on_evals<2>(n / (2 * flen), [&](size_t k, point rt) {
+                    k *= 2 * flen;
+                    vpoint vrt = {vz + real(rt), vz + imag(rt)};
+                    auto t = at(k + flen) * vrt;
+                    at(k + flen) = at(k) - t;
+                    at(k) += t;
+                });
+            }
             checkpoint("fft");
         }
         static constexpr size_t pre_evals = 1 << 16;

cp-algo/math/fft.hpp

@@ -13,13 +13,15 @@ namespace cp_algo::math::fft {
         static base factor, ifactor;
         using Int2 = base::Int2;
         static bool _init;
-        static int split;
+        static int split() {
+            static const int splt = int(std::sqrt(base::mod())) + 1;
+            return splt;
+        }
         static u64x4 mod, imod;
 
         void init() {
             if(!_init) {
                 factor = 1 + random::rng() % (base::mod() - 1);
-                split = int(std::sqrt(base::mod())) + 1;
                 ifactor = base(1) / factor;
                 mod = u64x4() + base::mod();
                 imod = u64x4() + inv2(-base::mod());
@@ -51,9 +53,9 @@ namespace cp_algo::math::fft {
                     };
                     au = montgomery_mul(au, mul, mod, imod);
                     au = au >= base::mod() ? au - base::mod() : au;
-                    auto ai = i64x4(au);
-                    ai = ai >= base::mod() / 2 ? ai - base::mod() : ai;
-                    return std::pair{to_double(ai % split), to_double(ai / split)};
+                    auto ai = to_double(i64x4(au >= base::mod() / 2 ? au - base::mod() : au));
+                    auto quo = round(ai / split());
+                    return std::pair{ai - quo * split(), quo};
                 };
                 auto [rai, qai] = splt(i, cur);
                 auto [rani, qani] = splt(n + i, montgomery_mul(cur, stepn, mod, imod));
@@ -101,7 +103,7 @@ namespace cp_algo::math::fft {
         void recover_mod(auto &&C, auto &res, size_t k) {
             res.assign((k / flen + 1) * flen, base(0));
             size_t n = A.size();
-            auto splitsplit = base(split * split).getr();
+            auto const splitsplit = base(split() * split()).getr();
             base b2x32 = bpow(base(2), 32);
             base b2x64 = bpow(base(2), 64);
             u64x4 cur = {
@@ -118,7 +120,7 @@ namespace cp_algo::math::fft {
                 auto [Cx, Cy] = C.at(i);
                 auto set_i = [&](size_t i, auto A, auto B, auto C, auto mul) {
                     auto A0 = lround(A), A1 = lround(C), A2 = lround(B);
-                    auto Ai = A0 + A1 * split + A2 * splitsplit + uint64_t(base::modmod());
+                    auto Ai = A0 + A1 * split() + A2 * splitsplit + uint64_t(base::modmod());
                     auto Au = montgomery_reduce(u64x4(Ai), mod, imod);
                     Au = montgomery_mul(Au, mul, mod, imod);
                     Au = Au >= base::mod() ? Au - base::mod() : Au;
@@ -174,7 +176,6 @@ namespace cp_algo::math::fft {
     template<modint_type base> base dft<base>::factor = 1;
     template<modint_type base> base dft<base>::ifactor = 1;
     template<modint_type base> bool dft<base>::_init = false;
-    template<modint_type base> int dft<base>::split = 1;
     template<modint_type base> u64x4 dft<base>::mod = {};
     template<modint_type base> u64x4 dft<base>::imod = {};