graph/src/dofpass.cpp

#include "dofpass.h"

#include <array>
#include <glm/glm.hpp>
#include <stb_image.h>

#include "renderer.h"
#include "camera.h"

DoFPass::DoFPass(Renderer& renderer) : renderer_(renderer) {
    createRenderPass();
    createDescriptorSetLayout();
    createPipeline();
    createBokehImage();
}

DoFPass::~DoFPass() {
    vkDestroySampler(renderer_.getDevice(), bokehSampler_, nullptr);
    vkDestroyImageView(renderer_.getDevice(), bokehImageView_, nullptr);
    vkFreeMemory(renderer_.getDevice(), bokehMemory_, nullptr);
    vkDestroyImage(renderer_.getDevice(), bokehImage_, nullptr);

    vkDestroyPipeline(renderer_.getDevice(), pipeline_, nullptr);
    vkDestroyPipelineLayout(renderer_.getDevice(), pipelineLayout_, nullptr);

    vkDestroyDescriptorSetLayout(renderer_.getDevice(), setLayout_, nullptr);

    vkDestroyRenderPass(renderer_.getDevice(), renderPass_, nullptr);
}

void DoFPass::render(VkCommandBuffer commandBuffer, Camera& camera, RenderTarget* target) {
    VkViewport viewport = {};
    viewport.width = target->extent.width / dofFramebufferDownscaleFactor;
    viewport.height = target->extent.height / dofFramebufferDownscaleFactor;
    viewport.maxDepth = 1.0f;

    vkCmdSetViewport(commandBuffer, 0, 1, &viewport);

    VkRect2D scissor = {};
    scissor.extent.width = target->extent.width / dofFramebufferDownscaleFactor;
    scissor.extent.height = target->extent.height / dofFramebufferDownscaleFactor;

    vkCmdSetScissor(commandBuffer, 0, 1, &scissor);

    VkClearValue clearColor = {};

    VkRenderPassBeginInfo renderPassBeginInfo = {};
    renderPassBeginInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    renderPassBeginInfo.framebuffer = target->farFieldFramebuffers[target->currentResource];
    renderPassBeginInfo.renderPass = renderPass_;
    renderPassBeginInfo.renderArea.extent.width = target->extent.width / dofFramebufferDownscaleFactor;
    renderPassBeginInfo.renderArea.extent.height = target->extent.height / dofFramebufferDownscaleFactor;
    renderPassBeginInfo.clearValueCount = 1;
    renderPassBeginInfo.pClearValues = &clearColor;

    // far field
    vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

    glm::vec4 dpack;
    dpack[0] = camera.aperture;
    dpack[1] = (camera.far - camera.focusDistance) / camera.far;
    dpack[2] = target->extent.width / dofFramebufferDownscaleFactor;
    dpack[3] = target->extent.height / dofFramebufferDownscaleFactor;

    vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_);
    vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout_, 0, 1, &target->dofSets[target->currentResource], 0, nullptr);

    vkCmdPushConstants(commandBuffer, pipelineLayout_, VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(glm::vec4), &dpack);

    if(camera.aperture > 0.0f)
        vkCmdDraw(commandBuffer, 3, (target->extent.width / dofFramebufferDownscaleFactor) * (target->extent.height / dofFramebufferDownscaleFactor), 0, 0);

    vkCmdEndRenderPass(commandBuffer);

    //near field
    renderPassBeginInfo.framebuffer = target->nearFieldFramebuffers[target->currentResource];

    vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

    vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_);
    vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout_, 0, 1, &target->dofSets[target->currentResource], 0, nullptr);

    vkCmdPushConstants(commandBuffer, pipelineLayout_, VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(glm::vec4), &dpack);

    //FIXME: near field is bugged
    //vkCmdDraw(commandBuffer, 3, (target->extent.width / 2) * (target->extent.height / 2), 0, 0);

    vkCmdEndRenderPass(commandBuffer);
}

