extension field muls & torus benchmarks

src/benchmarks.py

@@ -0,0 +1,111 @@
+from src.definitions import STARK, CurveID, CURVES, PyFelt, Curve
+from random import randint
+from src.extension_field_modulo_circuit import ExtensionFieldModuloCircuit
+from src.precompiled_circuits.final_exp import FinalExpCircuit
+
+
+def test_extf_mul_circuit_amortized(curve_id: CurveID, extension_degree: int):
+    curve: Curve = CURVES[curve_id.value]
+    p = curve.p
+    X = [PyFelt(randint(0, p - 1), p) for _ in range(extension_degree)]
+    circuit = ExtensionFieldModuloCircuit(
+        f"{curve_id.name}_mul_amortized_Fp{extension_degree}",
+        curve_id=curve_id.value,
+        extension_degree=extension_degree,
+    )
+    circuit.create_powers_of_Z(PyFelt(2, STARK), mock=True)
+    X = circuit.write_elements(X)
+    circuit.extf_mul(X, X, extension_degree)
+    return circuit.summarize()
+
+
+def test_extf_square_amortized(curve_id: CurveID, extension_degree: int):
+    curve: Curve = CURVES[curve_id.value]
+    p = curve.p
+    X = [PyFelt(randint(0, p - 1), p) for _ in range(extension_degree)]
+    circuit = ExtensionFieldModuloCircuit(
+        f"{curve_id.name}_square_amortized_Fp{extension_degree}",
+        curve_id=curve_id.value,
+        extension_degree=extension_degree,
+    )
+    circuit.create_powers_of_Z(PyFelt(2, STARK), mock=True)
+    X = circuit.write_elements(X)
+    circuit.extf_square(X, extension_degree)
+    return circuit.summarize()
+
+
+def test_extf_mul_circuit_full(curve_id: CurveID, extension_degree: int):
+    curve: Curve = CURVES[curve_id.value]
+    p = curve.p
+    X = [PyFelt(randint(0, p - 1), p) for _ in range(extension_degree)]
+    circuit = ExtensionFieldModuloCircuit(
+        f"{curve_id.name}_mul_zpow_verif_Fp{extension_degree}",
+        curve_id=curve_id.value,
+        extension_degree=extension_degree,
+    )
+    circuit.create_powers_of_Z(PyFelt(2, STARK), mock=False)
+    X = circuit.write_elements(X)
+    circuit.extf_mul(X, X, extension_degree)
+    circuit.finalize_circuit()
+    return circuit.summarize()
+
+
+def test_square_torus_amortized(curve_id: CurveID, extension_degree: int):
+    curve: Curve = CURVES[curve_id.value]
+    p = curve.p
+    X = [PyFelt(randint(0, p - 1), p) for _ in range(extension_degree)]
+    circuit = FinalExpCircuit(
+        f"{curve_id.name}_square_torus_amortized_Fp{extension_degree}",
+        curve_id=curve_id.value,
+        extension_degree=extension_degree,
+    )
+    circuit.create_powers_of_Z(PyFelt(2, STARK), mock=True)
+    X = circuit.write_elements(X)
+    circuit.square_torus(X)
+    return circuit.summarize()
+
+
+def test_mul_torus_amortized(curve_id: CurveID, extension_degree: int):
+    curve: Curve = CURVES[curve_id.value]
+    p = curve.p
+    X = [PyFelt(randint(0, p - 1), p) for _ in range(extension_degree)]
+    circuit = FinalExpCircuit(
+        f"{curve_id.name}_mul_torus_amortized_Fp{extension_degree}",
+        curve_id=curve_id.value,
+        extension_degree=extension_degree,
+    )
+    circuit.create_powers_of_Z(PyFelt(2, STARK), mock=True)
+    X = circuit.write_elements(X)
+    circuit.mul_torus(X, X)
+    return circuit.summarize()
+
+
+if __name__ == "__main__":
+    data = []
+    for curveID in [CurveID.BLS12_381]:
+        data.append(test_extf_mul_circuit_amortized(curveID, 6))
+        data.append(test_extf_mul_circuit_full(curveID, 6))
+        data.append(test_extf_mul_circuit_amortized(curveID, 12))
+        data.append(test_extf_mul_circuit_full(curveID, 12))
+        data.append(test_square_torus_amortized(curveID, 6))
+        data.append(test_extf_square_amortized(curveID, 12))
+        data.append(test_mul_torus_amortized(curveID, 6))
+
+    import pandas as pd
+
+    df = pd.DataFrame(data)
+
+    costs = {"MULMOD": 8, "ADDMOD": 4, "POSEIDON": 7}  # , "ASSERT_EQ": 3}
+    acc = pd.Series([0] * len(df))
+
+    for column, multiplier in costs.items():
+        tmp = df[column] * multiplier
+        df[f"{column}_w"] = tmp
+        acc += tmp
+
+    df[f"Total with weights"] = acc
+
+    print("\n\n")
+    print(f"Weights: {costs}")
+    print(df.sort_values(by="Total with weights"))
+    print("\n")

