Merge pull request #884 from powdr-labs/memory-allow-bigger-diffs

Read/Write memory: Allow bigger diffs

executor/src/witgen/machines/double_sorted_witness_machine.rs

@@ -7,6 +7,7 @@ use num_traits::Zero;
 
 use super::{FixedLookup, Machine};
 use crate::witgen::affine_expression::AffineExpression;
+use crate::witgen::global_constraints::GlobalConstraints;
 use crate::witgen::util::is_simple_poly_of_name;
 use crate::witgen::{EvalResult, FixedData, MutableState, QueryCallback};
 use crate::witgen::{EvalValue, IncompleteCause};
@@ -30,6 +31,9 @@ pub struct DoubleSortedWitnesses<T> {
     data: BTreeMap<T, T>,
     namespace: String,
     name: String,
+    /// If the machine has the `m_diff_upper` and `m_diff_lower` columns, this is the base of the
+    /// two digits.
+    diff_columns_base: Option<u64>,
 }
 
 struct Operation<T> {
@@ -47,6 +51,7 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
         fixed_data: &FixedData<T>,
         _identities: &[&Identity<Expression<T>>],
         witness_cols: &HashSet<PolyID>,
+        global_range_constraints: &GlobalConstraints<T>,
     ) -> Option<Self> {
         // get the namespaces and column names
         let (mut namespaces, columns): (HashSet<_>, HashSet<_>) = witness_cols
@@ -77,19 +82,47 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
         ]
         .into_iter()
         .collect();
-        if expected_witnesses
+        let difference = expected_witnesses
             .symmetric_difference(&columns)
-            .next()
-            .is_none()
-        {
-            Some(Self {
+            .collect::<HashSet<_>>();
+        match difference.len() {
+            0 => Some(Self {
                 name,
                 namespace,
                 degree: fixed_data.degree,
+                // No diff columns.
+                diff_columns_base: None,
                 ..Default::default()
-            })
-        } else {
-            None
+            }),
+            2 if difference == ["m_diff_upper", "m_diff_lower"].iter().collect() => {
+                // We have the `m_diff_upper` and `m_diff_lower` columns.
+                // Now, we check that they both have the same range constraint and use it to determine
+                // the base of the two digits.
+                let upper_poly_id =
+                    fixed_data.try_column_by_name(&format!("{namespace}.m_diff_upper"))?;
+                let upper_range_constraint =
+                    global_range_constraints.witness_constraints[&upper_poly_id].as_ref()?;
+                let lower_poly_id =
+                    fixed_data.try_column_by_name(&format!("{namespace}.m_diff_lower"))?;
+                let lower_range_constraint =
+                    global_range_constraints.witness_constraints[&lower_poly_id].as_ref()?;
+
+                let (min, max) = upper_range_constraint.range();
+
+                if upper_range_constraint == lower_range_constraint && min == T::zero() {
+                    let diff_columns_base = Some(max.to_degree() + 1);
+                    Some(Self {
+                        name,
+                        namespace,
+                        degree: fixed_data.degree,
+                        diff_columns_base,
+                        ..Default::default()
+                    })
+                } else {
+                    None
+                }
+            }
+            _ => None,
         }
     }
 }
@@ -126,13 +159,24 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
         let mut value = vec![];
         let mut is_write = vec![];
         let mut is_read = vec![];