void DoFPass::createDescriptorSet(RenderTarget* target) {
    VkDescriptorSetAllocateInfo allocInfo = {};
    allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = renderer_.getDescriptorPool();
    allocInfo.descriptorSetCount = target->numImages;

    // FIXME: lol what
    auto layouts = new VkDescriptorSetLayout[target->numImages];
    for(uint32_t i = 0; i < target->numImages; i++)
        layouts[i] = setLayout_;

    allocInfo.pSetLayouts = layouts;

    target->dofSets = new VkDescriptorSet[target->numImages];
    vkAllocateDescriptorSets(renderer_.getDevice(), &allocInfo, target->dofSets);

    delete[] layouts;

    for(uint32_t i = 0; i < target->numImages; i++) {
        VkDescriptorImageInfo bokehImageInfo = {};
        bokehImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        bokehImageInfo.imageView = bokehImageView_;
        bokehImageInfo.sampler = bokehSampler_;

        VkDescriptorImageInfo sceneImageInfo = {};
        sceneImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        sceneImageInfo.imageView = target->offscreenColorImageViews[i];
        sceneImageInfo.sampler = bokehSampler_;

        VkDescriptorImageInfo depthImageInfo = {};
        depthImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;;
        depthImageInfo.imageView = target->offscreenDepthImageViews[i];
        depthImageInfo.sampler = bokehSampler_;

        VkWriteDescriptorSet bokehDescriptorWrite = {};
        bokehDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        bokehDescriptorWrite.descriptorCount = 1;
        bokehDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        bokehDescriptorWrite.dstSet = target->dofSets[i];
        bokehDescriptorWrite.pImageInfo = &bokehImageInfo;

        VkWriteDescriptorSet sceneDescriptorWrite = {};
        sceneDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        sceneDescriptorWrite.descriptorCount = 1;
        sceneDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        sceneDescriptorWrite.dstBinding = 1;
        sceneDescriptorWrite.dstSet = target->dofSets[i];
        sceneDescriptorWrite.pImageInfo = &sceneImageInfo;

        VkWriteDescriptorSet depthDescriptorWrite = {};
        depthDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        depthDescriptorWrite.descriptorCount = 1;
        depthDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        depthDescriptorWrite.dstBinding = 2;
        depthDescriptorWrite.dstSet = target->dofSets[i];
        depthDescriptorWrite.pImageInfo = &depthImageInfo;

        const std::array<VkWriteDescriptorSet, 3> descriptorWrites = {
            bokehDescriptorWrite,
            sceneDescriptorWrite,
            depthDescriptorWrite
        };

        vkUpdateDescriptorSets(renderer_.getDevice(), descriptorWrites.size(), descriptorWrites.data(), 0, nullptr);
    }
}

void DoFPass::createRenderPass() {
    VkAttachmentDescription colorAttachment = {};
    colorAttachment.format = VK_FORMAT_R32G32B32A32_SFLOAT;
    colorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
    colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
    colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
    colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    colorAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    colorAttachment.finalLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

    VkAttachmentReference colorAttachmentRef = {};
    colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    VkSubpassDescription subpass = {};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = 1;
    subpass.pColorAttachments = &colorAttachmentRef;

    VkRenderPassCreateInfo renderPassInfo = {};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
    renderPassInfo.attachmentCount = 1;
    renderPassInfo.pAttachments = &colorAttachment;
    renderPassInfo.subpassCount = 1;
    renderPassInfo.pSubpasses = &subpass;

    vkCreateRenderPass(renderer_.getDevice(), &renderPassInfo, nullptr, &renderPass_);
}

void DoFPass::createDescriptorSetLayout() {
    VkDescriptorSetLayoutBinding bokehSamplerBinding = {};
    bokehSamplerBinding.descriptorCount = 1;
    bokehSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    bokehSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding sceneSamplerBinding = {};
    sceneSamplerBinding.binding = 1;
    sceneSamplerBinding.descriptorCount = 1;
    sceneSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    sceneSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding depthSamplerBinding = {};
    depthSamplerBinding.binding = 2;
    depthSamplerBinding.descriptorCount = 1;
    depthSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    depthSamplerBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;

    const std::array<VkDescriptorSetLayoutBinding, 3> bindings = {
        bokehSamplerBinding,
        sceneSamplerBinding,
        depthSamplerBinding
    };

    VkDescriptorSetLayoutCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    createInfo.bindingCount = bindings.size();
    createInfo.pBindings = bindings.data();

    vkCreateDescriptorSetLayout(renderer_.getDevice(), &createInfo, nullptr, &setLayout_);
}

