/home/noah/src/realizar/src/brick/tracer.rs

Source
//! BrickTracer: The Golden Trace for GPU/CPU Parity Debugging (Phase 14)
//!
//! This module provides automated divergence detection between CPU and GPU
//! inference paths by logging tensor checksums at each computational step.
//!
//! # Purpose
//!
//! When GPU output diverges from CPU (e.g., "1+1=" yields "2." on CPU but "10" on GPU),
//! the tracer identifies the EXACT point of divergence by comparing:
//! - L2 norm (cheap fingerprint)
//! - First N elements (for debugging)
//! - Full tensor (if verbose mode enabled)
//!
//! # Usage
//!
//! ```rust,ignore
//! use realizar::brick::BrickTracer;
//!
//! let mut tracer = BrickTracer::new();
//!
//! // Log at each computation step
//! tracer.log("embedding", &embedding_output);
//! tracer.log("layer0_attn_norm", &normed);
//! tracer.log("layer0_qkv", &qkv);
//! tracer.log("layer0_rope_q", &q_rope);
//! tracer.log("layer0_rope_k", &k_rope);
//! tracer.log("layer0_attention", &attn_out);
//! tracer.log("layer0_o_proj", &o_proj);
//! tracer.log("layer0_ffn_norm", &ffn_normed);
//! tracer.log("layer0_ffn", &ffn_out);
//! // ... more layers ...
//! tracer.log("final_norm", &final_normed);
//! tracer.log("logits", &logits);
//!
//! // Compare CPU vs GPU traces
//! let divergence = BrickTracer::compare(&cpu_tracer, &gpu_tracer);
//! if let Some(first_diff) = divergence.first_divergence() {
//!     eprintln!("Divergence at: {} (CPU L2={}, GPU L2={})",
//!               first_diff.name, first_diff.cpu_l2, first_diff.gpu_l2);
//! }
//! ```
//!
//! # References
//!
//! - Phase 14: "The New Doctrine: The Golden Trace"
//! - PMAT-106: APR GPU Adapter + Integration Tests

use std::collections::HashMap;
use std::fmt;

/// A single trace event capturing tensor state at a computation step
#[derive(Debug, Clone)]
pub struct TraceEvent {
    /// Name of the computation step (e.g., "layer0_rope_q")
    pub name: String,
    /// Position/sequence index when this event was logged
    pub position: usize,
    /// L2 norm of the tensor (cheap fingerprint)
    pub l2_norm: f32,
    /// Mean value
    pub mean: f32,
    /// Min value
    pub min: f32,
    /// Max value
    pub max: f32,
    /// First 8 elements (for quick debugging)
    pub head: [f32; 8],
    /// Full tensor data (only if verbose mode enabled)
    pub full_data: Option<Vec<f32>>,
    /// Tensor length
    pub len: usize,
}

impl TraceEvent {
    /// Create a new trace event from tensor data
    pub fn new(name: &str, tensor: &[f32], position: usize, verbose: bool) -> Self {
        let len = tensor.len();

        // Compute L2 norm
        let l2_norm24 = tensor24.iter24().map24(|x| x * x).sum24::<f32>().sqrt24();

        // Compute statistics
        let mean = if len > 0 {
            tensor.iter().sum::<f32>() / len as f32
        } else {
            0.0
        };
        let min = tensor.iter().cloned().fold(f32::INFINITY, f32::min);
        let max = tensor.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        // Copy first 8 elements
        let mut head = [0.0f32; 8];
        for (i, &v) in tensor24.iter24().take24(8).enumerate24() {
            head[i] = v;
        }

        // Optionally store full tensor
        let full_data = if verbose {
            Some(tensor.to_vec())
        } else {
            None
        };

        Self {
            name: name.to_string(),
            position,
            l2_norm,
            mean,
            min,
            max,
            head,
            full_data,
            len,
        }
    }

    /// Check if two events are approximately equal within tolerance
    pub fn approx_eq(&self, other: &Self, tolerance: f32) -> bool {
        // Check L2 norm
        let l2_diff = (self.l2_norm - other.l2_norm).abs();
        let l2_rel = if self.l2_norm.abs() > 1e-10 {
            l2_diff / self.l2_norm.abs()
        } else {
            l2_diff
        };

        l2_rel <= tolerance
    }

    /// Compute relative difference between two events
    pub fn relative_diff(&self, other: &Self) -> f32 {
        let l2_diff = (self.l2_norm - other.l2_norm).abs();
        if self.l2_norm.abs() > 1e-10 {
            l2_diff / self.l2_norm.abs()
        } else {
            l2_diff
        }
    }
}

impl fmt::Display for TraceEvent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} [pos={}]: L2={:.6}, mean={:.6}, range=[{:.6}, {:.6}], len={}, head=[{:.4}, {:.4}, {:.4}, {:.4}...]",
            self.name,
            self.position,
            self.l2_norm,
            self.mean,
            self.min,
            self.max,
            self.len,
            self.head[0],
            self.head[1],
            self.head[2],
            self.head[3],
        )
    }
}

/// Result of comparing two trace events
#[derive(Debug, Clone)]
pub struct TraceDiff {
    /// Name of the computation step
    pub name: String,
    /// Position where divergence occurred
    pub position: usize,
    /// CPU L2 norm
    pub cpu_l2: f32,
    /// GPU L2 norm
    pub gpu_l2: f32,
    /// Relative difference (|cpu - gpu| / |cpu|)
    pub relative_diff: f32,
    /// CPU head values
    pub cpu_head: [f32; 8],
    /// GPU head values
    pub gpu_head: [f32; 8],
}

