Patch things for cooperative binary search test

Author: deprilula28

72_CooperativeBinarySearch/app_resources/binarySearch.comp.hlsl

@@ -5,16 +5,115 @@
 #pragma wave shader_stage(compute)
 
 #include "common.h"
+#include "nbl/builtin/hlsl/glsl_compat/subgroup_ballot.hlsl"
 using namespace nbl::hlsl;
 
-[[vk::push_constant]] ConstantBuffer<PushConstants> Constants;
+[[vk::push_constant]] PushConstants Constants;
 [[vk::binding(0)]] StructuredBuffer<uint> Histogram;
 [[vk::binding(1)]] RWStructuredBuffer<uint> Output;
 
 static const uint32_t GroupsharedSize = 256;
 
+uint getNextPowerOfTwo(uint number) {
+	return 2 << firstbithigh(number - 1);
+}
+
+uint getLaneWithFirstBitSet(bool condition) {
+	uint4 ballot = WaveActiveBallot(condition);
+	if (all(ballot == 0)) {
+		return WaveGetLaneCount();
+	}
+	return nbl::hlsl::glsl::subgroupBallotFindLSB(ballot);
+}
+
+// findValue must be the same across the entire wave
+// Could use something like WaveReadFirstLane to be fully sure
+uint binarySearchLowerBoundFindValue(uint findValue, StructuredBuffer<uint> searchBuffer, uint searchBufferSize) {
+	uint lane = WaveGetLaneIndex();
+	
+	uint left = 0;
+	uint right = searchBufferSize - 1;
+
+	uint32_t range = getNextPowerOfTwo(right - left);
+	// do pivots as long as we can't coalesced load
+	while (range > WaveGetLaneCount())
+	{
+		// there must be at least 1 gap between subsequent pivots 
+		const uint32_t step = range / WaveGetLaneCount(); 
+		const uint32_t halfStep = step >> 1;
+		const uint32_t pivotOffset = lane * step+halfStep;
+		const uint32_t pivotIndex = left + pivotOffset;
+
+		uint4 notGreaterPivots = WaveActiveBallot(pivotIndex < right && !(findValue < searchBuffer[pivotIndex]));
+		uint partition = nbl::hlsl::glsl::subgroupBallotBitCount(notGreaterPivots);
+		// only move left if needed
+		if (partition != 0)
+			left += partition * step - halfStep;
+		// if we go into final half partition, the range becomes less too
+		range = partition != WaveGetLaneCount() ? step : halfStep;
+	}
+
+	uint threadSearchIndex = left + lane;
+	bool laneValid = threadSearchIndex < searchBufferSize;
+	uint histAtIndex = laneValid ? searchBuffer[threadSearchIndex] : -1;
+	uint firstLaneGreaterThan = getLaneWithFirstBitSet(histAtIndex > findValue);
+
+	return left + firstLaneGreaterThan - 1;
+}
+
+groupshared uint shared_groupSearchBufferMinIndex;
+groupshared uint shared_groupSearchBufferMaxIndex;
+groupshared uint shared_groupSearchValues[GroupsharedSize];
+
+// Binary search using the entire workgroup, making it log32 or log64 (every iteration, the possible set of 
+// values is divided by the number of lanes in a wave)
+uint binarySearchLowerBoundCooperative(uint groupIndex, uint groupThread, StructuredBuffer<uint> searchBuffer, uint searchBufferSize) {
+	uint minSearchValue = groupIndex.x * GroupsharedSize;
+	uint maxSearchValue = ((groupIndex.x + 1) * GroupsharedSize) - 1;
+
+	// On each workgroup, two subgroups do the search
+	// - One searches for the minimum, the other searches for the maximum
+	// - Store the minimum and maximum on groupshared memory, then do a barrier
+	uint wave = groupThread / WaveGetLaneCount();
+	if (wave < 2) {
+		uint search = wave == 0 ? minSearchValue : maxSearchValue;
+		uint searchResult = binarySearchLowerBoundFindValue(search, searchBuffer, searchBufferSize);
+		if (WaveIsFirstLane()) {
+			if (wave == 0) shared_groupSearchBufferMinIndex = searchResult;
+			else shared_groupSearchBufferMaxIndex = searchResult;
+		}
+	}
+	GroupMemoryBarrierWithGroupSync();
+
+	// Since every instance has at least one triangle, we know that having workgroup values 
+	// for each value in the range of minimum to maximum will suffice.
+
+	// Write every value in the range to groupshared memory and barrier.
+	uint idx = shared_groupSearchBufferMinIndex + groupThread.x;
+	if (idx <= shared_groupSearchBufferMaxIndex) {
+		shared_groupSearchValues[groupThread.x] = searchBuffer[idx];
+	}
+	GroupMemoryBarrierWithGroupSync();
+
+	uint maxValueIndex = shared_groupSearchBufferMaxIndex - shared_groupSearchBufferMinIndex;
+
+	uint searchValue = minSearchValue + groupThread;
+	uint currentSearchValueIndex = 0;
+	uint laneValue = shared_groupSearchBufferMaxIndex;
+	while (currentSearchValueIndex <= maxValueIndex) {
+		uint curValue = shared_groupSearchValues[currentSearchValueIndex];
+		if (curValue > searchValue) {
+			laneValue = shared_groupSearchBufferMinIndex + currentSearchValueIndex - 1;
+			break;
+		}
+		currentSearchValueIndex ++;
+	}
+
+	return laneValue;
+}
+
 [numthreads(256, 1, 1)]
 void main(const uint3 thread : SV_DispatchThreadID, const uint3 groupThread : SV_GroupThreadID, const uint3 group : SV_GroupID)
 {
-
+    Output[thread.x] = binarySearchLowerBoundCooperative(group.x, groupThread.x, Histogram, Constants.EntityCount);
 }

