mixed_dim — Per-Call Dimension

VpeContext separates two units: buffers are allocated in bytes (use bytesForDimension() to size for an element count), while every load/store and op takes an explicit per-call dimension in elements. So one context can score the same buffers at several dimensions and hold differently sized buffers side by side. mixed_dim/main.cpp holds both kernels next to the host driver that launches them.

Multi-resolution dot — score one loaded pair at the full dimension and again over a leading prefix, with no extra load (a Matryoshka / truncated-embedding pattern):

__pxl_kernel__ void mixedDimMultiResolutionDot(const float* vec1, const float* vec2, float* results,
                                               uint64_t fullDim, uint64_t prefixDim)
{
    mu::vdma::VpeContext ctx{};

    uint64_t bufs[2];
    // 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(2, ctx.bytesForDimension(fullDim), bufs))
        return;
    ctx.load(vec1, bufs[0], fullDim);   // load the full vectors once
    ctx.load(vec2, bufs[1], fullDim);

    results[0] = ctx.dot(bufs[0], bufs[1], fullDim);    // full-dim dot
    results[1] = ctx.dot(bufs[0], bufs[1], prefixDim);  // coarse dot over the first prefixDim elements
}

A per-call dimension must be ≤ the capacity the buffer was allocated with, so here prefixDim <= fullDim.

Mixed-size buffers — the second kernel, mixedDimBuffers, reserves a long and a short buffer in one all-or-nothing call via the per-buffer sizes overload, then operates each at its own dimension:

const uint64_t sizes[2] = {ctx.bytesForDimension(longDim), ctx.bytesForDimension(shortDim)};
uint64_t bufs[2];
// 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.waitAllocateBuffersWithSizes(2, sizes, bufs))
    return;
// bufs[0] holds longDim elements, bufs[1] holds shortDim;
// each load/add/store below carries its own per-call dimension.
ctx.load(longVec, bufs[0], longDim);
ctx.add(bufs[0], bufs[0], bufs[0], longDim);   // longBuf = 2 * longVec
ctx.store(bufs[0], longOut, longDim);

This proves differently sized buffers coexist on the shared chunk allocator and that every operation honors its per-call size.

./run.sh mixed_dim      # or, from the build/ directory:
./bin/xarith_mixed_dim