Clean up profiling a little.

executor/src/witgen/machines/fixed_lookup_machine.rs

@@ -12,11 +12,12 @@ use powdr_number::FieldElement;
 
 use crate::witgen::affine_expression::AffineExpression;
 use crate::witgen::global_constraints::{GlobalConstraints, RangeConstraintSet};
+use crate::witgen::machines::profiling::{record_end_identity, record_start_identity};
 use crate::witgen::machines::record_start;
 use crate::witgen::range_constraints::RangeConstraint;
 use crate::witgen::rows::RowPair;
 use crate::witgen::util::try_to_simple_poly_ref;
-use crate::witgen::{EvalError, EvalValue, IncompleteCause};
+use crate::witgen::{profiling::IDENTITY_LOOKUP_CACHE, EvalError, EvalValue, IncompleteCause};
 use crate::witgen::{EvalResult, FixedData};
 
 use super::record_end;
@@ -83,6 +84,8 @@ impl<T: FieldElement> IndexedColumns<T> {
             return;
         }
 
+        record_start_identity(IDENTITY_LOOKUP_CACHE);
+
         let (sorted_input_fixed_columns, sorted_output_fixed_columns) = &sorted_fixed_columns;
 
         // create index for this lookup
@@ -160,6 +163,7 @@ impl<T: FieldElement> IndexedColumns<T> {
             ),
             index,
         );
+        record_end_identity(IDENTITY_LOOKUP_CACHE)
     }
 }
 

executor/src/witgen/machines/profiling.rs

@@ -5,16 +5,44 @@ use std::{
     time::{Duration, Instant},
 };
 
