Skip to content

Adding Operators

This guide walks through the practical steps required to add a brand-new LiteRT operator to heliaAOT. Follow the sequence end-to-end to avoid the common pitfalls that typically derail new operators.

1. Ensure an NS-CMSIS-NN Reference Exists

  • heliaAOT ultimately emits kernels that call into ns-cmsis-nn, so every LiteRT operator you introduce must have a matching reference there.
  • If an implementation does not exist yet, upstream it first; otherwise any generated kernel will have nothing to invoke.

2. Build a Minimal LiteRT Test Model

  • Create a small LiteRT/TFLite flatbuffer that contains a single instance of the new op.
  • This model serves two purposes: it validates your parser and provides golden data for regression tests later.

3. Lower the Operator into AIR

  1. Enumerations - Add the new op to helia_aot/air/enums.py. - Search existing enums (op type, padding mode, activations, etc.) for required additions.

  2. Options schema - Define the operator-specific options class in helia_aot/air/options.py. - Export it from helia_aot/air/__init__.py. - Mirror LiteRT fields from helia_aot/litert/schema_py_generated.py.

  3. Parser wiring - Add the parser function in helia_aot/converters/litert/op_parsers.py. - Built-ins are registered via register_default_litert_parsers(...). - Custom parser additions are injected through RegistryContext customizers. - Validate attribute ranges early so conversion fails fast.

4. Implement the AOT Operator Class

  • Create helia_aot/aot/operators/<op>.py, subclassing AotOperator.

Responsibilities: - Validate Check tensor rank/layout/dtype, quantization constraints, and platform requirements. - Compute Values Populate template inputs (sub_values) including tensor metadata, quant params, and scratch sizing. - Registration Include class in built-in registration path (register_default_aot_operators(...)) or inject via RegistryContext customizer for plugin-style custom ops. - Declare the cmsis_nn_context.buf contract Set CTX_BUF_USAGE. This is mandatory for built-in operators — see below.

Required: Declare How You Use cmsis_nn_context.buf

Operator templates wire ctx.buf to a tensor when the operator allocated one and hardcode .buf = NULL when it did not, then pass &ctx to the kernel either way. So an operator whose sizing returns 0 bytes for some dtype/target combination that still reaches a buf-reading kernel emits a guaranteed NULL dereference. Resolve time is the only place with enough information to catch it, so every built-in operator declares its contract and resolve() checks it.

Set CTX_BUF_USAGE to one of the four CtxBufUsage values (from helia_aot.aot.operators import CtxBufUsage):

Value Meaning
NO_CTX The template passes no cmsis_nn_context to any kernel.
IGNORED A context is passed, but no dispatchable kernel dereferences buf (typically an explicit (void)ctx; upstream).
REQUIRED At least one dispatchable kernel dereferences ctx->buf in every configuration the operator supports.
CONDITIONAL Whether buf is dereferenced depends on dtype, target capability, or operator options.

