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 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 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 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 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}