From 27b808504879ae41abc8fccaa169892661b235aa Mon Sep 17 00:00:00 2001
From: Dzmitry Malyshau <kvarkus@gmail.com>
Date: Thu, 8 Jul 2021 23:11:20 -0400
Subject: [PATCH] hal/dx12: queries

---
 wgpu-hal/src/dx12/adapter.rs |  4 +++-
 wgpu-hal/src/dx12/command.rs | 32 ++++++++++++++++++++++++++++----
 wgpu-hal/src/dx12/device.rs  | 19 ++++++++++++++-----
 wgpu-hal/src/dx12/mod.rs     |  2 +-
 wgpu-types/src/lib.rs        |  4 ++--
 5 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/wgpu-hal/src/dx12/adapter.rs b/wgpu-hal/src/dx12/adapter.rs
index 17dc1a3910..8b17767747 100644
--- a/wgpu-hal/src/dx12/adapter.rs
+++ b/wgpu-hal/src/dx12/adapter.rs
@@ -145,7 +145,9 @@ impl super::Adapter {
             | wgt::Features::MULTI_DRAW_INDIRECT_COUNT
             | wgt::Features::ADDRESS_MODE_CLAMP_TO_BORDER
             | wgt::Features::NON_FILL_POLYGON_MODE
-            |wgt::Features::VERTEX_WRITABLE_STORAGE;
+            | wgt::Features::VERTEX_WRITABLE_STORAGE
+            | wgt::Features::TIMESTAMP_QUERY
+            | wgt::Features::PIPELINE_STATISTICS_QUERY;
 
         features.set(
             wgt::Features::CONSERVATIVE_RASTERIZATION,
diff --git a/wgpu-hal/src/dx12/command.rs b/wgpu-hal/src/dx12/command.rs
index bdf3bc22ba..628ed8a9c6 100644
--- a/wgpu-hal/src/dx12/command.rs
+++ b/wgpu-hal/src/dx12/command.rs
@@ -338,10 +338,26 @@ impl crate::CommandEncoder<super::Api> for super::CommandEncoder {
         }
     }
 
-    unsafe fn begin_query(&mut self, set: &super::QuerySet, index: u32) {}
-    unsafe fn end_query(&mut self, set: &super::QuerySet, index: u32) {}
-    unsafe fn write_timestamp(&mut self, set: &super::QuerySet, index: u32) {}
-    unsafe fn reset_queries(&mut self, set: &super::QuerySet, range: Range<u32>) {}
+    unsafe fn begin_query(&mut self, set: &super::QuerySet, index: u32) {
+        self.list
+            .unwrap()
+            .BeginQuery(set.raw.as_mut_ptr(), set.raw_ty, index);
+    }
+    unsafe fn end_query(&mut self, set: &super::QuerySet, index: u32) {
+        self.list
+            .unwrap()
+            .EndQuery(set.raw.as_mut_ptr(), set.raw_ty, index);
+    }
+    unsafe fn write_timestamp(&mut self, set: &super::QuerySet, index: u32) {
+        self.list.unwrap().EndQuery(
+            set.raw.as_mut_ptr(),
+            d3d12::D3D12_QUERY_TYPE_TIMESTAMP,
+            index,
+        );
+    }
+    unsafe fn reset_queries(&mut self, _set: &super::QuerySet, _range: Range<u32>) {
+        // nothing to do here
+    }
     unsafe fn copy_query_results(
         &mut self,
         set: &super::QuerySet,
@@ -350,6 +366,14 @@ impl crate::CommandEncoder<super::Api> for super::CommandEncoder {
         offset: wgt::BufferAddress,
         stride: wgt::BufferSize,
     ) {
+        self.list.unwrap().ResolveQueryData(
+            set.raw.as_mut_ptr(),
+            set.raw_ty,
+            range.start,
+            range.end - range.start,
+            buffer.resource.as_mut_ptr(),
+            offset,
+        );
     }
 
     // render
diff --git a/wgpu-hal/src/dx12/device.rs b/wgpu-hal/src/dx12/device.rs
index a466ac820c..59e0322052 100644
--- a/wgpu-hal/src/dx12/device.rs
+++ b/wgpu-hal/src/dx12/device.rs
@@ -1059,10 +1059,19 @@ impl crate::Device<super::Api> for super::Device {
         &self,
         desc: &wgt::QuerySetDescriptor<crate::Label>,
     ) -> Result<super::QuerySet, crate::DeviceError> {
-        let heap_ty = match desc.ty {
-            wgt::QueryType::Occlusion => native::QueryHeapType::Occlusion,
-            wgt::QueryType::PipelineStatistics(_) => native::QueryHeapType::PipelineStatistics,
-            wgt::QueryType::Timestamp => native::QueryHeapType::Timestamp,
+        let (heap_ty, raw_ty) = match desc.ty {
+            wgt::QueryType::Occlusion => (
+                native::QueryHeapType::Occlusion,
+                d3d12::D3D12_QUERY_TYPE_BINARY_OCCLUSION,
+            ),
+            wgt::QueryType::PipelineStatistics(_) => (
+                native::QueryHeapType::PipelineStatistics,
+                d3d12::D3D12_QUERY_TYPE_TIMESTAMP,
+            ),
+            wgt::QueryType::Timestamp => (
+                native::QueryHeapType::Timestamp,
+                d3d12::D3D12_QUERY_TYPE_PIPELINE_STATISTICS,
+            ),
         };
 
         let raw = self
@@ -1070,7 +1079,7 @@ impl crate::Device<super::Api> for super::Device {
             .create_query_heap(heap_ty, desc.count, 0)
             .into_device_result("Query heap creation")?;
 
-        Ok(super::QuerySet { raw, ty: desc.ty })
+        Ok(super::QuerySet { raw, raw_ty })
     }
     unsafe fn destroy_query_set(&self, set: super::QuerySet) {
         set.raw.destroy();
diff --git a/wgpu-hal/src/dx12/mod.rs b/wgpu-hal/src/dx12/mod.rs
index 71eb7299ea..b8d320beea 100644
--- a/wgpu-hal/src/dx12/mod.rs
+++ b/wgpu-hal/src/dx12/mod.rs
@@ -302,7 +302,7 @@ unsafe impl Sync for Sampler {}
 #[derive(Debug)]
 pub struct QuerySet {
     raw: native::QueryHeap,
-    ty: wgt::QueryType,
+    raw_ty: d3d12::D3D12_QUERY_TYPE,
 }
 
 unsafe impl Send for QuerySet {}
diff --git a/wgpu-types/src/lib.rs b/wgpu-types/src/lib.rs
index 3b17b82be6..bda8bc7e9b 100644
--- a/wgpu-types/src/lib.rs
+++ b/wgpu-types/src/lib.rs
@@ -204,7 +204,7 @@ bitflags::bitflags! {
         ///
         /// Supported Platforms:
         /// - Vulkan (works)
-        /// - DX12 (future)
+        /// - DX12 (works)
         ///
         /// This is a web and native feature.
         const TIMESTAMP_QUERY = 0x0000_0000_0000_0004;
@@ -219,7 +219,7 @@ bitflags::bitflags! {
         ///
         /// Supported Platforms:
         /// - Vulkan (works)
-        /// - DX12 (future)
+        /// - DX12 (works)
         ///
         /// This is a web and native feature.
         const PIPELINE_STATISTICS_QUERY = 0x0000_0000_0000_0008;