spec: KECCAK accelerator

spec/about_ecalls.typ

@@ -22,3 +22,13 @@ When `ECALL` is executed, it is assumed that:
 - the return value is written to `A0`,
 where `A0`-`A7` are symbolic names for the registers `x10`-`x17`
 #footnote([RISC-V - Register sets; en.wikipedia.org, #link("https://web.archive.org/web/20260209053447/https://en.wikipedia.org/wiki/RISC-V#Register_sets")[[src]]]).
+
+= ECALL number overview
+
+We provide a list of supported ECALL numbers.
+Negative numbers (represented as 2s complement 64-bit numbers), are used for our own custom accelerators/extensions.
+
+/ 64: `write` (@commit)
+/ 93: `exit` (@halt)
+/ -1: `SHA256` (@sha256)
+/ -2: `KECCAK` (@keccak)

spec/book.typ

@@ -43,6 +43,7 @@
       ("halt.typ", [`HALT` chip], <halt>),
       ("commit.typ", [`COMMIT` chip], <commit>),
       ("sha256.typ", [`SHA256` accelerator], <sha256>),
+      ("keccak.typ", [`KECCAK` accelerator], <keccak>),
     ))
   )
 )

spec/chip.typ

@@ -14,12 +14,15 @@
   .map(pair => pair.at(1))
   .flatten()
   .map(var => {
-    let (label, factor) = if type(var.type) == array {
-      (var.type.at(0), var.type.at(1))
-    } else {
-      (var.type, 1)
+    let (factor, var_type) = (1, var.type)
+    while type(var_type) == array {
+      assert(var_type.len() == 2, message: "invalid var (sub)type length: " + repr(var.type))
+      assert(type(var_type.at(1)) == int, message: "invalid var (sub)type length: " + repr(var.type))
+      factor *= var_type.at(1)
+      var_type = var_type.at(0)
     }
-    config.variables.types.filter(type => type.label == label).first().subtypes.len() * factor
+
+    config.variables.types.filter(type => type.label == var_type).first().subtypes.len() * factor
   })
   .sum()
 }
@@ -290,7 +293,7 @@
   } else if type(groups) == str {
     groups = (groups,)
   }
-  assert(groups.all(group => group in all_groups), message: "unknown group")
+  assert(groups.all(group => group in all_groups), message: "unknown group: " + repr(groups))
   let selected_constraints = groups.map(g => ((g): chip.constraints.at(g))).join()
 
   // Find the group definition in the constraint_groups

spec/expr.typ

@@ -1,20 +1,26 @@
 // Types and array types
 // <type> ::= str
-//          | [str, int]
+//          | [<type>, int]
 
 // Check that a type expression is structurally valid, without validating against a set of known base types
 #let check_array_type(typ) = {
-  assert(type(typ.at(0)) == str, message: "Array types need to have a regular type as base")
-  assert(type(typ.at(1)) == int, message: "Array types need to have a constant dimension")
+  while type(typ) == array {
+    assert(typ.len() == 2, message: "Array types must specify two parameters")
+    assert(type(typ.at(1)) == int, message: "Array types need to have a constant dimension")
+    typ = typ.at(0)
+  }
+  assert(type(typ) == str, message: "Array types need to have a regular type as base")
 }
 
 // Render a type to code
 #let type_to_code(typ) = {
-  if type(typ) == array {
-    check_array_type(typ)
-    return raw(typ.at(0) + "[" + str(typ.at(1)) + "]")
-  } else if type(typ) == str {
-    return raw(typ)
+  let label = ""
+  while type(typ) == array {
+    label += "[" + str(typ.at(1)) + "]"
+    typ = typ.at(0)
+  }
+  if type(typ) == str {
+    return raw(typ + label)
   } else {
     assert(false, message: "Unknown format for type: " + repr(typ))
   }
@@ -54,12 +60,13 @@
   "cast": 3, // cast
   "mul": 4,  // *
   "div": 5,  // /
-  "sum": 6,  // Σ
-  "not": 7,  // not
-  "sub": 8,  // -
-  "add": 9,  // +  
-  "eq": 10,   // = and :=
-  "MAX": 11, // <the void outside every expression>
+  "mod": 6,  // mod
+  "sum": 7,  // Σ
+  "not": 8,  // not
+  "sub": 9,  // -
+  "add": 10,  // +  
+  "eq": 11,   // = and :=
+  "MAX": 12, // <the void outside every expression>
 )
 
 // Mutual recursion through a trick from https://github.com/typst/typst/issues/744
