Replace "witness only" with "non native felt op" attribute on felt arith ops (#288)

Author: tim-hoffman

changelogs/unreleased/th__non_native_felt_attr.yaml

@@ -0,0 +1,5 @@
+added:
+  - "Add `function.allow_non_native_field_ops` attribute"
+changed:
+  - "Require `function.allow_non_native_field_ops` attribute on functions that use
+    non-native field operations, instead of requiring `function.allow_witness` attribute."

include/llzk-c/Dialect/Function.h

@@ -84,6 +84,12 @@ LLZK_DECLARE_OP_PREDICATE(FuncDefOp, HasAllowWitnessAttr);
 /// Sets the allow_witness attribute in the FuncDefOp operation.
 MLIR_CAPI_EXPORTED void llzkFuncDefOpSetAllowWitnessAttr(MlirOperation op, bool value);
 
+/// Returns true if the FuncDefOp has the allow_non_native_field_ops attribute.
+LLZK_DECLARE_OP_PREDICATE(FuncDefOp, HasAllowNonNativeFieldOpsAttr);
+
+/// Sets the allow_non_native_field_ops attribute in the FuncDefOp operation.
+MLIR_CAPI_EXPORTED void llzkFuncDefOpSetAllowNonNativeFieldOpsAttr(MlirOperation op, bool value);
+
 /// Returns true if the `idx`-th argument has the Pub attribute.
 LLZK_DECLARE_NARY_OP_PREDICATE(FuncDefOp, HasArgIsPub, unsigned arg);
 

include/llzk/Dialect/Bool/IR/Ops.td

@@ -39,7 +39,8 @@ class BoolUnaryOpBase<string mnemonic, Type resultType, list<Trait> traits = []>
 // Boolean operators
 //===------------------------------------------------------------------===//
 
-def LLZK_AndBoolOp : BoolBinaryOpBase<"and", I1, [Commutative]> {
+def LLZK_AndBoolOp
+    : BoolBinaryOpBase<"and", I1, [NotFieldNative, Commutative]> {
   let summary = "logical AND operator";
   let description = [{
     This operation computes the logical AND (i.e., conjunction) of two `i1` (i.e., boolean)
@@ -49,7 +50,7 @@ def LLZK_AndBoolOp : BoolBinaryOpBase<"and", I1, [Commutative]> {
   }];
 }
 
-def LLZK_OrBoolOp : BoolBinaryOpBase<"or", I1, [Commutative]> {
+def LLZK_OrBoolOp : BoolBinaryOpBase<"or", I1, [NotFieldNative, Commutative]> {
   let summary = "logical OR operator";
   let description = [{
     This operation computes the logical OR (i.e., disjunction) of two `i1` (i.e., boolean)
@@ -59,7 +60,8 @@ def LLZK_OrBoolOp : BoolBinaryOpBase<"or", I1, [Commutative]> {
   }];
 }
 
-def LLZK_XorBoolOp : BoolBinaryOpBase<"xor", I1, [Commutative]> {
+def LLZK_XorBoolOp
+    : BoolBinaryOpBase<"xor", I1, [NotFieldNative, Commutative]> {
   let summary = "logical XOR operator";
   let description = [{
     This operation computes the logical XOR (i.e., exclusive disjunction) of two `i1` (i.e., boolean)
@@ -69,7 +71,7 @@ def LLZK_XorBoolOp : BoolBinaryOpBase<"xor", I1, [Commutative]> {
   }];
 }
 
-def LLZK_NotBoolOp : BoolUnaryOpBase<"not", I1, []> {
+def LLZK_NotBoolOp : BoolUnaryOpBase<"not", I1, [NotFieldNative]> {
   let summary = "logical NOT operator";
   let description = [{
     This operation computes the logical NOT (i.e., negation) of an `i1` (i.e., boolean)

include/llzk/Dialect/Felt/IR/Ops.td

@@ -103,7 +103,8 @@ def LLZK_MulFeltOp : FeltDialectBinaryOp<"mul", LLZK_FeltType, [Commutative]> {
   let description = [{}];
 }
 
-def LLZK_PowFeltOp : FeltDialectBinaryOp<"pow", LLZK_FeltType, [WitnessGen]> {
+def LLZK_PowFeltOp
+    : FeltDialectBinaryOp<"pow", LLZK_FeltType, [NotFieldNative]> {
   let summary = "exponentiation operator for field elements";
   let description = [{
 
@@ -122,7 +123,7 @@ def LLZK_DivFeltOp : FeltDialectBinaryOp<"div", LLZK_FeltType> {
 }
 
 def LLZK_UnsignedIntDivFeltOp
-    : FeltDialectBinaryOp<"uintdiv", LLZK_FeltType, [WitnessGen]> {
+    : FeltDialectBinaryOp<"uintdiv", LLZK_FeltType, [NotFieldNative]> {
   let summary = "unsigned integer division operator for field elements";
   let description = [{
     Treats the operands as if they were unsigned integers with bitwidth
@@ -131,7 +132,7 @@ def LLZK_UnsignedIntDivFeltOp
 }
 
 def LLZK_SignedIntDivFeltOp
-    : FeltDialectBinaryOp<"sintdiv", LLZK_FeltType, [WitnessGen]> {
+    : FeltDialectBinaryOp<"sintdiv", LLZK_FeltType, [NotFieldNative]> {
   let summary = "signed integer division operator for field elements";
   let description = [{
     Treats the operands as if they were signed integers with bitwidth
@@ -148,7 +149,7 @@ def LLZK_SignedIntDivFeltOp
 }
 
 def LLZK_UnsignedModFeltOp
-    : FeltDialectBinaryOp<"umod", LLZK_FeltType, [WitnessGen]> {
+    : FeltDialectBinaryOp<"umod", LLZK_FeltType, [NotFieldNative]> {
   let summary =
       "unsigned integer modulus/remainder operator for field elements";
   let description = [{
@@ -158,7 +159,7 @@ def LLZK_UnsignedModFeltOp
 }
 
 def LLZK_SignedModFeltOp
-    : FeltDialectBinaryOp<"smod", LLZK_FeltType, [WitnessGen]> {
+    : FeltDialectBinaryOp<"smod", LLZK_FeltType, [NotFieldNative]> {
   let summary = "signed integer modulus/remainder operator for field elements";
   let description = [{
     Computes the remainder that would result from the division operation performed
@@ -171,31 +172,35 @@ def LLZK_NegFeltOp : FeltDialectUnaryOp<"neg", LLZK_FeltType> {
   let description = [{}];
 }
 
-def LLZK_InvFeltOp : FeltDialectUnaryOp<"inv", LLZK_FeltType, [WitnessGen]> {
+def LLZK_InvFeltOp
+    : FeltDialectUnaryOp<"inv", LLZK_FeltType, [NotFieldNative]> {
   let summary = "inverse operator for field elements";
   let description = [{}];
 }
 
 def LLZK_AndFeltOp
-    : FeltDialectBinaryOp<"bit_and", LLZK_FeltType, [WitnessGen, Commutative]> {
+    : FeltDialectBinaryOp<"bit_and",
+                          LLZK_FeltType, [NotFieldNative, Commutative]> {
   let summary = "bitwise AND operator for field elements";
   let description = [{}];
 }
 
 def LLZK_OrFeltOp
-    : FeltDialectBinaryOp<"bit_or", LLZK_FeltType, [WitnessGen, Commutative]> {
+    : FeltDialectBinaryOp<"bit_or",
+                          LLZK_FeltType, [NotFieldNative, Commutative]> {
   let summary = "bitwise OR operator for field elements";
   let description = [{}];
 }
 
 def LLZK_XorFeltOp
-    : FeltDialectBinaryOp<"bit_xor", LLZK_FeltType, [WitnessGen, Commutative]> {
+    : FeltDialectBinaryOp<"bit_xor",
+                          LLZK_FeltType, [NotFieldNative, Commutative]> {
   let summary = "bitwise XOR operator for field elements";
   let description = [{}];
 }
 
 def LLZK_NotFeltOp
-    : FeltDialectUnaryOp<"bit_not", LLZK_FeltType, [WitnessGen]> {
+    : FeltDialectUnaryOp<"bit_not", LLZK_FeltType, [NotFieldNative]> {
   let summary = "integer complement (bitwise-not) operator for field elements";
   let description = [{
     Treats the operand as an integer with a bitwidth equal to the bitwidth
@@ -204,12 +209,14 @@ def LLZK_NotFeltOp
   }];
 }
 
-def LLZK_ShlFeltOp : FeltDialectBinaryOp<"shl", LLZK_FeltType, [WitnessGen]> {
+def LLZK_ShlFeltOp
+    : FeltDialectBinaryOp<"shl", LLZK_FeltType, [NotFieldNative]> {
   let summary = "left shift operator for field elements";
   let description = [{}];
 }
 
-def LLZK_ShrFeltOp : FeltDialectBinaryOp<"shr", LLZK_FeltType, [WitnessGen]> {
+def LLZK_ShrFeltOp
+    : FeltDialectBinaryOp<"shr", LLZK_FeltType, [NotFieldNative]> {
   let summary = "right shift operator for field elements";
   let description = [{}];
 }

include/llzk/Dialect/Function/IR/Attrs.td

@@ -22,19 +22,31 @@ class FunctionDialectAttr<string name, string attrMnemonic,
 
 def LLZK_AllowConstraintAttr
     : FunctionDialectAttr<"AllowConstraint", "allow_constraint"> {
-  let summary = "Marks functions to allow `ConstraintGen` ops";
+  let summary =
+      "Marks functions to allow `ConstraintGen` ops within their body";
   let description = [{
     A unit attribute that can be attached to a `FuncDefOp` to indicate that ops
-    marked with `ConstraintGen` (ex: emit ops) are allowed within its body.
+    marked with `ConstraintGen` (e.g., emit ops) are allowed within its body.
   }];
 }
 
 def LLZK_AllowWitnessAttr
     : FunctionDialectAttr<"AllowWitness", "allow_witness"> {
-  let summary = "Marks functions to allow `WitnessGen` ops";
+  let summary = "Marks functions to allow `WitnessGen` ops within their body";
   let description = [{
-    A unit attribute that can be attached to a `FuncDefOp` to indicate that ops
-    marked with `WitnessGen` (i.e., boolean and bitwise ops) are allowed within its body.
+    A unit attribute that can be attached to a `FuncDefOp` to indicate that ops marked
+    with `WitnessGen` (e.g., construct and write ops) are allowed within its body.
+  }];
+}
+
+def LLZK_AllowNonNativeFieldOpsAttr
+    : FunctionDialectAttr<"AllowNonNativeFieldOps",
+                          "allow_non_native_field_ops"> {
+  let summary =
+      "Marks functions to allow `NotFieldNative` ops within their body";
+  let description = [{
+    A unit attribute that can be attached to a `FuncDefOp` to indicate that ops marked
+    with `NotFieldNative` (e.g., boolean and bitwise ops) are allowed within its body.
   }];
 }
 

include/llzk/Dialect/Function/IR/OpTraits.h

@@ -17,6 +17,7 @@ namespace llzk::function {
 
 mlir::LogicalResult verifyConstraintGenTraitImpl(mlir::Operation *op);
 mlir::LogicalResult verifyWitnessGenTraitImpl(mlir::Operation *op);
+mlir::LogicalResult verifyNotFieldNativeTraitImpl(mlir::Operation *op);
 
 /// Marker for ops that are specific to constraint generation.
 /// Verifies that the surrounding function is marked with the `AllowConstraintAttr`.
@@ -38,4 +39,14 @@ class WitnessGen : public mlir::OpTrait::TraitBase<TypeClass, WitnessGen> {
   }
 };
 
+/// Marker for ops over `llzk.felt` type operands that are not native to finite field arithmetic.
+/// Verifies that the surrounding function is marked with the `AllowNonNativeFieldOpsAttr`.
+template <typename TypeClass>
+class NotFieldNative : public mlir::OpTrait::TraitBase<TypeClass, NotFieldNative> {
+public:
+  inline static mlir::LogicalResult verifyTrait(mlir::Operation *op) {
+    return verifyNotFieldNativeTraitImpl(op);
+  }
+};
+
 } // namespace llzk::function

include/llzk/Dialect/Function/IR/OpTraits.td

@@ -24,4 +24,12 @@ def ConstraintGen : NativeOpTrait<"ConstraintGen">, StructuralOpTrait {
   string cppNamespace = "::llzk::function";
 }
 
+/// Marker for ops over `llzk.felt` type operands that are not native to finite
+/// field arithmetic and thus must be rewritten or folded into a constant after
+/// the flattening and unrolling pass. Op classes with this trait are valid only
+/// in functions that have the `AllowNonNativeFieldOpsAttr`.
+def NotFieldNative : NativeOpTrait<"NotFieldNative">, StructuralOpTrait {
+  string cppNamespace = "::llzk::function";
+}
+
 #endif // LLZK_FUNCTION_TRAITS

include/llzk/Dialect/Function/IR/Ops.td

@@ -150,6 +150,14 @@ def FuncDefOp
     /// Add (resp. remove) the `allow_witness` attribute to (resp. from) the function def.
     void setAllowWitnessAttr(bool newValue = true);
 
+    /// Return `true` iff the function def has the `allow_non_native_field_ops` attribute.
+    inline bool hasAllowNonNativeFieldOpsAttr() {
+      return getOperation()->hasAttr(llzk::function::AllowNonNativeFieldOpsAttr::name);
+    }
+
+    /// Add (resp. remove) the `allow_non_native_field_ops` attribute to (resp. from) the function def.
+    void setAllowNonNativeFieldOpsAttr(bool newValue = true);
+
     /// Return `true` iff the argument at the given index has `pub` attribute.
     bool hasArgPublicAttr(unsigned index);
 

lib/CAPI/Dialect/Function.cpp

@@ -79,6 +79,14 @@ void llzkFuncDefOpSetAllowWitnessAttr(MlirOperation op, bool value) {
   unwrap_cast<FuncDefOp>(op).setAllowWitnessAttr(value);
 }
 
+bool llzkFuncDefOpGetHasAllowNonNativeFieldOpsAttr(MlirOperation op) {
+  return unwrap_cast<FuncDefOp>(op).hasAllowNonNativeFieldOpsAttr();
+}
+
+void llzkFuncDefOpSetAllowNonNativeFieldOpsAttr(MlirOperation op, bool value) {
+  unwrap_cast<FuncDefOp>(op).setAllowNonNativeFieldOpsAttr(value);
+}
+
 bool llzkFuncDefOpGetHasArgIsPub(MlirOperation op, unsigned argNo) {
   return unwrap_cast<FuncDefOp>(op).hasArgPublicAttr(argNo);
 }

lib/Dialect/Function/IR/OpTraits.cpp

@@ -43,4 +43,12 @@ LogicalResult verifyWitnessGenTraitImpl(Operation *op) {
   return success();
 }
 
+LogicalResult verifyNotFieldNativeTraitImpl(Operation *op) {
+  if (!parentFuncDefOpHasAttr(op, &FuncDefOp::hasAllowNonNativeFieldOpsAttr)) {
+    return op->emitOpError() << "only valid within a '" << FuncDefOp::getOperationName()
+                             << "' with '" << AllowNonNativeFieldOpsAttr::name << "' attribute";
+  }
+  return success();
+}
+
 } // namespace llzk::function