From 645d45f227a2c1537529a544f625ede6ca964bc2 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Thu, 10 Feb 2022 12:05:04 -0800 Subject: [PATCH] Column definitions for addition, range checks & lookups (#477) * Column definitions for addition, range checks & lookups I implemented addition (unsigned for now) as an example of how the arithmetic unit can interact with the 16-bit range check unit. Range checks and lookups aren't implemented yet. * Missing constraints * Tweaks to get tests passing * Reorg registers into files * Minor --- field/src/field_types.rs | 6 + plonky2/src/plonk/plonk_common.rs | 2 +- starky/src/lib.rs | 2 - system_zero/src/arithmetic/addition.rs | 70 ++++++++++++ system_zero/src/arithmetic/division.rs | 31 +++++ system_zero/src/arithmetic/mod.rs | 75 ++++++++++++ system_zero/src/arithmetic/multiplication.rs | 31 +++++ system_zero/src/arithmetic/subtraction.rs | 31 +++++ system_zero/src/column_layout.rs | 108 ------------------ system_zero/src/core_registers.rs | 39 +++++-- system_zero/src/lib.rs | 5 +- system_zero/src/permutation_unit.rs | 11 +- system_zero/src/registers/arithmetic.rs | 37 ++++++ system_zero/src/registers/boolean.rs | 10 ++ system_zero/src/registers/core.rs | 20 ++++ system_zero/src/registers/logic.rs | 3 + system_zero/src/registers/lookup.rs | 21 ++++ system_zero/src/registers/memory.rs | 3 + system_zero/src/registers/mod.rs | 20 ++++ system_zero/src/registers/permutation.rs | 57 +++++++++ system_zero/src/registers/range_check_16.rs | 11 ++ .../src/registers/range_check_degree.rs | 11 ++ system_zero/src/system_zero.rs | 20 +++- 23 files changed, 489 insertions(+), 135 deletions(-) create mode 100644 system_zero/src/arithmetic/addition.rs create mode 100644 system_zero/src/arithmetic/division.rs create mode 100644 system_zero/src/arithmetic/mod.rs create mode 100644 system_zero/src/arithmetic/multiplication.rs create mode 100644 system_zero/src/arithmetic/subtraction.rs delete mode 100644 system_zero/src/column_layout.rs create mode 100644 system_zero/src/registers/arithmetic.rs create mode 100644 system_zero/src/registers/boolean.rs create mode 100644 system_zero/src/registers/core.rs create mode 100644 system_zero/src/registers/logic.rs create mode 100644 system_zero/src/registers/lookup.rs create mode 100644 system_zero/src/registers/memory.rs create mode 100644 system_zero/src/registers/mod.rs create mode 100644 system_zero/src/registers/permutation.rs create mode 100644 system_zero/src/registers/range_check_16.rs create mode 100644 system_zero/src/registers/range_check_degree.rs diff --git a/field/src/field_types.rs b/field/src/field_types.rs index 95696475..83826b9f 100644 --- a/field/src/field_types.rs +++ b/field/src/field_types.rs @@ -278,6 +278,12 @@ pub trait Field: Self::from_canonical_u64(n as u64) } + /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()`. + // TODO: Should probably be unsafe. + fn from_canonical_u16(n: u16) -> Self { + Self::from_canonical_u64(n as u64) + } + /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()`. // TODO: Should probably be unsafe. fn from_canonical_usize(n: usize) -> Self { diff --git a/plonky2/src/plonk/plonk_common.rs b/plonky2/src/plonk/plonk_common.rs index 519593b3..09cf2652 100644 --- a/plonky2/src/plonk/plonk_common.rs +++ b/plonky2/src/plonk/plonk_common.rs @@ -138,7 +138,7 @@ where sum } -pub(crate) fn reduce_with_powers_ext_recursive, const D: usize>( +pub fn reduce_with_powers_ext_recursive, const D: usize>( builder: &mut CircuitBuilder, terms: &[ExtensionTarget], alpha: Target, diff --git a/starky/src/lib.rs b/starky/src/lib.rs index dc61e7e7..eefab529 100644 --- a/starky/src/lib.rs +++ b/starky/src/lib.rs @@ -1,8 +1,6 @@ // TODO: Remove these when crate is closer to being finished. #![allow(dead_code)] #![allow(unused_variables)] -#![allow(unreachable_code)] -#![allow(clippy::diverging_sub_expression)] #![allow(incomplete_features)] #![feature(generic_const_exprs)] diff --git a/system_zero/src/arithmetic/addition.rs b/system_zero/src/arithmetic/addition.rs new file mode 100644 index 00000000..653d533b --- /dev/null +++ b/system_zero/src/arithmetic/addition.rs @@ -0,0 +1,70 @@ +use plonky2::field::extension_field::Extendable; +use plonky2::field::field_types::Field; +use plonky2::field::packed_field::PackedField; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::ext_target::ExtensionTarget; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2::plonk::plonk_common::reduce_with_powers_ext_recursive; +use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; + +use crate::registers::arithmetic::*; +use crate::registers::NUM_COLUMNS; + +pub(crate) fn generate_addition(values: &mut [F; NUM_COLUMNS]) { + let in_1 = values[COL_ADD_INPUT_1].to_canonical_u64(); + let in_2 = values[COL_ADD_INPUT_2].to_canonical_u64(); + let in_3 = values[COL_ADD_INPUT_3].to_canonical_u64(); + let output = in_1 + in_2 + in_3; + + values[COL_ADD_OUTPUT_1] = F::from_canonical_u16(output as u16); + values[COL_ADD_OUTPUT_2] = F::from_canonical_u16((output >> 16) as u16); + values[COL_ADD_OUTPUT_3] = F::from_canonical_u16((output >> 32) as u16); +} + +pub(crate) fn eval_addition>( + local_values: &[P; NUM_COLUMNS], + yield_constr: &mut ConstraintConsumer

, +) { + let is_add = local_values[IS_ADD]; + let in_1 = local_values[COL_ADD_INPUT_1]; + let in_2 = local_values[COL_ADD_INPUT_2]; + let in_3 = local_values[COL_ADD_INPUT_3]; + let out_1 = local_values[COL_ADD_OUTPUT_1]; + let out_2 = local_values[COL_ADD_OUTPUT_2]; + let out_3 = local_values[COL_ADD_OUTPUT_3]; + + let weight_2 = F::from_canonical_u64(1 << 16); + let weight_3 = F::from_canonical_u64(1 << 32); + // Note that this can't overflow. Since each output limb has been range checked as 16-bits, + // this sum can be around 48 bits at most. + let out = out_1 + out_2 * weight_2 + out_3 * weight_3; + + let computed_out = in_1 + in_2 + in_3; + + yield_constr.constraint_wrapping(is_add * (out - computed_out)); +} + +pub(crate) fn eval_addition_recursively, const D: usize>( + builder: &mut CircuitBuilder, + local_values: &[ExtensionTarget; NUM_COLUMNS], + yield_constr: &mut RecursiveConstraintConsumer, +) { + let is_add = local_values[IS_ADD]; + let in_1 = local_values[COL_ADD_INPUT_1]; + let in_2 = local_values[COL_ADD_INPUT_2]; + let in_3 = local_values[COL_ADD_INPUT_3]; + let out_1 = local_values[COL_ADD_OUTPUT_1]; + let out_2 = local_values[COL_ADD_OUTPUT_2]; + let out_3 = local_values[COL_ADD_OUTPUT_3]; + + let limb_base = builder.constant(F::from_canonical_u64(1 << 16)); + // Note that this can't overflow. Since each output limb has been range checked as 16-bits, + // this sum can be around 48 bits at most. + let out = reduce_with_powers_ext_recursive(builder, &[out_1, out_2, out_3], limb_base); + + let computed_out = builder.add_many_extension(&[in_1, in_2, in_3]); + + let diff = builder.sub_extension(out, computed_out); + let filtered_diff = builder.mul_extension(is_add, diff); + yield_constr.constraint_wrapping(builder, filtered_diff); +} diff --git a/system_zero/src/arithmetic/division.rs b/system_zero/src/arithmetic/division.rs new file mode 100644 index 00000000..2f15b233 --- /dev/null +++ b/system_zero/src/arithmetic/division.rs @@ -0,0 +1,31 @@ +use plonky2::field::extension_field::Extendable; +use plonky2::field::field_types::Field; +use plonky2::field::packed_field::PackedField; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::ext_target::ExtensionTarget; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; + +use crate::registers::arithmetic::*; +use crate::registers::NUM_COLUMNS; + +pub(crate) fn generate_division(values: &mut [F; NUM_COLUMNS]) { + // TODO +} + +pub(crate) fn eval_division>( + local_values: &[P; NUM_COLUMNS], + yield_constr: &mut ConstraintConsumer

, +) { + let is_div = local_values[IS_DIV]; + // TODO +} + +pub(crate) fn eval_division_recursively, const D: usize>( + builder: &mut CircuitBuilder, + local_values: &[ExtensionTarget; NUM_COLUMNS], + yield_constr: &mut RecursiveConstraintConsumer, +) { + let is_div = local_values[IS_DIV]; + // TODO +} diff --git a/system_zero/src/arithmetic/mod.rs b/system_zero/src/arithmetic/mod.rs new file mode 100644 index 00000000..c635d58d --- /dev/null +++ b/system_zero/src/arithmetic/mod.rs @@ -0,0 +1,75 @@ +use plonky2::field::extension_field::Extendable; +use plonky2::field::field_types::Field; +use plonky2::field::packed_field::PackedField; +use plonky2::hash::hash_types::RichField; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; +use starky::vars::StarkEvaluationTargets; +use starky::vars::StarkEvaluationVars; + +use crate::arithmetic::addition::{eval_addition, eval_addition_recursively, generate_addition}; +use crate::arithmetic::division::{eval_division, eval_division_recursively, generate_division}; +use crate::arithmetic::multiplication::{ + eval_multiplication, eval_multiplication_recursively, generate_multiplication, +}; +use crate::arithmetic::subtraction::{ + eval_subtraction, eval_subtraction_recursively, generate_subtraction, +}; +use crate::public_input_layout::NUM_PUBLIC_INPUTS; +use crate::registers::arithmetic::*; +use crate::registers::NUM_COLUMNS; + +mod addition; +mod division; +mod multiplication; +mod subtraction; + +pub(crate) fn generate_arithmetic_unit(values: &mut [F; NUM_COLUMNS]) { + if values[IS_ADD].is_one() { + generate_addition(values); + } else if values[IS_SUB].is_one() { + generate_subtraction(values); + } else if values[IS_MUL].is_one() { + generate_multiplication(values); + } else if values[IS_DIV].is_one() { + generate_division(values); + } +} + +pub(crate) fn eval_arithmetic_unit>( + vars: StarkEvaluationVars, + yield_constr: &mut ConstraintConsumer

, +) { + let local_values = &vars.local_values; + + // Check that the operation flag values are binary. + for col in [IS_ADD, IS_SUB, IS_MUL, IS_DIV] { + let val = local_values[col]; + yield_constr.constraint_wrapping(val * val - val); + } + + eval_addition(local_values, yield_constr); + eval_subtraction(local_values, yield_constr); + eval_multiplication(local_values, yield_constr); + eval_division(local_values, yield_constr); +} + +pub(crate) fn eval_arithmetic_unit_recursively, const D: usize>( + builder: &mut CircuitBuilder, + vars: StarkEvaluationTargets, + yield_constr: &mut RecursiveConstraintConsumer, +) { + let local_values = &vars.local_values; + + // Check that the operation flag values are binary. + for col in [IS_ADD, IS_SUB, IS_MUL, IS_DIV] { + let val = local_values[col]; + let constraint = builder.mul_add_extension(val, val, val); + yield_constr.constraint_wrapping(builder, constraint); + } + + eval_addition_recursively(builder, local_values, yield_constr); + eval_subtraction_recursively(builder, local_values, yield_constr); + eval_multiplication_recursively(builder, local_values, yield_constr); + eval_division_recursively(builder, local_values, yield_constr); +} diff --git a/system_zero/src/arithmetic/multiplication.rs b/system_zero/src/arithmetic/multiplication.rs new file mode 100644 index 00000000..2eefad38 --- /dev/null +++ b/system_zero/src/arithmetic/multiplication.rs @@ -0,0 +1,31 @@ +use plonky2::field::extension_field::Extendable; +use plonky2::field::field_types::Field; +use plonky2::field::packed_field::PackedField; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::ext_target::ExtensionTarget; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; + +use crate::registers::arithmetic::*; +use crate::registers::NUM_COLUMNS; + +pub(crate) fn generate_multiplication(values: &mut [F; NUM_COLUMNS]) { + // TODO +} + +pub(crate) fn eval_multiplication>( + local_values: &[P; NUM_COLUMNS], + yield_constr: &mut ConstraintConsumer

, +) { + let is_mul = local_values[IS_MUL]; + // TODO +} + +pub(crate) fn eval_multiplication_recursively, const D: usize>( + builder: &mut CircuitBuilder, + local_values: &[ExtensionTarget; NUM_COLUMNS], + yield_constr: &mut RecursiveConstraintConsumer, +) { + let is_mul = local_values[IS_MUL]; + // TODO +} diff --git a/system_zero/src/arithmetic/subtraction.rs b/system_zero/src/arithmetic/subtraction.rs new file mode 100644 index 00000000..3613dee6 --- /dev/null +++ b/system_zero/src/arithmetic/subtraction.rs @@ -0,0 +1,31 @@ +use plonky2::field::extension_field::Extendable; +use plonky2::field::field_types::Field; +use plonky2::field::packed_field::PackedField; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::ext_target::ExtensionTarget; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; + +use crate::registers::arithmetic::*; +use crate::registers::NUM_COLUMNS; + +pub(crate) fn generate_subtraction(values: &mut [F; NUM_COLUMNS]) { + // TODO +} + +pub(crate) fn eval_subtraction>( + local_values: &[P; NUM_COLUMNS], + yield_constr: &mut ConstraintConsumer

, +) { + let is_sub = local_values[IS_SUB]; + // TODO +} + +pub(crate) fn eval_subtraction_recursively, const D: usize>( + builder: &mut CircuitBuilder, + local_values: &[ExtensionTarget; NUM_COLUMNS], + yield_constr: &mut RecursiveConstraintConsumer, +) { + let is_sub = local_values[IS_SUB]; + // TODO +} diff --git a/system_zero/src/column_layout.rs b/system_zero/src/column_layout.rs deleted file mode 100644 index fa5d627a..00000000 --- a/system_zero/src/column_layout.rs +++ /dev/null @@ -1,108 +0,0 @@ -//// CORE REGISTERS - -/// A cycle counter. Starts at 0; increments by 1. -pub(crate) const COL_CLOCK: usize = 0; - -/// A column which contains the values `[0, ... 2^16 - 1]`, potentially with duplicates. Used for -/// 16-bit range checks. -/// -/// For ease of verification, we enforce that it must begin with 0 and end with `2^16 - 1`, and each -/// delta must be either 0 or 1. -pub(crate) const COL_RANGE_16: usize = COL_CLOCK + 1; - -/// Pointer to the current instruction. -pub(crate) const COL_INSTRUCTION_PTR: usize = COL_RANGE_16 + 1; -/// Pointer to the base of the current call's stack frame. -pub(crate) const COL_FRAME_PTR: usize = COL_INSTRUCTION_PTR + 1; -/// Pointer to the tip of the current call's stack frame. -pub(crate) const COL_STACK_PTR: usize = COL_FRAME_PTR + 1; - -//// PERMUTATION UNIT -pub(crate) mod permutation { - use plonky2::hash::hashing::SPONGE_WIDTH; - use plonky2::hash::poseidon; - - const START_UNIT: usize = super::COL_STACK_PTR + 1; - - const START_FULL_FIRST: usize = START_UNIT + SPONGE_WIDTH; - - pub const fn col_full_first_mid_sbox(round: usize, i: usize) -> usize { - debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS); - debug_assert!(i < SPONGE_WIDTH); - START_FULL_FIRST + 2 * round * SPONGE_WIDTH + i - } - - pub const fn col_full_first_after_mds(round: usize, i: usize) -> usize { - debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS); - debug_assert!(i < SPONGE_WIDTH); - START_FULL_FIRST + (2 * round + 1) * SPONGE_WIDTH + i - } - - const START_PARTIAL: usize = - col_full_first_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, SPONGE_WIDTH - 1) + 1; - - pub const fn col_partial_mid_sbox(round: usize) -> usize { - debug_assert!(round < poseidon::N_PARTIAL_ROUNDS); - START_PARTIAL + 2 * round - } - - pub const fn col_partial_after_sbox(round: usize) -> usize { - debug_assert!(round < poseidon::N_PARTIAL_ROUNDS); - START_PARTIAL + 2 * round + 1 - } - - const START_FULL_SECOND: usize = col_partial_after_sbox(poseidon::N_PARTIAL_ROUNDS - 1) + 1; - - pub const fn col_full_second_mid_sbox(round: usize, i: usize) -> usize { - debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS); - debug_assert!(i < SPONGE_WIDTH); - START_FULL_SECOND + 2 * round * SPONGE_WIDTH + i - } - - pub const fn col_full_second_after_mds(round: usize, i: usize) -> usize { - debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS); - debug_assert!(i < SPONGE_WIDTH); - START_FULL_SECOND + (2 * round + 1) * SPONGE_WIDTH + i - } - - pub const fn col_input(i: usize) -> usize { - debug_assert!(i < SPONGE_WIDTH); - START_UNIT + i - } - - pub const fn col_output(i: usize) -> usize { - debug_assert!(i < SPONGE_WIDTH); - col_full_second_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, i) - } - - pub(super) const END_UNIT: usize = col_output(SPONGE_WIDTH - 1); -} - -//// MEMORY UNITS - -//// DECOMPOSITION UNITS -pub(crate) mod decomposition { - - const START_UNITS: usize = super::permutation::END_UNIT + 1; - - const NUM_UNITS: usize = 4; - /// The number of bits associated with a single decomposition unit. - const UNIT_BITS: usize = 32; - /// One column for the value being decomposed, plus one column per bit. - const UNIT_COLS: usize = 1 + UNIT_BITS; - - pub const fn col_input(unit: usize) -> usize { - debug_assert!(unit < NUM_UNITS); - START_UNITS + unit * UNIT_COLS - } - - pub const fn col_bit(unit: usize, bit: usize) -> usize { - debug_assert!(unit < NUM_UNITS); - debug_assert!(bit < UNIT_BITS); - START_UNITS + unit * UNIT_COLS + 1 + bit - } - - pub(super) const END_UNITS: usize = START_UNITS + UNIT_COLS * NUM_UNITS; -} - -pub(crate) const NUM_COLUMNS: usize = decomposition::END_UNITS; diff --git a/system_zero/src/core_registers.rs b/system_zero/src/core_registers.rs index 21faa288..03e7fa04 100644 --- a/system_zero/src/core_registers.rs +++ b/system_zero/src/core_registers.rs @@ -6,10 +6,9 @@ use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsume use starky::vars::StarkEvaluationTargets; use starky::vars::StarkEvaluationVars; -use crate::column_layout::{ - COL_CLOCK, COL_FRAME_PTR, COL_INSTRUCTION_PTR, COL_RANGE_16, COL_STACK_PTR, NUM_COLUMNS, -}; use crate::public_input_layout::NUM_PUBLIC_INPUTS; +use crate::registers::core::*; +use crate::registers::NUM_COLUMNS; use crate::system_zero::SystemZero; impl, const D: usize> SystemZero { @@ -35,11 +34,11 @@ impl, const D: usize> SystemZero { let next_range_16 = (prev_range_16 + 1).min((1 << 16) - 1); next_values[COL_RANGE_16] = F::from_canonical_u64(next_range_16); - next_values[COL_INSTRUCTION_PTR] = todo!(); + // next_values[COL_INSTRUCTION_PTR] = todo!(); - next_values[COL_FRAME_PTR] = todo!(); + // next_values[COL_FRAME_PTR] = todo!(); - next_values[COL_STACK_PTR] = todo!(); + // next_values[COL_STACK_PTR] = todo!(); } #[inline] @@ -64,9 +63,9 @@ impl, const D: usize> SystemZero { let delta_range_16 = next_range_16 - local_range_16; yield_constr.constraint_first_row(local_range_16); yield_constr.constraint_last_row(local_range_16 - FE::from_canonical_u64((1 << 16) - 1)); - yield_constr.constraint(delta_range_16 * (delta_range_16 - FE::ONE)); + yield_constr.constraint(delta_range_16 * delta_range_16 - delta_range_16); - todo!() + // TODO constraints for stack etc. } pub(crate) fn eval_core_registers_recursively( @@ -75,6 +74,28 @@ impl, const D: usize> SystemZero { vars: StarkEvaluationTargets, yield_constr: &mut RecursiveConstraintConsumer, ) { - todo!() + let one_ext = builder.one_extension(); + let max_u16 = builder.constant(F::from_canonical_u64((1 << 16) - 1)); + let max_u16_ext = builder.convert_to_ext(max_u16); + + // The clock must start with 0, and increment by 1. + let local_clock = vars.local_values[COL_CLOCK]; + let next_clock = vars.next_values[COL_CLOCK]; + let delta_clock = builder.sub_extension(next_clock, local_clock); + yield_constr.constraint_first_row(builder, local_clock); + let constraint = builder.sub_extension(delta_clock, one_ext); + yield_constr.constraint(builder, constraint); + + // The 16-bit table must start with 0, end with 2^16 - 1, and increment by 0 or 1. + let local_range_16 = vars.local_values[COL_RANGE_16]; + let next_range_16 = vars.next_values[COL_RANGE_16]; + let delta_range_16 = builder.sub_extension(next_range_16, local_range_16); + yield_constr.constraint_first_row(builder, local_range_16); + let constraint = builder.sub_extension(local_range_16, max_u16_ext); + yield_constr.constraint_last_row(builder, constraint); + let constraint = builder.mul_add_extension(delta_range_16, delta_range_16, delta_range_16); + yield_constr.constraint(builder, constraint); + + // TODO constraints for stack etc. } } diff --git a/system_zero/src/lib.rs b/system_zero/src/lib.rs index 029c2abd..1c097573 100644 --- a/system_zero/src/lib.rs +++ b/system_zero/src/lib.rs @@ -1,12 +1,11 @@ // TODO: Remove these when crate is closer to being finished. #![allow(dead_code)] #![allow(unused_variables)] -#![allow(unreachable_code)] -#![allow(clippy::diverging_sub_expression)] -mod column_layout; +mod arithmetic; mod core_registers; mod memory; mod permutation_unit; mod public_input_layout; +mod registers; pub mod system_zero; diff --git a/system_zero/src/permutation_unit.rs b/system_zero/src/permutation_unit.rs index e15474e4..2681f2d9 100644 --- a/system_zero/src/permutation_unit.rs +++ b/system_zero/src/permutation_unit.rs @@ -8,12 +8,9 @@ use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsume use starky::vars::StarkEvaluationTargets; use starky::vars::StarkEvaluationVars; -use crate::column_layout::permutation::{ - col_full_first_after_mds, col_full_first_mid_sbox, col_full_second_after_mds, - col_full_second_mid_sbox, col_input, col_partial_after_sbox, col_partial_mid_sbox, -}; -use crate::column_layout::NUM_COLUMNS; use crate::public_input_layout::NUM_PUBLIC_INPUTS; +use crate::registers::permutation::*; +use crate::registers::NUM_COLUMNS; use crate::system_zero::SystemZero; fn constant_layer( @@ -272,10 +269,10 @@ mod tests { use starky::constraint_consumer::ConstraintConsumer; use starky::vars::StarkEvaluationVars; - use crate::column_layout::permutation::{col_input, col_output}; - use crate::column_layout::NUM_COLUMNS; use crate::permutation_unit::SPONGE_WIDTH; use crate::public_input_layout::NUM_PUBLIC_INPUTS; + use crate::registers::permutation::{col_input, col_output}; + use crate::registers::NUM_COLUMNS; use crate::system_zero::SystemZero; #[test] diff --git a/system_zero/src/registers/arithmetic.rs b/system_zero/src/registers/arithmetic.rs new file mode 100644 index 00000000..92c0d2c3 --- /dev/null +++ b/system_zero/src/registers/arithmetic.rs @@ -0,0 +1,37 @@ +//! Arithmetic unit. + +pub(crate) const IS_ADD: usize = super::START_ARITHMETIC; +pub(crate) const IS_SUB: usize = IS_ADD + 1; +pub(crate) const IS_MUL: usize = IS_SUB + 1; +pub(crate) const IS_DIV: usize = IS_MUL + 1; + +const START_SHARED_COLS: usize = IS_DIV + 1; + +/// Within the arithmetic unit, there are shared columns which can be used by any arithmetic +/// circuit, depending on which one is active this cycle. +// Can be increased as needed as other operations are implemented. +const NUM_SHARED_COLS: usize = 3; + +const fn shared_col(i: usize) -> usize { + debug_assert!(i < NUM_SHARED_COLS); + START_SHARED_COLS + i +} + +/// The first value to be added; treated as an unsigned u32. +pub(crate) const COL_ADD_INPUT_1: usize = shared_col(0); +/// The second value to be added; treated as an unsigned u32. +pub(crate) const COL_ADD_INPUT_2: usize = shared_col(1); +/// The third value to be added; treated as an unsigned u32. +pub(crate) const COL_ADD_INPUT_3: usize = shared_col(2); + +// Note: Addition outputs three 16-bit chunks, and since these values need to be range-checked +// anyway, we might as well use the range check unit's columns as our addition outputs. So the +// three proceeding columns are basically aliases, not columns owned by the arithmetic unit. +/// The first 16-bit chunk of the output, based on little-endian ordering. +pub(crate) const COL_ADD_OUTPUT_1: usize = super::range_check_16::col_rc_16_input(0); +/// The second 16-bit chunk of the output, based on little-endian ordering. +pub(crate) const COL_ADD_OUTPUT_2: usize = super::range_check_16::col_rc_16_input(1); +/// The third 16-bit chunk of the output, based on little-endian ordering. +pub(crate) const COL_ADD_OUTPUT_3: usize = super::range_check_16::col_rc_16_input(2); + +pub(super) const END: usize = super::START_ARITHMETIC + NUM_SHARED_COLS; diff --git a/system_zero/src/registers/boolean.rs b/system_zero/src/registers/boolean.rs new file mode 100644 index 00000000..c59af8d4 --- /dev/null +++ b/system_zero/src/registers/boolean.rs @@ -0,0 +1,10 @@ +//! Boolean unit. Contains columns whose values must be 0 or 1. + +const NUM_BITS: usize = 128; + +pub const fn col_bit(index: usize) -> usize { + debug_assert!(index < NUM_BITS); + super::START_BOOLEAN + index +} + +pub(super) const END: usize = super::START_BOOLEAN + NUM_BITS; diff --git a/system_zero/src/registers/core.rs b/system_zero/src/registers/core.rs new file mode 100644 index 00000000..3fafab55 --- /dev/null +++ b/system_zero/src/registers/core.rs @@ -0,0 +1,20 @@ +//! Core registers. + +/// A cycle counter. Starts at 0; increments by 1. +pub(crate) const COL_CLOCK: usize = super::START_CORE; + +/// A column which contains the values `[0, ... 2^16 - 1]`, potentially with duplicates. Used for +/// 16-bit range checks. +/// +/// For ease of verification, we enforce that it must begin with 0 and end with `2^16 - 1`, and each +/// delta must be either 0 or 1. +pub(crate) const COL_RANGE_16: usize = COL_CLOCK + 1; + +/// Pointer to the current instruction. +pub(crate) const COL_INSTRUCTION_PTR: usize = COL_RANGE_16 + 1; +/// Pointer to the base of the current call's stack frame. +pub(crate) const COL_FRAME_PTR: usize = COL_INSTRUCTION_PTR + 1; +/// Pointer to the tip of the current call's stack frame. +pub(crate) const COL_STACK_PTR: usize = COL_FRAME_PTR + 1; + +pub(super) const END: usize = COL_STACK_PTR + 1; diff --git a/system_zero/src/registers/logic.rs b/system_zero/src/registers/logic.rs new file mode 100644 index 00000000..07f3f0e0 --- /dev/null +++ b/system_zero/src/registers/logic.rs @@ -0,0 +1,3 @@ +//! Logic unit. + +pub(super) const END: usize = super::START_LOGIC; diff --git a/system_zero/src/registers/lookup.rs b/system_zero/src/registers/lookup.rs new file mode 100644 index 00000000..eb773acf --- /dev/null +++ b/system_zero/src/registers/lookup.rs @@ -0,0 +1,21 @@ +//! Lookup unit. +//! See https://zcash.github.io/halo2/design/proving-system/lookup.html + +const START_UNIT: usize = super::START_LOOKUP; + +const NUM_LOOKUPS: usize = + super::range_check_16::NUM_RANGE_CHECKS + super::range_check_degree::NUM_RANGE_CHECKS; + +/// This column contains a permutation of the input values. +const fn col_permuted_input(i: usize) -> usize { + debug_assert!(i < NUM_LOOKUPS); + START_UNIT + 2 * i +} + +/// This column contains a permutation of the table values. +const fn col_permuted_table(i: usize) -> usize { + debug_assert!(i < NUM_LOOKUPS); + START_UNIT + 2 * i + 1 +} + +pub(super) const END: usize = START_UNIT + NUM_LOOKUPS; diff --git a/system_zero/src/registers/memory.rs b/system_zero/src/registers/memory.rs new file mode 100644 index 00000000..1373d0d8 --- /dev/null +++ b/system_zero/src/registers/memory.rs @@ -0,0 +1,3 @@ +//! Memory unit. + +pub(super) const END: usize = super::START_MEMORY; diff --git a/system_zero/src/registers/mod.rs b/system_zero/src/registers/mod.rs new file mode 100644 index 00000000..134a28bf --- /dev/null +++ b/system_zero/src/registers/mod.rs @@ -0,0 +1,20 @@ +pub(crate) mod arithmetic; +pub(crate) mod boolean; +pub(crate) mod core; +pub(crate) mod logic; +pub(crate) mod lookup; +pub(crate) mod memory; +pub(crate) mod permutation; +pub(crate) mod range_check_16; +pub(crate) mod range_check_degree; + +const START_ARITHMETIC: usize = 0; +const START_BOOLEAN: usize = arithmetic::END; +const START_CORE: usize = boolean::END; +const START_LOGIC: usize = core::END; +const START_LOOKUP: usize = logic::END; +const START_MEMORY: usize = lookup::END; +const START_PERMUTATION: usize = memory::END; +const START_RANGE_CHECK_16: usize = permutation::END; +const START_RANGE_CHECK_DEGREE: usize = range_check_16::END; +pub(crate) const NUM_COLUMNS: usize = range_check_degree::END; diff --git a/system_zero/src/registers/permutation.rs b/system_zero/src/registers/permutation.rs new file mode 100644 index 00000000..cde76af2 --- /dev/null +++ b/system_zero/src/registers/permutation.rs @@ -0,0 +1,57 @@ +//! Permutation unit. + +use plonky2::hash::hashing::SPONGE_WIDTH; +use plonky2::hash::poseidon; + +const START_FULL_FIRST: usize = super::START_PERMUTATION + SPONGE_WIDTH; + +pub const fn col_full_first_mid_sbox(round: usize, i: usize) -> usize { + debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS); + debug_assert!(i < SPONGE_WIDTH); + START_FULL_FIRST + 2 * round * SPONGE_WIDTH + i +} + +pub const fn col_full_first_after_mds(round: usize, i: usize) -> usize { + debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS); + debug_assert!(i < SPONGE_WIDTH); + START_FULL_FIRST + (2 * round + 1) * SPONGE_WIDTH + i +} + +const START_PARTIAL: usize = + col_full_first_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, SPONGE_WIDTH - 1) + 1; + +pub const fn col_partial_mid_sbox(round: usize) -> usize { + debug_assert!(round < poseidon::N_PARTIAL_ROUNDS); + START_PARTIAL + 2 * round +} + +pub const fn col_partial_after_sbox(round: usize) -> usize { + debug_assert!(round < poseidon::N_PARTIAL_ROUNDS); + START_PARTIAL + 2 * round + 1 +} + +const START_FULL_SECOND: usize = col_partial_after_sbox(poseidon::N_PARTIAL_ROUNDS - 1) + 1; + +pub const fn col_full_second_mid_sbox(round: usize, i: usize) -> usize { + debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS); + debug_assert!(i < SPONGE_WIDTH); + START_FULL_SECOND + 2 * round * SPONGE_WIDTH + i +} + +pub const fn col_full_second_after_mds(round: usize, i: usize) -> usize { + debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS); + debug_assert!(i < SPONGE_WIDTH); + START_FULL_SECOND + (2 * round + 1) * SPONGE_WIDTH + i +} + +pub const fn col_input(i: usize) -> usize { + debug_assert!(i < SPONGE_WIDTH); + super::START_PERMUTATION + i +} + +pub const fn col_output(i: usize) -> usize { + debug_assert!(i < SPONGE_WIDTH); + col_full_second_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, i) +} + +pub(super) const END: usize = col_output(SPONGE_WIDTH - 1) + 1; diff --git a/system_zero/src/registers/range_check_16.rs b/system_zero/src/registers/range_check_16.rs new file mode 100644 index 00000000..c44db494 --- /dev/null +++ b/system_zero/src/registers/range_check_16.rs @@ -0,0 +1,11 @@ +//! Range check unit which checks that values are in `[0, 2^16)`. + +pub(super) const NUM_RANGE_CHECKS: usize = 5; + +/// The input of the `i`th range check, i.e. the value being range checked. +pub(crate) const fn col_rc_16_input(i: usize) -> usize { + debug_assert!(i < NUM_RANGE_CHECKS); + super::START_RANGE_CHECK_16 + i +} + +pub(super) const END: usize = super::START_RANGE_CHECK_16 + NUM_RANGE_CHECKS; diff --git a/system_zero/src/registers/range_check_degree.rs b/system_zero/src/registers/range_check_degree.rs new file mode 100644 index 00000000..6d61e6e2 --- /dev/null +++ b/system_zero/src/registers/range_check_degree.rs @@ -0,0 +1,11 @@ +//! Range check unit which checks that values are in `[0, degree)`. + +pub(super) const NUM_RANGE_CHECKS: usize = 5; + +/// The input of the `i`th range check, i.e. the value being range checked. +pub(crate) const fn col_rc_degree_input(i: usize) -> usize { + debug_assert!(i < NUM_RANGE_CHECKS); + super::START_RANGE_CHECK_DEGREE + i +} + +pub(super) const END: usize = super::START_RANGE_CHECK_DEGREE + NUM_RANGE_CHECKS; diff --git a/system_zero/src/system_zero.rs b/system_zero/src/system_zero.rs index 70d3bbca..780b1d38 100644 --- a/system_zero/src/system_zero.rs +++ b/system_zero/src/system_zero.rs @@ -9,9 +9,12 @@ use starky::stark::Stark; use starky::vars::StarkEvaluationTargets; use starky::vars::StarkEvaluationVars; -use crate::column_layout::NUM_COLUMNS; +use crate::arithmetic::{ + eval_arithmetic_unit, eval_arithmetic_unit_recursively, generate_arithmetic_unit, +}; use crate::memory::TransactionMemory; use crate::public_input_layout::NUM_PUBLIC_INPUTS; +use crate::registers::NUM_COLUMNS; /// We require at least 2^16 rows as it helps support efficient 16-bit range checks. const MIN_TRACE_ROWS: usize = 1 << 16; @@ -34,10 +37,16 @@ impl, const D: usize> SystemZero { loop { let mut next_row = [F::ZERO; NUM_COLUMNS]; self.generate_next_row_core_registers(&row, &mut next_row); + generate_arithmetic_unit(&mut next_row); Self::generate_permutation_unit(&mut next_row); trace.push(row); row = next_row; + + // TODO: Replace with proper termination condition. + if trace.len() == (1 << 16) - 1 { + break; + } } trace.push(row); @@ -66,8 +75,9 @@ impl, const D: usize> Stark for SystemZero, { self.eval_core_registers(vars, yield_constr); + eval_arithmetic_unit(vars, yield_constr); Self::eval_permutation_unit(vars, yield_constr); - todo!() + // TODO: Other units } fn eval_ext_recursively( @@ -77,8 +87,9 @@ impl, const D: usize> Stark for SystemZero, ) { self.eval_core_registers_recursively(builder, vars, yield_constr); + eval_arithmetic_unit_recursively(builder, vars, yield_constr); Self::eval_permutation_unit_recursively(builder, vars, yield_constr); - todo!() + // TODO: Other units } fn constraint_degree(&self) -> usize { @@ -103,7 +114,7 @@ mod tests { use crate::system_zero::SystemZero; #[test] - #[ignore] // TODO + #[ignore] // A bit slow. fn run() -> Result<()> { type F = GoldilocksField; type C = PoseidonGoldilocksConfig; @@ -121,7 +132,6 @@ mod tests { } #[test] - #[ignore] // TODO fn degree() -> Result<()> { type F = GoldilocksField; type C = PoseidonGoldilocksConfig;