feat(field): add RootOfUnityAttr

Author: batzor

tests/Dialect/Field/prime_field_to_mod_arith.mlir

@@ -2,6 +2,8 @@
 !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>

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",

zkir/Dialect/Field/IR/FieldAttributes.cpp

@@ -2,6 +2,7 @@
 
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "zkir/Utils/APIntUtils.h"
 
 namespace mlir::zkir::field {
 
@@ -17,4 +18,20 @@ LogicalResult PrimeFieldAttr::verify(
   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();
+}
+
 }  // namespace mlir::zkir::field

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_