@@ -97,6 +104,13 @@
     "not": (pp, rec, e) => cwrap(rec(PREC.not, 1) + ` - ` + rec(PREC.not, e.at(1)), pp < PREC.not),
     "+": (pp, rec, e) => cwrap(e.slice(1).map(rec.with(PREC.add)).join(` + `), pp < PREC.add),
     "sum": (pp, rec, e) => assert(false, message: "sum is unsupported in code."),
+    "mod": (pp, rec, e) => {
+      assert(e.len() == 3 and type(e.at(2)) == int, message: "Invalid mod expr: " + repr(e))
+      cwrap(
+        rec(PREC.mod, e.at(1)) + ` % ` + rec(PREC.mod, e.at(2)), 
+        pp <= PREC.mod
+      ) 
+    },
     "*": (pp, rec, e) => {
       if e.len() == 3 and type(e.at(1)) == int and type(e.at(2)) == str and e.at(2).len() == 1 {
         // multiplication of a constant with one-letter variable. 
@@ -165,6 +179,13 @@
         pp <= PREC.sub
       )
     },
+    "mod": (pp, rec, e) => {
+      assert(e.len() == 3 and type(e.at(2)) == int, message: "Invalid mod expr: " + repr(e))
+      mwrap(
+        $#rec(PREC.mod, e.at(1)) mod #rec(PREC.mod, e.at(2))$, 
+        pp <= PREC.mod
+      )
+    },
     "*": (pp, rec, e) => {
       if e.len() == 3 and type(e.at(1)) == int and type(e.at(2)) == str and e.at(2).len() == 1 {
         // multiplication of a constant with one-letter variable. 

spec/keccak.typ

@@ -0,0 +1,122 @@
+#import "/book.typ": book-page
+#import "/src.typ": load_config, load_chip
+#import "/chip.typ": (
+  compute_nr_interactions,
+  render_chip_assumptions,
+  render_chip_variable_table,
+  total_nr_variables,
+  total_nr_instantiated_columns,
+  render_constraint_table,
+  render_chip_padding_table,
+)
+
+#let config = load_config()
+#let chip = load_chip("src/keccak.toml", config)
+
+#show: book-page(chip.name)
+#let keccak = raw(chip.name)
+
+The #keccak chip applies the keccak permutation $kappa$ to a given memory range;
+other aspects of keccak hashing (such as repeated permutation invocation, 
+input padding and state initialization) fall outside the scope of this accelerator.
+
+This permutation $kappa: FF_2^1600 -> FF_2^1600$ operates on 1600 bits and is composed of 24 applications of round-permutation $Lambda: FF_2^1600 times NN -> FF_2^1600$, where the additional parameter is the round constant.
+$Lambda$ is defined as the composition $iota compose chi compose pi compose rho compose theta$, where only $iota$ depends on the round constant.
+#footnote("More details on the KECCAK permutation: FIPS 202, NIST, " + link("https://csrc.nist.gov/pubs/fips/202/final"))
+
+The keccak accelerator comprises two chips: a core chip that interacts with the memory --- loading the input and writing the output, and a round chip that applies the round permutation.
+
+
+= Core chip
+== Columns
+#let nr_variables = total_nr_variables(chip)
+#let nr_columns = total_nr_instantiated_columns(chip, config)
+#let nr_interactions = compute_nr_interactions(chip)
+
+The #keccak chip is comprised of #nr_variables variables that are expressed using #nr_columns columns and leverages #nr_interactions interaction(s):
+#render_chip_variable_table(chip, config)
+
+== Constraints
+In this VM, we assign syscall number -2 to the #keccak accelerator.
+The chip therefore contributes the following interaction to the lookup-argument:
+#render_constraint_table(chip, config, groups: "output")
+
+The address containing the state to be permuted is passed in as argument `A0 = x10`.
+The following constraints describe that this address is read into `addr` (@keccak:c:read_addr), from which `state_ptr` --- the collection of pointers to all lanes of the state --- is derived (@keccak:c:state_ptr).
+The state is then read into `input_state`, while the `output_state` is written back to the indicated address (@keccak:c:load_store_state).
+#render_constraint_table(chip, config, groups: "mem")
+
+Lastly, the input state is pushed to the Keccak-round function, while the output after 24 rounds is taken off the bus:
+#render_constraint_table(chip, config, groups: "round")
+
+== Padding
+The #keccak table can be padded to the next power of two with the following value assignments:
+#render_chip_padding_table(chip, config)
+
+= Round chip
+#let round_chip = load_chip("src/keccak_round.toml", config)
+#let keccak_rnd = raw(round_chip.name)
+
+== Columns
+#let nr_variables = total_nr_variables(round_chip)
+#let nr_columns = total_nr_instantiated_columns(round_chip, config)
+#let nr_interactions = compute_nr_interactions(round_chip)
+
+The #keccak_rnd chip is comprised of #nr_variables variables that are expressed using #nr_columns columns and leverages #nr_interactions interaction(s):
+#render_chip_variable_table(round_chip, config)
+
+#strong("Note on " + raw("start") + ".")
+`start` contains the state to which the permutation should be applied.
+Its three-dimensional array mimics the specification's three-dimensional state
+#footnote("FIPS 202, NIST, Section 3.1 (" + link("https://csrc.nist.gov/pubs/fips/202/final") + ")")
+and orders the bits as prescribed.
+#footnote("FIPS 202, NIST, Section B.1, Algorithm 10 (" + link("https://csrc.nist.gov/pubs/fips/202/final") + ")")
+
+#strong("Note on " + raw("rnc") + " and " + raw("rbc") + ".")
+Rho rotates every lane by a rotation offset in $[0, 64)$.
+These offsets are identical for every round.
+#footnote("FIPS 202, NIST, page 13, Table 2 (" + link("https://csrc.nist.gov/pubs/fips/202/final") + ")")
+We decompose each offset in three components: the lower nibble (4 bits) are represented by `rnc`, while the upper two bits are represented by as `Bit`s in `rbc`.
+That is, $#`rho_offset[x][y]` = #`rnc[x][y]` + 16 dot #`rbc[x][y][0]` + 32 dot #`rbc[x][y][1]`$.
+
+
+== Constraints
+
+The following constraints ensure that `theta` captures the state after applying the first subpermutation of the round-permutation: $theta$.
+Note here that `Cxz_left` and `Cxz_right` do have to be range-checked; it cannot be assumed that this implicitly follows from @keccak:c:Dxz combined with `rotated_Cxz`'s definition.
+#render_constraint_table(round_chip, config, groups: "theta")
+
+Next, we constrain that `rho` captures the state after applying subpermutation $rho$.
+Note here as well that `rot_left` and `rot_right` do have to be range-checked; it cannot be assumed that this implicitly follows from later constraints.
+#render_constraint_table(round_chip, config, groups: "rho")
+
+Observe that the lane-permutation performed by $pi$ is absorbed in `pi`'s definition.
+The next permutation that is constrained in $chi$:
+#render_constraint_table(round_chip, config, groups: "chi")
+
+Lastly, the round constants are added to one of the lanes in the state.
+`iota` contains the updated lane.
+In the definition of `out`, the output of `chi` and `iota` is combined to construct the output of the permutation.
+#render_constraint_table(round_chip, config, groups: "iota")
+
+Lastly, the round chip contributes the following interactions to the lookup:
+#render_constraint_table(round_chip, config, groups: "io")
+
+== Notes/potential optimizations
+- one does not have to repeat `addr` in `state_ptr`; this saves 4 columns and 4 `IS_HALF` checks.
+- step $rho$ does not need to be applied to `state[0][0]`; its has a zero-shift. This saves 16 columns and 4 `HWSL` interactions.
+- $#`rc[2]` = #`rc[4]` = #`rc[5]` = #`rc[6]` = 0$. As such, those elements need not be stored in `rc`, and need not be XORed into the state in the $iota$-step. This saves 8 columns and 4 `XOR_BYTE` interactions.
+- when executed in large volumnes, `KECCAK_RND` could benefit from having a three-way XOR lookup table. With this in place, the 80 interactions in @keccak:c:theta_cxz_start and @keccak:c:theta_cxz could be dropped.
+  Likewise, 80 columns could be removed from the chip (a \~5% savings).
+
+= Round constant lookup
+#let rc_chip = load_chip("src/keccak_rc.toml", config)
+#let keccak_rc = raw(rc_chip.name)
+
+== Columns
+#let nr_variables = total_nr_variables(rc_chip)
+#let nr_columns = total_nr_instantiated_columns(rc_chip, config)
+
+We provide the round constants through a short precomputed lookup table: #keccak_rc.
+#render_chip_variable_table(rc_chip, config)
+#render_constraint_table(rc_chip, config)

spec/signatures.typ

@@ -1,5 +1,6 @@
 #import "/book.typ": book-page
 #import "/src.typ": load_signatures, load_config
+#import "/expr.typ": type_to_code
 
 #show: book-page("signatures.typ")
 
@@ -19,21 +20,11 @@
     raw(cond) + ` => `
   } else {``}
 
