lampshade 0.13.0

Fast, composable GPU primitives for Rust applications using wgpu and WGSL.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

struct Config {
    capacity_items: u32,
    sort_items_per_block: u32,
    reduction_pass_count: u32,
    reduction_items_per_block: u32,
    plan_stride_words: u32,
    max_workgroups_x: u32,
    _padding_0: u32,
    _padding_1: u32,
}

@group(0) @binding(0) var<storage, read> item_count: array<u32>;
@group(0) @binding(1) var<storage, read_write> sort_dispatch_args: array<u32>;
@group(0) @binding(2) var<storage, read_write> reduction_plans: array<u32>;
@group(0) @binding(3) var<storage, read_write> reduction_dispatch_args: array<u32>;
@group(0) @binding(4) var<uniform> config: Config;

fn divide_round_up(value: u32, divisor: u32) -> u32 {
    return value / divisor + select(0u, 1u, value % divisor != 0u);
}

@compute @workgroup_size(1)
fn main() {
    var input_items = min(item_count[0], config.capacity_items);
    let sort_groups = divide_round_up(input_items, config.sort_items_per_block);
    sort_dispatch_args[0] = min(sort_groups, config.max_workgroups_x);
    sort_dispatch_args[1] = divide_round_up(sort_groups, config.max_workgroups_x);
    sort_dispatch_args[2] = 1u;

    for (var level = 0u; level < config.reduction_pass_count; level++) {
        let output_items = divide_round_up(input_items, config.reduction_items_per_block);
        let plan_offset = level * config.plan_stride_words;
        reduction_plans[plan_offset] = input_items;
        reduction_plans[plan_offset + 1u] = output_items;
        let args_offset = level * 3u;
        reduction_dispatch_args[args_offset] = min(
            output_items,
            config.max_workgroups_x
        );
        reduction_dispatch_args[args_offset + 1u] = divide_round_up(
            output_items,
            config.max_workgroups_x
        );
        reduction_dispatch_args[args_offset + 2u] = 1u;
        input_items = output_items;
    }
}