github/powdr
1
1
Fork 0
mirror of https://github.com/powdr-labs/powdr.git synced 2026-04-20 03:03:25 -04:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #871 from powdr-labs/machine-stats

Browse Source 
Witness generation profiling
This commit is contained in:
Leo
2024-01-05 19:45:53 +00:00
committed by GitHub
parent de048f0b98 f1c674ee73
commit d897805afd
11 changed files with 263 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
executor/src/witgen/generator.rs
View File

@@ -6,6 +6,7 @@ use number::{DegreeType, FieldElement};
use std::collections::{HashMap, HashSet};
use crate::witgen::data_structures::finalizable_data::FinalizableData;
use crate::witgen::machines::profiling::{record_end, record_start};
use crate::witgen::processor::OuterQuery;
use crate::witgen::rows::CellValue;
use crate::witgen::EvalValue;
@@ -32,9 +33,14 @@ pub struct Generator<'a, T: FieldElement> {
global_range_constraints: GlobalConstraints<T>,
data: FinalizableData<'a, T>,
latch: Option<Expression<T>>,
name: String,
}
impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
fn name(&self) -> &str {
&self.name
}
fn process_plookup<Q: QueryCallback<T>>(
&mut self,
mutable_state: &mut MutableState<'a, '_, T, Q>,
@@ -103,6 +109,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
impl<'a, T: FieldElement> Generator<'a, T> {
pub fn new(
name: String,
fixed_data: &'a FixedData<'a, T>,
identities: &[&'a Identity<Expression<T>>],
witnesses: HashSet<PolyID>,
@@ -111,6 +118,7 @@ impl<'a, T: FieldElement> Generator<'a, T> {
) -> Self {
let data = FinalizableData::new(&witnesses);
Self {
name,
fixed_data,
identities: identities.to_vec(),
witnesses,
@@ -120,23 +128,13 @@ impl<'a, T: FieldElement> Generator<'a, T> {
}
}
pub fn name(&self) -> &str {
let first_witness = self.witnesses.iter().next().unwrap();
let first_witness_name = self.fixed_data.column_name(first_witness);
let namespace = first_witness_name
.rfind('.')
.map(|idx| &first_witness_name[..idx]);
// For machines compiled using Powdr ASM we'll always have a namespace, but as a last
// resort we'll use the first witness name.
namespace.unwrap_or(first_witness_name)
}
/// Runs the machine without any arguments from the first row.
pub fn run<'b, Q: QueryCallback<T>>(&mut self, mutable_state: &mut MutableState<'a, 'b, T, Q>) {
record_start(self.name());
assert!(self.data.is_empty());
let first_row = self.compute_partial_first_row(mutable_state);
self.data = self.process(first_row, 0, mutable_state, None).block;
record_end(self.name());
}
fn fill_remaining_rows<Q: QueryCallback<T>>(

11
executor/src/witgen/identity_processor.rs
View File

@@ -156,7 +156,7 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'b,
// TODO could it be that multiple machines match?
// query the fixed lookup "machine"
if let Some(result) = self.mutable_state.fixed_lookup.process_plookup(
if let Some(result) = self.mutable_state.fixed_lookup.process_plookup_timed(
self.fixed_data,
rows,
identity.kind,
@@ -174,9 +174,12 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'b,
query_callback: self.mutable_state.query_callback,
};
if let Some(result) =
current.process_plookup(&mut mutable_state, identity.kind, &left, &identity.right)
{
if let Some(result) = current.process_plookup_timed(
&mut mutable_state,
identity.kind,
&left,
&identity.right,
) {
return result;
}
}

19
executor/src/witgen/machines/block_machine.rs
View File

@@ -62,10 +62,12 @@ pub struct BlockMachine<'a, T: FieldElement> {
/// to make progress most quickly.
processing_sequence_cache: ProcessingSequenceCache,
fixed_data: &'a FixedData<'a, T>,
name: String,
}
impl<'a, T: FieldElement> BlockMachine<'a, T> {
pub fn try_new(
name: String,
fixed_data: &'a FixedData<'a, T>,
connecting_identities: &[&'a Identity<Expression<T>>],
identities: &[&'a Identity<Expression<T>>],
@@ -120,6 +122,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
(0..block_size).map(|i| row_factory.fresh_row(i as DegreeType)),
);
BlockMachine {
name,
block_size,
connecting_rhs,
identities: identities.to_vec(),
@@ -202,6 +205,10 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
})
}
fn name(&self) -> &str {
&self.name
}
fn take_witness_col_values<'b, Q: QueryCallback<T>>(
&mut self,
_fixed_lookup: &'b mut FixedLookup<T>,
@@ -289,18 +296,6 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
self.data.len() as DegreeType
}
fn name(&self) -> &str {
let first_witness = self.witness_cols.iter().next().unwrap();
let first_witness_name = self.fixed_data.column_name(first_witness);
let namespace = first_witness_name
.rfind('.')
.map(|idx| &first_witness_name[..idx]);
// For machines compiled using Powdr ASM we'll always have a namespace, but as a last
// resort we'll use the first witness name.
namespace.unwrap_or(first_witness_name)
}
fn process_plookup_internal<'b, Q: QueryCallback<T>>(
&mut self,
mutable_state: &mut MutableState<'a, 'b, T, Q>,

8
executor/src/witgen/machines/double_sorted_witness_machine.rs
View File

@@ -29,6 +29,7 @@ pub struct DoubleSortedWitnesses<T> {
trace: BTreeMap<(T, T), Operation<T>>,
data: BTreeMap<T, T>,
namespace: String,
name: String,
}
struct Operation<T> {
@@ -42,6 +43,7 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
}
pub fn try_new(
name: String,
fixed_data: &FixedData<T>,
_identities: &[&Identity<Expression<T>>],
witness_cols: &HashSet<PolyID>,
@@ -82,7 +84,7 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
.is_none()
{
Some(Self {
// store the namespace
name,
namespace,
degree: fixed_data.degree,
..Default::default()
@@ -94,6 +96,10 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
}
impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
fn name(&self) -> &str {
&self.name
}
fn process_plookup<Q: QueryCallback<T>>(
&mut self,
_mutable_state: &mut MutableState<'a, '_, T, Q>,

17
executor/src/witgen/machines/fixed_lookup_machine.rs
View File

@@ -11,12 +11,15 @@ use number::FieldElement;
use crate::witgen::affine_expression::AffineExpression;
use crate::witgen::global_constraints::{GlobalConstraints, RangeConstraintSet};
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::{EvalResult, FixedData};
use super::record_end;
type Application = (Vec<PolyID>, Vec<PolyID>);
type Index<T> = BTreeMap<Vec<T>, IndexValue>;
@@ -173,6 +176,20 @@ impl<T: FieldElement> FixedLookup<T> {
}
}
pub fn process_plookup_timed<'b>(
&mut self,
fixed_data: &FixedData<T>,
rows: &RowPair<'_, '_, T>,
kind: IdentityKind,
left: &[AffineExpression<&'b AlgebraicReference, T>],
right: &'b SelectedExpressions<Expression<T>>,
) -> Option<EvalResult<'b, T>> {
record_start("FixedLookup");
let result = self.process_plookup(fixed_data, rows, kind, left, right);
record_end("FixedLookup");
result
}
pub fn process_plookup<'b>(
&mut self,
fixed_data: &FixedData<T>,

43
executor/src/witgen/machines/machine_extractor.rs
View File

@@ -37,6 +37,7 @@ pub fn split_out_machines<'a, T: FieldElement>(
let all_witnesses = fixed.witness_cols.keys().collect::<HashSet<_>>();
let mut remaining_witnesses = all_witnesses.clone();
let mut base_identities = identities.clone();
let mut id_counter = 0;
for id in &identities {
// Extract all witness columns in the RHS of the lookup.
let lookup_witnesses = &refs_in_selected_expressions(&id.right) & (&remaining_witnesses);
@@ -98,22 +99,45 @@ pub fn split_out_machines<'a, T: FieldElement>(
.join("\n"),
);
if let Some(machine) =
SortedWitnesses::try_new(fixed, &machine_identities, &machine_witnesses)
{
let first_witness = machine_witnesses.iter().next().unwrap();
let first_witness_name = fixed.column_name(first_witness);
let namespace = first_witness_name
.rfind('.')
.map(|idx| &first_witness_name[..idx]);
// For machines compiled using Powdr ASM we'll always have a namespace, but as a last
// resort we'll use the first witness name.
let name = namespace.unwrap_or(first_witness_name);
let id = id_counter;
id_counter += 1;
let name_with_type = |t: &str| format!("Secondary machine {id}: {name} ({t})");
if let Some(machine) = SortedWitnesses::try_new(
name_with_type("SortedWitness"),
fixed,
&machine_identities,
&machine_witnesses,
) {
log::info!("Detected machine: sorted witnesses / write-once memory");
machines.push(KnownMachine::SortedWitnesses(machine));
} else if let Some(machine) =
DoubleSortedWitnesses::try_new(fixed, &machine_identities, &machine_witnesses)
{
} else if let Some(machine) = DoubleSortedWitnesses::try_new(
name_with_type("DoubleSortedWitnesses"),
fixed,
&machine_identities,
&machine_witnesses,
) {
log::info!("Detected machine: memory");
machines.push(KnownMachine::DoubleSortedWitnesses(machine));
} else if let Some(machine) =
WriteOnceMemory::try_new(fixed, &connecting_identities, &machine_identities)
{
} else if let Some(machine) = WriteOnceMemory::try_new(
name_with_type("WriteOnceMemory"),
fixed,
&connecting_identities,
&machine_identities,
) {
log::info!("Detected machine: write-once memory");
machines.push(KnownMachine::WriteOnceMemory(machine));
} else if let Some(machine) = BlockMachine::try_new(
name_with_type("BlockMachine"),
fixed,
&connecting_identities,
&machine_identities,
@@ -144,6 +168,7 @@ pub fn split_out_machines<'a, T: FieldElement>(
})
.unwrap();
machines.push(KnownMachine::Vm(Generator::new(
name_with_type("Vm"),
fixed,
&machine_identities,
machine_witnesses,

30
executor/src/witgen/machines/mod.rs
View File

@@ -8,6 +8,8 @@ use number::FieldElement;
use self::block_machine::BlockMachine;
use self::double_sorted_witness_machine::DoubleSortedWitnesses;
pub use self::fixed_lookup_machine::FixedLookup;
use self::profiling::record_end;
use self::profiling::record_start;
use self::sorted_witness_machine::SortedWitnesses;
use self::write_once_memory::WriteOnceMemory;
use ast::analyzed::IdentityKind;
@@ -23,12 +25,30 @@ mod block_machine;
mod double_sorted_witness_machine;
mod fixed_lookup_machine;
pub mod machine_extractor;
pub mod profiling;
mod sorted_witness_machine;
mod write_once_memory;
/// A machine is a set of witness columns and identities where the columns
/// are used on the right-hand-side of lookups. It can process plookups.
pub trait Machine<'a, T: FieldElement>: Send + Sync {
/// Like `process_plookup`, but also records the time spent in this machine.
fn process_plookup_timed<'b, Q: QueryCallback<T>>(
&mut self,
mutable_state: &'b mut MutableState<'a, 'b, T, Q>,
kind: IdentityKind,
left: &[AffineExpression<&'a AlgebraicReference, T>],
right: &'a SelectedExpressions<Expression<T>>,
) -> Option<EvalResult<'a, T>> {
record_start(self.name());
let result = self.process_plookup(mutable_state, kind, left, right);
record_end(self.name());
result
}
/// Returns a unique name for this machine.
fn name(&self) -> &str;
/// Process a plookup. Not all values on the LHS need to be available.
/// Can update internal data.
/// Only return an error if this machine is able to handle the query and
@@ -80,6 +100,16 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
}
}
fn name(&self) -> &str {
match self {
KnownMachine::SortedWitnesses(m) => m.name(),
KnownMachine::DoubleSortedWitnesses(m) => m.name(),
KnownMachine::WriteOnceMemory(m) => m.name(),
KnownMachine::BlockMachine(m) => m.name(),
KnownMachine::Vm(m) => m.name(),
}
}
fn take_witness_col_values<'b, Q: QueryCallback<T>>(
&mut self,
fixed_lookup: &'b mut FixedLookup<T>,

129
executor/src/witgen/machines/profiling.rs Normal file
View File

@@ -0,0 +1,129 @@
use std::{
cell::RefCell,
collections::BTreeMap,
time::{Duration, SystemTime},
};
#[derive(PartialEq, Debug, Copy, Clone)]
enum Event {
Start,
End,
}
thread_local! {
/// The event log is a list of (event, <ID>, time) tuples.
static EVENT_LOG: RefCell<Vec<(Event, usize, SystemTime)>> = 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>> = 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 {
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();
name_to_id.insert(name.to_string(), id);
id
})
})
}
/// 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, SystemTime::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, SystemTime::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<_, _>>()
});
// 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::info!("\n == Witgen 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];
for (i, &(event, id, time)) 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 {} (event: {:?}, name: {}, time: {:?})",
i, event, id, time
)
});
// Finish the execution of the currently running machine.
let duration = time.duration_since(current_time).unwrap();
*time_by_machine
.entry(current_machine_id)
.or_insert(Duration::default()) += duration;
current_time = 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();
}
}
}
assert!(
call_stack.is_empty(),
"Call stack is not empty: {:?}",
call_stack
);
// Sort by time, descending.
let mut time_by_machine = time_by_machine.into_iter().collect::<Vec<_>>();
time_by_machine.sort_by(|a, b| b.1.cmp(&a.1));
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)
.unwrap(),
total_time
);
for (id, duration) in time_by_machine {
let percentage = (duration.as_secs_f64() / total_time.as_secs_f64()) * 100.0;
log::info!(
" {:>5.1}% ({:>8.1?}): {}",
percentage,
duration,
id_to_name[&id]
);
}
log::info!(" ---------------------------");
log::info!(" ==> Total: {:?}", total_time);
log::info!("\n");
});
}

