perf(gpu): more efficient column evaluation claims at the end of GKR

sp1-gpu/crates/logup_gkr/benches/gkr.rs

@@ -18,19 +18,18 @@ use std::sync::Arc;
 use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use rand::{rngs::StdRng, Rng, SeedableRng};
 use slop_challenger::{FieldChallenger, IopCtx};
-use slop_multilinear::{MultilinearPcsChallenger, Point};
+use slop_multilinear::Point;
 use sp1_core_machine::riscv::RiscvAir;
-use sp1_gpu_cudart::{DevicePoint, TaskScope};
+use sp1_gpu_cudart::TaskScope;
 use sp1_gpu_jagged_tracegen::test_utils::bench_utils::{
     with_trace_source, FullKind, RealTraceData,
 };
 use sp1_gpu_jagged_tracegen::test_utils::tracegen_setup::CORE_MAX_LOG_ROW_COUNT;
-use sp1_gpu_logup_gkr::{
-    generate_gkr_circuit, prove_gkr_circuit, CudaLogUpGkrOptions, Interactions,
-};
+use sp1_gpu_logup_gkr::{generate_gkr_circuit, prove_logup_gkr, CudaLogUpGkrOptions, Interactions};
 use sp1_gpu_utils::{Ext, Felt, TestGC};
 use sp1_hypercube::air::MachineAir;
 use sp1_hypercube::Chip;
+use sp1_primitives::SP1GlobalContext;
 
 /// `prove_gkr_circuit` flag: when true, the prover recomputes the first layer on demand from
 /// the raw trace each time it walks back to the leaves; when false, it caches the materialized
