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 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 (Some(..), true) => {}
62 (Some(start), false) => {
64 run_start = None;
65 unsafe { raw_encoder.reset_queries(query_set.raw(), start..idx as u32) };
66 }
67 (None, true) => {
69 run_start = Some(idx as u32);
70 }
71 (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#[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#[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#[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 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 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 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 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 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 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 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}