Count shaded fragments via Vulkan query and manual counting

occlusion.cc

@@ -150,6 +150,23 @@ VkCommandBuffer CreateVkCommandBuffer(VkDevice device, VkCommandPool command_poo
   return command_buffer;
 }
 
+VkQueryPool CreateVkQueryPool(VkDevice device) {
+  VkQueryPoolCreateInfo query_pool_create_info = {};
+  query_pool_create_info.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO;
+  query_pool_create_info.pNext = nullptr;
+  query_pool_create_info.queryType = VK_QUERY_TYPE_OCCLUSION;
+  query_pool_create_info.queryCount = 1;
+  query_pool_create_info.flags = 0;
+  query_pool_create_info.pipelineStatistics = 0;
+
+  VkQueryPool query_pool;
+  VkResult result = vkCreateQueryPool(device, &query_pool_create_info, nullptr, &query_pool);
+  if (result != VK_SUCCESS) {
+    ERROR("Unable to create VkQueryPool: " << result);
+  }
+  return query_pool;
+}
+
 VkRenderPass CreateVkRenderPass(VkDevice device) {
   VkAttachmentDescription attachment_description = {};
   attachment_description.format = kVulkanFormat;
@@ -472,6 +489,15 @@ VkFramebuffer CreateVkFramebuffer(VkDevice device, VkRenderPass render_pass,
   return framebuffer;
 }
 
+void BeginQuery(VkCommandBuffer command_buffer, VkQueryPool query_pool) {
+  vkCmdResetQueryPool(command_buffer, query_pool, 0, 1);
+  vkCmdBeginQuery(command_buffer, query_pool, 0, 0);
+}
+
+void EndQuery(VkCommandBuffer command_buffer, VkQueryPool query_pool) {
+  vkCmdEndQuery(command_buffer, query_pool, 0);
+}
+
 void BeginCommandBuffer(VkCommandBuffer command_buffer) {
   VkCommandBufferBeginInfo command_buffer_begin_info = {};
   command_buffer_begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
@@ -509,8 +535,9 @@ void WaitForIdle(VkQueue queue) {
 }
 
 void Draw(VkDevice device, VkQueue queue, VkRenderPass render_pass, VkPipeline pipeline,
-    VkFramebuffer framebuffer, VkCommandBuffer command_buffer) {
+    VkFramebuffer framebuffer, VkCommandBuffer command_buffer, VkQueryPool query_pool) {
   BeginCommandBuffer(command_buffer);
+  BeginQuery(command_buffer, query_pool);
   VkRenderPassBeginInfo render_pass_begin_info = {};
   render_pass_begin_info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
   render_pass_begin_info.renderPass = render_pass;
@@ -521,9 +548,9 @@ void Draw(VkDevice device, VkQueue queue, VkRenderPass render_pass, VkPipeline p
   render_pass_begin_info.renderArea = render_area;
   render_pass_begin_info.clearValueCount = 1;
   VkClearValue clear_value = {};
-  clear_value.color.float32[0] = 1.0f;
-  clear_value.color.float32[1] = 1.0f;
-  clear_value.color.float32[2] = 0.0f;
+  clear_value.color.float32[0] = 0.0f;
+  clear_value.color.float32[1] = 0.0f;
+  clear_value.color.float32[2] = 1.0f;
   clear_value.color.float32[3] = 1.0f;
   render_pass_begin_info.pClearValues = &clear_value;
   vkCmdBeginRenderPass(command_buffer, &render_pass_begin_info, VK_SUBPASS_CONTENTS_INLINE);
@@ -531,6 +558,7 @@ void Draw(VkDevice device, VkQueue queue, VkRenderPass render_pass, VkPipeline p
   vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
   vkCmdDraw(command_buffer, 3, 1, 0, 0);
   vkCmdEndRenderPass(command_buffer);
+  EndQuery(command_buffer, query_pool);
 }
 
 void BlitToScanoutImage(VkQueue queue, VkCommandBuffer command_buffer, VkImage source_image,
@@ -594,6 +622,19 @@ uint8_t* MapVkDeviceMemory(VkDevice device, VkDeviceMemory image_memory) {
   return static_cast<uint8_t*>(mapped_memory);
 }
 
+uint32_t GetQueryPoolResults(VkDevice device, VkQueryPool query_pool) {
+  uint32_t passed_fragments = 0;
+  VkResult result = vkGetQueryPoolResults(device, query_pool, 0, 1,  sizeof(passed_fragments),
+      &passed_fragments, sizeof(uint32_t), VK_QUERY_RESULT_WAIT_BIT);
+  // TODO(brkho): Currently, we query for results with the VK_QUERY_RESULT_WAIT_BIT flag set which
+  // blocks execution until the results are ready. We should instead poll and check that result
+  // isn't VK_NOT_READY.
+  if (result != VK_SUCCESS) {
+    ERROR("Unable to get query pool results: " << result);
+  }
+  return passed_fragments;
+}
+
 void WriteImage(uint8_t* pixels) {
   FILE *f = fopen("occlusion.png", "wb");
 
@@ -609,6 +650,20 @@ void WriteImage(uint8_t* pixels) {
   fclose(f);
 }
 
+uint32_t GetShadedFragmentsCountManual(uint8_t* pixels) {
+  uint32_t count = 0;
+  for (size_t i = 0; i < kHeight * kWidth * 4; i += 4) {
+    uint8_t r = pixels[i];
+    uint8_t g = pixels[i + 1];
+    uint8_t b = pixels[i + 2];
+    uint8_t a = pixels[i + 3];
+    if (r == 255 && g == 0 && b == 0 && a == 255) {
+      count++;
+    }
+  }
+  return count;
+}
+
 int main() {
   VkInstance instance = CreateVkInstance();
   VkPhysicalDevice physical_device = ChooseVkPhysicalDevice(instance);
@@ -617,6 +672,7 @@ int main() {
   VkQueue queue = GetVkQueue(device, device_queue_family_index);
   VkCommandPool command_pool = CreateVkCommandPool(device, device_queue_family_index);
   VkCommandBuffer command_buffer = CreateVkCommandBuffer(device, command_pool);
+  VkQueryPool query_pool = CreateVkQueryPool(device);
 
   VkRenderPass render_pass = CreateVkRenderPass(device);
   VkPipeline pipeline = CreateVkPipeline(device, render_pass);
@@ -630,7 +686,7 @@ int main() {
   VkDeviceMemory scanout_image_memory =
       AllocateAndBindScanoutMemory(physical_device, device, scanout_image);
 
-  Draw(device, queue, render_pass, pipeline, framebuffer, command_buffer);
+  Draw(device, queue, render_pass, pipeline, framebuffer, command_buffer, query_pool);
   // Since the render target is created with VK_IMAGE_TILING_OPTIMAL, we need to copy it to a linear
   // format before scanning out.
   BlitToScanoutImage(queue, command_buffer, render_image, scanout_image);
@@ -639,5 +695,10 @@ int main() {
   uint8_t* pixels = MapVkDeviceMemory(device, scanout_image_memory);
   WriteImage(pixels);
 
+  uint32_t shaded_fragments_count_query = GetQueryPoolResults(device, query_pool);
+  uint32_t shaded_fragments_count_manual = GetShadedFragmentsCountManual(pixels);
+  std::cout << "[Occlusion query] shaded fragments: " << shaded_fragments_count_query << std::endl;
+  std::cout << "[Manual count] shaded fragments: " << shaded_fragments_count_manual << std::endl;
+
   return EXIT_SUCCESS;
 }

occlusion.frag

@@ -1,10 +1,8 @@
 #version 450
 #extension GL_ARB_separate_shader_objects : enable
 
-layout(location = 0) in vec3 fragColor;
-
 layout(location = 0) out vec4 outColor;
 
 void main() {
-    outColor = vec4(fragColor, 1.0);
+    outColor = vec4(1.0, 0.0, 0.0, 1.0);
 }

occlusion.vert

@@ -9,13 +9,6 @@ vec2 positions[3] = vec2[](
     vec2(-0.5, 0.5)
 );
 
-vec3 colors[3] = vec3[](
-    vec3(1.0, 0.0, 0.0),
-    vec3(0.0, 1.0, 0.0),
-    vec3(0.0, 0.0, 1.0)
-);
-
 void main() {
     gl_Position = vec4(positions[gl_VertexIndex], 0.0, 1.0);
-    fragColor = colors[gl_VertexIndex];
 }