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.

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.