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

#version 450

// ============================================================================
// Subgroup-accelerated per-block scan (Vulkan 1.1+ required).
//
// Uses subgroupInclusiveAdd() for intra-subgroup scan (no shared memory
// bank conflicts, no explicit tree traversal within a subgroup).
// Falls back to shared memory only for cross-subgroup communication.
//
// This is the MODERN approach -- significantly faster than the Blelloch tree
// when subgroup ops are available. Use this as the Level-1 block scan if
// your Vulkan driver supports VK_SUBGROUP_FEATURE_ARITHMETIC_BIT.
//
// Each thread processes 1 element. BLOCK_SIZE = workgroup size.
// Dispatch: ceil(N / BLOCK_SIZE) workgroups, 1, 1
//
// Requires: #extension GL_KHR_shader_subgroup_arithmetic : require
// ============================================================================

#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_basic : require

#define BLOCK_SIZE 256

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

layout(set = 0, binding = 3, std140) uniform Params {
    uint N;
    uint _pad0;
    uint _pad1;
    uint _pad2;
};

// Shared memory for subgroup totals. Max subgroups = 256/32 = 8.
// Allocate 32 to handle subgroup sizes as small as 8.
shared uint s_subgroup_sums[32];

void main() {
    uint lid = gl_LocalInvocationID.x;
    uint gid = gl_GlobalInvocationID.x;
    uint blockId = gl_WorkGroupID.x;

    // Load element (zero if out of bounds)
    uint val = (gid < N) ? input_data[gid] : 0u;

    // ---- Step 1: Intra-subgroup inclusive scan ----
    uint subgroup_scan = subgroupInclusiveAdd(val);

    // The last invocation in each subgroup writes its total.
    if (gl_SubgroupInvocationID == gl_SubgroupSize - 1) {
        s_subgroup_sums[gl_SubgroupID] = subgroup_scan;
    }

    memoryBarrierShared();
    barrier();

    // ---- Step 2: Scan subgroup totals (done by subgroup 0) ----
    if (gl_SubgroupID == 0) {
        uint num_subgroups = gl_NumSubgroups;
        uint sg_val = (gl_SubgroupInvocationID < num_subgroups)
                    ? s_subgroup_sums[gl_SubgroupInvocationID]
                    : 0u;
        uint sg_scan = subgroupInclusiveAdd(sg_val);
        s_subgroup_sums[gl_SubgroupInvocationID] = sg_scan;
    }

    memoryBarrierShared();
    barrier();

    // ---- Step 3: Add preceding subgroups' prefix ----
    uint block_prefix = 0u;
    if (gl_SubgroupID > 0) {
        block_prefix = s_subgroup_sums[gl_SubgroupID - 1];
    }

    // Convert inclusive to exclusive: exclusive[i] = inclusive[i] - val
    uint exclusive_scan = subgroup_scan - val + block_prefix;

    // ---- Write results ----
    if (gid < N) {
        output_data[gid] = exclusive_scan;
    }

    // Last thread in block writes block total sum.
    if (lid == BLOCK_SIZE - 1 || gid == N - 1) {
        // Block total = last inclusive scan + prefix from preceding subgroups
        block_sums[blockId] = subgroup_scan + block_prefix;
    }
}