wgpu_core/command/
query.rs

1use alloc::{sync::Arc, vec, vec::Vec};
2use core::{iter, mem};
3
4#[cfg(feature = "trace")]
5use crate::device::trace::Command as TraceCommand;
6use crate::{
7    command::{CommandBuffer, CommandEncoderError},
8    device::{DeviceError, MissingFeatures},
9    global::Global,
10    id,
11    init_tracker::MemoryInitKind,
12    resource::{
13        DestroyedResourceError, InvalidResourceError, MissingBufferUsageError, ParentDevice,
14        QuerySet, Trackable,
15    },
16    track::{StatelessTracker, TrackerIndex},
17    FastHashMap,
18};
19use thiserror::Error;
20use wgt::BufferAddress;
21
22#[derive(Debug)]
23pub(crate) struct QueryResetMap {
24    map: FastHashMap<TrackerIndex, (Vec<bool>, Arc<QuerySet>)>,
25}
26impl QueryResetMap {
27    pub fn new() -> Self {
28        Self {
29            map: FastHashMap::default(),
30        }
31    }
32
33    pub fn use_query_set(&mut self, query_set: &Arc<QuerySet>, query: u32) -> bool {
34        let vec_pair = self
35            .map
36            .entry(query_set.tracker_index())
37            .or_insert_with(|| {
38                (
39                    vec![false; query_set.desc.count as usize],
40                    query_set.clone(),
41                )
42            });
43
44        mem::replace(&mut vec_pair.0[query as usize], true)
45    }
46
47    pub fn reset_queries(&mut self, raw_encoder: &mut dyn hal::DynCommandEncoder) {
48        for (_, (state, query_set)) in self.map.drain() {
49            debug_assert_eq!(state.len(), query_set.desc.count as usize);
50
51            // Need to find all "runs" of values which need resets. If the state vector is:
52            // [false, true, true, false, true], we want to reset [1..3, 4..5]. This minimizes
53            // the amount of resets needed.
54            let mut run_start: Option<u32> = None;
55            for (idx, value) in state.into_iter().chain(iter::once(false)).enumerate() {
56                match (run_start, value) {
57                    // We're inside of a run, do nothing
58                    (Some(..), true) => {}
59                    // We've hit the end of a run, dispatch a reset
60                    (Some(start), false) => {
61                        run_start = None;
62                        unsafe { raw_encoder.reset_queries(query_set.raw(), start..idx as u32) };
63                    }
64                    // We're starting a run
65                    (None, true) => {
66                        run_start = Some(idx as u32);
67                    }
68                    // We're in a run of falses, do nothing.
69                    (None, false) => {}
70                }
71            }
72        }
73    }
74}
75
76#[derive(Debug, Copy, Clone, PartialEq, Eq)]
77pub enum SimplifiedQueryType {
78    Occlusion,
79    Timestamp,
80    PipelineStatistics,
81}
82impl From<wgt::QueryType> for SimplifiedQueryType {
83    fn from(q: wgt::QueryType) -> Self {
84        match q {
85            wgt::QueryType::Occlusion => SimplifiedQueryType::Occlusion,
86            wgt::QueryType::Timestamp => SimplifiedQueryType::Timestamp,
87            wgt::QueryType::PipelineStatistics(..) => SimplifiedQueryType::PipelineStatistics,
88        }
89    }
90}
91
92/// Error encountered when dealing with queries
93#[derive(Clone, Debug, Error)]
94#[non_exhaustive]
95pub enum QueryError {
96    #[error(transparent)]
97    Device(#[from] DeviceError),
98    #[error(transparent)]
99    Encoder(#[from] CommandEncoderError),
100    #[error(transparent)]
101    MissingFeature(#[from] MissingFeatures),
102    #[error("Error encountered while trying to use queries")]
103    Use(#[from] QueryUseError),
104    #[error("Error encountered while trying to resolve a query")]
105    Resolve(#[from] ResolveError),
106    #[error(transparent)]
107    DestroyedResource(#[from] DestroyedResourceError),
108    #[error(transparent)]
109    InvalidResource(#[from] InvalidResourceError),
110}
111
112/// Error encountered while trying to use queries
113#[derive(Clone, Debug, Error)]
114#[non_exhaustive]
115pub enum QueryUseError {
116    #[error(transparent)]
117    Device(#[from] DeviceError),
118    #[error("Query {query_index} is out of bounds for a query set of size {query_set_size}")]
119    OutOfBounds {
120        query_index: u32,
121        query_set_size: u32,
122    },
123    #[error("Query {query_index} has already been used within the same renderpass. Queries must only be used once per renderpass")]
124    UsedTwiceInsideRenderpass { query_index: u32 },
125    #[error("Query {new_query_index} was started while query {active_query_index} was already active. No more than one statistic or occlusion query may be active at once")]
126    AlreadyStarted {
127        active_query_index: u32,
128        new_query_index: u32,
129    },
130    #[error("Query was stopped while there was no active query")]
131    AlreadyStopped,
132    #[error("A query of type {query_type:?} was started using a query set of type {set_type:?}")]
133    IncompatibleType {
134        set_type: SimplifiedQueryType,
135        query_type: SimplifiedQueryType,
136    },
137}
138
139/// Error encountered while trying to resolve a query.
140#[derive(Clone, Debug, Error)]
141#[non_exhaustive]
142pub enum ResolveError {
143    #[error(transparent)]
144    MissingBufferUsage(#[from] MissingBufferUsageError),
145    #[error("Resolve buffer offset has to be aligned to `QUERY_RESOLVE_BUFFER_ALIGNMENT")]
146    BufferOffsetAlignment,
147    #[error("Resolving queries {start_query}..{end_query} would overrun the query set of size {query_set_size}")]
148    QueryOverrun {
149        start_query: u32,
150        end_query: u64,
151        query_set_size: u32,
152    },
153    #[error("Resolving queries {start_query}..{end_query} ({stride} byte queries) will end up overrunning the bounds of the destination buffer of size {buffer_size} using offsets {buffer_start_offset}..(<start> + {bytes_used})")]
154    BufferOverrun {
155        start_query: u32,
156        end_query: u32,
157        stride: u32,
158        buffer_size: BufferAddress,
159        buffer_start_offset: BufferAddress,
160        bytes_used: BufferAddress,
161    },
162}
163
164impl QuerySet {
165    pub(crate) fn validate_query(
166        self: &Arc<Self>,
167        query_type: SimplifiedQueryType,
168        query_index: u32,
169        reset_state: Option<&mut QueryResetMap>,
170    ) -> Result<(), QueryUseError> {
171        // We need to defer our resets because we are in a renderpass,
172        // add the usage to the reset map.
173        if let Some(reset) = reset_state {
174            let used = reset.use_query_set(self, query_index);
175            if used {
176                return Err(QueryUseError::UsedTwiceInsideRenderpass { query_index });
177            }
178        }
179
180        let simple_set_type = SimplifiedQueryType::from(self.desc.ty);
181        if simple_set_type != query_type {
182            return Err(QueryUseError::IncompatibleType {
183                query_type,
184                set_type: simple_set_type,
185            });
186        }
187
188        if query_index >= self.desc.count {
189            return Err(QueryUseError::OutOfBounds {
190                query_index,
191                query_set_size: self.desc.count,
192            });
193        }
194
195        Ok(())
196    }
197
198    pub(super) fn validate_and_write_timestamp(
199        self: &Arc<Self>,
200        raw_encoder: &mut dyn hal::DynCommandEncoder,
201        query_index: u32,
202        reset_state: Option<&mut QueryResetMap>,
203    ) -> Result<(), QueryUseError> {
204        let needs_reset = reset_state.is_none();
205        self.validate_query(SimplifiedQueryType::Timestamp, query_index, reset_state)?;
206
207        unsafe {
208            // If we don't have a reset state tracker which can defer resets, we must reset now.
209            if needs_reset {
210                raw_encoder.reset_queries(self.raw(), query_index..(query_index + 1));
211            }
212            raw_encoder.write_timestamp(self.raw(), query_index);
213        }
214
215        Ok(())
216    }
217}
218
219pub(super) fn validate_and_begin_occlusion_query(
220    query_set: Arc<QuerySet>,
221    raw_encoder: &mut dyn hal::DynCommandEncoder,
222    tracker: &mut StatelessTracker<QuerySet>,
223    query_index: u32,
224    reset_state: Option<&mut QueryResetMap>,
225    active_query: &mut Option<(Arc<QuerySet>, u32)>,
226) -> Result<(), QueryUseError> {
227    let needs_reset = reset_state.is_none();
228    query_set.validate_query(SimplifiedQueryType::Occlusion, query_index, reset_state)?;
229
230    tracker.insert_single(query_set.clone());
231
232    if let Some((_old, old_idx)) = active_query.take() {
233        return Err(QueryUseError::AlreadyStarted {
234            active_query_index: old_idx,
235            new_query_index: query_index,
236        });
237    }
238    let (query_set, _) = &active_query.insert((query_set, query_index));
239
240    unsafe {
241        // If we don't have a reset state tracker which can defer resets, we must reset now.
242        if needs_reset {
243            raw_encoder.reset_queries(query_set.raw(), query_index..(query_index + 1));
244        }
245        raw_encoder.begin_query(query_set.raw(), query_index);
246    }
247
248    Ok(())
249}
250
251pub(super) fn end_occlusion_query(
252    raw_encoder: &mut dyn hal::DynCommandEncoder,
253    active_query: &mut Option<(Arc<QuerySet>, u32)>,
254) -> Result<(), QueryUseError> {
255    if let Some((query_set, query_index)) = active_query.take() {
256        unsafe { raw_encoder.end_query(query_set.raw(), query_index) };
257        Ok(())
258    } else {
259        Err(QueryUseError::AlreadyStopped)
260    }
261}
262
263pub(super) fn validate_and_begin_pipeline_statistics_query(
264    query_set: Arc<QuerySet>,
265    raw_encoder: &mut dyn hal::DynCommandEncoder,
266    tracker: &mut StatelessTracker<QuerySet>,
267    cmd_buf: &CommandBuffer,
268    query_index: u32,
269    reset_state: Option<&mut QueryResetMap>,
270    active_query: &mut Option<(Arc<QuerySet>, u32)>,
271) -> Result<(), QueryUseError> {
272    query_set.same_device_as(cmd_buf)?;
273
274    let needs_reset = reset_state.is_none();
275    query_set.validate_query(
276        SimplifiedQueryType::PipelineStatistics,
277        query_index,
278        reset_state,
279    )?;
280
281    tracker.insert_single(query_set.clone());
282
283    if let Some((_old, old_idx)) = active_query.take() {
284        return Err(QueryUseError::AlreadyStarted {
285            active_query_index: old_idx,
286            new_query_index: query_index,
287        });
288    }
289    let (query_set, _) = &active_query.insert((query_set, query_index));
290
291    unsafe {
292        // If we don't have a reset state tracker which can defer resets, we must reset now.
293        if needs_reset {
294            raw_encoder.reset_queries(query_set.raw(), query_index..(query_index + 1));
295        }
296        raw_encoder.begin_query(query_set.raw(), query_index);
297    }
298
299    Ok(())
300}
301
302pub(super) fn end_pipeline_statistics_query(
303    raw_encoder: &mut dyn hal::DynCommandEncoder,
304    active_query: &mut Option<(Arc<QuerySet>, u32)>,
305) -> Result<(), QueryUseError> {
306    if let Some((query_set, query_index)) = active_query.take() {
307        unsafe { raw_encoder.end_query(query_set.raw(), query_index) };
308        Ok(())
309    } else {
310        Err(QueryUseError::AlreadyStopped)
311    }
312}
313
314impl Global {
315    pub fn command_encoder_write_timestamp(
316        &self,
317        command_encoder_id: id::CommandEncoderId,
318        query_set_id: id::QuerySetId,
319        query_index: u32,
320    ) -> Result<(), QueryError> {
321        let hub = &self.hub;
322
323        let cmd_buf = hub
324            .command_buffers
325            .get(command_encoder_id.into_command_buffer_id());
326        let mut cmd_buf_data = cmd_buf.data.lock();
327        let mut cmd_buf_data_guard = cmd_buf_data.record()?;
328        let cmd_buf_data = &mut *cmd_buf_data_guard;
329
330        cmd_buf
331            .device
332            .require_features(wgt::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS)?;
333
334        #[cfg(feature = "trace")]
335        if let Some(ref mut list) = cmd_buf_data.commands {
336            list.push(TraceCommand::WriteTimestamp {
337                query_set_id,
338                query_index,
339            });
340        }
341
342        let raw_encoder = cmd_buf_data.encoder.open()?;
343
344        let query_set = hub.query_sets.get(query_set_id).get()?;
345
346        query_set.validate_and_write_timestamp(raw_encoder, query_index, None)?;
347
348        cmd_buf_data.trackers.query_sets.insert_single(query_set);
349
350        cmd_buf_data_guard.mark_successful();
351        Ok(())
352    }
353
354    pub fn command_encoder_resolve_query_set(
355        &self,
356        command_encoder_id: id::CommandEncoderId,
357        query_set_id: id::QuerySetId,
358        start_query: u32,
359        query_count: u32,
360        destination: id::BufferId,
361        destination_offset: BufferAddress,
362    ) -> Result<(), QueryError> {
363        let hub = &self.hub;
364
365        let cmd_buf = hub
366            .command_buffers
367            .get(command_encoder_id.into_command_buffer_id());
368        let mut cmd_buf_data = cmd_buf.data.lock();
369        let mut cmd_buf_data_guard = cmd_buf_data.record()?;
370        let cmd_buf_data = &mut *cmd_buf_data_guard;
371
372        #[cfg(feature = "trace")]
373        if let Some(ref mut list) = cmd_buf_data.commands {
374            list.push(TraceCommand::ResolveQuerySet {
375                query_set_id,
376                start_query,
377                query_count,
378                destination,
379                destination_offset,
380            });
381        }
382
383        if destination_offset % wgt::QUERY_RESOLVE_BUFFER_ALIGNMENT != 0 {
384            return Err(QueryError::Resolve(ResolveError::BufferOffsetAlignment));
385        }
386
387        let query_set = hub.query_sets.get(query_set_id).get()?;
388
389        query_set.same_device_as(cmd_buf.as_ref())?;
390
391        let dst_buffer = hub.buffers.get(destination).get()?;
392
393        dst_buffer.same_device_as(cmd_buf.as_ref())?;
394
395        let snatch_guard = dst_buffer.device.snatchable_lock.read();
396        dst_buffer.check_destroyed(&snatch_guard)?;
397
398        let dst_pending = cmd_buf_data
399            .trackers
400            .buffers
401            .set_single(&dst_buffer, wgt::BufferUses::COPY_DST);
402
403        let dst_barrier = dst_pending.map(|pending| pending.into_hal(&dst_buffer, &snatch_guard));
404
405        dst_buffer
406            .check_usage(wgt::BufferUsages::QUERY_RESOLVE)
407            .map_err(ResolveError::MissingBufferUsage)?;
408
409        let end_query = u64::from(start_query)
410            .checked_add(u64::from(query_count))
411            .expect("`u64` overflow from adding two `u32`s, should be unreachable");
412        if end_query > u64::from(query_set.desc.count) {
413            return Err(ResolveError::QueryOverrun {
414                start_query,
415                end_query,
416                query_set_size: query_set.desc.count,
417            }
418            .into());
419        }
420        let end_query = u32::try_from(end_query)
421            .expect("`u32` overflow for `end_query`, which should be `u32`");
422
423        let elements_per_query = match query_set.desc.ty {
424            wgt::QueryType::Occlusion => 1,
425            wgt::QueryType::PipelineStatistics(ps) => ps.bits().count_ones(),
426            wgt::QueryType::Timestamp => 1,
427        };
428        let stride = elements_per_query * wgt::QUERY_SIZE;
429        let bytes_used: BufferAddress = u64::from(stride)
430            .checked_mul(u64::from(query_count))
431            .expect("`stride` * `query_count` overflowed `u32`, should be unreachable");
432
433        let buffer_start_offset = destination_offset;
434        let buffer_end_offset = buffer_start_offset
435            .checked_add(bytes_used)
436            .filter(|buffer_end_offset| *buffer_end_offset <= dst_buffer.size)
437            .ok_or(ResolveError::BufferOverrun {
438                start_query,
439                end_query,
440                stride,
441                buffer_size: dst_buffer.size,
442                buffer_start_offset,
443                bytes_used,
444            })?;
445
446        // TODO(https://github.com/gfx-rs/wgpu/issues/3993): Need to track initialization state.
447        cmd_buf_data.buffer_memory_init_actions.extend(
448            dst_buffer.initialization_status.read().create_action(
449                &dst_buffer,
450                buffer_start_offset..buffer_end_offset,
451                MemoryInitKind::ImplicitlyInitialized,
452            ),
453        );
454
455        let raw_dst_buffer = dst_buffer.try_raw(&snatch_guard)?;
456        let raw_encoder = cmd_buf_data.encoder.open()?;
457        unsafe {
458            raw_encoder.transition_buffers(dst_barrier.as_slice());
459            raw_encoder.copy_query_results(
460                query_set.raw(),
461                start_query..end_query,
462                raw_dst_buffer,
463                destination_offset,
464                wgt::BufferSize::new_unchecked(stride as u64),
465            );
466        }
467
468        if matches!(query_set.desc.ty, wgt::QueryType::Timestamp) {
469            // Timestamp normalization is only needed for timestamps.
470            cmd_buf
471                .device
472                .timestamp_normalizer
473                .get()
474                .unwrap()
475                .normalize(
476                    &snatch_guard,
477                    raw_encoder,
478                    &mut cmd_buf_data.trackers.buffers,
479                    dst_buffer
480                        .timestamp_normalization_bind_group
481                        .get(&snatch_guard)
482                        .unwrap(),
483                    &dst_buffer,
484                    destination_offset,
485                    query_count,
486                );
487        }
488
489        cmd_buf_data.trackers.query_sets.insert_single(query_set);
490
491        cmd_buf_data_guard.mark_successful();
492        Ok(())
493    }
494}