From 1d9b20c4b8eac4b5eba2ea69514b123999086575 Mon Sep 17 00:00:00 2001
From: Nate Morrical <natemorrical@gmail.com>
Date: Fri, 9 Oct 2026 09:29:52 -0700
Subject: [PATCH 1/6] Cache sort bindings and wait before rebinding or
 releasing buffers

---
 gprt/gprt.cpp              | 24 +++++++++++++++++++++++-
 tests/CMakeLists.txt       |  2 ++
 tests/t10-sortLifetime.cpp | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 tests/t10-sortLifetime.cpp

diff --git a/gprt/gprt.cpp b/gprt/gprt.cpp
index 52bb751..193eec9 100644
--- a/gprt/gprt.cpp
+++ b/gprt/gprt.cpp
@@ -653,6 +653,8 @@ struct Context {
 
   // TODO, we can probably refactor this...
   struct SortStages {
+    VkBuffer boundKeys = VK_NULL_HANDLE, boundValues = VK_NULL_HANDLE, boundScratch = VK_NULL_HANDLE;
+    VkDeviceSize boundSize = 0;
     Stage Count;
     Stage CountReduce;
     Stage Scan;
@@ -939,6 +941,7 @@ struct Buffer {
 
   /*! Calls vkDestroy on the buffer, and frees underlying memory */
   void destroy() {
+    invalidateSortBindings();
     // Free sampler slot for use by subsequently made buffers
     context->buffers[virtualAddress] = nullptr;
 
@@ -957,6 +960,15 @@ struct Buffer {
     }
   }
 
+  void invalidateSortBindings() {
+    auto &sort = context->sortStages;
+    if (buffer && (buffer == sort.boundKeys || buffer == sort.boundValues || buffer == sort.boundScratch)) {
+      VK_CHECK_RESULT(context->synchronizeCompute());
+      sort.boundKeys = sort.boundValues = sort.boundScratch = VK_NULL_HANDLE;
+      sort.boundSize = 0;
+    }
+  }
+
   /* Sets all bytes to 0 */
   void clear() {
     if (hostVisible) {
@@ -975,6 +987,8 @@ struct Buffer {
     if (size == bytes)
       return;
 
+    invalidateSortBindings();
+
     uint32_t queueFamilies[] = {context->queueFamilyIndices.graphics, context->queueFamilyIndices.compute,
                                context->queueFamilyIndices.transfer};
     std::sort(queueFamilies, queueFamilies + 3);
@@ -9030,7 +9044,11 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
   };
 
   // Do binding setups
-  {
+  auto &sort = context->sortStages;
+  VkBuffer valueBuffer = values ? values->buffer : VK_NULL_HANDLE;
+  if (sort.boundKeys != keys->buffer || sort.boundValues != valueBuffer ||
+      sort.boundScratch != scratch->buffer || sort.boundSize != keys->size) {
+    VK_CHECK_RESULT(context->synchronizeCompute());
     VkBuffer BufferMaps[4];
     VkDeviceSize Offsets1[4] = {0, 0, 0, 0};
 
@@ -9073,6 +9091,10 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
     BufferMaps[1] = scratch->buffer;
     VkDeviceSize Offsets4[4] = {scratchOffset, reducedScratchOffset, 0, 0};
     BindUAVBuffer(BufferMaps, Offsets4, context->sortStages.m_SortDescriptorSetScratch, 0, 2);
+    sort.boundKeys = keys->buffer;
+    sort.boundValues = valueBuffer;
+    sort.boundScratch = scratch->buffer;
+    sort.boundSize = keys->size;
   }
 
   // Transition barrier
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..43ece3e
--- /dev/null
+++ b/tests/t10-sortLifetime.cpp
@@ -0,0 +1,34 @@
+#include <gprt.h>
+#include <stdexcept>
+#include <vector>
+
+int main() {
+  auto context = gprtContextCreate();
+  auto scratch = gprtDeviceBufferCreate<uint64_t>(context);
+  for (size_t count : {1024, 2048, 512}) {
+    std::vector<uint64_t> keys(count), values(count);
+    for (size_t i = 0; i < count; ++i) {
+      keys[i] = count - i - 1;
+      values[i] = keys[i] ^ 0xabcdef;
+    }
+    auto keyBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, keys.data());
+    auto valueBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, values.data());
+    for (int i = 0; i < 8; ++i) gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    gprtComputeSynchronize(context);
+    gprtBufferMap(keyBuffer);
+    gprtBufferMap(valueBuffer);
+    for (size_t i = 0; i < count; ++i)
+      if (gprtBufferGetHostPointer(keyBuffer)[i] != i ||
+          gprtBufferGetHostPointer(valueBuffer)[i] != (i ^ 0xabcdef))
+        throw std::runtime_error("Sort key or payload mismatch");
+    gprtBufferUnmap(keyBuffer);
+    gprtBufferUnmap(valueBuffer);
+    gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    gprtBufferResize(context, scratch, 1, false);
+    gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    gprtBufferDestroy(valueBuffer);
+    gprtBufferDestroy(keyBuffer);
+  }
+  gprtBufferDestroy(scratch);
+  gprtContextDestroy(context);
+}

From d7db29012cec65ed50315b9f94176cf54b9061d6 Mon Sep 17 00:00:00 2001
From: Nate Morrical <natemorrical@gmail.com>
Date: Fri, 9 Oct 2026 13:43:59 -0700
Subject: [PATCH 2/6] Preserve scratch capacity and retire sort descriptors per
 command buffer

---
 gprt/gprt.cpp              | 99 ++++++++++++++++++--------------------
 gprt/gprt_host.h           | 12 ++---
 tests/t10-sortLifetime.cpp | 18 +++++--
 3 files changed, 67 insertions(+), 62 deletions(-)

diff --git a/gprt/gprt.cpp b/gprt/gprt.cpp
index 193eec9..dd8d9af 100644
--- a/gprt/gprt.cpp
+++ b/gprt/gprt.cpp
@@ -653,8 +653,6 @@ struct Context {
 
   // TODO, we can probably refactor this...
   struct SortStages {
-    VkBuffer boundKeys = VK_NULL_HANDLE, boundValues = VK_NULL_HANDLE, boundScratch = VK_NULL_HANDLE;
-    VkDeviceSize boundSize = 0;
     Stage Count;
     Stage CountReduce;
     Stage Scan;
@@ -667,9 +665,11 @@ struct Context {
     VkDescriptorSetLayout m_SortDescriptorSetLayoutScan;
     VkDescriptorSetLayout m_SortDescriptorSetLayoutScratch;
 
-    VkDescriptorSet m_SortDescriptorSetInputOutput[2];
-    VkDescriptorSet m_SortDescriptorSetScanSets[2];
-    VkDescriptorSet m_SortDescriptorSetScratch;
+    struct Bindings {
+      VkDescriptorSet m_SortDescriptorSetInputOutput[2];
+      VkDescriptorSet m_SortDescriptorSetScanSets[2];
+      VkDescriptorSet m_SortDescriptorSetScratch;
+    } bindings[64]; // One set per compute command buffer; reuse follows its completion.
 
     VkPipelineLayout layout;
 
@@ -853,6 +853,7 @@ struct Buffer {
 
   VkDeviceSize size = 0;
   VkDeviceSize alignment = 16;
+  uint64_t lastSortSubmission = 0;
   void *mapped = nullptr;
 
   VkResult map(VkDeviceSize mapSize = VK_WHOLE_SIZE, VkDeviceSize offset = 0) {
@@ -941,7 +942,7 @@ struct Buffer {
 
   /*! Calls vkDestroy on the buffer, and frees underlying memory */
   void destroy() {
-    invalidateSortBindings();
+    waitForSort();
     // Free sampler slot for use by subsequently made buffers
     context->buffers[virtualAddress] = nullptr;
 
@@ -960,13 +961,14 @@ struct Buffer {
     }
   }
 
-  void invalidateSortBindings() {
-    auto &sort = context->sortStages;
-    if (buffer && (buffer == sort.boundKeys || buffer == sort.boundValues || buffer == sort.boundScratch)) {
-      VK_CHECK_RESULT(context->synchronizeCompute());
-      sort.boundKeys = sort.boundValues = sort.boundScratch = VK_NULL_HANDLE;
-      sort.boundSize = 0;
-    }
+  void waitForSort() {
+    if (!lastSortSubmission) return;
+    VkSemaphoreWaitInfo wait{VK_STRUCTURE_TYPE_SEMAPHORE_WAIT_INFO};
+    wait.semaphoreCount = 1;
+    wait.pSemaphores = &context->CETimelineSemaphore;
+    wait.pValues = &lastSortSubmission;
+    VK_CHECK_RESULT(vkWaitSemaphores(context->logicalDevice, &wait, requestedFeatures.syncTDR));
+    lastSortSubmission = 0;
   }
 
   /* Sets all bytes to 0 */
@@ -987,7 +989,7 @@ struct Buffer {
     if (size == bytes)
       return;
 
-    invalidateSortBindings();
+    waitForSort();
 
     uint32_t queueFamilies[] = {context->queueFamilyIndices.graphics, context->queueFamilyIndices.compute,
                                context->queueFamilyIndices.transfer};
@@ -6251,13 +6253,13 @@ Context::setupInternalPrograms() {
 
     VkDescriptorPoolSize poolSize;
     poolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
-    poolSize.descriptorCount = 16;
+    poolSize.descriptorCount = 16 * 64;
 
     VkDescriptorPoolCreateInfo descriptorPoolInfo{};
     descriptorPoolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
     descriptorPoolInfo.poolSizeCount = 1;
     descriptorPoolInfo.pPoolSizes = &poolSize;
-    descriptorPoolInfo.maxSets = 5;
+    descriptorPoolInfo.maxSets = 5 * 64;
     descriptorPoolInfo.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
     VK_CHECK_RESULT(vkCreateDescriptorPool(logicalDevice, &descriptorPoolInfo, nullptr, &sortStages.pool));
 
@@ -6311,28 +6313,26 @@ Context::setupInternalPrograms() {
     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]);
+    for (auto &binding : sortStages.bindings)
+      for (auto &set : binding.m_SortDescriptorSetInputOutput)
+        AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutInputOutputs, &set);
 
     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]);
+    for (auto &binding : sortStages.bindings)
+      for (auto &set : binding.m_SortDescriptorSetScanSets)
+        AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScan, &set);
 
     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);
+    for (auto &binding : sortStages.bindings)
+      AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScratch, &binding.m_SortDescriptorSetScratch);
 
     // Create constant range representing our static constant
     VkPushConstantRange constant_range;
@@ -6604,9 +6604,7 @@ 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);
+    // Destroying the pool below releases all command-buffer descriptor sets.
 
     vkDestroyDescriptorSetLayout(logicalDevice, sortStages.m_SortDescriptorSetLayoutScratch, nullptr);
     vkDestroyDescriptorSetLayout(logicalDevice, sortStages.m_SortDescriptorSetLayoutScan, nullptr);
@@ -8977,7 +8975,7 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
   Context *context = (Context *) _context;
   Buffer *keys = (Buffer *) _keys;
   Buffer *values = (Buffer *) _values;
-  Buffer *scratch = (Buffer *) _scratch;
+  Buffer *scratch = _scratch ? (Buffer *)_scratch : (Buffer *)gprtDeviceBufferCreate<uint64_t>(_context);
 
   bool bHasPayload = false;
   if (values) {
@@ -9006,8 +9004,8 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
   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);
+  size_t requiredScratch = keysSize + valuesSize + scratchBufferSize + reducedScratchBufferSize;
+  if (scratch->size < requiredScratch) scratch->resize(requiredScratch, false);
   // All offsets must be a multiple of device limit VkPhysicalDeviceLimits::minStorageBufferOffseteAlignment
   size_t valuesOffset = keysSize;
   size_t scratchOffset = keysSize + valuesSize;
@@ -9043,12 +9041,10 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
     vkUpdateDescriptorSets(context->logicalDevice, 1, &write_set, 0, nullptr);
   };
 
-  // Do binding setups
-  auto &sort = context->sortStages;
-  VkBuffer valueBuffer = values ? values->buffer : VK_NULL_HANDLE;
-  if (sort.boundKeys != keys->buffer || sort.boundValues != valueBuffer ||
-      sort.boundScratch != scratch->buffer || sort.boundSize != keys->size) {
-    VK_CHECK_RESULT(context->synchronizeCompute());
+  // beginComputeCommands waits only when this command-buffer slot is still in use.
+  VkCommandBuffer commandList = context->beginComputeCommands();
+  auto &bindings = context->sortStages.bindings[context->CETimelineCounter % 64];
+  {
     VkBuffer BufferMaps[4];
     VkDeviceSize Offsets1[4] = {0, 0, 0, 0};
 
@@ -9061,7 +9057,7 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
       Offsets1[2] = 0;
       Offsets1[3] = valuesOffset;
     }
-    BindUAVBuffer(BufferMaps, Offsets1, context->sortStages.m_SortDescriptorSetInputOutput[0], 0,
+    BindUAVBuffer(BufferMaps, Offsets1, bindings.m_SortDescriptorSetInputOutput[0], 0,
                   (bHasPayload) ? 4 : 2);
 
     BufferMaps[0] = scratch->buffer;
@@ -9072,29 +9068,25 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
       Offsets1[2] = valuesOffset;
       Offsets1[3] = 0;
     }
-    BindUAVBuffer(BufferMaps, Offsets1, context->sortStages.m_SortDescriptorSetInputOutput[1], 0,
+    BindUAVBuffer(BufferMaps, Offsets1, bindings.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);
+    BindUAVBuffer(BufferMaps, Offsets2, bindings.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);
+    BindUAVBuffer(BufferMaps, Offsets3, bindings.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);
-    sort.boundKeys = keys->buffer;
-    sort.boundValues = valueBuffer;
-    sort.boundScratch = scratch->buffer;
-    sort.boundSize = keys->size;
+    BindUAVBuffer(BufferMaps, Offsets4, bindings.m_SortDescriptorSetScratch, 0, 2);
   }
 
   // Transition barrier
@@ -9114,11 +9106,9 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
     return bufferBarrier;
   };
 
-  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);
+                          &bindings.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),
@@ -9135,7 +9125,7 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
 
     // 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);
+                            &bindings.m_SortDescriptorSetInputOutput[inputSet], 0, nullptr);
 
     // Sort Count
     {
@@ -9166,7 +9156,7 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
     {
       // 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);
+                              &bindings.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");
@@ -9181,7 +9171,7 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
 
       // 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);
+                              &bindings.m_SortDescriptorSetScanSets[1], 0, nullptr);
       vkCmdBindPipeline(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.ScanAdd.pipeline);
       vkCmdDispatch(commandList, NumReducedThreadgroupsToRun, 1, 1);
     }
@@ -9221,6 +9211,9 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
     inputSet = !inputSet;
   }
   context->endComputeCommands(commandList);
