Extend Poseidon2 Chip for MULTI_OBSERVE (#5)

* Establish entry

* Add opcode MULTI_OBSERVE

* feat: rewrite sample_ext using new IR instruction

* chore(doc): update comment

* feat(openvm_circuit): store span info in cycle tracker for convenient debug info

* Fix annotation, address spaces and register assignments

* Consolidate MultiObserve chip debugging

* Preserve testing environment

* Fix constraints

* Fix constraints

* Fix constraints

* Fix constraints

* Fix constraint

* Fix constraints

* Complete constraints

* Remove debug code

* Remove multi observe logic from poseidon2 chip

* Resolve minor build issue

* Resolve minor build issue

* set log_blowup=1

* Debug constraints

* Debug constraints

* Debug constraints

* Debug constraints

* Remove standalone chip

* Populate poseidon2 columns

* Reduce columns

* Add more constraints

* Remove debug flags

* Grammatic correction

* Add constraints

* remove unrelated log file

---------

Co-authored-by: Rohit Narurkar <[email protected]>
Co-authored-by: kunxian xia <[email protected]>

Author: 3 people

crates/vm/src/arch/segment.rs

@@ -282,16 +282,18 @@ impl<F: PrimeField32, VC: VmConfig<F>> ExecutionSegment<F, VC> {
                         Some(SysPhantom::CtStart) =>
                         {
                             #[cfg(feature = "bench-metrics")]
-                            metrics
-                                .cycle_tracker
-                                .start(dsl_instr.cloned().unwrap_or("Default".to_string()))
+                            metrics.cycle_tracker.start(
+                                dsl_instr.cloned().unwrap_or("Default".to_string()),
+                                metrics.cycle_count,
+                            )
                         }
                         Some(SysPhantom::CtEnd) =>
                         {
                             #[cfg(feature = "bench-metrics")]
-                            metrics
-                                .cycle_tracker
-                                .end(dsl_instr.cloned().unwrap_or("Default".to_string()))
+                            metrics.cycle_tracker.end(
+                                dsl_instr.cloned().unwrap_or("Default".to_string()),
+                                metrics.cycle_count,
+                            )
                         }
                         _ => {}
                     }

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

@@ -1,7 +1,18 @@
+/// 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,
+}
+
 #[derive(Clone, Debug, Default)]
 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,
 }
 
 impl CycleTracker {
@@ -10,29 +21,41 @@ impl CycleTracker {
     }
 
     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) {
+    /// 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, cycles_count: usize) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
-        self.stack.push(name);
+        self.stack.push(SpanInfo {
+            tag: name.clone(),
+            start: cycles_count,
+        });
+        let padding = "│ ".repeat(self.depth);
+        tracing::info!("{}┌╴{}", padding, name);
+        self.depth += 1;
     }
 
     /// 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) {
         // 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");
+        let SpanInfo { tag, start } = self.stack.pop().unwrap();
+        assert_eq!(tag, name, "Stack top does not match name");
+        self.depth -= 1;
+        let padding = "│ ".repeat(self.depth);
+        let span_cycles = cycles_count - start;
+        tracing::info!("{}└╴{} cycles", padding, span_cycles);
     }
 
     /// Ends the current cycle tracker span.
@@ -42,7 +65,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

@@ -117,7 +117,7 @@ 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 {

crates/vm/src/system/memory/controller/mod.rs

@@ -59,7 +59,7 @@ pub const MERKLE_AIR_OFFSET: usize = 1;
 pub const BOUNDARY_AIR_OFFSET: usize = 0;
 
 #[repr(C)]
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
 pub struct RecordId(pub usize);
 
 pub type MemoryImage<F> = AddressMap<F, PAGE_SIZE>;

extensions/native/circuit/Cargo.toml

@@ -19,7 +19,6 @@ openvm-instructions = { workspace = true }
 openvm-rv32im-circuit = { workspace = true }
 openvm-native-compiler = { workspace = true }
 
-
 strum.workspace = true
 itertools.workspace = true
 tracing.workspace = true

extensions/native/circuit/src/extension.rs

@@ -1,4 +1,4 @@
-use air::VerifyBatchBus;
+use poseidon2::air::VerifyBatchBus;
 use alu_native_adapter::AluNativeAdapterChip;
 use branch_native_adapter::BranchNativeAdapterChip;
 use derive_more::derive::From;
@@ -30,7 +30,7 @@ use strum::IntoEnumIterator;
 
 use crate::{
     adapters::{convert_adapter::ConvertAdapterChip, *},
-    chip::NativePoseidon2Chip,
+    poseidon2::chip::NativePoseidon2Chip,
     phantom::*,
     *,
 };
@@ -203,6 +203,7 @@ impl<F: PrimeField32> VmExtension<F> for Native {
                 VerifyBatchOpcode::VERIFY_BATCH.global_opcode(),
                 Poseidon2Opcode::PERM_POS2.global_opcode(),
                 Poseidon2Opcode::COMP_POS2.global_opcode(),
+                Poseidon2Opcode::MULTI_OBSERVE.global_opcode(),
             ],
         )?;