NetDEM v1.0
Loading...
Searching...
No Matches
radix_reorder.wgsl.h
Go to the documentation of this file.
1#pragma once
2
3#include <string>
4
6// Define the binding group layout
7@group(0) @binding(0) var<storage, read_write> data_keys: array<u32>;
8@group(0) @binding(1) var<storage, read_write> data_vals: array<u32>;
9@group(0) @binding(2) var<storage, read_write> data_keys_bak: array<u32>;
10@group(0) @binding(3) var<storage, read_write> data_vals_bak: array<u32>;
11@group(0) @binding(4) var<storage, read_write> bit_id: u32;
12
13@group(1) @binding(0) var<storage, read_write> bit_0_flag: array<u32>;
14@group(1) @binding(1) var<storage, read_write> bit_0_prefix_sum: array<u32>;
15
16@group(1) @binding(3) var<storage, read_write> bit_1_flag: array<u32>;
17@group(1) @binding(4) var<storage, read_write> bit_1_prefix_sum: array<u32>;
18
19const WORKGROUP_SIZE : u32 = 256u;
20
21@compute @workgroup_size(WORKGROUP_SIZE)
22fn RadixReorder(
23 @builtin(global_invocation_id) global_id: vec3<u32>,
24 @builtin(local_invocation_id) local_id: vec3<u32>,
25 @builtin(workgroup_id) group_id: vec3<u32>
26) {
27 let gid = global_id.x;
28 let array_size = arrayLength(&data_keys);
29 if gid < array_size {
30 let gid_max = array_size - 1u;
31 let prefix_sum_0 = bit_0_prefix_sum[gid_max] + bit_0_flag[gid_max];
32 let id: u32 = select(
33 bit_0_prefix_sum[gid],
34 bit_1_prefix_sum[gid] + prefix_sum_0,
35 bit_0_flag[gid] == 0u
36 );
37
38 data_keys[id] = data_keys_bak[gid];
39 data_vals[id] = data_vals_bak[gid];
40 }
41
42 if gid == 0 {
43 bit_id ++;
44 }
45}
46)";
shader_source
Definition radix_reorder.wgsl.h:5