src/definitions.py

@@ -1,6 +1,7 @@
 from src.algebra import Polynomial
 from src.algebra import BaseField, PyFelt, ModuloCircuitElement
 from dataclasses import dataclass
+from enum import Enum
 
 N_LIMBS = 4
 BASE = 2**96
@@ -9,6 +10,11 @@
 BLS12_381_ID = int.from_bytes(b"bls12_381", "big")
 
 
+class CurveID(Enum):
+    BN254 = int.from_bytes(b"bn254", "big")
+    BLS12_381 = int.from_bytes(b"bls12_381", "big")
+
+
 @dataclass(slots=True, frozen=True)
 class Curve:
     id: int
@@ -47,10 +53,10 @@ def get_base_field(curve_id: int) -> BaseField:
 
 
 def get_irreducible_poly(curve_id: int, extension_degree: int) -> Polynomial:
+    field = BaseField(CURVES[curve_id].p)
     return Polynomial(
         coefficients=[
-            PyFelt(x, CURVES[curve_id].p)
-            for x in CURVES[curve_id].irreducible_polys[extension_degree]
+            field(x) for x in CURVES[curve_id].irreducible_polys[extension_degree]
         ],
         raw_init=True,
     )

src/extension_field_modulo_circuit.py

@@ -3,9 +3,12 @@
 from src.poseidon_transcript import CairoPoseidonTranscript
 from src.hints.extf_mul import (
     nondeterministic_extension_field_mul_divmod,
+    nondeterministic_extension_field_div,
 )
 from dataclasses import dataclass
-from src.definitions import get_irreducible_poly
+from src.definitions import get_irreducible_poly, CurveID, STARK, CURVES, Curve
+from pprint import pprint
+from random import randint
 
 
 # Accumulates equations of the form c_i * X_i(Z)*Y_i(z) = Q_i*P + R_i
