Writing XArith Kernels

XArith provides VPE vector primitives — distance, dot product, and element-wise operations — for MU device kernels. This page covers the programming model for writing your own kernel. For ready-to-build sample projects, see XArith Example Projects.

Programming Model

A device kernel allocates SRAM buffers from a VpeContext, loads vectors into them from DRAM, runs VPE operations, and stores results back. Annotate the kernel with __pxl_kernel__ and place it in the same main.cpp as its host driver: pxcc++ splits the file, compiles the kernel for the device (RISC-V), and embeds the device blob into the host executable, which launches the kernel by symbol with pxl::Launcher().execute<kernel>(...). To build one, copy the maintained CMake setup from any example project — it drives the build through pxcc++ and forwards the xarith include/lib flags to the device chain via pxcc’s -Xmu= flag forwarding.

Buffer Allocation

VpeContext manages SRAM as a bitmap of fixed-size chunks. Allocation works in bytes: a buffer occupies enough contiguous chunks to cover its byte size, so a larger payload consumes more chunks and fewer buffers fit concurrently; using FP16 instead of FP32 halves the per-buffer payload and can free additional chunks. Use bytesForDimension(dim) to size a buffer for an element count, then waitAllocateBuffer(byteSize) for a single buffer or waitAllocateBuffersWithSizes(count, byteSizes, out) to reserve mixed-size buffers all-or-nothing. Every operation takes an explicit per-call dimension in elements (≤ the count the buffer was sized for), and the context tracks each buffer’s size so freeBuffer() needs no size argument. There is no count-query API — allocate and act on the result rather than checking availability up front.

When multiple buffers are needed, use waitAllocateBuffers(). It is atomic all-or-nothing, which prevents deadlock when multiple tasks compete for limited SRAM chunks, and it waits for room instead of failing:

uint64_t bufs[3];
if (!ctx.waitAllocateBuffers(3, ctx.bytesForDimension(dimension), bufs)) {
    return;  // dimension too large to ever fit this VPE's SRAM
}

// use bufs[0], bufs[1], bufs[2] ...

ctx.freeBuffers(3, bufs);  // release all at once

Do not write while (!ctx.tryAllocateBuffers(...)) {}. Each failed attempt issues allocator atomics, and on MX1 those are handled by the same L3 units the tasks holding the buffers need in order to run and free — so a flat-out retry loop slows down the very peers it is waiting on. waitAllocateBuffers() paces its retries instead.

The wait is unbounded, and deliberately so: waitAllocate* reports failure only for a request that could not fit even an empty pool (a zero byte size, or a rounded-up chunk total above the 80 KB per-VPE pool). A pool that is merely full right now is always a reason to wait.

That is why the check above matters whenever the size comes from a runtime value: on the failure path the output array is left untouched, so an unchecked call goes on to use whatever was on the stack as SRAM offsets. The single-buffer form reports the same condition as INVALID_BUFFER:

uint64_t buf = ctx.waitAllocateBuffer(ctx.bytesForDimension(dimension));
if (buf == mu::vdma::VpeContext::INVALID_BUFFER) {
    return;  // dimension too large to ever fit this VPE's SRAM
}

count == 0 is a vacuous success (returns true, nothing reserved), matching tryAllocateBuffers().

allocateBuffer(byteSize) and tryAllocateBuffers(...) remain available as the non-blocking forms, for a caller with other work to do while the pool is full.

For complete, runnable kernels built on these primitives — L2 distance, inner product, fused dot2, vector add/sub — see the XArith Example Projects, starting with basic.

Performance

  • Minimize DRAM transfers – Keep data in SRAM buffers as long as possible; batch operations before storing results back
  • Choose VpeIdStrategy – Use ByThreadId for parallelism within an MU, ByClusterId for parallelism across clusters

Next Steps