hyper-tps / distribute / hypertps-win64 / shaders / scan_propagate.glsl

#version 450

// ============================================================================
// Multi-level prefix sum -- Level 3: Propagate block prefix sums.
//
// After scan_blocks.glsl produced per-block scans and block_sums[],
// and block_sums[] has been scanned (by scan_single.glsl or recursively),
// this shader adds block_sums[blockId] to every element in that block.
//
// Dispatch: ceil(N / BLOCK_SIZE) workgroups, 1, 1
// ============================================================================

#define BLOCK_SIZE 1024
#define THREADS_PER_GROUP 256

layout(local_size_x = THREADS_PER_GROUP, local_size_y = 1, local_size_z = 1) in;

layout(set = 0, binding = 0, std430) buffer DataBuf      { uint data[];       };
layout(set = 0, binding = 1, std430) buffer BlockSumsBuf  { uint block_sums[]; };

// Phase B: Push constant replaces UBO — no buffer_update() needed between dispatches.
layout(push_constant) uniform ScanPushConstants {
    uint N;          // total number of elements
    uint _pad0;
    uint _pad1;
    uint _pad2;
};

shared uint s_blockSum;

void main() {
    uint blockId = gl_WorkGroupID.x;

    // Block 0 needs no adjustment (its prefix is 0).
    if (blockId == 0) return;

    // Load the scanned block sum once per workgroup.
    if (gl_LocalInvocationID.x == 0) {
        s_blockSum = block_sums[blockId];
    }
    memoryBarrierShared();
    barrier();

    // Each thread adds the block prefix to multiple elements.
    // We launch THREADS_PER_GROUP threads but each block has BLOCK_SIZE elements.
    uint elemsPerThread = BLOCK_SIZE / THREADS_PER_GROUP;  // = 4
    uint base = blockId * BLOCK_SIZE + gl_LocalInvocationID.x * elemsPerThread;

    for (uint i = 0; i < elemsPerThread; i++) {
        uint idx = base + i;
        if (idx < N) {
            data[idx] += s_blockSum;
        }
    }
}