@@ -31,37 +34,46 @@ def __init__(
         self.transcript: CairoPoseidonTranscript = CairoPoseidonTranscript(
             init_hash=int.from_bytes(name.encode(), "big")
         )
+        self.z_powers: list[ModuloCircuitElement] = []
 
     def _init_accumulator(self):
         return EuclideanPolyAccumulator(
-            xy=self.write_element(self.field.zero()),
+            xy=self.constants["ZERO"],
             nondeterministic_Q=Polynomial([self.field.zero()]),
-            R=[
-                self.write_element(self.field.zero())
-                for _ in range(self.extension_degree)
-            ],
+            R=[self.constants["ZERO"]] * self.extension_degree,
         )
 
-    def write_commitment(self, elmt: PyFelt) -> ModuloCircuitElement:
+    def write_commitment(
+        self, elmt: PyFelt, add_to_transcript: bool = True
+    ) -> ModuloCircuitElement:
         """
         1) Add the commitment to the list of commitments
         2) Add to transcript to precompute Z
         3) Write the commitment to the values segment and return the ModuloElement
         """
         self.commitments.append(elmt)
-        self.transcript.hash_limbs(elmt)
+        if add_to_transcript:
+            self.transcript.hash_limbs(elmt)
         return self.write_element(elmt, WriteOps.COMMIT)
 
     def write_commitments(
-        self, elmts: list[PyFelt]
+        self, elmts: list[PyFelt], add_to_transcript: bool = True
     ) -> (list[ModuloCircuitElement], PyFelt, PyFelt):
-        vals = [self.write_commitment(elmt) for elmt in elmts]
+        vals = [self.write_commitment(elmt, add_to_transcript) for elmt in elmts]
         return vals, self.transcript.continuable_hash, self.transcript.s1
 
-    def create_powers_of_Z(self, Z: PyFelt) -> list[ModuloCircuitElement]:
-        powers = [self.write_cairo_native_felt(Z.value)]
-        for _ in range(1, self.extension_degree):
-            powers.append(self.mul(powers[-1], powers[0]))
+    def create_powers_of_Z(
+        self, Z: PyFelt, mock: bool = False
+    ) -> list[ModuloCircuitElement]:
+        powers = [self.write_cairo_native_felt(Z)]
+        if not mock:
+            for _ in range(1, self.extension_degree + 1):
+                powers.append(self.mul(powers[-1], powers[0]))
+        else:
+            powers = powers + [
+                self.write_cairo_native_felt(self.field(Z.value**i))
+                for i in range(1, self.extension_degree + 1)
+            ]
         self.z_powers = powers
         return powers
 
@@ -75,7 +87,7 @@ def eval_poly_in_precomputed_Z(
         assert len(X) <= len(self.z_powers), f"{len(X)} > {len(self.z_powers)}"
         X_of_z = X[0]
         for i in range(1, len(X)):
-            X_of_z = self.add(X_of_z, self.mul(X[i], self.z_powers[i]))
+            X_of_z = self.add(X_of_z, self.mul(X[i], self.z_powers[i - 1]))
         return X_of_z
 
     def extf_add(
@@ -102,24 +114,36 @@ def extf_scalar_mul(
         assert type(c) == ModuloCircuitElement
         return [self.mul(x_i, c) for x_i in X]
 
+    def extf_neg(self, X: list[ModuloCircuitElement]) -> list[ModuloCircuitElement]:
+        """
+        Negates a polynomial with coefficients `X`.
+        Returns R = [-x0, -x1, -x2, ... -xn-1] mod p
+        """
+        return [self.neg(x_i) for x_i in X]
+
+    def extf_sub(
+        self, X: list[ModuloCircuitElement], Y: list[ModuloCircuitElement]
+    ) -> list[ModuloCircuitElement]:
+        return self.extf_add(X, self.extf_neg(Y))
+
     def extf_mul(
         self,
         X: list[ModuloCircuitElement],
         Y: list[ModuloCircuitElement],
+        extension_degree: int,
     ) -> list[ModuloCircuitElement]:
         """
         Multiply in the extension field X * Y mod irreducible_poly
         Commit to R and accumulates Q.
         """
+        assert len(X) == len(Y) == extension_degree
+        Q, R = nondeterministic_extension_field_mul_divmod(
+            X, Y, self.curve_id, self.extension_degree
+        )
+        R, s0, s1 = self.write_commitments(R)
 
-        def commit_to_R_and_accumulate_Q(X, Y, Q, R):
-            Q, R = nondeterministic_extension_field_mul_divmod(
-                X, Y, self.curve_id, self.extension_degree
-            )
-            R = self.write_commitments(R)
-
-        (_, s1) = self.transcript.hash_limbs(R)
-        s1 = PyFelt(s1, self.field)
+        Q = Polynomial(Q)
+        s1 = self.field(s1)
         Q_acc: Polynomial = self.acc.nondeterministic_Q + s1 * Q
         s1 = self.write_cairo_native_felt(s1)
 
@@ -138,23 +162,101 @@ def commit_to_R_and_accumulate_Q(X, Y, Q, R):
         )
         return R
 
+    def extf_square(
+        self,
+        X: list[ModuloCircuitElement],
+        extension_degree: int,
+    ) -> list[ModuloCircuitElement]:
+        """
+        Multiply in the extension field X * Y mod irreducible_poly
+        Commit to R and accumulates Q.
+        """
+        assert len(X) == extension_degree
+        Q, R = nondeterministic_extension_field_mul_divmod(
+            X, X, self.curve_id, self.extension_degree
+        )
+        R, s0, s1 = self.write_commitments(R)
+
+        Q = Polynomial(Q)
+        s1 = self.field(s1)
+        Q_acc: Polynomial = self.acc.nondeterministic_Q + s1 * Q
+        s1 = self.write_cairo_native_felt(s1)
+
+        # Evaluate polynomials X(z), Y(z) inside circuit.
+        X_of_z = self.eval_poly_in_precomputed_Z(X)
+        XY_of_z = self.mul(X_of_z, X_of_z)
+        ci_XY_of_z = self.mul(s1, XY_of_z)
+        XY_acc = self.add(self.acc.xy, ci_XY_of_z)
+
+        # Computes R_acc = R_acc + s1 * R as a Polynomial inside circuit
+        R_acc = [self.add(r_acc, self.mul(s1, r)) for r_acc, r in zip(self.acc.R, R)]
+
+        self.acc = EuclideanPolyAccumulator(
+            xy=XY_acc, nondeterministic_Q=Q_acc, R=R_acc
+        )
+        return R
+
+    def extf_div(
+        self,
+        X: list[ModuloCircuitElement],
+        Y: list[ModuloCircuitElement],
+        extension_degree: int,
+    ) -> list[ModuloCircuitElement]:
+        assert len(X) == len(Y) == extension_degree
+        x_over_y = nondeterministic_extension_field_div(
+            X, Y, self.curve_id, extension_degree
+        )
+        add_to_transcript = True
+        x_over_y, _, _ = self.write_commitments(
+            x_over_y, add_to_transcript
+        )  # Is it really necessary to hash this in addition to what is hashed in extf_mul right after ?
+        should_be_X = self.extf_mul(x_over_y, Y, extension_degree)
+        self.extf_assert_eq(should_be_X, X)
+
+        return x_over_y
+
+    def extf_assert_eq(
+        self, X: list[ModuloCircuitElement], Y: list[ModuloCircuitElement]
+    ):
+        assert len(X) == len(Y)
+        for x, y in zip(X, Y):
+            self.assert_eq(x, y)
+
     def finalize_circuit(self):
+        # print("\n Finalize Circuit")
         Q = self.acc.nondeterministic_Q.get_coeffs()
         Q = Q + [self.field.zero()] * (self.extension_degree - 1 - len(Q))
-        Q = self.write_commitments(Q)
+        Q, s0, s1 = self.write_commitments(Q)
         Q_of_Z = self.eval_poly_in_precomputed_Z(Q)
         P, sparsity = self.write_sparse_elements(
-            get_irreducible_poly(self.curve_id, self.extension_degree).get_coeffs()
+            get_irreducible_poly(self.curve_id, self.extension_degree).get_coeffs(),
+            WriteOps.CONSTANT,
         )
-
         P_of_z = P[0]
-        for i in range(1, len(P)):
+        sparse_p_index = 1
+        for i in range(1, len(sparsity)):
             if sparsity[i] == 1:
-                P_of_z = self.add(P_of_z, self.mul(P[i], self.z_powers[i]))
+                P_of_z = self.add(
+                    P_of_z, self.mul(P[sparse_p_index], self.z_powers[i - 1])
+                )
+                sparse_p_index += 1
 
         R_of_Z = self.eval_poly_in_precomputed_Z(self.acc.R)
 
         lhs = self.acc.xy
-        rhs = self.add(Q_of_Z, self.mul(P_of_z, R_of_Z))
-        assert lhs.emulated_felt == rhs.emulated_felt
+        rhs = self.add(self.mul(Q_of_Z, P_of_z), R_of_Z)
+        assert lhs.value == rhs.value, f"{lhs.value} != {rhs.value}"
         return rhs
+
+    def summarize(self):
+        add_count, mul_count = self.values_segment.summarize()
+        summary = {
+            "circuit": self.name,
+            "MULMOD": mul_count,
+            "ADDMOD": add_count,
+            "POSEIDON": self.transcript.permutations_count,
+        }
+        # TODO : add Number of poseidon.
+
+        # pprint(summary, )
+        return summary