spec: Draft FMA extension field accelerator

spec/about_ecalls.typ

@@ -34,3 +34,6 @@ Negative numbers (represented as 2s complement 64-bit numbers), are used for our
 / -2: `KECCAK` (@keccak)
 / -11: `ECSM`/`secp256k1` (@ecsm)
 / -12: `ECSM`/`secp256r1` (@ecsm)
+/ -20: `FEXT_LOAD` (@fext)
+/ -21: `FEXT_FMA` (@fext)
+/ -22: `FEXT_ZERO` (@fext)

spec/book.typ

@@ -53,6 +53,7 @@
       ("sha256.typ", [`SHA256` accelerator], <sha256>),
       ("keccak.typ", [`KECCAK` accelerator], <keccak>),
       ("ecsm.typ", [`ECSM` accelerator], <ecsm>),
+      ("fext.typ", [Extension field accelerator], <fext>),
     )),
     ("MATHEMATICS", (
       ("limbs_and_carries.typ", [On limb decomposition and carries], <limbs>),

spec/fext.typ

@@ -0,0 +1,99 @@
+#import "/book.typ": book-page, aside
+#import "/src.typ": load_config, load_chip
+#import "/chip.typ": (
+  render_chip_variable_table,
+  total_nr_variables,
+  total_nr_instantiated_columns,
+  compute_nr_interactions,
+  render_constraint_table,
+  render_chip_assumptions,
+  render_chip_padding_table,
+)
+
+#show: book-page("fext.typ")
+
+#let config = load_config()
+#let loadchip = load_chip("src/fext_load.toml", config)
+#let load = raw(loadchip.name)
+#let fmachip = load_chip("src/fext_fma.toml", config)
+#let fma = raw(fmachip.name)
+#let zerochip = load_chip("src/fext_zero.toml", config)
+#let zero = raw(zerochip.name)
+
+We introduce a set of chips for faster processing of numbers mod the native goldilocks prime,
+or a degree three extension field thereof.
+Our approach is to offer an arithmetic black box (ABB), consisting of the #fma and #zero chips,
+that operates on a separate memory domain, and the #load chip to bridge the gap
+from normal byte-addressed RAM memory to this separate field-storage.
+As noted in @memory, we reserve the domain separator values $3$, $4$ and $5$ for field-storage.
+
+= The #load chip
+#let nr_variables = total_nr_variables(loadchip)
+#let nr_columns = total_nr_instantiated_columns(loadchip, config)
+#let nr_interactions = compute_nr_interactions(loadchip)
+
+We use the #load chip to load the three composing coefficients from registers A1-A3 (in little-endian),
+verify that all of them are in the correct range for a field element,
+and then write them as field elements into field-storage.
+We do this using #nr_variables variables spanning #nr_columns columns and #nr_interactions interactions.
+
+== Variables
+
+#render_chip_variable_table(loadchip, config)
+
+== Constraints
+
+#render_constraint_table(loadchip, config)
+
+== Padding
+
+#render_chip_padding_table(loadchip, config)
+
+= The #fma chip
+#let nr_variables = total_nr_variables(fmachip)
+#let nr_columns = total_nr_instantiated_columns(fmachip, config)
+#let nr_interactions = compute_nr_interactions(fmachip)
+
+To compute a fused multiply-add (FMA) operation `output = a * b + c`, we first load
+everything from memory, passing the input (ABB) addresses in A0-A2 and the output address in A3.
+Then we constrain the extension field operation, and write back to (ABB) memory.
+
+The extension field is expressed through three constant columns, $alpha$, $beta$ and $gamma$,
+such that the defining polynomial is $X^3 - alpha X^2 - beta X - gamma = 0$, or alternatively,
+$X^3 = alpha X^2 + beta X + gamma$.
+
+== Variables
+
+We express this chip using #nr_variables variables spanning #nr_columns columns and #nr_interactions interactions.
+
+#render_chip_variable_table(fmachip, config)
+
+== Constraints
+
+#render_constraint_table(fmachip, config)
+
+== Padding
+
+#render_chip_padding_table(fmachip, config)
+
+= The #zero chip
+#let nr_variables = total_nr_variables(zerochip)
+#let nr_columns = total_nr_instantiated_columns(zerochip, config)
+#let nr_interactions = compute_nr_interactions(zerochip)
+
+This chip asserts that a field extension element in the ABB, passed in the `A0` register, is zero.
+It uses #nr_variables variables spanning #nr_columns columns and #nr_interactions interactions.
+
+== Variables
+
+#render_chip_variable_table(zerochip, config)
+
+== Constraints
+
+Read the field element address from the `A0` register, and then assert that the field memory reads a zero.
+
+#render_constraint_table(zerochip, config)
+
+== Padding
+
+#render_chip_padding_table(zerochip, config)

spec/memory.typ

@@ -30,6 +30,7 @@ For specific domains and domain separators, we use the following assignment.
 / RAM memory: $0$
 / Registers: $1$
 / Committed values: $2$
+/ (Extension) field values: $3$, $4$, $5$
 
 On a high level, we ensure memory consistency by an interacting system of
 reads and writes to a lookup argument, combined with an initialization and finalization scheme.

spec/src/fext_fma.toml

@@ -0,0 +1,144 @@
+name = "FEXT_FMA"
+code = "FMA"
+
+[[variables.input]]
+name = "input_addrs"
+type = ["DWordWL", 3]
+desc = "The addresses of `values`"
+pad = 0
+
+[[variables.input]]
+name = "output_addr"
+type = "DWordWL"
+desc = "The address of `output`"
+pad = 0
+
+[[variables.input]]
+name = "timestamp"
+type = "Word"
+desc = "The timestamp"
+pad = 0
+
+[[variables.auxiliary]]
+name = "values"
+type = [["BaseField", 3], 3]
+desc = "The value to operate on: `a`, `b`, `c` in order to compute `a * b + c`"
+pad = 0
+
+[[variables.output]]
+name = "output"
+type = ["BaseField", 3]
+desc = "The output value"
+pad = 0
+
+[[variables.constant]]
+name = "α"
+type = "BaseField"
+desc = "The $#`X`^2$ coefficient of the definining polynomial"
+
+[[variables.constant]]
+name = "β"
+type = "BaseField"
+desc = "The $#`X`^1$ coefficient of the definining polynomial"
+
+[[variables.constant]]
+name = "γ"
+type = "BaseField"
+desc = "The $#`X`^0$ coefficient of the definining polynomial"
+
+[[variables.virtual]]
+name = "a"
+type = ["BaseField", 3]
+def = ["idx", "values", 0]
+desc = ""
+
+[[variables.virtual]]
+name = "b"
+type = ["BaseField", 3]
+def = ["idx", "values", 1]
+desc = ""
+
+[[variables.virtual]]
+name = "c"
+type = ["BaseField", 3]
+def = ["idx", "values", 2]
+desc = ""
+
+[[variables.multiplicity]]
+name = "μ"
+type = "Bit"
+desc = ""
+pad = 0
+
+[[constraint_groups]]
+name = "all"
+
+[[constraints.all]]
+kind = "template"
+tag = "REG"
+input = [["+", 10, "i"], ["idx", "input_addrs", "i"], "timestamp"]
+output = ["idx", "input_addrs", "i"]
+iter = ["i", 0, 2]
+cond = "μ"
+
+[[constraints.all]]
+kind = "template"
+tag = "REG"
+input = [13, "output_addr", "timestamp"]
+output = "output_addr"
+cond = "μ"
+
+[[constraints.all]]
+kind = "interaction"
+tag = "MEMW"
+input = [["+", 3, "d"], ["idx", "input_addrs", "i"], ["arr", ["idx", ["idx", "values", "i"], "d"], 0, 0, 0, 0, 0, 0, 0], ["+", "timestamp", "i"], 0, 0, 0]
+output = ["arr", ["idx", ["idx", "values", "i"], "d"], 0, 0, 0, 0, 0, 0, 0]
+iters = [["d", 0, 2], ["i", 0, 2]]
+multiplicity = "μ"
+
+[[constraints.all]]
+kind = "interaction"
+tag = "MEMW"
+input = [["+", 3, "d"], "output_addr", ["arr", ["idx", "output", "d"], 0, 0, 0, 0, 0, 0, 0], ["+", "timestamp", 3], 0, 0, 0]
+multiplicity = "μ"
+iter = ["d", 0, 2]
+
+[[constraints.all]]
+kind = "arith"
+constraint = "$#`output[0]` = a_0 b_0 + gamma (a_1 b_2 + a_2 b_1 + alpha dot a_2 b_2) + c_0$"
+poly = ["-", ["idx", "output", 0],
+  ["*", ["idx", "a", 0], ["idx", "b", 0]],
+  ["*", "γ", ["+",
+      ["*", ["idx", "a", 1], ["idx", "b", 2]],
+      ["*", ["idx", "a", 2], ["idx", "b", 1]],
+      ["*", "α", ["idx", "a", 2], ["idx", "b", 2]]]],
+  ["idx", "c", 0]]
+
+[[constraints.all]]
+kind = "arith"
+constraint = "$#`output[1]` = a_0 b_1 + a_1 b_0 + beta (a_1 b_2 + a_2 b_1) + (alpha beta + gamma) a_2 b_2 + c_1$"
+poly = ["-", ["idx", "output", 1],
+  ["*", ["idx", "a", 0], ["idx", "b", 1]],
+  ["*", ["idx", "a", 1], ["idx", "b", 0]],
+  ["*", "β", ["idx", "a", 1], ["idx", "b", 2]],
+  ["*", "β", ["idx", "a", 2], ["idx", "b", 1]],
+  ["*", ["+", ["*", "α", "β"], "γ"], ["idx", "a", 2], ["idx", "b", 2]],
+  ["idx", "c", 1]]
+
+[[constraints.all]]
+kind = "arith"
+constraint = "$#`output[2]` = a_0 b_2 + a_1 b_1 + a_2 b_0 + alpha (a_1 b_2 + a_2 b_1) + (alpha^2 + beta) a_2 b_2 + c_2$"
+poly = ["-", ["idx", "output", 2],
+  ["*", ["idx", "a", 0], ["idx", "b", 2]],
+  ["*", ["idx", "a", 1], ["idx", "b", 1]],
+  ["*", ["idx", "a", 2], ["idx", "b", 0]],
+  ["*", "α", ["idx", "a", 1], ["idx", "b", 2]],
+  ["*", "α", ["idx", "a", 2], ["idx", "b", 1]],
+  ["*", ["+", ["*", "α", "α"], "β"], ["idx", "a", 2], ["idx", "b", 2]],
+  ["idx", "c", 2]]
+
+[[constraints.all]]
+kind = "interaction"
+tag = "ECALL"
+input = ["timestamp", ["arr", ["-", ["^", 2, 32], 21], ["-", ["^", 2, 32], 1]]]
+multiplicity = ["-", "μ"]

spec/src/fext_load.toml

@@ -0,0 +1,67 @@
+name = "FEXT_LOAD"
+code = "FXL"
+
+[[variables.input]]
+name = "timestamp"
+type = "Word"
+desc = "The timestamp"
+pad = 0
+
+[[variables.input]]
+name = "addr"
+type = "DWordWL"
+desc = "The address at which to store the field element"
+pad = 0
+
+[[variables.input]]
+name = "coeffs"
+type = ["DWordWL", 3]
+desc = "The extension field element coefficients, in native form"
+pad = 0
+
+[[variables.multiplicity]]
+name = "μ"
+type = "Bit"
+desc = ""
+pad = 0
+
+[[constraint_groups]]
+name = "all"
+
+[[constraints.all]]
+kind = "template"
+tag = "REG"
+input = [10, "addr", "timestamp"]
+output = "addr"
+cond = "μ"
+
+[[constraints.all]]
+kind = "template"
+tag = "REG"
+input = [["+", 11, "i"], ["idx", "coeffs", "i"], "timestamp"]
+output = ["idx", "coeffs", "i"]
+cond = "μ"
+iter = ["i", 0, 2]
+
+[[constraints.all]]
+kind = "interaction"
+tag = "ALU"
+input = [["idx", "coeffs", "i"], ["arr", 1, ["-", ["^", 2, 32], 1]], ["opsel", "LT"]]
+output = ["cast", 1, "DWordWL"]
+iter = ["i", 0, 2]
+multiplicity = "μ"
+
+# This can either happen directly, or we introduce another dedicated MEMW accelerator?
+[[constraints.all]]
+kind = "interaction"
+tag = "MEMW"
+input = [["+", 3, "i"], "addr", ["arr", ["+", ["*", ["^", 2, 32], ["idx", ["idx", "coeffs", "i"], 1]], ["idx", ["idx", "coeffs", "i"], 0]], 0, 0, 0, 0, 0, 0, 0], "timestamp", 0, 0, 0]
+output = ["arr", ["+", ["*", ["^", 2, 32], ["idx", ["idx", "coeffs", "i"], 1]], ["idx", ["idx", "coeffs", "i"], 0]], 0, 0, 0, 0, 0, 0, 0]
+iter = ["i", 0, 2]
+multiplicity = "μ"
+
+[[constraints.all]]
+kind = "interaction"
+tag = "ECALL"
+input = ["timestamp", ["arr", ["-", ["^", 2, 32], 20], ["-", ["^", 2, 32], 1]]]
+multiplicity = ["-", "μ"]

spec/src/fext_zero.toml

@@ -0,0 +1,45 @@
+name = "FEXT_ZERO"
+code = "FXZ"
+
+[[variables.input]]
+name = "timestamp"
+type = "Word"
+desc = "The timestamp"
+pad = 0
+
+[[variables.input]]
+name = "addr"
+type = "DWordWL"
+desc = "The address to assert equal to zero"
+pad = 0
+
+[[variables.multiplicity]]
+name = "μ"
+type = "Bit"
+desc = ""
+pad = 0
+
+[[constraint_groups]]
+name = "all"
+
+[[constraints.all]]
+kind = "template"
+tag = "REG"
+input = [10, "addr", "timestamp"]
+output = "addr"
+cond = "μ"
+
+[[constraints.all]]
+kind = "interaction"
+tag = "MEMW"
+input = [["+", 3, "d"], "addr", ["cast", 0, ["BaseField", 8]], "timestamp", 0, 0, 0]
+output = ["cast", 0, ["BaseField", 8]]
+multiplicity = "μ"
+iter = ["d", 0, 2]
+
+[[constraints.all]]
+kind = "interaction"
+tag = "ECALL"
+input = ["timestamp", ["arr", ["-", ["^", 2, 32], 22], ["-", ["^", 2, 32], 1]]]
+multiplicity = "μ"
+

spec/tooling/chip.py

@@ -811,7 +811,7 @@ class Assumption:
     iters: list[Iter]
 
     def __init__(self, config: Config, data: dict):
-        assert_no_unexpected(data, set(self.__annotations__.keys()) | {"iter", "iters", "ref"})
+        assert_no_unexpected(data, set(type(self).__annotations__.keys()) | {"iter", "iters", "ref"})
         self.desc = data["desc"]
         self.iters = iters_of(data, config)
 
@@ -824,7 +824,7 @@ class ArithConstraint:
     iters: list[Iter]
 
     def __init__(self, config: Config, data: dict):
-        assert_no_unexpected(data, set(self.__annotations__.keys()) | {"kind", "ref", "iter", "iters"})
+        assert_no_unexpected(data, set(type(self).__annotations__.keys()) | {"kind", "ref", "iter", "iters"})
         assert data["kind"] == "arith"
         self.constraint = data["constraint"]
         reporter.asserts(