+  keys->lastSortSubmission = scratch->lastSortSubmission = context->CETimelineCounter;
+  if (values) values->lastSortSubmission = context->CETimelineCounter;
+  if (!_scratch) gprtBufferDestroy((GPRTBuffer)scratch);
 }
 
 GPRT_API void
diff --git a/gprt/gprt_host.h b/gprt/gprt_host.h
index a0c417d..d3c60e5 100644
--- a/gprt/gprt_host.h
+++ b/gprt/gprt_host.h
@@ -2425,7 +2425,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 <typename T1, typename T2>
+template <typename T1, typename T2 = uint8_t>
 void
 gprtBufferSort(GPRTContext context, GPRTBufferOf<T1> buffer, GPRTBufferOf<T2> scratch GPRT_IF_CPP(= 0)) {
   gprtBufferSort(context, (GPRTBuffer) buffer, (GPRTBuffer) scratch);
@@ -2449,18 +2449,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 <typename T1, typename T2, typename T3>
+template <typename T1, typename T2, typename T3 = uint8_t>
 void
 gprtBufferSortPayload(GPRTContext context, GPRTBufferOf<T1> keys, GPRTBufferOf<T2> values,
                       GPRTBufferOf<T3> scratch GPRT_IF_CPP(= 0)) {
diff --git a/tests/t10-sortLifetime.cpp b/tests/t10-sortLifetime.cpp
index 43ece3e..b56942d 100644
--- a/tests/t10-sortLifetime.cpp
+++ b/tests/t10-sortLifetime.cpp
@@ -1,10 +1,13 @@
 #include <gprt.h>
+#include <iostream>
 #include <stdexcept>
 #include <vector>
 
-int main() {
+int main() try {
   auto context = gprtContextCreate();
-  auto scratch = gprtDeviceBufferCreate<uint64_t>(context);
+  auto scratch = gprtDeviceBufferCreate<uint64_t>(context, 131072);
+  auto scratchAddress = gprtBufferGetDevicePointer(scratch);
+  auto scratchSize = gprtBufferGetSize(scratch);
   for (size_t count : {1024, 2048, 512}) {
     std::vector<uint64_t> keys(count), values(count);
     for (size_t i = 0; i < count; ++i) {
@@ -13,7 +16,9 @@ int main() {
     }
     auto keyBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, keys.data());
     auto valueBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, values.data());
-    for (int i = 0; i < 8; ++i) gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    for (int i = 0; i < 80; ++i) gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    if (gprtBufferGetSize(scratch) != scratchSize || gprtBufferGetDevicePointer(scratch) != scratchAddress)
+      throw std::runtime_error("Sort reallocated sufficient caller-owned scratch");
     gprtComputeSynchronize(context);
     gprtBufferMap(keyBuffer);
     gprtBufferMap(valueBuffer);
@@ -26,9 +31,16 @@ int main() {
     gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
     gprtBufferResize(context, scratch, 1, false);
     gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    gprtBufferSortPayload(context, keyBuffer, valueBuffer);
+    gprtBufferResize(context, scratch, 131072, false);
+    scratchAddress = gprtBufferGetDevicePointer(scratch);
+    scratchSize = gprtBufferGetSize(scratch);
     gprtBufferDestroy(valueBuffer);
     gprtBufferDestroy(keyBuffer);
   }
   gprtBufferDestroy(scratch);
   gprtContextDestroy(context);
+} catch (const std::exception &error) {
+  std::cerr << error.what() << '\n';
+  return 1;
 }

From 48b29fc378fb4b3149f100369e7cc496af6360eb Mon Sep 17 00:00:00 2001
From: Nate Morrical <natemorrical@gmail.com>
Date: Fri, 9 Oct 2026 13:54:54 -0700
Subject: [PATCH 3/6] Document the supported sort key width

---
 gprt/gprt_host.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gprt/gprt_host.h b/gprt/gprt_host.h
index d3c60e5..0219d67 100644
--- a/gprt/gprt_host.h
+++ b/gprt/gprt_host.h
@@ -2416,7 +2416,7 @@ GPRT_API void gprtBufferSort(GPRTContext context, GPRTBuffer buffer, GPRTBuffer
  * @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

From 214f8e98a1d216c73440f857e2724cd46ae3d97d Mon Sep 17 00:00:00 2001
From: Nate Morrical <natemorrical@gmail.com>
Date: Fri, 9 Oct 2026 13:57:13 -0700
Subject: [PATCH 4/6] Exercise pending sort rebinding between key-only and
 payload batches

---
 tests/t10-sortLifetime.cpp | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/t10-sortLifetime.cpp b/tests/t10-sortLifetime.cpp
index b56942d..db84440 100644
--- a/tests/t10-sortLifetime.cpp
+++ b/tests/t10-sortLifetime.cpp
@@ -16,18 +16,25 @@ int main() try {
     }
     auto keyBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, keys.data());
     auto valueBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, values.data());
-    for (int i = 0; i < 80; ++i) gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
+    auto otherKeys = gprtDeviceBufferCreate<uint64_t>(context, count, keys.data());
+    for (int i = 0; i < 80; ++i) {
+      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] != i ||
-          gprtBufferGetHostPointer(valueBuffer)[i] != (i ^ 0xabcdef))
+          gprtBufferGetHostPointer(valueBuffer)[i] != (i ^ 0xabcdef) ||
+          gprtBufferGetHostPointer(otherKeys)[i] != i)
         throw std::runtime_error("Sort key or payload mismatch");
     gprtBufferUnmap(keyBuffer);
     gprtBufferUnmap(valueBuffer);
