freedreno-occlusion-vulkan/occlusion.cc

#include <png.h>
#include <vulkan/vulkan.h>

#include <array>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

const uint32_t kWidth = 256;
const uint32_t kHeight = 256;
const VkFormat kVulkanFormat = VK_FORMAT_A8B8G8R8_UNORM_PACK32;
const VkFormat kVulkanDepthFormat = VK_FORMAT_D32_SFLOAT;
const std::string kVertexShaderPath = "occlusion.vert.spv";
const std::string kFragmentShaderPath = "occlusion.frag.spv";

#define ERROR(message) \
    std::cerr << message << std::endl; \
    std::exit(EXIT_FAILURE)

VkInstance CreateVkInstance() {
  VkApplicationInfo app_info = {};
  app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
  app_info.apiVersion = VK_API_VERSION_1_1;

  VkInstanceCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
  create_info.pApplicationInfo = &app_info;
  create_info.enabledExtensionCount = 0;
  create_info.ppEnabledExtensionNames = nullptr;
  create_info.enabledLayerCount = 0;
  create_info.ppEnabledLayerNames = nullptr;

  VkInstance instance;
  VkResult result = vkCreateInstance(&create_info, nullptr, &instance);
  if (result != VK_SUCCESS) {
    ERROR("Error creating VkInstance: " << result);
  }
  return instance;
}

VkPhysicalDevice ChooseVkPhysicalDevice(VkInstance instance) {
  uint32_t device_count = 0;
  VkResult result = vkEnumeratePhysicalDevices(instance, &device_count, nullptr);
  if (result != VK_SUCCESS) {
    ERROR("Error enumerating VkPhysicalDevices: " << result);
  }
  if (device_count == 0) {
    ERROR("No available VkPhysicalDevices");
  }

  std::vector<VkPhysicalDevice> devices(device_count);
  result = vkEnumeratePhysicalDevices(instance, &device_count, devices.data());
  if (result != VK_SUCCESS) {
    ERROR("Error fetching VkPhysicalDevices: " << result);
  }

  std::cout << "Found " << device_count << " device(s):" << std::endl;
  VkPhysicalDevice chosen_device = VK_NULL_HANDLE;
  for (VkPhysicalDevice device : devices) {
    VkPhysicalDeviceProperties device_props;
    vkGetPhysicalDeviceProperties(device, &device_props);
    std::cout << "\t- " << device_props.deviceName << " [V: " <<
        VK_VERSION_MAJOR(device_props.apiVersion) << "." <<
        VK_VERSION_MINOR(device_props.apiVersion) << "." <<
        VK_VERSION_PATCH(device_props.apiVersion) << "]" << std::endl;
    // Currently, any device with Vulkan API version 1.1 is fine.
    if (chosen_device == VK_NULL_HANDLE && device_props.apiVersion >= VK_API_VERSION_1_1) {
      chosen_device = device;
    }
  }

  if (chosen_device == VK_NULL_HANDLE) {
    ERROR("Unable to find suitable VkPhysicalDevice");
  }
  return chosen_device;
}

uint32_t ChooseDeviceQueueFamilyIndex(VkPhysicalDevice physical_device) {
  uint32_t props_count;
  vkGetPhysicalDeviceQueueFamilyProperties(physical_device, &props_count, nullptr);
  std::vector<VkQueueFamilyProperties> props(props_count);
  vkGetPhysicalDeviceQueueFamilyProperties(physical_device, &props_count, props.data());

  // Simply choose the first graphics queue.
  for (size_t i = 0; i < props_count; i++) {
    const VkQueueFamilyProperties& prop = props[i];
    if ((prop.queueFlags & VK_QUEUE_GRAPHICS_BIT) && prop.queueCount > 0) {
      return i;
    }
  }

  ERROR("Unable to find suitable queue family");
}

