fidget-wgpu 0.5.0

// Interval evaluation stage for raymarching shader
//
// This must be combined with opcode definitions and the generic interpreter
// from `tape_interpreter.wgsl`

/// Per-state IO bindings
@group(2) @binding(0) var<storage, read> tiles_in: TileListInput;
@group(2) @binding(1) var<storage, read> tile_zmin: array<u32>;

@group(2) @binding(2) var<storage, read_write> subtiles_out: TileListOutput;
@group(2) @binding(3) var<storage, read_write> subtile_zmin: array<atomic<u32>>;
@group(2) @binding(4) var<storage, read_write> subtile_zhist: array<atomic<u32>>;

/// Input tile size; one input tile maps to a 4x4x4 workgroup
override TILE_SIZE: u32;

/// Output tile size, must be TILE_SIZE / 4; one output tile maps to one thread
override SUBTILE_SIZE: u32;

@compute @workgroup_size(4, 4, 4)
fn interval_tile_main(
    @builtin(workgroup_id) workgroup_id: vec3u,
    @builtin(num_workgroups) num_workgroups: vec3u,
    @builtin(local_invocation_id) local_id: vec3u
) {
    // We dispatch with workgroups only on the X axis
    for (var i=workgroup_id.x; i < tiles_in.count; i += num_workgroups.x) {
        interval_tile_worker(i, local_id);
    }
}

fn interval_tile_worker(
    active_tile_index: u32,
    local_id: vec3u
) {
    // Convert to a size in tile units
    let size64 = config.render_size / 64;
    let size_tiles = size64 * (64 / TILE_SIZE);
    let size_subtiles = size_tiles * 4u;

    // Get global tile position, in tile coordinates.  The top bit indicates
    // that the tile is filled.
    let t_raw = tiles_in.active_tiles[active_tile_index];
    let t_filled = (t_raw & (1 << 31u)) != 0;
    let t = t_raw & 0x7FFFFFFF;
    let tx = t % size_tiles.x;
    let ty = (t / size_tiles.x) % size_tiles.y;
    let tz = (t / (size_tiles.x * size_tiles.y)) % size_tiles.z;
    let tile_corner = vec3u(tx, ty, tz);

    // Subtile corner position
    let subtile_corner = tile_corner * 4 + local_id;
    let subtile_index_xy = subtile_corner.x + subtile_corner.y * size_subtiles.x;

    // Subtile corner position, in voxels
    let corner_pos = subtile_corner * SUBTILE_SIZE;

    // Special handling for uniformly filled tiles
    if t_filled {
        // Snap down to the larger tile size
        let z = (corner_pos.z / TILE_SIZE) * TILE_SIZE;
        atomicMax(&subtile_zmin[subtile_index_xy], z + TILE_SIZE - 1);
        return;
    }

    // Check for Z masking from tile
    let tile_index_xy = tile_corner.x + tile_corner.y * size_tiles.x;
    if tile_zmin[tile_index_xy] >= corner_pos.z + SUBTILE_SIZE {
        atomicMax(&subtile_zmin[subtile_index_xy], tile_zmin[tile_index_xy]);
        return;
    }

    // Last-minute check to see if anyone filled out this tile
    if atomicLoad(&subtile_zmin[subtile_index_xy]) >= corner_pos.z + SUBTILE_SIZE {
        return;
    }

    // Compute transformed interval regions
    let m = interval_inputs(subtile_corner, SUBTILE_SIZE);

    // Do the actual interpreter work
    var stack = Stack();
    let tape_offset = get_tape_offset_for_level(corner_pos, TILE_SIZE);
    var tape_start = tile_tape[tape_offset];
    let out = run_tape(tape_start, m, &stack);

    let v = out.value.v;
    if v[1] < 0.0 {
        // Full, write to subtile_zmin but don't return yet (because we want to
        // store a simplified tape for normal evaluation)
        atomicMax(&subtile_zmin[subtile_index_xy], corner_pos.z + SUBTILE_SIZE - 1);
    } else if v[0] > 0.0 {
        return;
        // Empty, nothing to do here
    } else {
        // Push this subtile to the output list
        let offset = atomicAdd(&subtiles_out.count, 1u);
        let subtile_index_xyz = subtile_corner.x +
            (subtile_corner.y * size_subtiles.x) +
            (subtile_corner.z * size_subtiles.x * size_subtiles.y);
        subtiles_out.active_tiles[offset] = subtile_index_xyz;

        // Update the indirect dispatch count.  We'll divide by 64 here
        // (rounding up) because each thread in the sorting pass handles a
        // single tile, and we dispatch that pass with [64, 1, 1] workgroups
        let count = ((offset + 1u) + 63u) / 64u;
        let wg_dispatch_x = min(count, 32768u);
        let wg_dispatch_y = (count + 32767u) / 32768u;
        atomicMax(&subtiles_out.wg_size[0], wg_dispatch_x);
        atomicMax(&subtiles_out.wg_size[1], wg_dispatch_y);
        atomicMax(&subtiles_out.wg_size[2], 1u);

        // Update the cumsum histogram.  This is a little inefficient, because
        // we're building the cumsum by iterating over the prefix (instead of
        // doing a proper prefix sum), but the worst-case is iterating over 16
        // values.
        let stz = (corner_pos.z % 64u) / SUBTILE_SIZE;
        for (var i=0u; i <= stz; i += 1u) {
            atomicAdd(&subtile_zhist[i], 1u);
        }
    }

    let next = simplify_tape(out.pos, out.count, &stack);
    if next != 0 {
        tape_start = next;
    }
    let next_tape_offset = get_tape_offset_for_level(corner_pos, SUBTILE_SIZE);
    tile_tape[next_tape_offset] = tape_start;
}

/// Allocates a new chunk, returning a past-the-end pointer
fn alloc(chunk_size: u32) -> u32 {
    return atomicAdd(&config.tape_data_offset, chunk_size);
}

fn dealloc(chunk_size: u32) {
    atomicSub(&config.tape_data_offset, chunk_size);
}