+        let mut diff = vec![];
 
         for ((a, s), o) in std::mem::take(&mut self.trace) {
             if let Some(prev_address) = addr.last() {
                 assert!(a >= *prev_address, "Expected addresses to be sorted");
-                if (a - *prev_address).to_degree() >= self.degree {
+                if self.diff_columns_base.is_none()
+                    && (a - *prev_address).to_degree() >= self.degree
+                {
                     log::error!("Jump in memory accesses between {prev_address:x} and {a:x} is larger than or equal to the degree {}! This will violate the constraints.", self.degree);
                 }
+
+                let current_diff = if a != *prev_address {
+                    a - *prev_address
+                } else {
+                    s - *step.last().unwrap()
+                };
+                assert!(current_diff > T::zero());
+                diff.push(current_diff.to_degree() - 1);
             }
 
             addr.push(a);
@@ -153,11 +197,16 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
         while addr.len() < self.degree as usize {
             addr.push(*addr.last().unwrap());
             step.push(*step.last().unwrap() + T::from(1));
+            diff.push(0);
             value.push(*value.last().unwrap());
             is_write.push(0.into());
             is_read.push(0.into());
         }
 
+        // We have all diffs, except from the last to the first element, which is unconstrained.
+        assert_eq!(diff.len(), self.degree as usize - 1);
+        diff.push(0);
+
         let change = addr
             .iter()
             .tuple_windows()
@@ -166,6 +215,23 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
             .collect::<Vec<_>>();
         assert_eq!(change.len(), addr.len());
 
+        let diff_columns = if let Some(diff_columns_base) = self.diff_columns_base {
+            let diff_upper = diff
+                .iter()
+                .map(|d| T::from(*d / diff_columns_base))
+                .collect::<Vec<_>>();
+            let diff_lower = diff
+                .iter()
+                .map(|d| T::from(*d % diff_columns_base))
+                .collect::<Vec<_>>();
+            vec![
+                (self.namespaced("m_diff_upper"), diff_upper),
+                (self.namespaced("m_diff_lower"), diff_lower),
+            ]
+        } else {
+            vec![]
+        };
+
         [
             (self.namespaced("m_value"), value),
             (self.namespaced("m_addr"), addr),
@@ -175,6 +241,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
             (self.namespaced("m_is_read"), is_read),
         ]
         .into_iter()
+        .chain(diff_columns)
         .collect()
     }
 }

executor/src/witgen/machines/machine_extractor.rs

@@ -125,6 +125,7 @@ pub fn split_out_machines<'a, T: FieldElement>(
             fixed,
             &machine_identities,
             &machine_witnesses,
