Allow Hint Space Ids for NativeSumcheck Inputs

crates/vm/src/system/memory/offline_checker/columns.rs

@@ -62,7 +62,7 @@ impl<const N: usize, T> MemoryWriteAuxCols<T, N> {
 #[repr(C)]
 #[derive(Clone, Copy, Debug, AlignedBorrow)]
 pub struct MemoryReadAuxCols<T> {
-    pub(in crate::system::memory) base: MemoryBaseAuxCols<T>,
+    pub base: MemoryBaseAuxCols<T>,
 }
 
 impl<F: PrimeField32> MemoryReadAuxCols<F> {

extensions/native/circuit/cuda/include/native/sumcheck.cuh

@@ -7,10 +7,10 @@ using namespace native;
 
 template <typename T> struct HeaderSpecificCols {
     T pc;
-    T registers[3];
-    T prod_id;
-    T logup_id;
-    MemoryReadAuxCols<T> read_records[6];
+    T registers[5];
+    T prod_evals_id;
+    T logup_evals_id;
+    MemoryReadAuxCols<T> read_records[8];
     MemoryWriteAuxCols<T, EXT_DEG> write_records;
 };
 
@@ -63,7 +63,7 @@ template <typename T> struct NativeSumcheckCols {
     T start_timestamp;
     T last_timestamp;
 
-    T register_ptrs[3];
+    T register_ptrs[5];
 
     T ctx[EXT_DEG * 2];
 
@@ -82,6 +82,8 @@ template <typename T> struct NativeSumcheckCols {
 
     T eval_acc[EXT_DEG];
 
+    T is_hint_src_id;
+
     T specific[COL_SPECIFIC_WIDTH];
 };
 

extensions/native/circuit/cuda/src/sumcheck.cu

@@ -11,7 +11,7 @@ __device__ void fill_sumcheck_specific(RowSlice row, MemoryAuxColsFactory &mem_h
     uint32_t start_timestamp = row[COL_INDEX(NativeSumcheckCols, start_timestamp)].asUInt32();
 
     if (row[COL_INDEX(NativeSumcheckCols, header_row)] == Fp::one()) {
-        for (uint32_t i = 0; i < 6; ++i) {
+        for (uint32_t i = 0; i < 8; ++i) {
             mem_fill_base(
                 mem_helper,
                 start_timestamp + i,
@@ -26,32 +26,32 @@ __device__ void fill_sumcheck_specific(RowSlice row, MemoryAuxColsFactory &mem_h
         );
     } else if (row[COL_INDEX(NativeSumcheckCols, prod_row)] == Fp::one()) {
         if (row[COL_INDEX(NativeSumcheckCols, within_round_limit)] == Fp::one()) {
-            // mem_fill_base(
-            //     mem_helper,
-            //     start_timestamp,
-            //     specific.slice_from(COL_INDEX(ProdSpecificCols, read_records[0].base))
-            // );
             mem_fill_base(
                 mem_helper,
                 start_timestamp,
+                specific.slice_from(COL_INDEX(ProdSpecificCols, ps_record.base))
+            );
+            mem_fill_base(
+                mem_helper,
+                start_timestamp + 1,
                 specific.slice_from(COL_INDEX(ProdSpecificCols, write_record.base))
             );
         }
     } else if (row[COL_INDEX(NativeSumcheckCols, logup_row)] == Fp::one()) {
         if (row[COL_INDEX(NativeSumcheckCols, within_round_limit)] == Fp::one()) {
-            // mem_fill_base(
-            //     mem_helper,
-            //     start_timestamp,
-            //     specific.slice_from(COL_INDEX(LogupSpecificCols, read_records[0].base))
-            // );
             mem_fill_base(
                 mem_helper,
                 start_timestamp,
-                specific.slice_from(COL_INDEX(LogupSpecificCols, write_records[0].base))
+                specific.slice_from(COL_INDEX(LogupSpecificCols, pqs_record.base))
             );
             mem_fill_base(
                 mem_helper,
                 start_timestamp + 1,
+                specific.slice_from(COL_INDEX(LogupSpecificCols, write_records[0].base))
+            );
+            mem_fill_base(
+                mem_helper,
+                start_timestamp + 2,
                 specific.slice_from(COL_INDEX(LogupSpecificCols, write_records[1].base))
             );
         }

extensions/native/circuit/src/sumcheck/air.rs

@@ -2,7 +2,7 @@ use std::borrow::Borrow;
 
 use openvm_circuit::{
     arch::{ExecutionBridge, ExecutionState},
-    system::memory::{offline_checker::MemoryBridge, MemoryAddress},
+    system::memory::{MemoryAddress, offline_checker::{MemoryBridge, MemoryReadAuxCols}},
 };
 use openvm_circuit_primitives::utils::{and, assert_array_eq, not};
 use openvm_instructions::{LocalOpcode, NATIVE_AS};
@@ -23,9 +23,8 @@ use crate::{
     },
 };
 
-pub const TOPLEVEL_TIMESTAMP_DIFF: usize = 6;
-pub const NUM_RWS_FOR_PRODUCT: usize = 1;
-pub const NUM_RWS_FOR_LOGUP: usize = 2;
+pub const NUM_RWS_FOR_PRODUCT: usize = 2;
+pub const NUM_RWS_FOR_LOGUP: usize = 3;
 
 #[derive(Clone, Debug)]
 pub struct NativeSumcheckAir {
@@ -103,6 +102,7 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             within_round_limit,
             should_acc,
             eval_acc,
+            is_hint_src_id,
             specific,
         } = local;
 
@@ -230,7 +230,7 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .when(next.prod_row + next.logup_row)
             .assert_eq(
                 next.start_timestamp,
-                start_timestamp + AB::F::from_canonical_usize(TOPLEVEL_TIMESTAMP_DIFF),
+                start_timestamp + AB::F::from_canonical_usize(8),
             );
         builder
             .when(prod_row)
@@ -294,31 +294,27 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .execute_and_increment_pc(
                 AB::Expr::from_canonical_usize(SUMCHECK_LAYER_EVAL.global_opcode().as_usize()),
                 [
-                    registers[2].into(),
+                    registers[4].into(),
                     registers[0].into(),
                     registers[1].into(),
                     native_as.into(),
                     native_as.into(),
-                    header_row_specific.prod_id.into(),
-                    header_row_specific.logup_id.into(),
+                    registers[2].into(),
+                    registers[3].into(),
                 ],
                 ExecutionState::new(header_row_specific.pc, first_timestamp),
                 last_timestamp - first_timestamp,
             )
             .eval(builder, header_row);
 
-        let mut header_timestamp_diff = (0..TOPLEVEL_TIMESTAMP_DIFF)
-            .into_iter()
-            .map(|i| AB::F::from_canonical_usize(i));
-        let mut header_read_records_iter = header_row_specific.read_records.iter();
         // Read registers
-        for i in 0..3usize {
+        for i in 0..5usize {
             self.memory_bridge
                 .read(
                     MemoryAddress::new(native_as, registers[i]),
                     [register_ptrs[i]],
-                    first_timestamp + header_timestamp_diff.next().unwrap(),
-                    header_read_records_iter.next().unwrap(),
+                    first_timestamp + AB::F::from_canonical_usize(i),
+                    &header_row_specific.read_records[i],
                 )
                 .eval(builder, header_row);
         }
@@ -328,8 +324,8 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .read(
                 MemoryAddress::new(native_as, register_ptrs[0]),
                 ctx,
-                first_timestamp + header_timestamp_diff.next().unwrap(),
-                header_read_records_iter.next().unwrap(),
+                first_timestamp + AB::F::from_canonical_usize(5),
+                &header_row_specific.read_records[5],
             )
             .eval(builder, header_row);
 
@@ -338,8 +334,8 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .read(
                 MemoryAddress::new(native_as, register_ptrs[1]),
                 challenges,
-                first_timestamp + header_timestamp_diff.next().unwrap(),
-                header_read_records_iter.next().unwrap(),
+                first_timestamp + AB::F::from_canonical_usize(6),
+                &header_row_specific.read_records[6],
             )
             .eval(builder, header_row);
 
@@ -350,16 +346,16 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
                     native_as,
                     register_ptrs[0] + AB::F::from_canonical_usize(CONTEXT_ARR_BASE_LEN),
                 ),
-                [max_round],
-                first_timestamp + header_timestamp_diff.next().unwrap(),
-                header_read_records_iter.next().unwrap(),
+                [max_round, is_hint_src_id, header_row_specific.prod_evals_id, header_row_specific.logup_evals_id],
+                first_timestamp + AB::F::from_canonical_usize(7),
+                &header_row_specific.read_records[7],
             )
             .eval(builder, header_row);
 
         // Write final result
         self.memory_bridge
             .write(
-                MemoryAddress::new(native_as, register_ptrs[2]),
+                MemoryAddress::new(native_as, register_ptrs[4]),
                 eval_acc,
                 last_timestamp - AB::F::ONE,
                 &header_row_specific.write_records,
@@ -393,17 +389,33 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
         );
         builder.assert_eq(prod_row * should_acc, prod_acc);
 
-        let mut prod_timestamp_diff = (0..1).into_iter().map(|i| AB::F::from_canonical_usize(i));
-
-        // TODO: read from hint space then write to memory
-        // self.memory_bridge
-        //     .read(
-        //         MemoryAddress::new(native_as, register_ptrs[2] + prod_row_specific.data_ptr),
-        //         prod_row_specific.p,
-        //         start_timestamp,
-        //         &prod_row_specific.read_records[0],
-        //     )
-        //     .eval(builder, prod_row * within_round_limit);
+        // Read p1, p2 from witness arrays
+        self.memory_bridge
+            .read(
+                MemoryAddress::new(native_as, register_ptrs[2] + prod_row_specific.data_ptr),
+                prod_row_specific.p,
+                start_timestamp,
+                &MemoryReadAuxCols {
+                    base: prod_row_specific.ps_record.base,
+                },
+            )
+            .eval(
+                builder,
+                (prod_in_round_evaluation + prod_next_round_evaluation) * not(is_hint_src_id),
+            );
+
+        // Obtain p1, p2 from hint space and write back to witness arrays
+        self.memory_bridge
+            .write(
+                MemoryAddress::new(native_as, register_ptrs[2] + prod_row_specific.data_ptr),
+                prod_row_specific.p,
+                start_timestamp,
+                &prod_row_specific.ps_record,
+            )
+            .eval(
+                builder,
+                (prod_in_round_evaluation + prod_next_round_evaluation) * is_hint_src_id,
+            );
 
         let p1: [AB::Var; EXT_DEG] = prod_row_specific.p[0..EXT_DEG].try_into().unwrap();
         let p2: [AB::Var; EXT_DEG] = prod_row_specific.p[EXT_DEG..(EXT_DEG * 2)]
@@ -414,16 +426,14 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .write(
                 MemoryAddress::new(
                     native_as,
-                    register_ptrs[2] + curr_prod_n * AB::F::from_canonical_usize(EXT_DEG),
+                    register_ptrs[4] + curr_prod_n * AB::F::from_canonical_usize(EXT_DEG),
                 ),
                 prod_row_specific.p_evals,
-                start_timestamp + prod_timestamp_diff.next().unwrap(),
+                start_timestamp + AB::F::ONE,
                 &prod_row_specific.write_record,
             )
             .eval(builder, prod_row * within_round_limit);
 
-        assert!(prod_timestamp_diff.next().is_none());
-
         // Calculate evaluations
         let next_round_p_evals = FieldExtension::add(
             FieldExtension::multiply::<AB::Var, AB::Expr>(p1, c1),
@@ -486,17 +496,33 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
         );
         builder.assert_eq(logup_row * should_acc, logup_acc);
 
-        let mut logup_timestamp_diff = (0..2).into_iter().map(|i| AB::F::from_canonical_usize(i));
-
-        // self.memory_bridge
-        //     .read(
-        //         MemoryAddress::new(native_as, register_ptrs[3] + logup_row_specific.data_ptr),
-        //         logup_row_specific.pq,
-        //         start_timestamp,
-        //         &logup_row_specific.read_records[0],
-        //     )
-        //     .eval(builder, logup_row * within_round_limit);
-
+        // Read p1, p2, q1, q2 from witness arrays
+        self.memory_bridge
+            .read(
+                MemoryAddress::new(native_as, register_ptrs[3] + logup_row_specific.data_ptr),
+                logup_row_specific.pq,
+                start_timestamp,
+                &MemoryReadAuxCols {
+                    base: logup_row_specific.pqs_record.base,
+                },
+            )
+            .eval(
+                builder,
+                (logup_in_round_evaluation + logup_next_round_evaluation) * not(is_hint_src_id),
+            );
+
+        // Obtain p1, p2, q1, q2 from hint space
+        self.memory_bridge
+            .write(
+                MemoryAddress::new(native_as, register_ptrs[3] + logup_row_specific.data_ptr),
+                logup_row_specific.pq,
+                start_timestamp,
+                &logup_row_specific.pqs_record,
+            )
+            .eval(
+                builder,
+                (logup_in_round_evaluation + logup_next_round_evaluation) * is_hint_src_id,
+            );
         let p1: [_; EXT_DEG] = logup_row_specific.pq[0..EXT_DEG].try_into().unwrap();
         let p2: [_; EXT_DEG] = logup_row_specific.pq[EXT_DEG..(EXT_DEG * 2)]
             .try_into()
@@ -513,11 +539,11 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .write(
                 MemoryAddress::new(
                     native_as,
-                    register_ptrs[2]
+                    register_ptrs[4]
                         + (num_prod_spec + curr_logup_n) * AB::F::from_canonical_usize(EXT_DEG),
                 ),
                 logup_row_specific.p_evals,
-                start_timestamp + logup_timestamp_diff.next().unwrap(),
+                start_timestamp + AB::F::ONE,
                 &logup_row_specific.write_records[0],
             )
             .eval(builder, logup_row * within_round_limit);
@@ -527,12 +553,12 @@ impl<AB: InteractionBuilder> Air<AB> for NativeSumcheckAir {
             .write(
                 MemoryAddress::new(
                     native_as,
-                    register_ptrs[2]
+                    register_ptrs[4]
                         + (num_prod_spec + num_logup_spec + curr_logup_n)
                             * AB::F::from_canonical_usize(EXT_DEG),
                 ),
                 logup_row_specific.q_evals,
-                start_timestamp + logup_timestamp_diff.next().unwrap(),
+                start_timestamp + AB::F::TWO,
                 &logup_row_specific.write_records[1],
             )
             .eval(builder, logup_row * within_round_limit);