wonfy-tools 0.2.0

// Scores every (anchor row, horizontal offset) candidate: sum of absolute
// differences between part1 rows [anchor, anchor + window) and part2's window
// rows, plus a penalty for non-overlapping margin pixels (their raw values),
// matching the CPU algorithm this replaced. Lower score = better match.
// Pixels flagged invalid (0x100, transparent composite padding) cost the
// maximum 255 inside the overlap so empty regions can't fake a match, but
// stay free in the margins where absent content is legitimate.
// Each workgroup reduces its 16x16 candidates to one (score, index) pair.

struct Params {
    width1: u32,
    width2: u32,
    window: u32,
    anchor_count: u32,
    offset_start: i32,
    offset_count: u32,
    _pad0: u32,
    _pad1: u32,
}

const SENTINEL: u32 = 0xffffffffu;
const INVALID: u32 = 0x100u;

fn margin_cost(raw: u32) -> u32 {
    return select(raw & 0xffu, 0u, (raw & INVALID) != 0u);
}

fn diff_cost(a: u32, b: u32) -> u32 {
    if ((a | b) & INVALID) != 0u {
        return 255u;
    }
    let va = a & 0xffu;
    let vb = b & 0xffu;
    return max(va, vb) - min(va, vb);
}

@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read> part1: array<u32>;
@group(0) @binding(2) var<storage, read> part2: array<u32>;
@group(0) @binding(3) var<storage, read_write> results: array<u32>;

var<workgroup> wg_score: array<u32, 256>;
var<workgroup> wg_index: array<u32, 256>;

fn candidate_score(anchor: u32, offset: i32) -> u32 {
    var score = 0u;

    for (var i = 0u; i < params.window; i++) {
        let base1 = (anchor + i) * params.width1;
        let base2 = i * params.width2;

        if offset >= 0 {
            let off = u32(offset);
            let margin1 = min(off, params.width1);
            let margin2 = min(off, params.width2);
            let overlap = min(params.width1 - margin1, params.width2);

            for (var x = 0u; x < margin1; x++) {
                score += margin_cost(part1[base1 + x]);
            }
            for (var x = params.width2 - margin2; x < params.width2; x++) {
                score += margin_cost(part2[base2 + x]);
            }
            for (var x = 0u; x < overlap; x++) {
                score += diff_cost(part1[base1 + off + x], part2[base2 + x]);
            }
        } else {
            let off = u32(-offset);
            let margin1 = min(off, params.width1);
            let margin2 = min(off, params.width2);
            let overlap = min(params.width1, params.width2 - margin2);

            for (var x = params.width1 - margin1; x < params.width1; x++) {
                score += margin_cost(part1[base1 + x]);
            }
            for (var x = 0u; x < margin2; x++) {
                score += margin_cost(part2[base2 + x]);
            }
            for (var x = 0u; x < overlap; x++) {
                score += diff_cost(part1[base1 + x], part2[base2 + off + x]);
            }
        }
    }

    return score;
}

@compute @workgroup_size(16, 16)
fn main(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_index) local_index: u32,
    @builtin(workgroup_id) wid: vec3<u32>,
    @builtin(num_workgroups) wg_count: vec3<u32>,
) {
    let offset_index = gid.x;
    let anchor_index = gid.y;

    var score = SENTINEL;
    var index = SENTINEL;

    if offset_index < params.offset_count && anchor_index < params.anchor_count {
        score = candidate_score(anchor_index, params.offset_start + i32(offset_index));
        index = anchor_index * params.offset_count + offset_index;
    }

    wg_score[local_index] = score;
    wg_index[local_index] = index;
    workgroupBarrier();

    for (var stride = 128u; stride > 0u; stride >>= 1u) {
        if local_index < stride {
            let other_score = wg_score[local_index + stride];
            let other_index = wg_index[local_index + stride];
            let mine_score = wg_score[local_index];
            let mine_index = wg_index[local_index];

            if other_score < mine_score || (other_score == mine_score && other_index < mine_index) {
                wg_score[local_index] = other_score;
                wg_index[local_index] = other_index;
            }
        }
        workgroupBarrier();
    }

    if local_index == 0u {
        let out = (wid.y * wg_count.x + wid.x) * 2u;
        results[out] = wg_score[0];
        results[out + 1u] = wg_index[0];
    }
}