trustformers_debug/profiler/
gpu.rs1#![allow(dead_code)]
6
7use anyhow::Result;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::time::Duration;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct GpuKernelProfile {
15 pub kernel_name: String,
16 pub grid_size: (u32, u32, u32),
17 pub block_size: (u32, u32, u32),
18 pub shared_memory_bytes: usize,
19 pub registers_per_thread: u32,
20 pub occupancy: f64,
21 pub execution_time: Duration,
22 pub memory_bandwidth_gb_s: f64,
23 pub compute_utilization: f64,
24 pub stream_id: i32,
25}
26
27#[derive(Debug)]
29pub struct GpuProfiler {
30 device_count: i32,
34 pub(crate) active_streams: HashMap<i32, Vec<GpuKernelProfile>>,
35 memory_pools: HashMap<i32, GpuMemoryPool>,
36}
37
38#[derive(Debug)]
39pub struct GpuMemoryPool {
40 device_id: i32,
41 total_memory: usize,
42 free_memory: usize,
43 fragmentation_score: f64,
44}
45
46impl GpuProfiler {
47 pub fn new() -> Result<Self> {
55 Ok(Self {
56 device_count: 0,
57 active_streams: HashMap::new(),
58 memory_pools: HashMap::new(),
59 })
60 }
61
62 pub fn profile_kernel(&mut self, kernel_profile: GpuKernelProfile) {
63 self.active_streams
64 .entry(kernel_profile.stream_id)
65 .or_default()
66 .push(kernel_profile);
67 }
68
69 pub fn get_gpu_utilization(&self, device_id: i32) -> f64 {
70 if let Some(kernels) = self.active_streams.get(&device_id) {
72 if kernels.is_empty() {
73 0.0
74 } else {
75 kernels.iter().map(|k| k.compute_utilization).sum::<f64>() / kernels.len() as f64
76 }
77 } else {
78 0.0
79 }
80 }
81}
82
83#[derive(Debug, Serialize, Deserialize)]
84pub struct GpuKernelSummary {
85 pub total_kernels: usize,
86 pub total_execution_time: Duration,
87 pub avg_occupancy: f64,
88 pub avg_compute_utilization: f64,
89 pub slowest_kernels: Vec<String>,
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95
96 #[test]
97 fn test_gpu_profiler_new() {
98 let profiler = GpuProfiler::new();
99 assert!(profiler.is_ok());
100 }
101
102 #[test]
103 fn test_gpu_profiler_utilization_empty() {
104 let profiler = GpuProfiler::new().expect("should create profiler");
105 assert!((profiler.get_gpu_utilization(0) - 0.0).abs() < 1e-9);
106 }
107
108 #[test]
109 fn test_gpu_profiler_profile_kernel() {
110 let mut profiler = GpuProfiler::new().expect("should create profiler");
111 let kernel = GpuKernelProfile {
112 kernel_name: "matmul".to_string(),
113 grid_size: (128, 1, 1),
114 block_size: (256, 1, 1),
115 shared_memory_bytes: 4096,
116 registers_per_thread: 32,
117 occupancy: 0.85,
118 execution_time: Duration::from_micros(500),
119 memory_bandwidth_gb_s: 300.0,
120 compute_utilization: 0.9,
121 stream_id: 0,
122 };
123 profiler.profile_kernel(kernel);
124 let util = profiler.get_gpu_utilization(0);
125 assert!((util - 0.9).abs() < 1e-9);
126 }
127
128 #[test]
129 fn test_gpu_profiler_multiple_kernels_avg() {
130 let mut profiler = GpuProfiler::new().expect("should create profiler");
131 for util in [0.8, 0.6] {
132 profiler.profile_kernel(GpuKernelProfile {
133 kernel_name: "kern".to_string(),
134 grid_size: (1, 1, 1),
135 block_size: (1, 1, 1),
136 shared_memory_bytes: 0,
137 registers_per_thread: 0,
138 occupancy: 0.5,
139 execution_time: Duration::from_micros(100),
140 memory_bandwidth_gb_s: 0.0,
141 compute_utilization: util,
142 stream_id: 0,
143 });
144 }
145 let avg = profiler.get_gpu_utilization(0);
146 assert!((avg - 0.7).abs() < 1e-9);
147 }
148}