Architecture

Architecture

VKNN loads an ONNX model, lowers it to a backend-agnostic NCHW IR, optimizes it with graph passes, partitions it into backend-specific segments, and executes those segments — primarily on a Vulkan compute backend tuned for a mobile GPU, with a CPU backend that doubles as the numeric reference and the fallback. Every executable operator has a Vulkan kernel; only data-dependent control flow (Loop / If / NonMaxSuppression) and const-folded import-time shape ops stay off the GPU.

End-to-end pipeline

model.onnx
   -> IMPORT            hand-rolled protobuf parser (src/import/onnx/onnx_parser.cpp)
                        importOnnx(path) -> Graph (backend-agnostic NCHW IR)
   -> GRAPH PASSES      runStandardPasses(g): dequantizeGraph, inferShapes,
                        foldBatchNorm/lowerBatchNorm, constFold+inferShapes to convergence,
                        subpixelConvTranspose, lowerConv, eliminateDeadNodes,
                        fusePointwiseChains, pruneDeadInitializers
   -> SESSION::plan()   instantiate backends in priority order, allocate the RtTensor pool,
                        assign each node to the first backend that supports it, partition into
                        maximal same-backend segments, compileSegment() each, finalize() caches
   -> RUN               bind inputs -> run segments in order -> reconcile residency at boundaries
                        -> collect graph outputs (host, NCHW, fp32)

The canonical IR layout is NCHW throughout. Only the Vulkan backend re-packs into its internal layouts; the rest of the engine never sees them.

The two backends

Backend	Role	Notes
VulkanBackend	Primary	NC4HW4 packed layout, one pre-recorded command buffer per static segment, push descriptors, timestamp queries, fp16 storage / fp32 accumulate, on-device autotuning.
CpuBackend	Reference + fallback	Scalar reference kernels for every op, NEON kernels for Add and Gemm; operates directly on host NCHW. Cosine 1.000000 vs onnxruntime.

The CPU oracle

The CPU backend is not just a fallback — it is the numeric ground truth. Its host buffers are canonical NCHW fp32, and it produces cosine 1.000000 against onnxruntime. Registering a CPU op and registering a Vulkan op are independent, so the CPU op is always available as the oracle for A/B validation: forcing any op back to CPU at runtime (Config::disableVkOps = "Foo") re-partitions the graph and lets a new Vulkan kernel be diffed byte-for-byte against its reference without touching code.

The ONNX importer and standard passes

The importer is a dependency-free protobuf parser (src/import/onnx/onnx_parser.cpp) that lowers ONNX to the Graph IR. The graph passes live under src/import/, one .cpp per pass, declared in passes.h and ordered by runStandardPasses():

• dequantizeGraph — QDQ / QLinear / dynamic-quant clusters lowered to float (saturation clamps preserved; default on, --no-dequantize opts out).
• inferShapes — resolves dynamic dims from declared shapes; a dynamic batch axis defaults to 1, and a dynamic non-batch axis with no declaration is a hard error (never a silent 1x1 plan).
• foldBatchNorm / lowerBatchNorm — BN folded into the preceding Conv, or lowered to a per-channel Mul+Add for the pointwise fusion to absorb.
• constFold + inferShapes to convergence — resolves shape-path subgraphs (Shape → Gather → Reshape head chains) one block per round.
• subpixelConvTranspose / lowerConv — geometry rewrites (ConvTranspose → Conv + DepthToSpace; K×K Conv → an implicit-GEMM ConvGemm, opt-in).
• fusePointwiseChains — the one general region-growing fusion pass: grows each maximal same-shape per-element region and emits it as a single fused unit, folded into the producing kernel's store epilogue or one standalone FusedPointwise kernel.
• eliminateDeadNodes / pruneDeadInitializers — dead-code and dead-weight elimination.

Tensor layouts: NC4HW4 and flat

The Vulkan backend has two internal GPU layouts:

• NC4HW4 (TensorFormat::NC4HW4) — the CNN-default. Channels are packed into vec4 blocks: a logical NCHW tensor with C channels becomes ceil(C/4) channel-blocks of 4. vec4 is the GPU's natural width; packing 4 channels per element keeps memory accesses coalesced and makes a conv's inner loop a tidy vec4 reduction. The device exposes no VK_KHR_cooperative_matrix, so GEMM and conv run on subgroup + vec4 math rather than matrix-core tiles.
• Flat row-major — a generic N-D path (rank up to 8) for transformer-shaped tensors (attention, RoPE, gather/broadcast/reduce/matmul). A pure decoder is entirely flat, so it needs no NC4HW4↔flat converts.

The layout pass splices ConvertLayout nodes at the NC4HW4↔flat boundary. Host data is always plain NCHW fp32; the conversion to/from a device layout happens only at segment boundaries.

fp16 storage, fp32 accumulate. With precision = low, packed buffers are 16-bit but kernels accumulate in fp32. normal keeps a named geometry-tail set in fp32 storage under fp16 compute (a no-op for models without those tensors); high is full fp32. Every fp16 narrowing store rounds to nearest even (the RoundingModeRTE SPIR-V execution mode) and saturates a finite result to ±65504 rather than overflowing to ±inf.

Segments: pre-recorded command buffers and CPU fallback

