vello/piet-gpu/shader/binning.comp

// The binning stage of the pipeline.
//
// Each workgroup processes N_TILE paths.
// Each thread processes one path and calculates a N_TILE_X x N_TILE_Y coverage mask
// based on the path bounding box to bin the paths.

#version 450
#extension GL_GOOGLE_include_directive : enable

#include "setup.h"

layout(local_size_x = N_TILE, local_size_y = 1) in;

layout(set = 0, binding = 0) buffer AnnotatedBuf {
    uint[] annotated;
};

layout(set = 0, binding = 1) buffer AllocBuf {
    uint n_elements; // paths
    // Will be incremented atomically to claim tiles
    uint tile_ix;
    uint alloc;
};

layout(set = 0, binding = 2) buffer BinsBuf {
    uint[] bins;
};

#include "annotated.h"
#include "bins.h"

// scale factors useful for converting coordinates to bins
#define SX (1.0 / float(N_TILE_X * TILE_WIDTH_PX))
#define SY (1.0 / float(N_TILE_Y * TILE_HEIGHT_PX))

#define TSY (1.0 / float(TILE_HEIGHT_PX))

// Constant not available in GLSL. Also consider uintBitsToFloat(0x7f800000)
#define INFINITY (1.0 / 0.0)

// Note: cudaraster has N_TILE + 1 to cut down on bank conflicts.
// Bitmaps are sliced (256bit into 8 (N_SLICE) 32bit submaps)
shared uint bitmaps[N_SLICE][N_TILE];
shared uint count[N_SLICE][N_TILE];
shared uint sh_chunk_start[N_TILE];

shared float sh_right_edge[N_TILE];

void main() {
    uint chunk_n = 0;
    uint my_n_elements = n_elements;
    uint my_partition = gl_WorkGroupID.x;

    for (uint i = 0; i < N_SLICE; i++) {
        bitmaps[i][gl_LocalInvocationID.x] = 0;
    }
    barrier();

    // Read inputs and determine coverage of bins
    uint element_ix = my_partition * N_TILE + gl_LocalInvocationID.x;
    AnnotatedRef ref = AnnotatedRef(element_ix * Annotated_size);
    uint tag = Annotated_Nop;
    if (element_ix < my_n_elements) {
        tag = Annotated_tag(ref);
    }
    int x0 = 0, y0 = 0, x1 = 0, y1 = 0;
    float my_right_edge = INFINITY;
    // bool crosses_edge = false;
    switch (tag) {
    // case Annotated_FillLine:
    // case Annotated_StrokeLine:
    //     AnnoStrokeLineSeg line = Annotated_StrokeLine_read(ref);
    //     x0 = int(floor((min(line.p0.x, line.p1.x) - line.stroke.x) * SX));
    //     y0 = int(floor((min(line.p0.y, line.p1.y) - line.stroke.y) * SY));
    //     x1 = int(ceil((max(line.p0.x, line.p1.x) + line.stroke.x) * SX));
    //     y1 = int(ceil((max(line.p0.y, line.p1.y) + line.stroke.y) * SY));
    //     crosses_edge = tag == Annotated_FillLine && ceil(line.p0.y * TSY) != ceil(line.p1.y * TSY);
    //     break;
    case Annotated_Fill:
    case Annotated_Stroke:
        // Note: we take advantage of the fact that fills and strokes
        // have compatible layout.
        AnnoFill fill = Annotated_Fill_read(ref);
        x0 = int(floor(fill.bbox.x * SX));
        y0 = int(floor(fill.bbox.y * SY));
        x1 = int(ceil(fill.bbox.z * SX));
        y1 = int(ceil(fill.bbox.w * SY));
        // It probably makes more sense to track x1, to avoid having to redo
        // the rounding to tile coords.
        my_right_edge = fill.bbox.z;
        break;
    }

    // At this point, we run an iterator over the coverage area,
    // trying to keep divergence low.
    // Right now, it's just a bbox, but we'll get finer with
    // segments.
    x0 = clamp(x0, 0, N_TILE_X);
    x1 = clamp(x1, x0, N_TILE_X);
    y0 = clamp(y0, 0, N_TILE_Y);
    y1 = clamp(y1, y0, N_TILE_Y);
    if (x0 == x1) y1 = y0;
    int x = x0, y = y0;
    uint my_slice = gl_LocalInvocationID.x / 32;
    uint my_mask = 1 << (gl_LocalInvocationID.x & 31);
    while (y < y1) {
        atomicOr(bitmaps[my_slice][y * N_TILE_X + x], my_mask);
        x++;
        if (x == x1) {
            x = x0;
            y++;
        }
    }

    barrier();
    // Allocate output segments.
    uint element_count = 0;
    for (uint i = 0; i < N_SLICE; i++) {
        element_count += bitCount(bitmaps[i][gl_LocalInvocationID.x]);
        count[i][gl_LocalInvocationID.x] = element_count;
    }
    // element_count is number of elements covering bin for this invocation.
    uint chunk_start = 0;
    if (element_count != 0) {
        // TODO: aggregate atomic adds (subgroup is probably fastest)
        chunk_start = atomicAdd(alloc, element_count * BinInstance_size);
        sh_chunk_start[gl_LocalInvocationID.x] = chunk_start;
    }
    // Note: it might be more efficient for reading to do this in the
    // other order (each bin is a contiguous sequence of partitions)
    uint out_ix = (my_partition * N_TILE + gl_LocalInvocationID.x) * 2;
    bins[out_ix] = element_count;
    bins[out_ix + 1] = chunk_start;

    barrier();
    // Use similar strategy as Laine & Karras paper; loop over bbox of bins
    // touched by this element
    x = x0;
    y = y0;
    while (y < y1) {
        uint bin_ix = y * N_TILE_X + x;
        uint out_mask = bitmaps[my_slice][bin_ix];
        if ((out_mask & my_mask) != 0) {
            uint idx = bitCount(out_mask & (my_mask - 1));
            if (my_slice > 0) {
                idx += count[my_slice - 1][bin_ix];
            }
            uint out_offset = sh_chunk_start[bin_ix] + idx * BinInstance_size;
            BinInstance_write(BinInstanceRef(out_offset), BinInstance(element_ix, my_right_edge));
        }
        x++;
        if (x == x1) {
            x = x0;
            y++;
        }
    }
}