Future optimization for plonky3 keccak memory (16 bits limbs)

[qwang98]

This draft PR documents the one final optimization that we can apply to the existing optimizations for Plonky3 memory Keccak (16-bits).

The Issue
The current optimization (branch compared against) moves inverse and carry constraints for incrementing 4 to two 16-bit limbs addresses to the main machine. These constraints require degree-4 constraints, which our current Plonky3 parameters limit to degree-3, which our current optimization solves by creating 50 additional helper[50] columns that are used to reduce the degrees.

This Future Optmization
This optimization removes the need for the 50 helper columns while reducing degree-4 constraints to degree-3, by populating intermediate columns for constraining the increment 4 logic, to the second and third row of each block (for input addresses) and to the 22nd and 23rd row of each block (for output addresses). Currently constraint is only implemented by input addresses but doesn't work yet according to the output.txt witgen trace.

Most specifically, this optimization has the following column layout for row 1 - 3:

addr_l || addr_h
addr_l || carry
addr_l * inverse || inverse

The degree-4 constraint is due to a first_step * inverse * addr_l * addr_l term, which we can now move to row two (which can reference row three via rotation). I was able to mathematically prove that two rows won't be sufficient to accommodate for the degree-4 constraint without adding additional columns, so we need to use at least 3 rows and this should be one way to make it work. The other way is to make addr_l * inverse in row 3 addr_l * addr_l, but the idea is similar.

Where This Optimization Doesn't Work
According to witgen trace from output.txt, output addresses increment 4 are calculated correctly (Appendix 1), input address is copied from the latch row to the first row, but then input addresses increment 4 aren't calculated as intended (Appendix 2).

How We Might Fix This Optimization In the Future
Currently witgen can't figure out the value of these three rows, likely due to the fact that it only has context of two rows while we have three rows that express the logic of increment 4 here. We might be able to add hints and calculate the increment 4 logic using if else for all addr_l[50] and addr_h[50] columns in the first step rather than letting witgen do it. I'm not sure if rotations in prover functions are allowed, to express logics like those between addr_l and carry.

When Is This Optimization Worth It
Again, this optimization will shave off another 50 columns, out of a total of 2700+, so it might be worth or not worth it depending on the time it takes.

Appendix 1

