Pass sort buffer addresses through push constants

gprt/gprt_host.h

@@ -2410,13 +2410,16 @@ 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));
 
 /**
  * @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 <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);
@@ -2442,25 +2445,26 @@ 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));
 
 /**
  * @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)) {

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;

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

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)

tests/t10-sortLifetime.cpp

@@ -0,0 +1,66 @@
+#include <gprt.h>
+#include <algorithm>
+#include <iostream>
+#include <stdexcept>
+#include <vector>
+
+int main() try {
+  auto context = gprtContextCreate();
+  auto scratch = gprtDeviceBufferCreate<uint64_t>(context, 2097152);
+  auto scratchAddress = gprtBufferGetDevicePointer(scratch);
+  auto scratchSize = gprtBufferGetSize(scratch);
+  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] = (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());
+    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;
+}