Adding an Operator

Adding an Operator

Adding an operator to VKNN is a set of self-registering files — no edits to any core dispatch code. The CMake globs use CONFIGURE_DEPENDS, so new .cpp / .comp files are picked up on the next ./build.sh, and the static library is linked whole-archive so the VKNN_REGISTER_* registrars survive.

Each piece is independently shippable. Under the capability / fallback model, a CPU-only op runs on the CPU backend and the Vulkan path falls back to it automatically; adding the Vulkan kernel then promotes the node to the GPU.

Pointwise shortcut (usually this)

A new pointwise op — an activation or an elementwise binary — does not need a new OpType. The elementwise families are sub-codes: a unary imports as OpType::Unary with a UnaryType in Node::subOp, a binary as OpType::Binary with a BinaryType. So a new activation is four small edits:

1. Add the value to UnaryType (or BinaryType) in include/vknn/op_type.h.
2. Map its ONNX name in unaryFromOnnx() / binaryFromOnnx() (src/core/op.cpp).
3. Add a case to the CPU unary/binary switch.
4. Add a case to the GLSL unary/binary switch.

No gate, no descriptor row, no shape rule. (SiLU already exists this way: UnaryType::SiLU.)

Full recipe for a genuinely new OpType

For an op that is not a pointwise family member (e.g. a native RMSNorm, or ArgMax):

#	File	Edit / Create	Content
1	include/vknn/op_type.h	Edit	Append the value at the END of enum class OpType. Append-only — model_io serializes OpType as a raw int into .vxm, so a mid-enum insert corrupts every existing model.
2	src/core/op.cpp	Edit	Add case OpType::Foo: return "Foo"; to opTypeName() and {"Foo", OpType::Foo} to opTypeFromOnnx(). The string must equal the ONNX op_type exactly.
3	src/import/infer_shapes.cpp	Edit (only if shape-changing)	Add a case OpType::Foo: producing the exact ONNX output shape for every shape-affecting attribute (use readI64Param() for attr-or-initializer params). Identity-shape ops need nothing. A missing rule leaves the output shape unresolved and planning fails; a wrong rule is worse (a Shape node const-folds the lie).
4	src/backend/cpu/ops/foo.cpp	Create	struct FooCpu : CpuOp { void run(const Node&, ExecContext&) override {…} }; + VKNN_REGISTER_CPU_OP(OpType::Foo, FooCpu);. Host buffers are canonical NCHW fp32; size the output with cpu::allocOut (or cpu::allocOutI64 for index outputs). This is the numeric oracle — write it first.
5	shaders/foo.comp	Create	GLSL targeting vulkan1.3. #include "precision.glsl", type SSBOs as STORE (read float(x[i]), compute fp32, write y[i] = TO_STORE(v)), and a layout(push_constant) block byte-matching the C++ struct. #include "common.glsl" for vx_act(). Globbed automatically — no CMake edit.
6	src/backend/vulkan/ops/foo.cpp	Create	#include "vk_op_common.h". A PC struct byte-matching the shader; struct FooOp : VulkanOp with prepare() (build the pipeline via env.pipeline(shader("foo", env.useFp16), …), upload const weights) and record() (env.devBuf(id) → pipe->dispatch(…)). VKNN_REGISTER_VK_OP(OpType::Foo, FooOp);.
7	src/core/vk_gates.cpp	Edit (conditional)	Add a vkNodeGate arm only if the kernel cannot handle every shape/attribute the op admits (refuse via refuse(whyNot, "Foo: <reason>")). A pure pointwise op falls through to the final return true. Touch vkKernelDeclared() only to declare an op unsupported on the GPU.
8	src/core/op_descriptor.cpp	Edit (only if non-default)	One set(OpType::Foo, LayoutClass::…, pwMember, pwEpilogue); row. Default (unlisted) = {Nc4, false, false}. Use LayoutClass::Flat for a generic N-D op; set pwEpilogue = true to let a fused pointwise chain ride the store.
9	src/import/insert_layout_converts.cpp	Edit (only if ShapeDependent)	Add the matching arm to gpuFlatNode, kept in sync with the vkNodeGate case (a gtest asserts they never disagree).
10	tests/…	Create	A gtest diffing the Vulkan output against the CPU oracle (and optionally an onnxruntime golden). Confirm placement with vknn_compile <m>.onnx out.vxm --support-report r.json.

