[Mlir-commits] [mlir] Upstream polynomial.ntt and polynomial.intt (PR #90992)

[Oleksandr Alex Zinenko]

================
@@ -104,3 +104,80 @@ LogicalResult MulScalarOp::verify() {
 
   return success();
 }
+
+// Test if a value is a primitive nth root of unity modulo cmod
+bool isPrimitiveNthRootOfUnity(const APInt &root, const unsigned n,
+                               const APInt &cmod) {
+  // root bitwidth may be 1 less then cmod
+  APInt r = APInt(root).zext(cmod.getBitWidth());
+  assert(r.ule(cmod) && "root must be less than cmod");
+
+  APInt a = r;
+  for (size_t k = 1; k < n; k++) {
+    if (a.isOne())
+      return false;
+    a = (a * r).urem(cmod);
+  }
+  return a.isOne();
+}
+
----------------
ftynse wrote:

Nit: please document top-level functions.

https://github.com/llvm/llvm-project/pull/90992