Merge pull request #19 from a41-official/feat/add-root-of-unity

batzor · web-flow · commit 238a5dbdb49c · 2025-05-14T15:09:40.000+09:00
feat: add root of unity attribute
diff --git a/benchmark/field/mul_benchmark_test.cc b/benchmark/field/mul_benchmark_test.cc
@@ -5,7 +5,7 @@
 namespace zkir {
 namespace {
 
-using ::zkir::benchmark::Memref;
+using benchmark::Memref;
 
 struct i256 {
   uint64_t limbs[4];  // 4 x 64 = 256 bits
diff --git a/benchmark/ntt/ntt_benchmark.mlir b/benchmark/ntt/ntt_benchmark.mlir
@@ -3,11 +3,12 @@
 !memref_ty = memref<1048576x!coeff_ty>
 
 #root_elem = #field.pf_elem<17220337697351015657950521176323262483320249231368149235373741788599650842711:i256> : !coeff_ty
-#root = #poly.primitive_root<root=#root_elem, degree=1048576:i256>
+#root_of_unity = #field.root_of_unity<#root_elem, 1048576:i256>
+#root = #poly.primitive_root<root_of_unity=#root_of_unity>
 
 !mod = !mod_arith.int<21888242871839275222246405745257275088548364400416034343698204186575808495617 : i256>
 #mont = #mod_arith.montgomery<!mod>
-#root_mont = #poly.primitive_root<root=#root_elem, degree=1048576:i256, montgomery=#mont>
+#root_mont = #poly.primitive_root<root_of_unity=#root_of_unity, montgomery=#mont>
 
 func.func @ntt(%arg0 : !memref_ty) attributes { llvm.emit_c_interface } {
   %t = bufferization.to_tensor %arg0 restrict writable : !memref_ty to !coefft_ty
diff --git a/benchmark/ntt/ntt_benchmark_test.cc b/benchmark/ntt/ntt_benchmark_test.cc
@@ -10,7 +10,7 @@
 namespace zkir {
 namespace {
 
-using ::zkir::benchmark::Memref;
+using benchmark::Memref;
 
 using i256 = benchmark::BigInt<4>;
 
diff --git a/tests/Dialect/Field/prime_field_to_mod_arith.mlir b/tests/Dialect/Field/prime_field_to_mod_arith.mlir
@@ -1,7 +1,8 @@
 // RUN: zkir-opt -prime-field-to-mod-arith --split-input-file %s | FileCheck %s --enable-var-scope
 !PF1 = !field.pf<3:i32>
 !PFv = tensor<4x!PF1>
-#elem = #field.pf_elem<31:i32> : !PF1
+#root_elem = #field.pf_elem<2:i32> : !PF1
+#root = #field.root_of_unity<#root_elem, 2>
 
 !mod = !mod_arith.int<3 : i32>
 #mont = #mod_arith.montgomery<!mod>
diff --git a/tests/Dialect/Poly/poly_canonicalization.mlir b/tests/Dialect/Poly/poly_canonicalization.mlir
@@ -2,7 +2,8 @@
 
 !coeff_ty = !field.pf<7681:i32>
 #elem = #field.pf_elem<3383:i32>  : !coeff_ty
-#root = #poly.primitive_root<root=#elem, degree=4 :i32>
+#root_of_unity = #field.root_of_unity<#elem, 4:i32>
+#root = #poly.primitive_root<root_of_unity=#root_of_unity>
 !poly_ty = !poly.polynomial<!coeff_ty, 3>
 !tensor_ty = tensor<4x!coeff_ty>
 
diff --git a/tests/Dialect/Poly/poly_ntt_runner.mlir b/tests/Dialect/Poly/poly_ntt_runner.mlir
@@ -6,7 +6,8 @@
 !coeff_ty = !field.pf<7681:i32>
 #elem = #field.pf_elem<3383:i32>  : !coeff_ty
 #inv_elem = #field.pf_elem<4298:i32>  : !coeff_ty
-#root = #poly.primitive_root<root=#elem, degree=4 :i32>
+#root_of_unity = #field.root_of_unity<#elem, 4:i32>
+#root = #poly.primitive_root<root_of_unity=#root_of_unity>
 !poly_ty = !poly.polynomial<!coeff_ty, 3>
 
 func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface }
diff --git a/tests/Dialect/Poly/poly_syntax.mlir b/tests/Dialect/Poly/poly_syntax.mlir
@@ -6,7 +6,8 @@
 !poly_ty2 = !poly.polynomial<!PF2, 32>
 #uni_poly = #poly.univariate_polynomial<x**6 + 1> : !poly_ty2
 #elem = #field.pf_elem<2:i32>  : !PF1
-#root = #poly.primitive_root<root=#elem, degree=3>
+#root_of_unity = #field.root_of_unity<#elem, 3:i32>
+#root = #poly.primitive_root<root_of_unity=#root_of_unity>
 
 // CHECK-LABEL: @test_poly_syntax
 func.func @test_poly_syntax() {
diff --git a/tests/Dialect/Poly/poly_to_field.mlir b/tests/Dialect/Poly/poly_to_field.mlir
@@ -3,8 +3,9 @@
 !PF1 = !field.pf<7:i255>
 !poly_ty1 = !poly.polynomial<!PF1, 3>
 !poly_ty2 = !poly.polynomial<!PF1, 4>
-#elem = #field.pf_elem<2:i255>  : !PF1
-#root = #poly.primitive_root<root=#elem, degree=4>
+#elem = #field.pf_elem<6:i255>  : !PF1
+#root_of_unity = #field.root_of_unity<#elem, 2:i255>
+#root = #poly.primitive_root<root_of_unity=#root_of_unity>
 
 // FIXME(batzor): without this line, the test will fail with the following error:
 // LLVM ERROR: can't create Attribute 'mlir::polynomial::IntPolynomialAttr' because storage uniquer isn't initialized: the dialect was likely not loaded, or the attribute wasn't added with addAttributes<...>() in the Dialect::initialize() method.
@@ -68,16 +69,16 @@ func.func @test_lower_from_tensor(%t : tensor<4x!PF1>) -> !poly_ty1 {
 
 // CHECK-LABEL: @test_lower_ntt
 // CHECK-SAME: (%[[INPUT:.*]]: [[T:.*]]) -> [[T]] {
-func.func @test_lower_ntt(%input : tensor<4x!PF1>) -> tensor<4x!PF1> {
+func.func @test_lower_ntt(%input : tensor<2x!PF1>) -> tensor<2x!PF1> {
   // CHECK-NOT: poly.ntt
-  %res = poly.ntt %input {root=#root} : tensor<4x!PF1>
-  return %res: tensor<4x!PF1>
+  %res = poly.ntt %input {root=#root} : tensor<2x!PF1>
+  return %res: tensor<2x!PF1>
 }
 
 // CHECK-LABEL: @test_lower_intt
 // CHECK-SAME: (%[[INPUT:.*]]: [[T:.*]]) -> [[P:.*]] {
-func.func @test_lower_intt(%input : tensor<4x!PF1>) -> tensor<4x!PF1> {
+func.func @test_lower_intt(%input : tensor<2x!PF1>) -> tensor<2x!PF1> {
   // CHECK-NOT: poly.intt
-  %res = poly.intt %input {root=#root} : tensor<4x!PF1>
-  return %res: tensor<4x!PF1>
+  %res = poly.intt %input {root=#root} : tensor<2x!PF1>
+  return %res: tensor<2x!PF1>
 }
diff --git a/zkir/Dialect/Arith/Conversions/ArithToModArith/ArithToModArith.cpp b/zkir/Dialect/Arith/Conversions/ArithToModArith/ArithToModArith.cpp
@@ -29,10 +29,10 @@ namespace mlir::zkir::arith {
 static mod_arith::ModArithType convertArithType(Type type) {
   auto modulusBitSize = static_cast<uint64_t>(type.getIntOrFloatBitWidth());
   auto modulus = (1L << (modulusBitSize - 1L));
-  auto newType = mlir::IntegerType::get(type.getContext(), modulusBitSize + 1);
+  auto newType = IntegerType::get(type.getContext(), modulusBitSize + 1);
 
   return mod_arith::ModArithType::get(type.getContext(),
-                                      mlir::IntegerAttr::get(newType, modulus));
+                                      IntegerAttr::get(newType, modulus));
 }
 
 static Type convertArithLikeType(ShapedType type) {
@@ -129,14 +129,14 @@ struct ConvertExtUI : public OpConversionPattern<mlir::arith::ExtUIOp> {
   }
 };
 
-struct ConvertLoadOp : public OpConversionPattern<mlir::memref::LoadOp> {
+struct ConvertLoadOp : public OpConversionPattern<memref::LoadOp> {
   explicit ConvertLoadOp(MLIRContext *context)
-      : OpConversionPattern<mlir::memref::LoadOp>(context) {}
+      : OpConversionPattern<memref::LoadOp>(context) {}
 
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(
-      ::mlir::memref::LoadOp op, OpAdaptor adaptor,
+      memref::LoadOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter) const override {
     ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
diff --git a/zkir/Dialect/Field/IR/BUILD.bazel b/zkir/Dialect/Field/IR/BUILD.bazel
@@ -23,6 +23,7 @@ cc_library(
         ":ops_inc_gen",
         ":types_inc_gen",
         "//zkir/Dialect/ModArith/IR:ModArith",
+        "//zkir/Utils:APIntUtils",
         "//zkir/Utils:OpUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
diff --git a/zkir/Dialect/Field/IR/FieldAttributes.cpp b/zkir/Dialect/Field/IR/FieldAttributes.cpp
@@ -2,18 +2,47 @@
 
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "zkir/Utils/APIntUtils.h"
 
 namespace mlir::zkir::field {
 
 LogicalResult PrimeFieldAttr::verify(
     ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
-    PrimeFieldType type, IntegerAttr value) {
-  if (type.getModulus().getValue().getBitWidth() !=
-      value.getValue().getBitWidth()) {
+    PrimeFieldType type, IntegerAttr _value) {
+  APInt modulus = type.getModulus().getValue();
+  APInt value = _value.getValue();
+
+  // check if storage type is same
+  if (modulus.getBitWidth() != value.getBitWidth()) {
     emitError()
         << "prime field modulus bitwidth does not match the value bitwidth";
     return failure();
   }
+
+  // check if value is in the field defined by modulus
+  if (value.uge(modulus)) {
+    emitError() << value.getZExtValue()
+                << " is not in the field defined by modulus "
+                << modulus.getZExtValue();
+    return failure();
+  }
+
+  return success();
+}
+
+LogicalResult RootOfUnityAttr::verify(
+    ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
+    PrimeFieldAttr root, IntegerAttr degree) {
+  APInt modulus = root.getType().getModulus().getValue();
+  APInt rootOfUnity = root.getValue().getValue();
+  unsigned degreeValue = degree.getValue().getZExtValue();
+
+  if (!expMod(rootOfUnity, degreeValue, modulus).isOne()) {
+    emitError() << rootOfUnity.getZExtValue()
+                << " is not a root of unity of degree " << degreeValue;
+    return failure();
+  }
+
   return success();
 }
 
diff --git a/zkir/Dialect/Field/IR/FieldAttributes.td b/zkir/Dialect/Field/IR/FieldAttributes.td
@@ -37,4 +37,16 @@ def Field_PrimeFieldAttr : Field_Attr<"PrimeField", "pf_elem", [TypedAttrInterfa
   let genVerifyDecl = 1;
 }
 
+def Field_RootOfUnityAttr : Field_Attr<"RootOfUnity", "root_of_unity", [TypedAttrInterface]> {
+  let summary = "An attribute representing a root of unity in a prime field";
+  let parameters = (ins Field_PrimeFieldAttr:$root, "IntegerAttr":$degree);
+  let assemblyFormat = "`<` $root `,` $degree `>`";
+  let extraClassDeclaration = [{
+    PrimeFieldType getType() const {
+      return getRoot().getType();
+    }
+  }];
+  let genVerifyDecl = 1;
+}
+
 #endif  // ZKIR_DIALECT_FIELD_IR_FIELDATTRS_TD_
diff --git a/zkir/Dialect/ModArith/Conversions/ModArithToArith/ModArithToArith.cpp b/zkir/Dialect/ModArith/Conversions/ModArithToArith/ModArithToArith.cpp
@@ -72,7 +72,7 @@ static TypedAttr modulusAttr(Op op, bool mul = false) {
   auto intType = IntegerType::get(op.getContext(), width);
   auto truncmod = modulus.zextOrTrunc(width);
 
-  if (auto st = mlir::dyn_cast<ShapedType>(type)) {
+  if (auto st = dyn_cast<ShapedType>(type)) {
     auto containerType = st.cloneWith(st.getShape(), intType);
     return DenseElementsAttr::get(containerType, truncmod);
   }
@@ -330,7 +330,7 @@ struct ConvertInverse : public OpConversionPattern<InverseOp> {
       InverseOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter) const override {
     // TODO(batzor): Support tensor input.
-    if (mlir::isa<ShapedType>(op.getInput().getType())) {
+    if (isa<ShapedType>(op.getInput().getType())) {
       return op->emitError("tensor input not supported");
     }
 
diff --git a/zkir/Dialect/Poly/Conversions/PolyToField/PolyToField.cpp b/zkir/Dialect/Poly/Conversions/PolyToField/PolyToField.cpp
@@ -234,7 +234,7 @@ static std::pair<Value, Value> bflyGS(ImplicitLocOpBuilder &b, Value A, Value B,
 }
 
 struct ConvertBitReverse : public OpConversionPattern<BitReverseOp> {
-  explicit ConvertBitReverse(mlir::MLIRContext *context)
+  explicit ConvertBitReverse(MLIRContext *context)
       : OpConversionPattern<BitReverseOp>(context) {}
 
   using OpConversionPattern::OpConversionPattern;
diff --git a/zkir/Dialect/Poly/IR/PolyAttributes.cpp b/zkir/Dialect/Poly/IR/PolyAttributes.cpp
@@ -59,17 +59,23 @@ static void precomputeRoots(APInt root, const APInt &mod, unsigned degree,
   pool.wait();
 }
 
-zkir::field::PrimeFieldAttr PrimitiveRootAttr::getRoot() const {
-  return getImpl()->root;
+field::RootOfUnityAttr PrimitiveRootAttr::getRootOfUnity() const {
+  return getImpl()->rootOfUnity;
 }
 
-zkir::field::PrimeFieldAttr PrimitiveRootAttr::getInvRoot() const {
+field::PrimeFieldAttr PrimitiveRootAttr::getRoot() const {
+  return getImpl()->rootOfUnity.getRoot();
+}
+
+field::PrimeFieldAttr PrimitiveRootAttr::getInvRoot() const {
   return getImpl()->invRoot;
 }
 
-IntegerAttr PrimitiveRootAttr::getDegree() const { return getImpl()->degree; }
+IntegerAttr PrimitiveRootAttr::getDegree() const {
+  return getImpl()->rootOfUnity.getDegree();
+}
 
-zkir::field::PrimeFieldAttr PrimitiveRootAttr::getInvDegree() const {
+field::PrimeFieldAttr PrimitiveRootAttr::getInvDegree() const {
   return getImpl()->invDegree;
 }
 
@@ -81,7 +87,7 @@ DenseElementsAttr PrimitiveRootAttr::getInvRoots() const {
   return getImpl()->invRoots;
 }
 
-zkir::mod_arith::MontgomeryAttr PrimitiveRootAttr::getMontgomery() const {
+mod_arith::MontgomeryAttr PrimitiveRootAttr::getMontgomery() const {
   return getImpl()->montgomery;
 }
 
@@ -90,12 +96,13 @@ namespace detail {
 PrimitiveRootAttrStorage *PrimitiveRootAttrStorage::construct(
     AttributeStorageAllocator &allocator, KeyTy &&key) {
   // Extract the root and degree from the key.
-  zkir::field::PrimeFieldAttr root = std::get<0>(key);
-  IntegerAttr degree = std::get<1>(key);
-  zkir::mod_arith::MontgomeryAttr montgomery = std::get<2>(key);
+  field::RootOfUnityAttr rootOfUnity = std::get<0>(key);
+  field::PrimeFieldAttr root = rootOfUnity.getRoot();
+  IntegerAttr degree = rootOfUnity.getDegree();
+  mod_arith::MontgomeryAttr montgomery = std::get<1>(key);
 
   std::optional<IntegerAttr> montgomeryR;
-  if (montgomery != zkir::mod_arith::MontgomeryAttr()) {
+  if (montgomery != mod_arith::MontgomeryAttr()) {
     montgomeryR = montgomery.getR();
   }
 
@@ -111,8 +118,8 @@ PrimitiveRootAttrStorage *PrimitiveRootAttrStorage::construct(
 
   field::PrimeFieldAttr invDegree =
       field::PrimeFieldAttr::get(root.getType(), invDegreeVal);
-  zkir::field::PrimeFieldAttr invRoot =
-      zkir::field::PrimeFieldAttr::get(root.getType(), invRootVal);
+  field::PrimeFieldAttr invRoot =
+      field::PrimeFieldAttr::get(root.getType(), invRootVal);
 
   // Compute the exponent table.
   SmallVector<APInt> roots, invRoots;
@@ -125,10 +132,9 @@ PrimitiveRootAttrStorage *PrimitiveRootAttrStorage::construct(
   DenseElementsAttr rootsAttr = DenseElementsAttr::get(tensorType, roots);
   DenseElementsAttr invRootsAttr = DenseElementsAttr::get(tensorType, invRoots);
   return new (allocator.allocate<PrimitiveRootAttrStorage>())
-      PrimitiveRootAttrStorage(std::move(degree), std::move(invDegree),
-                               std::move(root), std::move(invRoot),
-                               std::move(rootsAttr), std::move(invRootsAttr),
-                               std::move(montgomery));
+      PrimitiveRootAttrStorage(std::move(rootOfUnity), std::move(invDegree),
+                               std::move(invRoot), std::move(rootsAttr),
+                               std::move(invRootsAttr), std::move(montgomery));
 }
 
 }  // namespace detail
diff --git a/zkir/Dialect/Poly/IR/PolyAttributes.h b/zkir/Dialect/Poly/IR/PolyAttributes.h
@@ -13,43 +13,38 @@
 namespace mlir::zkir::poly::detail {
 
 struct PrimitiveRootAttrStorage : public AttributeStorage {
-  using KeyTy = std::tuple<zkir::field::PrimeFieldAttr, IntegerAttr,
-                           mod_arith::MontgomeryAttr>;
-  PrimitiveRootAttrStorage(IntegerAttr degree,
-                           zkir::field::PrimeFieldAttr invDegree,
-                           zkir::field::PrimeFieldAttr root,
-                           zkir::field::PrimeFieldAttr invRoot,
+  using KeyTy = std::tuple<field::RootOfUnityAttr, mod_arith::MontgomeryAttr>;
+  PrimitiveRootAttrStorage(field::RootOfUnityAttr rootOfUnity,
+                           field::PrimeFieldAttr invDegree,
+                           field::PrimeFieldAttr invRoot,
                            DenseElementsAttr roots, DenseElementsAttr invRoots,
-                           zkir::mod_arith::MontgomeryAttr montgomery)
-      : degree(std::move(degree)),
+                           mod_arith::MontgomeryAttr montgomery)
+      : rootOfUnity(std::move(rootOfUnity)),
         invDegree(std::move(invDegree)),
-        root(std::move(root)),
         invRoot(std::move(invRoot)),
         roots(std::move(roots)),
         invRoots(std::move(invRoots)),
         montgomery(std::move(montgomery)) {}
 
-  KeyTy getAsKey() const { return KeyTy(root, degree, montgomery); }
+  KeyTy getAsKey() const { return KeyTy(rootOfUnity, montgomery); }
 
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(root, degree, montgomery);
+    return key == KeyTy(rootOfUnity, montgomery);
   }
 
   static llvm::hash_code hashKey(const KeyTy &key) {
-    return llvm::hash_combine(std::get<0>(key), std::get<1>(key),
-                              std::get<2>(key));
+    return llvm::hash_combine(std::get<0>(key), std::get<1>(key));
   }
 
   static PrimitiveRootAttrStorage *construct(
       AttributeStorageAllocator &allocator, KeyTy &&key);
 
-  IntegerAttr degree;
-  zkir::field::PrimeFieldAttr invDegree;
-  zkir::field::PrimeFieldAttr root;
-  zkir::field::PrimeFieldAttr invRoot;
+  field::RootOfUnityAttr rootOfUnity;
+  field::PrimeFieldAttr invDegree;
+  field::PrimeFieldAttr invRoot;
   DenseElementsAttr roots;
   DenseElementsAttr invRoots;
-  zkir::mod_arith::MontgomeryAttr montgomery;
+  mod_arith::MontgomeryAttr montgomery;
 };
 
 }  // namespace mlir::zkir::poly::detail
diff --git a/zkir/Dialect/Poly/IR/PolyAttributes.td b/zkir/Dialect/Poly/IR/PolyAttributes.td
diff --git a/zkir/Dialect/Poly/IR/PolyOps.cpp b/zkir/Dialect/Poly/IR/PolyOps.cpp
diff --git a/zkir/Pipelines/PipelineRegistration.cpp b/zkir/Pipelines/PipelineRegistration.cpp
diff --git a/zkir/Pipelines/PipelineRegistration.h b/zkir/Pipelines/PipelineRegistration.h
diff --git a/zkir/Utils/ConversionUtils.h b/zkir/Utils/ConversionUtils.h