VkDevice CreateVkDevice(VkPhysicalDevice physical_device, uint32_t device_queue_family_index) {
  VkDeviceQueueCreateInfo queue_create_info = {};
  queue_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
  queue_create_info.queueFamilyIndex = device_queue_family_index;
  queue_create_info.queueCount = 1;
  float queue_priority = 1.0f;
  queue_create_info.pQueuePriorities = &queue_priority;

  VkDeviceCreateInfo device_create_info = {};
  device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
  device_create_info.queueCreateInfoCount = 1;
  device_create_info.pQueueCreateInfos = &queue_create_info;
  // Let's not use any device extensions for now.
  device_create_info.enabledExtensionCount = 0;
  device_create_info.ppEnabledExtensionNames = nullptr;

  VkDevice device;
  VkResult result = vkCreateDevice(physical_device, &device_create_info, nullptr, &device);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create logical device: " << result);
  }
  return device;
}

VkQueue GetVkQueue(VkDevice device, uint32_t device_queue_family_index) {
  VkQueue queue;
  vkGetDeviceQueue(device, device_queue_family_index, /*queueIndex*/ 0, &queue);
  return queue;
}

VkCommandPool CreateVkCommandPool(VkDevice device, uint32_t device_queue_family_index) {
  VkCommandPool command_pool;
  VkCommandPoolCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
  create_info.flags = 0;
  create_info.queueFamilyIndex = device_queue_family_index;
  VkResult result = vkCreateCommandPool(device, &create_info, nullptr, &command_pool);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create command pool: " << result);
  }
  return command_pool;
}

VkCommandBuffer CreateVkCommandBuffer(VkDevice device, VkCommandPool command_pool) {
  VkCommandBufferAllocateInfo allocate_info = {};
  allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  allocate_info.commandPool = command_pool;
  allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
  allocate_info.commandBufferCount = 1;

  VkCommandBuffer command_buffer;
  VkResult result = vkAllocateCommandBuffers( device, &allocate_info, &command_buffer);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create VkCommandBuffer: " << result);
  }
  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) {
  std::array<VkAttachmentDescription, 2> attachment_descriptions = {};

  VkAttachmentDescription* color_attachment_description = &attachment_descriptions[0];
  color_attachment_description->format = kVulkanFormat;
  color_attachment_description->samples = VK_SAMPLE_COUNT_1_BIT;
  color_attachment_description->loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
  color_attachment_description->storeOp = VK_ATTACHMENT_STORE_OP_STORE;
  color_attachment_description->stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
  color_attachment_description->stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
  color_attachment_description->initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
  color_attachment_description->finalLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;

  VkAttachmentReference color_attachment_reference = {};
  color_attachment_reference.attachment = 0;
  color_attachment_reference.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

  VkAttachmentDescription* depth_attachment_description = &attachment_descriptions[1];
  depth_attachment_description->format = kVulkanDepthFormat;
  depth_attachment_description->samples = VK_SAMPLE_COUNT_1_BIT;
  depth_attachment_description->loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
  depth_attachment_description->storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
  depth_attachment_description->stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
  depth_attachment_description->stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
  depth_attachment_description->initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
  depth_attachment_description->finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

  VkAttachmentReference depth_attachment_reference = {};
  depth_attachment_reference.attachment = 1;
  depth_attachment_reference.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

  VkSubpassDescription subpass_description = {};
  subpass_description.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
  subpass_description.colorAttachmentCount = 1;
  subpass_description.pColorAttachments = &color_attachment_reference;
  subpass_description.pDepthStencilAttachment = &depth_attachment_reference;

  VkSubpassDependency subpass_dependency = {};
  subpass_dependency.srcSubpass = VK_SUBPASS_EXTERNAL;
  subpass_dependency.dstSubpass = 0;
  subpass_dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
  subpass_dependency.srcAccessMask = 0;
  subpass_dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
  subpass_dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT |
      VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;

  VkRenderPassCreateInfo render_pass_info = {};
  render_pass_info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
  render_pass_info.attachmentCount = attachment_descriptions.size();
  render_pass_info.pAttachments = attachment_descriptions.data();
  render_pass_info.subpassCount = 1;
  render_pass_info.pSubpasses = &subpass_description;
  render_pass_info.dependencyCount = 1;
  render_pass_info.pDependencies = &subpass_dependency;

  VkRenderPass render_pass;
  VkResult result = vkCreateRenderPass(device, &render_pass_info, nullptr, &render_pass);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create render pass: " << result);
  }
  return render_pass;
}