-#[derive(PartialEq, Debug, Copy, Clone)]
-enum Event {
-    Start,
-    End,
+use log::log_enabled;
+use powdr_number::FieldElement;
+
+use crate::witgen::FixedData;
+
+pub static UNUSED_IDENTITY_ID: u64 = u64::MAX;
+pub static IDENTITY_LOOKUP_CACHE: u64 = UNUSED_IDENTITY_ID - 1;
+pub static PROCESS_PROVER_QUERIES_ID: u64 = UNUSED_IDENTITY_ID - 2;
+
+#[derive(Debug)]
+struct Event {
+    is_start: bool,
+    id: u64,
+    time: Instant,
+}
+
+impl Event {
+    fn start(id: u64) -> Self {
+        Self {
+            is_start: true,
+            id,
+            time: Instant::now(),
+        }
+    }
+
+    fn end(id: u64) -> Self {
+        Self {
+            is_start: false,
+            id,
+            time: Instant::now(),
+        }
+    }
 }
 
 thread_local! {
-    /// The event log is a list of (event, <ID>, time) tuples.
-    static EVENT_LOG: RefCell<Vec<(Event, u64, Instant)>> = const { RefCell::new(Vec::new()) };
-    static EVENT_LOG_BY_IDENTITY: RefCell<Vec<(Event, u64, Instant)>> = const { RefCell::new(Vec::new()) };
+    /// The event log is a list of events.
+    static EVENT_LOG: RefCell<Vec<Event>> = const { RefCell::new(Vec::new()) };
+    static EVENT_LOG_BY_IDENTITY: RefCell<Vec<Event>> = const { RefCell::new(Vec::new()) };
     /// Maps a machine name (assumed to be globally unique) to an ID.
     /// This is done so that we can use a usize in the event log.
     static NAME_TO_ID: RefCell<BTreeMap<String, u64>> = const { RefCell::new(BTreeMap::new()) };
@@ -35,80 +63,107 @@ fn id_from_name(name: &str) -> u64 {
 /// Adds the start of a computation to the event log.
 pub fn record_start(name: &str) {
     let id = id_from_name(name);
-    EVENT_LOG.with(|s| s.borrow_mut().push((Event::Start, id, Instant::now())));
+    EVENT_LOG.with(|s| s.borrow_mut().push(Event::start(id)));
 }
 
 pub fn record_start_identity(identity: u64) {
-    EVENT_LOG_BY_IDENTITY.with(|s| {
-        s.borrow_mut()
-            .push((Event::Start, identity, Instant::now()))
-    });
+    if log_enabled!(log::Level::Debug) {
+        EVENT_LOG_BY_IDENTITY.with(|s| s.borrow_mut().push(Event::start(identity)));
+    }
 }
 
 /// Adds the end of a computation to the event log.
 pub fn record_end(name: &str) {
     let id = id_from_name(name);
-    EVENT_LOG.with(|s| s.borrow_mut().push((Event::End, id, Instant::now())));
+    EVENT_LOG.with(|s| s.borrow_mut().push(Event::end(id)));
 }
 
 pub fn record_end_identity(identity: u64) {
-    EVENT_LOG_BY_IDENTITY.with(|s| s.borrow_mut().push((Event::End, identity, Instant::now())));
+    if log_enabled!(log::Level::Debug) {
+        EVENT_LOG_BY_IDENTITY.with(|s| s.borrow_mut().push(Event::end(identity)));
+    }
 }
+
 pub fn reset_and_print_profile_summary() {
+    if !log_enabled!(log::Level::Debug) {
+        return;
+    }
     let id_to_name = NAME_TO_ID.with(|name_to_id| {
         let name_to_id = name_to_id.borrow();
         name_to_id
             .iter()
             .map(|(name, id)| (*id, name.clone()))
             .collect::<BTreeMap<_, _>>()
     });
-    reset_and_print_profile_summary_impl(&EVENT_LOG, id_to_name)
+    reset_and_print_profile_summary_impl("Per machine", &EVENT_LOG, id_to_name)
 }
-pub fn reset_and_print_profile_summary_identity(id_to_name: BTreeMap<u64, String>) {
-    reset_and_print_profile_summary_impl(&EVENT_LOG_BY_IDENTITY, id_to_name)
+
+pub fn reset_and_print_profile_summary_identity<T: FieldElement>(fixed_data: &FixedData<T>) {
+    if !log_enabled!(log::Level::Debug) {
+        return;
+    }
+    let id_to_name = fixed_data
+        .analyzed
+        .identities
+        .iter()
+        .map(|identity| (identity.id, format!("{identity}")))
+        .chain([(UNUSED_IDENTITY_ID, "other".to_string())])
+        .chain([(IDENTITY_LOOKUP_CACHE, "lookup cache creation".to_string())])
+        .chain([(
+            PROCESS_PROVER_QUERIES_ID,
+            "process prover queries".to_string(),
+        )])
+        .collect();
+
+    reset_and_print_profile_summary_impl(
+        "Identities processing",
+        &EVENT_LOG_BY_IDENTITY,
+        id_to_name,
+    )
 }
 
 fn reset_and_print_profile_summary_impl(
-    event_log: &'static LocalKey<RefCell<Vec<(Event, u64, Instant)>>>,
+    name: &str,
+    event_log: &'static LocalKey<RefCell<Vec<Event>>>,
     id_to_name: BTreeMap<u64, String>,
 ) {
     event_log.with(|event_log| {
         // Taking the events out is actually important, because there might be
         // multiple (consecutive) runs of witgen in the same thread.
         let event_log = std::mem::take(&mut (*event_log.borrow_mut()));
-        log::debug!("\n == Witgen profile ({} events)", event_log.len());
+        log::debug!("\n == {name} profile ({} events)", event_log.len());
 
         // Aggregate time spent in each machine.
         let mut time_by_machine = BTreeMap::new();
-        assert_eq!(event_log[0].0, Event::Start);
-        let mut current_time = event_log[0].2;
-        let mut call_stack = vec![event_log[0].1];
+        let mut machine_invocations = BTreeMap::new();
+        assert!(event_log[0].is_start);
+        let mut current_time = event_log[0].time;
+        let mut call_stack = vec![event_log[0].id];
 
-        for (i, &(event, id, time)) in event_log.iter().enumerate().skip(1) {
+        for (i, event) in event_log.iter().enumerate().skip(1) {
             // We expect one top-level call, so we should never have an empty call stack.
-            let current_machine_id = *call_stack.last().unwrap_or_else(|| {
-                panic!(
-                    "Call stack is empty at index {i} (event: {event:?}, name: {id}, time: {time:?})"
-                )
-            });
+            let current_machine_id = *call_stack
+                .last()
+                .unwrap_or_else(|| panic!("Call stack is empty at index {i}, event: {event:?}"));
 
             // Finish the execution of the currently running machine.
-            let duration = time.duration_since(current_time);
+            let duration = event.time.duration_since(current_time);
             *time_by_machine
                 .entry(current_machine_id)
                 .or_insert(Duration::default()) += duration;
-            current_time = time;
+            current_time = event.time;
 
             // Update the call stack.
-            match event {
-                Event::Start => {
-                    assert!(current_machine_id != id, "Unexpected recursive call!");
-                    call_stack.push(id);
-                }
-                Event::End => {
-                    assert_eq!(current_machine_id, id, "Unexpected end of call!");
-                    call_stack.pop().unwrap();
-                }
+            if event.is_start {
+                assert!(current_machine_id != event.id, "Unexpected recursive call!");
+                call_stack.push(event.id);
+            } else {
+                assert_eq!(current_machine_id, event.id, "Unexpected end of call!");
+                machine_invocations
+                    .entry(event.id)
+                    .and_modify(|c| *c += 1)
+                    .or_insert(1);
+                call_stack.pop().unwrap();
             }
         }
 
@@ -123,20 +178,28 @@ fn reset_and_print_profile_summary_impl(
 
         let total_time = time_by_machine.iter().map(|(_, d)| *d).sum::<Duration>();
         assert_eq!(
-            event_log.last().unwrap().2.duration_since(event_log[0].2),
+            event_log
+                .last()
+                .unwrap()
+                .time
+                .duration_since(event_log[0].time),
             total_time
         );
 
         for (i, (id, duration)) in time_by_machine.iter().enumerate() {
-            if i > 10 {
+            if i > 100 {
+                log::debug!("... skipping {} more items", time_by_machine.len() - i);
                 break;
             }
             let percentage = (duration.as_secs_f64() / total_time.as_secs_f64()) * 100.0;
             log::debug!(
-                "  {:>5.1}% ({:>8.1?}): {}",
-                percentage,
-                duration,
+                "  {percentage:>5.5}% ({duration:>8.5?}), {} runs: {}",
+                machine_invocations.get(id).cloned().unwrap_or_default(),
                 id_to_name[&id]
+                    .to_string()
+                    .chars()
+                    .take(200)
+                    .collect::<String>()
             );
         }
         log::debug!("  ---------------------------");

executor/src/witgen/mod.rs

@@ -3,7 +3,7 @@ use std::rc::Rc;
 use std::sync::Arc;
 
 use machines::profiling::{
-    record_end_identity, record_start_identity, reset_and_print_profile_summary_identity,
+    self, record_end_identity, record_start_identity, reset_and_print_profile_summary_identity,
 };
 use powdr_ast::analyzed::{
     AlgebraicExpression, AlgebraicReference, Analyzed, Expression, FunctionValueDefinition, PolyID,
@@ -46,7 +46,6 @@ mod util;
 mod vm_processor;
 
 static OUTER_CODE_NAME: &str = "witgen (outer code)";
-static UNUSED_IDENTITY_ID: u64 = 1234567;
 
 pub trait QueryCallback<T>: Fn(&str) -> Result<Option<T>, String> + Send + Sync {}
 impl<T, F> QueryCallback<T> for F where F: Fn(&str) -> Result<Option<T>, String> + Send + Sync {}
@@ -156,7 +155,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
     /// @returns the values (in source order) and the degree of the polynomials.
     pub fn generate(self) -> Vec<(String, Vec<T>)> {
         record_start(OUTER_CODE_NAME);
-        record_start_identity(UNUSED_IDENTITY_ID);
+        record_start_identity(profiling::UNUSED_IDENTITY_ID);
         let fixed = FixedData::new(
             self.analyzed,
             self.fixed_col_values,
@@ -169,7 +168,6 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
             .identities_with_inlined_intermediate_polynomials()
             .into_iter()
             .filter(|identity| {
-                assert_ne!(identity.id, UNUSED_IDENTITY_ID);
                 let discard = identity.expr_any(|expr| {
                     if let AlgebraicExpression::Challenge(challenge) = expr {
                         challenge.stage >= self.stage.into()
@@ -231,14 +229,9 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
             .collect::<BTreeMap<_, _>>();
 
         record_end(OUTER_CODE_NAME);
-        record_end_identity(UNUSED_IDENTITY_ID);
+        record_end_identity(profiling::UNUSED_IDENTITY_ID);
         reset_and_print_profile_summary();
-        let name_to_id = identities
-            .iter()
-            .map(|identity| (identity.id, format!("{identity}")))
-            .chain([(UNUSED_IDENTITY_ID, "other".to_string())])
-            .collect();
-        reset_and_print_profile_summary_identity(name_to_id);
+        reset_and_print_profile_summary_identity(&fixed);
 
         // Order columns according to the order of declaration.
         let witness_cols = self

executor/src/witgen/processor.rs

@@ -501,6 +501,7 @@ Known values in current row (local: {row_index}, global {global_row_index}):
         // This could be computed from the identity, but should be pre-computed for performance reasons.
         has_next_reference: bool,
     ) -> bool {
+        record_start_identity(identity.id);
         let mut identity_processor = IdentityProcessor::new(self.fixed_data, self.mutable_state);
         let row_pair = match has_next_reference {
             // Check whether identities with a reference to the next row are satisfied
@@ -526,16 +527,15 @@ Known values in current row (local: {row_index}, global {global_row_index}):
             ),
         };
 
-        if identity_processor
+        let is_err = identity_processor
             .process_identity(identity, &row_pair)
-            .is_err()
-        {
+            .is_err();
+        if is_err {
             log::debug!("Previous {:?}", &self.data[row_index - 1]);
             log::debug!("Proposed {:?}", proposed_row);
             log::debug!("Failed on identity: {}", identity);
-
-            return false;
         }
-        true
+        record_end_identity(identity.id);
+        !is_err
     }
 }

executor/src/witgen/vm_processor.rs

@@ -10,9 +10,11 @@ use std::collections::HashSet;
 use std::time::Instant;
 
 use crate::witgen::identity_processor::{self};
+use crate::witgen::machines::profiling;
 use crate::witgen::IncompleteCause;
 
 use super::data_structures::finalizable_data::FinalizableData;
+use super::machines::profiling::{record_end_identity, record_start_identity};
 use super::processor::{OuterQuery, Processor};
 
 use super::rows::{Row, RowIndex, UnknownStrategy};
@@ -155,6 +157,9 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> VmProcessor<'a, 'b, 'c, T
                         loop_detection_log_level,
                         "Found loop with period {p} starting at row {row_index}"
                     );
+                    record_end_identity(profiling::UNUSED_IDENTITY_ID);
+                    profiling::reset_and_print_profile_summary_identity(self.fixed_data);
+                    record_start_identity(profiling::UNUSED_IDENTITY_ID);
                 }
             }
             if let Some(period) = looping_period {
@@ -348,10 +353,13 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> VmProcessor<'a, 'b, 'c, T
             }
 
             progress |= self.processor.set_inputs_if_unset(row_index);
-            progress |= self
-                .processor
-                .process_queries(row_index)
-                .map_err(|e| vec![e])?;
+
+            record_start_identity(profiling::PROCESS_PROVER_QUERIES_ID);
+            progress |= self.processor.process_queries(row_index).map_err(|e| {
+                record_end_identity(profiling::PROCESS_PROVER_QUERIES_ID);
+                vec![e]
+            })?;
+            record_end_identity(profiling::PROCESS_PROVER_QUERIES_ID);
 
             if !progress {
                 break;