candle_graph/trace/
memory.rs

1//! Memory timeline analysis — TensorFlow Memory Profile model.
2
3use std::collections::{BTreeMap, HashMap};
4
5use serde::{Deserialize, Serialize};
6
7use super::document::TraceDocument;
8
9/// TensorFlow-profiler-style memory category for timeline breakdown.
10#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum MemoryCategory {
13    Parameter,
14    Activation,
15    Gradient,
16    Optimizer,
17    #[default]
18    Other,
19}
20
21/// Allocation or deallocation recorded during a probe run.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
23#[serde(rename_all = "snake_case")]
24pub enum MemoryAction {
25    Alloc,
26    Free,
27}
28
29/// Derive element size in bytes from a Candle-style dtype label (`f32`, `F32`, …).
30pub fn dtype_size_bytes(dtype: &str) -> Option<usize> {
31    match dtype.trim().to_ascii_lowercase().as_str() {
32        "u8" | "i8" | "bool" => Some(1),
33        "u16" | "i16" | "f16" | "bf16" => Some(2),
34        "u32" | "i32" | "f32" => Some(4),
35        "u64" | "i64" | "f64" => Some(8),
36        "f8" => Some(1),
37        _ => None,
38    }
39}
40
41/// Product of shape dimensions; returns `0` for empty shape.
42pub fn elem_count(shape: &[usize]) -> u64 {
43    shape
44        .iter()
45        .fold(1u64, |acc, dim| acc.saturating_mul(*dim as u64))
46}
47
48/// Storage bytes for a dense tensor (`elem_count × dtype_bytes`).
49pub fn storage_bytes(shape: &[usize], dtype: &str) -> u64 {
50    let count = elem_count(shape);
51    if count == 0 {
52        return 0;
53    }
54    dtype_size_bytes(dtype)
55        .map(|sz| count.saturating_mul(sz as u64))
56        .unwrap_or(0)
57}
58
59/// Resolve explicit bytes or derive from shape/dtype.
60pub fn resolve_storage_bytes(explicit: Option<u64>, shape: &[usize], dtype: &str) -> u64 {
61    explicit.unwrap_or_else(|| storage_bytes(shape, dtype))
62}
63
64/// Map a training step (+ tensor flags) to a PyTorch-style memory category.
65pub fn category_for_step(
66    step: Option<crate::phase::ExecutionStep>,
67    requires_grad: bool,
68) -> MemoryCategory {
69    match step {
70        Some(crate::phase::ExecutionStep::Backward) => MemoryCategory::Gradient,
71        Some(crate::phase::ExecutionStep::Optimizer) => MemoryCategory::Optimizer,
72        Some(crate::phase::ExecutionStep::Forward) => MemoryCategory::Activation,
73        None if requires_grad => MemoryCategory::Parameter,
74        None => MemoryCategory::Activation,
75    }
76}
77
78fn category_key(category: MemoryCategory) -> String {
79    format!("{category:?}").to_ascii_lowercase()
80}
81
82/// Aggregated memory statistics for one probe run.
83#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
84pub struct MemorySummary {
85    pub alloc_count: u64,
86    pub free_count: u64,
87    /// Sum of all allocation sizes (can exceed peak when tensors are freed).
88    pub total_alloc_bytes: u64,
89    pub peak_bytes: u64,
90    pub peak_timestamp_ns: u64,
91    pub peak_device: String,
92    /// Live bytes by category at global peak (PyTorch memory timeline breakdown).
93    #[serde(default)]
94    pub peak_by_category: BTreeMap<String, u64>,
95    /// Activation bytes still live when backward starts (autograd retention hint).
96    #[serde(default)]
97    pub autograd_retained_bytes: u64,
98}
99
100/// One point on the memory-vs-time curve for a device.
101#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
102pub struct MemoryTimelinePoint {
103    pub timestamp_ns: u64,
104    pub device: String,
105    pub live_bytes: u64,
106    pub heap_bytes: u64,
107    pub free_bytes: u64,
108    pub by_category: BTreeMap<String, u64>,
109}
110
111/// Active allocation at the global peak (TensorFlow breakdown table).
112#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
113pub struct LiveAllocation {
114    pub tensor_id: String,
115    pub span_id: String,
116    pub op_name: Option<String>,
117    pub bytes: u64,
118    pub shape: Vec<usize>,
119    pub dtype: String,
120    pub device: String,
121    pub category: MemoryCategory,
122}
123
124/// Per-device memory stats.
125#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
126pub struct DeviceMemoryStats {
127    pub device: String,
128    pub capacity_bytes: u64,
129    pub peak_bytes: u64,
130    pub peak_timestamp_ns: u64,
131    pub alloc_count: u64,
132    pub free_count: u64,
133}
134
135/// Full memory profile reconstructed from trace evidence.
136#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
137pub struct MemoryProfile {
138    pub summary: MemorySummary,
139    pub timeline: Vec<MemoryTimelinePoint>,
140    pub peak_breakdown: Vec<LiveAllocation>,
141    pub by_device: Vec<DeviceMemoryStats>,
142}
143
144#[derive(Debug, Clone)]
145struct TimelineEvent {
146    timestamp_ns: u64,
147    device: String,
148    action: MemoryAction,
149    bytes: u64,
150    category: MemoryCategory,
151    tensor_id: String,
152    span_id: String,
153    op_name: Option<String>,
154    shape: Vec<usize>,
155    dtype: String,
156}
157
158/// Build a TensorFlow-style memory profile from a parsed trace document.
159pub fn analyze_memory(doc: &TraceDocument) -> MemoryProfile {
160    let events = collect_timeline_events(doc);
161    let span_steps: HashMap<String, crate::phase::ExecutionStep> = doc
162        .spans
163        .iter()
164        .filter_map(|span| span.step.map(|step| (span.id.clone(), step)))
165        .collect();
166    if events.is_empty() {
167        return MemoryProfile {
168            summary: MemorySummary::default(),
169            timeline: Vec::new(),
170            peak_breakdown: Vec::new(),
171            by_device: Vec::new(),
172        };
173    }
174
175    let mut by_device_live: HashMap<String, u64> = HashMap::new();
176    let mut by_device_category: HashMap<String, BTreeMap<String, u64>> = HashMap::new();
177    let mut live_tensors: HashMap<String, LiveAllocation> = HashMap::new();
178
179    let mut global_peak = 0u64;
180    let mut global_peak_ts = 0u64;
181    let mut global_peak_device = String::new();
182    let mut peak_breakdown = Vec::new();
183
184    let mut device_peaks: HashMap<String, (u64, u64)> = HashMap::new();
185    let mut device_capacity: HashMap<String, u64> = HashMap::new();
186    let mut device_alloc_count: HashMap<String, u64> = HashMap::new();
187    let mut device_free_count: HashMap<String, u64> = HashMap::new();
188
189    let mut timeline = Vec::new();
190    let mut alloc_count = 0u64;
191    let mut free_count = 0u64;
192    let mut total_alloc_bytes = 0u64;
193    let mut peak_by_category: BTreeMap<String, u64> = BTreeMap::new();
194    let mut autograd_retained_bytes = 0u64;
195    let mut backward_seen = false;
196
197    for event in &events {
198        let device = event.device.clone();
199        let live = by_device_live.entry(device.clone()).or_insert(0);
200        let categories = by_device_category.entry(device.clone()).or_default();
201        let cat_key = category_key(event.category);
202
203        match event.action {
204            MemoryAction::Alloc => {
205                alloc_count += 1;
206                *device_alloc_count.entry(device.clone()).or_insert(0) += 1;
207                total_alloc_bytes = total_alloc_bytes.saturating_add(event.bytes);
208                *live = live.saturating_add(event.bytes);
209                {
210                    let cat = categories.entry(cat_key).or_insert(0);
211                    *cat = cat.saturating_add(event.bytes);
212                }
213                live_tensors.insert(
214                    event.tensor_id.clone(),
215                    LiveAllocation {
216                        tensor_id: event.tensor_id.clone(),
217                        span_id: event.span_id.clone(),
218                        op_name: event.op_name.clone(),
219                        bytes: event.bytes,
220                        shape: event.shape.clone(),
221                        dtype: event.dtype.clone(),
222                        device: device.clone(),
223                        category: event.category,
224                    },
225                );
226            }
227            MemoryAction::Free => {
228                free_count += 1;
229                *device_free_count.entry(device.clone()).or_insert(0) += 1;
230                if let Some(alloc) = live_tensors.remove(&event.tensor_id) {
231                    *live = live.saturating_sub(alloc.bytes);
232                    let key = category_key(alloc.category);
233                    if let Some(cat_live) = categories.get_mut(&key) {
234                        *cat_live = cat_live.saturating_sub(alloc.bytes);
235                    }
236                } else {
237                    *live = live.saturating_sub(event.bytes);
238                }
239            }
240        }
241
242        let heap = *live;
243        let capacity = device_capacity.get(&device).copied().unwrap_or(0);
244        let free = capacity.saturating_sub(heap);
245
246        timeline.push(MemoryTimelinePoint {
247            timestamp_ns: event.timestamp_ns,
248            device: device.clone(),
249            live_bytes: heap,
250            heap_bytes: heap,
251            free_bytes: free,
252            by_category: categories.clone(),
253        });
254
255        let (peak, peak_ts) = device_peaks.entry(device.clone()).or_insert((0, 0));
256        if heap > *peak {
257            *peak = heap;
258            *peak_ts = event.timestamp_ns;
259        }
260
261        if heap > global_peak {
262            global_peak = heap;
263            global_peak_ts = event.timestamp_ns;
264            global_peak_device = device.clone();
265            peak_breakdown = live_tensors.values().cloned().collect();
266            peak_breakdown.sort_by(|a, b| {
267                b.bytes
268                    .cmp(&a.bytes)
269                    .then_with(|| a.tensor_id.cmp(&b.tensor_id))
270            });
271            peak_by_category = categories.clone();
272        }
273
274        if !backward_seen {
275            let in_backward = span_steps
276                .get(&event.span_id)
277                .is_some_and(|step| *step == crate::phase::ExecutionStep::Backward);
278            if in_backward {
279                backward_seen = true;
280                autograd_retained_bytes = categories
281                    .get(&category_key(MemoryCategory::Activation))
282                    .copied()
283                    .unwrap_or(0);
284            }
285        }
286    }
287
288    for snapshot in &doc.device_memory {
289        let entry = device_capacity.entry(snapshot.device.clone()).or_insert(0);
290        *entry = (*entry).max(snapshot.used_bytes.saturating_add(snapshot.free_bytes));
291    }
292
293    let mut by_device: Vec<DeviceMemoryStats> = device_peaks
294        .into_iter()
295        .map(
296            |(device, (peak_bytes, peak_timestamp_ns))| DeviceMemoryStats {
297                device: device.clone(),
298                capacity_bytes: device_capacity.get(&device).copied().unwrap_or(0),
299                peak_bytes,
300                peak_timestamp_ns,
301                alloc_count: device_alloc_count.get(&device).copied().unwrap_or(0),
302                free_count: device_free_count.get(&device).copied().unwrap_or(0),
303            },
304        )
305        .collect();
306    by_device.sort_by(|a, b| a.device.cmp(&b.device));
307
308    MemoryProfile {
309        summary: MemorySummary {
310            alloc_count,
311            free_count,
312            total_alloc_bytes,
313            peak_bytes: global_peak,
314            peak_timestamp_ns: global_peak_ts,
315            peak_device: global_peak_device,
316            peak_by_category,
317            autograd_retained_bytes,
318        },
319        timeline,
320        peak_breakdown,
321        by_device,
322    }
323}
324
325fn collect_timeline_events(doc: &TraceDocument) -> Vec<TimelineEvent> {
326    let mut events: Vec<TimelineEvent> = Vec::new();
327
328    for mem in &doc.memory {
329        events.push(TimelineEvent {
330            timestamp_ns: mem.timestamp_ns,
331            device: mem.device.clone(),
332            action: mem.action,
333            bytes: mem.bytes,
334            category: mem.category,
335            tensor_id: mem.tensor_id.clone(),
336            span_id: mem.span_id.clone(),
337            op_name: mem.op_name.clone(),
338            shape: mem.shape.clone(),
339            dtype: mem.dtype.clone(),
340        });
341    }
342
343    events.sort_by(|a, b| {
344        a.timestamp_ns
345            .cmp(&b.timestamp_ns)
346            .then_with(|| format!("{:?}", a.action).cmp(&format!("{:?}", b.action)))
347            .then_with(|| a.tensor_id.cmp(&b.tensor_id))
348    });
349    events
350}
351
352/// Per-span / per-op TensorFlow memory metrics.
353#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
354pub struct NodeMemoryMetrics {
355    /// Total bytes requested by ops directly on this node.
356    pub bytes: u64,
357    /// Peak live bytes in this node's subtree during execution.
358    pub peak_bytes: u64,
359    /// Bytes still live when this node finishes (not deallocated).
360    pub residual_bytes: u64,
361    /// Output storage for op leaf nodes.
362    pub storage_bytes: u64,
363}
364
365/// Attribute memory metrics to graph node ids (span ids and `{span}/op/{n}` op ids).
366pub fn node_memory_metrics(
367    doc: &TraceDocument,
368    op_node_ids: &HashMap<(String, usize), String>,
369) -> HashMap<String, NodeMemoryMetrics> {
370    let profile = analyze_memory(doc);
371    let mut metrics: HashMap<String, NodeMemoryMetrics> = HashMap::new();
372
373    let mut per_span_index: HashMap<String, usize> = HashMap::new();
374    for op in &doc.ops {
375        let index = *per_span_index.entry(op.span_id.clone()).or_insert(0);
376        per_span_index.insert(op.span_id.clone(), index + 1);
377
378        let bytes = op
379            .storage_bytes
380            .unwrap_or_else(|| resolve_storage_bytes(None, &op.shape, &op.dtype));
381        if bytes == 0 {
382            continue;
383        }
384
385        let node_id = op_node_ids
386            .get(&(op.span_id.clone(), index))
387            .cloned()
388            .unwrap_or_else(|| format!("{}/op/{}", op.span_id, index));
389
390        let entry = metrics.entry(node_id).or_default();
391        entry.storage_bytes = bytes;
392        entry.bytes = bytes;
393
394        let span_entry = metrics.entry(op.span_id.clone()).or_default();
395        span_entry.bytes = span_entry.bytes.saturating_add(bytes);
396    }
397
398    for live in &profile.peak_breakdown {
399        if let Some(entry) = metrics.get_mut(&live.span_id) {
400            entry.peak_bytes = entry.peak_bytes.max(live.bytes);
401        }
402    }
403
404    for span in &doc.spans {
405        rollup_span_memory(span.id.as_str(), &doc.spans, &mut metrics);
406    }
407
408    metrics
409}
410
411fn rollup_span_memory(
412    span_id: &str,
413    spans: &[super::schema::SpanRecord],
414    metrics: &mut HashMap<String, NodeMemoryMetrics>,
415) {
416    let children: Vec<_> = spans
417        .iter()
418        .filter(|s| s.parent_id.as_deref() == Some(span_id))
419        .collect();
420
421    let mut child_peak = 0u64;
422    for child in &children {
423        rollup_span_memory(&child.id, spans, metrics);
424        if let Some(m) = metrics.get(&child.id) {
425            child_peak = child_peak.max(m.peak_bytes);
426        }
427    }
428
429    let own = metrics.get(span_id).cloned().unwrap_or_default();
430    let entry = metrics.entry(span_id.to_string()).or_default();
431    entry.peak_bytes = entry.peak_bytes.max(own.bytes).max(child_peak);
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437    use crate::trace::events::{MemoryEvent, OpEvent};
438    use crate::trace::schema::{SpanKind, SpanRecord, TraceRunMeta, SCHEMA};
439
440    fn doc_with_ops_and_memory() -> TraceDocument {
441        TraceDocument {
442            schema: SCHEMA.to_string(),
443            run: TraceRunMeta {
444                run_id: "r".into(),
445                correlation_id: "test/update-1".into(),
446                entrypoint: "test".into(),
447                phase: crate::phase::ExecutionPhase::Train,
448                timestamp: "2026-01-01T00:00:00Z".into(),
449                capture_step: 1,
450                warmup_steps: 0,
451                device: "cpu".into(),
452                timing_mode: crate::trace::TimingMode::Host,
453                tags: Default::default(),
454                candle_version: None,
455            },
456            spans: vec![SpanRecord {
457                id: "s1".into(),
458                parent_id: None,
459                name: "forward".into(),
460                kind: SpanKind::Function,
461                measured: true,
462                start_ns: 0,
463                closed: true,
464                duration_ns: 1000,
465                step: None,
466            }],
467            ops: vec![OpEvent {
468                span_id: "s1".into(),
469                op_name: "matmul".into(),
470                inputs: vec![],
471                output: Some("out".into()),
472                shape: vec![100, 100],
473                dtype: "f32".into(),
474                device: "cpu".into(),
475                duration_ns: 500,
476                timestamp_ns: 500,
477                storage_bytes: None,
478                input_storage_bytes: 0,
479            }],
480            tensors: vec![],
481            memory: vec![
482                MemoryEvent {
483                    timestamp_ns: 500,
484                    tensor_id: "out".into(),
485                    span_id: "s1".into(),
486                    op_name: Some("matmul".into()),
487                    device: "cpu".into(),
488                    bytes: 100 * 100 * 4,
489                    action: MemoryAction::Alloc,
490                    shape: vec![100, 100],
491                    dtype: "f32".into(),
492                    category: MemoryCategory::Activation,
493                },
494                MemoryEvent {
495                    timestamp_ns: 1000,
496                    tensor_id: "out".into(),
497                    span_id: "s1".into(),
498                    op_name: Some("matmul".into()),
499                    device: "cpu".into(),
500                    bytes: 100 * 100 * 4,
501                    action: MemoryAction::Free,
502                    shape: vec![100, 100],
503                    dtype: "f32".into(),
504                    category: MemoryCategory::Activation,
505                },
506            ],
507            device_memory: vec![],
508            gradients: vec![],
509            edges: vec![],
510        }
511    }
512
513    #[test]
514    fn explicit_memory_events_compute_peak() {
515        let profile = analyze_memory(&doc_with_ops_and_memory());
516        assert_eq!(profile.summary.peak_bytes, 100 * 100 * 4);
517        assert_eq!(profile.summary.alloc_count, 1);
518        assert_eq!(profile.summary.free_count, 1);
519        assert_eq!(profile.peak_breakdown.len(), 1);
520    }
521
522    #[test]
523    fn does_not_invent_lifetimes_from_storage_metadata() {
524        let mut doc = doc_with_ops_and_memory();
525        doc.memory.clear();
526        let profile = analyze_memory(&doc);
527        assert_eq!(profile.summary.peak_bytes, 0);
528        assert!(profile.timeline.is_empty());
529    }
530}
candle_graph/trace/memory.rs

candle_graph/trace/
memory.rs