VkShaderModule CreateVkShaderModule(VkDevice device, std::string path) {
  std::ifstream fstream(path);
  if (!fstream) {
    ERROR("Unable to open: " << path);
  }
  std::stringstream buffer;
  buffer << fstream.rdbuf();
  std::string spirv_source = buffer.str();

  VkShaderModuleCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
  create_info.codeSize = spirv_source.length();
  create_info.pCode = (const uint32_t*)spirv_source.c_str();

  VkShaderModule shader_module;
  VkResult result = vkCreateShaderModule(device, &create_info, nullptr, &shader_module);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create shader module for " << path << ": ");
  }
  return shader_module;
}

VkPipeline CreateVkPipeline(VkDevice device, VkRenderPass render_pass) {
  VkShaderModule vertex_shader_module = CreateVkShaderModule(device, kVertexShaderPath);
  VkShaderModule fragment_shader_module = CreateVkShaderModule(device, kFragmentShaderPath);

  VkPipelineShaderStageCreateInfo vertex_shader_stage_info = {};
  vertex_shader_stage_info.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  vertex_shader_stage_info.stage = VK_SHADER_STAGE_VERTEX_BIT;
  vertex_shader_stage_info.module = vertex_shader_module;
  vertex_shader_stage_info.pName = "main";

  VkPipelineShaderStageCreateInfo fragment_shader_stage_info = {};
  fragment_shader_stage_info.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  fragment_shader_stage_info.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
  fragment_shader_stage_info.module = fragment_shader_module;
  fragment_shader_stage_info.pName = "main";

  std::vector<VkPipelineShaderStageCreateInfo> shader_stages =
      {vertex_shader_stage_info, fragment_shader_stage_info};

  VkPipelineVertexInputStateCreateInfo vertex_input_info = {};
  vertex_input_info.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
  vertex_input_info.vertexBindingDescriptionCount = 0;
  vertex_input_info.pVertexBindingDescriptions = nullptr;
  vertex_input_info.vertexAttributeDescriptionCount = 0;
  vertex_input_info.pVertexAttributeDescriptions = nullptr;

  VkPipelineInputAssemblyStateCreateInfo input_assembly_info = {};
  input_assembly_info.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
  input_assembly_info.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
  input_assembly_info.primitiveRestartEnable = VK_FALSE;

  VkViewport viewport = {};
  viewport.x = 0.0f;
  viewport.y = 0.0f;
  viewport.width = (float)kWidth;
  viewport.height = (float)kHeight;
  viewport.minDepth = 0.0f;
  viewport.maxDepth = 1.0f;

  VkExtent2D extent = {};
  extent.width = kWidth;
  extent.height = kHeight;
  VkRect2D scissor = {};
  scissor.offset = {0, 0};
  scissor.extent = extent;

  VkPipelineViewportStateCreateInfo viewport_state_info = {};
  viewport_state_info.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
  viewport_state_info.viewportCount = 1;
  viewport_state_info.pViewports = &viewport;
  viewport_state_info.scissorCount = 1;
  viewport_state_info.pScissors = &scissor;

  VkPipelineRasterizationStateCreateInfo rasterization_state_info = {};
  rasterization_state_info.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
  rasterization_state_info.depthClampEnable = VK_FALSE;
  rasterization_state_info.rasterizerDiscardEnable = VK_FALSE;
  rasterization_state_info.polygonMode = VK_POLYGON_MODE_FILL;
  rasterization_state_info.lineWidth = 1.0f;
  rasterization_state_info.cullMode = VK_CULL_MODE_BACK_BIT;
  rasterization_state_info.frontFace = VK_FRONT_FACE_CLOCKWISE;
  rasterization_state_info.depthBiasEnable = VK_FALSE;

  VkPipelineMultisampleStateCreateInfo multisampling_state_info = {};
  multisampling_state_info.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
  multisampling_state_info.sampleShadingEnable = VK_FALSE;
  multisampling_state_info.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

  VkPipelineDepthStencilStateCreateInfo depth_stencil_create_info = {};
  depth_stencil_create_info.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
  depth_stencil_create_info.depthTestEnable = VK_TRUE;
  depth_stencil_create_info.depthWriteEnable = VK_TRUE;
  depth_stencil_create_info.depthCompareOp = VK_COMPARE_OP_LESS;
  // TODO(brkho): Test depth bounds functionality.
  depth_stencil_create_info.depthBoundsTestEnable = VK_FALSE;
  depth_stencil_create_info.minDepthBounds = 0.0f;
  depth_stencil_create_info.maxDepthBounds = 1.0f;
  depth_stencil_create_info.stencilTestEnable = VK_FALSE;

  VkPipelineColorBlendAttachmentState color_blend_attachment_state = {};
  color_blend_attachment_state.colorWriteMask = VK_COLOR_COMPONENT_R_BIT |
      VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
  color_blend_attachment_state.blendEnable = VK_FALSE;

  VkPipelineColorBlendStateCreateInfo color_blend_state = {};
  color_blend_state.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
  color_blend_state.logicOpEnable = VK_FALSE;
  color_blend_state.attachmentCount = 1;
  color_blend_state.pAttachments = &color_blend_attachment_state;

  VkPipelineLayoutCreateInfo pipeline_layout_create_info = {};
  pipeline_layout_create_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;

  VkPipelineLayout pipeline_layout;
  VkResult result = vkCreatePipelineLayout(device, &pipeline_layout_create_info, nullptr,
      &pipeline_layout);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create VkPipelineLayout: " << result);
  }

  VkGraphicsPipelineCreateInfo pipeline_create_info = {};
  pipeline_create_info.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
  pipeline_create_info.stageCount = 2;
  pipeline_create_info.pStages = shader_stages.data();
  pipeline_create_info.pVertexInputState = &vertex_input_info;
  pipeline_create_info.pInputAssemblyState = &input_assembly_info;
  pipeline_create_info.pViewportState = &viewport_state_info;
  pipeline_create_info.pRasterizationState = &rasterization_state_info;
  pipeline_create_info.pMultisampleState = &multisampling_state_info;
  pipeline_create_info.pDepthStencilState = &depth_stencil_create_info;
  pipeline_create_info.pColorBlendState = &color_blend_state;
  pipeline_create_info.pDynamicState = nullptr;
  pipeline_create_info.layout = pipeline_layout;
  pipeline_create_info.renderPass = render_pass;
  pipeline_create_info.subpass = 0;
  pipeline_create_info.basePipelineHandle = VK_NULL_HANDLE;
  pipeline_create_info.basePipelineIndex = -1;

  VkPipeline pipeline;
  result = vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipeline_create_info, nullptr,
      &pipeline);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create VkPipeline: " << result);
  }
  return pipeline;
}

