candle_graph/

activation.rs

//! Capability-qualified rankings for activation-producing operations.
//!
//! This module never collapses host time, device time, dense tensor footprint, logical storage,
//! or physical allocator state into one score. Each metric keeps its own qualification.

use std::collections::{BTreeSet, HashMap, HashSet};

use serde::{Deserialize, Deserializer, Serialize, Serializer};

use crate::capability::{CapabilityLevel, CapabilityState, CoverageLevel, EvidenceCapabilities};
use crate::timing::TimingProfile;
use crate::trace::memory::{resolve_dense_tensor_bytes, MemoryCategory, MemoryProfile};
use crate::trace::{SpanKind, TraceDocument};

/// What directly classified one observed operation output as an activation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ActivationClassificationSource {
    TensorOutput,
    LogicalStorage,
}

/// Exact storage space-time cost. It serializes as a decimal string because realistic
/// byte-nanosecond products can exceed JSON's `u64` number range.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct ByteNanoseconds(pub u128);

impl Serialize for ByteNanoseconds {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.0.to_string())
    }
}

impl<'de> Deserialize<'de> for ByteNanoseconds {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = String::deserialize(deserializer)?;
        value
            .parse::<u128>()
            .map(Self)
            .map_err(serde::de::Error::custom)
    }
}

/// Device-busy evidence attributable to a dedicated operation span.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ActivationDeviceTiming {
    pub device: String,
    pub clock_id: String,
    pub backends: Vec<String>,
    pub streams: Vec<String>,
    pub interval_count: usize,
    /// Union of intervals on this device clock; it is never added across clocks.
    pub busy_ns: u64,
}

/// One observed activation-producing operation and its deliberately separate cost metrics.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ActivationOperation {
    /// Stable graph identity within this trace: the dedicated operation span ID when available,
    /// otherwise `{span_id}/op/{op_index}`.
    pub id: String,
    pub span_id: String,
    pub op_index: usize,
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub output_tensor_id: Option<String>,
    pub shape: Vec<usize>,
    pub dtype: String,
    pub device: String,
    /// Direct `OpEvent` host duration. On asynchronous devices this is commonly launch time.
    pub observed_host_duration_ns: u64,
    /// Present only when this event is the sole operation in a dedicated `SpanKind::Op` span.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub device_timings: Vec<ActivationDeviceTiming>,
    /// Dense output footprint (`shape × dtype`), not backing storage or allocator usage.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub dense_output_bytes: Option<u64>,
    /// Unique logical activation storage directly attributable to this operation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub logical_allocated_bytes: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub logical_storage_count: Option<usize>,
    /// Sum of measured-region-clipped live durations across attributed storages.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub logical_storage_live_duration_ns: Option<u64>,
    /// Sum of `storage bytes × measured-region-clipped live duration`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub logical_byte_nanoseconds: Option<ByteNanoseconds>,
    /// Attributed logical storage still live at the measured region's end.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub logical_retained_bytes_at_measured_end: Option<u64>,
    pub classification_sources: Vec<ActivationClassificationSource>,
}

/// Activation hotspot evidence for one profile run.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ActivationProfile {
    pub definition: String,
    pub coverage: CapabilityState,
    pub host_time: CapabilityState,
    pub device_time: CapabilityState,
    pub dense_output: CapabilityState,
    pub logical_memory: CapabilityState,
    /// Physical samples do not carry operation or activation identity.
    pub physical_memory_attribution: CapabilityState,
    pub unattributed_activation_storages: usize,
    pub operations: Vec<ActivationOperation>,
}

