Update bitonic_sort_shader.comp.hlsl

Author: CrabExtra

13_BitonicSort/app_resources/bitonic_sort_shader.comp.hlsl

@@ -1,112 +1,87 @@
-#include "nbl/builtin/hlsl/bda/bda_accessor.hlsl"
+#include "common.hlsl"
+#include "nbl/builtin/hlsl/workgroup/basic.hlsl"
+#include "nbl/builtin/hlsl/workgroup/bitonic_sort.hlsl"
 
-struct BitonicPushData
+[[vk::push_constant]] PushConstantData pushConstants;
+
+using namespace nbl::hlsl;
+
+using BitonicSortConfig = workgroup::bitonic_sort::bitonic_sort_config<ElementsPerThreadLog2, WorkgroupSizeLog2, uint32_t, uint32_t, less<uint32_t> >;
+
+NBL_CONSTEXPR uint32_t WorkgroupSize = BitonicSortConfig::WorkgroupSize;
+
+groupshared uint32_t sharedmem[2 * WorkgroupSize];
+
+uint32_t3 glsl::gl_WorkGroupSize() { return uint32_t3(uint32_t(BitonicSortConfig::WorkgroupSize), 1, 1); }
+
+struct SharedMemoryAccessor
 {
-    uint64_t inputKeyAddress;
-    uint64_t inputValueAddress;
-    uint64_t outputKeyAddress;
-    uint64_t outputValueAddress;
-    uint32_t dataElementCount;
+	template <typename AccessType, typename IndexType>
+	void set(IndexType idx, AccessType value)
+	{
+		sharedmem[idx] = value;
+	}
+
+	template <typename AccessType, typename IndexType>
+	void get(IndexType idx, NBL_REF_ARG(AccessType) value)
+	{
+		value = sharedmem[idx];
+	}
+
+	void workgroupExecutionAndMemoryBarrier()
+	{
+		glsl::barrier();
+	}
+
 };
 
-using namespace nbl::hlsl;
+struct Accessor
+{
+	static Accessor create(const uint64_t address)
+	{
+		Accessor accessor;
+		accessor.address = address;
+		return accessor;
+	}
 
-[[vk::push_constant]] BitonicPushData pushData;
+	template <typename AccessType, typename IndexType>
+	void get(const IndexType index, NBL_REF_ARG(AccessType) value)
+	{
+		value = vk::RawBufferLoad<AccessType>(address + index * sizeof(AccessType));
+	}
 
-using DataPtr = bda::__ptr<uint32_t>;
-using DataAccessor = BdaAccessor<uint32_t>;
+	template <typename AccessType, typename IndexType>
+	void set(const IndexType index, const AccessType value)
+	{
+		vk::RawBufferStore<AccessType>(address + index * sizeof(AccessType), value);
+	}
 
-groupshared uint32_t sharedKeys[ElementCount];
-groupshared uint32_t sharedValues[ElementCount];
+	uint64_t address;
+};
 
-[numthreads(WorkgroupSize, 1, 1)]
+[numthreads(BitonicSortConfig::WorkgroupSize, 1, 1)]
 [shader("compute")]
-void main(uint32_t3 dispatchId : SV_DispatchThreadID, uint32_t3 localId : SV_GroupThreadID)
+void main()
 {
-    const uint32_t threadId = localId.x;
-    const uint32_t dataSize = pushData.dataElementCount;
-    
-    DataAccessor inputKeys = DataAccessor::create(DataPtr::create(pushData.inputKeyAddress));
-    DataAccessor inputValues = DataAccessor::create(DataPtr::create(pushData.inputValueAddress));
-    
-    for (uint32_t i = threadId; i < dataSize; i += WorkgroupSize)
-    {
-        inputKeys.get(i, sharedKeys[i]);
-        inputValues.get(i, sharedValues[i]);
-    }
-    
-    // Synchronize all threads after loading
-    GroupMemoryBarrierWithGroupSync();
-    
-
-    for (uint32_t stage = 0; stage < Log2ElementCount; stage++)
-    {
-        for (uint32_t pass = 0; pass <= stage; pass++)
-        {
-            const uint32_t compareDistance = 1 << (stage - pass);
-            
-            for (uint32_t i = threadId; i < dataSize; i += WorkgroupSize)
-            {
-                const uint32_t partnerId = i ^ compareDistance;
-                
-                if (partnerId >= dataSize)
-                    continue;
-               
-                const uint32_t waveSize = WaveGetLaneCount();
-                const uint32_t myWaveId = i / waveSize;
-                const uint32_t partnerWaveId = partnerId / waveSize;
-                const bool sameWave = (myWaveId == partnerWaveId);
-
-                uint32_t myKey, myValue, partnerKey, partnerValue;
-                [branch]
-                if (sameWave && compareDistance < waveSize)
-                {
-                    // WAVE INTRINSIC
-                    myKey = sharedKeys[i];
-                    myValue = sharedValues[i];
-
-                    const uint32_t partnerLane = partnerId % waveSize;
-                    partnerKey = WaveReadLaneAt(myKey, partnerLane);
-                    partnerValue = WaveReadLaneAt(myValue, partnerLane);
-                }
-                else
-                {
-                    // SHARED MEM
-                    myKey = sharedKeys[i];
-                    myValue = sharedValues[i];
-                    partnerKey = sharedKeys[partnerId];
-                    partnerValue = sharedValues[partnerId];
-                }
-
-                const uint32_t sequenceSize = 1 << (stage + 1);
-                const uint32_t sequenceIndex = i / sequenceSize;
-                const bool sequenceAscending = (sequenceIndex % 2) == 0;
-                const bool ascending = true;
-                const bool finalDirection = sequenceAscending == ascending;
-                
-                const bool swap = (myKey > partnerKey) == finalDirection;
-                
-                // WORKGROUP COORDINATION: Only lower-indexed element writes both
-                if (i < partnerId && swap)
-                {
-                    sharedKeys[i] = partnerKey;
-                    sharedKeys[partnerId] = myKey;
-                    sharedValues[i] = partnerValue;
-                    sharedValues[partnerId] = myValue;
-                }
-            }
-            
-            GroupMemoryBarrierWithGroupSync();
-        }
-    }
-    
-
-    DataAccessor outputKeys = DataAccessor::create(DataPtr::create(pushData.outputKeyAddress));
-    DataAccessor outputValues = DataAccessor::create(DataPtr::create(pushData.outputValueAddress));
-    
-    for (uint32_t i = threadId; i < dataSize; i += WorkgroupSize)
-    {
-        outputKeys.set(i, sharedKeys[i]);
-        outputValues.set(i, sharedValues[i]);
-    }
-}
+	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);
+}