/* * Copyright (c) 2014-2021, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION * SPDX-License-Identifier: Apache-2.0 */ #include "raytrace.hpp" #include "nvh/fileoperations.hpp" #include "nvvk/descriptorsets_vk.hpp" #include "nvh/alignment.hpp" #include "nvvk/shaders_vk.hpp" #include "obj_loader.h" #include "nvvk/buffers_vk.hpp" extern std::vector defaultSearchPaths; void Raytracer::setup(const VkDevice& device, const VkPhysicalDevice& physicalDevice, nvvk::ResourceAllocator* allocator, uint32_t queueFamily) { m_device = device; m_physicalDevice = physicalDevice; m_alloc = allocator; m_graphicsQueueIndex = queueFamily; // Requesting ray tracing properties VkPhysicalDeviceProperties2 prop2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2}; prop2.pNext = &m_rtProperties; vkGetPhysicalDeviceProperties2(m_physicalDevice, &prop2); m_rtBuilder.setup(m_device, allocator, m_graphicsQueueIndex); m_sbtWrapper.setup(device, queueFamily, allocator, m_rtProperties); m_debug.setup(device); } void Raytracer::destroy() { m_sbtWrapper.destroy(); m_rtBuilder.destroy(); vkDestroyDescriptorPool(m_device, m_rtDescPool, nullptr); vkDestroyDescriptorSetLayout(m_device, m_rtDescSetLayout, nullptr); vkDestroyPipeline(m_device, m_rtPipeline, nullptr); vkDestroyPipelineLayout(m_device, m_rtPipelineLayout, nullptr); m_alloc->destroy(m_rtSBTBuffer); } //-------------------------------------------------------------------------------------------------- // Converting a OBJ primitive to the ray tracing geometry used for the BLAS // auto Raytracer::objectToVkGeometryKHR(const ObjModel& model) { // Building part VkDeviceAddress vertexAddress = nvvk::getBufferDeviceAddress(m_device, model.vertexBuffer.buffer); VkDeviceAddress indexAddress = nvvk::getBufferDeviceAddress(m_device, model.indexBuffer.buffer); VkAccelerationStructureGeometryTrianglesDataKHR triangles{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR}; triangles.vertexFormat = VK_FORMAT_R32G32B32_SFLOAT; triangles.vertexData.deviceAddress = vertexAddress; triangles.vertexStride = sizeof(VertexObj); triangles.indexType = VK_INDEX_TYPE_UINT32; triangles.indexData.deviceAddress = indexAddress; triangles.transformData = {}; triangles.maxVertex = model.nbVertices; // Setting up the build info of the acceleration VkAccelerationStructureGeometryKHR asGeom{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR}; asGeom.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR; asGeom.flags = VK_GEOMETRY_NO_DUPLICATE_ANY_HIT_INVOCATION_BIT_KHR; // For AnyHit asGeom.geometry.triangles = triangles; VkAccelerationStructureBuildRangeInfoKHR offset; offset.firstVertex = 0; offset.primitiveCount = model.nbIndices / 3; // Nb triangles offset.primitiveOffset = 0; offset.transformOffset = 0; nvvk::RaytracingBuilderKHR::BlasInput input; input.asGeometry.emplace_back(asGeom); input.asBuildOffsetInfo.emplace_back(offset); return input; } //-------------------------------------------------------------------------------------------------- // Returning the ray tracing geometry used for the BLAS, containing all spheres // auto Raytracer::implicitToVkGeometryKHR(const ImplInst& implicitObj) { VkDeviceAddress dataAddress = nvvk::getBufferDeviceAddress(m_device, implicitObj.implBuf.buffer); VkAccelerationStructureGeometryAabbsDataKHR aabbs{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR}; aabbs.data.deviceAddress = dataAddress; aabbs.stride = sizeof(ObjImplicit); // Setting up the build info of the acceleration VkAccelerationStructureGeometryKHR asGeom{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR}; asGeom.geometryType = VK_GEOMETRY_TYPE_AABBS_KHR; asGeom.flags = VK_GEOMETRY_NO_DUPLICATE_ANY_HIT_INVOCATION_BIT_KHR; // For AnyHit asGeom.geometry.aabbs = aabbs; VkAccelerationStructureBuildRangeInfoKHR offset; offset.firstVertex = 0; offset.primitiveCount = static_cast(implicitObj.objImpl.size()); // Nb aabb offset.primitiveOffset = 0; offset.transformOffset = 0; nvvk::RaytracingBuilderKHR::BlasInput input; input.asGeometry.emplace_back(asGeom); input.asBuildOffsetInfo.emplace_back(offset); return input; } void Raytracer::createBottomLevelAS(std::vector& models, ImplInst& implicitObj) { // BLAS - Storing each primitive in a geometry std::vector allBlas; allBlas.reserve(models.size()); for(const auto& obj : models) { auto blas = objectToVkGeometryKHR(obj); // We could add more geometry in each BLAS, but we add only one for now allBlas.emplace_back(blas); } // Adding implicit if(!implicitObj.objImpl.empty()) { auto blas = implicitToVkGeometryKHR(implicitObj); allBlas.emplace_back(blas); implicitObj.blasId = static_cast(allBlas.size() - 1); // remember blas ID for tlas } m_rtBuilder.buildBlas(allBlas, VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR | VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR); } void Raytracer::createTopLevelAS(std::vector& instances, ImplInst& implicitObj) { std::vector tlas; auto nbObj = static_cast(instances.size()) - 1; // minus the implicit (for material) tlas.reserve(instances.size()); for(uint32_t i = 0; i < nbObj; i++) { VkAccelerationStructureInstanceKHR rayInst{}; rayInst.transform = nvvk::toTransformMatrixKHR(instances[i].transform); // Position of the instance rayInst.instanceCustomIndex = instances[i].objIndex; // gl_InstanceCustomIndexEXT rayInst.accelerationStructureReference = m_rtBuilder.getBlasDeviceAddress(instances[i].objIndex); rayInst.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR; rayInst.mask = 0xFF; // Only be hit if rayMask & instance.mask != 0 rayInst.instanceShaderBindingTableRecordOffset = 0; // We will use the same hit group for all objects tlas.emplace_back(rayInst); } // Add the blas containing all implicit if(!implicitObj.objImpl.empty()) { VkAccelerationStructureInstanceKHR rayInst{}; rayInst.transform = nvvk::toTransformMatrixKHR(implicitObj.transform); // Position of the instance rayInst.instanceCustomIndex = instances[nbObj].objIndex; rayInst.accelerationStructureReference = m_rtBuilder.getBlasDeviceAddress(static_cast(implicitObj.blasId)); rayInst.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR; rayInst.mask = 0xFF; // Only be hit if rayMask & instance.mask != 0 rayInst.instanceShaderBindingTableRecordOffset = 1; // We will use the same hit group for all objects (the second one) tlas.emplace_back(rayInst); } m_rtBuilder.buildTlas(tlas, VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR); } //-------------------------------------------------------------------------------------------------- // This descriptor set holds the Acceleration structure and the output image // void Raytracer::createRtDescriptorSet(const VkImageView& outputImage) { using vkDSLB = VkDescriptorSetLayoutBinding; m_rtDescSetLayoutBind.addBinding(RtxBindings::eTlas, VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 1, VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR); // TLAS m_rtDescSetLayoutBind.addBinding(RtxBindings::eOutImage, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1, VK_SHADER_STAGE_RAYGEN_BIT_KHR); // Output image m_rtDescPool = m_rtDescSetLayoutBind.createPool(m_device); m_rtDescSetLayout = m_rtDescSetLayoutBind.createLayout(m_device); VkDescriptorSetAllocateInfo allocateInfo{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO}; allocateInfo.descriptorPool = m_rtDescPool; allocateInfo.descriptorSetCount = 1; allocateInfo.pSetLayouts = &m_rtDescSetLayout; vkAllocateDescriptorSets(m_device, &allocateInfo, &m_rtDescSet); VkAccelerationStructureKHR tlas = m_rtBuilder.getAccelerationStructure(); VkWriteDescriptorSetAccelerationStructureKHR descASInfo{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR}; descASInfo.accelerationStructureCount = 1; descASInfo.pAccelerationStructures = &tlas; VkDescriptorImageInfo imageInfo{{}, outputImage, VK_IMAGE_LAYOUT_GENERAL}; std::vector writes; writes.emplace_back(m_rtDescSetLayoutBind.makeWrite(m_rtDescSet, RtxBindings::eTlas, &descASInfo)); writes.emplace_back(m_rtDescSetLayoutBind.makeWrite(m_rtDescSet, RtxBindings::eOutImage, &imageInfo)); vkUpdateDescriptorSets(m_device, static_cast(writes.size()), writes.data(), 0, nullptr); } //-------------------------------------------------------------------------------------------------- // Writes the output image to the descriptor set // - Required when changing resolution // void Raytracer::updateRtDescriptorSet(const VkImageView& outputImage) { VkDescriptorImageInfo imageInfo{{}, outputImage, VK_IMAGE_LAYOUT_GENERAL}; VkWriteDescriptorSet wds = m_rtDescSetLayoutBind.makeWrite(m_rtDescSet, RtxBindings::eOutImage, &imageInfo); vkUpdateDescriptorSets(m_device, 1, &wds, 0, nullptr); } //-------------------------------------------------------------------------------------------------- // Pipeline for the ray tracer: all shaders, raygen, chit, miss // void Raytracer::createRtPipeline(VkDescriptorSetLayout& sceneDescLayout) { enum StageIndices { eRaygen, eMiss, eMiss2, eClosestHit, eAnyHit, eClosestHit1, eAnyHit1, eIntersect, eCall0, eCall1, eCall2, eShaderGroupCount }; // All stages std::array stages{}; VkPipelineShaderStageCreateInfo stage{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO}; stage.pName = "main"; // All the same entry point // Raygen stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rgen.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_RAYGEN_BIT_KHR; stages[eRaygen] = stage; // Miss stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rmiss.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_MISS_BIT_KHR; stages[eMiss] = stage; // The second miss shader is invoked when a shadow ray misses the geometry. It simply indicates that no occlusion has been found stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytraceShadow.rmiss.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_MISS_BIT_KHR; stages[eMiss2] = stage; // Hit Group - Closest Hit stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rchit.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR; stages[eClosestHit] = stage; // Hit Group - Closest Hit stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rahit.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_ANY_HIT_BIT_KHR; stages[eAnyHit] = stage; // Hit Group - 1 stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace2.rchit.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR; stages[eClosestHit1] = stage; // Hit stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace2.rahit.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_ANY_HIT_BIT_KHR; stages[eAnyHit1] = stage; // Hit stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rint.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_INTERSECTION_BIT_KHR; stages[eIntersect] = stage; // Call0 stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/light_point.rcall.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_CALLABLE_BIT_KHR; stages[eCall0] = stage; // Call1 stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/light_spot.rcall.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_CALLABLE_BIT_KHR; stages[eCall1] = stage; // Call2 stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/light_inf.rcall.spv", true, defaultSearchPaths, true)); stage.stage = VK_SHADER_STAGE_CALLABLE_BIT_KHR; stages[eCall2] = stage; // Shader groups VkRayTracingShaderGroupCreateInfoKHR group{VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR}; group.anyHitShader = VK_SHADER_UNUSED_KHR; group.closestHitShader = VK_SHADER_UNUSED_KHR; group.generalShader = VK_SHADER_UNUSED_KHR; group.intersectionShader = VK_SHADER_UNUSED_KHR; // Raygen group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; group.generalShader = eRaygen; m_rtShaderGroups.push_back(group); // Miss group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; group.generalShader = eMiss; m_rtShaderGroups.push_back(group); // Shadow Miss group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; group.generalShader = eMiss2; m_rtShaderGroups.push_back(group); // closest hit shader group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR; group.generalShader = VK_SHADER_UNUSED_KHR; group.closestHitShader = eClosestHit; group.anyHitShader = eAnyHit; m_rtShaderGroups.push_back(group); // closest hit shader group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR; group.generalShader = VK_SHADER_UNUSED_KHR; group.closestHitShader = eClosestHit1; group.anyHitShader = eAnyHit1; group.intersectionShader = eIntersect; m_rtShaderGroups.push_back(group); // Callable shaders group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; group.closestHitShader = VK_SHADER_UNUSED_KHR; group.anyHitShader = VK_SHADER_UNUSED_KHR; group.intersectionShader = VK_SHADER_UNUSED_KHR; group.generalShader = eCall0; m_rtShaderGroups.push_back(group); group.generalShader = eCall1; m_rtShaderGroups.push_back(group); group.generalShader = eCall2; m_rtShaderGroups.push_back(group); // Push constant: we want to be able to update constants used by the shaders VkPushConstantRange pushConstant{VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_MISS_BIT_KHR | VK_SHADER_STAGE_CALLABLE_BIT_KHR, 0, sizeof(PushConstantRay)}; VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; pipelineLayoutCreateInfo.pushConstantRangeCount = 1; pipelineLayoutCreateInfo.pPushConstantRanges = &pushConstant; // Descriptor sets: one specific to ray tracing, and one shared with the rasterization pipeline std::vector rtDescSetLayouts = {m_rtDescSetLayout, sceneDescLayout}; pipelineLayoutCreateInfo.setLayoutCount = static_cast(rtDescSetLayouts.size()); pipelineLayoutCreateInfo.pSetLayouts = rtDescSetLayouts.data(); vkCreatePipelineLayout(m_device, &pipelineLayoutCreateInfo, nullptr, &m_rtPipelineLayout); // Assemble the shader stages and recursion depth info into the ray tracing pipeline VkRayTracingPipelineCreateInfoKHR rayPipelineInfo{VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR}; rayPipelineInfo.stageCount = static_cast(stages.size()); // Stages are shaders rayPipelineInfo.pStages = stages.data(); rayPipelineInfo.groupCount = static_cast(m_rtShaderGroups.size()); rayPipelineInfo.pGroups = m_rtShaderGroups.data(); // The ray tracing process can shoot rays from the camera, and a shadow ray can be shot from the // hit points of the camera rays, hence a recursion level of 2. This number should be kept as low // as possible for performance reasons. Even recursive ray tracing should be flattened into a loop // in the ray generation to avoid deep recursion. rayPipelineInfo.maxPipelineRayRecursionDepth = 2; // Ray depth rayPipelineInfo.layout = m_rtPipelineLayout; vkCreateRayTracingPipelinesKHR(m_device, {}, {}, 1, &rayPipelineInfo, nullptr, &m_rtPipeline); m_sbtWrapper.create(m_rtPipeline, rayPipelineInfo); for(auto& s : stages) vkDestroyShaderModule(m_device, s.module, nullptr); } //-------------------------------------------------------------------------------------------------- // Ray Tracing the scene // void Raytracer::raytrace(const VkCommandBuffer& cmdBuf, const nvmath::vec4f& clearColor, VkDescriptorSet& sceneDescSet, VkExtent2D& size, PushConstantRaster& sceneConstants) { m_debug.beginLabel(cmdBuf, "Ray trace"); // Initializing push constant values m_pcRay.clearColor = clearColor; m_pcRay.lightPosition = sceneConstants.lightPosition; m_pcRay.lightIntensity = sceneConstants.lightIntensity; m_pcRay.lightDirection = sceneConstants.lightDirection; m_pcRay.lightSpotCutoff = sceneConstants.lightSpotCutoff; m_pcRay.lightSpotOuterCutoff = sceneConstants.lightSpotOuterCutoff; m_pcRay.lightType = sceneConstants.lightType; m_pcRay.frame = sceneConstants.frame; std::vector descSets{m_rtDescSet, sceneDescSet}; vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, m_rtPipeline); vkCmdBindDescriptorSets(cmdBuf, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, m_rtPipelineLayout, 0, (uint32_t)descSets.size(), descSets.data(), 0, nullptr); vkCmdPushConstants(cmdBuf, m_rtPipelineLayout, VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_MISS_BIT_KHR | VK_SHADER_STAGE_CALLABLE_BIT_KHR, 0, sizeof(PushConstantRay), &m_pcRay); auto& regions = m_sbtWrapper.getRegions(); vkCmdTraceRaysKHR(cmdBuf, ®ions[0], ®ions[1], ®ions[2], ®ions[3], size.width, size.height, 1); m_debug.endLabel(cmdBuf); }