current_round_count: 15
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::helper[8] - (main_keccakf16_memory::addr_l[8] * (main_keccakf16_memory::addr_l[8]' - 65532) - 1)) = 0;
      => main_keccakf16_memory::addr_l[8] (Row 22) = 14253508648738444382
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::addr_h[8] + main_keccakf16_memory::addr_l[8] * (main_keccakf16_memory::addr_l[8]' - 65532) - 1) = 0;
      => main_keccakf16_memory::addr_h[8] (Row 22) = 0
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::addr_h[8] * main_keccakf16_memory::addr_l[9]' + (1 - main_keccakf16_memory::addr_h[8]) * (main_keccakf16_memory::addr_l[9]' - main_keccakf16_memory::addr_l[8]' - 4)) = 0;
      => main_keccakf16_memory::addr_l[9] (Row 23) = 28
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::addr_h[9]' - main_keccakf16_memory::addr_h[8]' - main_keccakf16_memory::addr_h[8]) = 0;
      => main_keccakf16_memory::addr_h[9] (Row 23) = 50001
current_round_count: 16
    Updates from: main_keccakf16_memory::step_flags[22] * main_keccakf16_memory::helper[9] * (main_keccakf16_memory::addr_l[9]' - 65532) = 0;
      => main_keccakf16_memory::helper[9] (Row 22) = 0
current_round_count: 17
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::helper[9] - (main_keccakf16_memory::addr_l[9] * (main_keccakf16_memory::addr_l[9]' - 65532) - 1)) = 0;
      => main_keccakf16_memory::addr_l[9] (Row 22) = 12490358446976445078
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::addr_h[9] + main_keccakf16_memory::addr_l[9] * (main_keccakf16_memory::addr_l[9]' - 65532) - 1) = 0;
      => main_keccakf16_memory::addr_h[9] (Row 22) = 0
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::addr_h[9] * main_keccakf16_memory::addr_l[10]' + (1 - main_keccakf16_memory::addr_h[9]) * (main_keccakf16_memory::addr_l[10]' - main_keccakf16_memory::addr_l[9]' - 4)) = 0;
      => main_keccakf16_memory::addr_l[10] (Row 23) = 32
    Updates from: main_keccakf16_memory::step_flags[22] * (main_keccakf16_memory::addr_h[10]' - main_keccakf16_memory::addr_h[9]' - main_keccakf16_memory::addr_h[9]) = 0;
      => main_keccakf16_memory::addr_h[10] (Row 23) = 50001

Appendix 2

Memory read: addr=0x64fff0, step=27, value=0x0
    Updates from: main_keccakf16_memory::sel[0] * main_keccakf16_memory::step_flags[0] $ [0, main_keccakf16_memory::addr_h[0], main_keccakf16_memory::addr_l[0], main_keccakf16_memory::time_step, main_keccakf16_memory::preimage[3], main_keccakf16_memory::preimage[2]] is main_memory::selectors[2] $ [main_memory::m_is_write, main_memory::m_addr_high, main_memory::m_addr_low, main_memory::m_step_high * 65536 + main_memory::m_step_low, main_memory::m_value1, main_memory::m_value2];
      => main_keccakf16_memory::preimage[3] (Row 0) = 0
      => main_keccakf16_memory::preimage[2] (Row 0) = 0
current_round_count: 0
    Updates from: main_keccakf16_memory::step_flags[0] * (main_keccakf16_memory::addr_l[0]' - main_keccakf16_memory::addr_l[0]) = 0;
      => main_keccakf16_memory::addr_l[0] (Row 1) = 65520
    Updates from: main_keccakf16_memory::step_flags[0] * (main_keccakf16_memory::preimage[2] - main_keccakf16_memory::a[2]) = 0;
      => main_keccakf16_memory::a[2] (Row 0) = 0
    Updates from: main_keccakf16_memory::step_flags[0] * (main_keccakf16_memory::preimage[3] - main_keccakf16_memory::a[3]) = 0;
      => main_keccakf16_memory::a[3] (Row 0) = 0
    Updates from: (main_keccakf16_memory::preimage[2]' - main_keccakf16_memory::preimage[2]) * (1 - (main_keccakf16_memory::step_flags[23] + main_keccakf16_memory::is_last)) = 0;
      => main_keccakf16_memory::preimage[2] (Row 1) = 0
    Updates from: (main_keccakf16_memory::preimage[3]' - main_keccakf16_memory::preimage[3]) * (1 - (main_keccakf16_memory::step_flags[23] + main_keccakf16_memory::is_last)) = 0;
      => main_keccakf16_memory::preimage[3] (Row 1) = 0
current_round_count: 1
current_round_count: 2
current_round_count: 0
current_round_count: 0
    Updates from: (main_keccakf16_memory::preimage[2]' - main_keccakf16_memory::preimage[2]) * (1 - (main_keccakf16_memory::step_flags[23] + main_keccakf16_memory::is_last)) = 0;
      => main_keccakf16_memory::preimage[2] (Row 2) = 0
    Updates from: (main_keccakf16_memory::preimage[3]' - main_keccakf16_memory::preimage[3]) * (1 - (main_keccakf16_memory::step_flags[23] + main_keccakf16_memory::is_last)) = 0;
      => main_keccakf16_memory::preimage[3] (Row 2) = 0
current_round_count: 0
current_round_count: 1
    Updates from: (main_keccakf16_memory::preimage[2]' - main_keccakf16_memory::preimage[2]) * (1 - (main_keccakf16_memory::step_flags[23] + main_keccakf16_memory::is_last)) = 0;
      => main_keccakf16_memory::preimage[2] (Row 3) = 0
    Updates from: (main_keccakf16_memory::preimage[3]' - main_keccakf16_memory::preimage[3]) * (1 - (main_keccakf16_memory::step_flags[23] + main_keccakf16_memory::is_last)) = 0;
      => main_keccakf16_memory::preimage[3] (Row 3) = 0