fix primitive root overflow bug

PiperOrigin-RevId: 891818905

Author: j2kun

lib/Utils/APIntUtils.cpp

@@ -69,9 +69,9 @@ APInt modularExponentiation(const APInt& base, const APInt& exponent,
 
   while (e.ugt(0)) {
     if (e[0]) {
-      res = (res * b).urem(modulus);
+      res = modularMultiplication(res, b, modulus);
     }
-    b = (b * b).urem(modulus);
+    b = modularMultiplication(b, b, modulus);
     e = e.lshr(1);
   }
   return res;
@@ -100,7 +100,7 @@ bool isPrime(const APInt& n) {
     if (x.isOne() || x == n - 1) continue;
     bool composite = true;
     for (unsigned r = 1; r < s; ++r) {
-      x = (x * x).urem(n);
+      x = modularMultiplication(x, x, n);
       if (x == n - 1) {
         composite = false;
         break;
@@ -115,8 +115,12 @@ std::vector<APInt> factorize(APInt n) {
   std::vector<APInt> factors;
   if (n.ult(2)) return factors;
 
-  APInt d(n.getBitWidth(), 2);
-  while ((d * d).ule(n)) {
+  unsigned width = n.getBitWidth();
+  APInt d(width, 2);
+  while (true) {
+    APInt wide_d = d.zext(width * 2);
+    if ((wide_d * wide_d).ugt(n.zext(width * 2))) break;
+
     if (n.urem(d).isZero()) {
       factors.push_back(d);
       while (n.urem(d).isZero()) {

lib/Utils/APIntUtils.h

@@ -3,13 +3,23 @@
 
 #include <vector>
 
+#include "llvm/include/llvm/ADT/APInt.h"     // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
 
 namespace mlir {
 namespace heir {
 
 APInt multiplicativeInverse(const APInt& x, const APInt& modulo);
 
+inline APInt modularMultiplication(const APInt& a, const APInt& b,
+                                   const APInt& modulus) {
+  unsigned width = modulus.getBitWidth();
+  APInt wide_a = a.zext(width * 2);
+  APInt wide_b = b.zext(width * 2);
+  APInt wide_m = modulus.zext(width * 2);
+  return (wide_a * wide_b).urem(wide_m).trunc(width);
+}
+
 APInt modularExponentiation(const APInt& base, const APInt& exponent,
                             const APInt& modulus);
 

lib/Utils/MathUtilsTest.cpp

@@ -57,6 +57,22 @@ TEST(MathUtilsTest, FindPrimitive2nthRoot) {
   }
 }
 
+// A regression test for a bug
+TEST(MathUtilsTest, FindPrimitive2nthRootZeroRegression) {
+  // q = 1095233372161, n = 1024
+  auto root = findPrimitive2nthRoot(APInt(64, 1095233372161), 1024);
+  ASSERT_TRUE(root.has_value());
+  APInt q(64, 1095233372161);
+  APInt two_n_ap(64, 2048);
+  EXPECT_EQ(modularExponentiation(*root, two_n_ap, q), 1);
+
+  // Check it's PRIMITIVE 2n-th root
+  auto factors = factorize(two_n_ap);
+  for (const auto& p : factors) {
+    EXPECT_NE(modularExponentiation(*root, two_n_ap.udiv(p), q), 1);
+  }
+}
+
 void Primitive2nthRootProperty(uint64_t q_val, uint64_t n) {
   APInt q(64, q_val);
   if (!isPrime(q)) return;
@@ -78,11 +94,11 @@ void Primitive2nthRootProperty(uint64_t q_val, uint64_t n) {
 
 // Fuzz test with a set of known NTT-friendly primes and various degrees
 std::vector<uint64_t> ntt_primes = {
-    65537,      114689,     147457,     163841,     557057,
-    638977,     737281,     786433,     1032193,    1179649,
-    1769473,    1785857,    2277377,    2424833,    2572289,
-    2654209,    2752513,    2768897,    8380417,    2147565569,
-    2148155393, 2148384769, 3221225473, 3221241857, 3758161921};
+    65537,      114689,       147457,     163841,     557057,     638977,
+    737281,     786433,       1032193,    1179649,    1769473,    1785857,
+    2277377,    2424833,      2572289,    2654209,    2752513,    2768897,
+    8380417,    2147565569,   2148155393, 2148384769, 3221225473, 3221241857,
+    3758161921, 1095233372161};
 
 FUZZ_TEST(MathUtilsTest, Primitive2nthRootProperty)
     .WithDomains(fuzztest::ElementOf(ntt_primes),