VkImage CreateVkImage(VkDevice device, VkFormat format, VkImageTiling image_tiling,
    VkImageUsageFlags usage) {
  VkImageCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
  create_info.pNext = nullptr;
  create_info.imageType = VK_IMAGE_TYPE_2D;
  create_info.format = format;
  VkExtent3D extent  = {};
  extent.width = kWidth;
  extent.height = kHeight;
  extent.depth = 1;
  create_info.extent = extent;
  create_info.mipLevels = 1;
  create_info.arrayLayers = 1;
  create_info.samples = VK_SAMPLE_COUNT_1_BIT;
  create_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
  create_info.tiling = image_tiling;
  create_info.usage = usage;
  create_info.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;

  VkImage image;
  VkResult result = vkCreateImage(device, &create_info, nullptr, &image);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create VkImage: " << result);
  }
  return image;
}

inline VkImage CreateRenderVkImage(VkDevice device) {
  return CreateVkImage(device, kVulkanFormat, VK_IMAGE_TILING_OPTIMAL,
      VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT);
}

inline VkImage CreateDepthVkImage(VkDevice device) {
  return CreateVkImage(device, kVulkanDepthFormat, VK_IMAGE_TILING_OPTIMAL,
      VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT);
}

inline VkImage CreateScanoutVkImage(VkDevice device) {
  return CreateVkImage(device, kVulkanFormat, VK_IMAGE_TILING_LINEAR,
      VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT);
}