Each node is assigned to the highest-priority backend whose supportsNode(graph, node, dt) returns true (shape-aware). The topo-ordered node list is then sliced into maximal contiguous runs of the same backend — the segments. An all-Vulkan model is one segment.

Because each plan is a single fully-static shape, a VulkanSegment allocates device buffers for its activations, calls prepare() on each op (building pipelines and prepacking + uploading weights), and records the command buffer once, at compile time. At run time it only re-submits that command buffer — no re-recording, no per-op host round-trips. Ops bind their data with push descriptors, and each segment owns a timestamp query pool for the profiler.

Residency reconciliation happens only at segment boundaries. Because the host side of every RtTensor is canonical NCHW fp32, it is the universal handoff format. A Vulkan segment packs its boundary inputs host→device (NCHW→NC4HW4) only when the host is the sole valid copy, and unpacks its boundary outputs device→host so the next (possibly CPU) segment can read them. A CPU fallback segment in the middle of a Vulkan graph therefore needs no special-casing — the boundaries move data exactly once.

Plan buckets and dynamic shapes

The engine plans a graph for a fixed, fully-static shape. Dynamic shapes are supported by declaring the shapes ahead of time as plan buckets — each bucket is a full pass-and-plan product for one input-shape set. run() selects a bucket by the bound input shapes:

• A single bucket (a fixed-shape model, or one with only a dynamic batch) takes a zero-key fast path — no shape key, no lookup, no extra allocation on the hot path. The .vxm bytes and behavior are byte-identical to a model with no bucket feature at all.
• A multi-bucket .vxm (e.g. a prefill shape and a decode shape) selects by an exact-match shape key over all graph inputs; an unmatched shape is Status::InvalidArgument.

Compile-time buckets come from vknn_compile --bucket. An ONNX-loaded session can add a bucket at run time with Session::prepareShapes(shapes); a .vxm session cannot (recompile with --bucket).

The 0-CPU-fallback support gate

The support gate is the oracle for "does this model run entirely on the GPU?" It is one function, declared in core so the offline tool and the on-device engine cannot drift:

vknn_compile model.onnx out.vxm --support-report report.json

--support-report runs vkSupportSurvey(g) — the exact same vkNodeGate the device runs in VulkanBackend::supportsNode — and writes JSON: { nodes:[{name, op, backend[, reason]}], summary:{total, <backend>:count, reasons:{...}} }. "0 CPU ops" means the summary has no "cpu" / "none" key and every row is backend:"vulkan". Feed the report to tools/check_model_support.py --engine-report report.json.

At runtime the same facts are exposed through Session::fallbackReasons(), fallbackOps(), and nodeBackends(), and the benchmark runner prints fallbacks: N. To make any gap a hard error instead of a silent CPU landing, compile/run with Config::allowCpuFallback = false and an empty Config::fallback (a gap then throws Status::Unsupported), and keep tiny-GPU-island folding off (Hint::GpuIslandFold = Off) so the fallback count is meaningful.

Caches

One self-validating, multi-variant MessagePack cache file (Config::cacheFile, default <model>.cache) accelerates warm session creation. Loading it skips shader compilation, conv autotuning, and the Winograd weight transform. The whole file is guarded by a format version, a kernel hash (md5 of all embedded SPIR-V), the device (vendor/device/driver + pipeline-cache UUID), and a model hash — any mismatch discards and recomputes it. It holds one variant per cache-affecting configuration (precision, layout hints, fp32Tensors, conv-kernel hints), each bundling the Vulkan pipeline cache, the prepacked-weights cache, and the autotune table. It is written at teardown, and only when it actually changed during the session.

Extensibility

All three plugin points use self-registration (no edits to core dispatch), and rely on the whole-archive link of the static library:

• Add a CPU op: subclass vknn::CpuOp + VKNN_REGISTER_CPU_OP(OpType, Class).
• Add a Vulkan op: subclass vknn::VulkanOp + VKNN_REGISTER_VK_OP(OpType, Class) + a .comp shader.
• Add a backend: subclass vknn::Backend + VKNN_REGISTER_BACKEND(BackendKind, Class).

See Adding an Operator for the operator recipe.

Architecture Decision Records

The design decisions are recorded as ADRs under docs/adr/ in the repo:

ADR	Purpose
0001	C++17, CMake + NDK, and a static library as the build foundation.
0002	Link the NDK libvulkan directly and embed SPIR-V shaders into the static lib.
0003	UMA memory strategy: direct-mapped device-local + host-visible buffers, no staging copies.
0004	The internal NC4HW4 compute layout and the conv strategy without cooperative matrix.
0005	Caller-owned DMA-BUF I/O via VkImportMemoryFdInfoKHR and the dma_heap system heap.
0006	The segment-based execution model plus automatic per-op CPU fallback.
0008	The original disk caches (pipeline / weights / autotune) — superseded by 0009.
0009	The self-validating, multi-variant MessagePack cache and one tuning knob.
0010	One general pointwise-region fusion pass; the ConvGemm lowering.
0011	fp32-chained fused units (fusion v3): fp32 accumulation chained across a fused unit, rounded once per store.
0012	The generic, data-driven engine: dynamic-shape plan buckets, the quantized path, and full GPU op parity.