950 lines
28 KiB
Rust
950 lines
28 KiB
Rust
//! Profiling Integration
|
|
//!
|
|
//! This module provides profiling and performance analysis tools:
|
|
//! - NVTX/rocTX range markers for timeline visualization
|
|
//! - Automatic kernel timing and statistics
|
|
//! - Memory usage tracking
|
|
//! - Communication profiling
|
|
//! - Flamegraph generation support
|
|
//!
|
|
//! Integration with NVIDIA Nsight Systems and AMD ROCm profiler.
|
|
|
|
use crate::error::{DistributedError, Result};
|
|
use parking_lot::RwLock;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
use std::time::{Duration, Instant};
|
|
|
|
// =============================================================================
|
|
// Configuration
|
|
// =============================================================================
|
|
|
|
/// Profiling level
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ProfilingLevel {
|
|
/// No profiling
|
|
Off,
|
|
/// Basic profiling (major operations only)
|
|
Basic,
|
|
/// Detailed profiling (includes sub-operations)
|
|
Detailed,
|
|
/// Full profiling (all operations, high overhead)
|
|
Full,
|
|
}
|
|
|
|
/// Profiler backend
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ProfilerBackend {
|
|
/// NVIDIA NVTX (Nsight Systems)
|
|
Nvtx,
|
|
/// AMD rocTX (ROCm profiler)
|
|
RocTx,
|
|
/// Built-in timing (no external dependency)
|
|
Builtin,
|
|
/// Chrome trace format
|
|
ChromeTrace,
|
|
}
|
|
|
|
/// Configuration for profiling
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ProfilingConfig {
|
|
/// Profiling level
|
|
pub level: ProfilingLevel,
|
|
/// Profiler backend
|
|
pub backend: ProfilerBackend,
|
|
/// Enable memory profiling
|
|
pub profile_memory: bool,
|
|
/// Enable communication profiling
|
|
pub profile_comm: bool,
|
|
/// Enable kernel profiling
|
|
pub profile_kernels: bool,
|
|
/// Sample rate for detailed profiling (1 = every op, 10 = every 10th)
|
|
pub sample_rate: u32,
|
|
/// Output directory for traces
|
|
pub output_dir: Option<String>,
|
|
/// Enable automatic range nesting
|
|
pub auto_nest: bool,
|
|
/// Maximum range depth
|
|
pub max_depth: usize,
|
|
}
|
|
|
|
impl Default for ProfilingConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
level: ProfilingLevel::Basic,
|
|
backend: ProfilerBackend::Builtin,
|
|
profile_memory: true,
|
|
profile_comm: true,
|
|
profile_kernels: true,
|
|
sample_rate: 1,
|
|
output_dir: None,
|
|
auto_nest: true,
|
|
max_depth: 32,
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// NVTX Colors
|
|
// =============================================================================
|
|
|
|
/// NVTX color for range visualization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct NvtxColor(pub u32);
|
|
|
|
impl NvtxColor {
|
|
pub const RED: Self = Self(0xFF0000);
|
|
pub const GREEN: Self = Self(0x00FF00);
|
|
pub const BLUE: Self = Self(0x0000FF);
|
|
pub const YELLOW: Self = Self(0xFFFF00);
|
|
pub const CYAN: Self = Self(0x00FFFF);
|
|
pub const MAGENTA: Self = Self(0xFF00FF);
|
|
pub const ORANGE: Self = Self(0xFFA500);
|
|
pub const PURPLE: Self = Self(0x800080);
|
|
pub const WHITE: Self = Self(0xFFFFFF);
|
|
pub const GRAY: Self = Self(0x808080);
|
|
|
|
/// Color for forward pass
|
|
pub const FORWARD: Self = Self::GREEN;
|
|
/// Color for backward pass
|
|
pub const BACKWARD: Self = Self::BLUE;
|
|
/// Color for optimizer step
|
|
pub const OPTIMIZER: Self = Self::ORANGE;
|
|
/// Color for communication
|
|
pub const COMM: Self = Self::YELLOW;
|
|
/// Color for memory operations
|
|
pub const MEMORY: Self = Self::CYAN;
|
|
/// Color for data loading
|
|
pub const DATA: Self = Self::PURPLE;
|
|
}
|
|
|
|
// =============================================================================
|
|
// Profiling Range
|
|
// =============================================================================
|
|
|
|
/// A profiling range (NVTX range equivalent)
|
|
#[derive(Debug, Clone)]
|
|
pub struct ProfilingRange {
|
|
/// Range name
|
|
pub name: String,
|
|
/// Category
|
|
pub category: String,
|
|
/// Color for visualization
|
|
pub color: NvtxColor,
|
|
/// Start time
|
|
pub start_time: Instant,
|
|
/// End time (None if still active)
|
|
pub end_time: Option<Instant>,
|
|
/// Parent range ID (for nesting)
|
|
pub parent_id: Option<u64>,
|
|
/// Range ID
|
|
pub id: u64,
|
|
/// Thread ID
|
|
pub thread_id: u64,
|
|
/// Additional metadata
|
|
pub metadata: HashMap<String, String>,
|
|
}
|
|
|
|
impl ProfilingRange {
|
|
/// Create a new range
|
|
pub fn new(id: u64, name: String, category: String, color: NvtxColor) -> Self {
|
|
Self {
|
|
name,
|
|
category,
|
|
color,
|
|
start_time: Instant::now(),
|
|
end_time: None,
|
|
parent_id: None,
|
|
id,
|
|
thread_id: {
|
|
use std::hash::{Hash, Hasher};
|
|
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
|
std::thread::current().id().hash(&mut hasher);
|
|
hasher.finish()
|
|
},
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// End the range
|
|
pub fn end(&mut self) {
|
|
self.end_time = Some(Instant::now());
|
|
}
|
|
|
|
/// Get duration
|
|
pub fn duration(&self) -> Option<Duration> {
|
|
self.end_time.map(|end| end.duration_since(self.start_time))
|
|
}
|
|
|
|
/// Add metadata
|
|
pub fn with_metadata(mut self, key: &str, value: &str) -> Self {
|
|
self.metadata.insert(key.to_string(), value.to_string());
|
|
self
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Range Guard (RAII)
|
|
// =============================================================================
|
|
|
|
/// RAII guard for automatic range ending
|
|
pub struct RangeGuard {
|
|
profiler: Arc<Profiler>,
|
|
range_id: u64,
|
|
}
|
|
|
|
impl RangeGuard {
|
|
/// Create a new guard
|
|
pub fn new(profiler: Arc<Profiler>, range_id: u64) -> Self {
|
|
Self { profiler, range_id }
|
|
}
|
|
|
|
/// Add metadata to the range
|
|
pub fn add_metadata(&self, key: &str, value: &str) {
|
|
self.profiler.add_range_metadata(self.range_id, key, value);
|
|
}
|
|
}
|
|
|
|
impl Drop for RangeGuard {
|
|
fn drop(&mut self) {
|
|
self.profiler.end_range(self.range_id);
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Operation Statistics
|
|
// =============================================================================
|
|
|
|
/// Statistics for a profiled operation
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct OperationStats {
|
|
/// Operation name
|
|
pub name: String,
|
|
/// Number of calls
|
|
pub call_count: u64,
|
|
/// Total time spent
|
|
pub total_time: Duration,
|
|
/// Minimum time
|
|
pub min_time: Duration,
|
|
/// Maximum time
|
|
pub max_time: Duration,
|
|
/// Total bytes processed (if applicable)
|
|
pub total_bytes: u64,
|
|
}
|
|
|
|
impl OperationStats {
|
|
/// Create new stats for an operation
|
|
pub fn new(name: String) -> Self {
|
|
Self {
|
|
name,
|
|
min_time: Duration::MAX,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
/// Record a timing
|
|
pub fn record(&mut self, duration: Duration, bytes: u64) {
|
|
self.call_count += 1;
|
|
self.total_time += duration;
|
|
self.min_time = self.min_time.min(duration);
|
|
self.max_time = self.max_time.max(duration);
|
|
self.total_bytes += bytes;
|
|
}
|
|
|
|
/// Get average time
|
|
pub fn avg_time(&self) -> Duration {
|
|
if self.call_count == 0 {
|
|
Duration::ZERO
|
|
} else {
|
|
self.total_time / self.call_count as u32
|
|
}
|
|
}
|
|
|
|
/// Get throughput in GB/s
|
|
pub fn throughput_gbps(&self) -> f64 {
|
|
if self.total_time.as_secs_f64() == 0.0 {
|
|
0.0
|
|
} else {
|
|
(self.total_bytes as f64 / (1024.0 * 1024.0 * 1024.0)) / self.total_time.as_secs_f64()
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Memory Snapshot
|
|
// =============================================================================
|
|
|
|
/// Memory usage snapshot
|
|
#[derive(Debug, Clone)]
|
|
pub struct MemorySnapshot {
|
|
/// Timestamp
|
|
pub timestamp: Instant,
|
|
/// Device ID
|
|
pub device_id: i32,
|
|
/// Allocated bytes
|
|
pub allocated_bytes: u64,
|
|
/// Reserved bytes
|
|
pub reserved_bytes: u64,
|
|
/// Peak allocated bytes
|
|
pub peak_allocated_bytes: u64,
|
|
/// Number of allocations
|
|
pub num_allocations: u64,
|
|
}
|
|
|
|
// =============================================================================
|
|
// Communication Event
|
|
// =============================================================================
|
|
|
|
/// Communication event for profiling
|
|
#[derive(Debug, Clone)]
|
|
pub struct CommEvent {
|
|
/// Event name
|
|
pub name: String,
|
|
/// Start time
|
|
pub start_time: Instant,
|
|
/// End time
|
|
pub end_time: Option<Instant>,
|
|
/// Bytes transferred
|
|
pub bytes: u64,
|
|
/// Source rank
|
|
pub src_rank: Option<i32>,
|
|
/// Destination rank
|
|
pub dst_rank: Option<i32>,
|
|
/// Operation type
|
|
pub op_type: String,
|
|
}
|
|
|
|
// =============================================================================
|
|
// Profiler
|
|
// =============================================================================
|
|
|
|
/// Main profiler instance
|
|
pub struct Profiler {
|
|
/// Configuration
|
|
config: ProfilingConfig,
|
|
/// Whether profiling is enabled
|
|
enabled: AtomicBool,
|
|
/// Next range ID
|
|
next_range_id: AtomicU64,
|
|
/// Active ranges
|
|
active_ranges: RwLock<HashMap<u64, ProfilingRange>>,
|
|
/// Completed ranges
|
|
completed_ranges: RwLock<Vec<ProfilingRange>>,
|
|
/// Operation statistics
|
|
op_stats: RwLock<HashMap<String, OperationStats>>,
|
|
/// Memory snapshots
|
|
memory_snapshots: RwLock<Vec<MemorySnapshot>>,
|
|
/// Communication events
|
|
comm_events: RwLock<Vec<CommEvent>>,
|
|
/// Current range stack (per thread)
|
|
range_stack: RwLock<Vec<u64>>,
|
|
/// Profile start time
|
|
start_time: Instant,
|
|
}
|
|
|
|
impl Profiler {
|
|
/// Create a new profiler
|
|
pub fn new(config: ProfilingConfig) -> Self {
|
|
let enabled = config.level != ProfilingLevel::Off;
|
|
|
|
Self {
|
|
config,
|
|
enabled: AtomicBool::new(enabled),
|
|
next_range_id: AtomicU64::new(1),
|
|
active_ranges: RwLock::new(HashMap::new()),
|
|
completed_ranges: RwLock::new(Vec::new()),
|
|
op_stats: RwLock::new(HashMap::new()),
|
|
memory_snapshots: RwLock::new(Vec::new()),
|
|
comm_events: RwLock::new(Vec::new()),
|
|
range_stack: RwLock::new(Vec::new()),
|
|
start_time: Instant::now(),
|
|
}
|
|
}
|
|
|
|
/// Check if profiling is enabled
|
|
pub fn is_enabled(&self) -> bool {
|
|
self.enabled.load(Ordering::Relaxed)
|
|
}
|
|
|
|
/// Enable profiling
|
|
pub fn enable(&self) {
|
|
self.enabled.store(true, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Disable profiling
|
|
pub fn disable(&self) {
|
|
self.enabled.store(false, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Start a profiling range
|
|
pub fn start_range(&self, name: &str, category: &str, color: NvtxColor) -> u64 {
|
|
if !self.is_enabled() {
|
|
return 0;
|
|
}
|
|
|
|
let id = self.next_range_id.fetch_add(1, Ordering::SeqCst);
|
|
let mut range = ProfilingRange::new(id, name.to_string(), category.to_string(), color);
|
|
|
|
// Set parent if auto-nesting
|
|
if self.config.auto_nest {
|
|
let stack = self.range_stack.read();
|
|
range.parent_id = stack.last().copied();
|
|
}
|
|
|
|
// Add to active ranges
|
|
{
|
|
let mut active = self.active_ranges.write();
|
|
active.insert(id, range);
|
|
}
|
|
|
|
// Push to stack
|
|
{
|
|
let mut stack = self.range_stack.write();
|
|
if stack.len() < self.config.max_depth {
|
|
stack.push(id);
|
|
}
|
|
}
|
|
|
|
id
|
|
}
|
|
|
|
/// Start a range with RAII guard
|
|
pub fn range_guard(
|
|
self: &Arc<Self>,
|
|
name: &str,
|
|
category: &str,
|
|
color: NvtxColor,
|
|
) -> RangeGuard {
|
|
let id = self.start_range(name, category, color);
|
|
RangeGuard::new(self.clone(), id)
|
|
}
|
|
|
|
/// End a profiling range
|
|
pub fn end_range(&self, range_id: u64) {
|
|
if !self.is_enabled() || range_id == 0 {
|
|
return;
|
|
}
|
|
|
|
let range = {
|
|
let mut active = self.active_ranges.write();
|
|
active.remove(&range_id)
|
|
};
|
|
|
|
if let Some(mut range) = range {
|
|
range.end();
|
|
|
|
// Update operation stats
|
|
if let Some(duration) = range.duration() {
|
|
let mut stats = self.op_stats.write();
|
|
let op_stats = stats
|
|
.entry(range.name.clone())
|
|
.or_insert_with(|| OperationStats::new(range.name.clone()));
|
|
op_stats.record(duration, 0);
|
|
}
|
|
|
|
// Move to completed
|
|
let mut completed = self.completed_ranges.write();
|
|
completed.push(range);
|
|
}
|
|
|
|
// Pop from stack
|
|
{
|
|
let mut stack = self.range_stack.write();
|
|
if stack.last() == Some(&range_id) {
|
|
stack.pop();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Add metadata to an active range
|
|
pub fn add_range_metadata(&self, range_id: u64, key: &str, value: &str) {
|
|
if let Some(range) = self.active_ranges.write().get_mut(&range_id) {
|
|
range.metadata.insert(key.to_string(), value.to_string());
|
|
}
|
|
}
|
|
|
|
/// Record a memory snapshot
|
|
pub fn record_memory(
|
|
&self,
|
|
device_id: i32,
|
|
allocated: u64,
|
|
reserved: u64,
|
|
peak: u64,
|
|
num_allocs: u64,
|
|
) {
|
|
if !self.is_enabled() || !self.config.profile_memory {
|
|
return;
|
|
}
|
|
|
|
let snapshot = MemorySnapshot {
|
|
timestamp: Instant::now(),
|
|
device_id,
|
|
allocated_bytes: allocated,
|
|
reserved_bytes: reserved,
|
|
peak_allocated_bytes: peak,
|
|
num_allocations: num_allocs,
|
|
};
|
|
|
|
self.memory_snapshots.write().push(snapshot);
|
|
}
|
|
|
|
/// Record a communication event
|
|
pub fn record_comm_start(&self, name: &str, bytes: u64, op_type: &str) -> usize {
|
|
if !self.is_enabled() || !self.config.profile_comm {
|
|
return 0;
|
|
}
|
|
|
|
let event = CommEvent {
|
|
name: name.to_string(),
|
|
start_time: Instant::now(),
|
|
end_time: None,
|
|
bytes,
|
|
src_rank: None,
|
|
dst_rank: None,
|
|
op_type: op_type.to_string(),
|
|
};
|
|
|
|
let mut events = self.comm_events.write();
|
|
let idx = events.len();
|
|
events.push(event);
|
|
idx
|
|
}
|
|
|
|
/// End a communication event
|
|
pub fn record_comm_end(&self, event_idx: usize) {
|
|
let mut events = self.comm_events.write();
|
|
if let Some(event) = events.get_mut(event_idx) {
|
|
event.end_time = Some(Instant::now());
|
|
}
|
|
}
|
|
|
|
/// Get operation statistics
|
|
pub fn get_op_stats(&self) -> HashMap<String, OperationStats> {
|
|
self.op_stats.read().clone()
|
|
}
|
|
|
|
/// Get memory snapshots
|
|
pub fn get_memory_snapshots(&self) -> Vec<MemorySnapshot> {
|
|
self.memory_snapshots.read().clone()
|
|
}
|
|
|
|
/// Get communication events
|
|
pub fn get_comm_events(&self) -> Vec<CommEvent> {
|
|
self.comm_events.read().clone()
|
|
}
|
|
|
|
/// Get completed ranges
|
|
pub fn get_completed_ranges(&self) -> Vec<ProfilingRange> {
|
|
self.completed_ranges.read().clone()
|
|
}
|
|
|
|
/// Clear all profiling data
|
|
pub fn clear(&self) {
|
|
self.active_ranges.write().clear();
|
|
self.completed_ranges.write().clear();
|
|
self.op_stats.write().clear();
|
|
self.memory_snapshots.write().clear();
|
|
self.comm_events.write().clear();
|
|
self.range_stack.write().clear();
|
|
}
|
|
|
|
/// Generate summary report
|
|
pub fn summary(&self) -> ProfileSummary {
|
|
let stats = self.op_stats.read();
|
|
let completed = self.completed_ranges.read();
|
|
let memory = self.memory_snapshots.read();
|
|
let comm = self.comm_events.read();
|
|
|
|
let total_time = self.start_time.elapsed();
|
|
|
|
let mut op_summaries: Vec<_> = stats
|
|
.values()
|
|
.map(|s| OpSummary {
|
|
name: s.name.clone(),
|
|
calls: s.call_count,
|
|
total_ms: s.total_time.as_secs_f64() * 1000.0,
|
|
avg_ms: s.avg_time().as_secs_f64() * 1000.0,
|
|
percent: (s.total_time.as_secs_f64() / total_time.as_secs_f64()) * 100.0,
|
|
})
|
|
.collect();
|
|
|
|
op_summaries.sort_by(|a, b| b.total_ms.partial_cmp(&a.total_ms).unwrap());
|
|
|
|
let peak_memory = memory
|
|
.iter()
|
|
.map(|s| s.peak_allocated_bytes)
|
|
.max()
|
|
.unwrap_or(0);
|
|
|
|
let total_comm_bytes: u64 = comm.iter().map(|e| e.bytes).sum();
|
|
let total_comm_time: Duration = comm
|
|
.iter()
|
|
.filter_map(|e| e.end_time.map(|end| end.duration_since(e.start_time)))
|
|
.sum();
|
|
|
|
ProfileSummary {
|
|
total_time,
|
|
num_ranges: completed.len(),
|
|
num_ops: stats.len(),
|
|
op_summaries,
|
|
peak_memory_bytes: peak_memory,
|
|
total_comm_bytes,
|
|
total_comm_time,
|
|
}
|
|
}
|
|
|
|
/// Export to Chrome trace format
|
|
pub fn export_chrome_trace(&self) -> String {
|
|
let completed = self.completed_ranges.read();
|
|
let base_time = self.start_time;
|
|
|
|
let mut events = Vec::new();
|
|
|
|
for range in completed.iter() {
|
|
let start_us = range.start_time.duration_since(base_time).as_micros();
|
|
let dur_us = range.duration().map_or(0, |d| d.as_micros());
|
|
|
|
events.push(format!(
|
|
r#"{{"name":"{}","cat":"{}","ph":"X","ts":{},"dur":{},"pid":1,"tid":{}}}"#,
|
|
range.name, range.category, start_us, dur_us, range.thread_id
|
|
));
|
|
}
|
|
|
|
format!("[{}]", events.join(","))
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Profile Summary
|
|
// =============================================================================
|
|
|
|
/// Summary of profiling results
|
|
#[derive(Debug, Clone)]
|
|
pub struct ProfileSummary {
|
|
/// Total profiling time
|
|
pub total_time: Duration,
|
|
/// Number of ranges recorded
|
|
pub num_ranges: usize,
|
|
/// Number of unique operations
|
|
pub num_ops: usize,
|
|
/// Operation summaries (sorted by time)
|
|
pub op_summaries: Vec<OpSummary>,
|
|
/// Peak memory usage
|
|
pub peak_memory_bytes: u64,
|
|
/// Total communication bytes
|
|
pub total_comm_bytes: u64,
|
|
/// Total communication time
|
|
pub total_comm_time: Duration,
|
|
}
|
|
|
|
/// Summary for a single operation type
|
|
#[derive(Debug, Clone)]
|
|
pub struct OpSummary {
|
|
/// Operation name
|
|
pub name: String,
|
|
/// Number of calls
|
|
pub calls: u64,
|
|
/// Total time in milliseconds
|
|
pub total_ms: f64,
|
|
/// Average time in milliseconds
|
|
pub avg_ms: f64,
|
|
/// Percentage of total time
|
|
pub percent: f64,
|
|
}
|
|
|
|
impl ProfileSummary {
|
|
/// Format as string
|
|
pub fn to_string(&self) -> String {
|
|
let mut s = String::new();
|
|
s.push_str("=== Profile Summary ===\n");
|
|
s.push_str(&format!(
|
|
"Total time: {:.2}s\n",
|
|
self.total_time.as_secs_f64()
|
|
));
|
|
s.push_str(&format!("Ranges: {}\n", self.num_ranges));
|
|
s.push_str(&format!("Unique ops: {}\n", self.num_ops));
|
|
s.push_str(&format!(
|
|
"Peak memory: {:.2} MB\n",
|
|
self.peak_memory_bytes as f64 / 1024.0 / 1024.0
|
|
));
|
|
s.push_str(&format!(
|
|
"Total comm: {:.2} MB in {:.2}ms\n",
|
|
self.total_comm_bytes as f64 / 1024.0 / 1024.0,
|
|
self.total_comm_time.as_secs_f64() * 1000.0
|
|
));
|
|
s.push_str("\nTop operations:\n");
|
|
|
|
for (i, op) in self.op_summaries.iter().take(10).enumerate() {
|
|
s.push_str(&format!(
|
|
" {}. {} - {} calls, {:.2}ms total ({:.1}%)\n",
|
|
i + 1,
|
|
op.name,
|
|
op.calls,
|
|
op.total_ms,
|
|
op.percent
|
|
));
|
|
}
|
|
|
|
s
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Thread-Safe Wrapper
|
|
// =============================================================================
|
|
|
|
/// Thread-safe shared profiler
|
|
pub type SharedProfiler = Arc<Profiler>;
|
|
|
|
/// Create a shared profiler
|
|
pub fn shared_profiler(config: ProfilingConfig) -> SharedProfiler {
|
|
Arc::new(Profiler::new(config))
|
|
}
|
|
|
|
/// Global profiler instance
|
|
static GLOBAL_PROFILER: std::sync::OnceLock<SharedProfiler> = std::sync::OnceLock::new();
|
|
|
|
/// Get or initialize global profiler
|
|
pub fn global_profiler() -> &'static SharedProfiler {
|
|
GLOBAL_PROFILER.get_or_init(|| shared_profiler(ProfilingConfig::default()))
|
|
}
|
|
|
|
/// Initialize global profiler with config
|
|
pub fn init_global_profiler(config: ProfilingConfig) -> Result<()> {
|
|
GLOBAL_PROFILER
|
|
.set(shared_profiler(config))
|
|
.map_err(|_| DistributedError::runtime("Global profiler already initialized"))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Convenience Macros (as functions)
|
|
// =============================================================================
|
|
|
|
/// Profile a forward pass operation
|
|
pub fn profile_forward<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
|
|
where
|
|
F: FnOnce() -> T,
|
|
{
|
|
let _guard = profiler.range_guard(name, "forward", NvtxColor::FORWARD);
|
|
f()
|
|
}
|
|
|
|
/// Profile a backward pass operation
|
|
pub fn profile_backward<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
|
|
where
|
|
F: FnOnce() -> T,
|
|
{
|
|
let _guard = profiler.range_guard(name, "backward", NvtxColor::BACKWARD);
|
|
f()
|
|
}
|
|
|
|
/// Profile an optimizer step
|
|
pub fn profile_optimizer<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
|
|
where
|
|
F: FnOnce() -> T,
|
|
{
|
|
let _guard = profiler.range_guard(name, "optimizer", NvtxColor::OPTIMIZER);
|
|
f()
|
|
}
|
|
|
|
/// Profile a communication operation
|
|
pub fn profile_comm<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
|
|
where
|
|
F: FnOnce() -> T,
|
|
{
|
|
let _guard = profiler.range_guard(name, "comm", NvtxColor::COMM);
|
|
f()
|
|
}
|
|
|
|
// =============================================================================
|
|
// Tests
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_profiling_config_default() {
|
|
let config = ProfilingConfig::default();
|
|
assert_eq!(config.level, ProfilingLevel::Basic);
|
|
assert_eq!(config.backend, ProfilerBackend::Builtin);
|
|
assert!(config.profile_memory);
|
|
}
|
|
|
|
#[test]
|
|
fn test_nvtx_colors() {
|
|
assert_eq!(NvtxColor::RED.0, 0xFF0000);
|
|
assert_eq!(NvtxColor::GREEN.0, 0x00FF00);
|
|
assert_eq!(NvtxColor::FORWARD.0, NvtxColor::GREEN.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_profiler_creation() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
assert!(profiler.is_enabled());
|
|
}
|
|
|
|
#[test]
|
|
fn test_profiler_enable_disable() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
profiler.disable();
|
|
assert!(!profiler.is_enabled());
|
|
|
|
profiler.enable();
|
|
assert!(profiler.is_enabled());
|
|
}
|
|
|
|
#[test]
|
|
fn test_range_start_end() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
let id = profiler.start_range("test_op", "test", NvtxColor::GREEN);
|
|
assert!(id > 0);
|
|
|
|
std::thread::sleep(std::time::Duration::from_millis(10));
|
|
profiler.end_range(id);
|
|
|
|
let completed = profiler.get_completed_ranges();
|
|
assert_eq!(completed.len(), 1);
|
|
assert_eq!(completed[0].name, "test_op");
|
|
assert!(completed[0].duration().unwrap() >= Duration::from_millis(10));
|
|
}
|
|
|
|
#[test]
|
|
fn test_range_guard() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Arc::new(Profiler::new(config));
|
|
|
|
{
|
|
let _guard = profiler.range_guard("scoped_op", "test", NvtxColor::BLUE);
|
|
std::thread::sleep(std::time::Duration::from_millis(5));
|
|
}
|
|
|
|
let completed = profiler.get_completed_ranges();
|
|
assert_eq!(completed.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_stats() {
|
|
let mut stats = OperationStats::new("test".to_string());
|
|
|
|
stats.record(Duration::from_millis(10), 1024);
|
|
stats.record(Duration::from_millis(20), 2048);
|
|
|
|
assert_eq!(stats.call_count, 2);
|
|
assert_eq!(stats.total_time, Duration::from_millis(30));
|
|
assert_eq!(stats.avg_time(), Duration::from_millis(15));
|
|
assert_eq!(stats.total_bytes, 3072);
|
|
}
|
|
|
|
#[test]
|
|
fn test_memory_snapshot() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
profiler.record_memory(0, 1000, 2000, 1500, 10);
|
|
|
|
let snapshots = profiler.get_memory_snapshots();
|
|
assert_eq!(snapshots.len(), 1);
|
|
assert_eq!(snapshots[0].allocated_bytes, 1000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_comm_event() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
let idx = profiler.record_comm_start("allreduce", 1024, "AllReduce");
|
|
std::thread::sleep(std::time::Duration::from_millis(5));
|
|
profiler.record_comm_end(idx);
|
|
|
|
let events = profiler.get_comm_events();
|
|
assert_eq!(events.len(), 1);
|
|
assert_eq!(events[0].bytes, 1024);
|
|
assert!(events[0].end_time.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_summary() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
let id = profiler.start_range("op1", "test", NvtxColor::GREEN);
|
|
profiler.end_range(id);
|
|
|
|
let summary = profiler.summary();
|
|
assert_eq!(summary.num_ranges, 1);
|
|
assert!(summary.num_ops >= 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chrome_trace_export() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
let id = profiler.start_range("test_op", "category", NvtxColor::GREEN);
|
|
profiler.end_range(id);
|
|
|
|
let trace = profiler.export_chrome_trace();
|
|
assert!(trace.starts_with('['));
|
|
assert!(trace.ends_with(']'));
|
|
assert!(trace.contains("test_op"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_profiler_clear() {
|
|
let config = ProfilingConfig::default();
|
|
let profiler = Profiler::new(config);
|
|
|
|
let id = profiler.start_range("test", "test", NvtxColor::GREEN);
|
|
profiler.end_range(id);
|
|
|
|
profiler.clear();
|
|
|
|
assert!(profiler.get_completed_ranges().is_empty());
|
|
assert!(profiler.get_op_stats().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_profiler() {
|
|
let profiler = shared_profiler(ProfilingConfig::default());
|
|
|
|
let id = profiler.start_range("test", "test", NvtxColor::GREEN);
|
|
profiler.end_range(id);
|
|
|
|
assert_eq!(profiler.get_completed_ranges().len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_disabled_profiler() {
|
|
let config = ProfilingConfig {
|
|
level: ProfilingLevel::Off,
|
|
..Default::default()
|
|
};
|
|
let profiler = Profiler::new(config);
|
|
|
|
assert!(!profiler.is_enabled());
|
|
|
|
let id = profiler.start_range("test", "test", NvtxColor::GREEN);
|
|
assert_eq!(id, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_helpers() {
|
|
let profiler = shared_profiler(ProfilingConfig::default());
|
|
|
|
let result = profile_forward(&profiler, "forward", || 42);
|
|
|
|
assert_eq!(result, 42);
|
|
assert_eq!(profiler.get_completed_ranges().len(), 1);
|
|
}
|
|
}
|