-  let input_str = sig.input.map(elt => {
-    if type(elt) == array {
-      raw(elt.at(0)) + `[` + raw(str(elt.at(1))) + `]`
-    } else {
-      raw(elt)
-    }
-  }).join(`, `)
+  let input_str = sig.input.map(type_to_code).join(`, `)
 
   let output = sig.at("output", default: none)
   let output_str = if output != none {
-    if type(output) == array {
-      raw(output.at(0)) + `[` + raw(str(output.at(1))) + `]`
-    } else {
-      raw(output)
-    } + `; `
+    type_to_code(output) + `; `
   } else {``}
 
   return [#cond_str#raw(sig.tag)#lb#output_str#input_str#rb]
@@ -44,12 +35,13 @@
   let vars = sig.input + if "output" in sig { (sig.output, )} else {()}
 
   return vars.map(v => {
-    let (label, factor) = if type(v) == array {
-      (v.at(0), v.at(1))
-    } else {
-      (v, 1)
+    let factor = 1
+    while type(v) == array {
+      factor *= v.at(1)
+      v = v.at(0)
     }
-    config.variables.types.filter(type => type.label == label).first().subtypes.len() * factor
+    let lbl = v
+    config.variables.types.filter(type => type.label == lbl).first().subtypes.len() * factor
   })
   .sum()
 }

spec/src.typ

@@ -40,10 +40,11 @@
 
   // Verify that `var` is a valid variable.
   let verify_variable(var) = {
-    if type(var) == array {
-        assert(var.at(0) in var_labels, message: "Invalid var type: " + repr(var))
+    while type(var) == array {
         assert(type(var.at(1)) == int, message: "Invalid var type: " + repr(var))
-    } else if type(var) == str {
+        var = var.at(0)
+    }
+    if type(var) == str {
       assert(var in var_labels, message: "Invalid var type: " + repr(var))
     } else {
       assert(false, message: "Invalid var type: " + repr(var))
@@ -104,14 +105,13 @@
   let all_vars = chip.variables.values().flatten()
   let all_labels = config.variables.types.map(type => type.label);
   for var in all_vars {
-    let type_label = if type(var.type) == array {
-      var.type.at(0)
-    } else {
-      var.type
+    let type_label = var.type
+    while type(type_label) == array {
+      assert(type_label.len() == 2 and type(type_label.at(1)) == int, message: "invalid type: " + repr(var.type))
+      type_label = type_label.at(0)
     }
-
     // Check that all variable types are valid
-    assert(type_label in all_labels, message: "found invalid var type:" + repr(var.type))
+    assert(type_label in all_labels, message: "found invalid var type: " + repr(var.type))
   }
 }