uint32_t FindMemoryType(uint32_t valid_image_memory_types,
    VkPhysicalDeviceMemoryProperties device_memory_properties,
    VkMemoryPropertyFlags memory_property_flags) {
  for (uint32_t i = 0; i < device_memory_properties.memoryTypeCount; i++) {
    // We don't care about performance, so just choose the first mappable memory type.
    if ((valid_image_memory_types & (1 << i)) &&
        ((device_memory_properties.memoryTypes[i].propertyFlags & memory_property_flags) ==
        memory_property_flags)) {
      return i;
    }
  }
  ERROR("Unable to find suitable memory type index");
}

VkDeviceMemory AllocateAndBindMemory(VkPhysicalDevice physical_device, VkDevice device,
    VkImage image, VkMemoryPropertyFlags memory_property_flags) {
  VkMemoryRequirements image_memory_requirements;
  vkGetImageMemoryRequirements(device, image, &image_memory_requirements);

  VkPhysicalDeviceMemoryProperties device_memory_properties;
  vkGetPhysicalDeviceMemoryProperties(physical_device, &device_memory_properties);

  VkMemoryAllocateInfo memory_allocate_info = {};
  memory_allocate_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
  memory_allocate_info.allocationSize = image_memory_requirements.size;
  memory_allocate_info.memoryTypeIndex = FindMemoryType(
      image_memory_requirements.memoryTypeBits, device_memory_properties,
      memory_property_flags);

  VkDeviceMemory image_memory;
  VkResult result = vkAllocateMemory(device, &memory_allocate_info, nullptr, &image_memory);
  if (result != VK_SUCCESS) {
    ERROR("Unable to allocate image memory: " << result);
  }

  result = vkBindImageMemory(device, image, image_memory, 0);
  if (result != VK_SUCCESS) {
    ERROR("Unable to bind image memory: " << result);
  }
  return image_memory;
}

