Custom Op Starter (LiteRT + AOT)

This guide is a practical starting point for adding a custom LiteRT op to HeliaAOT.

It shows:

• a LiteRT parser function
• an AOT operator class
• a registry customizer that wires both into RegistryContext

The example uses a minimal pass-through custom op key: CUSTOM_NOOP.

1) Create a Plugin Module

Create a Python module in your project, for example 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.litert import schema_py_generated as litert
from helia_aot.registry.context 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 op into one AIR operator."""
    op = subgraph.operators[op_id]
    input_ids = [str(tid) for tid in op.inputs if tid >= 0]
    output_ids = [str(tid) for tid in op.outputs if tid >= 0]

    return AirOperator(
        id=str(op_id),
        op_type=CUSTOM_OP_KEY,
        input_ids=input_ids,
        output_ids=output_ids,
        named_tensors={},
        options=AirOperatorOptions(),
    )


class CustomNoopOperator(AotOperator):
    """Minimal AOT operator that emits a tiny stub implementation."""

    TYPE = CUSTOM_OP_KEY

    def on_resolve(self):
        # Add custom validation/mutations here if needed.
        if len(self.input_tensors) != 1 or len(self.output_tensors) != 1:
            raise ValueError("CUSTOM_NOOP expects exactly one input and one output")

    def emit(self, save_path: Path):
        # Keep output minimal. You can replace this with Jinja templates later.
        c_name = f"{self.prefix}_{self.name}"
        h_path = save_path / "includes-api" / f"{c_name}.h"
        c_path = save_path / "src" / f"{c_name}.c"

        h_path.write_text(
            f"#pragma once\n"
            f"#include <stdint.h>\n"
            f"int32_t {c_name}(void *ctx);\n"
        )
        c_path.write_text(
            f"#include \"{c_name}.h\"\n"
            f"int32_t {c_name}(void *ctx) {{ (void)ctx; return 0; }}\n"
        )


def customize_registry(context: RegistryContext) -> None:
    """Hook custom parser + custom AOT operator into conversion."""
    context.litert_parsers.register(CUSTOM_OP_KEY, parse_custom_noop, overwrite=True)
    context.aot_operator_classes.register(CUSTOM_OP_KEY, CustomNoopOperator, overwrite=True)

2) Build Registry Context With Your Customizer

from helia_aot.converter import AotConverter
from helia_aot.registry.context import build_default_registry_context
from my_plugin.custom_noop import customize_registry

registry_context = build_default_registry_context(
    customizers=[customize_registry],
    allow_override=True,
)

converter = AotConverter(config, registry_context=registry_context)
converter.convert()

3) Important Key-Matching Rule

Use helper custom_code_to_op_key(...) when you construct your registry key.

Example: - LiteRT customCode="custom-noop" maps to registry key CUSTOM_NOOP - custom_code_to_op_key("custom-noop") returns CUSTOM_NOOP

Current normalization behavior: - trims whitespace - converts to uppercase - replaces hyphens (-) with underscores (_) - does not currently rewrite other separators (for example .); do not use spaces in customCode, as spaces are not rewritten and will lead to invalid C identifiers in AOT-generated operator names

4) How This Flows at Runtime

1. LiteRT parser sees custom op and derives op_name from customCode.
2. context.litert_parsers resolves your parse_custom_noop.
3. AIR operator is created with op_type="CUSTOM_NOOP".
4. Resolve stage uses context.aot_operator_classes to pick CustomNoopOperator.
5. Emit stage writes your custom source/header files.

5) What To Extend Next

• Add a typed options class and parse custom options payload.
• Emit real kernel call(s) instead of a stub return.
• Add unit test for parser dispatch.
• Add unit test for operator resolve/emit.
• Add one end-to-end conversion test with your customizer.

6) Notes

• CLI default path builds default registries only.
• Custom registry customizers are currently supported via Python API only (not CLI flags/config).
• For custom registries/custom ops, use Python API and pass explicit registry_context.
• Operator/tensor YAML attributes (config.operators, memory.tensors) still apply independently.