procedural-3d-engine/examples/raytracingtextures/raytracingtextures.cpp
2024-03-19 21:59:25 +01:00

715 lines
33 KiB
C++

/*
* Vulkan Example - Texture mapping with transparency using accelerated ray tracing example
*
* Copyright (C) 2024 by Sascha Willems - www.saschawillems.de
*
* This code is licensed under the MIT license (MIT) (http://opensource.org/licenses/MIT)
*/
/*
* This hardware accelerated ray tracing sample renders a texture mapped quad with transparency
* The sample also makes use of buffer device addresses to pass references for vertex and index buffers
* to the shader, making data access a bit more straightforward than using descriptors.
* Buffer references themselves are then simply set at draw time using push constants.
* In addition to a closest hit shader, that now samples from the texture, an any hit shader is
* added to the closest hit shader group. We use this shader to check if the texel we want to
* sample at the currently hit ray position is transparent, and if that's the case the any hit
* shader will cancel the intersection.
*/
#include "VulkanRaytracingSample.h"
class VulkanExample : public VulkanRaytracingSample
{
public:
AccelerationStructure bottomLevelAS{};
AccelerationStructure topLevelAS{};
vks::Buffer vertexBuffer;
vks::Buffer indexBuffer;
uint32_t indexCount;
vks::Buffer transformBuffer;
std::vector<VkRayTracingShaderGroupCreateInfoKHR> shaderGroups{};
struct ShaderBindingTables {
ShaderBindingTable raygen;
ShaderBindingTable miss;
ShaderBindingTable hit;
} shaderBindingTables;
vks::Texture2D texture;
struct UniformData {
glm::mat4 viewInverse;
glm::mat4 projInverse;
} uniformData;
vks::Buffer ubo;
VkPipeline pipeline;
VkPipelineLayout pipelineLayout;
VkDescriptorSet descriptorSet;
VkDescriptorSetLayout descriptorSetLayout;
VulkanExample() : VulkanRaytracingSample()
{
title = "Ray tracing textures";
camera.type = Camera::CameraType::lookat;
camera.setPerspective(60.0f, (float)width / (float)height, 0.1f, 512.0f);
camera.setRotation(glm::vec3(45.0f, 0.0f, 0.0f));
camera.setTranslation(glm::vec3(0.0f, 0.0f, -1.0f));
enableExtensions();
// Buffer device address requires the 64-bit integer feature to be enabled
enabledFeatures.shaderInt64 = VK_TRUE;
}
~VulkanExample()
{
vkDestroyPipeline(device, pipeline, nullptr);
vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
deleteStorageImage();
deleteAccelerationStructure(bottomLevelAS);
deleteAccelerationStructure(topLevelAS);
vertexBuffer.destroy();
indexBuffer.destroy();
transformBuffer.destroy();
shaderBindingTables.raygen.destroy();
shaderBindingTables.miss.destroy();
shaderBindingTables.hit.destroy();
ubo.destroy();
texture.destroy();
}
void createAccelerationStructureBuffer(AccelerationStructure &accelerationStructure, VkAccelerationStructureBuildSizesInfoKHR buildSizeInfo)
{
VkBufferCreateInfo bufferCreateInfo{};
bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
bufferCreateInfo.size = buildSizeInfo.accelerationStructureSize;
bufferCreateInfo.usage = VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
VK_CHECK_RESULT(vkCreateBuffer(device, &bufferCreateInfo, nullptr, &accelerationStructure.buffer));
VkMemoryRequirements memoryRequirements{};
vkGetBufferMemoryRequirements(device, accelerationStructure.buffer, &memoryRequirements);
VkMemoryAllocateFlagsInfo memoryAllocateFlagsInfo{};
memoryAllocateFlagsInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
memoryAllocateFlagsInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT_KHR;
VkMemoryAllocateInfo memoryAllocateInfo{};
memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
memoryAllocateInfo.pNext = &memoryAllocateFlagsInfo;
memoryAllocateInfo.allocationSize = memoryRequirements.size;
memoryAllocateInfo.memoryTypeIndex = vulkanDevice->getMemoryType(memoryRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
VK_CHECK_RESULT(vkAllocateMemory(device, &memoryAllocateInfo, nullptr, &accelerationStructure.memory));
VK_CHECK_RESULT(vkBindBufferMemory(device, accelerationStructure.buffer, accelerationStructure.memory, 0));
}
/*
Create the bottom level acceleration structure contains the scene's actual geometry (vertices, triangles)
*/
void createBottomLevelAccelerationStructure()
{
// Setup vertices for a single triangle
struct Vertex {
float pos[3];
float normal[3];
float uv[2];
};
std::vector<Vertex> vertices = {
{ { 0.5f, 0.5f, 0.0f }, {.0f, .0f, -1.0f}, { 1.0f, 1.0f} },
{ { -.5f, 0.5f, 0.0f }, {.0f, .0f, -1.0f}, { 0.0f, 1.0f} },
{ { -.5f, -.5f, 0.0f }, {.0f, .0f, -1.0f}, { 0.0f, 0.0f} },
{ { 0.5f, -.5f, 0.0f }, {.0f, .0f, -1.0f}, { 1.0f, 0.0f} },
};
// Setup indices
std::vector<uint32_t> indices = { 0, 1, 2, 0, 3, 2 };
indexCount = static_cast<uint32_t>(indices.size());
// Setup identity transform matrix
VkTransformMatrixKHR transformMatrix = {
1.0f, 0.0f, 0.0f, 0.0f,
0.0f, 1.0f, 0.0f, 0.0f,
0.0f, 0.0f, 1.0f, 0.0f
};
// Create buffers
// For the sake of simplicity we won't stage the vertex data to the GPU memory
// Vertex buffer
VK_CHECK_RESULT(vulkanDevice->createBuffer(
VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
&vertexBuffer,
vertices.size() * sizeof(Vertex),
vertices.data()));
// Index buffer
VK_CHECK_RESULT(vulkanDevice->createBuffer(
VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
&indexBuffer,
indices.size() * sizeof(uint32_t),
indices.data()));
// Transform buffer
VK_CHECK_RESULT(vulkanDevice->createBuffer(
VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR,
VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
&transformBuffer,
sizeof(VkTransformMatrixKHR),
&transformMatrix));
VkDeviceOrHostAddressConstKHR vertexBufferDeviceAddress{};
VkDeviceOrHostAddressConstKHR indexBufferDeviceAddress{};
VkDeviceOrHostAddressConstKHR transformBufferDeviceAddress{};
vertexBufferDeviceAddress.deviceAddress = getBufferDeviceAddress(vertexBuffer.buffer);
indexBufferDeviceAddress.deviceAddress = getBufferDeviceAddress(indexBuffer.buffer);
transformBufferDeviceAddress.deviceAddress = getBufferDeviceAddress(transformBuffer.buffer);
// Build
VkAccelerationStructureGeometryKHR accelerationStructureGeometry{};
accelerationStructureGeometry.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR;
accelerationStructureGeometry.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
accelerationStructureGeometry.geometry.triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
accelerationStructureGeometry.geometry.triangles.vertexFormat = VK_FORMAT_R32G32B32_SFLOAT;
accelerationStructureGeometry.geometry.triangles.vertexData = vertexBufferDeviceAddress;
accelerationStructureGeometry.geometry.triangles.maxVertex = 3;
accelerationStructureGeometry.geometry.triangles.vertexStride = sizeof(Vertex);
accelerationStructureGeometry.geometry.triangles.indexType = VK_INDEX_TYPE_UINT32;
accelerationStructureGeometry.geometry.triangles.indexData = indexBufferDeviceAddress;
accelerationStructureGeometry.geometry.triangles.transformData.deviceAddress = 0;
accelerationStructureGeometry.geometry.triangles.transformData.hostAddress = nullptr;
accelerationStructureGeometry.geometry.triangles.transformData = transformBufferDeviceAddress;
// Get size info
VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
accelerationStructureBuildGeometryInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
accelerationStructureBuildGeometryInfo.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
accelerationStructureBuildGeometryInfo.geometryCount = 1;
accelerationStructureBuildGeometryInfo.pGeometries = &accelerationStructureGeometry;
const uint32_t numTriangles = static_cast<uint32_t>(indices.size() / 3);
VkAccelerationStructureBuildSizesInfoKHR accelerationStructureBuildSizesInfo{};
accelerationStructureBuildSizesInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR;
vkGetAccelerationStructureBuildSizesKHR(
device,
VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
&accelerationStructureBuildGeometryInfo,
&numTriangles,
&accelerationStructureBuildSizesInfo);
createAccelerationStructureBuffer(bottomLevelAS, accelerationStructureBuildSizesInfo);
VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfo{};
accelerationStructureCreateInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
accelerationStructureCreateInfo.buffer = bottomLevelAS.buffer;
accelerationStructureCreateInfo.size = accelerationStructureBuildSizesInfo.accelerationStructureSize;
accelerationStructureCreateInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
vkCreateAccelerationStructureKHR(device, &accelerationStructureCreateInfo, nullptr, &bottomLevelAS.handle);
// Create a small scratch buffer used during build of the bottom level acceleration structure
ScratchBuffer scratchBuffer = createScratchBuffer(accelerationStructureBuildSizesInfo.buildScratchSize);
VkAccelerationStructureBuildGeometryInfoKHR accelerationBuildGeometryInfo{};
accelerationBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
accelerationBuildGeometryInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
accelerationBuildGeometryInfo.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
accelerationBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR;
accelerationBuildGeometryInfo.dstAccelerationStructure = bottomLevelAS.handle;
accelerationBuildGeometryInfo.geometryCount = 1;
accelerationBuildGeometryInfo.pGeometries = &accelerationStructureGeometry;
accelerationBuildGeometryInfo.scratchData.deviceAddress = scratchBuffer.deviceAddress;
VkAccelerationStructureBuildRangeInfoKHR accelerationStructureBuildRangeInfo{};
accelerationStructureBuildRangeInfo.primitiveCount = numTriangles;
accelerationStructureBuildRangeInfo.primitiveOffset = 0;
accelerationStructureBuildRangeInfo.firstVertex = 0;
accelerationStructureBuildRangeInfo.transformOffset = 0;
std::vector<VkAccelerationStructureBuildRangeInfoKHR*> accelerationBuildStructureRangeInfos = { &accelerationStructureBuildRangeInfo };
// Build the acceleration structure on the device via a one-time command buffer submission
// Some implementations may support acceleration structure building on the host (VkPhysicalDeviceAccelerationStructureFeaturesKHR->accelerationStructureHostCommands), but we prefer device builds
VkCommandBuffer commandBuffer = vulkanDevice->createCommandBuffer(VK_COMMAND_BUFFER_LEVEL_PRIMARY, true);
vkCmdBuildAccelerationStructuresKHR(
commandBuffer,
1,
&accelerationBuildGeometryInfo,
accelerationBuildStructureRangeInfos.data());
vulkanDevice->flushCommandBuffer(commandBuffer, queue);
VkAccelerationStructureDeviceAddressInfoKHR accelerationDeviceAddressInfo{};
accelerationDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR;
accelerationDeviceAddressInfo.accelerationStructure = bottomLevelAS.handle;
bottomLevelAS.deviceAddress = vkGetAccelerationStructureDeviceAddressKHR(device, &accelerationDeviceAddressInfo);
deleteScratchBuffer(scratchBuffer);
}
/*
The top level acceleration structure contains the scene's object instances
*/
void createTopLevelAccelerationStructure()
{
VkTransformMatrixKHR transformMatrix = {
1.0f, 0.0f, 0.0f, 0.0f,
0.0f, 1.0f, 0.0f, 0.0f,
0.0f, 0.0f, 1.0f, 0.0f };
VkAccelerationStructureInstanceKHR instance{};
instance.transform = transformMatrix;
instance.instanceCustomIndex = 0;
instance.mask = 0xFF;
instance.instanceShaderBindingTableRecordOffset = 0;
instance.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR;
instance.accelerationStructureReference = bottomLevelAS.deviceAddress;
// Buffer for instance data
vks::Buffer instancesBuffer;
VK_CHECK_RESULT(vulkanDevice->createBuffer(
VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR,
VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
&instancesBuffer,
sizeof(VkAccelerationStructureInstanceKHR),
&instance));
VkDeviceOrHostAddressConstKHR instanceDataDeviceAddress{};
instanceDataDeviceAddress.deviceAddress = getBufferDeviceAddress(instancesBuffer.buffer);
VkAccelerationStructureGeometryKHR accelerationStructureGeometry{};
accelerationStructureGeometry.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR;
accelerationStructureGeometry.geometryType = VK_GEOMETRY_TYPE_INSTANCES_KHR;
accelerationStructureGeometry.flags = VK_GEOMETRY_OPAQUE_BIT_KHR;
accelerationStructureGeometry.geometry.instances.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_INSTANCES_DATA_KHR;
accelerationStructureGeometry.geometry.instances.arrayOfPointers = VK_FALSE;
accelerationStructureGeometry.geometry.instances.data = instanceDataDeviceAddress;
// Get size info
/*
The pSrcAccelerationStructure, dstAccelerationStructure, and mode members of pBuildInfo are ignored. Any VkDeviceOrHostAddressKHR members of pBuildInfo are ignored by this command, except that the hostAddress member of VkAccelerationStructureGeometryTrianglesDataKHR::transformData will be examined to check if it is NULL.*
*/
VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
accelerationStructureBuildGeometryInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR;
accelerationStructureBuildGeometryInfo.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
accelerationStructureBuildGeometryInfo.geometryCount = 1;
accelerationStructureBuildGeometryInfo.pGeometries = &accelerationStructureGeometry;
uint32_t primitive_count = 1;
VkAccelerationStructureBuildSizesInfoKHR accelerationStructureBuildSizesInfo{};
accelerationStructureBuildSizesInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR;
vkGetAccelerationStructureBuildSizesKHR(
device,
VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
&accelerationStructureBuildGeometryInfo,
&primitive_count,
&accelerationStructureBuildSizesInfo);
createAccelerationStructureBuffer(topLevelAS, accelerationStructureBuildSizesInfo);
VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfo{};
accelerationStructureCreateInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
accelerationStructureCreateInfo.buffer = topLevelAS.buffer;
accelerationStructureCreateInfo.size = accelerationStructureBuildSizesInfo.accelerationStructureSize;
accelerationStructureCreateInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR;
vkCreateAccelerationStructureKHR(device, &accelerationStructureCreateInfo, nullptr, &topLevelAS.handle);
// Create a small scratch buffer used during build of the top level acceleration structure
ScratchBuffer scratchBuffer = createScratchBuffer(accelerationStructureBuildSizesInfo.buildScratchSize);
VkAccelerationStructureBuildGeometryInfoKHR accelerationBuildGeometryInfo{};
accelerationBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
accelerationBuildGeometryInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR;
accelerationBuildGeometryInfo.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
accelerationBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR;
accelerationBuildGeometryInfo.dstAccelerationStructure = topLevelAS.handle;
accelerationBuildGeometryInfo.geometryCount = 1;
accelerationBuildGeometryInfo.pGeometries = &accelerationStructureGeometry;
accelerationBuildGeometryInfo.scratchData.deviceAddress = scratchBuffer.deviceAddress;
VkAccelerationStructureBuildRangeInfoKHR accelerationStructureBuildRangeInfo{};
accelerationStructureBuildRangeInfo.primitiveCount = 1;
accelerationStructureBuildRangeInfo.primitiveOffset = 0;
accelerationStructureBuildRangeInfo.firstVertex = 0;
accelerationStructureBuildRangeInfo.transformOffset = 0;
std::vector<VkAccelerationStructureBuildRangeInfoKHR*> accelerationBuildStructureRangeInfos = { &accelerationStructureBuildRangeInfo };
// Build the acceleration structure on the device via a one-time command buffer submission
// Some implementations may support acceleration structure building on the host (VkPhysicalDeviceAccelerationStructureFeaturesKHR->accelerationStructureHostCommands), but we prefer device builds
VkCommandBuffer commandBuffer = vulkanDevice->createCommandBuffer(VK_COMMAND_BUFFER_LEVEL_PRIMARY, true);
vkCmdBuildAccelerationStructuresKHR(
commandBuffer,
1,
&accelerationBuildGeometryInfo,
accelerationBuildStructureRangeInfos.data());
vulkanDevice->flushCommandBuffer(commandBuffer, queue);
VkAccelerationStructureDeviceAddressInfoKHR accelerationDeviceAddressInfo{};
accelerationDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR;
accelerationDeviceAddressInfo.accelerationStructure = topLevelAS.handle;
topLevelAS.deviceAddress = vkGetAccelerationStructureDeviceAddressKHR(device, &accelerationDeviceAddressInfo);
deleteScratchBuffer(scratchBuffer);
instancesBuffer.destroy();
}
/*
Create the Shader Binding Tables that binds the programs and top-level acceleration structure
SBT Layout used in this sample:
/-----------\
| raygen |
|-----------|
| miss |
|-----------|
| hit |
\-----------/
*/
void createShaderBindingTables() {
const uint32_t handleSize = rayTracingPipelineProperties.shaderGroupHandleSize;
const uint32_t handleSizeAligned = vks::tools::alignedSize(rayTracingPipelineProperties.shaderGroupHandleSize, rayTracingPipelineProperties.shaderGroupHandleAlignment);
const uint32_t groupCount = static_cast<uint32_t>(shaderGroups.size());
const uint32_t sbtSize = groupCount * handleSizeAligned;
std::vector<uint8_t> shaderHandleStorage(sbtSize);
VK_CHECK_RESULT(vkGetRayTracingShaderGroupHandlesKHR(device, pipeline, 0, groupCount, sbtSize, shaderHandleStorage.data()));
createShaderBindingTable(shaderBindingTables.raygen, 1);
createShaderBindingTable(shaderBindingTables.miss, 1);
createShaderBindingTable(shaderBindingTables.hit, 1);
// Copy handles
memcpy(shaderBindingTables.raygen.mapped, shaderHandleStorage.data(), handleSize);
memcpy(shaderBindingTables.miss.mapped, shaderHandleStorage.data() + handleSizeAligned, handleSize);
memcpy(shaderBindingTables.hit.mapped, shaderHandleStorage.data() + handleSizeAligned * 2, handleSize);
}
/*
Create the descriptor sets used for the ray tracing dispatch
*/
void createDescriptorSets()
{
std::vector<VkDescriptorPoolSize> poolSizes = {
{ VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 1 },
{ VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1 },
{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 },
{ VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1 }
};
VkDescriptorPoolCreateInfo descriptorPoolCreateInfo = vks::initializers::descriptorPoolCreateInfo(poolSizes, 1);
VK_CHECK_RESULT(vkCreateDescriptorPool(device, &descriptorPoolCreateInfo, nullptr, &descriptorPool));
VkDescriptorSetAllocateInfo descriptorSetAllocateInfo = vks::initializers::descriptorSetAllocateInfo(descriptorPool, &descriptorSetLayout, 1);
VK_CHECK_RESULT(vkAllocateDescriptorSets(device, &descriptorSetAllocateInfo, &descriptorSet));
VkWriteDescriptorSetAccelerationStructureKHR descriptorAccelerationStructureInfo = vks::initializers::writeDescriptorSetAccelerationStructureKHR();
descriptorAccelerationStructureInfo.accelerationStructureCount = 1;
descriptorAccelerationStructureInfo.pAccelerationStructures = &topLevelAS.handle;
VkWriteDescriptorSet accelerationStructureWrite{};
accelerationStructureWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
// The specialized acceleration structure descriptor has to be chained
accelerationStructureWrite.pNext = &descriptorAccelerationStructureInfo;
accelerationStructureWrite.dstSet = descriptorSet;
accelerationStructureWrite.dstBinding = 0;
accelerationStructureWrite.descriptorCount = 1;
accelerationStructureWrite.descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR;
VkDescriptorImageInfo storageImageDescriptor{ VK_NULL_HANDLE, storageImage.view, VK_IMAGE_LAYOUT_GENERAL };
std::vector<VkWriteDescriptorSet> writeDescriptorSets = {
// Binding 0: Top level acceleration structure
accelerationStructureWrite,
// Binding 1: Ray tracing result image
vks::initializers::writeDescriptorSet(descriptorSet, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1, &storageImageDescriptor),
// Binding 2: Uniform data
vks::initializers::writeDescriptorSet(descriptorSet, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 2, &ubo.descriptor),
// Binding 3: Texture image
vks::initializers::writeDescriptorSet(descriptorSet, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 3, &texture.descriptor),
};
vkUpdateDescriptorSets(device, static_cast<uint32_t>(writeDescriptorSets.size()), writeDescriptorSets.data(), 0, VK_NULL_HANDLE);
}
/*
Create our ray tracing pipeline
*/
void createRayTracingPipeline()
{
std::vector<VkDescriptorSetLayoutBinding> setLayoutBindings = {
// Binding 0: Top level acceleration structure
vks::initializers::descriptorSetLayoutBinding(VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, 0),
// Binding 1: Ray tracing result image
vks::initializers::descriptorSetLayoutBinding(VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, VK_SHADER_STAGE_RAYGEN_BIT_KHR, 1),
// Binding 2: Uniform buffer
vks::initializers::descriptorSetLayoutBinding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_MISS_BIT_KHR, 2),
// Binding 3: Texture image
vks::initializers::descriptorSetLayoutBinding(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_ANY_HIT_BIT_KHR, 3)
};
VkDescriptorSetLayoutCreateInfo descriptorSetLayoutCI = vks::initializers::descriptorSetLayoutCreateInfo(setLayoutBindings);
VK_CHECK_RESULT(vkCreateDescriptorSetLayout(device, &descriptorSetLayoutCI, nullptr, &descriptorSetLayout));
// We pass buffer references for vertex and index buffers via push constants
VkPushConstantRange pushConstantRange{};
pushConstantRange.stageFlags = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_ANY_HIT_BIT_KHR;
pushConstantRange.offset = 0;
pushConstantRange.size = sizeof(uint64_t) * 2;
VkPipelineLayoutCreateInfo pipelineLayoutCI = vks::initializers::pipelineLayoutCreateInfo(&descriptorSetLayout, 1);
pipelineLayoutCI.pPushConstantRanges = &pushConstantRange;
pipelineLayoutCI.pushConstantRangeCount = 1;
VK_CHECK_RESULT(vkCreatePipelineLayout(device, &pipelineLayoutCI, nullptr, &pipelineLayout));
/*
Setup ray tracing shader groups
*/
std::vector<VkPipelineShaderStageCreateInfo> shaderStages;
// Ray generation group
{
shaderStages.push_back(loadShader(getShadersPath() + "raytracingtextures/raygen.rgen.spv", VK_SHADER_STAGE_RAYGEN_BIT_KHR));
VkRayTracingShaderGroupCreateInfoKHR shaderGroup{};
shaderGroup.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
shaderGroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
shaderGroup.generalShader = static_cast<uint32_t>(shaderStages.size()) - 1;
shaderGroup.closestHitShader = VK_SHADER_UNUSED_KHR;
shaderGroup.anyHitShader = VK_SHADER_UNUSED_KHR;
shaderGroup.intersectionShader = VK_SHADER_UNUSED_KHR;
shaderGroups.push_back(shaderGroup);
}
// Miss group
{
shaderStages.push_back(loadShader(getShadersPath() + "raytracingtextures/miss.rmiss.spv", VK_SHADER_STAGE_MISS_BIT_KHR));
VkRayTracingShaderGroupCreateInfoKHR shaderGroup{};
shaderGroup.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
shaderGroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
shaderGroup.generalShader = static_cast<uint32_t>(shaderStages.size()) - 1;
shaderGroup.closestHitShader = VK_SHADER_UNUSED_KHR;
shaderGroup.anyHitShader = VK_SHADER_UNUSED_KHR;
shaderGroup.intersectionShader = VK_SHADER_UNUSED_KHR;
shaderGroups.push_back(shaderGroup);
}
// Closest hit group for doing texture lookups
{
shaderStages.push_back(loadShader(getShadersPath() + "raytracingtextures/closesthit.rchit.spv", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR));
VkRayTracingShaderGroupCreateInfoKHR shaderGroup{};
shaderGroup.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
shaderGroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR;
shaderGroup.generalShader = VK_SHADER_UNUSED_KHR;
shaderGroup.closestHitShader = static_cast<uint32_t>(shaderStages.size()) - 1;
shaderGroup.intersectionShader = VK_SHADER_UNUSED_KHR;
// This group also uses an anyhit shader for doing transparency (see anyhit.rahit for details)
shaderStages.push_back(loadShader(getShadersPath() + "raytracingtextures/anyhit.rahit.spv", VK_SHADER_STAGE_ANY_HIT_BIT_KHR));
shaderGroup.anyHitShader = static_cast<uint32_t>(shaderStages.size()) - 1;
shaderGroups.push_back(shaderGroup);
}
/*
Create the ray tracing pipeline
*/
VkRayTracingPipelineCreateInfoKHR rayTracingPipelineCI{};
rayTracingPipelineCI.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR;
rayTracingPipelineCI.stageCount = static_cast<uint32_t>(shaderStages.size());
rayTracingPipelineCI.pStages = shaderStages.data();
rayTracingPipelineCI.groupCount = static_cast<uint32_t>(shaderGroups.size());
rayTracingPipelineCI.pGroups = shaderGroups.data();
rayTracingPipelineCI.maxPipelineRayRecursionDepth = 1;
rayTracingPipelineCI.layout = pipelineLayout;
VK_CHECK_RESULT(vkCreateRayTracingPipelinesKHR(device, VK_NULL_HANDLE, VK_NULL_HANDLE, 1, &rayTracingPipelineCI, nullptr, &pipeline));
}
/*
Create the uniform buffer used to pass matrices to the ray tracing ray generation shader
*/
void createUniformBuffer()
{
VK_CHECK_RESULT(vulkanDevice->createBuffer(
VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
&ubo,
sizeof(uniformData),
&uniformData));
VK_CHECK_RESULT(ubo.map());
updateUniformBuffers();
}
/*
If the window has been resized, we need to recreate the storage image and it's descriptor
*/
void handleResize()
{
// Recreate image
createStorageImage(swapChain.colorFormat, { width, height, 1 });
// Update descriptor
VkDescriptorImageInfo storageImageDescriptor{ VK_NULL_HANDLE, storageImage.view, VK_IMAGE_LAYOUT_GENERAL };
VkWriteDescriptorSet resultImageWrite = vks::initializers::writeDescriptorSet(descriptorSet, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1, &storageImageDescriptor);
vkUpdateDescriptorSets(device, 1, &resultImageWrite, 0, VK_NULL_HANDLE);
resized = false;
}
/*
Command buffer generation
*/
void buildCommandBuffers()
{
if (resized)
{
handleResize();
}
VkCommandBufferBeginInfo cmdBufInfo = vks::initializers::commandBufferBeginInfo();
VkImageSubresourceRange subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
for (int32_t i = 0; i < drawCmdBuffers.size(); ++i)
{
VK_CHECK_RESULT(vkBeginCommandBuffer(drawCmdBuffers[i], &cmdBufInfo));
/*
Dispatch the ray tracing commands
*/
vkCmdBindPipeline(drawCmdBuffers[i], VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, pipeline);
vkCmdBindDescriptorSets(drawCmdBuffers[i], VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, pipelineLayout, 0, 1, &descriptorSet, 0, 0);
struct BufferReferences {
uint64_t vertices;
uint64_t indices;
} bufferReferences;
bufferReferences.vertices = getBufferDeviceAddress(vertexBuffer.buffer);
bufferReferences.indices = getBufferDeviceAddress(indexBuffer.buffer);
// We set the buffer references for the mesh to be rendered using a push constant
// If we wanted to render multiple objecets this would make it very easy to access their vertex and index buffers
vkCmdPushConstants(drawCmdBuffers[i], pipelineLayout, VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_ANY_HIT_BIT_KHR, 0, sizeof(uint64_t) * 2, &bufferReferences);
VkStridedDeviceAddressRegionKHR emptySbtEntry = {};
vkCmdTraceRaysKHR(
drawCmdBuffers[i],
&shaderBindingTables.raygen.stridedDeviceAddressRegion,
&shaderBindingTables.miss.stridedDeviceAddressRegion,
&shaderBindingTables.hit.stridedDeviceAddressRegion,
&emptySbtEntry,
width,
height,
1);
/*
Copy ray tracing output to swap chain image
*/
// Prepare current swap chain image as transfer destination
vks::tools::setImageLayout(
drawCmdBuffers[i],
swapChain.images[i],
VK_IMAGE_LAYOUT_UNDEFINED,
VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
subresourceRange);
// Prepare ray tracing output image as transfer source
vks::tools::setImageLayout(
drawCmdBuffers[i],
storageImage.image,
VK_IMAGE_LAYOUT_GENERAL,
VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
subresourceRange);
VkImageCopy copyRegion{};
copyRegion.srcSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 };
copyRegion.srcOffset = { 0, 0, 0 };
copyRegion.dstSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 };
copyRegion.dstOffset = { 0, 0, 0 };
copyRegion.extent = { width, height, 1 };
vkCmdCopyImage(drawCmdBuffers[i], storageImage.image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, swapChain.images[i], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copyRegion);
// Transition swap chain image back for presentation
vks::tools::setImageLayout(
drawCmdBuffers[i],
swapChain.images[i],
VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
subresourceRange);
// Transition ray tracing output image back to general layout
vks::tools::setImageLayout(
drawCmdBuffers[i],
storageImage.image,
VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
VK_IMAGE_LAYOUT_GENERAL,
subresourceRange);
drawUI(drawCmdBuffers[i], frameBuffers[i]);
VK_CHECK_RESULT(vkEndCommandBuffer(drawCmdBuffers[i]));
}
}
void updateUniformBuffers()
{
uniformData.projInverse = glm::inverse(camera.matrices.perspective);
uniformData.viewInverse = glm::inverse(camera.matrices.view);
memcpy(ubo.mapped, &uniformData, sizeof(uniformData));
}
void getEnabledFeatures()
{
// Enable features required for ray tracing using feature chaining via pNext
enabledBufferDeviceAddresFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES;
enabledBufferDeviceAddresFeatures.bufferDeviceAddress = VK_TRUE;
enabledRayTracingPipelineFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_FEATURES_KHR;
enabledRayTracingPipelineFeatures.rayTracingPipeline = VK_TRUE;
enabledRayTracingPipelineFeatures.pNext = &enabledBufferDeviceAddresFeatures;
enabledAccelerationStructureFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR;
enabledAccelerationStructureFeatures.accelerationStructure = VK_TRUE;
enabledAccelerationStructureFeatures.pNext = &enabledRayTracingPipelineFeatures;
deviceCreatepNextChain = &enabledAccelerationStructureFeatures;
}
void loadAssets()
{
texture.loadFromFile(getAssetPath() + "textures/gratefloor_rgba.ktx", VK_FORMAT_R8G8B8A8_UNORM, vulkanDevice, queue);
}
void prepare()
{
VulkanRaytracingSample::prepare();
loadAssets();
// Create the acceleration structures used to render the ray traced scene
createBottomLevelAccelerationStructure();
createTopLevelAccelerationStructure();
createStorageImage(swapChain.colorFormat, { width, height, 1 });
createUniformBuffer();
createRayTracingPipeline();
createShaderBindingTables();
createDescriptorSets();
buildCommandBuffers();
prepared = true;
}
void draw()
{
VulkanExampleBase::prepareFrame();
submitInfo.commandBufferCount = 1;
submitInfo.pCommandBuffers = &drawCmdBuffers[currentBuffer];
VK_CHECK_RESULT(vkQueueSubmit(queue, 1, &submitInfo, VK_NULL_HANDLE));
VulkanExampleBase::submitFrame();
}
virtual void render()
{
if (!prepared)
return;
draw();
if (camera.updated)
updateUniformBuffers();
}
};
VULKAN_EXAMPLE_MAIN()