void DoFPass::createPipeline() {
    VkShaderModule vertShaderModule = renderer_.createShader("shaders/dof.vert.spv");
    VkShaderModule fragShaderModule = renderer_.createShader("shaders/dof.frag.spv");

    VkPipelineShaderStageCreateInfo vertShaderStageInfo = {};
    vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
    vertShaderStageInfo.module = vertShaderModule;
    vertShaderStageInfo.pName = "main";

    VkPipelineShaderStageCreateInfo fragShaderStageInfo = {};
    fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
    fragShaderStageInfo.module = fragShaderModule;
    fragShaderStageInfo.pName = "main";

    const std::array<VkPipelineShaderStageCreateInfo, 2> shaderStages = {vertShaderStageInfo, fragShaderStageInfo};

    VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};
    vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;

    VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};
    inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;

    VkPipelineViewportStateCreateInfo viewportState = {};
    viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewportState.viewportCount = 1;
    viewportState.scissorCount = 1;

    VkPipelineRasterizationStateCreateInfo rasterizer = {};
    rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
    rasterizer.lineWidth = 1.0f;
    rasterizer.cullMode = VK_CULL_MODE_FRONT_BIT;
    rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;

    VkPipelineMultisampleStateCreateInfo multisampling = {};
    multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

    VkPipelineColorBlendAttachmentState colorBlendAttachment = {};
    colorBlendAttachment.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
    colorBlendAttachment.blendEnable = VK_TRUE;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;
    colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ONE;

    VkPipelineColorBlendStateCreateInfo colorBlending = {};
    colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlending.attachmentCount = 1;
    colorBlending.pAttachments = &colorBlendAttachment;

    const std::array<VkDynamicState, 2> dynamicStates = {
        VK_DYNAMIC_STATE_VIEWPORT,
        VK_DYNAMIC_STATE_SCISSOR
    };

    VkPipelineDynamicStateCreateInfo dynamicState = {};
    dynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicState.dynamicStateCount = dynamicStates.size();
    dynamicState.pDynamicStates = dynamicStates.data();

    VkPushConstantRange pushConstant = {};
    pushConstant.size = sizeof(glm::vec4);
    pushConstant.stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;

    VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = 1;
    pipelineLayoutInfo.pSetLayouts = &setLayout_;
    pipelineLayoutInfo.pushConstantRangeCount = 1;
    pipelineLayoutInfo.pPushConstantRanges = &pushConstant;

    vkCreatePipelineLayout(renderer_.getDevice(), &pipelineLayoutInfo, nullptr, &pipelineLayout_);

    VkGraphicsPipelineCreateInfo pipelineInfo = {};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = shaderStages.size();
    pipelineInfo.pStages = shaderStages.data();
    pipelineInfo.pVertexInputState = &vertexInputInfo;
    pipelineInfo.pInputAssemblyState = &inputAssembly;
    pipelineInfo.pViewportState = &viewportState;
    pipelineInfo.pRasterizationState = &rasterizer;
    pipelineInfo.pMultisampleState = &multisampling;
    pipelineInfo.pColorBlendState = &colorBlending;
    pipelineInfo.pDynamicState = &dynamicState;
    pipelineInfo.layout = pipelineLayout_;
    pipelineInfo.renderPass = renderPass_;

    vkCreateGraphicsPipelines(renderer_.getDevice(), nullptr, 1, &pipelineInfo, nullptr, &pipeline_);

    vkDestroyShaderModule(renderer_.getDevice(), fragShaderModule, nullptr);
    vkDestroyShaderModule(renderer_.getDevice(), vertShaderModule, nullptr);
}

void DoFPass::createBokehImage() {
    int width = 0, height = 0, channels = 0;
    stbi_uc* pixels = stbi_load("data/bokeh.png", &width, &height, &channels, STBI_rgb_alpha);
    if(pixels == nullptr)
        return; // haha what

    VkImageCreateInfo imageCreateInfo = {};
    imageCreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageCreateInfo.imageType = VK_IMAGE_TYPE_2D;
    imageCreateInfo.extent.width = width;
    imageCreateInfo.extent.height = height;
    imageCreateInfo.extent.depth = 1;
    imageCreateInfo.mipLevels = 1;
    imageCreateInfo.arrayLayers = 1;
    imageCreateInfo.format = VK_FORMAT_R8G8B8A8_UNORM;
    imageCreateInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
    imageCreateInfo.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
    imageCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    imageCreateInfo.samples = VK_SAMPLE_COUNT_1_BIT;

    vkCreateImage(renderer_.getDevice(), &imageCreateInfo, nullptr, &bokehImage_);

    VkMemoryRequirements memRequirements;
    vkGetImageMemoryRequirements(renderer_.getDevice(), bokehImage_, &memRequirements);

    VkMemoryAllocateInfo allocInfo = {};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = renderer_.findMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

    vkAllocateMemory(renderer_.getDevice(), &allocInfo, nullptr, &bokehMemory_);
    vkBindImageMemory(renderer_.getDevice(), bokehImage_, bokehMemory_, 0);

    renderer_.uploadImageData(
        bokehImage_,
        width,
        height,
        width * height * 4,
        pixels
    );

    stbi_image_free(pixels);

    VkImageViewCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    createInfo.image = bokehImage_;
    createInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
    createInfo.format = VK_FORMAT_R8G8B8A8_UNORM;
    createInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    createInfo.subresourceRange.levelCount = 1;
    createInfo.subresourceRange.layerCount = 1;

    vkCreateImageView(renderer_.getDevice(), &createInfo, nullptr, &bokehImageView_);

    VkSamplerCreateInfo samplerInfo = {};
    samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
    samplerInfo.magFilter = VK_FILTER_LINEAR;
    samplerInfo.minFilter = VK_FILTER_LINEAR;
    samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER;
    samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER;
    samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER;
    samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
    samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;

    vkCreateSampler(renderer_.getDevice(), &samplerInfo, nullptr, &bokehSampler_);
}