working, but requiring code cleanup

Author: erabinov

sp1-gpu/crates/logup_gkr/src/lib.rs

@@ -552,8 +552,8 @@ mod tests {
     #[serial]
     fn test_logup_gkr_e2e() {
         let rt = tokio::runtime::Runtime::new().unwrap();
-        let (machine, record, program) =
-            rt.block_on(tracegen_setup::setup(&test_artifacts::FIBONACCI_ELF, SP1Stdin::new()));
+        let (machine, record, program) = rt
+            .block_on(tracegen_setup::setup(&test_artifacts::SSZ_WITHDRAWALS_ELF, SP1Stdin::new()));
 
         run_sync_in_place(|scope| {
             // *********** Generate traces using the host tracegen. ***********

sp1-gpu/crates/shard_prover/src/prover.rs

@@ -810,7 +810,7 @@ mod tests {
             let buffer = PinnedBuffer::<Felt>::with_capacity(capacity);
             let queue = Arc::new(WorkerQueue::new(vec![buffer]));
             let buffer = queue.pop().await.unwrap();
-            let (_public_values, jagged_trace_data, _shard_chips, _permit) = full_tracegen(
+            let (_public_values, jagged_trace_data, shard_chips, _permit) = full_tracegen(
                 &machine,
                 program.clone(),
                 Arc::new(record),

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

@@ -1,17 +1,22 @@
 use crate::data::{InfoBuffer, JaggedDenseInfo};
+use slop_air::BaseAir;
 use slop_algebra::{AbstractField, ExtensionField, Field};
 use slop_alloc::{Backend, Buffer, CpuBackend, HasBackend, Slice};
 use slop_commit::Rounds;
 use slop_multilinear::{Evaluations, 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 sp1_hypercube::air::MachineAir;
+use sp1_hypercube::Chip;
+use std::collections::{BTreeMap, BTreeSet};
 use std::iter::once;
 
 pub trait JaggedFixLastVariableKernel<K: Field> {
@@ -123,20 +128,99 @@ 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()
+        .map(|index| index.num_polys)
+        .chain(trace_data.main_table_index.values().map(|index| index.num_polys))
+        .sum::<usize>();
+    let mut result_buffer =
+        DeviceBuffer::with_capacity_in(total_cols + 2, backend.clone()).into_inner();
+    let mut preprocessed_buffer =
+        DeviceBuffer::with_capacity_in(trace_data.preprocessed_cols, backend.clone()).into_inner();
+    let mut main_buffer = DeviceBuffer::with_capacity_in(total_cols, backend.clone()).into_inner();
+
+    // println!("Preprocessed table index: {:?}", trace_data.preprocessed_table_index);
+    // println!("Main table index: {:?}", trace_data.main_table_index);
+
+    let trace_ptr = trace_data.dense.as_ptr();
+    for index in trace_data.preprocessed_table_index.values() {
+        let chip_preprocessed_ptr = unsafe { trace_ptr.add(index.dense_offset.start) };
+        let chip_preprocessed_view = unsafe {
+            TensorView::from_raw_parts(
+                chip_preprocessed_ptr,
+                [index.num_polys, index.poly_size].try_into().unwrap(),
+                backend.clone(),
+            )
+        };
+        if index.dense_offset.start != index.dense_offset.end {
+            let result =
+                dot_along_dim_view(chip_preprocessed_view, partial_lagrange.guts().as_view(), 1);
+            // println!(
+            //     "Result dimensions: {:?}, chip num_polys: {}",
+            //     result.shape(),
+            //     index.num_polys
+            // );
+            preprocessed_buffer.extend_from_device_slice(result.as_buffer()).unwrap();
+        }
+    }
+    for index in trace_data.main_table_index.values() {
+        let chip_main_ptr = unsafe { trace_ptr.add(index.dense_offset.start) };
+        let chip_main_view = unsafe {
+            TensorView::from_raw_parts(
+                chip_main_ptr,
+                [index.num_polys, index.poly_size].try_into().unwrap(),
+                backend.clone(),
+            )
+        };
+        if index.dense_offset.start != index.dense_offset.end {
+            let result = dot_along_dim_view(chip_main_view, partial_lagrange.guts().as_view(), 1);
+            // println!(
+            //     "Result dimensions: {:?}, chip num_polys: {}",
+            //     result.shape(),
+            //     index.num_polys
+            // );
+            main_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<_>>()
+    result_buffer.extend_from_device_slice(&preprocessed_buffer).unwrap();
+    let padding = Tensor::zeros_in([1], backend.clone());
+    result_buffer.extend_from_device_slice(padding.as_buffer()).unwrap();
+    result_buffer.extend_from_device_slice(&main_buffer).unwrap();
+    let host_result = DeviceBuffer::from_raw(result_buffer).to_host().unwrap();
+
+    // println!("Host result first 100 values: {:?}", host_result[..10].to_vec());
+
+    // 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 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.
+    // let expected_host_dense = host_dense.into_iter().step_by(4).collect::<Vec<_>>();
+    // println!("Expected host dense first 100 values: {:?}", expected_host_dense[..10].to_vec());
+
+    // for (i, (elem, expected_elem)) in host_result.iter().zip(expected_host_dense.iter()).enumerate()
+    // {
+    //     assert_eq!(
+    //         *elem, *expected_elem,
+    //         "element mismatch: elem={} expected={}, index={}",
+    //         *elem, *expected_elem, i
+    //     );
+    // }
+    host_result
 }
 
 pub fn evaluate_jagged_columns(