OpenNPU v1.0: an open compiler and runtime for the RK3588 NPU

OpenNPU is an open-source compiler and runtime for the Rockchip RK3588’s Neural Processing Unit. It compiles models from ONNX, PyTorch, or JAX down to the NPU’s register commands and runs them through the public DRM ioctl interface. There is no closed toolkit and no librknnrt in the running path. The code is MIT licensed at github.com/poad42/opennpu_rk3588.

This is the technical background. It covers the hardware contract, the wire format, the compiler pipeline, and the runtime.

The hardware contract

The RK3588 NPU is a 6 TOPS device with three cores. Rockchip exposes it through a DRM render node at /dev/dri/card1. Six ioctls cover the whole accelerator:

Three properties of this hardware dominate every design decision.

The NPU has no on-chip SRAM. The RK3576 NPU has on-die SRAM; the RK3588 does not. Every operation moves data DRAM to NPU to DRAM. That makes the accelerator DMA-bound. Measured throughput on the matmul path is about 1 GB/s, and the model’s fastest full-context forward sits at that wall.

Buffer addressing is a per-file-descriptor IOMMU pool. Each process gets its own virtual address window into the NPU’s physical memory. The runtime allocates buffers against that window, and some paths are position-sensitive: the scratch buffer must sit at the top of the pool, and a submit is rejected when scratch falls below a fixed address. One consequence is that the hardware runs one large matmul scratch at a time; the 3.5 MB scratch window fits one matmul’s working set, not three concurrent ones. The three cores are therefore data-parallel across independent requests, not model-parallel within a single matmul.

Third, the native numeric types are int8 and float16. There is no fp4. Every op and every weight layout has to fit one of those two.

The wire format

The NPU is programmed by writing register-command entries. Each entry is eight bytes: a 16-bit register, a 16-bit value, and a 32-bit tag.

One 8-byte entry, a 640-byte command block, and the two structs that drive a submit.

A 32-bit DMA address does not fit the 16-bit value field, so it is split across two fields. The decode is dma = (tag32 & 0xffff) << 16 | val16. A command block holds 80 entries (640 bytes). Of those, 69 are real register writes and one is the chain pointer. Entry 69 holds next_offset; a zero means the chain ends. The NPU follows the chain and stops at the end, so the runtime must submit exactly the number of tasks the chain forms. A job that claims more tasks than the chain runs makes the kernel wait out the full count.

Two structs drive a submit. The task struct is 40 bytes and carries flags, op_idx, enable_mask, int_mask, int_clear, int_status, regcfg_amount, regcfg_offset, regcmd_addr. The submit struct is 104 bytes and carries the task-object pointer, the IOMMU domain id, and the per-core configuration. The three cores are addressed as 0x1001, 0x2001, and 0x3001.

The compiler pipeline

The compiler in src/opennpu/ is a pure-Python MLIR-style flow. It does not link LLVM or MLIR; it is a small dataclass hierarchy that emits MLIR text and lowers it. That keeps the whole lowering readable in a few hundred lines and makes the project genuinely dependency-free.

From a model graph to the NPU: importer, IR, two lowering passes, register-command bytes, and the DRM runtime.

The path is:

ONNX importer -> IR -> TosaToNPULowering -> NPUToRegCmd -> runtime

• The importer reads an ONNX graph and builds the IR.
• TosaToNPULowering maps TOSA ops to rk3588_npu ops. The dialect defines over thirty op types: elementwise add, subtract, multiply, divide, relu, sigmoid, tanh, gelu, silu, exp, sqrt; linear matmul, gemm, batch matmul; conv2d, depthwise, transposed conv; pooling; norm; fused conv-relu and add-relu patterns; reshape, transpose, concat, slice, pad; reduce; and softmax and SDPA.
• NPUToRegCmd lowers each op to the register-command bytes through a codegen pass. It carries the hardware constants and the DMA encoding from the section above.
• The runtime allocates buffers, patches the addresses, and issues the submit.

Three codegen strategies exist and coexist. The capture-based path keeps captured reference templates and patches their DMA addresses at the known entry points. The pure-synthesize path generates the command bytes from the reversed constants, with no reference file at all. The MLIR-style path drives the same synthesizer through the dialect. All three produce the same regcmd bytes; the pure-synthesize and MLIR paths have no dependency on the closed toolkit.

The supported operator set was not assumed. It was taken from Rockchip’s RKNN Compiler Support Operator List (v1.5.2), which lists int8 and float16 as the native types.

Matmul through the convolution engine

The NPU has no general GEMM primitive. The compiler implements Y = X @ W as a one-by-one “direct convolution,” which the hardware calls a CNA descriptor.

Y = X @ W runs as a CNA descriptor. A CNA matmul uses five buffers.

A CNA matmul uses five buffers: weights, input, output, the task array, and the register-command array. Two constraints matter. When the inner dimension exceeds 768, the weights must be pre-tiled to fit the descriptor. And the descriptor supports int8, int16, and fp16; there is no fp4.

