emx-onnx-cgen 0.2.0__py3-none-any.whl

emx_onnx_cgen/lowering/average_pool.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import AveragePoolOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class _AveragePoolSpec:
+    batch: int
+    channels: int
+    in_h: int
+    in_w: int
+    out_h: int
+    out_w: int
+    kernel_h: int
+    kernel_w: int
+    stride_h: int
+    stride_w: int
+    pad_top: int
+    pad_left: int
+    pad_bottom: int
+    pad_right: int
+    count_include_pad: bool
+
+
+def _value_shape(graph: Graph, name: str, node: Node) -> tuple[int, ...]:
+    try:
+        return graph.find_value(name).type.shape
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing shape for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _value_dtype(graph: Graph, name: str, node: Node) -> str:
+    try:
+        return graph.find_value(name).type.dtype
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing dtype for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _resolve_average_pool_spec(graph: Graph, node: Node) -> _AveragePoolSpec:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("AveragePool must have 1 input and 1 output")
+    supported_attrs = {
+        "auto_pad",
+        "ceil_mode",
+        "count_include_pad",
+        "kernel_shape",
+        "pads",
+        "strides",
+    }
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("AveragePool has unsupported attributes")
+    auto_pad = node.attrs.get("auto_pad", b"NOTSET")
+    if isinstance(auto_pad, bytes):
+        auto_pad = auto_pad.decode("utf-8", errors="ignore")
+    if auto_pad not in ("", "NOTSET"):
+        raise UnsupportedOpError("AveragePool supports auto_pad=NOTSET only")
+    ceil_mode = int(node.attrs.get("ceil_mode", 0))
+    if ceil_mode != 0:
+        raise UnsupportedOpError("AveragePool supports ceil_mode=0 only")
+    count_include_pad = int(node.attrs.get("count_include_pad", 0))
+    if count_include_pad not in (0, 1):
+        raise UnsupportedOpError("AveragePool supports count_include_pad 0 or 1")
+    kernel_shape = node.attrs.get("kernel_shape")
+    if kernel_shape is None:
+        raise UnsupportedOpError("AveragePool requires kernel_shape")
+    kernel_shape = tuple(int(value) for value in kernel_shape)
+    if len(kernel_shape) != 2:
+        raise UnsupportedOpError("AveragePool expects 2D kernel_shape")
+    kernel_h, kernel_w = kernel_shape
+    strides = tuple(int(value) for value in node.attrs.get("strides", (1, 1)))
+    if len(strides) != 2:
+        raise UnsupportedOpError("AveragePool expects 2D strides")
+    pads = tuple(int(value) for value in node.attrs.get("pads", (0, 0, 0, 0)))
+    if len(pads) != 4:
+        raise UnsupportedOpError("AveragePool expects 4D pads")
+    pad_top, pad_left, pad_bottom, pad_right = pads
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) != 4:
+        raise UnsupportedOpError("AveragePool supports NCHW 2D inputs only")
+    batch, channels, in_h, in_w = input_shape
+    stride_h, stride_w = strides
+    out_h = (in_h + pad_top + pad_bottom - kernel_h) // stride_h + 1
+    out_w = (in_w + pad_left + pad_right - kernel_w) // stride_w + 1
+    if out_h < 0 or out_w < 0:
+        raise ShapeInferenceError(
+            "AveragePool output shape must be non-negative"
+        )
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    expected_output_shape = (batch, channels, out_h, out_w)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "AveragePool output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    return _AveragePoolSpec(
+        batch=batch,
+        channels=channels,
+        in_h=in_h,
+        in_w=in_w,
+        out_h=out_h,
+        out_w=out_w,
+        kernel_h=kernel_h,
+        kernel_w=kernel_w,
+        stride_h=stride_h,
+        stride_w=stride_w,
+        pad_top=pad_top,
+        pad_left=pad_left,
+        pad_bottom=pad_bottom,
+        pad_right=pad_right,
+        count_include_pad=bool(count_include_pad),
+    )
+
+
+def _resolve_global_average_pool_spec(graph: Graph, node: Node) -> _AveragePoolSpec:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("GlobalAveragePool must have 1 input and 1 output")
+    if node.attrs:
+        raise UnsupportedOpError("GlobalAveragePool has unsupported attributes")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) != 4:
+        raise UnsupportedOpError("GlobalAveragePool supports NCHW 2D inputs only")
+    batch, channels, in_h, in_w = input_shape
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    expected_output_shape = (batch, channels, 1, 1)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "GlobalAveragePool output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    return _AveragePoolSpec(
+        batch=batch,
+        channels=channels,
+        in_h=in_h,
+        in_w=in_w,
+        out_h=1,
+        out_w=1,
+        kernel_h=in_h,
+        kernel_w=in_w,
+        stride_h=1,
+        stride_w=1,
+        pad_top=0,
+        pad_left=0,
+        pad_bottom=0,
+        pad_right=0,
+        count_include_pad=False,
+    )
+
+
+@register_lowering("AveragePool")
+def lower_average_pool(graph: Graph, node: Node) -> AveragePoolOp:
+    op_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if op_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "AveragePool expects matching input/output dtypes, "
+            f"got {op_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "AveragePool supports float16, float, and double inputs only"
+        )
+    spec = _resolve_average_pool_spec(graph, node)
+    return AveragePoolOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        batch=spec.batch,
+        channels=spec.channels,
+        in_h=spec.in_h,
+        in_w=spec.in_w,
+        out_h=spec.out_h,
+        out_w=spec.out_w,
+        kernel_h=spec.kernel_h,
+        kernel_w=spec.kernel_w,
+        stride_h=spec.stride_h,
+        stride_w=spec.stride_w,
+        pad_top=spec.pad_top,
+        pad_left=spec.pad_left,
+        pad_bottom=spec.pad_bottom,
+        pad_right=spec.pad_right,
+        count_include_pad=spec.count_include_pad,
+        dtype=op_dtype,
+    )
+
+
+@register_lowering("GlobalAveragePool")
+def lower_global_average_pool(graph: Graph, node: Node) -> AveragePoolOp:
+    op_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if op_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "GlobalAveragePool expects matching input/output dtypes, "
+            f"got {op_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "GlobalAveragePool supports float16, float, and double inputs only"
+        )
+    spec = _resolve_global_average_pool_spec(graph, node)
+    return AveragePoolOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        batch=spec.batch,
+        channels=spec.channels,
+        in_h=spec.in_h,
+        in_w=spec.in_w,
+        out_h=spec.out_h,
+        out_w=spec.out_w,
+        kernel_h=spec.kernel_h,
+        kernel_w=spec.kernel_w,
+        stride_h=spec.stride_h,
+        stride_w=spec.stride_w,
+        pad_top=spec.pad_top,
+        pad_left=spec.pad_left,
+        pad_bottom=spec.pad_bottom,
+        pad_right=spec.pad_right,
+        count_include_pad=spec.count_include_pad,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/batch_normalization.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import BatchNormOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class _BatchNormSpec:
+    shape: tuple[int, ...]
+    channels: int
+    epsilon: float
+
+
+def _value_shape(graph: Graph, name: str, node: Node) -> tuple[int, ...]:
+    try:
+        return graph.find_value(name).type.shape
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing shape for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _value_dtype(graph: Graph, name: str, node: Node) -> str:
+    try:
+        return graph.find_value(name).type.dtype
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing dtype for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _node_dtype(graph: Graph, node: Node, *names: str) -> str:
+    dtypes = {_value_dtype(graph, name, node) for name in names}
+    if len(dtypes) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} expects matching dtypes, got {', '.join(sorted(dtypes))}"
+        )
+    return next(iter(dtypes))
+
+
+def _resolve_batch_norm_spec(graph: Graph, node: Node) -> _BatchNormSpec:
+    if len(node.inputs) != 5 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "BatchNormalization must have 5 inputs and 1 output"
+        )
+    supported_attrs = {
+        "epsilon",
+        "is_test",
+        "momentum",
+        "spatial",
+        "training_mode",
+    }
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("BatchNormalization has unsupported attributes")
+    is_test = int(node.attrs.get("is_test", 1))
+    if is_test != 1:
+        raise UnsupportedOpError("BatchNormalization supports is_test=1 only")
+    training_mode = int(node.attrs.get("training_mode", 0))
+    if training_mode != 0:
+        raise UnsupportedOpError("BatchNormalization supports training_mode=0 only")
+    spatial = int(node.attrs.get("spatial", 1))
+    if spatial != 1:
+        raise UnsupportedOpError("BatchNormalization supports spatial=1 only")
+    epsilon = float(node.attrs.get("epsilon", 1e-5))
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) < 2:
+        raise UnsupportedOpError(
+            "BatchNormalization expects input rank of at least 2"
+        )
+    channels = input_shape[1]
+    for name in node.inputs[1:]:
+        shape = _value_shape(graph, name, node)
+        if shape != (channels,):
+            raise ShapeInferenceError(
+                "BatchNormalization parameter shape must be "
+                f"({channels},), got {shape}"
+            )
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    if output_shape != input_shape:
+        raise ShapeInferenceError(
+            "BatchNormalization output shape must match input shape, "
+            f"got {output_shape}"
+        )
+    return _BatchNormSpec(
+        shape=input_shape,
+        channels=channels,
+        epsilon=epsilon,
+    )
+
+
+@register_lowering("BatchNormalization")
+def lower_batch_normalization(graph: Graph, node: Node) -> BatchNormOp:
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "BatchNormalization supports float16, float, and double inputs only"
+        )
+    spec = _resolve_batch_norm_spec(graph, node)
+    return BatchNormOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        bias=node.inputs[2],
+        mean=node.inputs[3],
+        variance=node.inputs[4],
+        output=node.outputs[0],
+        shape=spec.shape,
+        channels=spec.channels,
+        epsilon=spec.epsilon,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/cast.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import onnx
+
+from ..codegen.c_emitter import CastOp
+from ..dtypes import scalar_type_from_onnx
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import ensure_supported_dtype, value_dtype, value_shape
+from .registry import register_lowering
+
+
+@register_lowering("Cast")
+def lower_cast(graph: Graph, node: Node) -> CastOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Cast must have 1 input and 1 output")
+    if "to" not in node.attrs:
+        raise UnsupportedOpError("Cast requires a 'to' attribute")
+    target_onnx_dtype = int(node.attrs["to"])
+    target_dtype = scalar_type_from_onnx(target_onnx_dtype)
+    if target_dtype is None:
+        name = onnx.TensorProto.DataType.Name(target_onnx_dtype)
+        raise UnsupportedOpError(
+            f"Cast 'to' dtype {target_onnx_dtype} ({name}) is not supported"
+        )
+    target_dtype = ensure_supported_dtype(target_dtype)
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if output_dtype != target_dtype:
+        raise UnsupportedOpError(
+            "Cast output dtype must match 'to' attribute, "
+            f"got {output_dtype.onnx_name} and {target_dtype.onnx_name}"
+        )
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if input_shape != output_shape:
+        raise ShapeInferenceError("Cast input and output shapes must match")
+    return CastOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        shape=output_shape,
+        input_dtype=input_dtype,
+        dtype=output_dtype,
+    )
+
+
+@register_lowering("CastLike")
+def lower_castlike(graph: Graph, node: Node) -> CastOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("CastLike must have 2 inputs and 1 output")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    like_dtype = value_dtype(graph, node.inputs[1], node)
+    target_dtype = ensure_supported_dtype(like_dtype)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if output_dtype != target_dtype:
+        raise UnsupportedOpError(
+            "CastLike output dtype must match like input dtype, "
+            f"got {output_dtype.onnx_name} and {target_dtype.onnx_name}"
+        )
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if input_shape != output_shape:
+        raise ShapeInferenceError("CastLike input and output shapes must match")
+    return CastOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        shape=output_shape,
+        input_dtype=input_dtype,
+        dtype=output_dtype,
+    )

