knn — VPE-Accelerated KNN

knn/main.cpp holds the device-kernel implementation of the same search the XVector C API exposes as a single call, next to the host driver that dispatches it. Read it to see the machinery the C API hides — manual data partitioning, per-task dispatch, and top-k merge.

Each task reads its own slice of the dataset (via mu::getTaskIdx() / mu::getTaskCount()), scores it against the query, and writes its top-k into [taskIdx*k, taskIdx*k + k); the host merges the per-task results. The L2 kernel loads the target once, then for each dataset vector computes L2² = (v − target)·(v − target) with sub + dot, keeping a bounded max-heap of the k smallest distances:

__pxl_kernel__ void knnVpeL2(const float* data, const float* target,
                             uint32_t vectorCount, uint32_t dim, uint32_t k,
                             float* distResult, uint32_t* labelResult)
{
    uint32_t taskIdx   = mu::getTaskIdx();
    uint32_t taskCount = mu::getTaskCount();
    uint32_t perTask = vectorCount / taskCount;
    uint32_t start = taskIdx * perTask;
    uint32_t end   = (taskIdx + 1 == taskCount) ? vectorCount : start + perTask;  // last task absorbs remainder

    mu::vdma::VpeContext ctx(mu::vdma::VpeIdStrategy::ByThreadId);
    uint64_t bufs[3];
    // false only if the request could never fit an empty pool -- with a
    // runtime dimension that is worth checking; the output array is untouched.
    if (!ctx.waitAllocateBuffers(3, ctx.bytesForDimension(dim), bufs))
        return;
    uint64_t queryBuf = bufs[0], targetBuf = bufs[1], diffBuf = bufs[2];

    ctx.load(target, targetBuf, dim);   // load the query target once, reuse across the slice

    std::vector<std::pair<float, uint32_t>> topk;   // bounded max-heap: root = current worst
    topk.reserve(k);
    for (uint32_t v = start; v < end; ++v)
    {
        ctx.load(data + v * dim, queryBuf, dim);
        ctx.sub(queryBuf, targetBuf, diffBuf, dim);
        float dist = ctx.dot(diffBuf, diffBuf, dim);

        if (topk.size() < k) { topk.emplace_back(dist, v); std::push_heap(topk.begin(), topk.end()); }
        else if (dist < topk.front().first)
        {
            std::pop_heap(topk.begin(), topk.end());
            topk.back() = {dist, v};
            std::push_heap(topk.begin(), topk.end());
        }
    }
    std::sort_heap(topk.begin(), topk.end());   // ascending: smallest distance first
    // write topk -> distResult/labelResult at [taskIdx*k, ...)
}

A companion knnVpeDot kernel does the same with inner-product similarity — no diff buffer, and a min-heap keeping the k largest scores. The host partitions the dataset so the smallest task still has ≥ k points, dispatches both kernels, merges the per-task top-k, and verifies against a CPU reference.

./run.sh knn                                    # run with defaults
./run.sh knn -- --dim 128 --vectors 4096 --k 10 # forward args to the example
# or, from the build/ directory:
./bin/xarith_knn --dim 128 --vectors 4096 --k 10

Options:

Option	Default	Description
--dim N	128	Vector dimension
--vectors N	960000	Number of dataset vectors
--k N	5	Number of nearest neighbors

Environment Variables:

Variable	Default	Description
DEVICEID	0	Device ID
NUMSUB	(device)	Number of subsystems. Unset asks the device, which differs per part
TASKCOUNT	48	Number of subtasks for data partitioning (must be ≥ NUMSUB)