The matmul runtime in C, cna_matmul.c, keeps weights in a persistent buffer and streams only the input on each call. That cached-weight path is one half of the single-token speed. The KV cache is the other half: the decoder reads one token’s worth of memory instead of reprocessing the full context. Together they put single-token decode at 28 ms per token. The same file exposes attention on the NPU by loading K and V into descriptor slots, and layer norm and an erf approximation for GELU. It also exposes the three-core path through a core mask.

The framework front ends

Three front ends reach the same runtime.

ONNX uses onnx_runner.py, which walks the ONNX graph and dispatches each node to either the NPU or the CPU. PyTorch uses torch_npu.py, which presents the NPU as a device. JAX uses a real PJRT C plugin.

The JAX plugin, libpjrt_npu.so, is the most involved of the three. It is pinned to PJRT API version 0.112 and jaxlib 0.10.2. The pin matters: jaxlib reads the plugin’s function table at the offsets of the exact version it was compiled against, so the plugin must export that table layout. The header for 0.112 comes from the pinned XLA commit, and with it jax.devices() reports the NPU as npu:0.

The plugin does not link the MLIR/XLA compiler. Instead it parses the raw stablehlo bytecode that JAX sends to Client_Compile, using a small varint reader. It accepts a strict linear chain of recognized ops: dot_general becomes a matmul, tanh and maximum become the corresponding activation, and an add or multiply chain becomes the elementwise op. The fast matmul path is also exposed as a custom_call that bypasses HLO parsing entirely. That is the path examples/jax_matmul.py and the 28 ms-per-token decoder use.

Buffer lifecycle uses the same raw-DRM path as the Python runtime. The plugin opens the device and pre-allocates the weight cache before the first host-to-device copy, so the buffers sit at the top of the address window and one open file can be reused. Host to NPU copy goes through BufferFromHostBuffer and Buffer_ToHostBuffer.

The C kernels and the hybrid split

The compute-heavy work runs as C kernels compiled with unrolled loops. Two forwards exist.

Matmuls go to the NPU; attention, activations, and layer norm run on the CPU.

lm_forward.c runs GPT-2 and LLaMA-family models from one file, switched by flags for RoPE, SwiGLU, and RMSNorm. All matmuls go to the NPU through the cached CNA path. Attention, the activations, and layer norm run on OpenMP threads over the CPU’s A76 cores. vision_forward.c is the same idea for the SigLIP ViT: all 48 matmuls on the NPU, attention, GELU, and layer norm on the CPU.

The split is deliberate and measured, not ideological. Layer norm stays on the CPU because a submit costs about 0.5 ms and the CPU computes it in about 0.05 ms for a 768-wide activation. The NPU is memory-bound, so small operations do not amortize the cost of a submit. The compiler routes by cost.

Correctness and limits

The claims above are verified on an Orange Pi 5 Max (RK3588, Armbian, kernel 6.1.115-vendor-rk35xx).

The SigLIP encoder runs all 48 matmuls on the NPU. Its output matches the HuggingFace reference to cosine similarity 0.9999. The GPT-2 decoder output matches PyTorch token for token. The NPU runs at 98 percent load on the vision encoder. Forty of the forty-one op and precision combinations pass on board, covering the elementwise ops, the activations, matmul, softmax, concat, transpose, and layer norm across fp16, w8a8, and w16a16i.

The honest limits are the hardware’s. The NPU is DMA-bound around 1 GB/s. One matmul scratch fits the address window, so true model-parallel across the three cores is not possible. Prompt caching is off for the hybrid attention architecture. These bounds are measured, and they are the reason the runtime is a hybrid that puts the memory-bound work on the CPU.

How to use it and extend it

The repo runs a set of examples.

pip install -e .
python examples/gpt2_generate.py
python examples/vision_encoder.py --framework torch
python examples/vision_encoder.py --framework jax

Two documented routes exist for a bespoke op. For a matmul-shaped op, use the CNA descriptor path documented in the CNA guide. For an elementwise or activation op, generate a C template and let the plugin dispatch to it. Either op is callable from both PyTorch and JAX once the PJRT client is present.

Where the approach comes from

The direct-ioctl execution and the CNA descriptor matmul build on prior work. mtx512/rk3588-npu, a GPL v3 project by Jasbir Matharu, first demonstrated the descriptor approach in early 2024 on the older 5.10 kernel (write-up). Their approach differed: they derived the CNA and DPU register descriptors from the chip’s technical reference manual and ran matmul on the NPU’s internal SRAM, which avoids the DRAM roundtrip. The register and submit formats changed in the newer 6.1 driver, so OpenNPU re-derived the descriptor on this kernel rather than porting it; the older code does not run correctly on 6.1. The CNA descriptor matmul is one part of the reverse engineering, not the whole of it. The larger part was capturing the register-command templates and decoding the wire format, which the compiler and the op coverage rest on. OpenNPU credits that project and reimplements the descriptor independently under MIT. The MIT userland is separate from the GPL kernel driver, so the two do not mix.