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.

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.