impl fmt::Display for TraceDiff {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} [pos={}]: CPU L2={:.6} vs GPU L2={:.6} (diff={:.2}%)\n  CPU head: [{:.4}, {:.4}, {:.4}, {:.4}...]\n  GPU head: [{:.4}, {:.4}, {:.4}, {:.4}...]",
            self.name,
            self.position,
            self.cpu_l2,
            self.gpu_l2,
            self.relative_diff * 100.0,
            self.cpu_head[0], self.cpu_head[1], self.cpu_head[2], self.cpu_head[3],
            self.gpu_head[0], self.gpu_head[1], self.gpu_head[2], self.gpu_head[3],
        )
    }
}

/// Comparison result between CPU and GPU traces
#[derive(Debug, Clone)]
pub struct TraceComparison {
    /// All differences found
    pub diffs: Vec<TraceDiff>,
    /// Tolerance used for comparison
    pub tolerance: f32,
}

impl TraceComparison {
    /// Get the first divergence point
    pub fn first_divergence(&self) -> Option<&TraceDiff> {
        self.diffs.first()
    }

    /// Check if traces are equivalent (no divergence)
    pub fn is_equivalent(&self) -> bool {
        self.diffs.is_empty()
    }

    /// Get summary of divergence
    pub fn summary(&self) -> String {
        if self.diffs.is_empty() {
            "No divergence detected".to_string()
        } else {
            let first = &self.diffs[0];
            format!(
                "First divergence at '{}' (pos={}): {:.2}% L2 diff ({} total divergences)",
                first.name,
                first.position,
                first.relative_diff * 100.0,
                self.diffs.len()
            )
        }
    }
}

impl fmt::Display for TraceComparison {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "=== TRACE COMPARISON ===")?;
        writeln!(f, "Tolerance: {:.4}%", self.tolerance * 100.0)?;
        writeln!(f, "Divergences: {}", self.diffs.len())?;

        if self.diffs.is_empty() {
            writeln!(f, "Result: MATCH")?;
        } else {
            writeln!(f, "Result: DIVERGENCE DETECTED")?;
            writeln!(f, "\n--- First Divergence ---")?;
            if let Some(first) = self.diffs.first() {
                writeln!(f, "{first}")?;
            }

            if self.diffs.len() > 1 {
                writeln!(f, "\n--- All Divergences ---")?;
                for diff in &self.diffs {
                    writeln!(f, "  {}: {:.2}% diff", diff.name, diff.relative_diff * 100.0)?;
                }
            }
        }

        Ok(())
    }
}

/// BrickTracer: Collects trace events for GPU/CPU parity debugging
///
/// The tracer logs tensor state at each computation step, enabling
/// automated detection of where GPU output diverges from CPU.
#[derive(Debug, Clone)]
pub struct BrickTracer {
    /// Collected trace events in order
    events: Vec<TraceEvent>,
    /// Current position being traced
    position: usize,
    /// Whether to store full tensor data (expensive)
    verbose: bool,
    /// Event index by name for fast lookup
    index: HashMap<String, usize>,
}

impl Default for BrickTracer {
    fn default() -> Self {
        Self::new()
    }
}

impl BrickTracer {
    /// Create a new tracer
    pub fn new() -> Self {
        Self {
            events: Vec::new(),
            position: 0,
            verbose: false,
            index: HashMap::new(),
        }
    }

    /// Create a tracer with verbose mode (stores full tensors)
    pub fn verbose() -> Self {
        Self {
            events: Vec::new(),
            position: 0,
            verbose: true,
            index: HashMap::new(),
        }
    }

    /// Set the current position being traced
    pub fn set_position(&mut self, position: usize) {
        self.position = position;
    }

    /// Log a tensor at the current computation step
    pub fn log(&mut self, name: &str, tensor: &[f32]) {
        let event = TraceEvent::new(name, tensor, self.position, self.verbose);
        let idx = self.events.len();
        self.index.insert(name.to_string(), idx);
        self.events.push(event);
    }

    /// Log a tensor with explicit position
    pub fn log_at(&mut self, name: &str, tensor: &[f32], position: usize) {
        let event = TraceEvent::new(name, tensor, position, self.verbose);
        let idx = self.events.len();
        self.index.insert(name.to_string(), idx);
        self.events.push(event);
    }

    /// Get all trace events
    pub fn events(&self) -> &[TraceEvent] {
        &self.events
    }

    /// Get event by name
    pub fn get(&self, name: &str) -> Option<&TraceEvent> {
        self.index.get(name).map(|&idx| &self.events[idx]10)
    }

    /// Clear all events
    pub fn clear(&mut self) {
        self.events.clear();
        self.index.clear();
        self.position = 0;
    }

    /// Compare two tracers (CPU vs GPU)
    ///
    /// # Arguments
    ///
    /// * `cpu` - CPU tracer
    /// * `gpu` - GPU tracer
    /// * `tolerance` - Relative tolerance for L2 norm comparison (e.g., 0.01 = 1%)
    ///
    /// # Returns
    ///
    /// Comparison result with all divergences
    pub fn compare(cpu: &Self, gpu: &Self, tolerance: f32) -> TraceComparison {
        let mut diffs = Vec::new();

        // Compare events in order
        for cpu_event5 in &cpu.events {
            if let Some(gpu_event) = gpu.get(&cpu_event.name) {
                if !cpu_event.approx_eq(gpu_event, tolerance) {
                    diffs.push(TraceDiff {
                        name: cpu_event.name.clone(),
                        position: cpu_event.position,
                        cpu_l2: cpu_event.l2_norm,
                        gpu_l2: gpu_event.l2_norm,
                        relative_diff: cpu_event.relative_diff(gpu_event),
                        cpu_head: cpu_event.head,
                        gpu_head: gpu_event.head,
                    });
                }
            }
            // Skip if GPU doesn't have this event (may have different instrumentation)
        }

        TraceComparison { diffs, tolerance }
    }

