125 lines
3.5 KiB
Plaintext
125 lines
3.5 KiB
Plaintext
|
|
[[vk::binding(0, 0)]]
|
|
RWStructuredBuffer<uint32_t> count_buffer;
|
|
|
|
[[vk::binding(0, 1)]]
|
|
RWStructuredBuffer<uint32_t> reduced_buffer;
|
|
[[vk::binding(1, 1)]]
|
|
RWStructuredBuffer<uint32_t> sum_buffer;
|
|
[[vk::binding(2, 1)]]
|
|
RWStructuredBuffer<uint32_t> compaction_buffer;
|
|
|
|
groupshared uint32_t local_data[256 * 2];
|
|
static uint32_t THREAD_WIDTH = 256;
|
|
static uint32_t DATA_WIDTH = THREAD_WIDTH * 2;
|
|
|
|
[numthreads(256, 1, 1)]
|
|
[shader("compute")]
|
|
void block_sum(
|
|
uint32_t3 workgroup_id: SV_GroupID,
|
|
uint32_t3 local_thread_id: SV_GroupThreadID,
|
|
uint32_t3 global_thread_id: SV_DispatchThreadID)
|
|
{
|
|
// Perform sum in current block
|
|
|
|
// Copy local_datainto LDS with predicate
|
|
let thread_index = global_thread_id.x;
|
|
let local_thread_index = local_thread_id.x;
|
|
let total = count_buffer.getCount();
|
|
|
|
if (thread_index * 2 < total)
|
|
{
|
|
local_data[local_thread_index * 2] = select(count_buffer[thread_index * 2] != 0, 1, 0);
|
|
}
|
|
else
|
|
{
|
|
local_data[local_thread_index * 2] = 0;
|
|
}
|
|
|
|
if (thread_index * 2 + 1 < total)
|
|
{
|
|
local_data[local_thread_index * 2 + 1] = select(count_buffer[thread_index * 2 + 1] != 0, 1, 0);
|
|
}
|
|
else
|
|
{
|
|
local_data[local_thread_index * 2 + 1] = 0;
|
|
}
|
|
|
|
GroupMemoryBarrierWithGroupSync();
|
|
|
|
var width : uint32_t = 2;
|
|
while (width <= DATA_WIDTH)
|
|
{
|
|
let dest_index = width * (thread_index + 1) - 1;
|
|
let get_index = dest_index - (width / 2);
|
|
// println!("{}, {}", get_index, dest_index);
|
|
if (dest_index < DATA_WIDTH)
|
|
{
|
|
local_data[dest_index] += local_data[get_index];
|
|
}
|
|
width *= 2;
|
|
GroupMemoryBarrierWithGroupSync();
|
|
}
|
|
|
|
local_data[DATA_WIDTH - 1] = 0;
|
|
while (width >= 2)
|
|
{
|
|
let dest_index = width * (thread_index + 1) - 1;
|
|
let get_index = dest_index - (width / 2);
|
|
// println!("{}, {}", get_index, dest_index);
|
|
if (dest_index < DATA_WIDTH)
|
|
{
|
|
let self_data = local_data[dest_index];
|
|
local_data[dest_index] += local_data[get_index];
|
|
local_data[get_index] = self_data;
|
|
}
|
|
width /= 2;
|
|
GroupMemoryBarrierWithGroupSync();
|
|
}
|
|
|
|
// Block now contains running local sum
|
|
// Dump back to sum buffer
|
|
sum_buffer[2 * thread_index] = local_data[2 * local_thread_index];
|
|
sum_buffer[2 * thread_index + 1] = local_data[2 * local_thread_index + 1];
|
|
|
|
// Write to reduced buffer
|
|
reduced_buffer[workgroup_id.x] = local_data[DATA_WIDTH - 1];
|
|
}
|
|
|
|
[numthreads(1, 1, 1)]
|
|
[shader("compute")]
|
|
void linear_reduced_sum(
|
|
uint32_t3 workgroup_id: SV_GroupID,
|
|
uint32_t3 local_thread_id: SV_GroupThreadID,
|
|
uint32_t3 global_thread_id: SV_DispatchThreadID)
|
|
{
|
|
let size = reduced_buffer.getCount();
|
|
|
|
// Perform exclusive sum
|
|
var running_sum : uint32_t = 0;
|
|
for (uint32_t i = 0; i < size; i++)
|
|
{
|
|
let value = reduced_buffer[i];
|
|
reduced_buffer[i] = running_sum;
|
|
running_sum += value;
|
|
}
|
|
}
|
|
|
|
[numthreads(256, 1, 1)]
|
|
[shader("compute")]
|
|
void uniform_add(
|
|
uint32_t3 workgroup_id: SV_GroupID,
|
|
uint32_t3 local_thread_id: SV_GroupThreadID,
|
|
uint32_t3 global_thread_id: SV_DispatchThreadID)
|
|
{
|
|
let thread_index = global_thread_id.x;
|
|
let local_thread_index = local_thread_id.x;
|
|
|
|
// Gather
|
|
let reduced_value = reduced_buffer[global_thread_id.x];
|
|
// Apply
|
|
sum_buffer[thread_index * 2] += reduced_value;
|
|
sum_buffer[thread_index * 2 + 1] += reduced_value;
|
|
}
|
|
|