paulistrings/engine/
trace.rs

1//! [`PartitionTrace`], the opt-in per-layer record of a run, partitioned or not, and the per-partition rows it is transposed from.
2
3/// What one layer did across the whole group; the per-partition vectors are indexed by rank.
4#[derive(Clone, Debug, Default, PartialEq, Eq)]
5pub struct PartitionLayerRecord {
6    /// Bucket bits every partition held for this layer.
7    pub bits: u8,
8    /// Remote deltas the layer had; zero means it made no transport call.
9    pub remote_deltas: u32,
10    /// Collective calls the layer issued besides the exchange: the scheduled bucket-count all-reduce plus the policy's collective finalization.
11    pub collectives: u32,
12    /// This layer's position in the [`Circuit`](crate::Circuit), independent of direction.
13    pub circuit_index: u32,
14    /// This layer's position in the propagation loop.
15    pub application_index: u32,
16    /// [`Channel::debug_name`](crate::Channel::debug_name) of the applied channel.
17    pub gate_name: &'static str,
18    /// Terms each partition held before the layer.
19    pub terms_in: Vec<usize>,
20    /// Terms each partition held after the layer and its finalization.
21    pub terms_out: Vec<usize>,
22    /// Rows sent, `rows_sent[from][to]`.
23    pub rows_sent: Vec<Vec<u64>>,
24    /// Wire bytes sent, `bytes_sent[from][to]`, for the same rows.
25    pub bytes_sent: Vec<Vec<u64>>,
26    /// Rows each partition received, summed over its partners.
27    pub rows_received: Vec<u64>,
28    /// Each partition's own wall time for the layer; ranks are not synchronized mid-layer, so `max` is only a critical-rank proxy.
29    pub nanos: Vec<u64>,
30    /// Passes each partition made over its terms to raise the bucket count for this layer: one per doubling on the host, one per counting pass of up to four bits on a device.
31    pub refine_passes: Vec<u32>,
32    /// Chunks each partition's received rows moved in; zero on a layer without an exchange.
33    pub receive_chunks: Vec<u32>,
34    /// Each partition's device counters, `None` on a host partition.
35    pub device: Vec<Option<DeviceLayerRecord>>,
36}
37
38/// What one device partition counted for a layer beyond the shared record.
39#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
40pub struct DeviceLayerRecord {
41    /// Pre-dedup records over every fused block; on the permutation path, its input rows.
42    pub records: u64,
43    /// The largest fused block's records.
44    pub max_block_records: u32,
45    /// Record capacity of the launched fused-layer variant.
46    pub record_capacity: u32,
47    /// Arena batches of the fused layer.
48    pub batches: u32,
49    /// Fused blocks that fell back to the high-word fingerprint passes, the sender-side merge's included.
50    pub high_word_fallbacks: u32,
51    /// Fused blocks that fell back to the full-key sort, the sender-side merge's included.
52    pub full_key_fallbacks: u32,
53    /// Exported rows the sender-side merge folded away before the exchange.
54    pub rows_premerged: u64,
55    /// The layer took the key-preserving rescale.
56    pub rescaled: bool,
57    /// The layer took the permutation path.
58    pub permuted: bool,
59    /// The fused layer reduced by the segmented scan rather than the head-serial walk.
60    pub dense: bool,
61}
62
63/// Per-layer records of propagations in application order, accumulated until drained by a driver's `take_trace`.
64///
65/// A plain [`propagate_with`](crate::propagate_with) run records one partition; term counts are post-truncation, so the transient expansion inside a layer is not captured.
66#[derive(Clone, Debug, Default, PartialEq, Eq)]
67pub struct PartitionTrace {
68    /// One record per layer applied.
69    pub layers: Vec<PartitionLayerRecord>,
70}
71
72impl PartitionTrace {
73    /// Peak resident term count between layers, summed over partitions, or `None` for a zero-layer trace.
74    pub fn peak_terms(&self) -> Option<usize> {
75        let first = self.layers.first()?;
76        let total = |counts: &[usize]| counts.iter().sum::<usize>();
77        let peak_out = self.layers.iter().map(|layer| total(&layer.terms_out)).max();
78        Some(total(&first.terms_in).max(peak_out.unwrap_or(0)))
79    }
80
81    /// Rows moved across partitions over the whole trace.
82    pub fn total_rows_exchanged(&self) -> u64 {
83        self.layers.iter().flat_map(|layer| layer.rows_sent.iter()).flat_map(|row| row.iter()).sum()
84    }
85
86    /// Per layer, the maximum input term count over partitions divided by the mean; `1.0` for a layer where every partition was empty.
87    pub fn imbalance(&self) -> Vec<f64> {
88        self.layers
89            .iter()
90            .map(|layer| {
91                let total: usize = layer.terms_in.iter().sum();
92                let max = layer.terms_in.iter().copied().max().unwrap_or(0);
93                if total == 0 {
94                    return 1.0;
95                }
96                let mean = total as f64 / layer.terms_in.len() as f64;
97                max as f64 / mean
98            })
99            .collect()
100    }
101
102    /// Layers that made no transport call.
103    pub fn local_layers(&self) -> usize {
104        self.layers.iter().filter(|layer| layer.remote_deltas == 0).count()
105    }
106
107    /// Layers that exchanged.
108    pub fn remote_layers(&self) -> usize {
109        self.layers.iter().filter(|layer| layer.remote_deltas > 0).count()
110    }
111
112    /// Collective calls over the whole trace, excluding the point-to-point exchanges.
113    pub fn total_collectives(&self) -> u64 {
114        self.layers.iter().map(|layer| u64::from(layer.collectives)).sum()
115    }
116}
117
118/// One layer as a single partition saw it, before [`assemble`] transposes it.
119#[derive(Clone, Debug, Default, PartialEq, Eq)]
120pub(crate) struct PartitionLayerRow {
121    pub bits: u8,
122    pub remote_deltas: u32,
123    pub collectives: u32,
124    pub circuit_index: u32,
125    pub application_index: u32,
126    pub gate_name: &'static str,
127    pub terms_in: usize,
128    pub terms_out: usize,
129    pub rows_sent: Vec<u64>,
130    pub bytes_sent: Vec<u64>,
131    pub rows_received: u64,
132    pub nanos: u64,
133    pub refine_passes: u32,
134    pub receive_chunks: u32,
135    pub device: Option<DeviceLayerRecord>,
136}
137
138/// Append one layer's record; cold and out of line to keep it from perturbing the inlined merge kernels' layout.
139#[cold]
140#[inline(never)]
141pub(super) fn record_layer_row(rows: &mut Vec<PartitionLayerRow>, row: PartitionLayerRow) {
142    rows.push(row);
143}
144
145/// Transposes the partitions' rows into per-layer records appended to `trace`, the per-partner vectors padded to `partners`; panics if the partitions disagree on a collective decision.
146pub(crate) fn assemble(
147    trace: &mut PartitionTrace,
148    per_partition: Vec<Vec<PartitionLayerRow>>,
149    partners: usize,
150) {
151    let size = per_partition.len();
152    let layers = per_partition[0].len();
153    for (rank, rows) in per_partition.iter().enumerate() {
154        assert_eq!(
155            rows.len(),
156            layers,
157            "partition {rank} traced {} layers, partition 0 traced {layers}",
158            rows.len(),
159        );
160    }
161
162    trace.layers.reserve(layers);
163    for k in 0..layers {
164        let head = &per_partition[0][k];
165        let mut record = PartitionLayerRecord {
166            bits: head.bits,
167            remote_deltas: head.remote_deltas,
168            collectives: head.collectives,
169            circuit_index: head.circuit_index,
170            application_index: head.application_index,
171            gate_name: head.gate_name,
172            terms_in: Vec::with_capacity(size),
173            terms_out: Vec::with_capacity(size),
174            rows_sent: Vec::with_capacity(size),
175            bytes_sent: Vec::with_capacity(size),
176            rows_received: Vec::with_capacity(size),
177            nanos: Vec::with_capacity(size),
178            refine_passes: Vec::with_capacity(size),
179            receive_chunks: Vec::with_capacity(size),
180            device: Vec::with_capacity(size),
181        };
182        for (rank, rows) in per_partition.iter().enumerate() {
183            let row = &rows[k];
184            assert_eq!(
185                row.bits, head.bits,
186                "layer {k}: partition {rank} has {} bucket bits, partition 0 has {}",
187                row.bits, head.bits,
188            );
189            assert_eq!(
190                row.remote_deltas, head.remote_deltas,
191                "layer {k}: partition {rank} saw {} remote deltas, partition 0 saw {}",
192                row.remote_deltas, head.remote_deltas,
193            );
194            assert_eq!(
195                row.collectives, head.collectives,
196                "layer {k}: partition {rank} issued {} collectives, partition 0 issued {}",
197                row.collectives, head.collectives,
198            );
199            assert_eq!(
200                row.circuit_index, head.circuit_index,
201                "layer {k}: partition {rank} ran circuit index {}, partition 0 ran {}",
202                row.circuit_index, head.circuit_index,
203            );
204            assert_eq!(
205                row.application_index, head.application_index,
206                "layer {k}: partition {rank} ran application index {}, partition 0 ran {}",
207                row.application_index, head.application_index,
208            );
209            assert_eq!(
210                row.gate_name, head.gate_name,
211                "layer {k}: partition {rank} ran gate {:?}, partition 0 ran {:?}",
212                row.gate_name, head.gate_name,
213            );
214            record.terms_in.push(row.terms_in);
215            record.terms_out.push(row.terms_out);
216            // A layer without an exchange reports empty per-partner vectors.
217            let mut rows_sent = row.rows_sent.clone();
218            rows_sent.resize(partners, 0);
219            let mut bytes_sent = row.bytes_sent.clone();
220            bytes_sent.resize(partners, 0);
221            record.rows_sent.push(rows_sent);
222            record.bytes_sent.push(bytes_sent);
223            record.rows_received.push(row.rows_received);
224            record.nanos.push(row.nanos);
225            record.refine_passes.push(row.refine_passes);
226            record.receive_chunks.push(row.receive_chunks);
227            record.device.push(row.device);
228        }
229        trace.layers.push(record);
230    }
231}
232
233#[cfg(test)]
234mod tests;
paulistrings/engine/trace.rs

paulistrings/engine/
trace.rs
