refac(poly): use RootOfUnityAttr

Author: batzor

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

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>
 

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 }

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() {

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>
 }

zkir/Dialect/Poly/IR/PolyAttributes.cpp

@@ -59,15 +59,21 @@ static void precomputeRoots(APInt root, const APInt &mod, unsigned degree,
   pool.wait();
 }
 
+field::RootOfUnityAttr PrimitiveRootAttr::getRootOfUnity() const {
+  return getImpl()->rootOfUnity;
+}
+
 field::PrimeFieldAttr PrimitiveRootAttr::getRoot() const {
-  return getImpl()->root;
+  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();
+}
 
 field::PrimeFieldAttr PrimitiveRootAttr::getInvDegree() const {
   return getImpl()->invDegree;
@@ -90,9 +96,10 @@ namespace detail {
 PrimitiveRootAttrStorage *PrimitiveRootAttrStorage::construct(
     AttributeStorageAllocator &allocator, KeyTy &&key) {
   // Extract the root and degree from the key.
-  field::PrimeFieldAttr root = std::get<0>(key);
-  IntegerAttr degree = std::get<1>(key);
-  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 != mod_arith::MontgomeryAttr()) {
@@ -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

zkir/Dialect/Poly/IR/PolyAttributes.h

@@ -13,39 +13,34 @@
 namespace mlir::zkir::poly::detail {
 
 struct PrimitiveRootAttrStorage : public AttributeStorage {
-  using KeyTy = std::tuple<field::PrimeFieldAttr, IntegerAttr,
-                           mod_arith::MontgomeryAttr>;
-  PrimitiveRootAttrStorage(IntegerAttr degree,
+  using KeyTy = std::tuple<field::RootOfUnityAttr, mod_arith::MontgomeryAttr>;
+  PrimitiveRootAttrStorage(field::RootOfUnityAttr rootOfUnity,
                            field::PrimeFieldAttr invDegree,
-                           field::PrimeFieldAttr root,
                            field::PrimeFieldAttr invRoot,
                            DenseElementsAttr roots, DenseElementsAttr invRoots,
                            mod_arith::MontgomeryAttr montgomery)
-      : degree(std::move(degree)),
+      : 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;
+  field::RootOfUnityAttr rootOfUnity;
   field::PrimeFieldAttr invDegree;
-  field::PrimeFieldAttr root;
   field::PrimeFieldAttr invRoot;
   DenseElementsAttr roots;
   DenseElementsAttr invRoots;

zkir/Dialect/Poly/IR/PolyAttributes.td

@@ -35,29 +35,28 @@ def Poly_UnivariatePolyAttr : Poly_Attr<"UnivariatePoly", "univariate_polynomial
 }
 
 def Poly_PrimitiveRootAttr: Poly_Attr<"PrimitiveRoot", "primitive_root"> {
-  let summary = "an attribute containing an integer and its degree as a root of unity";
+  let summary = "an attribute containing a primitive root of unity and its twiddles";
   let description = [{
-    A primitive root attribute stores an integer root `value` and an integer
-    `degree`, corresponding to a primitive root of unity of the given degree in
-    an unspecified ring.
-
-    This is used as an attribute on `poly.ntt` and `poly.intt` ops
-    to specify the root of unity used in lowering the transform.
+    A primitive root attribute stores a primitive root of unity specified by the `RootOfUnityAttr`
+    attribute and also some other necessary precomputed values for (I)NTT operations.
 
     Example:
     ```
     !PF1 = !field.pf<7:i32>
     #elem = #field.elem<2> : !PF1
-    #root = #poly.primitive_root<root=%generator, degree=3>
+    #root_of_unity = #field.root_of_unity<root=#elem, degree=3>
+    #root = #poly.primitive_root<root_of_unity=#root_of_unity>
     ```
   }];
   let extraClassDeclaration = [{
+    IntegerAttr getDegree() const;
+    field::PrimeFieldAttr getRoot() const;
     field::PrimeFieldAttr getInvRoot() const;
     field::PrimeFieldAttr getInvDegree() const;
     DenseElementsAttr getRoots() const;
     DenseElementsAttr getInvRoots() const;
   }];
-  let parameters = (ins Field_PrimeFieldAttr:$root, "IntegerAttr":$degree, OptionalParameter<"mod_arith::MontgomeryAttr">:$montgomery);
+  let parameters = (ins Field_RootOfUnityAttr:$root_of_unity, OptionalParameter<"mod_arith::MontgomeryAttr">:$montgomery);
   let assemblyFormat = "`<` struct(params) `>`";
   let genStorageClass = 0;
 }