Support Scheduled GKR

circuit-std-rs/src/logup.rs

@@ -326,6 +326,41 @@ impl LogUpSingleKeyTable {
             &alpha,
         );
 
+        assert_eq_rational(builder, &v_table, &v_query);
+    }
+    pub fn final_check_with_query_count<C: Config, B: RootAPI<C>>(
+        &mut self,
+        builder: &mut B,
+        query_count: &[Variable],
+    ) {
+        if self.table.is_empty() || self.query_keys.is_empty() {
+            panic!("empty table or empty query");
+        }
+
+        let value_len = self.table[0].len();
+
+        let alpha = builder.get_random_value();
+        let randomness = get_column_randomness(builder, value_len);
+
+        let table_combined = combine_columns(builder, &self.table, &randomness);
+        let mut inputs = vec![builder.constant(self.table.len() as u32)];
+        //append table keys
+        for i in 0..self.table.len() {
+            inputs.push(self.table[i][0]);
+        }
+        //append query keys
+        inputs.extend(self.query_keys.clone());
+        let v_table = logup_poly_val(builder, &table_combined, &query_count, &alpha);
+
+        let query_combined = combine_columns(builder, &self.query_results, &randomness);
+        let one = builder.constant(1);
+        let v_query = logup_poly_val(
+            builder,
+            &query_combined,
+            &vec![one; query_combined.len()],
+            &alpha,
+        );
+
         assert_eq_rational(builder, &v_table, &v_query);
     }
 }
@@ -455,6 +490,26 @@ impl LogUpRangeProofTable {
         );
         assert_eq_rational(builder, &v_table, &v_query);
     }
+
+    pub fn final_check_with_query_count<C: Config, B: RootAPI<C>>(
+        &mut self,
+        builder: &mut B,
+        query_count: &[Variable],
+    ) {
+        let alpha = builder.get_random_value();
+        let inputs = self.query_keys.clone();
+
+        let v_table = logup_poly_val(builder, &self.table_keys, &query_count, &alpha);
+
+        let one = builder.constant(1);
+        let v_query = logup_poly_val(
+            builder,
+            &self.query_keys,
+            &vec![one; self.query_keys.len()],
+            &alpha,
+        );
+        assert_eq_rational(builder, &v_table, &v_query);
+    }
 }
 
 pub fn query_count_hint<F: Field>(inputs: &[F], outputs: &mut [F]) -> Result<(), Error> {

expander_compiler/src/zkcuda/context.rs

@@ -577,7 +577,7 @@ impl<C: Config, H: HintCaller<CircuitField<C>>> Context<C, H> {
 
         let dm_shapes = self.propagate_and_get_shapes();
 
-        let (mut cg_kernels, cg_proof_templates, cg_commitments_lens) = if let Some(cg) = cg {
+        let (cg_kernels, cg_proof_templates, cg_commitments_lens) = if let Some(cg) = cg {
             for (i, kernel) in cg.kernels.iter().enumerate() {
                 assert_eq!(self.kernels.add(kernel), i);
             }
@@ -616,8 +616,9 @@ impl<C: Config, H: HintCaller<CircuitField<C>>> Context<C, H> {
                 .map(get_pad_shape)
                 .collect::<Vec<_>>();
             let kernel_primitive = self.kernel_primitives.get(kernel_call.kernel_id);
-            let kernel = if let Some(cg_kernels) = cg_kernels.as_mut() {
-                cg_kernels.drain(..1).next().unwrap()
+            let kernel = if cg_kernels.is_some() {
+                // Get kernel from loaded kernels by kernel_id
+                self.kernels.get(kernel_call.kernel_id).clone()
             } else {
                 let mut psi = Vec::new();
                 for (s, &ib) in pad_shapes_input.iter().zip(kernel_call.is_broadcast.iter()) {
@@ -708,8 +709,9 @@ impl<C: Config, H: HintCaller<CircuitField<C>>> Context<C, H> {
             });
         }
 
-        if let Some(cg_kernels) = cg_kernels {
-            assert!(cg_kernels.is_empty());
+        if cg_kernels.is_some() {
+            // No longer checking if cg_kernels is empty since we no longer consume it
+            // Kernels were already added earlier via self.kernels.add()
             assert_eq!(cg_proof_templates.unwrap(), self.proof_templates);
             assert_eq!(cg_commitments_lens.unwrap(), commitments_lens);
             Ok(None)

expander_compiler/src/zkcuda/proving_system/expander_no_oversubscribe/api_no_oversubscribe.rs

@@ -73,3 +73,15 @@ where
         wait_async(ClientHttpHelper::request_exit())
     }
 }
+
+impl<ZC: ZKCudaConfig> ExpanderNoOverSubscribe<ZC>
+where
+    <GetPCS<ZC> as ExpanderPCS<GetFieldConfig<ZC>>>::Commitment:
+        AsRef<<GetPCS<ZC> as ExpanderPCS<GetFieldConfig<ZC>>>::Commitment>,
+{
+    /// Lightweight prove that doesn't require computation_graph or prover_setup.
+    /// Use this after setup() to allow releasing those large data structures before proving.
+    pub fn prove_lightweight(device_memories: Vec<Vec<SIMDField<ZC::ECCConfig>>>) {
+        client_send_witness_and_prove::<ZC::GKRConfig, ZC::ECCConfig>(device_memories);
+    }
+}