checkpoints + better montgomery simd

adamant-pwn · adamant-pwn · commit 6ef98736fd2f · 2025-05-26T19:54:59.000+02:00
diff --git a/cp-algo/util/simd.hpp b/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)) |
diff --git a/verify/simd/many_facts.test.cpp b/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() {

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,8 @@`
`2`	`2`	`#define PROBLEM "https://judge.yosupo.jp/problem/many_factorials"`
`3`	`3`	`#pragma GCC optimize("Ofast,unroll-loops")`
`4`	`4`	`#include <bits/stdc++.h>`
	`5`	`+#define CP_ALGO_CHECKPOINT`
	`6`	`+#include "cp-algo/util/checkpoint.hpp"`
`5`	`7`	`#include "blazingio/blazingio.min.hpp"`
`6`	`8`	`#include "cp-algo/util/simd.hpp"`
`7`	`9`	`#include "cp-algo/math/common.hpp"`
`@@ -25,6 +27,7 @@ void facts_inplace(vector<int> &args) {`
`25`	`27`	`args_per_block[(mod - x - 1) / block].push_back(i);`
`26`	`28`	`}`
`27`	`29`	`}`
	`30`	`+ cp_algo::checkpoint("init");`
`28`	`31`	`uint32_t b2x32 = (1ULL << 32) % mod;`
`29`	`32`	`uint64_t fact = 1;`
`30`	`33`	`const int accum = 4;`
`@@ -49,6 +52,7 @@ void facts_inplace(vector<int> &args) {`
`49`	`52`	`prods[z][i] = montgomery_mul(prods[z][i - 1], cur[z], mod, imod);`
`50`	`53`	`}`
`51`	`54`	`}`
	`55`	`+ cp_algo::checkpoint("inner loop");`
`52`	`56`	`for(int z = 0; z < accum; z++) {`
`53`	`57`	`uint64_t bl = b + z * block;`
`54`	`58`	`for(auto i: args_per_block[bl / block]) {`
`@@ -75,6 +79,7 @@ void facts_inplace(vector<int> &args) {`
`75`	`79`	`fact = fact * prods[z].back()[j] % mod;`
`76`	`80`	`}`
`77`	`81`	`}`
	`82`	`+ cp_algo::checkpoint("write ans");`
`78`	`83`	`}`
`79`	`84`	`}`
`80`	`85`
`@@ -83,8 +88,11 @@ void solve() {`
`83`	`88`	`cin >> n;`
`84`	`89`	`vector<int> args(n);`
`85`	`90`	`for(auto &x : args) {cin >> x;}`
	`91`	`+ cp_algo::checkpoint("input read");`
`86`	`92`	`facts_inplace(args);`
`87`	`93`	`for(auto it: args) {cout << it << "\n";}`
	`94`	`+ cp_algo::checkpoint("output written");`
	`95`	`+ cp_algo::checkpoint<1>();`
`88`	`96`	`}`
`89`	`97`
`90`	`98`	`signed main() {`