VulkanRayTracing.h

浏览该文件的文档.

#pragma once
 
#include "common/Result.h"
#include "graphics/IRayTracing.h"
#include "graphics/RayTracingCaps.h"
#include "vkbuilder.hpp"
 
#include <glm/mat4x4.hpp>
#include <glm/vec3.hpp>
 
#include <cstdint>
#include <memory>
#include <vector>
 
namespace eve::graphics {
class Canvas;
class Graphics;
class Mesh;
class Texture;
}  // namespace eve::graphics
 
namespace eve::graphics::raytracing {
 
struct DeviceAddressBuffer {
    vk::Buffer        buffer{};
    vk::DeviceMemory  memory{};
    vk::DeviceSize    size    = 0;
    vk::DeviceAddress address = 0;
    vkb::Device*      device  = nullptr;
 
    DeviceAddressBuffer() = default;
    ~DeviceAddressBuffer() { release(); }
    DeviceAddressBuffer(const DeviceAddressBuffer&)            = delete;
    DeviceAddressBuffer& operator=(const DeviceAddressBuffer&) = delete;
    DeviceAddressBuffer(DeviceAddressBuffer&& o) noexcept { steal(o); }
    DeviceAddressBuffer& operator=(DeviceAddressBuffer&& o) noexcept {
        if (this != &o) {
            release();
            steal(o);
        }
        return *this;
    }
 
    void release();
    void steal(DeviceAddressBuffer& o) noexcept;
 
    [[nodiscard]] Result<void> allocate(vkb::Device& dev, vk::DeviceSize bytes, vk::BufferUsageFlags usage);
 
    [[nodiscard]] Result<void> upload(vk::CommandPool pool, vk::Queue queue, const void* data, vk::DeviceSize bytes);
};
 
struct AccelerationStructure {
    vk::AccelerationStructureKHR handle{};
    DeviceAddressBuffer          storage;
    vk::DeviceAddress            deviceAddress = 0;
    vkb::Device*                 device        = nullptr;
 
    AccelerationStructure() = default;
    ~AccelerationStructure() { release(); }
    AccelerationStructure(const AccelerationStructure&)            = delete;
    AccelerationStructure& operator=(const AccelerationStructure&) = delete;
    AccelerationStructure(AccelerationStructure&& o) noexcept { steal(o); }
    AccelerationStructure& operator=(AccelerationStructure&& o) noexcept {
        if (this != &o) {
            release();
            steal(o);
        }
        return *this;
    }
 
    void release();
    void steal(AccelerationStructure& o) noexcept;
};
 
struct TriangleMeshRecord {
    DeviceAddressBuffer   vertexBuffer;
    DeviceAddressBuffer   indexBuffer;
    AccelerationStructure blas;
    uint32_t              vertexCount = 0;
    uint32_t              indexCount  = 0;
    glm::mat4             transform{1.f};
};
 
class VulkanRayTracing final : public IRayTracing {
public:
    VulkanRayTracing() = default;
    ~VulkanRayTracing() override;
 
    VulkanRayTracing(const VulkanRayTracing&)            = delete;
    VulkanRayTracing& operator=(const VulkanRayTracing&) = delete;
 
    bool           isAvailable() const override;
    RayTracingCaps caps() const override;
 
    [[nodiscard]] Result<void> applyReflections(Graphics* gfx, Texture* sceneColor, Texture* hwDepth,
                                                Texture* worldNormal, Canvas* dest, const glm::mat4& invViewProj,
                                                const glm::mat4& viewProj, const glm::vec3& eyeWorld) override;
 
    [[nodiscard]] Result<void> rebuildScene() override;
 
    [[nodiscard]] Result<uint32_t> addTriangleMesh(const float* positionsXYZ, int vertexCount, const uint32_t* indices,
                                                   int indexCount, const glm::mat4& transform) override;
 
    [[nodiscard]] Result<uint32_t> addMesh(Mesh* mesh, const glm::mat4& transform) override;
 
    void clearScene() override;
 
    void detachFromGraphics() override { detachDevice(); }
 
    void attachDevice(vkb::Device* device, const RayTracingCaps& caps, vk::CommandPool uploadPool,
                      vk::Queue graphicsQueue);
 
    void detachDevice();
 
    [[nodiscard]] Result<void> ensureAttached();
 
private:
    [[nodiscard]] Result<void> ensurePipeline();
    [[nodiscard]] Result<void> ensureSampler();
    [[nodiscard]] Result<void> ensureOutputImage(uint32_t width, uint32_t height);
    [[nodiscard]] Result<void> ensureDescriptorSets(vk::ImageView outputView, vk::ImageView sceneView,
                                                    vk::ImageView depthView, vk::ImageView normalView);
    [[nodiscard]] Result<void> buildBlas(TriangleMeshRecord& mesh);
    [[nodiscard]] Result<void> buildTlas();
    [[nodiscard]] Result<void> createShaderBindingTable();
 
    struct StorageColorImage : vkb::GenericImage {
        void allocate(vkb::Device& device, uint32_t width, uint32_t height) {
            vk::ImageCreateInfo info{};
            info.imageType     = vk::ImageType::e2D;
            info.format        = vk::Format::eR16G16B16A16Sfloat;
            info.extent        = vk::Extent3D{width, height, 1};
            info.mipLevels     = 1;
            info.arrayLayers   = 1;
            info.samples       = vk::SampleCountFlagBits::e1;
            info.tiling        = vk::ImageTiling::eOptimal;
            info.usage         = vk::ImageUsageFlagBits::eStorage | vk::ImageUsageFlagBits::eTransferSrc |
                                 vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eSampled;
            info.sharingMode   = vk::SharingMode::eExclusive;
            info.initialLayout = vk::ImageLayout::eUndefined;
            create(device, info, vk::ImageViewType::e2D, vk::ImageAspectFlagBits::eColor, false);
        }
    };
 
    vkb::Device*    device_ = nullptr;
    vk::CommandPool uploadPool_{};
    vk::Queue       graphicsQueue_{};
    RayTracingCaps  caps_{};
 
    std::vector<std::unique_ptr<TriangleMeshRecord>> meshes_;
    AccelerationStructure                            tlas_;
 
    vk::DescriptorSetLayout setLayout_{};
    vk::PipelineLayout      pipelineLayout_{};
    vk::Pipeline            pipeline_{};
    vk::DescriptorPool      descriptorPool_{};
    vk::DescriptorSet       descriptorSet_{};
    vk::Sampler             sampler_{};
 
    StorageColorImage outputImage_{};
    uint32_t          outputWidth_  = 0;
    uint32_t          outputHeight_ = 0;
 
    DeviceAddressBuffer               sbtBuffer_;
    vk::StridedDeviceAddressRegionKHR raygenRegion_{};
    vk::StridedDeviceAddressRegionKHR missRegion_{};
    vk::StridedDeviceAddressRegionKHR hitRegion_{};
    vk::StridedDeviceAddressRegionKHR callableRegion_{};
 
    vk::ShaderModule raygenModule_{};
    vk::ShaderModule missModule_{};
    vk::ShaderModule closestHitModule_{};
 
    struct PushConstants {
        glm::mat4 invViewProj;
        glm::mat4 viewProj;
        glm::vec4 eye;  // xyz = eye, w unused
    };
};
 
VulkanRayTracing& vulkanRayTracing();
 
}  // namespace eve::graphics::raytracing