# Custom operators

A LiteRT custom op is lowered by two pieces of your code: a **parser** that
turns the LiteRT node into an AIR operator, and an **operator class** that
emits C for it. Both are registered on a `RegistryContext` that you build and
hand to the converter.

:::note
This is a Python API. The command line always builds the default, frozen
registry context, so there is no flag or configuration key for a custom
operator. Attribute rules under `operators` and `memory.tensors` are separate
from registry selection and apply either way.
:::

## Start with the maintained example

Use this guide when a model carries a LiteRT `customCode` that the built-in
registry cannot lower, or when you deliberately want to replace a built-in
lowering. Ordinary placement and operator attributes do not need a plugin.

Start from the [Custom operator example](https://ambiqai.github.io/helia-aot/examples/custom-operator/)
in a repository checkout with its Python dependencies installed. It includes
a model builder, conversion configuration, plugin and run script. Read its
setup before running the example; an isolated CLI installation does not supply
the Python environment used by the model builder.

Success at this stage means the custom node is parsed, resolved and emitted,
and its generated code can be built. To establish your operator's numerical
behavior, add trusted golden inputs and outputs and run the generated test.
A custom op may have no host-interpreter implementation; `skip_verification`
is only a runtime smoke test in that case.

## The registry context

`RegistryContext` holds every registry a conversion uses:

| Registry | Holds | Keyed by |
| --- | --- | --- |
| `litert_model_hooks` | Rewrites applied to the flatbuffer before parsing | A name |
| `litert_parsers` | One parser per operator | The canonical operator key |
| `air_options` | The typed options class per operator | The canonical operator key |
| `aot_operator_classes` | The class that emits C for an operator | The canonical operator key |
| `transforms` | Graph rewrites | The transform class's `NAME` |

Build one from the built-ins and pass it in:

```python
from helia_aot.converter import AotConverter
from helia_aot.registry import build_default_registry_context

registry_context = build_default_registry_context()
AotConverter(config, registry_context=registry_context).convert()
```

Omitting `registry_context` uses the built-ins implicitly, which is the same
thing the command line does.

## One key, three places

Use the same canonical string in all three places:

- the parser's key in `litert_parsers`,
- the `op_type` your parser writes onto the AIR operator,
- the operator class's key in `aot_operator_classes`.

For a LiteRT custom op, derive it from the `customCode` with
`custom_code_to_op_key`. That helper trims whitespace, uppercases, and replaces
hyphens with underscores, so `custom-noop` becomes `CUSTOM_NOOP`. It does not
rewrite anything else.

:::caution
Do not put spaces or dots in a `customCode`. Neither is rewritten, and both
produce invalid C identifiers in the generated operator names.
:::

## What an operator class has to emit

The starter below emits `<prefix>_<name>_init` and
`<prefix>_<name>_run`, each taking the module's context struct and returning an
`int32_t` status. The model propagates a nonzero status to its caller. The
function names and signatures must agree with the generated dispatch code.

An operator can declare `has_init` false when it needs no initialization; the
model then omits that init call. Keep the starter's no-op init while learning
the interface. Operators with internal state must arrange persistent storage
and initialization. An operator with explicit state inputs and outputs instead
needs a documented caller carry/reset contract; the supported GRU uses that
form and rejects persistent state tensors. Working memory must be reported
before planning. A local pointer in emitted code does not reserve space in the
model's arenas.

Two more rules come with it. Tensors are reached through
`ctx->tensor_ptrs[<prefix>_<tensor name>]` rather than through symbols of their
own, which is what lets the planner place the buffer wherever the rules say.
Keep `#include` lines outside the header's `extern "C"` guard; put the
emitted C function declarations inside it so C++ consumers use C linkage.

## A starter plugin

Put the parser, the operator class and the customizer in one module of your own
project. The operator below copies its input to its output, which is the
smallest thing that parses, emits, links and runs.

```python title="my_plugin/custom_noop.py"
from pathlib import Path

from helia_aot.air import AirOperator, AirOperatorOptions, custom_code_to_op_key
from helia_aot.air.model import AirModel
from helia_aot.aot.operators.operator import AotOperator
from helia_aot.converters.litert.schema_access import inputs_of, operator_at, outputs_of
from helia_aot.litert import schema_py_generated as litert
from helia_aot.registry import RegistryContext

CUSTOM_CODE = "custom-noop"
CUSTOM_OP_KEY = custom_code_to_op_key(CUSTOM_CODE)

def parse_custom_noop(
    op_id: int,
    model: litert.ModelT,
    subgraph: litert.SubGraphT,
    air_model: AirModel,
) -> AirOperator:
    """Parse one LiteRT custom node into one AIR operator."""
    del model, air_model
    # The generated LiteRT classes leave every absent vector as None, so read
    # them through the accessors rather than indexing them directly.
    op = operator_at(subgraph, op_id)
    return AirOperator(
        id=str(op_id),
        op_type=CUSTOM_OP_KEY,
        input_ids=[str(tid) for tid in inputs_of(op) if tid >= 0],
        output_ids=[str(tid) for tid in outputs_of(op) if tid >= 0],
        named_tensors={},
        options=AirOperatorOptions(),
    )

class CustomNoopOperator(AotOperator):
    """Emit the C for one CUSTOM_NOOP node: a byte copy from input to output."""

    TYPE = CUSTOM_OP_KEY

    def on_resolve(self) -> None:
        """Reject a node the emitted C could not serve."""
        if len(self.input_tensors) != 1 or len(self.output_tensors) != 1:
            raise ValueError("CUSTOM_NOOP expects exactly one input and one output")
        if self.input_tensors[0].nbytes != self.output_tensors[0].nbytes:
            raise ValueError("CUSTOM_NOOP expects input and output of the same size")

    def emit(self, save_path: Path) -> None:
        """Write the operator's header and source into the module tree."""
        stem = f"{self.prefix}_{self.name}"
        source = self.input_tensors[0]
        destination = self.output_tensors[0]
        (save_path / "includes-api" / f"{stem}.h").write_text(
            "#pragma once\n\n"
            f'#include "{self.prefix}_common.h"\n\n'
            '#ifdef __cplusplus\nextern "C" {\n#endif\n\n'
            f"int32_t {stem}_init({self.prefix}_model_context_t *ctx);\n"
            f"int32_t {stem}_run({self.prefix}_model_context_t *ctx);\n\n"
            "#ifdef __cplusplus\n}\n#endif\n",
            encoding="utf-8",
            newline="\n",
        )
        (save_path / "src" / f"{stem}.c").write_text(
            f'#include "{self.prefix}_common.h"\n'
            f'#include "{stem}.h"\n\n'
            f"int32_t\n{stem}_init({self.prefix}_model_context_t *ctx)\n"
            "{\n    (void)ctx;\n    return 0;\n}\n\n"
            f"int32_t\n{stem}_run({self.prefix}_model_context_t *ctx)\n"
            "{\n"
            f"    const int8_t *__restrict input = "
            f"(const int8_t *)ctx->tensor_ptrs[{self.prefix}_{source.name}];\n"
            f"    int8_t *__restrict output = "
            f"(int8_t *)ctx->tensor_ptrs[{self.prefix}_{destination.name}];\n"
            "    if (input == output) { return 0; }\n"
            f"    arm_memcpy_s8(output, input, {destination.nbytes});\n"
            "    return 0;\n}\n",
            encoding="utf-8",
            newline="\n",
        )

def customize_registry(context: RegistryContext) -> None:
    """Hook the parser and the operator class into a conversion."""
    context.litert_parsers.register(CUSTOM_OP_KEY, parse_custom_noop, overwrite=True)
    context.aot_operator_classes.register(CUSTOM_OP_KEY, CustomNoopOperator, overwrite=True)
```

Then build the context with your customizer and convert:

```python
from helia_aot import AotConverter, ConvertArgs
from helia_aot.registry import build_default_registry_context

from my_plugin.custom_noop import customize_registry

config = ConvertArgs.from_yaml("convert.yaml")
registry_context = build_default_registry_context(
    customizers=[customize_registry],
)
AotConverter(config, registry_context=registry_context).convert()
```

Run this Python entry point with the same configuration file you would use for
a normal conversion. The custom registry supplies the additional lowering.

The starter adds an operator and leaves `allow_override` at its default,
`False`. That protects existing built-ins even though the registration calls
request overwrite: replacing or removing a built-in causes context construction
to fail. Use `allow_override=True` only for the deliberate replacement below.

:::note
This example is the [Custom operator](https://ambiqai.github.io/helia-aot/examples/custom-operator/) page, backed by
`examples/custom-operator/` in the repository, with the
script that builds a model carrying the required `customCode`,
the configuration, and a run script. The repository's example checks
convert and host-compile it.
Those checks establish the integration contract; your kernel needs its own
numerical cases.
:::

## Replacing a built-in operator

The same registry replaces as well as adds. Register your own class over a
built-in key with `allow_override=True` and every node of that type is emitted
by your class instead. Pass this customizer to
`build_default_registry_context(customizers=[customize_registry], allow_override=True)`:

```python
def customize_registry(context):
    builtin = context.aot_operator_classes.get("CONV_2D")

    class MyConv2D(builtin):
        def emit(self, save_path):
            ...

    context.aot_operator_classes.register("CONV_2D", MyConv2D, overwrite=True)
```

Subclassing the built-in is usually what you want: you inherit its validation,
its kernel selection and its scratch sizing, and override only the part you are
changing. There is no configuration key that selects a kernel for an operator,
so this is the route when the kernel a built-in picks is not the one you want.

## What happens at conversion time

1. The LiteRT parser recognizes the custom op and derives its name from the
   `customCode`.
2. `litert_parsers` resolves your parser, which is called with the node.
3. The AIR operator is created with your `op_type`.
4. Resolve looks the type up in `aot_operator_classes` and instantiates your
   class.
5. Emit writes the files your class writes.

## Extending further

The same context extends beyond parsers and operator classes. `transforms`
takes graph rewrites, and `air_options` takes a typed options class so your
parser can carry a decoded options payload rather than an empty one.

From the byte copy above, the next steps are a typed options class so the
node's options survive parsing, real kernel calls in place of the copy, and
tests: one using a real model `customCode` to exercise parser dispatch, one
for resolve and emit, and one full conversion with the customizer applied.
Include rejected shapes/dtypes, a host build of the emitted interface, and
trusted numerical cases for the kernel. Attribute rules cannot substitute for
those admission checks.

The starter touches lower-level parser/schema helpers as well as the public
registry API. Keep your plugin tied to a tested heliaAOT version and use the
maintained example when updating those dependencies.

The classes and functions named here are documented in the
[Python API reference](https://ambiqai.github.io/helia-aot/reference/api/helia_aot/registry/context/).