72_CooperativeBinarySearch/main.cpp

@@ -22,6 +22,11 @@ using namespace nbl::examples;
 
 //using namespace glm;
 
+static constexpr uint32_t TestCaseIndices[] = {
+#include "testCaseData.h"
+};
+
+
 void cpu_tests();
 
 class CooperativeBinarySearch final : public application_templates::MonoDeviceApplication, public BuiltinResourcesApplication
@@ -101,14 +106,19 @@ class CooperativeBinarySearch final : public application_templates::MonoDeviceAp
 
             auto reqs = m_buffers[i]->getMemoryReqs();
             reqs.memoryTypeBits &= m_device->getPhysicalDevice()->getHostVisibleMemoryTypeBits();
-            m_device->allocate(reqs, m_buffers[i].get());
+
+            m_allocations[i] = m_device->allocate(reqs, m_buffers[i].get());
+            
+            auto allocationType = i == 0 ? IDeviceMemoryAllocation::EMCAF_WRITE : IDeviceMemoryAllocation::EMCAF_READ;
+            auto mapResult = m_allocations[i].memory->map({ 0ull,m_allocations[i].memory->getAllocationSize() }, allocationType);
+            assert(mapResult);
         }
 
 		smart_refctd_ptr<IDescriptorPool> descriptorPool = nullptr;
 		{
             IDescriptorPool::SCreateInfo createInfo = {};
             createInfo.maxSets = 1;
-            createInfo.maxDescriptorCount[static_cast<uint32_t>(IDescriptor::E_TYPE::ET_STORAGE_BUFFER)] = 1;
+            createInfo.maxDescriptorCount[static_cast<uint32_t>(IDescriptor::E_TYPE::ET_STORAGE_BUFFER)] = bindingCount;
             descriptorPool = m_device->createDescriptorPool(std::move(createInfo));
         }
 
@@ -130,6 +140,14 @@ class CooperativeBinarySearch final : public application_templates::MonoDeviceAp
 
         m_device->updateDescriptorSets(bindingCount, writeDescriptorSets, 0u, nullptr);
 
+        // Write test data to the m_buffers[0]
+        auto outPtr = m_allocations[0].memory->getMappedPointer();
+        assert(outPtr);
+        memcpy(
+            reinterpret_cast<void*>(outPtr), 
+            reinterpret_cast<const void*>(&TestCaseIndices[0]), 
+            sizeof(TestCaseIndices));
+
         // In contrast to fences, we just need one semaphore to rule all dispatches
         return true;
     }
@@ -196,9 +214,8 @@ class CooperativeBinarySearch final : public application_templates::MonoDeviceAp
             m_device->blockForSemaphores(waitInfos);
         }
 
-		auto mem = m_buffers[1]->getBoundMemory();
-		assert(mem.memory->isMappable());
-		auto* ptr = mem.memory->map({ .offset = 0, .length = mem.memory->getAllocationSize() });
+        auto ptr = m_allocations[1].memory->getMappedPointer();
+        assert(ptr);
         printf("readback ptr %p\n", ptr);
 
         m_keepRunning = false;
@@ -216,6 +233,7 @@ class CooperativeBinarySearch final : public application_templates::MonoDeviceAp
     smart_refctd_ptr<IGPUDescriptorSet> m_descriptorSet;
 
     smart_refctd_ptr<IGPUBuffer> m_buffers[2];
+	nbl::video::IDeviceMemoryAllocator::SAllocation m_allocations[2] = {};
     smart_refctd_ptr<IGPUCommandBuffer> m_cmdbuf = nullptr;
     IQueue* m_queue;
     smart_refctd_ptr<IGPUCommandPool> m_commandPool;