impl ActivationProfile {
    pub(crate) fn from_trace(
        document: &TraceDocument,
        timing: &TimingProfile,
        memory: &MemoryProfile,
    ) -> Self {
        let measured_ids = measured_subtree_ids(document);
        let measured_bounds = document
            .spans
            .iter()
            .find(|span| span.measured && span.closed)
            .map(|span| {
                (
                    span.start_ns,
                    span.start_ns.saturating_add(span.duration_ns),
                )
            });
        let activation_tensor_ids = document
            .tensors
            .iter()
            .filter(|tensor| tensor.category == MemoryCategory::Activation)
            .map(|tensor| tensor.tensor_id.as_str())
            .collect::<HashSet<_>>();

        let mut span_op_counts = HashMap::<&str, usize>::new();
        let mut ops_by_span = HashMap::<&str, Vec<usize>>::new();
        let mut op_indices = Vec::with_capacity(document.ops.len());
        for (global_index, op) in document.ops.iter().enumerate() {
            let count = span_op_counts.entry(op.span_id.as_str()).or_default();
            op_indices.push(*count);
            *count += 1;
            ops_by_span
                .entry(op.span_id.as_str())
                .or_default()
                .push(global_index);
        }

        let activation_lifetimes = memory
            .logical
            .as_ref()
            .map(|logical| {
                logical
                    .lifetimes
                    .iter()
                    .filter(|lifetime| {
                        lifetime.category == MemoryCategory::Activation
                            && measured_ids.contains(&lifetime.allocation_span_id)
                    })
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default();
        let mut attributed_lifetimes = HashMap::<usize, Vec<usize>>::new();
        let mut unresolved_spans = HashSet::<&str>::new();
        let mut unattributed_activation_storages = 0usize;
        for (lifetime_index, lifetime) in activation_lifetimes.iter().enumerate() {
            let span_ops = ops_by_span
                .get(lifetime.allocation_span_id.as_str())
                .map(Vec::as_slice)
                .unwrap_or_default();
            let output_candidates = span_ops
                .iter()
                .copied()
                .filter(|&index| {
                    document.ops[index].output.as_ref().is_some_and(|output| {
                        lifetime
                            .tensor_ids
                            .iter()
                            .any(|tensor_id| tensor_id == output)
                    })
                })
                .collect::<Vec<_>>();
            let named_candidates = lifetime
                .op_name
                .as_deref()
                .map(|op_name| {
                    span_ops
                        .iter()
                        .copied()
                        .filter(|&index| document.ops[index].op_name == op_name)
                        .collect::<Vec<_>>()
                })
                .unwrap_or_default();
            let resolved = match (output_candidates.as_slice(), named_candidates.as_slice()) {
                // A unique output-tensor link is the strongest evidence; accept it unless a
                // uniquely named candidate contradicts it.
                ([output], named) if named.is_empty() || named.contains(output) => Some(*output),
                (outputs, [named]) if outputs.is_empty() || outputs.contains(named) => Some(*named),
                _ => None,
            };
            if let Some(op_index) = resolved {
                attributed_lifetimes
                    .entry(op_index)
                    .or_default()
                    .push(lifetime_index);
            } else {
                unresolved_spans.insert(lifetime.allocation_span_id.as_str());
                unattributed_activation_storages += 1;
            }
        }

        let span_kinds = document
            .spans
            .iter()
            .map(|span| (span.id.as_str(), span.kind))
            .collect::<HashMap<_, _>>();
        let mut device_spans_by_id = HashMap::<&str, Vec<&crate::timing::DeviceSpanTiming>>::new();
        for item in &timing.device_spans {
            device_spans_by_id
                .entry(item.span_id.as_str())
                .or_default()
                .push(item);
        }
        let logical_complete =
            document.run.capture_contract.logical_memory == CoverageLevel::Complete;
        let mut operations = Vec::new();
        for (global_index, op) in document.ops.iter().enumerate() {
            if !measured_ids.contains(&op.span_id) {
                continue;
            }
            let tensor_classified = op
                .output
                .as_deref()
                .is_some_and(|output| activation_tensor_ids.contains(output));
            let lifetime_indices = attributed_lifetimes.get(&global_index);
            if !tensor_classified && lifetime_indices.is_none() {
                continue;
            }

            let mut sources = BTreeSet::new();
            if tensor_classified {
                sources.insert(ActivationClassificationSource::TensorOutput);
            }
            if lifetime_indices.is_some() {
                sources.insert(ActivationClassificationSource::LogicalStorage);
            }
            let (
                logical_allocated_bytes,
                logical_storage_count,
                logical_storage_live_duration_ns,
                logical_byte_nanoseconds,
                logical_retained_bytes_at_measured_end,
            ) = match (lifetime_indices, measured_bounds) {
                (Some(indices), Some((measured_start, measured_end))) => {
                    let mut allocated_bytes = 0u64;
                    let mut live_duration_ns = 0u64;
                    let mut byte_nanoseconds = 0u128;
                    let mut retained_bytes = 0u64;
                    for index in indices {
                        let lifetime = activation_lifetimes[*index];
                        allocated_bytes = allocated_bytes.saturating_add(lifetime.bytes);
                        let lifetime_end = lifetime.end_timestamp_ns.unwrap_or(measured_end);
                        let clipped_start = lifetime.start_timestamp_ns.max(measured_start);
                        let clipped_end = lifetime_end.min(measured_end);
                        let duration_ns = clipped_end.saturating_sub(clipped_start);
                        live_duration_ns = live_duration_ns.saturating_add(duration_ns);
                        byte_nanoseconds = byte_nanoseconds.saturating_add(
                            u128::from(lifetime.bytes).saturating_mul(u128::from(duration_ns)),
                        );
                        if lifetime.start_timestamp_ns < measured_end
                            && lifetime
                                .end_timestamp_ns
                                .is_none_or(|end| end > measured_end)
                        {
                            retained_bytes = retained_bytes.saturating_add(lifetime.bytes);
                        }
                    }
                    (
                        Some(allocated_bytes),
                        Some(indices.len()),
                        Some(live_duration_ns),
                        Some(ByteNanoseconds(byte_nanoseconds)),
                        Some(retained_bytes),
                    )
                }
                (Some(indices), None) => (
                    Some(indices.iter().fold(0u64, |total, index| {
                        total.saturating_add(activation_lifetimes[*index].bytes)
                    })),
                    Some(indices.len()),
                    None,
                    None,
                    None,
                ),
                // A confident zero requires complete logical coverage and no unresolved
                // activation storage in this operation's span; otherwise the storage cost is
                // unknown, never silently zero.
                (None, _)
                    if logical_complete && !unresolved_spans.contains(op.span_id.as_str()) =>
                {
                    (Some(0), Some(0), Some(0), Some(ByteNanoseconds(0)), Some(0))
                }
                (None, _) => (None, None, None, None, None),
            };
            let dedicated_op_span = span_kinds.get(op.span_id.as_str()) == Some(&SpanKind::Op)
                && span_op_counts.get(op.span_id.as_str()) == Some(&1);
            let device_timings = if dedicated_op_span {
                device_spans_by_id
                    .get(op.span_id.as_str())
                    .map(Vec::as_slice)
                    .unwrap_or_default()
                    .iter()
                    .map(|item| ActivationDeviceTiming {
                        device: item.device.clone(),
                        clock_id: item.clock_id.clone(),
                        backends: item.backends.clone(),
                        streams: item.streams.clone(),
                        interval_count: item.interval_count,
                        busy_ns: item.busy_ns,
                    })
                    .collect()
            } else {
                Vec::new()
            };
            let op_index = op_indices[global_index];
            operations.push(ActivationOperation {
                id: if dedicated_op_span {
                    op.span_id.clone()
                } else {
                    format!("{}/op/{op_index}", op.span_id)
                },
                span_id: op.span_id.clone(),
                op_index,
                name: op.op_name.clone(),
                output_tensor_id: op.output.clone(),
                shape: op.shape.clone(),
                dtype: op.dtype.clone(),
                device: op.device.clone(),
                observed_host_duration_ns: op.duration_ns,
                device_timings,
                dense_output_bytes: resolve_dense_tensor_bytes(
                    op.output_dense_bytes,
                    &op.shape,
                    &op.dtype,
                ),
                logical_allocated_bytes,
                logical_storage_count,
                logical_storage_live_duration_ns,
                logical_byte_nanoseconds,
                logical_retained_bytes_at_measured_end,
                classification_sources: sources.into_iter().collect(),
            });
        }

        Self {
            definition: "An activation-producing operation is an in-measured-subtree OpEvent whose output is category-linked as activation tensor metadata or uniquely attributable activation logical storage. Host time, per-clock device time, dense footprint, allocated logical bytes, and logical byte-nanoseconds remain separate evidence planes.".into(),
            coverage: CapabilityState::default(),
            host_time: CapabilityState::default(),
            device_time: CapabilityState::default(),
            dense_output: CapabilityState::default(),
            logical_memory: CapabilityState::default(),
            physical_memory_attribution: CapabilityState::unavailable(
                "physical device-memory samples have no operation or activation identity",
            ),
            unattributed_activation_storages,
            operations,
        }
    }

    pub(crate) fn qualify(&mut self, capabilities: &EvidenceCapabilities) {
        self.coverage = capabilities.activation_coverage.clone();
        self.host_time = combine_capabilities(
            "activation host-time ranking",
            &[&self.coverage, &capabilities.nested_host_time],
        );

        self.dense_output = self.coverage.clone();
        let missing_dense = self
            .operations
            .iter()
            .filter(|operation| operation.dense_output_bytes.is_none())
            .count();
        if missing_dense > 0 {
            downgrade(
                &mut self.dense_output,
                format!(
                    "{missing_dense} activation operations have unknown dense output footprint"
                ),
            );
        }

        self.logical_memory = combine_capabilities(
            "activation logical-memory ranking",
            &[&self.coverage, &capabilities.logical_memory_coverage],
        );
        if self.unattributed_activation_storages > 0 {
            downgrade(
                &mut self.logical_memory,
                format!(
                    "{} activation storages could not be attributed uniquely to an operation",
                    self.unattributed_activation_storages
                ),
            );
        }
        let missing_logical = self
            .operations
            .iter()
            .filter(|operation| operation.logical_allocated_bytes.is_none())
            .count();
        if missing_logical > 0 {
            downgrade(
                &mut self.logical_memory,
                format!(
                    "{missing_logical} activation operations have unknown logical allocation bytes"
                ),
            );
        }

        let mut device_base = combine_capabilities(
            "activation device-time ranking",
            &[&self.coverage, &capabilities.nested_device_time],
        );
        if matches!(device_base.level, CapabilityLevel::Invalid) {
            self.device_time = device_base;
            return;
        }
        let device_operations = self
            .operations
            .iter()
            .filter(|operation| !is_host_device(&operation.device))
            .collect::<Vec<_>>();
        if device_operations.is_empty() {
            self.device_time = CapabilityState::unavailable(
                "no activation-producing operations were observed on a non-host device",
            );
            return;
        }
        let attributed = device_operations
            .iter()
            .filter(|operation| !operation.device_timings.is_empty())
            .count();
        if attributed == 0 {
            self.device_time = CapabilityState::unavailable(
                "no activation operation has device timing on a dedicated operation span",
            );
            return;
        }
        if attributed < device_operations.len() {
            downgrade(
                &mut device_base,
                format!(
                    "device timing is attributable for {attributed} of {} non-host activation operations; use one OpEvent per SpanKind::Op span",
                    device_operations.len()
                ),
            );
        }
        self.device_time = device_base;
    }

    /// Largest operation by `key`. Ties resolve to the lexicographically smallest `id`, so every
    /// headline agrees with the query rankings derived from [`ActivationProfile::ranked_by`].
    pub fn top_by<K: Ord>(
        &self,
        key: impl Fn(&ActivationOperation) -> Option<K>,
    ) -> Option<&ActivationOperation> {
        self.ranked_by(key).into_iter().next()
    }

    /// Operations carrying the metric, ranked descending by `key`; ties resolve to the
    /// lexicographically smallest `id`. Operations without the metric are excluded.
    pub fn ranked_by<K: Ord>(
        &self,
        key: impl Fn(&ActivationOperation) -> Option<K>,
    ) -> Vec<&ActivationOperation> {
        let mut rows = self
            .operations
            .iter()
            .filter_map(|operation| key(operation).map(|value| (value, operation)))
            .collect::<Vec<_>>();
        rows.sort_by(|(left_value, left), (right_value, right)| {
            right_value
                .cmp(left_value)
                .then_with(|| left.id.cmp(&right.id))
        });
        rows.into_iter().map(|(_, operation)| operation).collect()
    }
}

fn measured_subtree_ids(document: &TraceDocument) -> HashSet<String> {
    let mut children = HashMap::<&str, Vec<&str>>::new();
    for span in &document.spans {
        if let Some(parent) = span.parent_id.as_deref() {
            children.entry(parent).or_default().push(span.id.as_str());
        }
    }
    let mut ids = HashSet::new();
    let mut queue = document
        .spans
        .iter()
        .filter(|span| span.measured)
        .map(|span| span.id.as_str())
        .collect::<Vec<_>>();
    while let Some(id) = queue.pop() {
        if ids.insert(id.to_owned()) {
            if let Some(child_ids) = children.get(id) {
                queue.extend(child_ids.iter().copied());
            }
        }
    }
    ids
}

fn is_host_device(device: &str) -> bool {
    device.eq_ignore_ascii_case("cpu")
        || device
            .get(..4)
            .is_some_and(|prefix| prefix.eq_ignore_ascii_case("cpu:"))
}

fn combine_capabilities(label: &str, states: &[&CapabilityState]) -> CapabilityState {
    if let Some(invalid) = states
        .iter()
        .find(|state| state.level == CapabilityLevel::Invalid)
    {
        return CapabilityState::invalid(
            invalid.source.clone(),
            format!("{label} is invalid: {}", invalid.reason),
        );
    }
    if let Some(unavailable) = states
        .iter()
        .find(|state| state.level == CapabilityLevel::Unavailable)
    {
        return CapabilityState::unavailable(format!(
            "{label} is unavailable: {}",
            unavailable.reason
        ));
    }
    let coverage = if states
        .iter()
        .all(|state| state.level == CapabilityLevel::Complete)
    {
        CoverageLevel::Complete
    } else {
        CoverageLevel::Partial
    };
    CapabilityState::from_coverage(
        coverage,
        states
            .iter()
            .map(|state| state.source.as_str())
            .collect::<Vec<_>>()
            .join(" + "),
        format!("{label} keeps its required evidence planes separate"),
    )
}

fn downgrade(state: &mut CapabilityState, reason: String) {
    match state.level {
        CapabilityLevel::Complete => state.level = CapabilityLevel::Partial,
        CapabilityLevel::Invalid | CapabilityLevel::Unavailable | CapabilityLevel::Partial => {}
    }
    state.reason = format!("{}; {reason}", state.reason);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::capability::{CaptureContract, MeasurementScope};
    use crate::evidence::EvidencePacket;
    use crate::nsight::NsightEvidence;
    use crate::phase::ExecutionPhase;
    use crate::trace::{
        DeviceIntervalEvent, MemoryAction, MemoryEvent, OpEvent, RunOutcome, SpanRecord,
        TensorEvent, TerminalEvent, TimingMode, TraceRunMeta, SCHEMA,
    };

    fn activation_document() -> TraceDocument {
        let span = |id: &str, parent_id: Option<&str>, name: &str, start_ns, duration_ns, kind| {
            SpanRecord {
                id: id.into(),
                parent_id: parent_id.map(str::to_owned),
                name: name.into(),
                kind,
                measured: parent_id.is_none(),
                start_ns,
                closed: true,
                duration_ns,
                step: None,
            }
        };
        let op = |span_id: &str,
                  name: &str,
                  output: &str,
                  elements: usize,
                  duration_ns: u64,
                  timestamp_ns: u64| OpEvent {
            span_id: span_id.into(),
            op_name: name.into(),
            inputs: Vec::new(),
            output: Some(output.into()),
            shape: vec![elements],
            dtype: "f32".into(),
            device: "cuda:0".into(),
            duration_ns,
            timestamp_ns,
            output_dense_bytes: None,
            input_dense_bytes: 0,
        };
        let tensor = |span_id: &str, tensor_id: &str, elements: usize| TensorEvent {
            span_id: span_id.into(),
            tensor_id: tensor_id.into(),
            label: None,
            shape: vec![elements],
            dtype: "f32".into(),
            device: "cuda:0".into(),
            requires_grad: false,
            dense_bytes: None,
            category: MemoryCategory::Activation,
        };
        let memory =
            |timestamp_ns, span_id: &str, op_name: &str, tensor_id: &str, bytes, action| {
                MemoryEvent {
                    timestamp_ns,
                    storage_id: format!("storage-{tensor_id}"),
                    tensor_id: tensor_id.into(),
                    span_id: span_id.into(),
                    op_name: Some(op_name.into()),
                    device: "cuda:0".into(),
                    bytes,
                    action,
                    shape: vec![bytes as usize],
                    dtype: "u8".into(),
                    category: MemoryCategory::Activation,
                }
            };

        TraceDocument {
            schema: SCHEMA.into(),
            run: TraceRunMeta {
                run_id: "activation-hotspots".into(),
                correlation_id: "infer/activation-hotspots".into(),
                entrypoint: "model::forward".into(),
                phase: ExecutionPhase::Infer,
                timestamp: "2026-08-30T00:00:00Z".into(),
                capture_step: 2,
                warmup_steps: 1,
                device: "cuda:0".into(),
                measured_region_device_synchronized: true,
                timing_mode: TimingMode::Host,
                capture_contract: CaptureContract {
                    measurement_scope: MeasurementScope::ProfiledWork,
                    operations: CoverageLevel::Complete,
                    activations: CoverageLevel::Complete,
                    tensors: CoverageLevel::Partial,
                    logical_memory: CoverageLevel::Complete,
                    device_timing: CoverageLevel::Complete,
                    ..CaptureContract::default()
                },
                comparison_identity: None,
                tags: Default::default(),
                candle_version: None,
            },
            spans: vec![
                span("root", None, "forward", 0, 1_000, SpanKind::Function),
                span("a", Some("root"), "wide", 10, 400, SpanKind::Op),
                span("b", Some("root"), "slow-gpu", 500, 400, SpanKind::Op),
            ],
            ops: vec![
                op("a", "wide", "ta", 100, 300, 10),
                op("b", "slow-gpu", "tb", 200, 100, 500),
            ],
            tensors: vec![tensor("a", "ta", 100), tensor("b", "tb", 200)],
            tensor_stats: Vec::new(),
            memory: vec![
                memory(20, "a", "wide", "ta", 1_000, MemoryAction::Alloc),
                memory(400, "a", "wide", "ta", 1_000, MemoryAction::Free),
                memory(520, "b", "slow-gpu", "tb", 500, MemoryAction::Alloc),
                memory(800, "b", "slow-gpu", "tb", 500, MemoryAction::Free),
            ],
            device_memory: Vec::new(),
            device_intervals: vec![
                DeviceIntervalEvent {
                    span_id: "a".into(),
                    device: "cuda:0".into(),
                    stream_id: "0".into(),
                    clock_id: "cuda-event".into(),
                    backend: "cuda-event".into(),
                    start_ns: 0,
                    duration_ns: 50,
                },
                DeviceIntervalEvent {
                    span_id: "b".into(),
                    device: "cuda:0".into(),
                    stream_id: "0".into(),
                    clock_id: "cuda-event".into(),
                    backend: "cuda-event".into(),
                    start_ns: 100,
                    duration_ns: 200,
                },
            ],
            gradients: Vec::new(),
            edges: Vec::new(),
            terminal: TerminalEvent {
                outcome: RunOutcome::Complete,
                timestamp_ns: 1_000,
                reason: None,
            },
        }
    }

    #[test]
    fn activation_metrics_remain_separate_and_capability_qualified() {
        let packet = EvidencePacket::from_document(
            activation_document(),
            NsightEvidence::unavailable("not captured"),
        )
        .unwrap();
        let profile = &packet.activations;
        assert_eq!(profile.coverage.level, CapabilityLevel::Complete);
        assert_eq!(profile.host_time.level, CapabilityLevel::Complete);
        assert_eq!(profile.device_time.level, CapabilityLevel::Complete);
        assert_eq!(profile.logical_memory.level, CapabilityLevel::Complete);
        assert_eq!(profile.operations.len(), 2);

        let wide = profile
            .operations
            .iter()
            .find(|operation| operation.name == "wide")
            .unwrap();
        assert_eq!(wide.observed_host_duration_ns, 300);
        assert_eq!(wide.device_timings[0].busy_ns, 50);
        assert_eq!(wide.dense_output_bytes, Some(400));
        assert_eq!(wide.logical_allocated_bytes, Some(1_000));
        assert_eq!(wide.logical_storage_live_duration_ns, Some(380));
        assert_eq!(
            wide.logical_byte_nanoseconds,
            Some(ByteNanoseconds(380_000))
        );

        let slow_gpu = profile
            .operations
            .iter()
            .find(|operation| operation.name == "slow-gpu")
            .unwrap();
        assert_eq!(slow_gpu.observed_host_duration_ns, 100);
        assert_eq!(slow_gpu.device_timings[0].busy_ns, 200);
        assert_eq!(slow_gpu.dense_output_bytes, Some(800));
        assert_eq!(slow_gpu.logical_allocated_bytes, Some(500));
        assert_eq!(
            slow_gpu.logical_byte_nanoseconds,
            Some(ByteNanoseconds(140_000))
        );
        assert_eq!(
            profile.physical_memory_attribution.level,
            CapabilityLevel::Unavailable
        );
    }

    fn push_duplicate_named_ops(document: &mut TraceDocument, storage_tensor_id: &str) {
        document.spans.push(SpanRecord {
            id: "c".into(),
            parent_id: Some("root".into()),
            name: "dup".into(),
            kind: SpanKind::Function,
            measured: false,
            start_ns: 900,
            closed: true,
            duration_ns: 80,
            step: None,
        });
        for (output, timestamp_ns) in [("tc1", 900), ("tc2", 910)] {
            document.ops.push(OpEvent {
                span_id: "c".into(),
                op_name: "dup".into(),
                inputs: Vec::new(),
                output: Some(output.into()),
                shape: vec![10],
                dtype: "f32".into(),
                device: "cuda:0".into(),
                duration_ns: 5,
                timestamp_ns,
                output_dense_bytes: None,
                input_dense_bytes: 0,
            });
            document.tensors.push(TensorEvent {
                span_id: "c".into(),
                tensor_id: output.into(),
                label: None,
                shape: vec![10],
                dtype: "f32".into(),
                device: "cuda:0".into(),
                requires_grad: false,
                dense_bytes: None,
                category: MemoryCategory::Activation,
            });
        }
        for (timestamp_ns, action) in [(905, MemoryAction::Alloc), (940, MemoryAction::Free)] {
            document.memory.push(MemoryEvent {
                timestamp_ns,
                storage_id: format!("storage-{storage_tensor_id}"),
                tensor_id: storage_tensor_id.into(),
                span_id: "c".into(),
                op_name: Some("dup".into()),
                device: "cuda:0".into(),
                bytes: 2_000,
                action,
                shape: vec![2_000],
                dtype: "u8".into(),
                category: MemoryCategory::Activation,
            });
        }
    }

    #[test]
    fn unique_output_link_attributes_storage_despite_ambiguous_op_names() {
        let mut document = activation_document();
        // The storage's tensor id uniquely identifies the first `dup` op even though the
        // recorded op name matches both ops in the span.
        push_duplicate_named_ops(&mut document, "tc1");
        let packet =
            EvidencePacket::from_document(document, NsightEvidence::unavailable("not captured"))
                .unwrap();
        assert_eq!(packet.activations.unattributed_activation_storages, 0);
        let first_dup = packet
            .activations
            .operations
            .iter()
            .find(|operation| operation.output_tensor_id.as_deref() == Some("tc1"))
            .unwrap();
        assert_eq!(first_dup.logical_allocated_bytes, Some(2_000));
    }

    #[test]
    fn ambiguous_activation_storage_never_reports_a_confident_zero() {
        let mut document = activation_document();
        // The storage links to neither op output and the op name matches both ops, so it is
        // unresolvable; the ops in that span must report unknown, not zero.
        push_duplicate_named_ops(&mut document, "tc-other");
        let packet =
            EvidencePacket::from_document(document, NsightEvidence::unavailable("not captured"))
                .unwrap();
        assert_eq!(packet.activations.unattributed_activation_storages, 1);
        for output in ["tc1", "tc2"] {
            let operation = packet
                .activations
                .operations
                .iter()
                .find(|operation| operation.output_tensor_id.as_deref() == Some(output))
                .unwrap();
            assert!(operation.logical_allocated_bytes.is_none());
            assert!(operation.logical_byte_nanoseconds.is_none());
        }
        assert_eq!(
            packet.activations.logical_memory.level,
            CapabilityLevel::Partial
        );
    }

    #[test]
    fn an_unclosed_measured_span_reports_unknown_not_zero_time_metrics() {
        let mut document = activation_document();
        document.spans[0].closed = false;
        let packet =
            EvidencePacket::from_document(document, NsightEvidence::unavailable("not captured"))
                .unwrap();
        let wide = packet
            .activations
            .operations
            .iter()
            .find(|operation| operation.name == "wide")
            .unwrap();
        assert_eq!(wide.logical_allocated_bytes, Some(1_000));
        assert!(wide.logical_storage_live_duration_ns.is_none());
        assert!(wide.logical_byte_nanoseconds.is_none());
        assert!(wide.logical_retained_bytes_at_measured_end.is_none());
    }

    #[test]
    fn complete_activation_contract_requires_operations_and_a_category_plane() {
        let missing_operations = CaptureContract {
            activations: CoverageLevel::Complete,
            tensors: CoverageLevel::Partial,
            ..CaptureContract::default()
        };
        assert!(missing_operations.validate().is_err());

        let missing_categories = CaptureContract {
            operations: CoverageLevel::Complete,
            activations: CoverageLevel::Complete,
            ..CaptureContract::default()
        };
        assert!(missing_categories.validate().is_err());
    }

    #[test]
    fn activation_profile_excludes_operations_outside_the_measured_subtree() {
        let mut document = activation_document();
        document.spans[0].measured = false;
        document.spans[0].parent_id = Some("session".into());
        document.spans.insert(
            0,
            SpanRecord {
                id: "session".into(),
                parent_id: None,
                name: "session".into(),
                kind: SpanKind::Function,
                measured: false,
                start_ns: 0,
                closed: true,
                duration_ns: 2_000,
                step: None,
            },
        );
        document.spans[1].measured = true;
        document.spans.push(SpanRecord {
            id: "setup".into(),
            parent_id: Some("session".into()),
            name: "setup-activation".into(),
            kind: SpanKind::Op,
            measured: false,
            start_ns: 1_200,
            closed: true,
            duration_ns: 100,
            step: None,
        });
        document.ops.push(OpEvent {
            span_id: "setup".into(),
            op_name: "setup-activation".into(),
            inputs: Vec::new(),
            output: Some("setup-output".into()),
            shape: vec![10_000],
            dtype: "f32".into(),
            device: "cuda:0".into(),
            duration_ns: 90,
            timestamp_ns: 1_200,
            output_dense_bytes: None,
            input_dense_bytes: 0,
        });
        document.tensors.push(TensorEvent {
            span_id: "setup".into(),
            tensor_id: "setup-output".into(),
            label: None,
            shape: vec![10_000],
            dtype: "f32".into(),
            device: "cuda:0".into(),
            requires_grad: false,
            dense_bytes: None,
            category: MemoryCategory::Activation,
        });
        document.terminal.timestamp_ns = 2_000;

        let packet =
            EvidencePacket::from_document(document, NsightEvidence::unavailable("not captured"))
                .unwrap();
        assert_eq!(packet.activations.operations.len(), 2);
        assert!(packet
            .activations
            .operations
            .iter()
            .all(|operation| operation.name != "setup-activation"));
    }
}