radix_count.wgsl.h

Go to the documentation of this file.

#pragma once
 
#include <string>
 
shader_source = R"(
// Define the binding group layout
@group(0) @binding(0) var<storage, read_write> data_keys: array<u32>;
@group(0) @binding(1) var<storage, read_write> data_vals: array<u32>;
@group(0) @binding(2) var<storage, read_write> data_keys_bak: array<u32>;
@group(0) @binding(3) var<storage, read_write> data_vals_bak: array<u32>;
@group(0) @binding(4) var<storage, read_write> bit_id: u32;
 
@group(1) @binding(0) var<storage, read_write> bit_0_flag: array<u32>;
@group(1) @binding(1) var<storage, read_write> bit_0_prefix_sum: array<u32>;
@group(1) @binding(2) var<storage, read_write> bit_0_aux_src: array<u32>;
 
@group(1) @binding(3) var<storage, read_write> bit_1_flag: array<u32>;
@group(1) @binding(4) var<storage, read_write> bit_1_prefix_sum: array<u32>;
@group(1) @binding(5) var<storage, read_write> bit_1_aux_src: array<u32>;
 
// Shared memory for the workgroup
const WORKGROUP_SIZE : u32 = 256u;
var<workgroup> shared_data_0: array<u32, WORKGROUP_SIZE>;
var<workgroup> shared_data_1: array<u32, WORKGROUP_SIZE>;
 
@compute @workgroup_size(WORKGROUP_SIZE)
fn RadixCount(
    @builtin(global_invocation_id) global_id: vec3<u32>,
    @builtin(local_invocation_id) local_id: vec3<u32>,
    @builtin(workgroup_id) group_id: vec3<u32>
) {
    // ids
    let tid = local_id.x;
    let gid = global_id.x;
 
    // Load data into shared memory
    if gid < arrayLength(&data_keys) {
        data_keys_bak[gid] = data_keys[gid];
        data_vals_bak[gid] = data_vals[gid];
 
        let extract_bits: u32 = (data_vals[gid] >> bit_id) & 0x1;
        bit_0_flag[gid] = 1 - extract_bits;
        bit_1_flag[gid] = extract_bits;
 
        shared_data_0[tid] = bit_0_flag[gid];
        shared_data_1[tid] = bit_1_flag[gid];
    }
    workgroupBarrier();
 
    // Up-sweep (reduce) phase
    var offset = 1u;
    while offset < WORKGROUP_SIZE {
        let idx = (tid + 1u) * offset * 2u - 1u;
        if idx < WORKGROUP_SIZE {
            shared_data_0[idx] += shared_data_0[idx - offset];
            shared_data_1[idx] += shared_data_1[idx - offset];
        }
        offset *= 2u;
        workgroupBarrier();
    }
 
    // Clear the last element
    if tid == 0u {
        shared_data_0[WORKGROUP_SIZE - 1u] = 0u;
        shared_data_1[WORKGROUP_SIZE - 1u] = 0u;
    }
    workgroupBarrier();
 
    // Down-sweep phase
    offset = WORKGROUP_SIZE / 2u;
    while offset > 0u {
        let idx = (tid + 1u) * offset * 2u - 1u;
        if idx < WORKGROUP_SIZE {
            let t_0 = shared_data_0[idx - offset];
            shared_data_0[idx - offset] = shared_data_0[idx];
            shared_data_0[idx] += t_0;
 
            let t_1 = shared_data_1[idx - offset];
            shared_data_1[idx - offset] = shared_data_1[idx];
            shared_data_1[idx] += t_1;
        }
        offset /= 2u;
        workgroupBarrier();
    }
 
    // Write the result back to the global memory
    if gid < arrayLength(&data_keys) {
        bit_0_prefix_sum[gid] = shared_data_0[tid];
        bit_1_prefix_sum[gid] = shared_data_1[tid];
 
        if tid == (WORKGROUP_SIZE - 1u) {
            bit_0_aux_src[group_id.x] = bit_0_flag[gid] + bit_0_prefix_sum[gid];
            bit_1_aux_src[group_id.x] = bit_1_flag[gid] + bit_1_prefix_sum[gid];
        }
    }
}
)";