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 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:

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.

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:

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 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):

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.