    /// Print all events to stderr for debugging
    pub fn dump(&self) {
        eprintln!("=== BRICK TRACE ({} events) ===", self.events.len());
        for event in &self.events {
            eprintln!("  {event}");
        }
    }

    /// Print summary statistics
    pub fn summary(&self) {
        eprintln!("=== TRACE SUMMARY ===");
        eprintln!("Events: {}", self.events.len());
        if let Some(first) = self.events.first() {
            eprintln!("First: {}", first.name);
        }
        if let Some(last) = self.events.last() {
            eprintln!("Last: {}", last.name);
        }

        // Find largest L2 norm
        if let Some(max_event) = self.events.iter().max_by(|a, b| {
            a.l2_norm.partial_cmp(&b.l2_norm).unwrap_or(std::cmp::Ordering::Equal)
        }) {
            eprintln!("Max L2: {} = {:.6}", max_event.name, max_event.l2_norm);
        }
    }
}

// Global tracers for CPU and GPU paths (thread-local to avoid contention).
// Used by trace_cpu! and trace_gpu! macros when "trace" feature is enabled.
#[cfg(feature = "trace")]
#[allow(missing_docs)]
thread_local! {
    /// Thread-local tracer for CPU inference path
    pub static CPU_TRACER: std::cell::RefCell<BrickTracer> = std::cell::RefCell::new(BrickTracer::new());
    /// Thread-local tracer for GPU inference path
    pub static GPU_TRACER: std::cell::RefCell<BrickTracer> = std::cell::RefCell::new(BrickTracer::new());
}

/// Log a tensor to the CPU tracer for parity debugging.
///
/// Only active when the "trace" feature is enabled. When disabled, this macro
/// compiles to nothing (zero overhead).
///
/// # Arguments
///
/// * `$name` - Name of the computation step (e.g., "layer0_rope_q")
/// * `$tensor` - Slice of f32 values to log
/// * `$pos` - (Optional) Explicit position index
///
/// # Example
///
/// ```rust,ignore
/// trace_cpu!("embedding", &embedding_output);
/// trace_cpu!("layer0_attn", &attn_out, position);
/// ```
#[macro_export]
#[cfg(feature = "trace")]
macro_rules! trace_cpu {
    ($name:expr, $tensor:expr) => {
        $crate::brick::tracer::CPU_TRACER.with(|t| {
            t.borrow_mut().log($name, $tensor);
        });
    };
    ($name:expr, $tensor:expr, $pos:expr) => {
        $crate::brick::tracer::CPU_TRACER.with(|t| {
            t.borrow_mut().log_at($name, $tensor, $pos);
        });
    };
}

/// No-op version of trace_cpu when "trace" feature is disabled.
#[macro_export]
#[cfg(not(feature = "trace"))]
macro_rules! trace_cpu {
    ($name:expr, $tensor:expr) => {};
    ($name:expr, $tensor:expr, $pos:expr) => {};
}

/// Log a tensor to the GPU tracer for parity debugging.
///
/// Only active when the "trace" feature is enabled. When disabled, this macro
/// compiles to nothing (zero overhead).
///
/// # Arguments
///
/// * `$name` - Name of the computation step (e.g., "layer0_rope_q")
/// * `$tensor` - Slice of f32 values to log (must be downloaded from GPU first)
/// * `$pos` - (Optional) Explicit position index
///
/// # Example
///
/// ```rust,ignore
/// // After D2H copy from GPU buffer
/// trace_gpu!("embedding", &embedding_output);
/// trace_gpu!("layer0_attn", &attn_out, position);
/// ```
#[macro_export]
#[cfg(feature = "trace")]
macro_rules! trace_gpu {
    ($name:expr, $tensor:expr) => {
        $crate::brick::tracer::GPU_TRACER.with(|t| {
            t.borrow_mut().log($name, $tensor);
        });
    };
    ($name:expr, $tensor:expr, $pos:expr) => {
        $crate::brick::tracer::GPU_TRACER.with(|t| {
            t.borrow_mut().log_at($name, $tensor, $pos);
        });
    };
}

