checkpoints + better montgomery simd

Author: adamant-pwn

cp-algo/util/simd.hpp

@@ -38,12 +38,16 @@ namespace cp_algo {
         };
     }
 
+    [[gnu::always_inline]] inline u64x4 low32(u64x4 x) {
+        return x & uint32_t(-1);
+    }
+
     [[gnu::always_inline]] inline u64x4 montgomery_reduce(u64x4 x, uint32_t mod, uint32_t imod) {
         auto x_ninv = u64x4(u32x8(x) * (u32x8() + imod));
 #ifdef __AVX2__
         x += u64x4(_mm256_mul_epu32(__m256i(x_ninv), __m256i() + mod));
 #else
-        x += x_ninv * mod;
+        x += low32(x_ninv) * mod;
 #endif
         return x >> 32;
     }
@@ -52,13 +56,13 @@ namespace cp_algo {
 #ifdef __AVX2__
         return montgomery_reduce(u64x4(_mm256_mul_epu32(__m256i(x), __m256i(y))), mod, imod);
 #else
-        return montgomery_reduce(x * y, mod, imod);
+        return montgomery_reduce(low32(x) * low32(y), mod, imod);
 #endif
     }
 
     [[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 x0246 = u64x4(x);
+        auto y0246 = u64x4(y);
         auto x1357 = u64x4(x) >> 32;
         auto y1357 = u64x4(y) >> 32;
         return u32x8(montgomery_mul(x0246, y0246, mod, imod)) |

verify/simd/many_facts.test.cpp

@@ -2,6 +2,8 @@
 #define PROBLEM "https://judge.yosupo.jp/problem/many_factorials"
 #pragma GCC optimize("Ofast,unroll-loops")
 #include <bits/stdc++.h>
+#define CP_ALGO_CHECKPOINT
+#include "cp-algo/util/checkpoint.hpp"
 #include "blazingio/blazingio.min.hpp"
 #include "cp-algo/util/simd.hpp"
 #include "cp-algo/math/common.hpp"
@@ -25,6 +27,7 @@ void facts_inplace(vector<int> &args) {
             args_per_block[(mod - x - 1) / block].push_back(i);
         }
     }
+    cp_algo::checkpoint("init");
     uint32_t b2x32 = (1ULL << 32) % mod;
     uint64_t fact = 1;
     const int accum = 4;
@@ -49,6 +52,7 @@ void facts_inplace(vector<int> &args) {
                 prods[z][i] = montgomery_mul(prods[z][i - 1], cur[z], mod, imod);
             }
         }
+        cp_algo::checkpoint("inner loop");
         for(int z = 0; z < accum; z++) {
             uint64_t bl = b + z * block;
             for(auto i: args_per_block[bl / block]) {
@@ -75,6 +79,7 @@ void facts_inplace(vector<int> &args) {
                 fact = fact * prods[z].back()[j] % mod;
             }
         }
+        cp_algo::checkpoint("write ans");
     }
 }
 
@@ -83,8 +88,11 @@ void solve() {
     cin >> n;
     vector<int> args(n);
     for(auto &x : args) {cin >> x;}
+    cp_algo::checkpoint("input read");
     facts_inplace(args);
     for(auto it: args) {cout << it << "\n";}
+    cp_algo::checkpoint("output written");
+    cp_algo::checkpoint<1>();
 }
 
 signed main() {