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