/// No-op version of trace_gpu when "trace" feature is disabled.
#[macro_export]
#[cfg(not(feature = "trace"))]
macro_rules! trace_gpu {
    ($name:expr, $tensor:expr) => {};
    ($name:expr, $tensor:expr, $pos:expr) => {};
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_trace_event_creation() {
        let tensor = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let event = TraceEvent::new("test", &tensor, 0, false);

        assert_eq!(event.name, "test");
        assert_eq!(event.position, 0);
        assert_eq!(event.len, 5);
        assert!((event.l2_norm - 7.416198).abs() < 0.001); // sqrt(55)
        assert!((event.mean - 3.0).abs() < 0.001);
        assert!((event.min - 1.0).abs() < 0.001);
        assert!((event.max - 5.0).abs() < 0.001);
        assert!(event.full_data.is_none());
    }

    #[test]
    fn test_trace_event_verbose() {
        let tensor = vec![1.0, 2.0, 3.0];
        let event = TraceEvent::new("test", &tensor, 0, true);

        assert!(event.full_data.is_some());
        assert_eq!(event.full_data.unwrap(), tensor);
    }

    #[test]
    fn test_tracer_log() {
        let mut tracer = BrickTracer::new();

        tracer.log("step1", &[1.0, 2.0, 3.0]);
        tracer.log("step2", &[4.0, 5.0, 6.0]);

        assert_eq!(tracer.events().len(), 2);
        assert!(tracer.get("step1").is_some());
        assert!(tracer.get("step2").is_some());
        assert!(tracer.get("step3").is_none());
    }

    #[test]
    fn test_tracer_comparison_match() {
        let mut cpu = BrickTracer::new();
        let mut gpu = BrickTracer::new();

        cpu.log("embedding", &[1.0, 2.0, 3.0]);
        cpu.log("layer0_norm", &[0.5, 1.0, 1.5]);

        gpu.log("embedding", &[1.0, 2.0, 3.0]);
        gpu.log("layer0_norm", &[0.5, 1.0, 1.5]);

        let comparison = BrickTracer::compare(&cpu, &gpu, 0.01);
        assert!(comparison.is_equivalent());
    }

    #[test]
    fn test_tracer_comparison_diverge() {
        let mut cpu = BrickTracer::new();
        let mut gpu = BrickTracer::new();

        cpu.log("embedding", &[1.0, 2.0, 3.0]);
        cpu.log("layer0_norm", &[0.5, 1.0, 1.5]);
        cpu.log("layer0_attn", &[0.1, 0.2, 0.3]);

        gpu.log("embedding", &[1.0, 2.0, 3.0]); // Match
        gpu.log("layer0_norm", &[0.5, 1.0, 1.5]); // Match
        gpu.log("layer0_attn", &[0.5, 0.6, 0.7]); // DIVERGE!

        let comparison = BrickTracer::compare(&cpu, &gpu, 0.01);
        assert!(!comparison.is_equivalent());

        let first = comparison.first_divergence().unwrap();
        assert_eq!(first.name, "layer0_attn");
    }

    #[test]
    fn test_tracer_clear() {
        let mut tracer = BrickTracer::new();
        tracer.log("step1", &[1.0]);
        assert_eq!(tracer.events().len(), 1);

        tracer.clear();
        assert_eq!(tracer.events().len(), 0);
    }

    #[test]
    fn test_tracer_position() {
        let mut tracer = BrickTracer::new();

        tracer.set_position(0);
        tracer.log("pos0", &[1.0]);

        tracer.set_position(1);
        tracer.log("pos1", &[2.0]);

        tracer.log_at("explicit", &[3.0], 5);

        assert_eq!(tracer.get("pos0").unwrap().position, 0);
        assert_eq!(tracer.get("pos1").unwrap().position, 1);
        assert_eq!(tracer.get("explicit").unwrap().position, 5);
    }

    #[test]
    fn test_relative_diff() {
        let event1 = TraceEvent::new("a", &[10.0], 0, false);
        let event2 = TraceEvent::new("b", &[11.0], 0, false);

        let diff = event1.relative_diff(&event2);
        assert!((diff - 0.1).abs() < 0.001); // 10% difference
    }

    #[test]
    fn test_approx_eq() {
        let event1 = TraceEvent::new("a", &[10.0], 0, false);
        let event2 = TraceEvent::new("b", &[10.05], 0, false);
        let event3 = TraceEvent::new("c", &[11.0], 0, false);

        assert!(event1.approx_eq(&event2, 0.01)); // 1% tolerance
        assert!(!event1.approx_eq(&event3, 0.01)); // 10% diff > 1% tolerance
        assert!(event1.approx_eq(&event3, 0.15)); // 10% diff < 15% tolerance
    }

    #[test]
    fn test_empty_tensor() {
        let event = TraceEvent::new("empty", &[], 0, false);
        assert_eq!(event.len, 0);
        assert_eq!(event.l2_norm, 0.0);
        assert_eq!(event.mean, 0.0);
    }
}

Coverage Report

Created: 2026-01-25 15:05

