Profile identities

executor/src/witgen/machines/profiling.rs

@@ -1,6 +1,7 @@
 use std::{
     cell::RefCell,
     collections::BTreeMap,
+    thread::LocalKey,
     time::{Duration, Instant},
 };
 
@@ -12,18 +13,19 @@ enum Event {
 
 thread_local! {
     /// The event log is a list of (event, <ID>, time) tuples.
-    static EVENT_LOG: RefCell<Vec<(Event, usize, Instant)>> = const { RefCell::new(Vec::new()) };
+    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()) };
     /// 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, usize>> = const { RefCell::new(BTreeMap::new()) };
+    static NAME_TO_ID: RefCell<BTreeMap<String, u64>> = const { RefCell::new(BTreeMap::new()) };
 }
 
 /// Returns the ID for a given machine name, creating a new one if necessary.
-fn id_from_name(name: &str) -> usize {
+fn id_from_name(name: &str) -> u64 {
     NAME_TO_ID.with(|name_to_id| {
         let mut name_to_id = name_to_id.borrow_mut();
         name_to_id.get(name).copied().unwrap_or_else(|| {
-            let id = name_to_id.len();
+            let id = name_to_id.len() as u64;
             name_to_id.insert(name.to_string(), id);
             id
         })
@@ -36,22 +38,41 @@ pub fn record_start(name: &str) {
     EVENT_LOG.with(|s| s.borrow_mut().push((Event::Start, id, Instant::now())));
 }
 
+pub fn record_start_identity(identity: u64) {
+    EVENT_LOG_BY_IDENTITY.with(|s| {
+        s.borrow_mut()
+            .push((Event::Start, identity, Instant::now()))
+    });
+}
+
 /// 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())));
 }
 
+pub fn record_end_identity(identity: u64) {
+    EVENT_LOG_BY_IDENTITY.with(|s| s.borrow_mut().push((Event::End, identity, Instant::now())));
+}
 pub fn reset_and_print_profile_summary() {
-    EVENT_LOG.with(|event_log| {
-        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<_, _>>()
-        });
+    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)
+}
+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)
+}
 
+fn reset_and_print_profile_summary_impl(
+    event_log: &'static LocalKey<RefCell<Vec<(Event, u64, Instant)>>>,
+    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()));
@@ -106,7 +127,10 @@ pub fn reset_and_print_profile_summary() {
             total_time
         );
 
-        for (id, duration) in time_by_machine {
+        for (i, (id, duration)) in time_by_machine.iter().enumerate() {
+            if i > 10 {
+                break;
+            }
             let percentage = (duration.as_secs_f64() / total_time.as_secs_f64()) * 100.0;
             log::debug!(
                 "  {:>5.1}% ({:>8.1?}): {}",

executor/src/witgen/mod.rs

@@ -2,6 +2,9 @@ use std::collections::{BTreeMap, HashMap};
 use std::rc::Rc;
 use std::sync::Arc;
 
+use machines::profiling::{
+    record_end_identity, record_start_identity, reset_and_print_profile_summary_identity,
+};
 use powdr_ast::analyzed::{
     AlgebraicExpression, AlgebraicReference, Analyzed, Expression, FunctionValueDefinition, PolyID,
     PolynomialType, SymbolKind, TypedExpression,
@@ -43,6 +46,7 @@ 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 {}
@@ -152,6 +156,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);
         let fixed = FixedData::new(
             self.analyzed,
             self.fixed_col_values,
@@ -164,6 +169,7 @@ 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()
@@ -225,7 +231,14 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
             .collect::<BTreeMap<_, _>>();
 
         record_end(OUTER_CODE_NAME);
+        record_end_identity(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);
 
         // Order columns according to the order of declaration.
         let witness_cols = self

executor/src/witgen/processor.rs

@@ -8,6 +8,7 @@ use powdr_number::{DegreeType, FieldElement};
 
 use crate::witgen::{query_processor::QueryProcessor, util::try_to_simple_poly, Constraint};
 
+use super::machines::profiling::{record_end_identity, record_start_identity};
 use super::{
     affine_expression::AffineExpression,
     data_structures::{
@@ -207,6 +208,7 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> Processor<'a, 'b, 'c, T,
         identity: &'a Identity<Expression<T>>,
         unknown_strategy: UnknownStrategy,
     ) -> Result<IdentityResult, EvalError<T>> {
+        record_start_identity(identity.id);
         // Create row pair
         let global_row_index = self.row_offset + row_index as u64;
         let row_pair = RowPair::new(
@@ -239,7 +241,14 @@ Known values in current row (local: {row_index}, global {global_row_index}):
                 }
                 error += &format!("   => Error: {e}");
                 error.into()
-            })?;
+            });
+        let updates = match updates {
+            Ok(updates) => updates,
+            Err(e) => {
+                record_end_identity(identity.id);
+                return Err(e);
+            }
+        };
 
         if unknown_strategy == UnknownStrategy::Zero {
             assert!(updates.constraints.is_empty());
@@ -250,11 +259,13 @@ Known values in current row (local: {row_index}, global {global_row_index}):
             });
         }
 
-        Ok(IdentityResult {
+        let res = Ok(IdentityResult {
             progress: self.apply_updates(row_index, &updates, || identity.to_string())
                 || updates.side_effect,
             is_complete: updates.is_complete(),
-        })
+        });
+        record_end_identity(identity.id);
+        res
     }
 
     pub fn process_outer_query(