From a8d128b117d70ddc9b843d00a98da224402f8904 Mon Sep 17 00:00:00 2001 From: Shahar Papini <43779613+spapinistarkware@users.noreply.github.com> Date: Mon, 15 Apr 2024 14:31:35 +0300 Subject: [PATCH] Add measurement logs (#493) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This change is [Reviewable](https://reviewable.io/reviews/starkware-libs/stwo/493) --- Cargo.lock | 194 +++++++++++++++++++++++- Cargo.toml | 3 + avx_benchmark.sh | 3 + src/core/air/accumulation.rs | 2 + src/core/commitment_scheme/prover.rs | 5 + src/core/commitment_scheme/quotients.rs | 2 + src/core/fri.rs | 2 + src/core/proof_of_work.rs | 2 + src/core/prover/mod.rs | 12 ++ src/examples/wide_fibonacci/avx.rs | 22 ++- src/examples/wide_fibonacci/mod.rs | 6 +- 11 files changed, 246 insertions(+), 7 deletions(-) create mode 100755 avx_benchmark.sh diff --git a/Cargo.lock b/Cargo.lock index 73efdcd71..6b7eb4239 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -268,6 +268,25 @@ version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" +[[package]] +name = "env_filter" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a009aa4810eb158359dda09d0c87378e4bbb89b5a801f016885a4707ba24f7ea" +dependencies = [ + "log", +] + +[[package]] +name = "env_logger" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b35839ba51819680ba087cd351788c9a3c476841207e0b8cee0b04722343b9" +dependencies = [ + "env_filter", + "log", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -355,6 +374,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + [[package]] name = "libc" version = "0.2.153" @@ -367,12 +392,31 @@ version = "0.4.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" +[[package]] +name = "matchers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +dependencies = [ + "regex-automata 0.1.10", +] + [[package]] name = "memchr" version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149" +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + [[package]] name = "num-traits" version = "0.2.18" @@ -394,6 +438,18 @@ version = "11.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + +[[package]] +name = "pin-project-lite" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58" + [[package]] name = "plotters" version = "0.3.5" @@ -504,8 +560,17 @@ checksum = "b62dbe01f0b06f9d8dc7d49e05a0785f153b00b2c227856282f671e0318c9b15" dependencies = [ "aho-corasick", "memchr", - "regex-automata", - "regex-syntax", + "regex-automata 0.4.6", + "regex-syntax 0.8.2", +] + +[[package]] +name = "regex-automata" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +dependencies = [ + "regex-syntax 0.6.29", ] [[package]] @@ -516,9 +581,15 @@ checksum = "86b83b8b9847f9bf95ef68afb0b8e6cdb80f498442f5179a29fad448fcc1eaea" dependencies = [ "aho-corasick", "memchr", - "regex-syntax", + "regex-syntax 0.8.2", ] +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + [[package]] name = "regex-syntax" version = "0.8.2" @@ -571,6 +642,21 @@ dependencies = [ "serde", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "smallvec" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7" + [[package]] name = "stwo" version = "0.1.1" @@ -583,7 +669,10 @@ dependencies = [ "itertools 0.12.1", "num-traits", "rand", + "test-log", "thiserror", + "tracing", + "tracing-subscriber", ] [[package]] @@ -603,6 +692,28 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "test-log" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b319995299c65d522680decf80f2c108d85b861d81dfe340a10d16cee29d9e6" +dependencies = [ + "env_logger", + "test-log-macros", + "tracing-subscriber", +] + +[[package]] +name = "test-log-macros" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8f546451eaa38373f549093fe9fd05e7d2bade739e2ddf834b9968621d60107" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "thiserror" version = "1.0.57" @@ -623,6 +734,16 @@ dependencies = [ "syn", ] +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + [[package]] name = "tinytemplate" version = "1.2.1" @@ -633,6 +754,67 @@ dependencies = [ "serde_json", ] +[[package]] +name = "tracing" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + [[package]] name = "typenum" version = "1.17.0" @@ -645,6 +827,12 @@ version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "version_check" version = "0.9.4" diff --git a/Cargo.toml b/Cargo.toml index 2efeac1f4..116bd5b86 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,10 +13,13 @@ itertools = "0.12.0" num-traits = "0.2.17" thiserror = "1.0.56" bytemuck = { version = "1.14.3", features = ["derive"] } +tracing = "0.1.40" [dev-dependencies] criterion = { version = "0.5.1", features = ["html_reports"] } rand = { version = "0.8.5", features = ["small_rng"] } +tracing-subscriber = "0.3.18" +test-log = { version = "0.2.15", features = ["trace"] } [lib] bench = false diff --git a/avx_benchmark.sh b/avx_benchmark.sh new file mode 100755 index 000000000..316e158b5 --- /dev/null +++ b/avx_benchmark.sh @@ -0,0 +1,3 @@ +RUST_LOG_SPAN_EVENTS=enter,close RUST_LOG=info RUST_BACKTRACE=1 \ + RUSTFLAGS="-C target-cpu=native -C target-feature=+avx512f -C opt-level=2" \ + cargo test test_avx_wide_fib_prove -- --nocapture diff --git a/src/core/air/accumulation.rs b/src/core/air/accumulation.rs index c3d9be479..dca4bc8b9 100644 --- a/src/core/air/accumulation.rs +++ b/src/core/air/accumulation.rs @@ -4,6 +4,7 @@ //! f(p) = sum_i alpha^{N-1-i} u_i(P). use itertools::Itertools; +use tracing::{span, Level}; use crate::core::backend::{Backend, CPUBackend}; use crate::core::fields::m31::BaseField; @@ -114,6 +115,7 @@ impl DomainEvaluationAccumulator { 0, "not all random coefficients were used" ); + let _span = span!(Level::INFO, "Constraints interpolation").entered(); let mut res_coeffs = SecureColumn::::zeros(1 << self.log_size()); let res_log_size = self.log_size(); diff --git a/src/core/commitment_scheme/prover.rs b/src/core/commitment_scheme/prover.rs index 34508e520..7a7f06a97 100644 --- a/src/core/commitment_scheme/prover.rs +++ b/src/core/commitment_scheme/prover.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; use itertools::Itertools; +use tracing::{span, Level}; use super::super::channel::Blake2sChannel; use super::super::circle::CirclePoint; @@ -147,6 +148,7 @@ impl> CommitmentTreeProver { log_blowup_factor: u32, channel: &mut ProofChannel, ) -> Self { + let span = span!(Level::INFO, "Commitment evaluation").entered(); let evaluations = polynomials .iter() .map(|poly| { @@ -156,6 +158,9 @@ impl> CommitmentTreeProver { }) .collect_vec(); + span.exit(); + + let _span = span!(Level::INFO, "Commitment merkle").entered(); let tree = MerkleProver::commit(evaluations.iter().map(|eval| &eval.values).collect()); channel.mix_digest(tree.root()); diff --git a/src/core/commitment_scheme/quotients.rs b/src/core/commitment_scheme/quotients.rs index 7cd42d68e..13e1d66de 100644 --- a/src/core/commitment_scheme/quotients.rs +++ b/src/core/commitment_scheme/quotients.rs @@ -3,6 +3,7 @@ use std::collections::BTreeMap; use std::iter::zip; use itertools::{izip, multiunzip, Itertools}; +use tracing::{span, Level}; use crate::core::backend::cpu::quotients::{accumulate_row_quotients, quotient_constants}; use crate::core::circle::CirclePoint; @@ -76,6 +77,7 @@ pub fn compute_fri_quotients( samples: &[Vec], random_coeff: SecureField, ) -> Vec> { + let _span = span!(Level::INFO, "Compute FRI quotients").entered(); zip(columns, samples) .sorted_by_key(|(c, _)| Reverse(c.domain.log_size())) .group_by(|(c, _)| c.domain.log_size()) diff --git a/src/core/fri.rs b/src/core/fri.rs index 4eb00a057..0ff83dd31 100644 --- a/src/core/fri.rs +++ b/src/core/fri.rs @@ -7,6 +7,7 @@ use std::ops::RangeInclusive; use itertools::Itertools; use num_traits::Zero; use thiserror::Error; +use tracing::{span, Level}; use super::backend::CPUBackend; use super::channel::Channel; @@ -149,6 +150,7 @@ impl, H: MerkleHasher> FriProver { columns: &[SecureEvaluation], twiddles: &TwiddleTree, ) -> Self { + let _span = span!(Level::INFO, "FRI commitment").entered(); assert!(!columns.is_empty(), "no columns"); assert!(columns.is_sorted_by_key(|e| Reverse(e.len())), "not sorted"); assert!(columns.iter().all(|e| e.domain.is_canonic()), "not canonic"); diff --git a/src/core/proof_of_work.rs b/src/core/proof_of_work.rs index 8484db026..b4b0472e5 100644 --- a/src/core/proof_of_work.rs +++ b/src/core/proof_of_work.rs @@ -1,4 +1,5 @@ use thiserror::Error; +use tracing::{span, Level}; use super::channel::Blake2sChannel; use crate::commitment_scheme::blake2_hash::{Blake2sHash, Blake2sHasher}; @@ -22,6 +23,7 @@ impl ProofOfWork { } pub fn prove(&self, channel: &mut Blake2sChannel) -> ProofOfWorkProof { + let _span = span!(Level::INFO, "Proof of work").entered(); let seed = channel.get_digest().as_ref().to_vec(); let proof = self.grind(seed); channel.mix_nonce(proof.nonce); diff --git a/src/core/prover/mod.rs b/src/core/prover/mod.rs index d618d4c20..1e6b15f35 100644 --- a/src/core/prover/mod.rs +++ b/src/core/prover/mod.rs @@ -1,5 +1,6 @@ use itertools::Itertools; use thiserror::Error; +use tracing::{span, Level}; use super::backend::Backend; use super::commitment_scheme::{CommitmentSchemeProof, TreeVec}; @@ -71,17 +72,28 @@ pub fn prove>( // Evaluate and commit on trace. // TODO(spapini): Commit on trace outside. + let span = span!(Level::INFO, "Trace interpolation").entered(); let trace_polys = trace.into_iter().map(|poly| poly.interpolate()).collect(); + span.exit(); + let mut commitment_scheme = CommitmentSchemeProver::new(LOG_BLOWUP_FACTOR); + let span = span!(Level::INFO, "Trace commitment").entered(); commitment_scheme.commit(trace_polys, channel); + span.exit(); // Evaluate and commit on composition polynomial. let random_coeff = channel.draw_felt(); + + let span = span!(Level::INFO, "Composition generation").entered(); let composition_polynomial_poly = air.compute_composition_polynomial( random_coeff, &air.component_traces(&commitment_scheme.trees[0].polynomials), ); + span.exit(); + + let span = span!(Level::INFO, "Composition commitment").entered(); commitment_scheme.commit(composition_polynomial_poly.to_vec(), channel); + span.exit(); // Draw OODS point. let oods_point = CirclePoint::::get_random_point(channel); diff --git a/src/examples/wide_fibonacci/avx.rs b/src/examples/wide_fibonacci/avx.rs index db928c6d3..76e84bd03 100644 --- a/src/examples/wide_fibonacci/avx.rs +++ b/src/examples/wide_fibonacci/avx.rs @@ -1,5 +1,6 @@ use itertools::Itertools; use num_traits::{One, Zero}; +use tracing::{span, Level}; use super::component::{WideFibAir, WideFibComponent}; use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator}; @@ -69,6 +70,7 @@ impl Component for WideFibComponent { trace: &ComponentTrace<'_, AVX512Backend>, evaluation_accumulator: &mut DomainEvaluationAccumulator, ) { + let span = span!(Level::INFO, "Constraint eval extension").entered(); assert_eq!(trace.columns.len(), N_COLS); // TODO(spapini): Steal evaluation from commitment. let eval_domain = CanonicCoset::new(self.log_size + 1).circle_domain(); @@ -77,8 +79,10 @@ impl Component for WideFibComponent { .iter() .map(|poly| poly.evaluate(eval_domain)) .collect_vec(); + span.exit(); // Denoms. + let span = span!(Level::INFO, "Constraint eval denominators").entered(); // TODO(spapini): Make this prettier. let zero_domain = CanonicCoset::new(self.log_size).coset; let mut denoms = @@ -86,6 +90,9 @@ impl Component for WideFibComponent { >::bit_reverse_column(&mut denoms); let mut denom_inverses = BaseFieldVec::zeros(denoms.len()); >::batch_inverse(&denoms, &mut denom_inverses); + span.exit(); + + let _span = span!(Level::INFO, "Constraint pointwise eval").entered(); let constraint_log_degree_bound = self.log_size + 1; let [accum] = evaluation_accumulator.columns([(constraint_log_degree_bound, N_COLS - 1)]); @@ -147,6 +154,8 @@ impl Component for WideFibComponent { #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] #[cfg(test)] mod tests { + use tracing::{span, Level}; + use crate::commitment_scheme::blake2_hash::Blake2sHasher; use crate::commitment_scheme::hasher::Hasher; use crate::core::channel::{Blake2sChannel, Channel}; @@ -156,13 +165,20 @@ mod tests { use crate::examples::wide_fibonacci::avx::{gen_trace, WideFibAir}; use crate::examples::wide_fibonacci::component::WideFibComponent; - #[test] + #[test_log::test] fn test_avx_wide_fib_prove() { - // Note: For benchmarks, increase to 17, to get 128MB of trace. - const LOG_SIZE: u32 = 12; + // Note: To see time measurement, run test with + // RUST_LOG_SPAN_EVENTS=enter,close RUST_LOG=info RUST_BACKTRACE=1 RUSTFLAGS=" + // -C target-cpu=native -C target-feature=+avx512f -C opt-level=2" cargo test + // test_avx_wide_fib_prove -- --nocapture + + // Note: 17 means 128MB of trace. + const LOG_SIZE: u32 = 17; let component = WideFibComponent { log_size: LOG_SIZE }; let air = WideFibAir { component }; + let span = span!(Level::INFO, "Trace generation").entered(); let trace = gen_trace(LOG_SIZE as usize); + span.exit(); let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[]))); let proof = prove(&air, channel, trace).unwrap(); diff --git a/src/examples/wide_fibonacci/mod.rs b/src/examples/wide_fibonacci/mod.rs index 8a7ec8590..9d2d676e7 100644 --- a/src/examples/wide_fibonacci/mod.rs +++ b/src/examples/wide_fibonacci/mod.rs @@ -98,8 +98,12 @@ mod tests { } } - #[test] + #[test_log::test] fn test_prove() { + // Note: To see time measurement, run test with + // RUST_LOG_SPAN_EVENTS=enter,close RUST_LOG=info RUST_BACKTRACE=1 cargo test + // test_prove -- --nocapture + let wide_fib = WideFibComponent { log_size: 7 }; let wide_fib_air = WideFibAir { component: wide_fib,