| // SPDX-License-Identifier: Apache-2.0 OR MIT OR Unlicense |
| |
| // The main phase for the bbox union computation |
| |
| #version 450 |
| |
| #define LG_N_SEQ 0 |
| #define N_SEQ (1 << LG_N_SEQ) |
| #define LG_WG_SIZE 9 |
| #define WG_SIZE (1 << LG_WG_SIZE) |
| #define PART_SIZE (WG_SIZE * N_SEQ) |
| |
| layout(local_size_x = WG_SIZE, local_size_y = 1) in; |
| |
| // The bicyclic monoid |
| struct Bic { |
| uint a; |
| uint b; |
| }; |
| |
| struct BicBbox { |
| Bic bic; |
| uint pad2; |
| uint pad3; |
| vec4 bbox; |
| }; |
| |
| Bic bic_combine(Bic x, Bic y) { |
| uint m = min(x.b, y.a); |
| return Bic(x.a + y.a - m, x.b + y.b - m); |
| } |
| |
| // Node type values |
| #define OPEN_PAREN 0 |
| #define CLOSE_PAREN 1 |
| |
| struct Node { |
| uint node_type; |
| uint pad1; |
| uint pad2; |
| uint pad3; |
| vec4 bbox; |
| }; |
| |
| #define BBOX_IDENTITY vec4(1e9, 1e9, -1e9, -1e9) |
| |
| vec4 bbox_union(vec4 a, vec4 b) { |
| return vec4(min(a.xy, b.xy), max(a.zw, b.zw)); |
| } |
| |
| layout(binding = 0) readonly buffer InBuf { |
| Node[] inbuf; |
| }; |
| |
| layout(binding = 1) readonly buffer BicBuf { |
| BicBbox[] bicbuf; |
| }; |
| |
| layout(binding = 2) readonly buffer StackBuf { |
| vec4[] stack; |
| }; |
| |
| layout(binding = 3) buffer OutBuf { |
| vec4[] outbuf; |
| }; |
| |
| shared Bic sh_bic[WG_SIZE * 2 - 2]; |
| shared vec4 sh_stack[WG_SIZE]; |
| shared vec4 sh_bbox[WG_SIZE * 2 - 2]; |
| |
| void main() { |
| uint th = gl_LocalInvocationID.x; |
| // materialize stack snapshot as of the start of this partition |
| Bic bic = Bic(0, 0); |
| vec4 bbox = BBOX_IDENTITY; |
| if (th < gl_WorkGroupID.x) { |
| bic = bicbuf[th].bic; |
| bbox = bicbuf[th].bbox; |
| } |
| sh_bic[th] = bic; |
| sh_bbox[th] = bbox; |
| for (uint i = 0; i < LG_WG_SIZE; i++) { |
| barrier(); |
| uint other_ix = th + (1u << i); |
| if (other_ix < WG_SIZE) { |
| bic = bic_combine(bic, sh_bic[other_ix]); |
| bbox = bbox_union(bbox, sh_bbox[other_ix]); |
| } |
| barrier(); |
| sh_bic[th] = bic; |
| sh_bbox[th] = bbox; |
| } |
| |
| barrier(); |
| uint size = sh_bic[0].b; |
| // binary search in stack to do stream compaction |
| uint sp = PART_SIZE - 1 - th; |
| uint ix = 0; |
| for (uint i = 0; i < LG_WG_SIZE; i++) { |
| uint probe = ix + (uint(PART_SIZE / 2) >> i); |
| if (sp < sh_bic[probe].b) { |
| ix = probe; |
| } |
| } |
| // ix is the largest value such that sp < sh_bic[ix].b (if any) |
| uint b = sh_bic[ix].b; |
| if (sp < b) { |
| vec4 bbox = stack[ix * PART_SIZE + b - sp - 1]; |
| if (ix + 1 < WG_SIZE) { |
| bbox = bbox_union(bbox, sh_bbox[ix + 1]); |
| } |
| sh_stack[th] = bbox; |
| } |
| barrier(); |
| |
| // Do tree reduction of bicyclic semigroups and bounding boxes |
| Node inp = inbuf[gl_GlobalInvocationID.x]; |
| uint node_type = inp.node_type; |
| bic = Bic(uint(node_type == CLOSE_PAREN), uint(node_type == OPEN_PAREN)); |
| sh_bic[th] = bic; |
| sh_bbox[th] = inp.bbox; |
| uint inbase = 0; |
| for (uint i = 0; i < LG_WG_SIZE - 1; i++) { |
| uint outbase = 2 * WG_SIZE - (1u << (LG_WG_SIZE - i)); |
| barrier(); |
| if (th < (1u << (LG_WG_SIZE - 1 - i))) { |
| sh_bic[outbase + th] = bic_combine(sh_bic[inbase + th * 2], sh_bic[inbase + th * 2 + 1]); |
| sh_bbox[outbase + th] = bbox_union(sh_bbox[inbase + th * 2], sh_bbox[inbase + th * 2 + 1]); |
| } |
| inbase = outbase; |
| } |
| barrier(); |
| |
| // Search for predecessor node and aggregate bounding boxes |
| ix = th; |
| bbox = inp.bbox; |
| bic = Bic(0, 0); |
| if (node_type == CLOSE_PAREN) { |
| // only aggregate bbox for nodes with children; ascribe to close paren |
| // Note: if we also maintained the link to the open paren, we could |
| // store the bbox there. This will depend on the use case. |
| uint j = 0; |
| while (j < LG_WG_SIZE) { |
| uint base = 2 * WG_SIZE - (2u << (LG_WG_SIZE - j)); |
| if (((ix >> j) & 1) != 0) { |
| Bic test = bic_combine(sh_bic[base + (ix >> j) - 1], bic); |
| if (test.b > 0) { |
| break; |
| } |
| bic = test; |
| bbox = bbox_union(sh_bbox[base + (ix >> j) - 1], bbox); |
| ix -= 1u << j; |
| } |
| j++; |
| } |
| if (ix > 0) { |
| while (j > 0) { |
| j--; |
| uint base = 2 * WG_SIZE - (2u << (LG_WG_SIZE - j)); |
| Bic test = bic_combine(sh_bic[base + (ix >> j) - 1], bic); |
| if (test.b == 0) { |
| bic = test; |
| bbox = bbox_union(sh_bbox[base + (ix >> j) - 1], bbox); |
| ix -= 1u << j; |
| } |
| } |
| } |
| // ix is the smallest value such that reduce(ix..th).b == 0 |
| if (ix == 0 && bic.a < size) { |
| bbox = bbox_union(sh_stack[PART_SIZE - 1 - bic.a], bbox); |
| } |
| } |
| outbuf[gl_GlobalInvocationID.x] = bbox; |
| } |