7
executor/src/witgen/machines/sorted_witness_machine.rs
View File

@@ -29,10 +29,12 @@ pub struct SortedWitnesses<'a, T> {
witness_positions: HashMap<PolyID, usize>,
data: BTreeMap<T, Vec<Option<T>>>,
fixed_data: &'a FixedData<'a, T>,
name: String,
}
impl<'a, T: FieldElement> SortedWitnesses<'a, T> {
pub fn try_new(
name: String,
fixed_data: &'a FixedData<T>,
identities: &[&Identity<Expression<T>>],
witnesses: &HashSet<PolyID>,
@@ -50,6 +52,7 @@ impl<'a, T: FieldElement> SortedWitnesses<'a, T> {
.collect();
SortedWitnesses {
name,
key_col,
witness_positions,
data: Default::default(),
@@ -123,6 +126,10 @@ fn check_constraint<T: FieldElement>(constraint: &Expression<T>) -> Option<PolyI
}
impl<'a, T: FieldElement> Machine<'a, T> for SortedWitnesses<'a, T> {
fn name(&self) -> &str {
&self.name
}
fn process_plookup<Q: QueryCallback<T>>(
&mut self,
_mutable_state: &mut MutableState<'a, '_, T, Q>,

7
executor/src/witgen/machines/write_once_memory.rs
View File

@@ -41,10 +41,12 @@ pub struct WriteOnceMemory<'a, T: FieldElement> {
key_to_index: BTreeMap<Vec<T>, DegreeType>,
/// The memory content
data: BTreeMap<DegreeType, Vec<Option<T>>>,
name: String,
}
impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
pub fn try_new(
name: String,
fixed_data: &'a FixedData<'a, T>,
connecting_identities: &[&'a Identity<Expression<T>>],
identities: &[&Identity<Expression<T>>],
@@ -103,6 +105,7 @@ impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
}
Some(Self {
name,
fixed_data,
rhs,
value_polys,
@@ -197,6 +200,10 @@ impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
}
impl<'a, T: FieldElement> Machine<'a, T> for WriteOnceMemory<'a, T> {
fn name(&self) -> &str {
&self.name
}
fn process_plookup<'b, Q: QueryCallback<T>>(
&mut self,
_mutable_state: &'b mut MutableState<'a, 'b, T, Q>,

8
executor/src/witgen/mod.rs
View File

@@ -14,6 +14,7 @@ use self::generator::Generator;
use self::identity_processor::Machines;
use self::machines::machine_extractor::ExtractionOutput;
use self::machines::profiling::{record_end, record_start, reset_and_print_profile_summary};
use self::machines::{FixedLookup, Machine};
mod affine_expression;
@@ -36,6 +37,8 @@ mod symbolic_witness_evaluator;
mod util;
mod vm_processor;
static OUTER_CODE_NAME: &str = "witgen (outer code)";
pub trait QueryCallback<T>: FnMut(&str) -> Result<Option<T>, String> + Send + Sync {}
impl<T, F> QueryCallback<T> for F where F: FnMut(&str) -> Result<Option<T>, String> + Send + Sync {}
@@ -92,6 +95,7 @@ impl<'a, 'b, T: FieldElement, Q: QueryCallback<T>> WitnessGenerator<'a, 'b, T, Q
/// Generates the committed polynomial values
/// @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);
let fixed = FixedData::new(
self.analyzed,
self.fixed_col_values,
@@ -124,6 +128,7 @@ impl<'a, 'b, T: FieldElement, Q: QueryCallback<T>> WitnessGenerator<'a, 'b, T, Q
query_callback: &mut query_callback,
};
let mut generator = Generator::new(
"Main Machine".to_string(),
&fixed,
&base_identities,
base_witnesses,
@@ -149,6 +154,9 @@ impl<'a, 'b, T: FieldElement, Q: QueryCallback<T>> WitnessGenerator<'a, 'b, T, Q
.chain(main_columns)
.collect::<BTreeMap<_, _>>();
record_end(OUTER_CODE_NAME);
reset_and_print_profile_summary();
// Order columns according to the order of declaration.
self.analyzed
.committed_polys_in_source_order()