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