+            global_range_constraints,
         ) {
             log::info!("Detected machine: memory");
             machines.push(KnownMachine::DoubleSortedWitnesses(machine));

pipeline/tests/asm.rs

@@ -192,6 +192,14 @@ fn test_mem_read_write() {
     gen_estark_proof(f, Default::default());
 }
 
+#[test]
+fn test_mem_read_write_large_diffs() {
+    let f = "asm/mem_read_write_large_diffs.asm";
+    verify_asm::<GoldilocksField>(f, Default::default());
+    gen_halo2_proof(f, Default::default());
+    gen_estark_proof(f, Default::default());
+}
+
 #[test]
 fn test_multi_assign() {
     let f = "asm/multi_assign.asm";

riscv/src/compiler.rs

@@ -504,12 +504,16 @@ fn preamble(degree: u64, coprocessors: &CoProcessors, with_bootloader: bool) ->
     // If the operation is a write operation.
     col witness m_is_write;
     col witness m_is_read;
+    col witness m_diff_lower;
+    col witness m_diff_upper;
 
-    // positive numbers (assumed to be much smaller than the field order)
-    col fixed POSITIVE(i) { i + 1 };
     col fixed FIRST = [1] + [0]*;
     col fixed LAST(i) { FIRST(i + 1) };
     col fixed STEP(i) { i };
+    col fixed BIT16(i) { i & 0xffff };
+
+    {m_diff_lower} in {BIT16};
+    {m_diff_upper} in {BIT16};
 
     m_change * (1 - m_change) = 0;
 
@@ -518,7 +522,11 @@ fn preamble(degree: u64, coprocessors: &CoProcessors, with_bootloader: bool) ->
 
     // Except for the last row, if m_change is 1, then m_addr has to increase,
     // if it is zero, m_step has to increase.
-    (1 - LAST) { m_change * (m_addr' - m_addr) + (1 - m_change) * (m_step' - m_step) } in POSITIVE;
+    // `m_diff_upper * 2**16 + m_diff_lower` has to be equal to the difference **minus one**.
+    // Since we know that both m_addr and m_step can only be 32-Bit, this enforces that
+    // the values are strictly increasing.
+    col diff = (m_change * (m_addr' - m_addr) + (1 - m_change) * (m_step' - m_step));
+    (1 - LAST) * (diff - 1 - m_diff_upper * 2**16 - m_diff_lower) = 0;
 
     // m_change has to be 1 in the last row, so that a first read on row zero is constrained to return 0
     (1 - m_change) * LAST = 0;

test_data/asm/mem_read_write_large_diffs.asm

@@ -0,0 +1,86 @@
+// A variant of the `mem_read_write` machine which does not have the limitation that
+// gaps between accessed memory cells must not be larger than the degree.
+// This test uses two 8-bit digits to represent the diff, so the diff has to be
+// representable in 16 bits.
+machine MemReadWrite {
+    reg pc[@pc];
+    reg X[<=];
+    reg A;
+    reg B;
+    reg I;
+    reg CNT;
+    reg ADDR;
+
+    col witness XInv;
+    col witness XIsZero;
+    XIsZero  = 1 - X * XInv;
+    XIsZero * X = 0;
+    XIsZero * (1 - XIsZero) = 0;
+
+    // Read-write memory. Columns are sorted by m_addr and
+    // then by m_step. m_change is 1 if and only if m_addr changes
+    // in the next row.
+    col witness m_addr;
+    col witness m_step;
+    col witness m_change;
+    col witness m_value;
+    // If the operation is a write operation.
+    col witness m_is_write;
+    col witness m_is_read;
+    col witness m_diff_lower;
+    col witness m_diff_upper;
+
+    col fixed FIRST = [1] + [0]*;
+    col fixed LAST  = [0]* + [1];
+    col fixed STEP(i) { i };
+    col fixed BYTE(i) { i & 0xff };
+
+    {m_diff_lower} in {BYTE};
+    {m_diff_upper} in {BYTE};
+
+    m_change * (1 - m_change) = 0;
+
+    // if m_change is zero, m_addr has to stay the same.
+    (m_addr' - m_addr) * (1 - m_change) = 0;
+
+    // Except for the last row, if m_change is 1, then m_addr has to increase,
+    // if it is zero, m_step has to increase.
+    // `m_diff_upper * 2**8 + m_diff_lower` has to be equal to the difference **minus one**.
+    col diff = (m_change * (m_addr' - m_addr) + (1 - m_change) * (m_step' - m_step));
+    (1 - LAST) * (diff - 1 - m_diff_upper * 2**8 - m_diff_lower) = 0;
+
+    // m_change has to be 1 in the last row, so that a first read on row zero is constrained to return 0
+    (1 - m_change) * LAST = 0;
+
+    m_is_write * (1 - m_is_write) = 0;
+    m_is_read * (1 - m_is_read) = 0;
+    m_is_read * m_is_write = 0;
+
+
+    // If the next line is a read and we stay at the same address, then the
+    // value cannot change.
+    (1 - m_is_write') * (1 - m_change) * (m_value' - m_value) = 0;
+
+    // If the next line is a read and we have an address change,
+    // then the value is zero.
+    (1 - m_is_write') * m_change * m_value' = 0;
+
+    instr assert_zero X { XIsZero = 1 }
+    instr mstore X { { ADDR, STEP, X } is m_is_write { m_addr, m_step, m_value } }
+    instr mload -> X { { ADDR, STEP, X } is m_is_read { m_addr, m_step, m_value } }
+
+    function main {
+        ADDR <=X= 4;
+        mstore 1;
+        // Write to the largest aligned memory cell.
+        // This wouldn't be possible in the `mem_read_write` machine.
+        ADDR <=X= 0xfffc;
+        mstore 4;
+        mload A;
+        assert_zero A - 4;
+        ADDR <=X= 4;
+        mload A;
+        assert_zero A - 1;
+        return;
+    }
+}