inline VkDeviceMemory AllocateAndBindRenderMemory(VkPhysicalDevice physical_device,
    VkDevice device, VkImage image) {
  return AllocateAndBindMemory(physical_device, device, image, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
}

inline VkDeviceMemory AllocateAndBindDepthMemory(VkPhysicalDevice physical_device,
    VkDevice device, VkImage image) {
  return AllocateAndBindMemory(physical_device, device, image, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
}

inline VkDeviceMemory AllocateAndBindScanoutMemory(VkPhysicalDevice physical_device,
    VkDevice device, VkImage image) {
  return AllocateAndBindMemory(physical_device, device, image, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT);
}

VkImageView CreateVkImageView(VkDevice device, VkImage image, VkFormat format,
    VkImageAspectFlags aspect) {
  VkImageViewCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
  create_info.image = image;
  create_info.viewType = VK_IMAGE_VIEW_TYPE_2D;
  create_info.format = format;

  VkComponentMapping component_mapping = {};
  component_mapping.r = VK_COMPONENT_SWIZZLE_IDENTITY;
  component_mapping.b = VK_COMPONENT_SWIZZLE_IDENTITY;
  component_mapping.g = VK_COMPONENT_SWIZZLE_IDENTITY;
  component_mapping.a = VK_COMPONENT_SWIZZLE_IDENTITY;
  create_info.components = component_mapping;

  VkImageSubresourceRange subresource_range = {};
  subresource_range.aspectMask = aspect;
  subresource_range.baseMipLevel = 0;
  subresource_range.levelCount = 1;
  subresource_range.baseArrayLayer = 0;
  subresource_range.layerCount = 1;
  create_info.subresourceRange = subresource_range;

  VkImageView image_view;
  VkResult result = vkCreateImageView(device, &create_info, nullptr, &image_view);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create VkImageView: " << result);
  }
  return image_view;
}

VkImageView CreateRenderVkImageView(VkDevice device, VkImage image) {
  return CreateVkImageView(device, image, kVulkanFormat, VK_IMAGE_ASPECT_COLOR_BIT);
}

VkImageView CreateDepthVkImageView(VkDevice device, VkImage image) {
  return CreateVkImageView(device, image, kVulkanDepthFormat, VK_IMAGE_ASPECT_DEPTH_BIT);
}

VkFramebuffer CreateVkFramebuffer(VkDevice device, VkRenderPass render_pass,
    VkImageView render_image_view, VkImageView depth_image_view) {
  std::array<VkImageView, 2> attachments = { render_image_view, depth_image_view };

  VkFramebufferCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
  create_info.renderPass = render_pass;
  create_info.attachmentCount = attachments.size();
  create_info.pAttachments = attachments.data();
  create_info.width = kWidth;
  create_info.height = kHeight;
  create_info.layers = 1;

  VkFramebuffer framebuffer;
  VkResult result = vkCreateFramebuffer(device, &create_info, nullptr, &framebuffer);
  if (result != VK_SUCCESS) {
    ERROR("Unable to create VkFramebuffer: " << result);
  }
  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;
  command_buffer_begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
  VkResult result = vkBeginCommandBuffer(command_buffer, &command_buffer_begin_info);
  if (result != VK_SUCCESS) {
    ERROR("Unable to begin command buffer recording: " << result);
  }
}

void EndCommandBufferAndSubmit(VkCommandBuffer command_buffer, VkQueue queue) {
  VkResult result = vkEndCommandBuffer(command_buffer);
  if (result != VK_SUCCESS) {
    ERROR("Unable to end command buffer recording: " << result);
  }

  VkSubmitInfo submit_info = {};
  submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
  submit_info.commandBufferCount = 1;
  submit_info.pCommandBuffers = &command_buffer;
  result = vkQueueSubmit(queue, 1, &submit_info, VK_NULL_HANDLE);
  if (result != VK_SUCCESS) {
    ERROR("Error in submitting command buffer to queue: " << result);
  }
}

void WaitForIdle(VkQueue queue) {
  // In a real application, we should use real synchronization primitives to figure out when the
  // command buffers have been executed. Waiting on an idle queue is simple, but it can cause
  // deadlock if we continue to submit to the queue while we wait for idle.
  VkResult result = vkQueueWaitIdle(queue);
  if (result != VK_SUCCESS) {
    ERROR("Error in waiting for graphics queue to reach idle state: " << result);
  }
}

void Draw(VkDevice device, VkQueue queue, VkRenderPass render_pass, VkPipeline pipeline,
    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;
  render_pass_begin_info.framebuffer = framebuffer;
  VkRect2D render_area = {};
  render_area.offset = { 0, 0 };
  render_area.extent = { kWidth, kHeight };
  render_pass_begin_info.renderArea = render_area;

  std::array<VkClearValue, 2> clear_values = {};
  clear_values[0].color.float32[0] = 0.0f;
  clear_values[0].color.float32[1] = 0.0f;
  clear_values[0].color.float32[2] = 1.0f;
  clear_values[0].color.float32[3] = 1.0f;
  clear_values[1].depthStencil.depth = 1.0f;
  clear_values[1].depthStencil.stencil = 0.0f;

  render_pass_begin_info.clearValueCount = clear_values.size();
  render_pass_begin_info.pClearValues = clear_values.data();
  vkCmdBeginRenderPass(command_buffer, &render_pass_begin_info, VK_SUBPASS_CONTENTS_INLINE);

  vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
  vkCmdDraw(command_buffer, 6, 1, 0, 0);
  vkCmdEndRenderPass(command_buffer);
  EndQuery(command_buffer, query_pool);
}

void BlitToScanoutImage(VkQueue queue, VkCommandBuffer command_buffer, VkImage source_image,
    VkImage dest_image) {
  // We need to make sure the writes of the render pass have executed before actually scanning out
  // the rendered image. Also, we need to transition the layout of the scanout image to
  // VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL.
  VkImageMemoryBarrier image_memory_barrier = {};
  image_memory_barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
  image_memory_barrier.pNext = nullptr;
  image_memory_barrier.srcAccessMask = 0;
  image_memory_barrier.dstAccessMask = 0;
  image_memory_barrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
  image_memory_barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
  image_memory_barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  image_memory_barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  image_memory_barrier.image = dest_image;
  VkImageSubresourceRange subresource_range = {};
  subresource_range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
  subresource_range.baseMipLevel = 0;
  subresource_range.levelCount = 1;
  subresource_range.baseArrayLayer = 0;
  subresource_range.layerCount = 1;
  image_memory_barrier.subresourceRange = subresource_range;
  vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
      VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &image_memory_barrier);

  VkOffset3D blit_start = {};
  blit_start.x = 0;
  blit_start.y = 0;
  blit_start.z = 0;
  VkOffset3D blit_end = {};
  blit_end.x = kWidth;
  blit_end.y = kHeight;
  blit_end.z = 1;
  VkImageSubresourceLayers subresource_layers = {};
  subresource_layers.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
  subresource_layers.mipLevel = 0;
  subresource_layers.baseArrayLayer = 0;
  subresource_layers.layerCount = 1;
  VkImageBlit image_blit = {};
  image_blit.srcSubresource = subresource_layers;
  image_blit.srcOffsets[0] = blit_start;
  image_blit.srcOffsets[1] = blit_end;
  image_blit.dstSubresource = subresource_layers;
  image_blit.dstOffsets[0] = blit_start;
  image_blit.dstOffsets[1] = blit_end;
  // TODO(brkho): Support multi-planar formats.
  vkCmdBlitImage(command_buffer, source_image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, dest_image,
      VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &image_blit, VK_FILTER_NEAREST);

  EndCommandBufferAndSubmit(command_buffer, queue);
}

