diff --git a/gprt/gprt.cpp b/gprt/gprt.cpp index 52bb751..fb90167 100644 --- a/gprt/gprt.cpp +++ b/gprt/gprt.cpp @@ -660,18 +660,7 @@ struct Context { Stage Scatter; Stage ScatterPayload; - // For now, copied over from sample - VkDescriptorSetLayout m_SortDescriptorSetLayoutInputOutputs; - VkDescriptorSetLayout m_SortDescriptorSetLayoutScan; - VkDescriptorSetLayout m_SortDescriptorSetLayoutScratch; - - VkDescriptorSet m_SortDescriptorSetInputOutput[2]; - VkDescriptorSet m_SortDescriptorSetScanSets[2]; - VkDescriptorSet m_SortDescriptorSetScratch; - VkPipelineLayout layout; - - VkDescriptorPool pool = VK_NULL_HANDLE; }; SortStages sortStages; @@ -6235,18 +6224,6 @@ Context::setupInternalPrograms() { // currently not using cache. VkPipelineCache cache = VK_NULL_HANDLE; - VkDescriptorPoolSize poolSize; - poolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; - poolSize.descriptorCount = 16; - - VkDescriptorPoolCreateInfo descriptorPoolInfo{}; - descriptorPoolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO; - descriptorPoolInfo.poolSizeCount = 1; - descriptorPoolInfo.pPoolSizes = &poolSize; - descriptorPoolInfo.maxSets = 5; - descriptorPoolInfo.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT; - VK_CHECK_RESULT(vkCreateDescriptorPool(logicalDevice, &descriptorPoolInfo, nullptr, &sortStages.pool)); - sortStages.Count.entryPoint = "Count"; sortStages.CountReduce.entryPoint = "CountReduce"; sortStages.Scan.entryPoint = "Scan"; @@ -6254,72 +6231,6 @@ Context::setupInternalPrograms() { sortStages.Scatter.entryPoint = "Scatter"; sortStages.ScatterPayload.entryPoint = "ScatterPayload"; - // Create binding for Radix sort passes - VkDescriptorSetLayoutBinding layout_bindings_set_InputOutputs[] = { - {0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // SrcBuffer (sort) - {1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // DstBuffer (sort) - {2, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // ScrPayload (sort only) - {3, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // DstPayload (sort only) - }; - - VkDescriptorSetLayoutBinding layout_bindings_set_Scan[] = { - {0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // ScanSrc - {1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // ScanDst - {2, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // ScanScratch - }; - - VkDescriptorSetLayoutBinding layout_bindings_set_Scratch[] = { - {0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // Scratch (sort only) - {1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}, // Scratch (reduced) - }; - - auto AllocDescriptor = [&](VkDescriptorPool pool, VkDescriptorSetLayout layout, VkDescriptorSet *descriptorSet) { - VkDescriptorSetAllocateInfo descriptorSetAllocateInfo{}; - descriptorSetAllocateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO; - descriptorSetAllocateInfo.descriptorPool = pool; - descriptorSetAllocateInfo.pSetLayouts = &layout; - descriptorSetAllocateInfo.descriptorSetCount = 1; - descriptorSetAllocateInfo.pNext = nullptr; - VkResult err = vkAllocateDescriptorSets(logicalDevice, &descriptorSetAllocateInfo, descriptorSet); - if (err != VK_SUCCESS) { - LOG_ERROR("failed to allocate descriptor! \n" + errorString(err)); - } - }; - - VkResult vkResult; - - VkDescriptorSetLayoutCreateInfo descriptor_set_layout_create_info = { - VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO}; - descriptor_set_layout_create_info.pNext = nullptr; - descriptor_set_layout_create_info.flags = 0; - descriptor_set_layout_create_info.pBindings = layout_bindings_set_InputOutputs; - descriptor_set_layout_create_info.bindingCount = 4; - vkResult = vkCreateDescriptorSetLayout(logicalDevice, &descriptor_set_layout_create_info, nullptr, - &sortStages.m_SortDescriptorSetLayoutInputOutputs); - VK_CHECK_RESULT(vkResult); - AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutInputOutputs, - &sortStages.m_SortDescriptorSetInputOutput[0]); - AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutInputOutputs, - &sortStages.m_SortDescriptorSetInputOutput[1]); - - descriptor_set_layout_create_info.pBindings = layout_bindings_set_Scan; - descriptor_set_layout_create_info.bindingCount = 3; - vkResult = vkCreateDescriptorSetLayout(logicalDevice, &descriptor_set_layout_create_info, nullptr, - &sortStages.m_SortDescriptorSetLayoutScan); - VK_CHECK_RESULT(vkResult); - AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScan, - &sortStages.m_SortDescriptorSetScanSets[0]); - AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScan, - &sortStages.m_SortDescriptorSetScanSets[1]); - - descriptor_set_layout_create_info.pBindings = layout_bindings_set_Scratch; - descriptor_set_layout_create_info.bindingCount = 2; - vkResult = vkCreateDescriptorSetLayout(logicalDevice, &descriptor_set_layout_create_info, nullptr, - &sortStages.m_SortDescriptorSetLayoutScratch); - VK_CHECK_RESULT(vkResult); - AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScratch, - &sortStages.m_SortDescriptorSetScratch); - // Create constant range representing our static constant VkPushConstantRange constant_range; constant_range.stageFlags = VK_SHADER_STAGE_ALL; @@ -6330,11 +6241,6 @@ Context::setupInternalPrograms() { VkPipelineLayoutCreateInfo layout_create_info = {VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; layout_create_info.pNext = nullptr; layout_create_info.flags = 0; - layout_create_info.setLayoutCount = 3; - VkDescriptorSetLayout layouts[] = {sortStages.m_SortDescriptorSetLayoutInputOutputs, - sortStages.m_SortDescriptorSetLayoutScan, - sortStages.m_SortDescriptorSetLayoutScratch}; - layout_create_info.pSetLayouts = layouts; layout_create_info.pushConstantRangeCount = 1; layout_create_info.pPushConstantRanges = &constant_range; @@ -6589,16 +6495,6 @@ Context::destroyInternalPrograms() { vkDestroyPipelineLayout(logicalDevice, sortStages.layout, nullptr); sortStages.layout = VK_NULL_HANDLE; - - vkFreeDescriptorSets(logicalDevice, sortStages.pool, 1, &sortStages.m_SortDescriptorSetScratch); - vkFreeDescriptorSets(logicalDevice, sortStages.pool, 2, sortStages.m_SortDescriptorSetScanSets); - vkFreeDescriptorSets(logicalDevice, sortStages.pool, 2, sortStages.m_SortDescriptorSetInputOutput); - - vkDestroyDescriptorSetLayout(logicalDevice, sortStages.m_SortDescriptorSetLayoutScratch, nullptr); - vkDestroyDescriptorSetLayout(logicalDevice, sortStages.m_SortDescriptorSetLayoutScan, nullptr); - vkDestroyDescriptorSetLayout(logicalDevice, sortStages.m_SortDescriptorSetLayoutInputOutputs, nullptr); - - vkDestroyDescriptorPool(logicalDevice, sortStages.pool, nullptr); } std::vector progNames; @@ -8963,242 +8859,83 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe Context *context = (Context *) _context; Buffer *keys = (Buffer *) _keys; Buffer *values = (Buffer *) _values; - Buffer *scratch = (Buffer *) _scratch; - - bool bHasPayload = false; - if (values) { - if (keys->getSize() != values->getSize()) - LOG_ERROR("Keys and Values buffers must be equal in size\n"); - - bHasPayload = true; + bool bHasPayload = values != nullptr; + if (values && keys->getSize() != values->getSize()) { + LOG_ERROR("Keys and Values buffers must be equal in size\n"); + return; } - uint32_t numKeys = uint32_t(keys->getSize() / sizeof(uint64_t)); - uint32_t maxNumThreadgroups = 800; - ParallelSortCB constantBufferData = {0}; - - // Allocate the scratch buffers needed for radix sort - auto alignedSize = [](size_t value, size_t alignment) -> size_t { - return (value + alignment - 1) & ~(alignment - 1); - }; - uint32_t offsetAlignment = (uint32_t) context->deviceProperties.limits.minStorageBufferOffsetAlignment; + if (!numKeys) return; + Buffer *scratch = _scratch ? (Buffer *)_scratch : (Buffer *)gprtDeviceBufferCreate(_context); - uint64_t scratchBufferSize; - uint64_t reducedScratchBufferSize; + auto alignedSize = [](size_t value) -> size_t { return (value + 7) & ~size_t(7); }; + uint64_t scratchBufferSize, reducedScratchBufferSize; ParallelSort_CalculateScratchResourceSize(numKeys, scratchBufferSize, reducedScratchBufferSize); - scratchBufferSize = alignedSize(scratchBufferSize, offsetAlignment); - reducedScratchBufferSize = alignedSize(reducedScratchBufferSize, offsetAlignment); - - uint64_t keysSize = alignedSize(keys->size, offsetAlignment); - uint64_t valuesSize = ((bHasPayload) ? alignedSize(values->size, offsetAlignment) : 0); - - scratch->resize(keysSize + valuesSize + scratchBufferSize + reducedScratchBufferSize, - /*don't transfer old contents*/ false); - // All offsets must be a multiple of device limit VkPhysicalDeviceLimits::minStorageBufferOffseteAlignment - size_t valuesOffset = keysSize; + uint64_t keysSize = alignedSize(keys->size); + uint64_t valuesSize = values ? alignedSize(values->size) : 0; size_t scratchOffset = keysSize + valuesSize; - size_t reducedScratchOffset = keysSize + valuesSize + scratchBufferSize; + size_t reducedScratchOffset = scratchOffset + alignedSize(scratchBufferSize); + size_t requiredScratch = reducedScratchOffset + alignedSize(reducedScratchBufferSize); + if (scratch->size < requiredScratch) { + // Complete compute uses before replacing an existing scratch allocation. + if (scratch->size && context->synchronizeCompute() != VK_SUCCESS) return; + scratch->resize(requiredScratch, false); + } - uint32_t NumThreadgroupsToRun; - uint32_t NumReducedThreadgroupsToRun; - ParallelSort_SetConstantAndDispatchData(numKeys, maxNumThreadgroups, constantBufferData, NumThreadgroupsToRun, + ParallelSortCB constantBufferData = {}; + uint32_t NumThreadgroupsToRun, NumReducedThreadgroupsToRun; + ParallelSort_SetConstantAndDispatchData(numKeys, 800, constantBufferData, NumThreadgroupsToRun, NumReducedThreadgroupsToRun); - - auto BindUAVBuffer = [&](VkBuffer *pBuffer, VkDeviceSize *Offsets, VkDescriptorSet &DescriptorSet, - uint32_t Binding /*=0*/, uint32_t Count /*=1*/) { - std::vector bufferInfos; - for (uint32_t i = 0; i < Count; i++) { - VkDescriptorBufferInfo bufferInfo; - bufferInfo.buffer = pBuffer[i]; - bufferInfo.offset = Offsets[i]; - bufferInfo.range = VK_WHOLE_SIZE; - bufferInfos.push_back(bufferInfo); - } - - VkWriteDescriptorSet write_set = {VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET}; - write_set.pNext = nullptr; - write_set.dstSet = DescriptorSet; - write_set.dstBinding = Binding; - write_set.dstArrayElement = 0; - write_set.descriptorCount = Count; - write_set.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; - write_set.pImageInfo = nullptr; - write_set.pBufferInfo = bufferInfos.data(); - write_set.pTexelBufferView = nullptr; - - vkUpdateDescriptorSets(context->logicalDevice, 1, &write_set, 0, nullptr); - }; - - // Do binding setups - { - VkBuffer BufferMaps[4]; - VkDeviceSize Offsets1[4] = {0, 0, 0, 0}; - - // Map inputs/outputs - BufferMaps[0] = keys->buffer; - BufferMaps[1] = scratch->buffer; - if (bHasPayload) { - BufferMaps[2] = values->buffer; - BufferMaps[3] = scratch->buffer; - Offsets1[2] = 0; - Offsets1[3] = valuesOffset; - } - BindUAVBuffer(BufferMaps, Offsets1, context->sortStages.m_SortDescriptorSetInputOutput[0], 0, - (bHasPayload) ? 4 : 2); - - BufferMaps[0] = scratch->buffer; - BufferMaps[1] = keys->buffer; - if (bHasPayload) { - BufferMaps[2] = scratch->buffer; - BufferMaps[3] = values->buffer; - Offsets1[2] = valuesOffset; - Offsets1[3] = 0; - } - BindUAVBuffer(BufferMaps, Offsets1, context->sortStages.m_SortDescriptorSetInputOutput[1], 0, - (bHasPayload) ? 4 : 2); - - // Map scan sets (reduced, scratch) - VkDeviceSize Offsets2[4] = {reducedScratchOffset, reducedScratchOffset, 0, 0}; - BufferMaps[0] = BufferMaps[1] = scratch->buffer; - BufferMaps[2] = scratch->buffer; - BindUAVBuffer(BufferMaps, Offsets2, context->sortStages.m_SortDescriptorSetScanSets[0], 0, 3); - - BufferMaps[0] = BufferMaps[1] = scratch->buffer; - BufferMaps[2] = scratch->buffer; - VkDeviceSize Offsets3[4] = {scratchOffset, scratchOffset, reducedScratchOffset, 0}; - BindUAVBuffer(BufferMaps, Offsets3, context->sortStages.m_SortDescriptorSetScanSets[1], 0, 3); - - // Map Scratch areas (fixed) - BufferMaps[0] = scratch->buffer; - BufferMaps[1] = scratch->buffer; - VkDeviceSize Offsets4[4] = {scratchOffset, reducedScratchOffset, 0, 0}; - BindUAVBuffer(BufferMaps, Offsets4, context->sortStages.m_SortDescriptorSetScratch, 0, 2); - } - - // Transition barrier - auto BufferTransition = [](VkBuffer buffer, VkAccessFlags before, VkAccessFlags after, VkDeviceSize offset, - VkDeviceSize size) { - VkBufferMemoryBarrier bufferBarrier = {}; - bufferBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - bufferBarrier.srcAccessMask = before; - bufferBarrier.dstAccessMask = after; - bufferBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - bufferBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - bufferBarrier.buffer = buffer; - bufferBarrier.size = VK_WHOLE_SIZE; - // bufferBarrier.offset = offset; - // bufferBarrier.size = size; - - return bufferBarrier; - }; + uint64_t scratchAddress = scratch->getDeviceAddress(); + constantBufferData.SrcBuffer = (uint64_t *)keys->getDeviceAddress(); + constantBufferData.DstBuffer = (uint64_t *)scratchAddress; + constantBufferData.SrcPayload = values ? (uint64_t *)values->getDeviceAddress() : nullptr; + constantBufferData.DstPayload = values ? (uint64_t *)(scratchAddress + keysSize) : nullptr; + constantBufferData.SumTable = (uint32_t *)(scratchAddress + scratchOffset); + constantBufferData.ReduceTable = (uint32_t *)(scratchAddress + reducedScratchOffset); VkCommandBuffer commandList = context->beginComputeCommands(); - - // Bind the scratch descriptor sets - vkCmdBindDescriptorSets(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.layout, 2, 1, - &context->sortStages.m_SortDescriptorSetScratch, 0, nullptr); - - // Push the data into the constant buffer and bind - vkCmdPushConstants(commandList, context->sortStages.layout, VK_SHADER_STAGE_ALL, 0, sizeof(ParallelSortCB), - &constantBufferData); - - // Perform Radix Sort (currently only support 64-bit key/payload sorting - uint32_t inputSet = 0; - VkBufferMemoryBarrier Barriers[3]; - for (uint64_t Shift = 0; Shift < 64u; Shift += PARALLELSORT_SORT_BITS_PER_PASS) { - // Update the bit shift - vkCmdPushConstants(commandList, context->sortStages.layout, VK_SHADER_STAGE_ALL, sizeof(ParallelSortCB) - 4, 4, - &Shift); - // vkCmdPushConstants(commandList, m_SortPipelineLayout, VK_SHADER_STAGE_ALL, 0, 4, &Shift); - - // Bind input/output for this pass - vkCmdBindDescriptorSets(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.layout, 0, 1, - &context->sortStages.m_SortDescriptorSetInputOutput[inputSet], 0, nullptr); - - // Sort Count - { - vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.Count.pipeline); - vkCmdDispatch(commandList, NumThreadgroupsToRun, 1, 1); - } - - // UAV barrier on the sum table - Barriers[0] = - BufferTransition(scratch->buffer, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, scratchOffset, scratchBufferSize); - vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, - nullptr, 1, Barriers, 0, nullptr); - - // Sort Reduce - { - vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.CountReduce.pipeline); - vkCmdDispatch(commandList, NumReducedThreadgroupsToRun, 1, 1); - } - // UAV barrier on the reduced sum table - Barriers[0] = BufferTransition(scratch->buffer, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, reducedScratchOffset, - reducedScratchBufferSize); - vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, - nullptr, 1, Barriers, 0, nullptr); - - // Sort Scan - { - // First do scan prefix of reduced values - vkCmdBindDescriptorSets(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.layout, 1, 1, - &context->sortStages.m_SortDescriptorSetScanSets[0], 0, nullptr); - vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.Scan.pipeline); - assert(NumReducedThreadgroupsToRun < PARALLELSORT_ELEMENTS_PER_THREAD * PARALLELSORT_THREADGROUP_SIZE && - "Need to account for bigger reduced histogram scan"); - vkCmdDispatch(commandList, 1, 1, 1); - - // UAV barrier on the reduced sum table - Barriers[0] = BufferTransition(scratch->buffer, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, reducedScratchOffset, - reducedScratchBufferSize); - vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, - nullptr, 1, Barriers, 0, nullptr); - - // Next do scan prefix on the histogram with partial sums that we just did - vkCmdBindDescriptorSets(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.layout, 1, 1, - &context->sortStages.m_SortDescriptorSetScanSets[1], 0, nullptr); - vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.ScanAdd.pipeline); - vkCmdDispatch(commandList, NumReducedThreadgroupsToRun, 1, 1); - } - - // UAV barrier on the sum table - Barriers[0] = - BufferTransition(scratch->buffer, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, scratchOffset, scratchBufferSize); - vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, - nullptr, 1, Barriers, 0, nullptr); - - // Sort Scatter - { - vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, - bHasPayload ? context->sortStages.ScatterPayload.pipeline - : context->sortStages.Scatter.pipeline); - vkCmdDispatch(commandList, NumThreadgroupsToRun, 1, 1); - } - - // Finish doing everything and barrier for the next pass - VkBuffer keysBuffer = (inputSet) ? scratch->buffer : keys->buffer; - Barriers[0] = BufferTransition(keysBuffer, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, 0, keysSize); - vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, - nullptr, 1, Barriers, 0, nullptr); - - if (bHasPayload) { - VkBuffer valsBuffer = (inputSet) ? scratch->buffer : values->buffer; - VkDeviceSize offset = (inputSet) ? valuesOffset : 0; - Barriers[0] = BufferTransition(valsBuffer, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT, offset, valuesSize); - vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, - nullptr, 1, Barriers, 0, nullptr); - } - - // Swap read/write sources - inputSet = !inputSet; - } - context->endComputeCommands(commandList); + auto computeBarrier = [&]() { + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; + barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; + vkCmdPipelineBarrier(commandList, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, + 0, 1, &barrier, 0, nullptr, 0, nullptr); + }; + // Order accesses after earlier compute submissions, including reused scratch. + computeBarrier(); + for (uint32_t shift = 0; shift < 64; shift += PARALLELSORT_SORT_BITS_PER_PASS) { + constantBufferData.CShiftBit = shift; + vkCmdPushConstants(commandList, context->sortStages.layout, VK_SHADER_STAGE_ALL, 0, + sizeof(ParallelSortCB), &constantBufferData); + + // Count keys in each radix bin, then reduce the counts. + vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.Count.pipeline); + vkCmdDispatch(commandList, NumThreadgroupsToRun, 1, 1); + computeBarrier(); + vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.CountReduce.pipeline); + vkCmdDispatch(commandList, NumReducedThreadgroupsToRun, 1, 1); + computeBarrier(); + + // Scan the reduced counts, then add those prefixes to the histogram. + vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.Scan.pipeline); + assert(NumReducedThreadgroupsToRun < PARALLELSORT_ELEMENTS_PER_THREAD * PARALLELSORT_THREADGROUP_SIZE); + vkCmdDispatch(commandList, 1, 1, 1); + computeBarrier(); + vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.ScanAdd.pipeline); + vkCmdDispatch(commandList, NumReducedThreadgroupsToRun, 1, 1); + computeBarrier(); + + vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, + bHasPayload ? context->sortStages.ScatterPayload.pipeline : context->sortStages.Scatter.pipeline); + vkCmdDispatch(commandList, NumThreadgroupsToRun, 1, 1); + computeBarrier(); + std::swap(constantBufferData.SrcBuffer, constantBufferData.DstBuffer); + std::swap(constantBufferData.SrcPayload, constantBufferData.DstPayload); + } + if (context->endComputeCommands(commandList) != VK_SUCCESS) return; + if (!_scratch && context->synchronizeCompute() == VK_SUCCESS) + gprtBufferDestroy((GPRTBuffer)scratch); } GPRT_API void diff --git a/gprt/gprt_host.h b/gprt/gprt_host.h index a0c417d..57331c2 100644 --- a/gprt/gprt_host.h +++ b/gprt/gprt_host.h @@ -2410,13 +2410,16 @@ gprtBufferSelect(GPRTContext context, GPRTBufferOf input, bool selectPositiv * internally. If a buffer is given, then if that buffer is undersized, the buffer will be allocated / resized and * returned by reference. Otherwise, the scratch buffer will be used directly without any device side allocations. */ +// Sorts run on the compute queue. Callers synchronize before reading, resizing, +// or destroying buffers used by pending work, and before access from other queues. +// Scratch growth and internal scratch release wait on the existing compute timeline. GPRT_API void gprtBufferSort(GPRTContext context, GPRTBuffer buffer, GPRTBuffer scratch GPRT_IF_CPP(= 0)); /** * @brief Sorts the input buffer using a GPU-parallel radix sorter. * Radix sort requires a temporary "scratch" space * - * @tparam T1 The template type of the given buffer (currently only uint32_t is supported) + * @tparam T1 The key type (uint64_t) * @tparam T2 The template type of the scratch buffer (uint8_t is assumed) * * @param context The GPRT context @@ -2425,7 +2428,7 @@ GPRT_API void gprtBufferSort(GPRTContext context, GPRTBuffer buffer, GPRTBuffer * internally. If a buffer is given, then if that buffer is undersized, the buffer will be allocated / resized and * returned by reference. Otherwise, the scratch buffer will be used directly without any device side allocations. */ -template +template void gprtBufferSort(GPRTContext context, GPRTBufferOf buffer, GPRTBufferOf scratch GPRT_IF_CPP(= 0)) { gprtBufferSort(context, (GPRTBuffer) buffer, (GPRTBuffer) scratch); @@ -2442,6 +2445,7 @@ gprtBufferSort(GPRTContext context, GPRTBufferOf buffer, GPRTBufferOf sc * internally. If a buffer is given, then if that buffer is undersized, the buffer will be allocated / resized and * returned by reference. Otherwise, the scratch buffer will be used directly without any device side allocations. */ +// Uses the same compute synchronization contract as gprtBufferSort. GPRT_API void gprtBufferSortPayload(GPRTContext context, GPRTBuffer keys, GPRTBuffer values, GPRTBuffer scratch GPRT_IF_CPP(= 0)); @@ -2449,18 +2453,18 @@ GPRT_API void gprtBufferSortPayload(GPRTContext context, GPRTBuffer keys, GPRTBu * @brief Sorts the input buffer using a GPU-parallel radix sorter. * Radix sort requires a temporary "scratch" space * - * @tparam T1 The template type of the given buffer (currently only uint32_t is supported) - * @tparam T2 The template type of the given buffer (currently only uint32_t is supported) + * @tparam T1 The key type (uint64_t) + * @tparam T2 The payload type (64 bits) * @tparam T3 The template type of the scratch buffer (uint8_t is assumed) * * @param context The GPRT context - * @param keys A buffer of 32-bit unsigned integer keys - * @param values A buffer of 32-bit values + * @param keys A buffer of 64-bit unsigned integer keys + * @param values A buffer of 64-bit values * @param scratch A scratch buffer to facilitate the sort. If null, scratch memory will be allocated and released * internally. If a buffer is given, then if that buffer is undersized, the buffer will be allocated / resized and * returned by reference. Otherwise, the scratch buffer will be used directly without any device side allocations. */ -template +template void gprtBufferSortPayload(GPRTContext context, GPRTBufferOf keys, GPRTBufferOf values, GPRTBufferOf scratch GPRT_IF_CPP(= 0)) { diff --git a/gprt/gprt_sort.h b/gprt/gprt_sort.h index 4ae87cc..5d62007 100644 --- a/gprt/gprt_sort.h +++ b/gprt/gprt_sort.h @@ -42,9 +42,18 @@ struct ParallelSortCB uint32_t NumScanValues; uint32_t CShiftBit; + uint32_t padding; + + uint64_t *SrcBuffer; + uint64_t *DstBuffer; + uint64_t *SrcPayload; + uint64_t *DstPayload; + uint32_t *SumTable; + uint32_t *ReduceTable; }; #ifndef __SLANG_COMPILER__ + static_assert(sizeof(ParallelSortCB) == 80, "Sort push constants must match the shader layout"); void ParallelSort_CalculateScratchResourceSize(uint32_t MaxNumKeys, uint64_t& ScratchBufferSize, uint64_t& ReduceScratchBufferSize) { uint64_t BlockSize = PARALLELSORT_ELEMENTS_PER_THREAD * PARALLELSORT_THREADGROUP_SIZE; diff --git a/gprt/gprt_sort.slang b/gprt/gprt_sort.slang index b3f01b8..3683efb 100644 --- a/gprt/gprt_sort.slang +++ b/gprt/gprt_sort.slang @@ -22,22 +22,8 @@ #define HLSL #include "gprt_sort.h" -// [[vk::binding(/*binding*/, /*set*/)]] - -[[vk::push_constant]] ParallelSortCB rootConstData; // Store the shift bit directly in the root signature - -[[vk::binding(0, 0)]] RWStructuredBuffer SrcBuffer ; // The unsorted keys or scan data -[[vk::binding(2, 0)]] RWStructuredBuffer SrcPayload ; // The payload data - -[[vk::binding(0, 2)]] RWStructuredBuffer SumTable ; // The sum table we will write sums to -[[vk::binding(1, 2)]] RWStructuredBuffer ReduceTable ; // The reduced sum table we will write sums to - -[[vk::binding(1, 0)]] RWStructuredBuffer DstBuffer ; // The sorted keys or prefixed data -[[vk::binding(3, 0)]] RWStructuredBuffer DstPayload ; // the sorted payload data - -[[vk::binding(0, 1)]] RWStructuredBuffer ScanSrc ; // Source for Scan Data -[[vk::binding(1, 1)]] RWStructuredBuffer ScanDst ; // Destination for Scan Data -[[vk::binding(2, 1)]] RWStructuredBuffer ScanScratch ; // Scratch data for Scan +// Device addresses are recorded with each dispatch's sort parameters. +[[vk::push_constant]] ParallelSortCB rootConstData; groupshared uint64_t gs_PARALLELSORT_LDSKeys[PARALLELSORT_THREADGROUP_SIZE]; @@ -109,6 +95,8 @@ uint32_t ParallelSort_BlockScanPrefix(uint32_t localSum, uint32_t localID) [numthreads(PARALLELSORT_THREADGROUP_SIZE, 1, 1)] [shader("compute")] void Count(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) { + uint64_t *SrcBuffer = rootConstData.SrcBuffer; + uint32_t *SumTable = rootConstData.SumTable; // Call the uint32_t version of the count part of the algorithm // ParallelSort_Count_uint(localID, groupID, CBuffer, rootConstData.CShiftBit, SrcBuffer, SumTable); @@ -179,6 +167,8 @@ void Count(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) [numthreads(PARALLELSORT_THREADGROUP_SIZE, 1, 1)] [shader("compute")] void CountReduce(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) { + uint32_t *SumTable = rootConstData.SumTable; + uint32_t *ReduceTable = rootConstData.ReduceTable; // Call the reduce part of the algorithm // ParallelSort_ReduceCount(localID, groupID, CBuffer, SumTable, ReduceTable); @@ -212,6 +202,9 @@ void CountReduce(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_Grou [numthreads(PARALLELSORT_THREADGROUP_SIZE, 1, 1)] [shader("compute")] void Scan(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) { + uint32_t *ScanSrc = rootConstData.ReduceTable; + uint32_t *ScanDst = rootConstData.ReduceTable; + uint32_t *ScanScratch = rootConstData.ReduceTable; uint32_t BaseIndex = PARALLELSORT_ELEMENTS_PER_THREAD * PARALLELSORT_THREADGROUP_SIZE * groupID; // ParallelSort_ScanPrefix(rootConstData.NumScanValues, localID, groupID, 0, BaseIndex, false, // CBuffer, ScanSrc, ScanDst, ScanScratch); @@ -279,6 +272,9 @@ void Scan(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) [numthreads(PARALLELSORT_THREADGROUP_SIZE, 1, 1)] [shader("compute")] void ScanAdd(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) { + uint32_t *ScanSrc = rootConstData.SumTable; + uint32_t *ScanDst = rootConstData.SumTable; + uint32_t *ScanScratch = rootConstData.ReduceTable; // When doing adds, we need to access data differently because reduce // has a more specialized access pattern to match optimized count // Access needs to be done similarly to reduce @@ -353,6 +349,9 @@ void ScanAdd(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) [numthreads(PARALLELSORT_THREADGROUP_SIZE, 1, 1)] [shader("compute")] void Scatter(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) { + uint64_t *SrcBuffer = rootConstData.SrcBuffer; + uint64_t *DstBuffer = rootConstData.DstBuffer; + uint32_t *SumTable = rootConstData.SumTable; uint32_t ShiftBit = rootConstData.CShiftBit; // Load the sort bin threadgroup offsets into LDS for faster referencing @@ -491,6 +490,11 @@ void Scatter(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) [numthreads(PARALLELSORT_THREADGROUP_SIZE, 1, 1)] [shader("compute")] void ScatterPayload(uint32_t localID : SV_GroupThreadID, uint32_t groupID : SV_GroupID) { + uint64_t *SrcBuffer = rootConstData.SrcBuffer; + uint64_t *DstBuffer = rootConstData.DstBuffer; + uint64_t *SrcPayload = rootConstData.SrcPayload; + uint64_t *DstPayload = rootConstData.DstPayload; + uint32_t *SumTable = rootConstData.SumTable; uint32_t ShiftBit = rootConstData.CShiftBit; // Load the sort bin threadgroup offsets into LDS for faster referencing diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3d65253..000438f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -37,5 +37,7 @@ target_link_libraries(t13_terminal PRIVATE gprt::gprt) embed_devicecode(OUTPUT_TARGET t15_deviceCode SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/t15-accelAddress.slang) add_executable(t15_accelAddress t15-accelAddress.cpp) target_link_libraries(t15_accelAddress PRIVATE gprt::gprt t15_deviceCode) +add_executable(t10_sortLifetime t10-sortLifetime.cpp) +target_link_libraries(t10_sortLifetime PRIVATE gprt::gprt) # add_subdirectory(t03-bufferScans) # add_subdirectory(t04-swBVH) diff --git a/tests/t10-sortLifetime.cpp b/tests/t10-sortLifetime.cpp new file mode 100644 index 0000000..c986ec7 --- /dev/null +++ b/tests/t10-sortLifetime.cpp @@ -0,0 +1,66 @@ +#include +#include +#include +#include +#include + +int main() try { + auto context = gprtContextCreate(); + auto scratch = gprtDeviceBufferCreate(context, 2097152); + auto scratchAddress = gprtBufferGetDevicePointer(scratch); + auto scratchSize = gprtBufferGetSize(scratch); + for (size_t count : {1, 511, 1024, 2049, 524289}) { + std::vector keys(count), values(count); + for (size_t i = 0; i < count; ++i) { + keys[i] = (uint64_t((i / 3) * 73 % count) << 33) ^ 0xfedcba9876543210ull; + values[i] = keys[i] ^ 0xabcdef; + } + keys[0] = ~uint64_t(0); + values[0] = keys[0] ^ 0xabcdef; + auto keyBuffer = gprtDeviceBufferCreate(context, count, keys.data()); + auto valueBuffer = gprtDeviceBufferCreate(context, count, values.data()); + auto otherKeys = gprtDeviceBufferCreate(context, count, keys.data()); + std::sort(keys.begin(), keys.end()); + for (int i = 0; i < (count == 511 ? 1 : count > 2049 ? 2 : 80); ++i) { + if (count == 511) { + gprtBufferSortPayload(context, keyBuffer, valueBuffer); + gprtBufferSort(context, otherKeys); + } else { + gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch); + gprtBufferSort(context, otherKeys, scratch); + } + } + if (gprtBufferGetSize(scratch) != scratchSize || gprtBufferGetDevicePointer(scratch) != scratchAddress) + throw std::runtime_error("Sort reallocated sufficient caller-owned scratch"); + gprtComputeSynchronize(context); + gprtBufferMap(keyBuffer); + gprtBufferMap(valueBuffer); + gprtBufferMap(otherKeys); + for (size_t i = 0; i < count; ++i) + if (gprtBufferGetHostPointer(keyBuffer)[i] != keys[i] || + gprtBufferGetHostPointer(valueBuffer)[i] != (keys[i] ^ 0xabcdef) || + gprtBufferGetHostPointer(otherKeys)[i] != keys[i]) + throw std::runtime_error("Sort key or payload mismatch"); + gprtBufferUnmap(keyBuffer); + gprtBufferUnmap(valueBuffer); + gprtBufferUnmap(otherKeys); + gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch); + gprtComputeSynchronize(context); + gprtBufferResize(context, scratch, 1, false); + gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch); + gprtBufferSortPayload(context, keyBuffer, valueBuffer); + gprtBufferSort(context, otherKeys); + gprtComputeSynchronize(context); + gprtBufferResize(context, scratch, 2097152, false); + scratchAddress = gprtBufferGetDevicePointer(scratch); + scratchSize = gprtBufferGetSize(scratch); + gprtBufferDestroy(valueBuffer); + gprtBufferDestroy(keyBuffer); + gprtBufferDestroy(otherKeys); + } + gprtBufferDestroy(scratch); + gprtContextDestroy(context); +} catch (const std::exception &error) { + std::cerr << error.what() << '\n'; + return 1; +}