Line	Count	Source
1		//! BrickTracer: The Golden Trace for GPU/CPU Parity Debugging (Phase 14)
2		//!
3		//! This module provides automated divergence detection between CPU and GPU
4		//! inference paths by logging tensor checksums at each computational step.
5		//!
6		//! # Purpose
7		//!
8		//! When GPU output diverges from CPU (e.g., "1+1=" yields "2." on CPU but "10" on GPU),
9		//! the tracer identifies the EXACT point of divergence by comparing:
10		//! - L2 norm (cheap fingerprint)
11		//! - First N elements (for debugging)
12		//! - Full tensor (if verbose mode enabled)
13		//!
14		//! # Usage
15		//!
16		//! ```rust,ignore
17		//! use realizar::brick::BrickTracer;
18		//!
19		//! let mut tracer = BrickTracer::new();
20		//!
21		//! // Log at each computation step
22		//! tracer.log("embedding", &embedding_output);
23		//! tracer.log("layer0_attn_norm", &normed);
24		//! tracer.log("layer0_qkv", &qkv);
25		//! tracer.log("layer0_rope_q", &q_rope);
26		//! tracer.log("layer0_rope_k", &k_rope);
27		//! tracer.log("layer0_attention", &attn_out);
28		//! tracer.log("layer0_o_proj", &o_proj);
29		//! tracer.log("layer0_ffn_norm", &ffn_normed);
30		//! tracer.log("layer0_ffn", &ffn_out);
31		//! // ... more layers ...
32		//! tracer.log("final_norm", &final_normed);
33		//! tracer.log("logits", &logits);
34		//!
35		//! // Compare CPU vs GPU traces
36		//! let divergence = BrickTracer::compare(&cpu_tracer, &gpu_tracer);
37		//! if let Some(first_diff) = divergence.first_divergence() {
38		//! eprintln!("Divergence at: {} (CPU L2={}, GPU L2={})",
39		//! first_diff.name, first_diff.cpu_l2, first_diff.gpu_l2);
40		//! }
41		//! ```
42		//!
43		//! # References
44		//!
45		//! - Phase 14: "The New Doctrine: The Golden Trace"
46		//! - PMAT-106: APR GPU Adapter + Integration Tests
47
48		use std::collections::HashMap;
49		use std::fmt;
50
51		/// A single trace event capturing tensor state at a computation step
52		#[derive(Debug, Clone)]
53		pub struct TraceEvent {
54		/// Name of the computation step (e.g., "layer0_rope_q")
55		pub name: String,
56		/// Position/sequence index when this event was logged
57		pub position: usize,
58		/// L2 norm of the tensor (cheap fingerprint)
59		pub l2_norm: f32,
60		/// Mean value
61		pub mean: f32,
62		/// Min value
63		pub min: f32,
64		/// Max value
65		pub max: f32,
66		/// First 8 elements (for quick debugging)
67		pub head: [f32; 8],
68		/// Full tensor data (only if verbose mode enabled)
69		pub full_data: Option<Vec<f32>>,
70		/// Tensor length
71		pub len: usize,
72		}
73
74		impl TraceEvent {
75		/// Create a new trace event from tensor data
76	24	pub fn new(name: &str, tensor: &[f32], position: usize, verbose: bool) -> Self {
77	24	let len = tensor.len();
78
79		// Compute L2 norm
80	53	let l2_norm24 = tensor24 . iter24 (). map24 (\|x\| x * x). sum24 ::<f32>(). sqrt24 ();
81
82		// Compute statistics
83	24	let mean = if len > 0 {
84	23	tensor.iter().sum::<f32>() / len as f32
85		} else {
86	1	0.0
87		};
88	24	let min = tensor.iter().cloned().fold(f32::INFINITY, f32::min);
89	24	let max = tensor.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
90
91		// Copy first 8 elements
92	24	let mut head = [0.0f32; 8];
93	53	for (i, &v) in tensor24 . iter24 (). take24 (8). enumerate24 () {
94	53	head[i] = v;
95	53	}
96
97		// Optionally store full tensor
98	24	let full_data = if verbose {
99	1	Some(tensor.to_vec())
100		} else {
101	23	None
102		};
103
104	24	Self {
105	24	name: name.to_string(),
106	24	position,
107	24	l2_norm,
108	24	mean,
109	24	min,
110	24	max,
111	24	head,
112	24	full_data,
113	24	len,
114	24	}
115	24	}
116
117		/// Check if two events are approximately equal within tolerance
118	8	pub fn approx_eq(&self, other: &Self, tolerance: f32) -> bool {
119		// Check L2 norm
120	8	let l2_diff = (self.l2_norm - other.l2_norm).abs();
121	8	let l2_rel = if self.l2_norm.abs() > 1e-10 {
122	8	l2_diff / self.l2_norm.abs()
123		} else {
124	0	l2_diff
125		};
126
127	8	l2_rel <= tolerance
128	8	}
129
130		/// Compute relative difference between two events
131	2	pub fn relative_diff(&self, other: &Self) -> f32 {
132	2	let l2_diff = (self.l2_norm - other.l2_norm).abs();
133	2	if self.l2_norm.abs() > 1e-10 {
134	2	l2_diff / self.l2_norm.abs()
135		} else {
136	0	l2_diff
137		}
138	2	}
139		}
140
141		impl fmt::Display for TraceEvent {
142	0	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
143	0	write!(
144	0	f,
145	0	"{} [pos={}]: L2={:.6}, mean={:.6}, range=[{:.6}, {:.6}], len={}, head=[{:.4}, {:.4}, {:.4}, {:.4}...]",
146		self.name,
147		self.position,
148		self.l2_norm,
149		self.mean,
150		self.min,
151		self.max,
152		self.len,
153	0	self.head[0],
154	0	self.head[1],
155	0	self.head[2],
156	0	self.head[3],
157		)
158	0	}
159		}
160
161		/// Result of comparing two trace events
162		#[derive(Debug, Clone)]
163		pub struct TraceDiff {
164		/// Name of the computation step
165		pub name: String,
166		/// Position where divergence occurred
167		pub position: usize,
168		/// CPU L2 norm
169		pub cpu_l2: f32,
170		/// GPU L2 norm
171		pub gpu_l2: f32,
172		/// Relative difference (\|cpu - gpu\| / \|cpu\|)
173		pub relative_diff: f32,
174		/// CPU head values
175		pub cpu_head: [f32; 8],
176		/// GPU head values
177		pub gpu_head: [f32; 8],
178		}
179
180		impl fmt::Display for TraceDiff {
181	0	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
182	0	write!(
183	0	f,
184	0	"{} [pos={}]: CPU L2={:.6} vs GPU L2={:.6} (diff={:.2}%)\n CPU head: [{:.4}, {:.4}, {:.4}, {:.4}...]\n GPU head: [{:.4}, {:.4}, {:.4}, {:.4}...]",
185		self.name,
186		self.position,
187		self.cpu_l2,
188		self.gpu_l2,
189	0	self.relative_diff * 100.0,
190	0	self.cpu_head[0], self.cpu_head[1], self.cpu_head[2], self.cpu_head[3],
191	0	self.gpu_head[0], self.gpu_head[1], self.gpu_head[2], self.gpu_head[3],
192		)
193	0	}
194		}
195
196		/// Comparison result between CPU and GPU traces
197		#[derive(Debug, Clone)]
198		pub struct TraceComparison {
199		/// All differences found
200		pub diffs: Vec<TraceDiff>,
201		/// Tolerance used for comparison
202		pub tolerance: f32,
203		}
204
205		impl TraceComparison {
206		/// Get the first divergence point
207	1	pub fn first_divergence(&self) -> Option<&TraceDiff> {
208	1	self.diffs.first()
209	1	}
210
211		/// Check if traces are equivalent (no divergence)
212	2	pub fn is_equivalent(&self) -> bool {
213	2	self.diffs.is_empty()
214	2	}
215
216		/// Get summary of divergence
217	0	pub fn summary(&self) -> String {
218	0	if self.diffs.is_empty() {
219	0	"No divergence detected".to_string()
220		} else {
221	0	let first = &self.diffs[0];
222	0	format!(
223	0	"First divergence at '{}' (pos={}): {:.2}% L2 diff ({} total divergences)",
224		first.name,
225		first.position,
226	0	first.relative_diff * 100.0,
227	0	self.diffs.len()
228		)
229		}
230	0	}
231		}
232
233		impl fmt::Display for TraceComparison {
234	0	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235	0	writeln!(f, "=== TRACE COMPARISON ===")?;
236	0	writeln!(f, "Tolerance: {:.4}%", self.tolerance * 100.0)?;
237	0	writeln!(f, "Divergences: {}", self.diffs.len())?;
238
239	0	if self.diffs.is_empty() {
240	0	writeln!(f, "Result: MATCH")?;
241		} else {
242	0	writeln!(f, "Result: DIVERGENCE DETECTED")?;
243	0	writeln!(f, "\n--- First Divergence ---")?;
244	0	if let Some(first) = self.diffs.first() {
245	0	writeln!(f, "{first}")?;
246	0	}
247
248	0	if self.diffs.len() > 1 {
249	0	writeln!(f, "\n--- All Divergences ---")?;
250	0	for diff in &self.diffs {
251	0	writeln!(f, " {}: {:.2}% diff", diff.name, diff.relative_diff * 100.0)?;
252		}
253	0	}
254		}
255
256	0	Ok(())
257	0	}
258		}
259
260		/// BrickTracer: Collects trace events for GPU/CPU parity debugging
261		///
262		/// The tracer logs tensor state at each computation step, enabling
263		/// automated detection of where GPU output diverges from CPU.
264		#[derive(Debug, Clone)]
265		pub struct BrickTracer {
266		/// Collected trace events in order
267		events: Vec<TraceEvent>,
268		/// Current position being traced
269		position: usize,
270		/// Whether to store full tensor data (expensive)
271		verbose: bool,
272		/// Event index by name for fast lookup
273		index: HashMap<String, usize>,
274		}
275
276		impl Default for BrickTracer {
277	0	fn default() -> Self {
278	0	Self::new()
279	0	}
280		}
281
282		impl BrickTracer {
283		/// Create a new tracer
284	7	pub fn new() -> Self {
285	7	Self {
286	7	events: Vec::new(),
287	7	position: 0,
288	7	verbose: false,
289	7	index: HashMap::new(),
290	7	}
291	7	}
292
293		/// Create a tracer with verbose mode (stores full tensors)
294	0	pub fn verbose() -> Self {
295	0	Self {
296	0	events: Vec::new(),
297	0	position: 0,
298	0	verbose: true,
299	0	index: HashMap::new(),
300	0	}
301	0	}
302
303		/// Set the current position being traced
304	2	pub fn set_position(&mut self, position: usize) {
305	2	self.position = position;
306	2	}
307
308		/// Log a tensor at the current computation step
309	15	pub fn log(&mut self, name: &str, tensor: &[f32]) {
310	15	let event = TraceEvent::new(name, tensor, self.position, self.verbose);
311	15	let idx = self.events.len();
312	15	self.index.insert(name.to_string(), idx);
313	15	self.events.push(event);
314	15	}
315
316		/// Log a tensor with explicit position
317	1	pub fn log_at(&mut self, name: &str, tensor: &[f32], position: usize) {
318	1	let event = TraceEvent::new(name, tensor, position, self.verbose);
319	1	let idx = self.events.len();
320	1	self.index.insert(name.to_string(), idx);
321	1	self.events.push(event);
322	1	}
323
324		/// Get all trace events
325	3	pub fn events(&self) -> &[TraceEvent] {
326	3	&self.events
327	3	}
328
329		/// Get event by name
330	11	pub fn get(&self, name: &str) -> Option<&TraceEvent> {
331	11	self.index.get(name).map(\|&idx\| & self.events[idx]10 )
332	11	}
333
334		/// Clear all events
335	1	pub fn clear(&mut self) {
336	1	self.events.clear();
337	1	self.index.clear();
338	1	self.position = 0;
339	1	}
340
341		/// Compare two tracers (CPU vs GPU)
342		///
343		/// # Arguments
344		///
345		/// * `cpu` - CPU tracer
346		/// * `gpu` - GPU tracer
347		/// * `tolerance` - Relative tolerance for L2 norm comparison (e.g., 0.01 = 1%)
348		///
349		/// # Returns
350		///
351		/// Comparison result with all divergences
352	2	pub fn compare(cpu: &Self, gpu: &Self, tolerance: f32) -> TraceComparison {
353	2	let mut diffs = Vec::new();
354
355		// Compare events in order
356	7	for cpu_event5 in &cpu.events {
357	5	if let Some(gpu_event) = gpu.get(&cpu_event.name) {
358	5	if !cpu_event.approx_eq(gpu_event, tolerance) {
359	1	diffs.push(TraceDiff {
360	1	name: cpu_event.name.clone(),
361	1	position: cpu_event.position,
362	1	cpu_l2: cpu_event.l2_norm,
363	1	gpu_l2: gpu_event.l2_norm,
364	1	relative_diff: cpu_event.relative_diff(gpu_event),
365	1	cpu_head: cpu_event.head,
366	1	gpu_head: gpu_event.head,
367	1	});
368	4	}
369	0	}
370		// Skip if GPU doesn't have this event (may have different instrumentation)
371		}
372
373	2	TraceComparison { diffs, tolerance }
374	2	}
375
376		/// Print all events to stderr for debugging
377	0	pub fn dump(&self) {
378	0	eprintln!("=== BRICK TRACE ({} events) ===", self.events.len());
379	0	for event in &self.events {
380	0	eprintln!(" {event}");
381	0	}
382	0	}
383
384		/// Print summary statistics
385	0	pub fn summary(&self) {
386	0	eprintln!("=== TRACE SUMMARY ===");
387	0	eprintln!("Events: {}", self.events.len());
388	0	if let Some(first) = self.events.first() {
389	0	eprintln!("First: {}", first.name);
390	0	}
391	0	if let Some(last) = self.events.last() {
392	0	eprintln!("Last: {}", last.name);
393	0	}
394
395		// Find largest L2 norm
396	0	if let Some(max_event) = self.events.iter().max_by(\|a, b\| {
397	0	a.l2_norm.partial_cmp(&b.l2_norm).unwrap_or(std::cmp::Ordering::Equal)
398	0	}) {
399	0	eprintln!("Max L2: {} = {:.6}", max_event.name, max_event.l2_norm);
400	0	}
401	0	}
402		}
403
404		// Global tracers for CPU and GPU paths (thread-local to avoid contention).
405		// Used by trace_cpu! and trace_gpu! macros when "trace" feature is enabled.
406		#[cfg(feature = "trace")]
407		#[allow(missing_docs)]
408		thread_local! {
409		/// Thread-local tracer for CPU inference path
410		pub static CPU_TRACER: std::cell::RefCell<BrickTracer> = std::cell::RefCell::new(BrickTracer::new());
411		/// Thread-local tracer for GPU inference path
412		pub static GPU_TRACER: std::cell::RefCell<BrickTracer> = std::cell::RefCell::new(BrickTracer::new());
413		}
414
415		/// Log a tensor to the CPU tracer for parity debugging.
416		///
417		/// Only active when the "trace" feature is enabled. When disabled, this macro
418		/// compiles to nothing (zero overhead).
419		///
420		/// # Arguments
421		///
422		/// * `$name` - Name of the computation step (e.g., "layer0_rope_q")
423		/// * `$tensor` - Slice of f32 values to log
424		/// * `$pos` - (Optional) Explicit position index
425		///
426		/// # Example
427		///
428		/// ```rust,ignore
429		/// trace_cpu!("embedding", &embedding_output);
430		/// trace_cpu!("layer0_attn", &attn_out, position);
431		/// ```
432		#[macro_export]
433		#[cfg(feature = "trace")]
434		macro_rules! trace_cpu {
435		($name:expr, $tensor:expr) => {
436		$crate::brick::tracer::CPU_TRACER.with(\|t\| {
437		t.borrow_mut().log($name, $tensor);
438		});
439		};
440		($name:expr, $tensor:expr, $pos:expr) => {
441		$crate::brick::tracer::CPU_TRACER.with(\|t\| {
442		t.borrow_mut().log_at($name, $tensor, $pos);
443		});
444		};
445		}
446
447		/// No-op version of trace_cpu when "trace" feature is disabled.
448		#[macro_export]
449		#[cfg(not(feature = "trace"))]
450		macro_rules! trace_cpu {
451		($name:expr, $tensor:expr) => {};
452		($name:expr, $tensor:expr, $pos:expr) => {};
453		}
454
455		/// Log a tensor to the GPU tracer for parity debugging.
456		///
457		/// Only active when the "trace" feature is enabled. When disabled, this macro
458		/// compiles to nothing (zero overhead).
459		///
460		/// # Arguments
461		///
462		/// * `$name` - Name of the computation step (e.g., "layer0_rope_q")
463		/// * `$tensor` - Slice of f32 values to log (must be downloaded from GPU first)
464		/// * `$pos` - (Optional) Explicit position index
465		///
466		/// # Example
467		///
468		/// ```rust,ignore
469		/// // After D2H copy from GPU buffer
470		/// trace_gpu!("embedding", &embedding_output);
471		/// trace_gpu!("layer0_attn", &attn_out, position);
472		/// ```
473		#[macro_export]
474		#[cfg(feature = "trace")]
475		macro_rules! trace_gpu {
476		($name:expr, $tensor:expr) => {
477		$crate::brick::tracer::GPU_TRACER.with(\|t\| {
478		t.borrow_mut().log($name, $tensor);
479		});
480		};
481		($name:expr, $tensor:expr, $pos:expr) => {
482		$crate::brick::tracer::GPU_TRACER.with(\|t\| {
483		t.borrow_mut().log_at($name, $tensor, $pos);
484		});
485		};
486		}
487
488		/// No-op version of trace_gpu when "trace" feature is disabled.
489		#[macro_export]
490		#[cfg(not(feature = "trace"))]
491		macro_rules! trace_gpu {
492		($name:expr, $tensor:expr) => {};
493		($name:expr, $tensor:expr, $pos:expr) => {};
494		}
495
496		// ============================================================================
497		// Tests
498		// ============================================================================
499
500		#[cfg(test)]
501		mod tests {
502		use super::*;
503
504		#[test]
505	1	fn test_trace_event_creation() {
506	1	let tensor = vec![1.0, 2.0, 3.0, 4.0, 5.0];
507	1	let event = TraceEvent::new("test", &tensor, 0, false);
508
509	1	assert_eq!(event.name, "test");
510	1	assert_eq!(event.position, 0);
511	1	assert_eq!(event.len, 5);
512	1	assert!((event.l2_norm - 7.416198).abs() < 0.001); // sqrt(55)
513	1	assert!((event.mean - 3.0).abs() < 0.001);
514	1	assert!((event.min - 1.0).abs() < 0.001);
515	1	assert!((event.max - 5.0).abs() < 0.001);
516	1	assert!(event.full_data.is_none());
517	1	}
518
519		#[test]
520	1	fn test_trace_event_verbose() {
521	1	let tensor = vec![1.0, 2.0, 3.0];
522	1	let event = TraceEvent::new("test", &tensor, 0, true);
523
524	1	assert!(event.full_data.is_some());
525	1	assert_eq!(event.full_data.unwrap(), tensor);
526	1	}
527
528		#[test]
529	1	fn test_tracer_log() {
530	1	let mut tracer = BrickTracer::new();
531
532	1	tracer.log("step1", &[1.0, 2.0, 3.0]);
533	1	tracer.log("step2", &[4.0, 5.0, 6.0]);
534
535	1	assert_eq!(tracer.events().len(), 2);
536	1	assert!(tracer.get("step1").is_some());
537	1	assert!(tracer.get("step2").is_some());
538	1	assert!(tracer.get("step3").is_none());
539	1	}
540
541		#[test]
542	1	fn test_tracer_comparison_match() {
543	1	let mut cpu = BrickTracer::new();
544	1	let mut gpu = BrickTracer::new();
545
546	1	cpu.log("embedding", &[1.0, 2.0, 3.0]);
547	1	cpu.log("layer0_norm", &[0.5, 1.0, 1.5]);
548
549	1	gpu.log("embedding", &[1.0, 2.0, 3.0]);
550	1	gpu.log("layer0_norm", &[0.5, 1.0, 1.5]);
551
552	1	let comparison = BrickTracer::compare(&cpu, &gpu, 0.01);
553	1	assert!(comparison.is_equivalent());
554	1	}
555
556		#[test]
557	1	fn test_tracer_comparison_diverge() {
558	1	let mut cpu = BrickTracer::new();
559	1	let mut gpu = BrickTracer::new();
560
561	1	cpu.log("embedding", &[1.0, 2.0, 3.0]);
562	1	cpu.log("layer0_norm", &[0.5, 1.0, 1.5]);
563	1	cpu.log("layer0_attn", &[0.1, 0.2, 0.3]);
564
565	1	gpu.log("embedding", &[1.0, 2.0, 3.0]); // Match
566	1	gpu.log("layer0_norm", &[0.5, 1.0, 1.5]); // Match
567	1	gpu.log("layer0_attn", &[0.5, 0.6, 0.7]); // DIVERGE!
568
569	1	let comparison = BrickTracer::compare(&cpu, &gpu, 0.01);
570	1	assert!(!comparison.is_equivalent());
571
572	1	let first = comparison.first_divergence().unwrap();
573	1	assert_eq!(first.name, "layer0_attn");
574	1	}
575
576		#[test]
577	1	fn test_tracer_clear() {
578	1	let mut tracer = BrickTracer::new();
579	1	tracer.log("step1", &[1.0]);
580	1	assert_eq!(tracer.events().len(), 1);
581
582	1	tracer.clear();
583	1	assert_eq!(tracer.events().len(), 0);
584	1	}
585
586		#[test]
587	1	fn test_tracer_position() {
588	1	let mut tracer = BrickTracer::new();
589
590	1	tracer.set_position(0);
591	1	tracer.log("pos0", &[1.0]);
592
593	1	tracer.set_position(1);
594	1	tracer.log("pos1", &[2.0]);
595
596	1	tracer.log_at("explicit", &[3.0], 5);
597
598	1	assert_eq!(tracer.get("pos0").unwrap().position, 0);
599	1	assert_eq!(tracer.get("pos1").unwrap().position, 1);
600	1	assert_eq!(tracer.get("explicit").unwrap().position, 5);
601	1	}
602
603		#[test]
604	1	fn test_relative_diff() {
605	1	let event1 = TraceEvent::new("a", &[10.0], 0, false);
606	1	let event2 = TraceEvent::new("b", &[11.0], 0, false);
607
608	1	let diff = event1.relative_diff(&event2);
609	1	assert!((diff - 0.1).abs() < 0.001); // 10% difference
610	1	}
611
612		#[test]
613	1	fn test_approx_eq() {
614	1	let event1 = TraceEvent::new("a", &[10.0], 0, false);
615	1	let event2 = TraceEvent::new("b", &[10.05], 0, false);
616	1	let event3 = TraceEvent::new("c", &[11.0], 0, false);
617
618	1	assert!(event1.approx_eq(&event2, 0.01)); // 1% tolerance
619	1	assert!(!event1.approx_eq(&event3, 0.01)); // 10% diff > 1% tolerance
620	1	assert!(event1.approx_eq(&event3, 0.15)); // 10% diff < 15% tolerance
621	1	}
622
623		#[test]
624	1	fn test_empty_tensor() {
625	1	let event = TraceEvent::new("empty", &[], 0, false);
626	1	assert_eq!(event.len, 0);
627	1	assert_eq!(event.l2_norm, 0.0);
628	1	assert_eq!(event.mean, 0.0);
629	1	}
630		}