Crafter.Graphics/shaders/ui-fused.comp.glsl

#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) ────────────────────────────────────
// One dispatch that composites up to four standard categories in canonical
// back-to-front order — quads → circles → images → text — into a single
// per-pixel register, loading the destination image ONCE and storing it ONCE.
// A run of consecutive Dispatch* calls would instead load+store the image per
// category and drain compute with a VkMemoryBarrier between each; this kernel
// collapses all of that to 1 load + 1 store + 0 inter-pass barriers.
//
// Each category keeps the exact cooperative shared-memory tile-cull + per-pixel
// accumulate of its standalone shader (see ui-quads/circles/images/text), so a
// fused category is pixel-identical to its standalone Dispatch* pass (the only
// difference is that intermediate results stay in full float precision between
// categories instead of round-tripping through the storage image — strictly
// more accurate). Categories run sequentially within each thread, so the VGPR
// high-water mark is ~max(per-category), not the sum.
//
// An absent category (itemCount 0) is a zero-trip, push-constant-uniform loop:
// no divergence, no memory traffic — using DispatchFused for only quads+text
// costs ~nothing for the unused circle/image phases.
//
// This is ADDITIVE: it has its own push-constant layout and does NOT touch the
// frozen 48-byte UIDispatchHeader or the per-element Dispatch* contract.

// Push-constant block. Mirrors Crafter::UIFusedHeader byte-for-byte: every
// member is vec4-aligned, no padding holes, exactly 128 bytes (the guaranteed
// push-constant minimum).
layout(push_constant) uniform PC {
    uvec4 itemBuffers;   // heap slots: (quads, circles, images, text)
    uvec4 itemCounts;    // item counts: (quads, circles, images, text)
    uvec4 misc;          // (outImage, fontTexture, fontSampler, flags)
    uvec4 surface;       // (surfaceWidth, surfaceHeight, frameIdx, _pad)
    vec4  clipQuads;
    vec4  clipCircles;
    vec4  clipImages;
    vec4  clipText;
} pc;

layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

// SDF tuning for the text phase — must match Crafter::FontAtlas constants
// (same values as ui-text.comp.glsl).
const float ON_EDGE    = 128.0 / 255.0;
const float DIST_SCALE = 32.0;

// Per-category clip-active bits packed into pc.misc.w (the fused analogue of
// ui-shared.glsl's UI_FLAG_CLIP). A bit is set only when that category's clip
// rect is narrower than the surface; when clear, the category skips its four
// per-pixel clip compares entirely (the common full-surface case). The branch
// is push-constant-uniform, so it never diverges.
const uint UI_FUSED_CLIP_QUADS   = 0x1u;
const uint UI_FUSED_CLIP_CIRCLES = 0x2u;
const uint UI_FUSED_CLIP_IMAGES  = 0x4u;
const uint UI_FUSED_CLIP_TEXT    = 0x8u;