@@ -125,66 +124,23 @@ fn run_prove<R: Rng>(
 ) {
     let RealTraceData { machine: _, cluster, public_values: _, device_mle } = data;
     let interactions = build_interactions(&cluster, scope);
-    let beta_dim = beta_seed_dim(&cluster);
-
-    // initial_number_of_variables = num_interaction_variables + 1, where
-    // num_interaction_variables is `log2(num_total_interactions).next_power_of_two()`.
-    let num_interactions: usize =
-        cluster.iter().map(|chip| chip.sends().len() + chip.receives().len()).sum();
-    let initial_number_of_variables = num_interactions.next_power_of_two().trailing_zeros() + 1;
 
     let mut group = c.benchmark_group("prove");
     group.sample_size(10);
     group.bench_with_input(id, &(), |b, _| {
         b.iter_batched(
             || {
-                // Per-iteration setup. Mirrors `prove_logup_gkr` lines 198-244 but skips the
-                // observe-output-claims step (which only affects challenger state, not the
-                // prove call's runtime). We still build the circuit and compute the initial
-                // numerator/denominator evaluations because `prove_gkr_circuit` consumes both.
-                let mut challenger = TestGC::default_challenger();
-                let alpha: Ext = challenger.sample_ext_element();
-                let beta_seed: Point<Ext> =
-                    (0..beta_dim).map(|_| challenger.sample_ext_element::<Ext>()).collect();
-
-                let (output, circuit) = generate_gkr_circuit(
+                let challenger = TestGC::default_challenger();
+                scope.synchronize_blocking().unwrap();
+                (interactions.clone(), challenger)
+            },
+            |(interactions, mut challenger)| {
+                let result = prove_logup_gkr::<SP1GlobalContext, _>(
                     &cluster,
-                    interactions.clone(),
+                    interactions,
                     &device_mle,
-                    alpha,
-                    beta_seed,
                     GKR_OPTIONS,
-                    scope.clone(),
-                );
-
-                let first_eval_point = challenger.sample_point::<Ext>(initial_number_of_variables);
-                let first_point =
-                    DevicePoint::from_host(&first_eval_point, output.numerator.backend())
-                        .unwrap()
-                        .into_inner();
-                let first_point_eq = DevicePoint::new(first_point).partial_lagrange();
-                let first_numerator_eval =
-                    output.numerator.eval_at_eq(&first_point_eq).to_host_vec().unwrap()[0];
-                let first_denominator_eval =
-                    output.denominator.eval_at_eq(&first_point_eq).to_host_vec().unwrap()[0];
-
-                scope.synchronize_blocking().unwrap();
-                (
-                    first_numerator_eval,
-                    first_denominator_eval,
-                    first_eval_point,
-                    circuit,
-                    challenger,
-                )
-            },
-            |(num_eval, den_eval, eval_point, circuit, mut challenger)| {
-                let result = prove_gkr_circuit(
-                    num_eval,
-                    den_eval,
-                    eval_point,
-                    circuit,
                     &mut challenger,
-                    RECOMPUTE_FIRST_LAYER,
                 );
                 scope.synchronize_blocking().unwrap();
                 black_box(result)

sp1-gpu/crates/zerocheck/src/primitives.rs

@@ -2,14 +2,18 @@ use crate::data::{InfoBuffer, JaggedDenseInfo};
 use slop_algebra::{AbstractField, ExtensionField, Field};
 use slop_alloc::{Backend, Buffer, CpuBackend, HasBackend, Slice};
 use slop_commit::Rounds;
+
 use slop_multilinear::{MleEval, Point};
-use slop_tensor::Tensor;
+use slop_tensor::{Tensor, TensorView};
+
 use sp1_gpu_cudart::sys::runtime::KernelPtr;
 use sp1_gpu_cudart::sys::v2_kernels::{
     fix_last_variable_jagged_ext, fix_last_variable_jagged_felt, fix_last_variable_jagged_info,
     initialize_jagged_info,
 };
-use sp1_gpu_cudart::{args, DeviceBuffer, DevicePoint, DeviceTensor, TaskScope};
+use sp1_gpu_cudart::{
+    args, dot_along_dim_view, DeviceBuffer, DevicePoint, DeviceTensor, TaskScope,
+};
 use sp1_gpu_utils::{Ext, Felt, JaggedMle, JaggedTraceMle, TraceDenseData, TraceOffset};
 use std::collections::BTreeMap;
 use std::iter::once;
@@ -123,20 +127,39 @@ where
 
 #[inline(always)]
 pub fn evaluate_traces(traces: &JaggedTraceMle<Felt, TaskScope>, point: &Point<Ext>) -> Vec<Ext> {
-    let mut next_input_jagged_trace_mle =
-        evaluate_jagged_fix_last_variable(traces, *point.last().unwrap());
-    for alpha in point.iter().rev().skip(1) {
-        next_input_jagged_trace_mle =
-            evaluate_jagged_fix_last_variable(&next_input_jagged_trace_mle, *alpha);
+    let trace_data = traces.dense();
+    let backend = traces.backend();
+    let device_point = DevicePoint::from_host(point, backend).unwrap();
+    let partial_lagrange = device_point.partial_lagrange();
+    let total_cols = trace_data
+        .preprocessed_table_index
+        .values()
+        .chain(trace_data.main_table_index.values())
+        .map(|index| index.num_polys)
+        .sum::<usize>();
+    let mut result_buffer =
+        DeviceBuffer::with_capacity_in(total_cols, backend.clone()).into_inner();
+
+    let trace_ptr = trace_data.dense.as_ptr();
+    let chip_indices =
+        trace_data.preprocessed_table_index.values().chain(trace_data.main_table_index.values());
+    for index in chip_indices {
+        if index.dense_offset.start == index.dense_offset.end {
+            continue;
+        }
+        let chip_ptr = unsafe { trace_ptr.add(index.dense_offset.start) };
+        let chip_view = unsafe {
+            TensorView::from_raw_parts(
+                chip_ptr,
+                [index.num_polys, index.poly_size].try_into().unwrap(),
+                backend.clone(),
+            )
+        };
+        let result = dot_along_dim_view(chip_view, partial_lagrange.guts().as_view(), 1);
+        result_buffer.extend_from_device_slice(result.as_buffer()).unwrap();
     }
 
-    let host_dense = DeviceBuffer::from_raw(next_input_jagged_trace_mle.dense_data.dense.clone())
-        .to_host()
-        .unwrap()
-        .to_vec();
-
-    // Only every four elements is not padding.
-    host_dense.into_iter().step_by(4).collect::<Vec<_>>()
+    DeviceBuffer::from_raw(result_buffer).to_host().unwrap()
 }
 
 pub fn evaluate_jagged_columns(
@@ -353,8 +376,6 @@ pub fn round_batch_evaluations(
     let preprocessed_host_evaluations =
         preprocessed_host_evaluations.into_iter().collect::<Vec<_>>();
 
-    // Skip the padding column, if it exists.
-    evals_so_far = jagged_trace_mle.dense().preprocessed_cols;
     let mut main_host_evaluations = Vec::new();
     for offset in jagged_trace_mle.dense().main_table_index.values() {
         if offset.poly_size == 0 {