emx_onnx_cgen/lowering/common.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+
+
+def ensure_supported_dtype(dtype: ScalarType) -> ScalarType:
+    if not isinstance(dtype, ScalarType):
+        raise UnsupportedOpError(f"Unsupported dtype {dtype}")
+    return dtype
+
+
+def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType:
+    try:
+        value = graph.find_value(name)
+    except KeyError as exc:
+        op_type = node.op_type if node is not None else "unknown"
+        raise ShapeInferenceError(
+            f"Missing dtype for value '{name}' in op {op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+    return ensure_supported_dtype(value.type.dtype)
+
+
+def value_shape(graph: Graph, name: str, node: Node | None = None) -> tuple[int, ...]:
+    try:
+        return graph.find_value(name).type.shape
+    except KeyError as exc:
+        op_type = node.op_type if node is not None else "unknown"
+        raise ShapeInferenceError(
+            f"Missing shape for value '{name}' in op {op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def node_dtype(graph: Graph, node: Node, *names: str) -> ScalarType:
+    filtered = [name for name in names if name]
+    if not filtered:
+        raise UnsupportedOpError(
+            f"{node.op_type} expects at least one typed input or output"
+        )
+    dtypes = {value_dtype(graph, name, node) for name in filtered}
+    if len(dtypes) != 1:
+        dtype_names = ", ".join(dtype.onnx_name for dtype in sorted(dtypes, key=str))
+        raise UnsupportedOpError(
+            f"{node.op_type} expects matching dtypes, got {dtype_names}"
+        )
+    return next(iter(dtypes))
+
+
+def shape_product(shape: tuple[int, ...]) -> int:
+    if not shape:
+        return 1
+    product = 1
+    for dim in shape:
+        if dim < 0:
+            raise ShapeInferenceError("Dynamic dims are not supported")
+        if dim == 0:
+            return 0
+        product *= dim
+    return product
+
+
+def optional_name(names: Sequence[str], index: int) -> str | None:
+    if index >= len(names):
+        return None
+    name = names[index]
+    return name or None

emx_onnx_cgen/lowering/concat.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import ConcatOp
+from ..errors import UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import node_dtype as _node_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+from ..validation import validate_concat_shapes
+
+
+@register_lowering("Concat")
+def lower_concat(graph: Graph, node: Node) -> ConcatOp:
+    if len(node.inputs) < 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Concat must have at least 1 input and 1 output")
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    input_shapes = tuple(_value_shape(graph, name, node) for name in node.inputs)
+    axis = validate_concat_shapes(
+        input_shapes,
+        output_shape,
+        int(node.attrs.get("axis", 0)),
+    )
+    return ConcatOp(
+        inputs=node.inputs,
+        output=node.outputs[0],
+        axis=axis,
+        input_shapes=input_shapes,
+        output_shape=output_shape,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/constant_of_shape.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from onnx import numpy_helper
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import ConstantOfShapeOp
+from ..dtypes import scalar_type_from_onnx
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .registry import register_lowering
+
+
+def _value_shape(graph: Graph, name: str, node: Node) -> tuple[int, ...]:
+    try:
+        return graph.find_value(name).type.shape
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing shape for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _value_dtype(graph: Graph, name: str, node: Node) -> ScalarType:
+    try:
+        return graph.find_value(name).type.dtype
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing dtype for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _parse_value_attr(node: Node) -> tuple[ScalarType, float | int | bool]:
+    value_attr = node.attrs.get("value")
+    if value_attr is None:
+        return ScalarType.F32, 0.0
+    dtype = scalar_type_from_onnx(value_attr.data_type)
+    if dtype is None:
+        raise UnsupportedOpError(
+            f"ConstantOfShape has unsupported value dtype {value_attr.data_type}"
+        )
+    data = numpy_helper.to_array(value_attr)
+    if data.size != 1:
+        raise UnsupportedOpError("ConstantOfShape value must be a scalar")
+    return dtype, data.reshape(-1)[0].item()
+
+
+@register_lowering("ConstantOfShape")
+def lower_constant_of_shape(graph: Graph, node: Node) -> ConstantOfShapeOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("ConstantOfShape must have 1 input and 1 output")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) != 1:
+        raise UnsupportedOpError("ConstantOfShape expects a 1D shape input")
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    if input_shape[0] != len(output_shape):
+        raise ShapeInferenceError(
+            "ConstantOfShape input length must match output rank"
+        )
+    for dim in output_shape:
+        if dim < 0:
+            raise ShapeInferenceError("Dynamic dims are not supported")
+    input_dtype = _value_dtype(graph, node.inputs[0], node)
+    if input_dtype != ScalarType.I64:
+        raise UnsupportedOpError(
+            "ConstantOfShape expects int64 shape input, "
+            f"got {input_dtype.onnx_name}"
+        )
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    value_dtype, value = _parse_value_attr(node)
+    if output_dtype != value_dtype:
+        raise UnsupportedOpError(
+            "ConstantOfShape output dtype must match value dtype, "
+            f"got {output_dtype.onnx_name} and {value_dtype.onnx_name}"
+        )
+    return ConstantOfShapeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        shape=output_shape,
+        value=value,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )