More timing code for FRI prover (#146)

src/fri/commitment.rs

@@ -147,7 +147,10 @@ impl<F: Field> PolynomialBatchCommitment<F> {
             );
         }
 
-        let os = OpeningSet::new(
+        let os = timed!(
+            timing,
+            "construct the opening set",
+            OpeningSet::new(
                 zeta,
                 g,
                 commitments[0],
@@ -155,6 +158,7 @@ impl<F: Field> PolynomialBatchCommitment<F> {
                 commitments[2],
                 commitments[3],
                 common_data,
+            )
         );
         challenger.observe_opening_set(&os);
 
@@ -181,21 +185,33 @@ impl<F: Field> PolynomialBatchCommitment<F> {
         .flat_map(|&p| &commitments[p.index].polynomials)
         .map(|p| p.to_extension())
         .chain(partial_products_polys);
-        let single_composition_poly = alpha.reduce_polys(single_polys);
+        let single_composition_poly = timed!(
+            timing,
+            "reduce single polys",
+            alpha.reduce_polys(single_polys)
+        );
 
         let single_quotient = Self::compute_quotient([zeta], single_composition_poly);
         final_poly += single_quotient;
         alpha.reset();
 
         // Zs polynomials are opened at `zeta` and `g*zeta`.
-        let zs_composition_poly = alpha.reduce_polys(zs_polys.into_iter());
+        let zs_composition_poly = timed!(
+            timing,
+            "reduce Z polys",
+            alpha.reduce_polys(zs_polys.into_iter())
+        );
 
         let zs_quotient = Self::compute_quotient([zeta, g * zeta], zs_composition_poly);
         alpha.shift_poly(&mut final_poly);
         final_poly += zs_quotient;
 
         let lde_final_poly = final_poly.lde(config.rate_bits);
-        let lde_final_values = lde_final_poly.coset_fft(F::coset_shift().into());
+        let lde_final_values = timed!(
+            timing,
+            &format!("perform final FFT {}", lde_final_poly.len()),
+            lde_final_poly.coset_fft(F::coset_shift().into())
+        );
 
         let fri_proof = fri_proof(
             &commitments

src/fri/prover.rs

@@ -29,11 +29,15 @@ pub fn fri_proof<F: Field + Extendable<D>, const D: usize>(
     assert_eq!(lde_polynomial_coeffs.coeffs.len(), n);
 
     // Commit phase
-    let (trees, final_coeffs) = fri_committed_trees(
+    let (trees, final_coeffs) = timed!(
+        timing,
+        "fold codewords in the commitment phase",
+        fri_committed_trees(
             lde_polynomial_coeffs,
             lde_polynomial_values,
             challenger,
             config,
+        )
     );
 
     // PoW phase