From da9f2a89e828cac8af08b354b37d7a31829ae90d Mon Sep 17 00:00:00 2001 From: catbot Date: Thu, 18 Jun 2026 13:26:37 +0000 Subject: [PATCH] perf(ui): subgroup prefix-sum stream compaction, not lane-0 serial scan (#124) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit uiCompactChunk ran the survivor scan entirely on gl_LocalInvocationIndex == 0: ~64 serial iterations with 63/64 lanes idle, between two barriers, 4x per chunk in the fused kernel. It must emit a stable in-order (buffer-order) exclusive prefix of the survivors so the per-pixel inner loop still sees items in draw order. Replace it with a per-subgroup subgroupExclusiveAdd of the keep bits plus a carry across subgroups: each subgroup publishes its survivor total to shared memory, then every lane sums the totals of all lower-id subgroups for its base. A survivor's slot in s_order[] therefore equals the number of survivors with a smaller local index — buffer (draw) order preserved exactly, unlike an atomicAdd which would scramble it. The carry makes the result correct for any subgroup width (2 subgroups on the 32-wide descriptor_heap target, up to 8/16 on narrower parts), relying only on the gl_LocalInvocationIndex <-> (gl_SubgroupID, gl_SubgroupInvocationID) linear mapping every Vulkan compute implementation provides. The feature is free at the device baseline (apiVersion 1.4 + VK_EXT_descriptor_heap implies Vulkan 1.1 subgroup arithmetic), so no correctness fallback is needed; WebGPU is unaffected (separate embedded WGSL). Applied to ui-fused.comp.glsl and the same pattern inlined in ui-quads/circles/images/text. Verified: all 23 tests pass (UIFusedShader recompiles + spirv-val + pins push-constant offsets); a 140k-trial CPU simulation of the algorithm matches the serial reference for subgroup sizes 1..64; HelloUI renders correctly on an RTX 4090 (DispatchFused quads+circles+text) with draw order intact. Co-Authored-By: Claude Opus 4.8 --- shaders/ui-circles.comp.glsl | 38 ++++++++++++++++++++----- shaders/ui-fused.comp.glsl | 55 +++++++++++++++++++++++++++--------- shaders/ui-images.comp.glsl | 38 ++++++++++++++++++++----- shaders/ui-quads.comp.glsl | 39 +++++++++++++++++++------ shaders/ui-text.comp.glsl | 38 ++++++++++++++++++++----- 5 files changed, 166 insertions(+), 42 deletions(-) diff --git a/shaders/ui-circles.comp.glsl b/shaders/ui-circles.comp.glsl index eea0c93..65fa935 100644 --- a/shaders/ui-circles.comp.glsl +++ b/shaders/ui-circles.comp.glsl @@ -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; @@ -52,13 +82,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) { diff --git a/shaders/ui-fused.comp.glsl b/shaders/ui-fused.comp.glsl index 998061b..b8834a7 100644 --- a/shaders/ui-fused.comp.glsl +++ b/shaders/ui-fused.comp.glsl @@ -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) ──────────────────────────────────── @@ -73,17 +75,44 @@ shared uvec4 s_v4[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() { @@ -120,7 +149,7 @@ void main() { } s_keep[lid] = keep ? 1u : 0u; barrier(); - uiCompactChunk(base, count); + uiCompactChunk(); barrier(); if (inClip) { @@ -180,7 +209,7 @@ void main() { } s_keep[lid] = keep ? 1u : 0u; barrier(); - uiCompactChunk(base, count); + uiCompactChunk(); barrier(); if (inClip) { @@ -237,7 +266,7 @@ void main() { } s_keep[lid] = keep ? 1u : 0u; barrier(); - uiCompactChunk(base, count); + uiCompactChunk(); barrier(); if (inClip) { @@ -287,7 +316,7 @@ void main() { } s_keep[lid] = keep ? 1u : 0u; barrier(); - uiCompactChunk(base, count); + uiCompactChunk(); barrier(); if (inClip) { diff --git a/shaders/ui-images.comp.glsl b/shaders/ui-images.comp.glsl index a4424a6..afe76ff 100644 --- a/shaders/ui-images.comp.glsl +++ b/shaders/ui-images.comp.glsl @@ -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 @@ -19,6 +21,34 @@ shared uvec4 s_slots[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; @@ -49,13 +79,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) { diff --git a/shaders/ui-quads.comp.glsl b/shaders/ui-quads.comp.glsl index 24e7daf..a4fa58d 100644 --- a/shaders/ui-quads.comp.glsl +++ b/shaders/ui-quads.comp.glsl @@ -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. @@ -54,14 +84,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) { diff --git a/shaders/ui-text.comp.glsl b/shaders/ui-text.comp.glsl index a64abbf..9f9bab3 100644 --- a/shaders/ui-text.comp.glsl +++ b/shaders/ui-text.comp.glsl @@ -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 @@ -26,6 +28,34 @@ shared vec4 s_color[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; @@ -55,13 +85,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) { -- 2.54.0