Native Extension For Ceno zkVM

.gitignore

@@ -9,6 +9,9 @@ Cargo.lock
 **/.env
 .DS_Store
 
+# Log outputs
+*.log
+
 .cache/
 rustc-*
 

crates/sdk/src/keygen/dummy.rs

@@ -54,7 +54,7 @@ use crate::{
 /// - trace heights ordered by AIR ID
 ///
 /// All trace heights are rounded to the next power of two (or 0 -> 0).
-pub(super) fn compute_root_proof_heights(
+pub fn compute_root_proof_heights(
     root_vm: &mut VirtualMachine<BabyBearPoseidon2RootEngine, NativeCpuBuilder>,
     root_committed_exe: &VmCommittedExe<BabyBearPoseidon2RootConfig>,
     dummy_internal_proof: &Proof<SC>,

crates/sdk/src/keygen/mod.rs

@@ -54,7 +54,7 @@ use crate::{
 };
 
 pub mod asm;
-pub(crate) mod dummy;
+pub mod dummy;
 pub mod perm;
 #[cfg(feature = "evm-prove")]
 pub mod static_verifier;

crates/sdk/src/keygen/perm.rs

@@ -4,7 +4,7 @@ use std::cmp::Reverse;
 use openvm_continuations::verifier::common::types::SpecialAirIds;
 
 /// Permutation of the AIR IDs to order them by forced trace heights.
-pub(crate) struct AirIdPermutation {
+pub struct AirIdPermutation {
     pub perm: Vec<usize>,
 }
 
@@ -47,7 +47,7 @@ impl AirIdPermutation {
         ret
     }
     /// arr[i] <- arr[perm[i]]
-    pub(crate) fn permute<T>(&self, arr: &mut [T]) {
+    pub fn permute<T>(&self, arr: &mut [T]) {
         debug_assert_eq!(arr.len(), self.perm.len());
         let mut perm = self.perm.clone();
         for i in 0..perm.len() {

crates/sdk/src/prover/agg.rs

@@ -27,12 +27,12 @@ where
     E: StarkFriEngine<SC = SC>,
     NativeBuilder: VmBuilder<E, VmConfig = NativeConfig>,
 {
-    leaf_prover: VmInstance<E, NativeBuilder>,
-    leaf_controller: LeafProvingController,
+    pub leaf_prover: VmInstance<E, NativeBuilder>,
+    pub leaf_controller: LeafProvingController,
 
     pub internal_prover: VmInstance<E, NativeBuilder>,
     #[cfg(feature = "evm-prove")]
-    root_prover: RootVerifierLocalProver,
+    pub root_prover: RootVerifierLocalProver,
     pub num_children_internal: usize,
     pub max_internal_wrapper_layers: usize,
 }

crates/vm/src/arch/testing/cpu.rs

@@ -37,16 +37,17 @@ use crate::{
         testing::{
             execution::air::ExecutionDummyAir,
             program::{air::ProgramDummyAir, ProgramTester},
-            ExecutionTester, MemoryTester, TestBuilder, TestChipHarness, EXECUTION_BUS, MEMORY_BUS,
-            MEMORY_MERKLE_BUS, POSEIDON2_DIRECT_BUS, RANGE_CHECKER_BUS, READ_INSTRUCTION_BUS,
+            ExecutionTester, MemoryTester, TestBuilder, TestChipHarness, EXECUTION_BUS, HINT_BUS,
+            MEMORY_BUS, MEMORY_MERKLE_BUS, POSEIDON2_DIRECT_BUS, RANGE_CHECKER_BUS,
+            READ_INSTRUCTION_BUS,
         },
         vm_poseidon2_config, Arena, ExecutionBridge, ExecutionBus, ExecutionState,
         MatrixRecordArena, MemoryConfig, PreflightExecutor, Streams, VmStateMut,
     },
     system::{
         memory::{
             adapter::records::arena_size_bound,
-            offline_checker::{MemoryBridge, MemoryBus},
+            offline_checker::{HintBridge, HintBus, MemoryBridge, MemoryBus},
             online::TracingMemory,
             MemoryAirInventory, MemoryController, SharedMemoryHelper, CHUNK,
         },
@@ -258,10 +259,13 @@ impl<F: PrimeField32> VmChipTestBuilder<F> {
     }
 
     pub fn system_port(&self) -> SystemPort {
+        let hint_bus = HintBus::new(HINT_BUS);
+        let hint_bridge = HintBridge::new(hint_bus);
         SystemPort {
             execution_bus: self.execution.bus,
             program_bus: self.program.bus,
             memory_bridge: self.memory_bridge(),
+            hint_bridge,
         }
     }
 

crates/vm/src/arch/testing/cuda.rs

@@ -50,7 +50,7 @@ use crate::{
             execution::{air::ExecutionDummyAir, DeviceExecutionTester},
             memory::DeviceMemoryTester,
             program::{air::ProgramDummyAir, DeviceProgramTester},
-            TestBuilder, TestChipHarness, EXECUTION_BUS, MEMORY_BUS, MEMORY_MERKLE_BUS,
+            TestBuilder, TestChipHarness, EXECUTION_BUS, HINT_BUS, MEMORY_BUS, MEMORY_MERKLE_BUS,
             POSEIDON2_DIRECT_BUS, READ_INSTRUCTION_BUS,
         },
         Arena, DenseRecordArena, ExecutionBridge, ExecutionBus, ExecutionState, MatrixRecordArena,
@@ -59,7 +59,7 @@ use crate::{
     system::{
         cuda::{poseidon2::Poseidon2PeripheryChipGPU, DIGEST_WIDTH},
         memory::{
-            offline_checker::{MemoryBridge, MemoryBus},
+            offline_checker::{HintBridge, HintBus, MemoryBridge, MemoryBus},
             MemoryAirInventory, SharedMemoryHelper,
         },
         poseidon2::air::Poseidon2PeripheryAir,
@@ -393,8 +393,14 @@ impl GpuChipTestBuilder {
             execution_bus: self.execution_bus(),
             program_bus: self.program_bus(),
             memory_bridge: self.memory_bridge(),
+            hint_bridge: self.hint_bridge(),
         }
     }
+
+    pub fn hint_bridge(&self) -> HintBridge {
+        let hint_bus = HintBus::new(HINT_BUS);
+        HintBridge::new(hint_bus)
+    }
     pub fn execution_bridge(&self) -> ExecutionBridge {
         ExecutionBridge::new(self.execution.bus(), self.program.bus())
     }

crates/vm/src/arch/testing/mod.rs

@@ -29,6 +29,7 @@ pub const BITWISE_OP_LOOKUP_BUS: BusIndex = 9;
 pub const BYTE_XOR_BUS: BusIndex = 10;
 pub const RANGE_TUPLE_CHECKER_BUS: BusIndex = 11;
 pub const MEMORY_MERKLE_BUS: BusIndex = 12;
+pub const HINT_BUS: BusIndex = 13;
 
 pub const RANGE_CHECKER_BUS: BusIndex = 4;
 

crates/vm/src/arch/vm.rs

@@ -910,7 +910,7 @@ where
     #[getset(get = "pub")]
     exe: Arc<VmExe<Val<E::SC>>>,
     #[getset(get = "pub", get_mut = "pub")]
-    state: Option<VmState<Val<E::SC>, GuestMemory>>,
+    pub state: Option<VmState<Val<E::SC>, GuestMemory>>,
 }
 
 impl<E, VB> VmInstance<E, VB>
@@ -1056,6 +1056,8 @@ where
         let mut trace_heights = trace_heights.to_vec();
         trace_heights[PUBLIC_VALUES_AIR_ID] = vm.config().as_ref().num_public_values as u32;
         let state = self.state.take().expect("State should always be present");
+        #[cfg(feature = "metrics")]
+        let debug_infos = state.metrics.debug_infos.clone();
         let num_custom_pvs = state.custom_pvs.len();
         let (proof, final_memory) = vm.prove(&mut self.interpreter, state, None, &trace_heights)?;
         let final_memory = final_memory.ok_or(ExecutionError::DidNotTerminate)?;
@@ -1068,6 +1070,10 @@ where
             DEFAULT_RNG_SEED,
             num_custom_pvs,
         ));
+        #[cfg(feature = "metrics")]
+        {
+            self.state.as_mut().unwrap().metrics.debug_infos = debug_infos;
+        }
         Ok(proof)
     }
 }

crates/vm/src/metrics/cycle_tracker/mod.rs

@@ -1,38 +1,98 @@
+use std::collections::BTreeMap;
+
+/// Stats for a nested span in the execution segment that is tracked by the [`CycleTracker`].
 #[derive(Clone, Debug, Default)]
+pub struct SpanInfo {
+    /// The name of the span.
+    pub tag: String,
+    /// The cycle count at which the span starts.
+    pub start: usize,
+    /// Maps (dsl_ir, opcode) to number of times opcode was executed
+    pub counts: BTreeMap<(Option<String>, String), usize>,
+}
+
+#[derive(Clone, Debug)]
 pub struct CycleTracker {
     /// Stack of span names, with most recent at the end
-    stack: Vec<String>,
+    stack: Vec<SpanInfo>,
+    /// Depth of the stack.
+    depth: usize,
+    max_depth: usize,
+}
+
+impl Default for CycleTracker {
+    fn default() -> Self {
+        Self {
+            stack: Vec::new(),
+            depth: 0,
+            max_depth: std::env::var("CYCLE_TRACKER_MAX_DEPTH")
+                .ok()
+                .and_then(|s| s.parse().ok())
+                .unwrap_or(2),
+        }
+    }
 }
 
 impl CycleTracker {
-    pub fn new() -> Self {
-        Self::default()
+    pub fn new(max_depth: usize) -> Self {
+        Self {
+            max_depth,
+            ..Default::default()
+        }
     }
 
     pub fn top(&self) -> Option<&String> {
-        self.stack.last()
+        match self.stack.last() {
+            Some(span) => Some(&span.tag),
+            _ => None,
+        }
     }
 
     /// Starts a new cycle tracker span for the given name.
     /// If a span already exists for the given name, it ends the existing span and pushes a new one
     /// to the vec.
-    pub fn start(&mut self, mut name: String) {
+    pub fn start(&mut self, mut name: String, cycles_count: usize, num_insns_by_dsl: &BTreeMap<(Option<String>, String), usize>) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
-        self.stack.push(name);
+        self.depth += 1;
+        if self.depth > self.max_depth {
+            return;
+        }
+        self.stack.push(SpanInfo {
+            tag: name.clone(),
+            start: cycles_count,
+            counts: num_insns_by_dsl.clone(),
+        });
+        let padding = "│ ".repeat(self.depth);
+        tracing::info!("{}┌╴{}", padding, name);
     }
 
     /// Ends the cycle tracker span for the given name.
     /// If no span exists for the given name, it panics.
-    pub fn end(&mut self, mut name: String) {
+    pub fn end(&mut self, mut name: String, cycles_count: usize, num_insns_by_dsl: &BTreeMap<(Option<String>, String), usize>) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
-        let stack_top = self.stack.pop();
-        assert_eq!(stack_top.unwrap(), name, "Stack top does not match name");
+        // keep padding info before pop
+        let padding = "│ ".repeat(self.depth);
+        self.depth -= 1;
+        if self.depth >= self.max_depth {
+            return;
+        }
+        let SpanInfo { tag, start, counts: num_insns_start } = self.stack.pop().unwrap();
+        assert_eq!(tag, name, "Stack top does not match name");
+        let span_cycles = cycles_count - start;
+        for (dsl_opcode, num_insns) in num_insns_by_dsl {
+            let start_count = num_insns_start.get(dsl_opcode).cloned().unwrap_or(0);
+            let span_count = num_insns - start_count;
+            if span_count > 0 {
+                tracing::info!("{}│   ({:?},{}): {} instructions", padding, dsl_opcode.0, dsl_opcode.1, span_count);
+            }
+        }
+        tracing::info!("{}└╴({}) {} cycles, abs: {}", padding, name, span_cycles, cycles_count);
     }
 
     /// Ends the current cycle tracker span.
@@ -42,7 +102,11 @@ impl CycleTracker {
 
     /// Get full name of span with all parent names separated by ";" in flamegraph format
     pub fn get_full_name(&self) -> String {
-        self.stack.join(";")
+        self.stack
+            .iter()
+            .map(|span_info| span_info.tag.clone())
+            .collect::<Vec<String>>()
+            .join(";")
     }
 }
 

crates/vm/src/metrics/mod.rs

@@ -40,6 +40,7 @@ pub struct VmMetrics {
     /// Metric collection tools. Only collected when "perf-metrics" feature is enabled.
     pub cycle_tracker: CycleTracker,
 
+    pub(crate) num_insns_executed: usize,
     pub(crate) current_trace_cells: Vec<usize>,
 
     /// Backtrace for guest debug panic display
@@ -68,6 +69,7 @@ pub fn update_instruction_metrics<F, RA, Executor>(
     {
         let pc = state.pc();
         state.metrics.update_backtrace(pc);
+        state.metrics.num_insns_executed += 1;
     }
 
     #[cfg(feature = "perf-metrics")]
@@ -224,13 +226,13 @@ impl VmMetrics {
                 .map(|(_, func)| (*func).clone())
                 .unwrap();
             if pc == self.current_fn.start {
-                self.cycle_tracker.start(self.current_fn.name.clone());
+                // self.cycle_tracker.start(self.current_fn.name.clone(), 0);
             } else {
                 while let Some(name) = self.cycle_tracker.top() {
                     if name == &self.current_fn.name {
                         break;
                     }
-                    self.cycle_tracker.force_end();
+                    // self.cycle_tracker.force_end();
                 }
             }
         };

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

@@ -10,7 +10,7 @@ use openvm_stark_backend::{
     interaction::InteractionBuilder, p3_air::AirBuilder, p3_field::FieldAlgebra,
 };
 
-use super::bus::MemoryBus;
+use super::bus::{HintBus, MemoryBus};
 use crate::system::memory::{
     offline_checker::columns::{
         MemoryBaseAuxCols, MemoryReadAuxCols, MemoryReadOrImmediateAuxCols, MemoryWriteAuxCols,
@@ -326,3 +326,52 @@ impl MemoryOfflineChecker {
             .eval(builder, enabled);
     }
 }
+
+/// The [HintBridge] is used to constrain hint space lookups.
+/// Consumer chips call `lookup` to verify that values they read from hint_space
+/// match what was originally loaded via the hint bus lookup table.
+#[derive(Clone, Copy, Debug)]
+pub struct HintBridge {
+    hint_bus: HintBus,
+}
+
+impl HintBridge {
+    /// Create a new [HintBridge] with the provided hint bus.
+    pub fn new(hint_bus: HintBus) -> Self {
+        Self { hint_bus }
+    }
+
+    pub fn hint_bus(&self) -> HintBus {
+        self.hint_bus
+    }
+
+    /// Perform a lookup on the hint bus for a single element.
+    ///
+    /// Constrains that `(hint_id, offset, value)` exists in the hint lookup table.
+    /// Caller must constrain that `enabled` is boolean.
+    pub fn lookup<AB: InteractionBuilder>(
+        &self,
+        builder: &mut AB,
+        hint_id: impl Into<AB::Expr>,
+        offset: impl Into<AB::Expr>,
+        value: impl Into<AB::Expr>,
+        enabled: impl Into<AB::Expr>,
+    ) {
+        self.hint_bus.lookup(builder, hint_id, offset, value, enabled);
+    }
+
+    /// Add a key to the hint lookup table.
+    ///
+    /// Provider chips call this to register that `(hint_id, offset, value)` is available.
+    pub fn provide<AB: InteractionBuilder>(
+        &self,
+        builder: &mut AB,
+        hint_id: impl Into<AB::Expr>,
+        offset: impl Into<AB::Expr>,
+        value: impl Into<AB::Expr>,
+        num_lookups: impl Into<AB::Expr>,
+    ) {
+        self.hint_bus
+            .provide(builder, hint_id, offset, value, num_lookups);
+    }
+}