diff --git a/lvk/vulkan/VulkanClasses.cpp b/lvk/vulkan/VulkanClasses.cpp index 426a2eabdf..bbf8b8f0b2 100644 --- a/lvk/vulkan/VulkanClasses.cpp +++ b/lvk/vulkan/VulkanClasses.cpp @@ -709,6 +709,11 @@ VkSurfaceFormatKHR chooseSwapSurfaceFormat(const std::vector return formats[0]; } +VkDeviceSize bufferSize(lvk::VulkanContext& ctx, const lvk::Holder& handle) { + lvk::VulkanBuffer* buffer = ctx.buffersPool_.get(handle); + return buffer ? buffer->bufferSize_ : 0; +} + } // namespace namespace lvk { @@ -2853,16 +2858,19 @@ void lvk::CommandBuffer::cmdUpdateTLAS(AccelStructHandle handle, BufferHandle in &accelerationStructureBuildGeometryInfo, &as->buildRangeInfo.primitiveCount, &accelerationStructureBuildSizesInfo); - - lvk::Holder scratchBuffer = ctx_->createBuffer( - lvk::BufferDesc{ - .usage = lvk::BufferUsageBits_Storage, - .storage = lvk::StorageType_Device, - .size = accelerationStructureBuildSizesInfo.buildScratchSize, - .debugName = "scratchBuffer", - }, - nullptr, - nullptr); + + if (!as->scratchBuffer.valid() || bufferSize(*ctx_, as->scratchBuffer) < accelerationStructureBuildSizesInfo.buildScratchSize) { + LLOGD("Recreating scratch buffer for TLAS update"); + as->scratchBuffer = ctx_->createBuffer( + lvk::BufferDesc{ + .usage = lvk::BufferUsageBits_Storage, + .storage = lvk::StorageType_Device, + .size = accelerationStructureBuildSizesInfo.buildScratchSize, + .debugName = "scratchBuffer", + }, + nullptr, + nullptr); + } const VkAccelerationStructureBuildGeometryInfoKHR accelerationBuildGeometryInfo = { .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR, @@ -2873,7 +2881,7 @@ void lvk::CommandBuffer::cmdUpdateTLAS(AccelStructHandle handle, BufferHandle in .dstAccelerationStructure = as->vkHandle, .geometryCount = 1, .pGeometries = &accelerationStructureGeometry, - .scratchData = {.deviceAddress = ctx_->gpuAddress(scratchBuffer)}, + .scratchData = {.deviceAddress = ctx_->gpuAddress(as->scratchBuffer)}, }; const VkAccelerationStructureBuildRangeInfoKHR* accelerationBuildStructureRangeInfos[] = {&as->buildRangeInfo}; @@ -4359,7 +4367,7 @@ lvk::AccelStructHandle lvk::VulkanContext::createTLAS(const AccelStructDesc& des }, nullptr, outResult); - + const VkAccelerationStructureBuildGeometryInfoKHR accelerationBuildGeometryInfo = { .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR, .type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR, @@ -4370,6 +4378,10 @@ lvk::AccelStructHandle lvk::VulkanContext::createTLAS(const AccelStructDesc& des .pGeometries = &accelerationStructureGeometry, .scratchData = {.deviceAddress = gpuAddress(scratchBuffer)}, }; + if (desc.buildFlags & lvk::AccelStructBuildFlagBits_AllowUpdate) { + // Store scratch buffer for future updates + accelStruct.scratchBuffer = std::move(scratchBuffer); + } const VkAccelerationStructureBuildRangeInfoKHR* accelerationBuildStructureRangeInfos[] = {&accelStruct.buildRangeInfo}; diff --git a/lvk/vulkan/VulkanClasses.h b/lvk/vulkan/VulkanClasses.h index 295ec205e6..ade4051e3f 100644 --- a/lvk/vulkan/VulkanClasses.h +++ b/lvk/vulkan/VulkanClasses.h @@ -321,6 +321,7 @@ struct AccelerationStructure { VkAccelerationStructureKHR vkHandle = VK_NULL_HANDLE; uint64_t deviceAddress = 0; lvk::Holder buffer; + lvk::Holder scratchBuffer; // Store only for TLAS }; class CommandBuffer final : public ICommandBuffer {