-
Notifications
You must be signed in to change notification settings - Fork 0
Operator Coverage
VKNN imports ONNX and lowers it to its IR. Each operator has a CPU oracle (the bit-exact reference and automatic fallback) and, for every op that carries compute, a Vulkan kernel. CNN-shaped tensors default to the NC4HW4 GPU layout (channels packed in vec4 blocks); transformer-shaped tensors (attention, RoPE, geometry) use a flat row-major GPU path (rank up to 8), and the layout pass splices NC4HW4↔flat converts in at the boundaries. Every operator is checked against onnxruntime (cosine ≥ 0.999 on the GPU fp16 path, 1.0 on the CPU fp32 path).
Every executable operator runs on the GPU. Across the benchmarked model zoo (CNN, detection, transformer encoder, and the autoregressive LLM), the nodes that carry compute all land on Vulkan. The only operators the engine does not run are:
-
Data-dependent control flow —
Loop,If,Scan,NonMaxSuppression— whose output shapes are not known at plan time. These are not in the op table and plan on no backend. -
Shape-arithmetic ops the passes const-fold away —
Shape,Constant,EyeLike— resolved at compile time, so no runtime kernel is dispatched at all.
A small number of ops have a GPU kernel but fall back to the CPU oracle on a specific input class the kernel cannot represent (an int64→narrow-integer Cast, a runtime-k TopK, an unresolved-shape ConstantOfShape / Range).
Two tools answer "will this model run, and where?":
-
tools/scan_unsupported_ops.py model.onnx— a fast, import-side regex scan that enumerates every unrecognized / potentially-unsupported op with its shapes and attributes (--jsonfor machine output). -
tools/check_model_support.py— consumes either the ONNX directly or, authoritatively, an engine support report.
The ground truth is the engine's own capability gate. Generate the per-node backend assignment with:
vknn_compile model.onnx out.vxm --support-report report.json
python3 tools/check_model_support.py --engine-report report.jsonThe report comes from vkSupportSurvey — the same vkNodeGate the device engine runs — so backend:"none" is a hard blocker and any backend != "vulkan" is a CPU fallback. See Architecture for the 0-CPU-fallback gate.
Every operator lives in its own file under src/backend/{cpu,vulkan}/ops/ (one op per file). The IR carries roughly 67 OpType values; the elementwise unary and binary families are sub-codes of OpType::Unary / OpType::Binary.
| Family | Operators | GPU | Notes |
|---|---|---|---|
| Convolution & pooling | Conv (group=1, depthwise, 1×1), ConvTranspose, GlobalAveragePool, AvgPool, MaxPool, BatchNorm | ✅ | NC4HW4; direct 3×3, split-K deep 1×1, tiled-GEMM Winograd, fused activation + residual in the epilogue. BatchNorm usually folded into Conv. |
| Unary | Sigmoid, Tanh, HardSwish, HardSigmoid, LeakyRelu, Elu, Abs, Neg, Exp, Log, Sqrt, Floor, Ceil, Relu, SiLU, Erf, Cos, Sin, Reciprocal, Softplus, Round | ✅ | The UnaryType family; a new activation extends the family enum + switches, no new OpType. |
| Binary | Mul, Sub, Div, Max, Min, Pow, Add | ✅ | Same-shape, channel-broadcast (SE), and general NumPy broadcast on the flat path. |
| Activation / compare | Relu, Relu6, Clip, PRelu, Where, Equal, Greater, GreaterEqual, Less, LessOrEqual, IsNaN, And | ✅ | Standalone and fused; flat broadcast (fp32 + int64). IsNaN / And added for the LLM path. |
| Transformer / attention | MatMul (batched N-D), Gemm, Softmax, LayerNorm, Einsum, Gather | ✅ | QKᵀ / AV / MLP matmuls, last-axis softmax, affine LayerNorm, RoPE outer-product Einsum, axis-aware Gather (const or runtime index). |
| Shape / data movement | Reshape, Flatten, Squeeze, Unsqueeze, Transpose, Slice, Concat, Split, Expand, Tile, DepthToSpace, ScatterND, Resize, Upsample, GridSample, Reduce (Mean/Sum/Max/Min/Prod/L2), Pad, Identity | ✅ | Metadata + flat copy / gather / scatter. Identity is rewired to its producer at import. |
| Cast | Cast | ✅ | float↔float and int→float on the GPU; an int64→narrow-integer target (INT16/UINT16/BOOL/wide-unsigned) keeps the exact CPU op. A BOOL-target int64 Cast runs on the GPU (added for the LLM path). |
| Generators | ConstantOfShape, Range, Shape, Constant, EyeLike | ✅ / const-fold | Resolved output sizes generate/fill on the GPU; unresolved (data-dependent) sizes keep the CPU op. Shape / Constant / EyeLike const-fold away before planning. |
| Selection | TopK | ✅ | k largest/smallest, values + int64 indices, ties to the lower index; GPU when k is a const int64 and the shape resolves, else the CPU op. |
| Eliminated / lowered | Dropout, InstanceNormalization | — | Dropout is removed as an inference-mode identity; InstanceNorm decomposes at import into ReduceMean + Sub/Mul/Add/Sqrt/Div + per-channel scale/bias, so it runs wherever those ops run. |
A quantized ONNX checkpoint runs dequantized to float via the import-time dequantize pass (src/import/dequantize_graph.cpp, default on; --no-dequantize disables). For static QDQ / QLinear the pass drops the 8-bit rounding but preserves the saturation clamp each quant hop encodes (so a ReLU folded into an activation quant range survives) — results match an int-exact runtime closely but not bit-exactly. For dynamic quantization it matches the canonical DynamicQuantizeLinear → MatMulInteger / ConvInteger → Cast → Mul cluster and folds it to a plain float MatMul / Conv. A cluster that does not match this shape stays intact and fails at planning. There is no int8 compute tier.
Lowered quantized ops: DequantizeLinear, QuantizeLinear, QLinearConv, QLinearMatMul, QGemm, QLinearAdd, QLinearGlobalAveragePool, DynamicQuantizeLinear, MatMulInteger, ConvInteger.
vknn_compile groups these behind an optimization level (-O0 = none, -O1 = the default production set, -O2/-O3 = + the experimental SE and depthwise+1×1 fusions); the --[no-]fuse-* / --[no-]lower-conv flags override a single pass on top of the level.
-
General pointwise fusion — the one fusion pass. It grows each maximal same-shape per-element region and emits it as a single fused unit: folded into the producing kernel's store epilogue (MatMul, Gemm, the Conv family, ConvGemm, Softmax, LayerNorm, Reduce, GridSample, Resize, ConvTranspose, pooling, Transpose/Slice, Concat) or one standalone
FusedPointwisekernel. Residual Adds, swish diamonds (x · sigmoid(x)), MatMul bias-Adds, and lone activations are all cases of this pass.--strict-fusekeeps every step rounded sofused == unfusedbyte-identical. - BatchNorm lowering — a BN the conv fold cannot absorb lowers to a per-channel Mul+Add.
- InstanceNorm lowering — decomposes unconditionally into spatial ReduceMean, centered Sub, squared-diff Mul, epsilon Add, Sqrt, Div, and a per-channel scale/bias Mul+Add.
-
Conv → ConvGemm lowering — a non-Winograd K×K Conv lowers to one LDS-tiled implicit-GEMM kernel (experimental, off by default;
--lower-conv). -
Squeeze-Excite and depthwise + 1×1-project fold to one kernel each (
-O2/--fuse-se/--fuse-dwpw, experimental). - Einsum → MatMul/Squeeze/Unsqueeze, and ConvTranspose → Conv + DepthToSpace (subpixel rewrite).
See Adding an Operator for the full recipe. In short: append an OpType value at the END of the enum in include/vknn/op_type.h (append-only — .vxm files store the raw integer), map its ONNX name in src/core/op.cpp, add a shape rule in src/import/infer_shapes.cpp if it changes shape, and add a CPU oracle. A GPU kernel additionally needs a Vulkan op + GLSL shader, a capability gate row in src/core/vk_gates.cpp, and (for a flat op) a descriptor row and a gpuFlatNode arm.