uint8_t* MapVkDeviceMemory(VkDevice device, VkDeviceMemory image_memory) {
  void* mapped_memory = nullptr;
  VkResult result = vkMapMemory(device, image_memory, 0, kWidth * kHeight * 4, 0, &mapped_memory);
  if (result != VK_SUCCESS) {
    ERROR("Unable to map device memory: " << result);
  }
  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");

  png_image image_info = {};
  image_info.version = PNG_IMAGE_VERSION;
  image_info.width = kWidth;
  image_info.height = kHeight;
  image_info.format = PNG_FORMAT_RGBA;
  if (png_image_write_to_stdio(&image_info, f, 0, pixels, kWidth * 4, nullptr) == 0) {
    ERROR("Error writing PNG: " << image_info.message);
  }

  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);
  uint32_t device_queue_family_index = ChooseDeviceQueueFamilyIndex(physical_device);
  VkDevice device = CreateVkDevice(physical_device, device_queue_family_index);
  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);

  VkImage render_image = CreateRenderVkImage(device);
  AllocateAndBindRenderMemory(physical_device, device, render_image);
  VkImageView render_image_view = CreateRenderVkImageView(device, render_image);
  VkImage depth_image = CreateDepthVkImage(device);
  AllocateAndBindDepthMemory(physical_device, device, depth_image);
  VkImageView depth_image_view = CreateDepthVkImageView(device, depth_image);
  VkFramebuffer framebuffer = CreateVkFramebuffer(device, render_pass, render_image_view,
      depth_image_view);

  VkImage scanout_image = CreateScanoutVkImage(device);
  VkDeviceMemory scanout_image_memory =
      AllocateAndBindScanoutMemory(physical_device, device, scanout_image);

  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);
  WaitForIdle(queue);

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