Merge pull request 'perf(ui): subgroup prefix-sum stream compaction in the UI compute shaders (#124)' (#145) from claude/issue-124 into master

This commit is contained in:
catbot 2026-06-18 15:27:30 +02:00
commit 9046562a90
5 changed files with 167 additions and 43 deletions

View file

@ -1,5 +1,7 @@
#version 460
#extension GL_GOOGLE_include_directive : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#include "ui-shared.glsl"
// One workgroup per 8×8 screen tile. The workgroup cooperatively streams the
@ -18,6 +20,34 @@ shared vec4 s_outline[UI_CHUNK];
shared uint s_keep[UI_CHUNK];
shared uint s_order[UI_CHUNK];
shared uint s_count;
shared uint s_subTotals[UI_CHUNK]; // per-subgroup survivor totals (carry step)
// Stable in-order compaction of one chunk's survivors via a per-subgroup
// exclusive prefix-sum plus a carry across subgroups, so a survivor's slot in
// s_order[] equals the number of survivors with a smaller local index — buffer
// (draw) order preserved. Replaces the old lane-0 serial scan. See
// ui-fused.comp.glsl for the full rationale; the body is identical.
void uiCompactChunk() {
uint lid = gl_LocalInvocationIndex;
uint keep = s_keep[lid];
uint subPrefix = subgroupExclusiveAdd(keep);
uint subTotal = subgroupAdd(keep);
if (subgroupElect())
s_subTotals[gl_SubgroupID] = subTotal;
barrier();
uint base = 0u;
uint total = 0u;
for (uint s = 0u; s < gl_NumSubgroups; ++s) {
uint t = s_subTotals[s];
if (s < gl_SubgroupID) base += t;
total += t;
}
if (lid == 0u) s_count = total;
if (keep != 0u)
s_order[base + subPrefix] = lid;
}
void main() {
ivec2 screenPx;
@ -56,13 +86,7 @@ void main() {
s_keep[lid] = keep ? 1u : 0u;
barrier();
if (lid == 0u) {
uint n = 0u;
uint lim = min(UI_CHUNK, pc.hdr.itemCount - base);
for (uint k = 0u; k < lim; ++k)
if (s_keep[k] != 0u) s_order[n++] = k;
s_count = n;
}
uiCompactChunk();
barrier();
if (valid) {

View file

@ -1,5 +1,7 @@
#version 460
#extension GL_GOOGLE_include_directive : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#include "ui-shared.glsl"
// ─── fused UI uber-kernel (issue #47) ────────────────────────────────────
@ -80,17 +82,44 @@ shared float s_band[UI_CHUNK];
shared uint s_keep[UI_CHUNK];
shared uint s_order[UI_CHUNK];
shared uint s_count;
// Per-subgroup survivor totals, published once per subgroup for the carry step
// of the compaction below. Sized for the worst case of one subgroup per lane.
shared uint s_subTotals[UI_CHUNK];
// Stable in-order compaction of one chunk's survivors. Lane 0 scans s_keep in
// buffer order so the inner per-pixel loop still sees items in draw order.
void uiCompactChunk(uint base, uint count) {
if (gl_LocalInvocationIndex == 0u) {
uint n = 0u;
uint lim = min(UI_CHUNK, count - base);
for (uint k = 0u; k < lim; ++k)
if (s_keep[k] != 0u) s_order[n++] = k;
s_count = n;
// Stable in-order compaction of one chunk's survivors. Each subgroup computes
// an exclusive prefix-sum of its keep bits at register speed (subgroupExclusiveAdd),
// publishes its survivor total to shared memory, then every lane carries in the
// totals of all lower-id subgroups. A survivor's slot in s_order[] therefore
// equals the number of survivors with a smaller local index, so the per-pixel
// inner loop still sees items in buffer (draw) order. This replaces the old
// lane-0 serial scan (~64 iterations with 63/64 lanes idle, run 4× per chunk).
//
// Order preservation relies on the gl_LocalInvocationIndex ↔ (gl_SubgroupID,
// gl_SubgroupInvocationID) linear mapping that every Vulkan compute
// implementation provides; the carry across subgroups makes the result correct
// for any subgroup width (2 subgroups on the 32-wide descriptor_heap target, up
// to 8 on 8-wide parts). Out-of-range lanes carry s_keep==0, so they neither
// add to a total nor write a slot — no separate bound is needed.
void uiCompactChunk() {
uint lid = gl_LocalInvocationIndex;
uint keep = s_keep[lid];
uint subPrefix = subgroupExclusiveAdd(keep); // rank within this subgroup
uint subTotal = subgroupAdd(keep); // survivors in this subgroup
if (subgroupElect())
s_subTotals[gl_SubgroupID] = subTotal;
barrier();
uint base = 0u;
uint total = 0u;
for (uint s = 0u; s < gl_NumSubgroups; ++s) {
uint t = s_subTotals[s];
if (s < gl_SubgroupID) base += t; // survivors in lower subgroups
total += t;
}
if (lid == 0u) s_count = total;
if (keep != 0u)
s_order[base + subPrefix] = lid;
}
void main() {
@ -127,7 +156,7 @@ void main() {
}
s_keep[lid] = keep ? 1u : 0u;
barrier();
uiCompactChunk(base, count);
uiCompactChunk();
barrier();
if (inClip) {
@ -187,7 +216,7 @@ void main() {
}
s_keep[lid] = keep ? 1u : 0u;
barrier();
uiCompactChunk(base, count);
uiCompactChunk();
barrier();
if (inClip) {
@ -245,7 +274,7 @@ void main() {
}
s_keep[lid] = keep ? 1u : 0u;
barrier();
uiCompactChunk(base, count);
uiCompactChunk();
barrier();
if (inClip) {
@ -302,7 +331,7 @@ void main() {
}
s_keep[lid] = keep ? 1u : 0u;
barrier();
uiCompactChunk(base, count);
uiCompactChunk();
barrier();
if (inClip) {

View file

@ -1,5 +1,7 @@
#version 460
#extension GL_GOOGLE_include_directive : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#include "ui-shared.glsl"
// One workgroup per 8×8 screen tile. The workgroup cooperatively streams the
@ -24,6 +26,34 @@ shared vec2 s_invRectSize[UI_CHUNK];
shared uint s_keep[UI_CHUNK];
shared uint s_order[UI_CHUNK];
shared uint s_count;
shared uint s_subTotals[UI_CHUNK]; // per-subgroup survivor totals (carry step)
// Stable in-order compaction of one chunk's survivors via a per-subgroup
// exclusive prefix-sum plus a carry across subgroups, so a survivor's slot in
// s_order[] equals the number of survivors with a smaller local index — buffer
// (draw) order preserved. Replaces the old lane-0 serial scan. See
// ui-fused.comp.glsl for the full rationale; the body is identical.
void uiCompactChunk() {
uint lid = gl_LocalInvocationIndex;
uint keep = s_keep[lid];
uint subPrefix = subgroupExclusiveAdd(keep);
uint subTotal = subgroupAdd(keep);
if (subgroupElect())
s_subTotals[gl_SubgroupID] = subTotal;
barrier();
uint base = 0u;
uint total = 0u;
for (uint s = 0u; s < gl_NumSubgroups; ++s) {
uint t = s_subTotals[s];
if (s < gl_SubgroupID) base += t;
total += t;
}
if (lid == 0u) s_count = total;
if (keep != 0u)
s_order[base + subPrefix] = lid;
}
void main() {
ivec2 screenPx;
@ -59,13 +89,7 @@ void main() {
s_keep[lid] = keep ? 1u : 0u;
barrier();
if (lid == 0u) {
uint n = 0u;
uint lim = min(UI_CHUNK, pc.hdr.itemCount - base);
for (uint k = 0u; k < lim; ++k)
if (s_keep[k] != 0u) s_order[n++] = k;
s_count = n;
}
uiCompactChunk();
barrier();
if (valid) {

View file

@ -1,5 +1,7 @@
#version 460
#extension GL_GOOGLE_include_directive : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#include "ui-shared.glsl"
// One workgroup per 8×8 screen tile. The workgroup cooperatively streams the
@ -21,6 +23,34 @@ shared vec4 s_outline[UI_CHUNK];
shared uint s_keep[UI_CHUNK];
shared uint s_order[UI_CHUNK];
shared uint s_count;
shared uint s_subTotals[UI_CHUNK]; // per-subgroup survivor totals (carry step)
// Stable in-order compaction of one chunk's survivors via a per-subgroup
// exclusive prefix-sum plus a carry across subgroups, so a survivor's slot in
// s_order[] equals the number of survivors with a smaller local index — buffer
// (draw) order preserved. Replaces the old lane-0 serial scan. See
// ui-fused.comp.glsl for the full rationale; the body is identical.
void uiCompactChunk() {
uint lid = gl_LocalInvocationIndex;
uint keep = s_keep[lid];
uint subPrefix = subgroupExclusiveAdd(keep);
uint subTotal = subgroupAdd(keep);
if (subgroupElect())
s_subTotals[gl_SubgroupID] = subTotal;
barrier();
uint base = 0u;
uint total = 0u;
for (uint s = 0u; s < gl_NumSubgroups; ++s) {
uint t = s_subTotals[s];
if (s < gl_SubgroupID) base += t;
total += t;
}
if (lid == 0u) s_count = total;
if (keep != 0u)
s_order[base + subPrefix] = lid;
}
void main() {
// NOTE: do not early-return — every thread must reach the barriers below.
@ -58,14 +88,7 @@ void main() {
s_keep[lid] = keep ? 1u : 0u;
barrier();
// Stable in-order compaction of this chunk's survivors.
if (lid == 0u) {
uint n = 0u;
uint lim = min(UI_CHUNK, pc.hdr.itemCount - base);
for (uint k = 0u; k < lim; ++k)
if (s_keep[k] != 0u) s_order[n++] = k;
s_count = n;
}
uiCompactChunk();
barrier();
if (valid) {

View file

@ -1,5 +1,7 @@
#version 460
#extension GL_GOOGLE_include_directive : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#include "ui-shared.glsl"
// One workgroup per 8×8 screen tile. The workgroup cooperatively streams the
@ -33,6 +35,34 @@ shared float s_band[UI_CHUNK];
shared uint s_keep[UI_CHUNK];
shared uint s_order[UI_CHUNK];
shared uint s_count;
shared uint s_subTotals[UI_CHUNK]; // per-subgroup survivor totals (carry step)
// Stable in-order compaction of one chunk's survivors via a per-subgroup
// exclusive prefix-sum plus a carry across subgroups, so a survivor's slot in
// s_order[] equals the number of survivors with a smaller local index — buffer
// (draw) order preserved. Replaces the old lane-0 serial scan. See
// ui-fused.comp.glsl for the full rationale; the body is identical.
void uiCompactChunk() {
uint lid = gl_LocalInvocationIndex;
uint keep = s_keep[lid];
uint subPrefix = subgroupExclusiveAdd(keep);
uint subTotal = subgroupAdd(keep);
if (subgroupElect())
s_subTotals[gl_SubgroupID] = subTotal;
barrier();
uint base = 0u;
uint total = 0u;
for (uint s = 0u; s < gl_NumSubgroups; ++s) {
uint t = s_subTotals[s];
if (s < gl_SubgroupID) base += t;
total += t;
}
if (lid == 0u) s_count = total;
if (keep != 0u)
s_order[base + subPrefix] = lid;
}
void main() {
ivec2 screenPx;
@ -75,13 +105,7 @@ void main() {
s_keep[lid] = keep ? 1u : 0u;
barrier();
if (lid == 0u) {
uint n = 0u;
uint lim = min(UI_CHUNK, pc.hdr.itemCount - base);
for (uint k = 0u; k < lim; ++k)
if (s_keep[k] != 0u) s_order[n++] = k;
s_count = n;
}
uiCompactChunk();
barrier();
if (valid) {