wgpu_primitives/
profiling.rs

1use std::time::Duration;
2
3use futures::channel::oneshot;
4
5use crate::Error;
6
7const TIMESTAMP_SIZE_BYTES: u64 = size_of::<u64>() as u64;
8
9/// GPU time measured for one labeled compute dispatch.
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct GpuTimestampSpan {
12    pub label: String,
13    pub duration: Duration,
14}
15
16/// Timestamp-query measurements for one primitive invocation.
17#[derive(Clone, Debug, Default, PartialEq, Eq)]
18pub struct GpuProfile {
19    /// Time from the beginning of the first dispatch through the end of the last dispatch.
20    pub gpu_elapsed: Duration,
21    /// Sum of the measured compute-pass durations.
22    pub dispatch_time: Duration,
23    pub spans: Vec<GpuTimestampSpan>,
24}
25
26impl GpuProfile {
27    pub(crate) fn empty() -> Self {
28        Self::default()
29    }
30}
31
32pub(crate) struct TimestampRecorder {
33    query_set: wgpu::QuerySet,
34    resolve_buffer: wgpu::Buffer,
35    readback_buffer: wgpu::Buffer,
36    labels: Vec<String>,
37    span_capacity: u32,
38    timestamp_period_ns: f64,
39}
40
41impl TimestampRecorder {
42    pub(crate) fn new(
43        device: &wgpu::Device,
44        queue: &wgpu::Queue,
45        span_capacity: u32,
46    ) -> Result<Self, Error> {
47        if !device.features().contains(wgpu::Features::TIMESTAMP_QUERY) {
48            return Err(Error::TimestampQueriesUnsupported);
49        }
50
51        let query_count = span_capacity.checked_mul(2).ok_or(Error::SizeOverflow)?;
52        let size_bytes = u64::from(query_count)
53            .checked_mul(TIMESTAMP_SIZE_BYTES)
54            .ok_or(Error::SizeOverflow)?;
55
56        let query_set = device.create_query_set(&wgpu::QuerySetDescriptor {
57            label: Some("Primitive Timestamp Queries"),
58            ty: wgpu::QueryType::Timestamp,
59            count: query_count,
60        });
61        let resolve_buffer = device.create_buffer(&wgpu::BufferDescriptor {
62            label: Some("Primitive Timestamp Resolve"),
63            size: size_bytes,
64            usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
65            mapped_at_creation: false,
66        });
67        let readback_buffer = device.create_buffer(&wgpu::BufferDescriptor {
68            label: Some("Primitive Timestamp Readback"),
69            size: size_bytes,
70            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
71            mapped_at_creation: false,
72        });
73
74        Ok(Self {
75            query_set,
76            resolve_buffer,
77            readback_buffer,
78            labels: Vec::with_capacity(span_capacity as usize),
79            span_capacity,
80            timestamp_period_ns: f64::from(queue.get_timestamp_period()),
81        })
82    }
83
84    fn reserve(&mut self, label: String) -> wgpu::ComputePassTimestampWrites<'_> {
85        assert!(
86            self.labels.len() < self.span_capacity as usize,
87            "timestamp span capacity must match the recorded dispatch count"
88        );
89        let beginning = self.labels.len() as u32 * 2;
90        self.labels.push(label);
91
92        wgpu::ComputePassTimestampWrites {
93            query_set: &self.query_set,
94            beginning_of_pass_write_index: Some(beginning),
95            end_of_pass_write_index: Some(beginning + 1),
96        }
97    }
98
99    pub(crate) fn resolve(&self, encoder: &mut wgpu::CommandEncoder) {
100        let query_count = self.labels.len() as u32 * 2;
101        let size_bytes = u64::from(query_count) * TIMESTAMP_SIZE_BYTES;
102        encoder.resolve_query_set(&self.query_set, 0..query_count, &self.resolve_buffer, 0);
103        encoder.copy_buffer_to_buffer(
104            &self.resolve_buffer,
105            0,
106            &self.readback_buffer,
107            0,
108            size_bytes,
109        );
110    }
111
112    pub(crate) async fn read(
113        self,
114        device: &wgpu::Device,
115        submission: wgpu::SubmissionIndex,
116    ) -> Result<GpuProfile, Error> {
117        let slice = self.readback_buffer.slice(..);
118        let (sender, receiver) = oneshot::channel();
119        slice.map_async(wgpu::MapMode::Read, move |result| {
120            let _ = sender.send(result);
121        });
122        device.poll(wgpu::PollType::Wait {
123            submission_index: Some(submission),
124            timeout: None,
125        })?;
126        receiver.await.map_err(|_| Error::ReadbackChannelClosed)??;
127
128        let timestamps: Vec<u64> = {
129            let data = slice.get_mapped_range();
130            bytemuck::cast_slice(&data).to_vec()
131        };
132        self.readback_buffer.unmap();
133
134        Ok(build_profile(
135            self.labels,
136            &timestamps,
137            self.timestamp_period_ns,
138        ))
139    }
140}
141
142pub(crate) fn record_compute_pass(
143    encoder: &mut wgpu::CommandEncoder,
144    pass_label: &'static str,
145    profile_label: Option<String>,
146    profiler: Option<&mut TimestampRecorder>,
147    record: impl FnOnce(&mut wgpu::ComputePass<'_>),
148) {
149    let timestamp_writes = profiler
150        .map(|profiler| profiler.reserve(profile_label.unwrap_or_else(|| pass_label.to_owned())));
151    let descriptor = wgpu::ComputePassDescriptor {
152        label: Some(pass_label),
153        timestamp_writes,
154    };
155    let mut pass = encoder.begin_compute_pass(&descriptor);
156    record(&mut pass);
157}
158
159fn build_profile(labels: Vec<String>, timestamps: &[u64], period_ns: f64) -> GpuProfile {
160    let spans: Vec<_> = labels
161        .into_iter()
162        .zip(timestamps.chunks_exact(2))
163        .map(|(label, timestamps)| GpuTimestampSpan {
164            label,
165            duration: ticks_to_duration(timestamps[1].saturating_sub(timestamps[0]), period_ns),
166        })
167        .collect();
168    let dispatch_time = spans
169        .iter()
170        .map(|span| span.duration)
171        .fold(Duration::ZERO, |total, duration| total + duration);
172    let gpu_elapsed = match (timestamps.first(), timestamps.last()) {
173        (Some(first), Some(last)) => ticks_to_duration(last.saturating_sub(*first), period_ns),
174        _ => Duration::ZERO,
175    };
176
177    GpuProfile {
178        gpu_elapsed,
179        dispatch_time,
180        spans,
181    }
182}
183
184fn ticks_to_duration(ticks: u64, period_ns: f64) -> Duration {
185    Duration::from_secs_f64(ticks as f64 * period_ns / 1_000_000_000.0)
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[test]
193    fn builds_labeled_profile_from_timestamp_pairs() {
194        let profile = build_profile(
195            vec!["reduce".to_owned(), "scatter".to_owned()],
196            &[10, 20, 25, 45],
197            2.0,
198        );
199
200        assert_eq!(profile.spans.len(), 2);
201        assert_eq!(profile.spans[0].duration, Duration::from_nanos(20));
202        assert_eq!(profile.spans[1].duration, Duration::from_nanos(40));
203        assert_eq!(profile.dispatch_time, Duration::from_nanos(60));
204        assert_eq!(profile.gpu_elapsed, Duration::from_nanos(70));
205    }
206}
wgpu_primitives/profiling.rs

wgpu_primitives/
profiling.rs
