Merge m4b/b1 @ 11d3e43 (b1 rebased on #1014 59e9890, the rest sink on) into m4b/b2

Author: MauroToscano

crypto/stark/src/multilinear_block.rs

@@ -741,15 +741,20 @@ where
                         .unwrap_or_else(|panic| std::panic::resume_unwind(panic));
                     let next = next.map(|(group, store, wide, upload, uploaded)| {
                         // What the committer waited for once its commit ended:
-                        // the group's arrival, then the rest of its upload.
+                        // the group's arrival, then the rest of its upload —
+                        // never more than the upload (a commit that ended before
+                        // the upload began paid all of it, and no more).
                         let ready = arrived.max(commit_end);
                         Ahead {
                             group,
                             store,
                             wide,
                             upload,
                             wait_paid: arrived.saturating_duration_since(commit_end).as_secs_f64(),
-                            upload_paid: uploaded.saturating_duration_since(ready).as_secs_f64(),
+                            upload_paid: uploaded
+                                .saturating_duration_since(ready)
+                                .as_secs_f64()
+                                .min(upload),
                         }
                     });
                     (committed, next)

prover/src/block_whir.rs

@@ -235,10 +235,13 @@ pub struct BlockOptions {
     pub drop_streamed_ops: bool,
     /// With windows: `0` lays each streamed chunk out on the layout thread as
     /// it arrives; `n > 0` lays them out on `n` threads, packed in arrival
-    /// order all the same, while the layout thread lays out the rest of the
-    /// run as soon as it is built. Production: 0. Three read −0.54 / −0.22 s of
-    /// base (FAST 421 / 422), but their extra host memory grows with the block
-    /// (BIG 390: +2.2 / +4.2 / +8.7 GiB at 1.0 / 1.3 / 1.8×).
+    /// order all the same (the same groups and proof bytes), while the layout
+    /// thread lays out the rest of the run as soon as it is built. Production:
+    /// 3. Once phase A uploads ahead, the inline layout closes the middle groups
+    /// after the card is free; three workers close them in time (FAST 835:
+    /// phase A −0.93 s, whole block −1.00 s, peak −0.68 GiB at 1×). Unbounded
+    /// they held more host memory as the block grew (BIG 390); `layout_ahead`
+    /// bounds them (BIG 392: flat).
     pub layout_workers: usize,
     /// With `layout_workers > 0`: `Some(k)` lets at most `k + 1` streamed
     /// chunks be laid out (or in the making) and not yet packed, where the
@@ -247,7 +250,12 @@ pub struct BlockOptions {
     pub layout_ahead: Option<usize>,
     /// With `layout_workers > 0`: pack the rest of the run as it is laid out,
     /// in AIR order, so a group closes once its own tables are ready instead
-    /// of after all of them; the packing order is the same.
+    /// of after all of them; the packing order is the same. Production: on,
+    /// with [`Self::rest_layout_bytes`]'s waves — laid out all at once, the
+    /// first table in AIR order was ready only with the last, and packing as
+    /// laid out read +0.21 s (FAST 422); in waves the AIR order's prefix is
+    /// laid out first, and since three layout workers (04984808d) phase A
+    /// waits on the groups the rest closes.
     pub pack_rest_as_laid_out: bool,
     /// How phase A holds each group's columns once committed
     /// ([`multilinear_block::Narrowing`]): production packs them narrow on the
@@ -302,9 +310,9 @@ impl BlockOptions {
             stream_keccak_rnd: false,
             stream_memw_lt: false,
             drop_streamed_ops: true,
-            layout_workers: 0,
+            layout_workers: 3,
             layout_ahead: Some(2),
-            pack_rest_as_laid_out: false,
+            pack_rest_as_laid_out: true,
             narrow: multilinear_block::Narrowing::CARD,
             upload_ahead: true,
             memlog: memlog::from_env(),
@@ -518,6 +526,9 @@ pub struct LayoutStamps {
     /// The rest's tables laid out narrow from the finish's packed columns
     /// ([`BlockOptions::pack_finished`]): how many, and their packed bytes.
     pub packed_rest: (usize, usize),
+    /// Each table of the rest in AIR order: its name, when its layout ended
+    /// (seconds since the prove started) and its group.
+    pub rest_tables: Vec<(String, f64, usize)>,
 }
 
 impl BlockStamps {
@@ -536,7 +547,7 @@ impl BlockStamps {
         );
         for (g, s) in self.groups.iter().enumerate() {
             out.push_str(&format!(
-                "BLOCK GROUP {g}: tables {} · polys {} · cells {:.1} M || A wait {:.3} upload {:.3} paid {:.3} commit {:.3} retire {:.3} done@{:.2} || B upload {:.3} argue {:.3} encode {:.3} open {:.3}\n",
+                "BLOCK GROUP {g}: tables {} · polys {} · cells {:.1} M || A wait {:.3} upload {:.3} paid {:.3} commit {:.3} retire {:.3} from@{:.2} done@{:.2} || B upload {:.3} argue {:.3} encode {:.3} open {:.3}\n",
                 s.tables,
                 s.polys,
                 s.cells as f64 / 1e6,
@@ -545,6 +556,7 @@ impl BlockStamps {
                 s.upload_paid,
                 s.commit,
                 s.retire,
+                s.committed_at - s.commit - s.retire,
                 s.committed_at,
                 s.upload_b,
                 s.argue,
@@ -607,6 +619,15 @@ impl BlockStamps {
                 l.packed_rest.0,
                 l.packed_rest.1 as f64 / (1u64 << 30) as f64,
             ));
+            let tables: Vec<String> = l
+                .rest_tables
+                .iter()
+                .map(|(name, at, group)| format!("{name}@{at:.2}·g{group}"))
+                .collect();
+            out.push_str(&format!(
+                "BLOCK REST TABLES (laid out at · group): {}\n",
+                tables.join(" ")
+            ));
         }
         out.push_str(&format!(
             "BLOCK RECOMMIT: 0 (a revived commitment builds no tree) · first-round paths from kept tops: {} calls · {} leaves re-hashed on the host · phase B on one thread + one upload helper\n",
@@ -2283,6 +2304,9 @@ struct RestLaid<'a> {
     busy: f64,
     /// [`LayoutStamps::packed_rest`].
     packed_rest: (usize, usize),
+    /// Each table's AIR index and when its layout ended (seconds since the
+    /// prove started), in AIR order.
+    laid_at: Vec<(usize, f64)>,
 }
 
 /// `prepared_now`: derive the prepared columns here, before the tables are
@@ -2351,8 +2375,9 @@ fn lay_out_rest<'a>(
         .enumerate()
         .filter(|(_, ((_, trace, _), _))| trace.main_table.width != 0)
         .collect();
-    // Each table's seconds, for the readout, and the tables laid out narrow.
+    // Each table's seconds and when its layout ended, for the readout.
     let timed = std::sync::Mutex::new(Vec::with_capacity(built.len()));
+    let laid_at = std::sync::Mutex::new(Vec::with_capacity(built.len()));
     let narrow_tables = std::sync::atomic::AtomicUsize::new(0);
     let narrow_bytes = std::sync::atomic::AtomicUsize::new(0);
     let lay_out = |(i, ((_, trace, _), air)): (usize, ((_, &mut TraceTable<F, E>, _), _))| {
@@ -2371,6 +2396,9 @@ fn lay_out_rest<'a>(
             table_of(air, trace, shape, true)
         }
         .map(|table| (i, shape, table));
+        if let Ok(mut laid_at) = laid_at.lock() {
+            laid_at.push((i, start.elapsed().as_secs_f64()));
+        }
         if let (Some(ledger), Ok((_, _, table))) = (ledger, &laid) {
             use std::sync::atomic::Ordering::Relaxed;
             ledger.rest.fetch_sub(rows, Relaxed);
@@ -2472,6 +2500,11 @@ fn lay_out_rest<'a>(
         slowest,
         busy: t.elapsed().as_secs_f64(),
         packed_rest: (narrow_tables.into_inner(), narrow_bytes.into_inner()),
+        laid_at: {
+            let mut laid_at = laid_at.into_inner().unwrap_or_default();
+            laid_at.sort_unstable_by_key(|&(i, _)| i);
+            laid_at
+        },
     })
 }
 
@@ -2801,6 +2834,7 @@ fn prove_streamed(
                     slowest,
                     busy: rest_busy,
                     packed_rest,
+                    laid_at,
                 } = rest;
                 let names: std::collections::HashMap<String, usize> = refs
                     .iter()
@@ -2877,6 +2911,16 @@ fn prove_streamed(
                     })
                     .collect::<Result<Vec<_>, Error>>()?;
                 let busy = chunks + rest_busy + t.elapsed().as_secs_f64();
+                let group_of = |table: usize| {
+                    groups
+                        .iter()
+                        .position(|g| g.contains(&(table as u32)))
+                        .unwrap_or(usize::MAX)
+                };
+                let rest_tables = laid_at
+                    .iter()
+                    .map(|&(i, at)| (refs[i].name().to_string(), at, group_of(i)))
+                    .collect();
                 Ok(Laid {
                     table_counts,
                     runtime_page_ranges,
@@ -2895,6 +2939,7 @@ fn prove_streamed(
                         slowest,
                         ahead_most,
                         packed_rest,
+                        rest_tables,
                     },
                 })
             });

prover/src/lfm/whir_block_tests.rs

@@ -1356,7 +1356,6 @@ fn the_whir_block_tree_on_a_real_block() {
     if let Some(n) = knob("BLOCK_WHIR_LAYOUT_WORKERS") {
         options.layout_workers = n;
     }
-    // `BLOCK_WHIR_PACK_REST=1`: the rest of the run packed as it is laid out.
     // `BLOCK_WHIR_LAYOUT_AHEAD=k` bounds the chunks unpacked at k + 1;
     // `none` lifts the bound (the reverted version, BIG 390).
     match std::env::var("BLOCK_WHIR_LAYOUT_AHEAD")
@@ -1371,8 +1370,17 @@ fn the_whir_block_tree_on_a_real_block() {
         }
         Err(_) => {}
     }
-    options.pack_rest_as_laid_out =
-        std::env::var("BLOCK_WHIR_PACK_REST").is_ok_and(|v| v.trim() == "1");
+    // `BLOCK_WHIR_PACK_REST=0|1` (production 1): the rest packed as it is laid
+    // out.
+    match std::env::var("BLOCK_WHIR_PACK_REST")
+        .as_deref()
+        .map(str::trim)
+    {
+        Ok("0") => options.pack_rest_as_laid_out = false,
+        Ok("1") => options.pack_rest_as_laid_out = true,
+        Ok(other) => panic!("BLOCK_WHIR_PACK_REST={other}: 0 or 1"),
+        Err(_) => {}
+    }
     // `BLOCK_WHIR_DROP_OPS=0`: the builder keeps the streamed chunks' ops
     // (production drops them).
     match std::env::var("BLOCK_WHIR_DROP_OPS")
@@ -1510,6 +1518,24 @@ fn the_whir_block_tree_on_a_real_block() {
             options.keccak_rows_log2,
             options.ecsm_rows_log2,
         );
+        // Each group's tables by AIR, in the group's order: which tables the
+        // card waited for, and where the layout's chunks went.
+        for (g, group) in proof.groups.iter().enumerate() {
+            let mut kinds: Vec<(String, usize)> = Vec::new();
+            for &index in group {
+                let name = refs[index as usize].name();
+                let kind = name.split('[').next().unwrap_or(name).to_string();
+                match kinds.last_mut() {
+                    Some((last, n)) if *last == kind => *n += 1,
+                    _ => kinds.push((kind, 1)),
+                }
+            }
+            let kinds: Vec<String> = kinds
+                .into_iter()
+                .map(|(kind, n)| if n == 1 { kind } else { format!("{kind}×{n}") })
+                .collect();
+            println!("W3 GROUP TABLES {g}: {}", kinds.join(" "));
+        }
     }
     println!(
         "W3 BASE: {base:.2}s · statement at {stated_at:.2}s · plan + {} leaves emitted by {ready_at:.2}s ({})",

prover/src/tests/block_whir_tests.rs

@@ -2102,7 +2102,7 @@ fn block_whir_on_a_real_block() {
     };
     let mut options = BlockOptions::production();
     // `BLOCK_WHIR_LAYOUT_WORKERS=n` (production 0, the inline layout) and
-    // `BLOCK_WHIR_PACK_REST=1`, as the tree's harness takes them;
+    // `BLOCK_WHIR_PACK_REST=0|1`, as the tree's harness takes them;
     // `BLOCK_WHIR_DROP_OPS=0`: the builder keeps the streamed chunks' ops
     // (production drops them).
     if let Some(n) = std::env::var("BLOCK_WHIR_LAYOUT_WORKERS")
@@ -2125,8 +2125,17 @@ fn block_whir_on_a_real_block() {
         }
         Err(_) => {}
     }
-    options.pack_rest_as_laid_out =
-        std::env::var("BLOCK_WHIR_PACK_REST").is_ok_and(|v| v.trim() == "1");
+    // `BLOCK_WHIR_PACK_REST=0|1` (production 1): the rest packed as it is laid
+    // out.
+    match std::env::var("BLOCK_WHIR_PACK_REST")
+        .as_deref()
+        .map(str::trim)
+    {
+        Ok("0") => options.pack_rest_as_laid_out = false,
+        Ok("1") => options.pack_rest_as_laid_out = true,
+        Ok(other) => panic!("BLOCK_WHIR_PACK_REST={other}: 0 or 1"),
+        Err(_) => {}
+    }
     match std::env::var("BLOCK_WHIR_DROP_OPS")
         .as_deref()
         .map(str::trim)