1/// Arguments for a block operation on a buffer.
2struct BufferOperationArgs {
3    /// Offset, as u32 count, where the operation starts in the source buffer.
4    src_offset: u32,
5    /// Stride, as u32 count, between elements in the source buffer.
6    src_stride: u32,
7    /// Offset, as u32 count, where the operation starts in the destination buffer.
8    dst_offset: u32,
9    /// Stride, as u32 count, between elements in the destination buffer.
10    dst_stride: u32,
11    /// Number of u32 elements to process for this operation.
12    count: u32,
13}
14
15@group(0) @binding(0) var<uniform> args : BufferOperationArgs;
16@group(0) @binding(1) var<storage, read> src_buffer : array<u32>;
17@group(0) @binding(2) var<storage, read_write> dst_buffer : array<u32>;
18
19/// Clear a buffer to zero. Each thread clears a single u32.
20@compute @workgroup_size(64)
21fn zero_buffer(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
22    let thread_index = global_invocation_id.x;
23    if (thread_index >= args.count) {
24        return;
25    }
26    let dst = args.dst_offset + thread_index;
27    dst_buffer[dst] = 0u;
28}
29
30/// Copy a source buffer into a destination buffer. Each thread copies a single u32.
31@compute @workgroup_size(64)
32fn copy_buffer(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
33    let thread_index = global_invocation_id.x;
34    if (thread_index >= args.count) {
35        return;
36    }
37    let src = args.src_offset + thread_index * args.src_stride;
38    let dst = args.dst_offset + thread_index * args.dst_stride;
39    let value = src_buffer[src];
40    dst_buffer[dst] = value;
41}
42
43/// Calculate the number of workgroups to dispatch for a given number of threads.
44fn calc_workground_count(thread_count: u32) -> u32 {
45    // We assume a workgroup size of 64. This is currently the case everywhere in Hanabi,
46    // but would need to be adapted if we decided to vary the workgroup size.
47    let workgroup_count = (thread_count + 63u) >> 6u;
48    return workgroup_count;
49}
50
51/// Fill indirect dispatch arguments from a raw element count, by copying the element count
52/// and rounding it up to the number of thread per workgroup. Each thread copies a single u32.
53@compute @workgroup_size(64)
54fn fill_dispatch_args(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
55    let thread_index = global_invocation_id.x;
56    if (thread_index >= args.count) {
57        return;
58    }
59
60    let src = args.src_offset + thread_index * args.src_stride;
61    let dst = args.dst_offset + thread_index * args.dst_stride;
62    let thread_count = src_buffer[src];
63    let workgroup_count = calc_workground_count(thread_count);
64    dst_buffer[dst] = workgroup_count;
65    dst_buffer[dst + 1u] = 1u;
66    dst_buffer[dst + 2u] = 1u;
67}
68
69/// Fill indirect dispatch arguments from a raw element count, by copying the element count
70/// and rounding it up to the number of thread per workgroup. Each thread copies a single
71/// single u32. Same as fill_dispatch_args(), but both read and write are from the destination
72/// buffer. The source buffer is ignored (no binding).
73@compute @workgroup_size(64)
74fn fill_dispatch_args_self(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
75    let thread_index = global_invocation_id.x;
76    if (thread_index >= args.count) {
77        return;
78    }
79    let src = args.src_offset + thread_index * args.src_stride;
80    // Note: we always assume 12-byte dispatch indirect structs, which is the canonical struct
81    // with its x/y/z values, without any extra field or padding.
82    let dst = args.dst_offset + thread_index * 3u;
83    // Note: only the dst_buffer is used
84    let thread_count = dst_buffer[src];
85    let workgroup_count = calc_workground_count(thread_count);
86    dst_buffer[dst] = workgroup_count;
87    dst_buffer[dst + 1u] = 1u;
88    dst_buffer[dst + 2u] = 1u;
89    // leave last entry untouched; sometimes it's used for unrelated things
90}