// Generic per-chunk cooperative-cull scratch, REUSED across the four phases
// (they composite sequentially with a barrier between, so the storage is free
// once a phase's last read completes). Keeping one shared set instead of four
// keeps the shared-memory footprint — and thus occupancy — at the per-category
// level. Member mapping per phase:
//   quads:   v0=rect          v1=color v2=corners v3=outline
//   circles: v0=centerRadius  v1=color v2=outline
//   images:  v0=rect          v1=uv    v2=tint     v4=slots
//   text:    v0=rect          v1=uv    v2=color
shared vec4  s_v0[UI_CHUNK];
shared vec4  s_v1[UI_CHUNK];
shared vec4  s_v2[UI_CHUNK];
shared vec4  s_v3[UI_CHUNK];
shared uvec4 s_v4[UI_CHUNK];
// Per-item constants precomputed once at cooperative-load time, REUSED across
// phases like the s_v* scratch (the per-pixel inner loops read them instead of
// recomputing per pixel × item):
//   images: s_inv = 1.0/rect.zw
//   text:   s_inv = 1.0/rect.zw, s_band = SDF AA scale (a divide + two maxes)
shared vec2  s_inv[UI_CHUNK];
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. 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() {
    // NOTE: do not early-return — every thread must reach the barriers below.
    uvec2 pxu       = gl_GlobalInvocationID.xy;
    bool  inSurface = pxu.x < pc.surface.x && pxu.y < pc.surface.y;
    ivec2 screenPx  = ivec2(pxu);
    vec2  sp        = vec2(screenPx) + 0.5;

    // Single load of the destination for the whole surface.
    vec4 dst = vec4(0.0);
    if (inSurface) dst = imageLoad(uiImages[pc.misc.x], screenPx);

    vec2 tileMin, tileMax;
    uiTileBounds(tileMin, tileMax);
    uint lid = gl_LocalInvocationIndex;

    // ─── QUADS ────────────────────────────────────────────────────────────
    {
        uint heap    = pc.itemBuffers.x;
        uint count   = pc.itemCounts.x;
        bool inClip  = inSurface &&
            ((pc.misc.w & UI_FUSED_CLIP_QUADS) == 0u || uiPixelInClipRect(pxu, pc.clipQuads));
        for (uint base = 0u; base < count; base += UI_CHUNK) {
            uint idx  = base + lid;
            bool keep = false;
            if (idx < count) {
                s_v0[lid] = uiQuadHeap[heap].items[idx].rect;
                s_v1[lid] = uiQuadHeap[heap].items[idx].color;
                s_v2[lid] = uiQuadHeap[heap].items[idx].corners;
                s_v3[lid] = uiQuadHeap[heap].items[idx].outline;
                keep = uiAabbOverlapsTile(s_v0[lid].xy, s_v0[lid].xy + s_v0[lid].zw,
                                          tileMin, tileMax);
            }
            s_keep[lid] = keep ? 1u : 0u;
            barrier();
            uiCompactChunk();
            barrier();

            if (inClip) {
                for (uint j = 0u; j < s_count; ++j) {
                    uint c = s_order[j];

                    vec2 lo = s_v0[c].xy;
                    vec2 hi = s_v0[c].xy + s_v0[c].zw;
                    if (sp.x < lo.x || sp.y < lo.y) continue;
                    if (sp.x >= hi.x || sp.y >= hi.y) continue;

                    vec2  halfSize = s_v0[c].zw * 0.5;
                    vec2  p        = sp - (s_v0[c].xy + halfSize);
                    float d        = uiSdRoundRect(p, halfSize, s_v2[c]);

                    vec4  outline = s_v3[c];
                    float bodyA   = clamp(0.5 - d, 0.0, 1.0);
                    if (bodyA <= 0.0 && outline.x <= 0.0) continue;

                    vec4 col = s_v1[c];
                    vec4 src = vec4(col.rgb, col.a * bodyA);

                    if (outline.x > 0.0) {
                        float t = abs(d + outline.x * 0.5) - outline.x * 0.5;
                        float outlineA = clamp(0.5 - t, 0.0, 1.0);
                        src.rgb = mix(src.rgb, outline.yzw, outlineA);
                        src.a   = max(src.a, outlineA);
                    }

                    if (src.a <= 0.0) continue;
                    dst = uiBlendOver(dst, src);
                }
            }
            barrier(); // done reading shared for this chunk before reuse
        }
    }

    // ─── CIRCLES ──────────────────────────────────────────────────────────
    {
        uint heap   = pc.itemBuffers.y;
        uint count  = pc.itemCounts.y;
        bool inClip = inSurface &&
            ((pc.misc.w & UI_FUSED_CLIP_CIRCLES) == 0u || uiPixelInClipRect(pxu, pc.clipCircles));
        for (uint base = 0u; base < count; base += UI_CHUNK) {
            uint idx  = base + lid;
            bool keep = false;
            if (idx < count) {
                s_v0[lid] = uiCircleHeap[heap].items[idx].centerRadius;
                s_v1[lid] = uiCircleHeap[heap].items[idx].color;
                s_v2[lid] = uiCircleHeap[heap].items[idx].outline;
                float radius = s_v0[lid].z;
                if (radius > 0.0) {
                    vec2 cen = s_v0[lid].xy;
                    vec2 r   = vec2(radius + 1.0);
                    keep = uiAabbOverlapsTile(cen - r, cen + r, tileMin, tileMax);
                }
            }
            s_keep[lid] = keep ? 1u : 0u;
            barrier();
            uiCompactChunk();
            barrier();

            if (inClip) {
                for (uint j = 0u; j < s_count; ++j) {
                    uint c = s_order[j];

                    vec2  center = s_v0[c].xy;
                    float radius = s_v0[c].z;
                    if (radius <= 0.0) continue;

                    if (abs(sp.x - center.x) > radius + 1.0) continue;
                    if (abs(sp.y - center.y) > radius + 1.0) continue;

                    float d = length(sp - center) - radius;

                    vec4  outline = s_v2[c];
                    float bodyA   = clamp(0.5 - d, 0.0, 1.0);
                    if (bodyA <= 0.0 && outline.x <= 0.0) continue;

                    vec4 col = s_v1[c];
                    vec4 src = vec4(col.rgb, col.a * bodyA);

                    if (outline.x > 0.0) {
                        float t = abs(d + outline.x * 0.5) - outline.x * 0.5;
                        float outlineA = clamp(0.5 - t, 0.0, 1.0);
                        src.rgb = mix(src.rgb, outline.yzw, outlineA);
                        src.a   = max(src.a, outlineA);
                    }

                    if (src.a <= 0.0) continue;
                    dst = uiBlendOver(dst, src);
                }
            }
            barrier();
        }
    }

    // ─── IMAGES ───────────────────────────────────────────────────────────
    {
        uint heap   = pc.itemBuffers.z;
        uint count  = pc.itemCounts.z;
        bool inClip = inSurface &&
            ((pc.misc.w & UI_FUSED_CLIP_IMAGES) == 0u || uiPixelInClipRect(pxu, pc.clipImages));
        for (uint base = 0u; base < count; base += UI_CHUNK) {
            uint idx  = base + lid;
            bool keep = false;
            if (idx < count) {
                s_v0[lid] = uiImageHeap[heap].items[idx].rect;
                s_v1[lid] = uiImageHeap[heap].items[idx].uv;
                s_v2[lid] = uiImageHeap[heap].items[idx].tint;
                s_v4[lid] = uiImageHeap[heap].items[idx].slots;
                s_inv[lid] = 1.0 / s_v0[lid].zw;
                keep = uiAabbOverlapsTile(s_v0[lid].xy, s_v0[lid].xy + s_v0[lid].zw,
                                          tileMin, tileMax);
            }
            s_keep[lid] = keep ? 1u : 0u;
            barrier();
            uiCompactChunk();
            barrier();

            if (inClip) {
                for (uint j = 0u; j < s_count; ++j) {
                    uint c = s_order[j];

                    vec2 lo = s_v0[c].xy;
                    vec2 hi = s_v0[c].xy + s_v0[c].zw;
                    if (sp.x < lo.x || sp.y < lo.y) continue;
                    if (sp.x >= hi.x || sp.y >= hi.y) continue;

                    vec2 t  = (sp - s_v0[c].xy) * s_inv[c];
                    vec2 uv = mix(s_v1[c].xy, s_v1[c].zw, t);

                    uint texSlot  = s_v4[c].x;
                    uint sampSlot = s_v4[c].y;

                    vec4 sampled = texture(
                        sampler2D(uiTextures[nonuniformEXT(texSlot)],
                                  uiSamplers[nonuniformEXT(sampSlot)]),
                        uv
                    );
                    vec4 src = sampled * s_v2[c];
                    if (src.a <= 0.0) continue;
                    dst = uiBlendOver(dst, src);
                }
            }
            barrier();
        }
    }

    // ─── TEXT ─────────────────────────────────────────────────────────────
    {
        uint heap   = pc.itemBuffers.w;
        uint count  = pc.itemCounts.w;
        bool inClip = inSurface &&
            ((pc.misc.w & UI_FUSED_CLIP_TEXT) == 0u || uiPixelInClipRect(pxu, pc.clipText));
        for (uint base = 0u; base < count; base += UI_CHUNK) {
            uint idx  = base + lid;
            bool keep = false;
            if (idx < count) {
                s_v0[lid] = uiGlyphHeap[heap].items[idx].rect;
                s_v1[lid] = uiGlyphHeap[heap].items[idx].uv;
                s_v2[lid] = uiGlyphHeap[heap].items[idx].color;
                s_inv[lid] = 1.0 / s_v0[lid].zw;
                // SDF AA band — atlas-px per screen-px, constant per glyph
                // (uvSpan * kAtlasSize(1024) / screenSpan), max'd to 1px floor.
                vec2  uvSpan         = s_v1[lid].zw - s_v1[lid].xy;
                vec2  atlasPerScreen = (uvSpan * 1024.0) * s_inv[lid];
                float scalePx        = max(atlasPerScreen.x, atlasPerScreen.y);
                s_band[lid]          = max(scalePx, 0.0001);
                keep = uiAabbOverlapsTile(s_v0[lid].xy, s_v0[lid].xy + s_v0[lid].zw,
                                          tileMin, tileMax);
            }
            s_keep[lid] = keep ? 1u : 0u;
            barrier();
            uiCompactChunk();
            barrier();

            if (inClip) {
                for (uint j = 0u; j < s_count; ++j) {
                    uint c = s_order[j];

                    vec2 lo = s_v0[c].xy;
                    vec2 hi = s_v0[c].xy + s_v0[c].zw;
                    if (sp.x < lo.x || sp.y < lo.y) continue;
                    if (sp.x >= hi.x || sp.y >= hi.y) continue;

                    vec2 t  = (sp - s_v0[c].xy) * s_inv[c];
                    vec2 uv = mix(s_v1[c].xy, s_v1[c].zw, t);

                    // Font slots are push constants — provably dynamically uniform,
                    // so no nonuniformEXT (same as ui-text.comp.glsl).
                    float sdf = texture(
                        sampler2D(uiTextures[pc.misc.y],
                                  uiSamplers[pc.misc.z]),
                        uv
                    ).r;

                    float dAtlas = (ON_EDGE - sdf) * DIST_SCALE;

                    // band (the AA scale) was precomputed once per glyph at load.
                    float a = clamp(0.5 - dAtlas / s_band[c], 0.0, 1.0);
                    if (a <= 0.0) continue;

                    vec4 col = s_v2[c];
                    vec4 src = vec4(col.rgb, col.a * a);
                    dst = uiBlendOver(dst, src);
                }
            }
            barrier();
        }
    }

    if (inSurface) imageStore(uiImages[pc.misc.x], screenPx, dst);
}