The two capability functions

The Vulkan backend's supportsNode() delegates to two pure functions in src/core/vk_gates.cpp — this file lives in core, so vknn_compile --support-report and the host tests evaluate the exact gate the device runs, with no chance of the two drifting:

• vkKernelDeclared(OpType) mirrors the VKNN_REGISTER_VK_OP set. It returns true by default and lists (as return false) only the ops with no GPU kernel. A new op with a Vulkan kernel is already covered by default: return true.
• vkNodeGate(graph, node, whyNot) is the shape/attribute gate. It returns true when the kernel accepts this node's shapes/attributes/operand-constness, and on refusal fills whyNot with a short "<Op>: <reason>" that shows up in the fallback warning and the support report.

The shader / SPIR-V embed flow is automatic

The build compiles foo.comp → foo.spv, doubles it to foo_fp16.spv if it includes precision.glsl, and derives foo_epi[_rx].spv if it includes pw_epilogue.glsl. tools/embed_spirv.py embeds every .spv into embeddedShaders() keyed by basename and folds an MD5 into embeddedShadersHash() (the cache-version guard). env.pipeline(shader("foo", useFp16), …) resolves the key at runtime. Configure-time tools/check_epi_sync.py and tools/check_shader_contracts.py FATAL-error on epilogue / store-contract drift.

Hosting a pointwise-chain epilogue (producer ops only)

A producer op (Conv / MatMul / pool / reduce / gather-shaped kernel) can host a fused pointwise chain at its store instead of leaving it a standalone FusedPointwise node. Three things wire it up, and tools/check_epi_sync.py keeps them in sync:

1. Shader — under #ifdef PW_EPI, #include "pw_epilogue.glsl" and apply v = pw_apply(v, idx) (or pw_apply_nc4) before STORE. Including that header is what marks the kernel epilogue-capable; the build derives the _epi / _epi_rx variants.
2. Op — wire PwEpi into the kernel (src/backend/vulkan/ops/pw_plan.h): request the variant with shader((std::string("<stem>") + epi.suffix()).c_str(), …) sized nbuf + epi.extraBufs(), and epi.append(bufs, …) after the kernel's own buffers.
3. Capability — set the descriptor row's pwEpilogue so fusePointwiseChains attaches units to this producer.

The capability / fallback model

Backend::supports(OpType, DType) is a thin wrapper over the op registries: registering the CPU op makes CpuOpRegistry::has() true (the oracle / fallback target); registering the Vulkan op makes VkOpRegistry::has() true (promotes the node to the GPU when vkNodeGate accepts). They are independent, so the CPU op is always available as the numeric reference. The fallback path can be exercised without touching code by forcing an op back to CPU at runtime:

cfg.disableVkOps = "Foo";   // re-partitions the graph; output stays cosine 1.000000

Checklist

• include/vknn/op_type.h: append at the END of the enum (append-only).
• src/core/op.cpp: opTypeName() + opTypeFromOnnx().
• src/import/infer_shapes.cpp: a shape rule if the op changes shape (cross-check vs onnx.shape_inference over the attribute matrix).
• src/backend/cpu/ops/foo.cpp: CPU oracle + VKNN_REGISTER_CPU_OP (one op per file).
• shaders/foo.comp: precision.glsl + STORE buffers (the build emits the _fp16 variant).
• src/backend/vulkan/ops/foo.cpp: Vulkan op + VKNN_REGISTER_VK_OP (one op per file).
• src/core/vk_gates.cpp: a vkNodeGate case only if the kernel can't take every shape/attribute; leave vkKernelDeclared alone unless the op has no GPU kernel.
• src/core/op_descriptor.cpp: the layout class + fusion roles row (pure pointwise keeps the all-default row).
• A gtest under tests/; confirm placement with --support-report.