grandchildrice · grandchildrice · Dec 2, 2024 · Dec 2, 2024 · Dec 2, 2024 · Dec 3, 2024
diff --git a/.gitignore b/.gitignore
@@ -21,4 +21,8 @@ temp_file
 
 # ide
 .idea
-.vscode
+.vscode
+
+tmp_register_init.bin
+/jolt-core/instruction_polynomials.json
+/jolt-core/memory_polynomials.json
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/common/src/constants.rs b/common/src/constants.rs
@@ -22,3 +22,5 @@ pub const fn virtual_register_index(index: u64) -> u64 {
 //     peripheral devices || inputs || outputs || panic || termination || padding || RAM
 // Notably, we want to be able to map the VM memory address space to witness indices
 // using a constant shift, namely (RAM_WITNESS_OFFSET + RAM_START_ADDRESS)
+
+pub const TRACE_SEGMENTATION_SIZE: u64 = 100; // this size is the number of instructions in a segment
diff --git a/common/src/lib.rs b/common/src/lib.rs
@@ -1,4 +1,4 @@
 pub mod attributes;
 pub mod constants;
 pub mod rv_trace;
-pub mod serializable;
+pub mod serializable;
diff --git a/jolt-core/Cargo.toml b/jolt-core/Cargo.toml
@@ -41,7 +41,7 @@ rand = "0.7.3"
 rand_chacha = { version = "0.3.0", default-features = false }
 rand_core = { version = "0.6.4", default-features = false }
 rayon = { version = "^1.8.0", optional = true }
-serde = { version = "1.0.*", default-features = false }
+serde = { version = "1.0.*", default-features = false, features = ["derive"] }
 sha3 = "0.10.8"
 strum = "0.26.3"
 strum_macros = "0.26.4"
@@ -67,6 +67,9 @@ tokio = { version = "1.38.0", optional = true }
 alloy-primitives = "0.7.6"
 alloy-sol-types = "0.7.6"
 
+log = "0.4"
+env_logger = "0.10"
+
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }
 iai-callgrind = "0.10.2"
@@ -109,6 +112,9 @@ default = [
 ]
 host = ["dep:reqwest", "dep:tokio"]
 
+para = ["tracer/trace-segmentation"]
+ignore-all-io = []
+
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 memory-stats = "1.0.0"
 tokio = { version = "1.38.0", optional = true, features = ["rt-multi-thread"] }

diff --git a/jolt-core/src/host/mod.rs b/jolt-core/src/host/mod.rs
@@ -8,13 +8,15 @@ use std::{
     process::Command,
 };
 
+use log::debug;
 use postcard;
 use rayon::prelude::*;
 use serde::Serialize;
 
 use common::{
     constants::{
         DEFAULT_MAX_INPUT_SIZE, DEFAULT_MAX_OUTPUT_SIZE, DEFAULT_MEMORY_SIZE, DEFAULT_STACK_SIZE,
+        TRACE_SEGMENTATION_SIZE,
     },
     rv_trace::JoltDevice,
 };
@@ -225,6 +227,67 @@ impl Program {
         (io_device, trace)
     }
 
+    // TODO(moodlezoup): Make this generic over InstructionSet
+    #[tracing::instrument(skip_all, name = "Program::trace")]
+    pub fn segment_trace(
+        &mut self,
+    ) -> (
+        JoltDevice,
+        Vec<([i64; 32], Vec<u64>, u64)>,
+        Vec<Vec<JoltTraceStep<RV32I>>>,
+    ) {
+        self.build();
+        let elf = self.elf.clone().unwrap();
+        let (raw_trace, io_device, snapshots) =
+            tracer::segment_trace(&elf, &self.input, self.max_input_size, self.max_output_size);
+
+        let traces = raw_trace
+            .chunks(TRACE_SEGMENTATION_SIZE as usize)
+            .map(|segmented_raw_trace| {
+                let segmented_raw_trace = segmented_raw_trace.to_vec();
+                debug!("segmented_raw_trace len: {}", segmented_raw_trace.len());
+                let trace: Vec<_> = segmented_raw_trace
+                    .into_par_iter()
+                    .flat_map(|row| match row.instruction.opcode {
+                        tracer::RV32IM::MULH => MULHInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::MULHSU => MULHSUInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::DIV => DIVInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::DIVU => DIVUInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::REM => REMInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::REMU => REMUInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::SH => SHInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::SB => SBInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::LBU => LBUInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::LHU => LHUInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::LB => LBInstruction::<32>::virtual_trace(row),
+                        tracer::RV32IM::LH => LHInstruction::<32>::virtual_trace(row),
+                        _ => vec![row],
+                    })
+                    .map(|row| {
+                        let instruction_lookup = if let Ok(jolt_instruction) = RV32I::try_from(&row)
+                        {
+                            Some(jolt_instruction)
+                        } else {
+                            // Instruction does not use lookups
+                            None
+                        };
+
+                            JoltTraceStep {
+                            instruction_lookup,
+                            bytecode_row: BytecodeRow::from_instruction::<RV32I>(&row.instruction),
+                            memory_ops: (&row).into(),
+                            circuit_flags: row.instruction.to_circuit_flags(),
+                        }
+                    })
+                    .collect();
+                debug!("trace with virtual sequence: {}", trace.len());
+                trace
+            })
+            .collect();
+
+        (io_device, snapshots, traces)
+    }
+
     pub fn trace_analyze<F: JoltField>(mut self) -> ProgramSummary {
         self.build();
         let elf = self.elf.as_ref().unwrap();