impl BrickTracer {
pub fn new() -> Self {
Self {
events: Vec::new(),
position: 0,
verbose: false,
index: HashMap::new(),
}
}
pub fn verbose() -> Self {
Self {
events: Vec::new(),
position: 0,
verbose: true,
index: HashMap::new(),
}
}
pub fn set_position(&mut self, position: usize) {
self.position = position;
}
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);
}
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);
}
pub fn events(&self) -> &[TraceEvent] {
&self.events
}
pub fn get(&self, name: &str) -> Option<&TraceEvent> {
self.index.get(name).map(|&idx| &self.events[idx])
}
pub fn clear(&mut self) {
self.events.clear();
self.index.clear();
self.position = 0;
}
pub fn compare(cpu: &Self, gpu: &Self, tolerance: f32) -> TraceComparison {
let mut diffs = Vec::new();
for cpu_event 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,
});
}
}
}
TraceComparison { diffs, tolerance }
}
pub fn dump(&self) {
eprintln!("=== BRICK TRACE ({} events) ===", self.events.len());
for event in &self.events {
eprintln!(" {event}");
}
}
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);
}
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);
}
}
}
#[cfg(feature = "trace")]
thread_local! {
pub static CPU_TRACER: std::cell::RefCell<BrickTracer> = std::cell::RefCell::new(BrickTracer::new());
pub static GPU_TRACER: std::cell::RefCell<BrickTracer> = std::cell::RefCell::new(BrickTracer::new());
}
#[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);
});
};
}
#[macro_export]
#[cfg(not(feature = "trace"))]
macro_rules! trace_cpu {
($name:expr, $tensor:expr) => {};
($name:expr, $tensor:expr, $pos:expr) => {};
}
#[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);
});
};
}
#[macro_export]
#[cfg(not(feature = "trace"))]
macro_rules! trace_gpu {
($name:expr, $tensor:expr) => {};
($name:expr, $tensor:expr, $pos:expr) => {};
}