spec: Draft FMA extension field accelerator

spec/about_ecalls.typ

@@ -31,4 +31,6 @@ Negative numbers (represented as 2s complement 64-bit numbers), are used for our
 / 64: `write` (@commit)
 / 93: `exit` (@halt)
 / -1: `SHA256` (@sha256)
-/ -2: `KECCAK` (@keccak)
+/ -2: `KECCAK` (@keccak)
+/-20: `FEXT_LOAD` (@fext)
+/-21: `FEXT_FMA` (@fext)

spec/book.typ

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

spec/fext.typ

@@ -0,0 +1,75 @@
+#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)
+
+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 off an arithmetic black box, consisting of the *TODO* 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)

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,143 @@
+name = "FEXT_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 = "interaction"
+tag = "MEMW"
+input = [1, ["cast", ["*", 2, ["+", 10, "i"]], "DWordWL"], ["arr", ["idx", ["idx", "input_addrs", "i"], 0], ["idx", ["idx", "input_addrs", "i"], 1], 0, 0, 0, 0, 0, 0], "timestamp", 1, 0, 0]
+output  = ["arr", ["idx", ["idx", "input_addrs", "i"], 0], ["idx", ["idx", "input_addrs", "i"], 1], 0, 0, 0, 0, 0, 0]
+iter = ["i", 0, 2]
+multiplicity = "μ"
+
+[[constraints.all]]
+kind = "interaction"
+tag = "MEMW"
+input = [1, ["cast", ["*", 2, 13], "DWordWL"], ["arr", ["idx", "output_addr", 0], ["idx", "output_addr", 1], 0, 0, 0, 0, 0, 0], "timestamp", 1, 0, 0]
+output  = ["arr", ["idx", "output_addr", 0], ["idx", "output_addr", 1], 0, 0, 0, 0, 0, 0]
+multiplicity = "μ"
+
+[[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", 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", 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 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$"
+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$"
+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,66 @@
+name = "FEXT_LOAD"
+
+[[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 = "interaction"
+tag = "MEMW"
+input = [1, ["cast", ["*", 2, 10], "DWordWL"], ["arr", ["idx", "addr", 0], ["idx", "addr", 1], 0, 0, 0, 0, 0, 0], "timestamp", 1, 0, 0]
+output = ["arr", ["idx", "addr", 0], ["idx", "addr", 1], 0, 0, 0, 0, 0, 0]
+multiplicity = "μ"
+
+[[constraints.all]]
+kind = "interaction"
+tag = "MEMW"
+input = [1, ["cast", ["*", 2, ["+", 11, "i"]], "DWordWL"], ["arr", ["idx", ["idx", "coeffs", "i"], 0], ["idx", ["idx", "coeffs", "i"], 1], 0, 0, 0, 0, 0, 0], "timestamp", 1, 0, 0]
+output = ["arr", ["idx", ["idx", "coeffs", "i"], 0], ["idx", ["idx", "coeffs", "i"], 1], 0, 0, 0, 0, 0, 0]
+multiplicity = "μ"
+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 = ["-", "μ"]