#version 450
// ============================================================================
// Single-pass decoupled look-back prefix sum.
//
// Based on Merrill & Garland 2016: "Single-pass Parallel Prefix Scan with
// Decoupled Look-back" (NVIDIA Technical Report NVR-2016-002).
//
// ADVANTAGES:
// - Single dispatch for ANY array size (no multi-level passes)
// - Near-memcpy bandwidth for large arrays
// - O(N) work, O(log N) span
//
// REQUIREMENTS:
// - Vulkan 1.1+ with VK_SUBGROUP_FEATURE_ARITHMETIC_BIT
// - Forward thread progress guarantees (NVIDIA, AMD RDNA2+)
// - NOT safe on: Mali, older Adreno, Apple GPUs, WebGPU/Metal
// - Requires VK_KHR_vulkan_memory_model for correct acquire/release
//
// CAUTION: If your target GPU lacks forward progress guarantees, use the
// multi-level approach (scan_blocks + scan_single + scan_propagate) instead.
//
// Partition tile status flags (packed into upper 2 bits of uint):
// X = 0: not yet available
// A = 1: aggregate (local reduction) available
// P = 2: prefix (inclusive prefix) available
//
// state_flags[partitionId] stores: (status << 30) | (value & 0x3FFFFFFF)
// This limits values to 30 bits (max ~1 billion), fine for flag counts.
//
// Dispatch: ceil(N / TILE_SIZE) workgroups, 1, 1
// PRE-DISPATCH: clear state_flags[] to all zeros
// clear partition_counter[0] to 0
// ============================================================================
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_basic : require
#define TILE_SIZE 1024
#define THREADS 256
#define ELEMS_PER_THREAD 4 // TILE_SIZE / THREADS
#define FLAG_X 0u
#define FLAG_A 1u
#define FLAG_P 2u
#define PACK(status, value) (((status) << 30u) | ((value) & 0x3FFFFFFFu))
#define UNPACK_STATUS(packed) ((packed) >> 30u)
#define UNPACK_VALUE(packed) ((packed) & 0x3FFFFFFFu)
layout(local_size_x = THREADS, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0, std430) buffer InputBuf { uint input_data[]; };
layout(set = 0, binding = 1, std430) buffer OutputBuf { uint output_data[]; };
layout(set = 0, binding = 2, std430) coherent buffer StateBuf { uint state_flags[]; };
layout(set = 0, binding = 3, std430) coherent buffer CounterBuf { uint partition_counter[]; };
layout(set = 0, binding = 4, std140) uniform Params {
uint N; // total elements
uint _pad0;
uint _pad1;
uint _pad2;
};
// Shared memory for intra-workgroup scan
shared uint s_subgroup_sums[32]; // max 256/8 = 32 subgroups
shared uint s_partition_id;
shared uint s_exclusive_prefix; // prefix from all preceding partitions
void main() {
uint lid = gl_LocalInvocationID.x;
// ---- Dynamic partition assignment via atomic counter ----
if (lid == 0) {
s_partition_id = atomicAdd(partition_counter[0], 1u);
}
memoryBarrierShared();
barrier();
uint partId = s_partition_id;
uint tileOffset = partId * TILE_SIZE;
// ---- Phase 1: Load and compute local serial sum per thread ----
uint thread_vals[ELEMS_PER_THREAD];
uint thread_sum = 0u;
for (uint e = 0; e < ELEMS_PER_THREAD; e++) {
uint idx = tileOffset + lid * ELEMS_PER_THREAD + e;
uint v = (idx < N) ? input_data[idx] : 0u;
thread_vals[e] = v;
thread_sum += v;
}
// ---- Phase 2: Intra-workgroup exclusive scan of thread_sum ----
// Step 2a: Subgroup inclusive scan
uint sg_inclusive = subgroupInclusiveAdd(thread_sum);
if (gl_SubgroupInvocationID == gl_SubgroupSize - 1) {
s_subgroup_sums[gl_SubgroupID] = sg_inclusive;
}
memoryBarrierShared();
barrier();
// Step 2b: Scan subgroup totals (subgroup 0 only)
if (gl_SubgroupID == 0) {
uint v = (gl_SubgroupInvocationID < gl_NumSubgroups)
? s_subgroup_sums[gl_SubgroupInvocationID] : 0u;
uint scanned = subgroupInclusiveAdd(v);
s_subgroup_sums[gl_SubgroupInvocationID] = scanned;
}
memoryBarrierShared();
barrier();
// Step 2c: Combine for exclusive scan
uint preceding_subgroups = (gl_SubgroupID > 0)
? s_subgroup_sums[gl_SubgroupID - 1] : 0u;
uint thread_exclusive = sg_inclusive - thread_sum + preceding_subgroups;
// Tile aggregate = last subgroup's scanned total
uint tile_aggregate = s_subgroup_sums[gl_NumSubgroups - 1];
// ---- Phase 3: Publish aggregate / look-back ----
if (partId == 0) {
// First partition: prefix = aggregate
if (lid == 0) {
atomicExchange(state_flags[0], PACK(FLAG_P, tile_aggregate));
s_exclusive_prefix = 0u;
}
} else {
// Publish aggregate
if (lid == 0) {
atomicExchange(state_flags[partId], PACK(FLAG_A, tile_aggregate));
}
// Look-back: thread 0 walks backward
if (lid == 0) {
uint running_prefix = 0u;
int lookback = int(partId) - 1;
while (lookback >= 0) {
uint their_state = atomicAdd(state_flags[lookback], 0u); // atomic load
uint status = UNPACK_STATUS(their_state);
if (status == FLAG_P) {
// They have full inclusive prefix -- done
running_prefix += UNPACK_VALUE(their_state);
break;
} else if (status == FLAG_A) {
// They have aggregate only -- accumulate and keep looking
running_prefix += UNPACK_VALUE(their_state);
lookback--;
}
// else FLAG_X: spin (not ready yet)
}
// Publish our inclusive prefix
uint our_prefix = running_prefix + tile_aggregate;
atomicExchange(state_flags[partId], PACK(FLAG_P, our_prefix));
s_exclusive_prefix = running_prefix;
}
}
memoryBarrierShared();
barrier();
uint global_prefix = s_exclusive_prefix;
// ---- Phase 4: Write final exclusive scan output ----
uint running = global_prefix + thread_exclusive;
for (uint e = 0; e < ELEMS_PER_THREAD; e++) {
uint idx = tileOffset + lid * ELEMS_PER_THREAD + e;
if (idx < N) {
output_data[idx] = running;
running += thread_vals[e];
}
}
}