+    gprtBufferUnmap(otherKeys);
     gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
     gprtBufferResize(context, scratch, 1, false);
     gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
@@ -37,6 +44,7 @@ int main() try {
     scratchSize = gprtBufferGetSize(scratch);
     gprtBufferDestroy(valueBuffer);
     gprtBufferDestroy(keyBuffer);
+    gprtBufferDestroy(otherKeys);
   }
   gprtBufferDestroy(scratch);
   gprtContextDestroy(context);

From d2af6594d681d4e63c3912bd2a40247f1358720a Mon Sep 17 00:00:00 2001
From: Nate Morrical <natemorrical@gmail.com>
Date: Fri, 9 Oct 2026 18:32:37 -0700
Subject: [PATCH 5/6] Stop buffer release when a pending sort wait fails

---
 gprt/gprt.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/gprt/gprt.cpp b/gprt/gprt.cpp
index dd8d9af..6738cbc 100644
--- a/gprt/gprt.cpp
+++ b/gprt/gprt.cpp
@@ -32,6 +32,7 @@
 #include <numeric>
 #include <set>
 #include <sstream>
+#include <stdexcept>
 
 #include <regex>
 
@@ -967,7 +968,8 @@ struct Buffer {
     wait.semaphoreCount = 1;
     wait.pSemaphores = &context->CETimelineSemaphore;
     wait.pValues = &lastSortSubmission;
-    VK_CHECK_RESULT(vkWaitSemaphores(context->logicalDevice, &wait, requestedFeatures.syncTDR));
+    if (vkWaitSemaphores(context->logicalDevice, &wait, requestedFeatures.syncTDR) != VK_SUCCESS)
+      throw std::runtime_error("Failed to wait for the buffer's pending sort");
     lastSortSubmission = 0;
   }
 
