dc-terrain / shaders / scan_single.glsl

#version 450

// ============================================================================
// Single-workgroup exclusive prefix sum (Blelloch up-sweep / down-sweep).
// Handles arrays up to BLOCK_SIZE elements (default 1024) in one dispatch.
//
// Usage: dispatch(1, 1, 1) with N <= BLOCK_SIZE.
// Result: output[i] = sum of input[0..i-1]  (exclusive scan).
//         output[N] (if within buffer) holds total sum.
//
// Bank-conflict-free shared memory via padding.
// ============================================================================

// Each thread handles 2 elements => 512 threads scan 1024 elements.
#define BLOCK_SIZE 1024
#define HALF_BLOCK 512

// Bank conflict avoidance: 32 banks on modern GPUs.
// Pad every 32 elements with 1 slot.
#define NUM_BANKS 32
#define LOG_NUM_BANKS 5
#define CONFLICT_FREE_OFFSET(n) ((n) >> LOG_NUM_BANKS)

// Shared memory: BLOCK_SIZE + padding slots
// Max padding = BLOCK_SIZE / NUM_BANKS = 32 extra slots
shared uint s_data[BLOCK_SIZE + (BLOCK_SIZE >> LOG_NUM_BANKS)];

layout(local_size_x = HALF_BLOCK, 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[]; };

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

void main() {
    uint thid = gl_LocalInvocationID.x;

    // ---- Load into shared memory with bank-conflict-free indexing ----
    uint ai = thid;
    uint bi = thid + HALF_BLOCK;

    uint bankOffsetA = CONFLICT_FREE_OFFSET(ai);
    uint bankOffsetB = CONFLICT_FREE_OFFSET(bi);

    // Load from global; zero-pad if beyond N
    s_data[ai + bankOffsetA] = (ai < N) ? input_data[ai] : 0u;
    s_data[bi + bankOffsetB] = (bi < N) ? input_data[bi] : 0u;

    // ---- Up-sweep (reduce) phase ----
    uint offset = 1;
    for (uint d = HALF_BLOCK; d > 0; d >>= 1) {
        memoryBarrierShared();
        barrier();

        if (thid < d) {
            uint ai_idx = offset * (2 * thid + 1) - 1;
            uint bi_idx = offset * (2 * thid + 2) - 1;
            ai_idx += CONFLICT_FREE_OFFSET(ai_idx);
            bi_idx += CONFLICT_FREE_OFFSET(bi_idx);
            s_data[bi_idx] += s_data[ai_idx];
        }
        offset <<= 1;
    }

    // ---- Store total sum and clear last element (exclusive scan) ----
    if (thid == 0) {
        uint lastIdx = BLOCK_SIZE - 1 + CONFLICT_FREE_OFFSET(BLOCK_SIZE - 1);
        // Store total sum at output_data[N] if caller allocated space
        if (N <= BLOCK_SIZE) {
            output_data[N] = s_data[lastIdx];
        }
        s_data[lastIdx] = 0u;
    }

    // ---- Down-sweep phase ----
    for (uint d = 1; d < BLOCK_SIZE; d <<= 1) {
        offset >>= 1;
        memoryBarrierShared();
        barrier();

        if (thid < d) {
            uint ai_idx = offset * (2 * thid + 1) - 1;
            uint bi_idx = offset * (2 * thid + 2) - 1;
            ai_idx += CONFLICT_FREE_OFFSET(ai_idx);
            bi_idx += CONFLICT_FREE_OFFSET(bi_idx);

            uint temp = s_data[ai_idx];
            s_data[ai_idx] = s_data[bi_idx];
            s_data[bi_idx] += temp;
        }
    }

    memoryBarrierShared();
    barrier();

    // ---- Write results back to global memory ----
    if (ai < N) output_data[ai] = s_data[ai + bankOffsetA];
    if (bi < N) output_data[bi] = s_data[bi + bankOffsetB];
}