operator

Classes

AotOperator

AotOperator(op: AirOperator, model: AirModel, platform: SocPlatform, prefix: str = 'aot', attributes: dict[str, str] | None = None)

Base class for all AOT operators.

Execution flow: 1. Initialize the operator. 2. Call resolve() to validate and prepare the operator. 3. Call plan() to set up any necessary planning. 4. Call emit() to generate the operator's code.

Parameters:

• op

  (AirOperator) –

  The AIR operator to wrap.

• model

  (AirModel) –

  The AIR model.

• platform

  (SocPlatform) –

  The target platform for code generation.

• prefix

  (str, default: 'aot' ) –

  Prefix for generated code files. Defaults to "aot".

• attributes

  (dict[str, str] | None, default: None ) –

  Attributes for template values.

Attributes

id property

id: str

Return the operator ID.

name property

name: str

Return the operator name.

input_tensors property

input_tensors: list[AirTensor]

Return the input tensors.

output_tensors property

output_tensors: list[AirTensor]

Return the output tensors.

local_tensors property

local_tensors: list[AirTensor]

Get all local tensors used by the operator.

has_init property

has_init: bool

Whether this operator emits a per-instance _init function.

Operators whose initialization is a no-op may opt out by returning False. When an operator opts out, it must not emit an _init function/declaration in its templates, and the dispatch table records a NULL init slot (the model init loop skips the call while preserving the init callback contract). Defaults to True so existing operators keep emitting their own init unchanged.

Returns:

• bool –

  True if a per-instance _init function is emitted.

capabilities property

capabilities: OpCapability

Return the operator's effective optimization capabilities.

The effective set is the declared :attr:CAPABILITIES widened by the capabilities implied by the operator's runtime behavior, so the legacy seams remain the source of truth during migration and no operator has to declare a capability twice:

• not self.has_init implies :attr:OpCapability.STATELESS.
• A non-empty :meth:shared_kernel_helpers or :meth:shared_activation_helpers implies :attr:OpCapability.DESCRIPTOR_DRIVEN and :attr:OpCapability.SHARED_KERNEL.

Because some seams are dtype-dependent (e.g. fully-connected only routes through a shared kernel for int8 per-channel quantization), this is an instance property rather than a class constant: it reflects the concrete resolved operator. It performs no model mutation and does not change emitted code.

Returns:

• OpCapability –

  The effective :class:OpCapability flag set.

direct_dispatch property

direct_dispatch: tuple[str, str] | None

Return the direct shared-kernel call target for static scheduling.

Under static scheduling the model schedule can call a descriptor driven operator's shared kernel helper directly, eliding the thin per-node _run thunk (the main .text reclaim described in RFC 0003 Stage 4). This is only valid for operators whose _run body is exactly return helper(ctx, &desc); -- i.e. operators that route through a single shared kernel helper with the uniform helper(ctx, const desc_t *) signature (ADD/MUL/per-channel int8 FULLY_CONNECTED). Shared activation helpers marshal flattened arguments and keep mutable file-scope state, so they are intentionally excluded and keep their _run thunk.

Returns:

• tuple[str, str] | None –

  A (helper_symbol, descriptor_symbol) pair naming the C symbols

• tuple[str, str] | None –

  the schedule should call as helper_symbol(ctx, &descriptor_symbol),

• tuple[str, str] | None –

  or None when the operator is not directly dispatchable.

Functions

get_local_tensor

get_local_tensor(name: str) -> AirTensor

Get a local tensor by name.

tensor_zero_point staticmethod

tensor_zero_point(tensor: AirTensor, index: int = 0) -> int

Return a tensor's quantization zero point, defaulting to 0.

Parameters:

• tensor

  (AirTensor) –

  The tensor to inspect.

• index

  (int, default: 0 ) –

  The zero-point index to read (per-tensor quant uses 0).

Returns:

• int –

  The integer zero point, or 0 when the tensor is unquantized.

intern_constant_array

intern_constant_array(name: str, data: ndarray, *, role: AirTensorKind = AirTensorKind.CONSTANT, alignment: int | None = None) -> TensorId

Promote an operator-local constant array into a model tensor.

This routes operator metadata (e.g. per-channel quantization tables) through the same machinery as ordinary constants: memory planning, duplicate-constant interning (storage dedup), cold/staged residency, and uniform tensor-pointer resolution. The tensor is registered under op.named_tensors[name] so templates can reference it via ctx->tensor_ptrs[...] like any other tensor.

The operation is idempotent: a previously interned tensor with the same deterministic id is replaced.

Parameters:

• name

  (str) –

  Local tensor name; also used to build the deterministic id.

• data

  (ndarray) –

  Array payload, stored as a contiguous copy.

• role

  (AirTensorKind, default: CONSTANT ) –

  Tensor kind/bucket. Defaults to CONSTANT.

• alignment

  (int | None, default: None ) –

  Optional alignment hint in bytes.

Returns:

• TensorId ( TensorId ) –

  The id of the registered tensor.

validate

validate()

Validate the operator configuration.

resolve

resolve()

Public entry point—only runs once, even if called repeatedly.

This method validates the AIR operator and performs any model mutations needed.

on_resolve

on_resolve()

Override this in subclasses to do the actual work.

plan

plan()

Operator planning step.

shared_kernel_helpers

shared_kernel_helpers() -> list[KernelHelper]

Return the shared operator-kernel helpers this operator requires.

This mirrors :meth:shared_activation_helpers but for general operator kernels (e.g. elementwise ADD/MUL or FULLY_CONNECTED). An operator that routes its _run body through a shared runtime helper (instead of an inlined per-node call site) declares the helper descriptors here so the code generator emits each distinct helper exactly once. Helpers are keyed by signature, so distinct operator variants map to distinct helpers and never collide. The base implementation requires no shared kernels.

Returns:

• list[KernelHelper] –

  A list of :class:KernelHelper descriptors (see

• list[KernelHelper] –

  mod:helia_aot.aot.operators.kernel_runtime). Empty for operators

• list[KernelHelper] –

  that do not use a shared kernel helper.

shared_activation_helpers

shared_activation_helpers() -> list[dict[str, object]]

Return the shared LUT-activation helpers this operator requires.

Operators that emit their hot loop via a shared runtime helper (rather than an inlined per-node body) declare the helper descriptors here so the code generator can emit each distinct helper exactly once. The base implementation requires no shared helpers.

Returns:

• list[dict[str, object]] –

  A list of helper descriptor dicts (see

• list[dict[str, object]] –

  func:helia_aot.aot.operators.activation_runtime.activation_lut_helper).

• list[dict[str, object]] –

  Empty for operators that do not use a shared activation helper.

compute_values

compute_values() -> dict[str, str]

Compute the values for the operator source code template.

Args:

Returns:

• dict[str, str] –

  dict[str, str]: Dictionary of values for the operator template.

print_info

print_info(verbose: int = 0)

Debug print the template values for the operator.

emit

emit(save_path: Path)

Generate the source code for the operator.

This method should be overridden by subclasses.

Parameters:

• save_path

  (Path) –

  Path to save the generated source code.

Functions

run_once

run_once(method)

Decorator to run a method only once per instance.