@@ -9041,7 +9043,7 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
     vkUpdateDescriptorSets(context->logicalDevice, 1, &write_set, 0, nullptr);
   };
 
-  // beginComputeCommands waits only when this command-buffer slot is still in use.
+  // Retire this command-buffer slot before updating its descriptor sets.
   VkCommandBuffer commandList = context->beginComputeCommands();
   auto &bindings = context->sortStages.bindings[context->CETimelineCounter % 64];
   {

From ef4bece60451aa040c4c89e2915a7f32c3923425 Mon Sep 17 00:00:00 2001
From: Nate Morrical <natemorrical@gmail.com>
Date: Sat, 10 Oct 2026 07:52:54 -0700
Subject: [PATCH 6/6] Pass sort buffer addresses through push constants and use
 compute synchronization

---
 gprt/gprt.cpp              | 414 ++++++-------------------------------
 gprt/gprt_host.h           |   4 +
 gprt/gprt_sort.h           |   9 +
 gprt/gprt_sort.slang       |  36 ++--
 tests/t10-sortLifetime.cpp |  32 ++-
 5 files changed, 122 insertions(+), 373 deletions(-)

diff --git a/gprt/gprt.cpp b/gprt/gprt.cpp
index 6738cbc..fb90167 100644
--- a/gprt/gprt.cpp
+++ b/gprt/gprt.cpp
@@ -32,7 +32,6 @@
 #include <numeric>
 #include <set>
 #include <sstream>
-#include <stdexcept>
 
 #include <regex>
 
@@ -661,20 +660,7 @@ struct Context {
     Stage Scatter;
     Stage ScatterPayload;
 
-    // For now, copied over from sample
-    VkDescriptorSetLayout m_SortDescriptorSetLayoutInputOutputs;
-    VkDescriptorSetLayout m_SortDescriptorSetLayoutScan;
-    VkDescriptorSetLayout m_SortDescriptorSetLayoutScratch;
-
-    struct Bindings {
-      VkDescriptorSet m_SortDescriptorSetInputOutput[2];
-      VkDescriptorSet m_SortDescriptorSetScanSets[2];
-      VkDescriptorSet m_SortDescriptorSetScratch;
-    } bindings[64]; // One set per compute command buffer; reuse follows its completion.
-
     VkPipelineLayout layout;
-
-    VkDescriptorPool pool = VK_NULL_HANDLE;
   };
   SortStages sortStages;
 
@@ -854,7 +840,6 @@ struct Buffer {
 
   VkDeviceSize size = 0;
   VkDeviceSize alignment = 16;
-  uint64_t lastSortSubmission = 0;
   void *mapped = nullptr;
 
   VkResult map(VkDeviceSize mapSize = VK_WHOLE_SIZE, VkDeviceSize offset = 0) {
@@ -943,7 +928,6 @@ struct Buffer {
 
   /*! Calls vkDestroy on the buffer, and frees underlying memory */
   void destroy() {
-    waitForSort();
     // Free sampler slot for use by subsequently made buffers
     context->buffers[virtualAddress] = nullptr;
 
@@ -962,17 +946,6 @@ struct Buffer {
     }
   }
 
-  void waitForSort() {
-    if (!lastSortSubmission) return;
-    VkSemaphoreWaitInfo wait{VK_STRUCTURE_TYPE_SEMAPHORE_WAIT_INFO};
-    wait.semaphoreCount = 1;
-    wait.pSemaphores = &context->CETimelineSemaphore;
-    wait.pValues = &lastSortSubmission;
-    if (vkWaitSemaphores(context->logicalDevice, &wait, requestedFeatures.syncTDR) != VK_SUCCESS)
-      throw std::runtime_error("Failed to wait for the buffer's pending sort");
-    lastSortSubmission = 0;
-  }
-
   /* Sets all bytes to 0 */
   void clear() {
     if (hostVisible) {
@@ -991,8 +964,6 @@ struct Buffer {
     if (size == bytes)
       return;
 
-    waitForSort();
-
     uint32_t queueFamilies[] = {context->queueFamilyIndices.graphics, context->queueFamilyIndices.compute,
                                context->queueFamilyIndices.transfer};
     std::sort(queueFamilies, queueFamilies + 3);
@@ -6253,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 * 64;
-
-    VkDescriptorPoolCreateInfo descriptorPoolInfo{};
-    descriptorPoolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
-    descriptorPoolInfo.poolSizeCount = 1;
-    descriptorPoolInfo.pPoolSizes = &poolSize;
-    descriptorPoolInfo.maxSets = 5 * 64;
-    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";
@@ -6272,70 +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);
-    for (auto &binding : sortStages.bindings)
-      for (auto &set : binding.m_SortDescriptorSetInputOutput)
-        AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutInputOutputs, &set);
-
-    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);
-    for (auto &binding : sortStages.bindings)
-      for (auto &set : binding.m_SortDescriptorSetScanSets)
-        AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScan, &set);
-
-    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);
-    for (auto &binding : sortStages.bindings)
-      AllocDescriptor(sortStages.pool, sortStages.m_SortDescriptorSetLayoutScratch, &binding.m_SortDescriptorSetScratch);
-
     // Create constant range representing our static constant
     VkPushConstantRange constant_range;
     constant_range.stageFlags = VK_SHADER_STAGE_ALL;
@@ -6346,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;
 
@@ -6605,14 +6495,6 @@ Context::destroyInternalPrograms() {
 
     vkDestroyPipelineLayout(logicalDevice, sortStages.layout, nullptr);
     sortStages.layout = VK_NULL_HANDLE;
-
-    // Destroying the pool below releases all command-buffer descriptor sets.
-
-    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<std::string> progNames;
@@ -8977,245 +8859,83 @@ bufferSort(GPRTContext _context, GPRTBuffer _keys, GPRTBuffer _values, GPRTBuffe
   Context *context = (Context *) _context;
   Buffer *keys = (Buffer *) _keys;
   Buffer *values = (Buffer *) _values;
-  Buffer *scratch = _scratch ? (Buffer *)_scratch : (Buffer *)gprtDeviceBufferCreate<uint64_t>(_context);
-
-  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<uint64_t>(_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);
-
-  size_t requiredScratch = keysSize + valuesSize + scratchBufferSize + reducedScratchBufferSize;
-  if (scratch->size < requiredScratch) scratch->resize(requiredScratch, 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);
+  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);
 
-  auto BindUAVBuffer = [&](VkBuffer *pBuffer, VkDeviceSize *Offsets, VkDescriptorSet &DescriptorSet,
-                           uint32_t Binding /*=0*/, uint32_t Count /*=1*/) {
-    std::vector<VkDescriptorBufferInfo> 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);
-  };
-
-  // Retire this command-buffer slot before updating its descriptor sets.
   VkCommandBuffer commandList = context->beginComputeCommands();
-  auto &bindings = context->sortStages.bindings[context->CETimelineCounter % 64];
-  {
-    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, bindings.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, bindings.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, bindings.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, bindings.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, bindings.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;
+  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);
   };
-
-  // Bind the scratch descriptor sets
-  vkCmdBindDescriptorSets(commandList, VK_PIPELINE_BIND_POINT_COMPUTE, context->sortStages.layout, 2, 1,
-                          &bindings.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,
-                            &bindings.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,
-                              &bindings.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,
-                              &bindings.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);
-  keys->lastSortSubmission = scratch->lastSortSubmission = context->CETimelineCounter;
-  if (values) values->lastSortSubmission = context->CETimelineCounter;
-  if (!_scratch) gprtBufferDestroy((GPRTBuffer)scratch);
+  // 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 0219d67..57331c2 100644
--- a/gprt/gprt_host.h
+++ b/gprt/gprt_host.h
@@ -2410,6 +2410,9 @@ gprtBufferSelect(GPRTContext context, GPRTBufferOf<T1> 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));
 
 /**
@@ -2442,6 +2445,7 @@ gprtBufferSort(GPRTContext context, GPRTBufferOf<T1> buffer, GPRTBufferOf<T2> 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));
 
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<uint64_t>	SrcBuffer		;					// The unsorted keys or scan data
-[[vk::binding(2, 0)]] RWStructuredBuffer<uint64_t>	SrcPayload		;					// The payload data
-				 
-[[vk::binding(0, 2)]] RWStructuredBuffer<uint32_t>	SumTable		;					// The sum table we will write sums to
-[[vk::binding(1, 2)]] RWStructuredBuffer<uint32_t>	ReduceTable		;					// The reduced sum table we will write sums to
-				 
-[[vk::binding(1, 0)]] RWStructuredBuffer<uint64_t>	DstBuffer		;					// The sorted keys or prefixed data
-[[vk::binding(3, 0)]] RWStructuredBuffer<uint64_t>	DstPayload		;					// the sorted payload data
-				 
-[[vk::binding(0, 1)]] RWStructuredBuffer<uint32_t>	ScanSrc			;					// Source for Scan Data
-[[vk::binding(1, 1)]] RWStructuredBuffer<uint32_t>	ScanDst			;					// Destination for Scan Data
-[[vk::binding(2, 1)]] RWStructuredBuffer<uint32_t>	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/t10-sortLifetime.cpp b/tests/t10-sortLifetime.cpp
index db84440..c986ec7 100644
--- a/tests/t10-sortLifetime.cpp
+++ b/tests/t10-sortLifetime.cpp
@@ -1,25 +1,34 @@
 #include <gprt.h>
+#include <algorithm>
 #include <iostream>
 #include <stdexcept>
 #include <vector>
 
 int main() try {
   auto context = gprtContextCreate();
-  auto scratch = gprtDeviceBufferCreate<uint64_t>(context, 131072);
+  auto scratch = gprtDeviceBufferCreate<uint64_t>(context, 2097152);
   auto scratchAddress = gprtBufferGetDevicePointer(scratch);
   auto scratchSize = gprtBufferGetSize(scratch);
-  for (size_t count : {1024, 2048, 512}) {
+  for (size_t count : {1, 511, 1024, 2049, 524289}) {
     std::vector<uint64_t> keys(count), values(count);
     for (size_t i = 0; i < count; ++i) {
-      keys[i] = count - i - 1;
+      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<uint64_t>(context, count, keys.data());
     auto valueBuffer = gprtDeviceBufferCreate<uint64_t>(context, count, values.data());
     auto otherKeys = gprtDeviceBufferCreate<uint64_t>(context, count, keys.data());
-    for (int i = 0; i < 80; ++i) {
-      gprtBufferSortPayload(context, keyBuffer, valueBuffer, scratch);
-      gprtBufferSort(context, otherKeys, scratch);
+    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");
@@ -28,18 +37,21 @@ int main() try {
     gprtBufferMap(valueBuffer);
     gprtBufferMap(otherKeys);
     for (size_t i = 0; i < count; ++i)
-      if (gprtBufferGetHostPointer(keyBuffer)[i] != i ||
-          gprtBufferGetHostPointer(valueBuffer)[i] != (i ^ 0xabcdef) ||
-          gprtBufferGetHostPointer(otherKeys)[i] != 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);
-    gprtBufferResize(context, scratch, 131072, false);
+    gprtBufferSort(context, otherKeys);
+    gprtComputeSynchronize(context);
+    gprtBufferResize(context, scratch, 2097152, false);
     scratchAddress = gprtBufferGetDevicePointer(scratch);
     scratchSize = gprtBufferGetSize(scratch);
     gprtBufferDestroy(valueBuffer);