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.

Custom-operator plugins are experimental in 1.x. The names this guide imports from the API manifest, such as AotConverter, ConvertArgs, RegistryContext, build_default_registry_context and custom_code_to_op_key, keep the compatibility promise. The operator base class AotOperator and its companions (CtxBufUsage, OpCapability, register_default_aot_operators) are provisional: importable from helia_aot.aot.operators, but a minor release can change them. The rest of the plugin surface is outside the promise and can change in any release: the LiteRT schema helpers the starter imports (helia_aot.converters.litert.schema_access and helia_aot.litert.schema_py_generated), the parser signature, and the C contract an operator’s emitted code follows. A stable plugin API is tracked in #829.

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 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 replaces hyphens and spaces with underscores, then trims any remaining surrounding whitespace and uppercases, so custom-noop becomes CUSTOM_NOOP and my op becomes MY_OP. Because spaces are replaced before the trim, a leading or trailing space becomes an underscore: ' custom-noop ' becomes _CUSTOM_NOOP_. It does not rewrite anything else.

Caution

Do not put dots or other punctuation in a customCode. They are not rewritten, and they 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. Return 0 on success and, on failure, a code from <prefix>_status_custom_first to <prefix>_status_custom_last (1000 to 1999): the generated module never returns one of those, so the caller can tell your operator’s failure from the module’s. 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.

Attributes of your own

An operator reads its operators[].attributes from self.attributes, an instance of its ATTRIBUTES model. A class that keeps the base model takes only code_placement. Declare the attributes your operator takes as fields of a pydantic dataclass that subclasses OperatorAttributes and set it as ATTRIBUTES; the conversion checks every rule against the models of the registry it uses, so a key no reached operator declares is an error that names the nearest valid key. To keep keys you do not declare, as a plugin that passes them to its own templates may, declare the model with extra="allow":

from pydantic import ConfigDict
from pydantic.dataclasses import dataclass

from helia_aot.aot.operators import AotOperator, OperatorAttributes


@dataclass(config=ConfigDict(extra="allow"))
class MyAttributes(OperatorAttributes):
    pass


class MyOperator(AotOperator):
    ATTRIBUTES = MyAttributes

In a rule for your operator’s type, keys a free-form model does not declare are not checked: a misspelled one reaches your code as written. A type: "*" rule is still checked against the operators that do not take free-form keys, and a removed key such as scratch_placement is an error either way. A field without a default is a required attribute. The inherited code_placement has a default, so declare the dataclass with kw_only=True for a required field to follow it:

from pydantic.dataclasses import dataclass

from helia_aot.aot.operators import OperatorAttributes


@dataclass(kw_only=True)
class GainAttributes(OperatorAttributes):
    gain: int

Any rule that reaches the operator may set it, and if none does, the conversion stops with a configuration error that names it.

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 AirModel, AirOperator, AirOperatorOptions, custom_code_to_op_key
from helia_aot.aot.operators 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)
    context.aot_operator_classes.register(CUSTOM_OP_KEY, CustomNoopOperator)

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 a new key, so it needs neither overwrite=True on its registrations nor allow_override=True on the context. Leaving allow_override at its default, False, protects the built-ins: a customizer that replaces or removes one makes context construction 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.

Built-in CORE plans

When extending a built-in operator, preserve its validation, resolved kernel plan and resource bindings. Rebuild the conversion context when changing the model or target; lowered constants may no longer contain the original values. See Backend plans for provider ownership and the resolution lifecycle.

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. A transform declares the options it takes as OPTIONS, a pydantic model or dataclass. Without one it takes no options, and OPTIONS = None leaves them for the transform to check.

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.