Two further rules:

  1. CONDITIONAL must override ctx_buf_required(). Derive the answer from the upstream preprocessor structure, not from your own compute_scratch_size(). The guard exists to cross-check two independent derivations; deriving one from the other makes it a tautology that always passes.
  2. Name your backing roles if ctx.buf is not backed by a plain scratch tensor. CTX_BUF_BACKING_TENSOR_NAMES lists the op.named_tensors keys that may reach a context buffer. Leave it empty only when "any non-empty SCRATCH tensor this operator allocated" is exactly right. Declare it when the backing is a precomputed constant (CONV_2D, FULLY_CONNECTED point at a weight-sum constant) or when the operator allocates scratch tensors that are not context buffers (BATCH_MATMUL's lhs_tp/rhs_tp are data operands to arm_transpose_*), since the fallback would count those and silently satisfy the check.

Cite your evidence. These declarations are only worth having if they are backed by the source. Every declaration in the tree carries a comment naming the upstream file, line, and preprocessor gate it was read from, at a pinned ns-cmsis-nn version. Follow that convention: a reviewer must be able to check your claim without guessing which kernel you meant.

Coverage is pinned by tests/unit/aot/test_ctx_buf_guard.py (test_every_registered_operator_declares_ctx_buf_usage), which fails for any class in _AOT_OPERATOR_CLASSES that leaves CTX_BUF_USAGE unset, and test_conditional_operators_override_ctx_buf_required, which fails for a CONDITIONAL operator that inherits the base hook. Add a "guard fires" test there too: force your sizing to 0 for a configuration where ctx_buf_required() is true and assert resolve() raises — a declaration that no test can make fire is not yet evidence of anything.

Custom operators are not covered by that meta-test

AotOperator.CTX_BUF_USAGE defaults to None, which is an undeclared sentinel rather than a usage. _assert_ctx_buf_backed() returns immediately when it sees None, before any backing check runs — so an operator that never sets it skips the guard entirely, silently.

Completeness is enforced statically rather than at runtime, and the meta-test parametrizes over _AOT_OPERATOR_CLASSES: the hardcoded list of built-in classes in helia_aot/aot/operators/__init__.py. A custom operator injected through a RegistryContext customizer is not in that list, so it is not covered — it defaults to None and opts out of the guard by doing nothing at all.

That is deliberate (a half-migrated tree has to stay runnable), but it means the safety net is yours to opt into. If your custom operator's template passes a cmsis_nn_context to any kernel, declare CTX_BUF_USAGE anyway: the guard is the only thing standing between a sizing bug and a NULL-buf kernel call, and that call is either silent memory corruption or a stale output — neither of which the FVP will catch.

5. Author the Code-Gen Templates

  • Add helia_aot/aot/templates/<op>.c.j2 and <op>.h.j2.
  • Keep conventions aligned with existing templates (headers, naming, const placement).
  • Pass required params through Jinja vars; avoid hidden recomputation inside templates.
  • Preserve layout and platform feature branches (#ifdef) where needed.

Optional: Route Through a Shared Operator Kernel

Most operators inline their CMSIS-NN call site per node. When many nodes of the same operator share an identical call signature, you can collapse that call site into a single shared runtime (emitted once into {prefix}_kernels.c/.h) and have each node emit only a compact rodata descriptor plus a thin _run wrapper. ADD, MUL and per-channel int8 FULLY_CONNECTED already do this.

This seam has a strict three-part contract — all three must agree or generated code breaks:

  1. Describe the helper. Add a factory in helia_aot/aot/operators/kernel_runtime.py returning a KernelHelper (a frozen dataclass keyed by signature). Use a single source-of-truth method on the operator (e.g. _kernel_helper()) so the helper is computed once.
  2. Declare it. Override shared_kernel_helpers() to return the helper(s), and pass the helper into compute_values() for the _run body.
  3. Match has_init. If the shared helper makes per-node init a no-op, override has_init to return False and omit the _init function/declaration from your .c.j2/.h.j2. Derive has_init from the same predicate as the helper so they cannot drift.

Failure modes if the parts disagree:

  • has_init=True but the template omits _init → aot_model.c references a missing {prefix}_{name}_init symbol → link error.
  • has_init=False but the template still emits _init → a dead init function compiles in silently, defeating the size optimization.

The invariant is pinned by tests/unit/aot/test_shared_kernels.py (test_has_init_matches_rendered_init); add your operator there when you opt in. The kernel descriptor struct in kernels.h.j2 and its implementation branch in kernels.c.j2 are selected by the helper's struct_kind; the struct field names must match the .field = {{ value }} literals in your operator .c.j2.

6. Testing

  • Unit / functional tests Add focused conversion tests for the new op model and generated values.
  • End-to-end tests Add coverage in tests/ to catch integration issues (buffers, quantization, runtime wiring).
  • Re-run relevant platform-specific suites (e.g., Ethos-U, MVE).

Common Pitfalls Checklist

  • ns-CMSIS-NN signature mismatch (dims, offsets, activation clamps).
  • Missing enum/options export causing delayed parser/runtime failures.
  • Wrong tensor role mapping (input/output/named/scratch).
  • Missing per-channel quant arrays where required.
  • Inadequate edge-shape test coverage (broadcast, dilation, unusual ranks).

Note on Config Overrides

Operator/tensor YAML attribute rules (config.operators, memory.tensors) are independent of RegistryContext and continue to work as before.