wgpu_core/command/
query.rs

1use alloc::{sync::Arc, vec, vec::Vec};
2use core::{iter, mem};
3
4#[cfg(feature = "trace")]
5use crate::command::Command as TraceCommand;
6use crate::{
7    command::{CommandBufferMutable, 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            write_timestamp(cmd_buf_data, hub, &cmd_enc, query_set_id, query_index)
370        })
371    }
372
373    pub fn command_encoder_resolve_query_set(
374        &self,
375        command_encoder_id: id::CommandEncoderId,
376        query_set_id: id::QuerySetId,
377        start_query: u32,
378        query_count: u32,
379        destination: id::BufferId,
380        destination_offset: BufferAddress,
381    ) -> Result<(), EncoderStateError> {
382        let hub = &self.hub;
383
384        let cmd_enc = hub.command_encoders.get(command_encoder_id);
385        let mut cmd_buf_data = cmd_enc.data.lock();
386        cmd_buf_data.record_with(|cmd_buf_data| -> Result<(), QueryError> {
387            resolve_query_set(
388                cmd_buf_data,
389                hub,
390                &cmd_enc,
391                query_set_id,
392                start_query,
393                query_count,
394                destination,
395                destination_offset,
396            )
397        })
398    }
399}
400
401pub(super) fn write_timestamp(
402    cmd_buf_data: &mut CommandBufferMutable,
403    hub: &crate::hub::Hub,
404    cmd_enc: &Arc<CommandEncoder>,
405    query_set_id: id::QuerySetId,
406    query_index: u32,
407) -> Result<(), QueryError> {
408    #[cfg(feature = "trace")]
409    if let Some(ref mut list) = cmd_buf_data.trace_commands {
410        list.push(TraceCommand::WriteTimestamp {
411            query_set_id,
412            query_index,
413        });
414    }
415
416    cmd_enc.device.check_is_valid()?;
417
418    cmd_enc
419        .device
420        .require_features(wgt::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS)?;
421
422    let raw_encoder = cmd_buf_data.encoder.open()?;
423
424    let query_set = hub.query_sets.get(query_set_id).get()?;
425    query_set.same_device_as(cmd_enc.as_ref())?;
426
427    query_set.validate_and_write_timestamp(raw_encoder, query_index, None)?;
428
429    cmd_buf_data.trackers.query_sets.insert_single(query_set);
430
431    Ok(())
432}
433
434pub(super) fn resolve_query_set(
435    cmd_buf_data: &mut CommandBufferMutable,
436    hub: &crate::hub::Hub,
437    cmd_enc: &Arc<CommandEncoder>,
438    query_set_id: id::QuerySetId,
439    start_query: u32,
440    query_count: u32,
441    destination: id::BufferId,
442    destination_offset: BufferAddress,
443) -> Result<(), QueryError> {
444    #[cfg(feature = "trace")]
445    if let Some(ref mut list) = cmd_buf_data.trace_commands {
446        list.push(TraceCommand::ResolveQuerySet {
447            query_set_id,
448            start_query,
449            query_count,
450            destination,
451            destination_offset,
452        });
453    }
454
455    cmd_enc.device.check_is_valid()?;
456
457    if destination_offset % wgt::QUERY_RESOLVE_BUFFER_ALIGNMENT != 0 {
458        return Err(QueryError::Resolve(ResolveError::BufferOffsetAlignment));
459    }
460
461    let query_set = hub.query_sets.get(query_set_id).get()?;
462
463    query_set.same_device_as(cmd_enc.as_ref())?;
464
465    let dst_buffer = hub.buffers.get(destination).get()?;
466
467    dst_buffer.same_device_as(cmd_enc.as_ref())?;
468
469    let snatch_guard = dst_buffer.device.snatchable_lock.read();
470    dst_buffer.check_destroyed(&snatch_guard)?;
471
472    let dst_pending = cmd_buf_data
473        .trackers
474        .buffers
475        .set_single(&dst_buffer, wgt::BufferUses::COPY_DST);
476    let dst_barrier = dst_pending.map(|pending| pending.into_hal(&dst_buffer, &snatch_guard));
477
478    dst_buffer
479        .check_usage(wgt::BufferUsages::QUERY_RESOLVE)
480        .map_err(ResolveError::MissingBufferUsage)?;
481
482    let end_query = u64::from(start_query)
483        .checked_add(u64::from(query_count))
484        .expect("`u64` overflow from adding two `u32`s, should be unreachable");
485    if end_query > u64::from(query_set.desc.count) {
486        return Err(ResolveError::QueryOverrun {
487            start_query,
488            end_query,
489            query_set_size: query_set.desc.count,
490        }
491        .into());
492    }
493    let end_query =
494        u32::try_from(end_query).expect("`u32` overflow for `end_query`, which should be `u32`");
495
496    let elements_per_query = match query_set.desc.ty {
497        wgt::QueryType::Occlusion => 1,
498        wgt::QueryType::PipelineStatistics(ps) => ps.bits().count_ones(),
499        wgt::QueryType::Timestamp => 1,
500    };
501    let stride = elements_per_query * wgt::QUERY_SIZE;
502    let bytes_used: BufferAddress = u64::from(stride)
503        .checked_mul(u64::from(query_count))
504        .expect("`stride` * `query_count` overflowed `u32`, should be unreachable");
505
506    let buffer_start_offset = destination_offset;
507    let buffer_end_offset = buffer_start_offset
508        .checked_add(bytes_used)
509        .filter(|buffer_end_offset| *buffer_end_offset <= dst_buffer.size)
510        .ok_or(ResolveError::BufferOverrun {
511            start_query,
512            end_query,
513            stride,
514            buffer_size: dst_buffer.size,
515            buffer_start_offset,
516            bytes_used,
517        })?;
518
519    // TODO(https://github.com/gfx-rs/wgpu/issues/3993): Need to track initialization state.
520    cmd_buf_data.buffer_memory_init_actions.extend(
521        dst_buffer.initialization_status.read().create_action(
522            &dst_buffer,
523            buffer_start_offset..buffer_end_offset,
524            MemoryInitKind::ImplicitlyInitialized,
525        ),
526    );
527
528    let raw_dst_buffer = dst_buffer.try_raw(&snatch_guard)?;
529    let raw_encoder = cmd_buf_data.encoder.open()?;
530    unsafe {
531        raw_encoder.transition_buffers(dst_barrier.as_slice());
532        raw_encoder.copy_query_results(
533            query_set.raw(),
534            start_query..end_query,
535            raw_dst_buffer,
536            destination_offset,
537            wgt::BufferSize::new_unchecked(stride as u64),
538        );
539    }
540
541    if matches!(query_set.desc.ty, wgt::QueryType::Timestamp) {
542        // Timestamp normalization is only needed for timestamps.
543        cmd_enc
544            .device
545            .timestamp_normalizer
546            .get()
547            .unwrap()
548            .normalize(
549                &snatch_guard,
550                raw_encoder,
551                &mut cmd_buf_data.trackers.buffers,
552                dst_buffer
553                    .timestamp_normalization_bind_group
554                    .get(&snatch_guard)
555                    .unwrap(),
556                &dst_buffer,
557                destination_offset,
558                query_count,
559            );
560    }
561
562    cmd_buf_data.trackers.query_sets.insert_single(query_set);
563
564    Ok(())
565}