Update bitonic_sort_shader.comp.hlsl

Author: CrabExtra

13_BitonicSort/app_resources/bitonic_sort_shader.comp.hlsl

@@ -10,7 +10,7 @@ using BitonicSortConfig = workgroup::bitonic_sort::bitonic_sort_config<ElementsP
 
 NBL_CONSTEXPR uint32_t WorkgroupSize = BitonicSortConfig::WorkgroupSize;
 
-groupshared uint32_t sharedmem[2 * WorkgroupSize];
+groupshared uint32_t sharedmem[BitonicSortConfig::SharedmemDWORDs];
 
 uint32_t3 glsl::gl_WorkGroupSize() { return uint32_t3(uint32_t(BitonicSortConfig::WorkgroupSize), 1, 1); }
 
@@ -19,20 +19,21 @@ struct SharedMemoryAccessor
 	template <typename AccessType, typename IndexType>
 	void set(IndexType idx, AccessType value)
 	{
-		sharedmem[idx] = value;
+		sharedmem[idx * 2] = value.first;
+		sharedmem[idx * 2 + 1] = value.second;
 	}
 
 	template <typename AccessType, typename IndexType>
 	void get(IndexType idx, NBL_REF_ARG(AccessType) value)
 	{
-		value = sharedmem[idx];
+		value.first = sharedmem[idx * 2];
+		value.second = sharedmem[idx * 2 + 1];
 	}
 
 	void workgroupExecutionAndMemoryBarrier()
 	{
 		glsl::barrier();
 	}
-
 };
 
 struct Accessor
@@ -47,13 +48,17 @@ struct Accessor
 	template <typename AccessType, typename IndexType>
 	void get(const IndexType index, NBL_REF_ARG(AccessType) value)
 	{
-		value = vk::RawBufferLoad<AccessType>(address + index * sizeof(AccessType));
+		const uint64_t offset = address + index * sizeof(AccessType);
+		value.first = vk::RawBufferLoad<uint32_t>(offset);
+		value.second = vk::RawBufferLoad<uint32_t>(offset + sizeof(uint32_t));
 	}
 
 	template <typename AccessType, typename IndexType>
 	void set(const IndexType index, const AccessType value)
 	{
-		vk::RawBufferStore<AccessType>(address + index * sizeof(AccessType), value);
+		const uint64_t offset = address + index * sizeof(AccessType);
+		vk::RawBufferStore<uint32_t>(offset, value.first);
+		vk::RawBufferStore<uint32_t>(offset + sizeof(uint32_t), value.second);
 	}
 
 	uint64_t address;
@@ -66,22 +71,6 @@ void main()
 	Accessor accessor = Accessor::create(pushConstants.deviceBufferAddress);
 	SharedMemoryAccessor sharedmemAccessor;
 
-	const uint32_t threadID = glsl::gl_LocalInvocationID().x;
-
-	// Each thread handles 2 ADJACENT elements: lo and hi
-	// Following bitonic sort pattern: thread i handles elements [2*i] and [2*i + 1]
-	const uint32_t loIdx = threadID * 2;
-	const uint32_t hiIdx = threadID * 2 + 1;
-
-	uint32_t loKey, hiKey;
-	accessor.get(loIdx, loKey);
-	accessor.get(hiIdx, hiKey);
-
-	uint32_t loVal = loIdx;
-	uint32_t hiVal = hiIdx;
-
-	workgroup::BitonicSort<BitonicSortConfig>::template __call<Accessor, SharedMemoryAccessor>(accessor, sharedmemAccessor, loKey, hiKey, loVal, hiVal);
-
-	accessor.set(loIdx, loKey);
-	accessor.set(hiIdx, hiKey);
+	// The sort handles load/store internally
+	workgroup::BitonicSort<BitonicSortConfig>::template __call<Accessor, SharedMemoryAccessor>(accessor, sharedmemAccessor);
 }