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

#version 450

// ============================================================================
// Multi-level prefix sum -- Level 1: Per-block scan.
//
// Scans blocks of BLOCK_SIZE elements each, writing:
//   output_data[i]      = exclusive prefix sum within each block
//   block_sums[blockId] = total sum of that block
//
// After this pass, scan block_sums[] with scan_single.glsl (or recursively),
// then run scan_propagate.glsl to add block prefixes back.
//
// Dispatch: ceil(N / BLOCK_SIZE) workgroups, 1, 1
// ============================================================================

#define BLOCK_SIZE 1024
#define HALF_BLOCK 512

#define NUM_BANKS 32
#define LOG_NUM_BANKS 5
#define CONFLICT_FREE_OFFSET(n) ((n) >> LOG_NUM_BANKS)

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[]; };
layout(set = 0, binding = 2, std430) buffer BlockSumsBuf { uint block_sums[];  };

// Phase B: Push constant replaces UBO — allows chaining scan dispatches
// in a single compute list without buffer_update() syncs between them.
layout(push_constant) uniform ScanPushConstants {
    uint N;          // total number of elements
    uint _pad0;
    uint _pad1;
    uint _pad2;
};

void main() {
    uint thid = gl_LocalInvocationID.x;
    uint blockId = gl_WorkGroupID.x;
    uint blockOffset = blockId * BLOCK_SIZE;

    // ---- Load ----
    uint ai = thid;
    uint bi = thid + HALF_BLOCK;

    uint globalAi = blockOffset + ai;
    uint globalBi = blockOffset + bi;

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

    s_data[ai + bankOffsetA] = (globalAi < N) ? input_data[globalAi] : 0u;
    s_data[bi + bankOffsetB] = (globalBi < N) ? input_data[globalBi] : 0u;

    // ---- Up-sweep ----
    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;
    }

    // ---- Save block total, set root to zero ----
    if (thid == 0) {
        uint lastIdx = BLOCK_SIZE - 1 + CONFLICT_FREE_OFFSET(BLOCK_SIZE - 1);
        block_sums[blockId] = s_data[lastIdx];
        s_data[lastIdx] = 0u;
    }

    // ---- Down-sweep ----
    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 per-block exclusive scan ----
    if (globalAi < N) output_data[globalAi] = s_data[ai + bankOffsetA];
    if (globalBi < N) output_data[globalBi] = s_data[bi + bankOffsetB];
}