use u32x8 in many_facts

Author: adamant-pwn

cp-algo/util/simd.hpp

@@ -56,29 +56,13 @@ namespace cp_algo {
 #endif
     }
 
-    u32x8 montgomery_mul(u32x8 x, u32x8 y, uint32_t mod, uint32_t imod) {
+    [[gnu::always_inline]] inline u32x8 montgomery_mul(u32x8 x, u32x8 y, uint32_t mod, uint32_t imod) {
         auto x0246 = u64x4(x) & uint32_t(-1);
         auto y0246 = u64x4(y) & uint32_t(-1);
         auto x1357 = u64x4(x) >> 32;
         auto y1357 = u64x4(y) >> 32;
-#ifdef __AVX2__
-        auto xy0246 = u64x4(_mm256_mul_epu32(__m256i(x0246), __m256i(y0246)));
-        auto xy1357 = u64x4(_mm256_mul_epu32(__m256i(x1357), __m256i(y1357)));
-#else
-        u64x4 xy0246 = x0246 * y0246;
-        u64x4 xy1357 = x1357 * y1357;
-#endif
-        auto xy_inv = u64x4(u32x8(xy0246 | (xy1357 << 32)) * (u32x8() + imod));
-        auto xy_inv0246 = xy_inv & uint32_t(-1);
-        auto xy_inv1357 = xy_inv >> 32;
-#ifdef __AVX2__
-        xy0246 += u64x4(_mm256_mul_epu32(__m256i(xy_inv0246), __m256i() + mod));
-        xy1357 += u64x4(_mm256_mul_epu32(__m256i(xy_inv1357), __m256i() + mod));
-#else
-        xy0246 += xy_inv0246 * mod;
-        xy1357 += xy_inv1357 * mod;
-#endif
-        return u32x8((xy0246 >> 32) | (xy1357 & -1ULL << 32));
+        return u32x8(montgomery_mul(x0246, y0246, mod, imod)) |
+               u32x8(montgomery_mul(x1357, y1357, mod, imod) << 32);
     }
 
     [[gnu::always_inline]] inline dx4 rotate_right(dx4 x) {

verify/simd/many_facts.test.cpp

@@ -10,8 +10,7 @@ using namespace std;
 using namespace cp_algo;
 
 constexpr int mod = 998244353;
-constexpr auto mod4 = u64x4() + mod;
-constexpr auto imod4 = u64x4() - math::inv2(mod);
+constexpr int imod = -math::inv2(mod);
 
 void facts_inplace(vector<int> &args) {
     constexpr int block = 1 << 16;
@@ -26,39 +25,40 @@ void facts_inplace(vector<int> &args) {
             args_per_block[(mod - x - 1) / block].push_back(i);
         }
     }
-    uint64_t b2x32 = (1ULL << 32) % mod;
+    uint32_t b2x32 = (1ULL << 32) % mod;
     uint64_t fact = 1;
-    const int K = 4;
-    for(uint64_t b = 0; b <= limit; b += K * block) {
-        u64x4 cur[K];
-        static array<u64x4, block / 4> prods[K];
-        for(int z = 0; z < K; z++) {
-            for(int j = 0; j < 4; j++) {
-                cur[z][j] = b + z * block + j * block / 4;
+    const int accum = 4;
+    const int simd_size = 8;
+    for(uint64_t b = 0; b <= limit; b += accum * block) {
+        u32x8 cur[accum];
+        static array<u32x8, block / simd_size> prods[accum];
+        for(int z = 0; z < accum; z++) {
+            for(int j = 0; j < simd_size; j++) {
+                cur[z][j] = uint32_t(b + z * block + j * (block / simd_size));
                 prods[z][0][j] = cur[z][j] + !(b || z || j);
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
-                cur[z][j] = cur[z][j] * b2x32 % mod;
+                cur[z][j] = uint32_t(uint64_t(cur[z][j]) * b2x32 % mod);
 #pragma GCC diagnostic pop
             }
         }
-        for(int i = 1; i < block / 4; i++) {
-            for(int z = 0; z < K; z++) {
+        for(int i = 1; i < block / simd_size; i++) {
+            for(int z = 0; z < accum; z++) {
                 cur[z] += b2x32;
                 cur[z] = cur[z] >= mod ? cur[z] - mod : cur[z];
-                prods[z][i] = montgomery_mul(prods[z][i - 1], cur[z], mod4, imod4);
+                prods[z][i] = montgomery_mul(prods[z][i - 1], cur[z], mod, imod);
             }
         }
-        for(int z = 0; z < K; z++) {
+        for(int z = 0; z < accum; z++) {
             uint64_t bl = b + z * block;
             for(auto i: args_per_block[bl / block]) {
                 size_t x = args[i];
                 if(x >= mod / 2) {
                     x = mod - x - 1;
                 }
                 x -= bl;
-                auto pre_blocks = x / (block / 4);
-                auto in_block = x % (block / 4);
+                auto pre_blocks = x / (block / simd_size);
+                auto in_block = x % (block / simd_size);
                 auto ans = fact * prods[z][in_block][pre_blocks] % mod;
                 for(size_t j = 0; j < pre_blocks; j++) {
                     ans = ans * prods[z].back()[j] % mod;
@@ -71,7 +71,7 @@ void facts_inplace(vector<int> &args) {
                 }
             }
             args_per_block[bl / block].clear();
-            for(int j = 0; j < 4; j++) {
+            for(int j = 0; j < simd_size; j++) {
                 fact = fact * prods[z].back()[j] % mod;
             }
         }