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 u32.
71/// Specialized variant for the indirect init pass, where the destination offset is computed
72/// via an indirection read from the source buffer.
73@compute @workgroup_size(64)
74fn init_fill_dispatch_args(@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
80    // Note: we ignore both src_offset and dst_offset, as we:
81    // - bind the proper view of the ChildInfo array (src), which otherwise doesn't have a constant stride
82    // - bind the entire init indirect dispatch buffer (dst) and retrieve its offset here in the shader
83    
84    // Retrieve the source ChildInfo array buffer
85    let src = thread_index * args.src_stride;
86    let dispatch_indirect_index = src_buffer[src];
87    let event_count = src_buffer[src + 1u];
88
89    // Note: we always assume 12-byte dispatch indirect structs, which is the canonical struct
90    // with its x/y/z values, without any extra field or padding.
91    let dst = dispatch_indirect_index * 3u;
92    let workgroup_count = calc_workground_count(event_count);
93    dst_buffer[dst] = workgroup_count;
94    dst_buffer[dst + 1u] = 1u;
95    dst_buffer[dst + 2u] = 1u;
96}
97
98/// Fill indirect dispatch arguments from a raw element count, by copying the element count
99/// and rounding it up to the number of thread per workgroup. Each thread copies a single
100/// single u32. Same as fill_dispatch_args(), but both read and write are from the destination
101/// buffer. The source buffer is ignored (no binding).
102@compute @workgroup_size(64)
103fn fill_dispatch_args_self(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
104    let thread_index = global_invocation_id.x;
105    if (thread_index >= args.count) {
106        return;
107    }
108    let src = args.src_offset + thread_index * args.src_stride;
109    // Note: we always assume 12-byte dispatch indirect structs, which is the canonical struct
110    // with its x/y/z values, without any extra field or padding.
111    let dst = args.dst_offset + thread_index * 3u;
112    // Note: only the dst_buffer is used
113    let thread_count = dst_buffer[src];
114    let workgroup_count = calc_workground_count(thread_count);
115    dst_buffer[dst] = workgroup_count;
116    dst_buffer[dst + 1u] = 1u;
117    dst_buffer[dst + 2u] = 1u;
118    // leave last entry untouched; sometimes it's used for unrelated things
119}