Skip to content

Commit

Permalink
Add queries
Browse files Browse the repository at this point in the history
  • Loading branch information
z2oh committed Jul 19, 2020
1 parent 8e4d001 commit 8c53da9
Show file tree
Hide file tree
Showing 4 changed files with 356 additions and 7 deletions.
37 changes: 37 additions & 0 deletions examples/hello-compute/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,13 @@ async fn execute_gpu(numbers: Vec<u32>) -> Vec<u32> {
.await
.unwrap();

let pipeline_statistics_queries = vec![wgt::PipelineStatisticName::ComputeShaderInvocations];

let query_set = device.create_query_set(&wgpu::QuerySetDescriptor {
type_: wgpu::QueryType::PipelineStatistics(&pipeline_statistics_queries),
count: 1,
});

let cs_module = device.create_shader_module(wgpu::include_spirv!("shader.comp.spv"));

let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
Expand All @@ -52,6 +59,13 @@ async fn execute_gpu(numbers: Vec<u32>) -> Vec<u32> {
mapped_at_creation: false,
});

let query_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: None,
size: 16,
usage: wgpu::BufferUsage::MAP_READ | wgpu::BufferUsage::COPY_DST,
mapped_at_creation: false,
});

let storage_buffer = device.create_buffer_with_data(
bytemuck::cast_slice(&numbers),
wgpu::BufferUsage::STORAGE | wgpu::BufferUsage::COPY_DST | wgpu::BufferUsage::COPY_SRC,
Expand Down Expand Up @@ -96,23 +110,46 @@ async fn execute_gpu(numbers: Vec<u32>) -> Vec<u32> {
device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
{
let mut cpass = encoder.begin_compute_pass();
cpass.begin_pipeline_statistics_query(&query_set, 0);
cpass.set_pipeline(&compute_pipeline);
cpass.set_bind_group(0, &bind_group, &[]);
cpass.dispatch(numbers.len() as u32, 1, 1);
cpass.end_pipeline_statistics_query();
}
encoder.copy_buffer_to_buffer(&storage_buffer, 0, &staging_buffer, 0, size);
encoder.resolve_query_set(&query_set, 0, 1, &query_buffer, 0);

queue.submit(Some(encoder.finish()));

// Note that we're not calling `.await` here.
let buffer_slice = staging_buffer.slice(..);
let buffer_future = buffer_slice.map_async(wgpu::MapMode::Read);

let query_buffer_slice = query_buffer.slice(..);
let query_buffer_future = query_buffer_slice.map_async(wgpu::MapMode::Read);

// Poll the device in a blocking manner so that our future resolves.
// In an actual application, `device.poll(...)` should
// be called in an event loop or on another thread.
device.poll(wgpu::Maintain::Wait);

if let Ok(()) = query_buffer_future.await {
let data = query_buffer_slice.get_mapped_range();
let result: Vec<_> = data
.chunks_exact(8)
.map(|b| u64::from_ne_bytes(b.try_into().unwrap()))
.collect();

println!("raw query data: {:?}", result);

// With the current interface, we have to make sure all mapped views are
// dropped before we unmap the buffer.
drop(data);
query_buffer.unmap();
} else {
panic!("failed to retrieve query information from gpu");
}

if let Ok(()) = buffer_future.await {
let data = buffer_slice.get_mapped_range();
let result = data
Expand Down
76 changes: 76 additions & 0 deletions src/backend/direct.rs
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,7 @@ impl crate::Context for Context {
type CommandBufferId = wgc::id::CommandBufferId;
type RenderBundleEncoderId = wgc::command::RenderBundleEncoder;
type RenderBundleId = wgc::id::RenderBundleId;
type QuerySetId = wgc::id::QuerySetId;
type SurfaceId = wgc::id::SurfaceId;
type SwapChainId = wgc::id::SwapChainId;

Expand Down Expand Up @@ -775,6 +776,18 @@ impl crate::Context for Context {
wgc::command::RenderBundleEncoder::new(desc, *device, None).unwrap()
}

fn device_create_query_set(
&self,
device: &Self::DeviceId,
desc: &wgt::QuerySetDescriptor,
) -> Self::QuerySetId {
gfx_select!(*device => self.device_create_query_set(
*device,
&desc,
PhantomData
))
}

fn device_drop(&self, device: &Self::DeviceId) {
#[cfg(not(target_arch = "wasm32"))]
gfx_select!(*device => self.device_poll(*device, true));
Expand Down Expand Up @@ -925,6 +938,9 @@ impl crate::Context for Context {
fn render_bundle_drop(&self, render_bundle: &Self::RenderBundleId) {
gfx_select!(*render_bundle => self.render_bundle_destroy(*render_bundle))
}
fn query_set_drop(&self, query_set: &Self::QuerySetId) {
gfx_select!(*query_set => self.query_set_destroy(*query_set))
}
fn compute_pipeline_drop(&self, pipeline: &Self::ComputePipelineId) {
gfx_select!(*pipeline => self.compute_pipeline_destroy(*pipeline))
}
Expand Down Expand Up @@ -1000,6 +1016,66 @@ impl crate::Context for Context {
.unwrap()
}

fn command_encoder_write_timestamp(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
query_index: u32,
pipeline_stage: wgc::PipelineStage,
) {
gfx_select!(*encoder => self.command_encoder_write_timestamp(
*encoder,
*query_set,
query_index,
pipeline_stage
))
}

fn command_encoder_begin_pipeline_statistics_query(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
query_index: u32,
) {
gfx_select!(*encoder => self.command_encoder_begin_pipeline_statistics_query(
*encoder,
*query_set,
query_index
))
}

fn command_encoder_end_pipeline_statistics_query(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
query_index: u32,
) {
gfx_select!(*encoder => self.command_encoder_end_pipeline_statistics_query(
*encoder,
*query_set,
query_index
))
}

fn command_encoder_resolve_query_set(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
first_query: u32,
query_count: u32,
destination: &Self::BufferId,
destination_offset: wgt::BufferAddress,
) {
gfx_select!(*encoder => self.command_encoder_resolve_query_set(
*encoder,
*query_set,
first_query,
query_count,
*destination,
destination_offset
))
}

fn command_encoder_begin_compute_pass(
&self,
encoder: &Self::CommandEncoderId,
Expand Down
59 changes: 59 additions & 0 deletions src/backend/web.rs
Original file line number Diff line number Diff line change
Expand Up @@ -688,6 +688,7 @@ impl crate::Context for Context {
type CommandBufferId = Sendable<web_sys::GpuCommandBuffer>;
type RenderBundleEncoderId = RenderPass; //web_sys::GpuRenderBundleEncoder;
type RenderBundleId = (); //web_sys::GpuRenderBundle;
type QuerySetId = ();
type SurfaceId = Sendable<web_sys::GpuCanvasContext>;
type SwapChainId = Sendable<web_sys::GpuSwapChain>;

Expand Down Expand Up @@ -1226,6 +1227,9 @@ impl crate::Context for Context {
fn render_bundle_drop(&self, _render_bundle: &Self::RenderBundleId) {
// Dropped automatically
}
fn query_set_drop(&self, _query_set: &Self::QuerySetId) {
// Dropped automatically
}
fn compute_pipeline_drop(&self, _pipeline: &Self::ComputePipelineId) {
// Dropped automatically
}
Expand Down Expand Up @@ -1293,6 +1297,61 @@ impl crate::Context for Context {
)
}

fn command_encoder_write_timestamp(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
query_index: u32,
pipeline_stage: wgc::PipelineStage,
) {
encoder.write_timestamp(
query_set,
query_index,
)
}

fn command_encoder_begin_pipeline_statistics_query(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
query_index: u32,
) {
encoder.begin_pipeline_statistics_query(
query_set,
query_index,
)
}

fn command_encoder_end_pipeline_statistics_query(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
query_index: u32,
) {
encoder.end_pipeline_statistics_query(
query_set,
query_index,
)
}

fn command_encoder_resolve_query_set(
&self,
encoder: &Self::CommandEncoderId,
query_set: &Self::QuerySetId,
first_query: u32, // JTODO: QueryId instead of u32 ?
query_count: u32,
destination: &Self::BufferId,
destination_offset: BufferAddress,
) {
encoder.resolve_query_set(
query_set,
first_query,
query_count,
destination,
destination_offset,
)
}

fn command_encoder_begin_compute_pass(
&self,
encoder: &Self::CommandEncoderId,
Expand Down
Loading

0 comments on commit 8c53da9

Please sign in to comment.