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

emx_onnx_cgen/lowering/unsqueeze.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import ReshapeOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, 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 _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _validate_shape(shape: tuple[int, ...], node: Node, label: str) -> None:
+    for dim in shape:
+        if dim < 0:
+            raise ShapeInferenceError(
+                f"{node.op_type} does not support dynamic dims in {label}"
+            )
+
+
+def _normalize_axes(
+    axes: list[int], output_rank: int, node: Node
+) -> tuple[int, ...]:
+    normalized: list[int] = []
+    for axis in axes:
+        if axis < 0:
+            axis += output_rank
+        if axis < 0 or axis >= output_rank:
+            raise ShapeInferenceError(
+                f"{node.op_type} axis {axis} is out of range for rank {output_rank}"
+            )
+        normalized.append(axis)
+    if len(set(normalized)) != len(normalized):
+        raise ShapeInferenceError(f"{node.op_type} axes must be unique")
+    return tuple(sorted(normalized))
+
+
+def _resolve_axes(graph: Graph, node: Node) -> tuple[int, ...] | None:
+    axes_attr = node.attrs.get("axes")
+    axes_values: list[int] | None = None
+    if len(node.inputs) == 2:
+        axes_initializer = _find_initializer(graph, node.inputs[1])
+        if axes_initializer is not None:
+            if axes_initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
+                raise UnsupportedOpError(
+                    "Unsqueeze axes input must be int64 or int32, "
+                    f"got {axes_initializer.type.dtype.onnx_name}"
+                )
+            axes_values = [int(value) for value in axes_initializer.data.reshape(-1)]
+    elif axes_attr is not None:
+        axes_values = [int(value) for value in axes_attr]
+    if axes_values is None and axes_attr is None and len(node.inputs) != 2:
+        raise UnsupportedOpError("Unsqueeze requires axes")
+    if axes_values is None:
+        return None
+    if not axes_values:
+        raise UnsupportedOpError("Unsqueeze requires non-empty axes")
+    return tuple(axes_values)
+
+
+def _expected_output_shape(
+    input_shape: tuple[int, ...], axes: tuple[int, ...], node: Node
+) -> tuple[int, ...]:
+    output_rank = len(input_shape) + len(axes)
+    normalized_axes = _normalize_axes(list(axes), output_rank, node)
+    output_dims: list[int] = []
+    input_index = 0
+    for axis in range(output_rank):
+        if axis in normalized_axes:
+            output_dims.append(1)
+        else:
+            output_dims.append(input_shape[input_index])
+            input_index += 1
+    return tuple(output_dims)
+
+
+@register_lowering("Unsqueeze")
+def lower_unsqueeze(graph: Graph, node: Node) -> ReshapeOp:
+    if len(node.outputs) != 1 or len(node.inputs) not in {1, 2}:
+        raise UnsupportedOpError("Unsqueeze must have 1 or 2 inputs and 1 output")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    _validate_shape(input_shape, node, "input")
+    _validate_shape(output_shape, node, "output")
+    input_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Unsqueeze expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    axes = _resolve_axes(graph, node)
+    if axes is None:
+        if len(node.inputs) == 2:
+            axes_dtype = _value_dtype(graph, node.inputs[1], node)
+            if axes_dtype not in {ScalarType.I64, ScalarType.I32}:
+                raise UnsupportedOpError(
+                    "Unsqueeze axes input must be int64 or int32, "
+                    f"got {axes_dtype.onnx_name}"
+                )
+        if len(output_shape) <= len(input_shape):
+            raise ShapeInferenceError(
+                "Unsqueeze output rank must exceed input rank"
+            )
+        input_index = 0
+        for dim in output_shape:
+            if input_index < len(input_shape) and dim == input_shape[input_index]:
+                input_index += 1
+                continue
+            if dim != 1:
+                raise ShapeInferenceError(
+                    "Unsqueeze output shape must insert ones only"
+                )
+        if input_index != len(input_shape):
+            raise ShapeInferenceError(
+                "Unsqueeze output shape must contain input shape in order"
+            )
+    else:
+        expected_shape = _expected_output_shape(input_shape, axes, node)
+        if expected_shape != output_shape:
+            raise ShapeInferenceError(
+                "Unsqueeze output shape must be "
+                f"{expected_shape}, got {output_shape}"
+            )
+    return ReshapeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        dtype=input_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/variadic.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from shared.scalar_functions import ScalarFunction
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import MultiInputBinaryOp
+from ..errors import UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..lowering.common import node_dtype, value_dtype, value_shape
+from ..lowering.registry import register_lowering
+from ..ops import binary_op_symbol
+
+VARIADIC_OP_FUNCTIONS: dict[str, ScalarFunction] = {
+    "Sum": ScalarFunction.ADD,
+    "Mean": ScalarFunction.MEAN,
+    "Max": ScalarFunction.MAXIMUM,
+    "Min": ScalarFunction.MINIMUM,
+    "And": ScalarFunction.LOGICAL_AND,
+    "Or": ScalarFunction.LOGICAL_OR,
+    "Xor": ScalarFunction.LOGICAL_XOR,
+    "BitwiseAnd": ScalarFunction.BITWISE_AND,
+    "BitwiseOr": ScalarFunction.BITWISE_OR,
+    "BitwiseXor": ScalarFunction.BITWISE_XOR,
+}
+
+BINARY_ONLY_OPS = {
+    "And",
+    "Or",
+    "Xor",
+    "BitwiseAnd",
+    "BitwiseOr",
+    "BitwiseXor",
+}
+
+
+def _validate_inputs(
+    graph: Graph, node: Node, *, function: ScalarFunction
+) -> tuple[ScalarType, tuple[int, ...]]:
+    if len(node.outputs) != 1:
+        raise UnsupportedOpError(f"{node.op_type} must have 1 output")
+    if node.op_type in BINARY_ONLY_OPS:
+        if len(node.inputs) != 2:
+            raise UnsupportedOpError(
+                f"{node.op_type} must have exactly 2 inputs"
+            )
+    elif len(node.inputs) < 2:
+        raise UnsupportedOpError(
+            f"{node.op_type} must have at least 2 inputs"
+        )
+    for name in node.inputs:
+        if not name:
+            raise UnsupportedOpError(f"{node.op_type} input must be provided")
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if op_dtype != output_dtype:
+        raise UnsupportedOpError(
+            f"{node.op_type} expects matching input/output dtypes, "
+            f"got {op_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    output_shape = value_shape(graph, node.outputs[0], node)
+    for name in node.inputs:
+        input_shape = value_shape(graph, name, node)
+        if input_shape != output_shape:
+            raise UnsupportedOpError(
+                f"{node.op_type} expects identical input/output shapes"
+            )
+    op_spec = binary_op_symbol(function, dtype=op_dtype, validate_attrs=False)
+    if op_spec is None:
+        raise UnsupportedOpError(
+            f"{node.op_type} does not support dtype {op_dtype.onnx_name}"
+        )
+    return op_dtype, output_shape
+
+
+def _lower_variadic(graph: Graph, node: Node) -> MultiInputBinaryOp:
+    function = VARIADIC_OP_FUNCTIONS[node.op_type]
+    op_dtype, output_shape = _validate_inputs(graph, node, function=function)
+    op_spec = binary_op_symbol(function, dtype=op_dtype, validate_attrs=False)
+    if op_spec is None:
+        raise UnsupportedOpError(
+            f"{node.op_type} does not support dtype {op_dtype.onnx_name}"
+        )
+    return MultiInputBinaryOp(
+        inputs=tuple(node.inputs),
+        output=node.outputs[0],
+        function=function,
+        operator_kind=op_spec.kind,
+        shape=output_shape,
+        dtype=op_dtype,
+        input_dtype=op_dtype,
+    )
+
+
+for _op_type in VARIADIC_OP_FUNCTIONS:
+    register_lowering(_op_type)(_lower_variadic)

emx_onnx_cgen/lowering/where.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import WhereOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+def _broadcast_shape(shapes: tuple[tuple[int, ...], ...], node: Node) -> tuple[int, ...]:
+    if not shapes:
+        return ()
+    max_rank = max(len(shape) for shape in shapes)
+    padded = [
+        (1,) * (max_rank - len(shape)) + shape
+        for shape in shapes
+    ]
+    broadcast: list[int] = []
+    for dims in zip(*padded):
+        dim = max(dims)
+        if any(item not in (1, dim) for item in dims):
+            raise ShapeInferenceError(
+                f"{node.op_type} inputs must be broadcastable, got {shapes}"
+            )
+        broadcast.append(dim)
+    return tuple(broadcast)
+
+
+@register_lowering("Where")
+def lower_where(graph: Graph, node: Node) -> WhereOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Where must have 3 inputs and 1 output")
+    condition_name, x_name, y_name = node.inputs
+    output_name = node.outputs[0]
+    condition_dtype = _value_dtype(graph, condition_name, node)
+    if condition_dtype != ScalarType.BOOL:
+        raise UnsupportedOpError(
+            f"Where expects bool condition, got {condition_dtype.onnx_name}"
+        )
+    x_dtype = _value_dtype(graph, x_name, node)
+    y_dtype = _value_dtype(graph, y_name, node)
+    output_dtype = _value_dtype(graph, output_name, node)
+    if x_dtype != y_dtype or output_dtype != x_dtype:
+        raise UnsupportedOpError(
+            "Where expects matching input/output dtypes, "
+            f"got {x_dtype.onnx_name}, {y_dtype.onnx_name}, {output_dtype.onnx_name}"
+        )
+    condition_shape = _value_shape(graph, condition_name, node)
+    x_shape = _value_shape(graph, x_name, node)
+    y_shape = _value_shape(graph, y_name, node)
+    output_shape = _value_shape(graph, output_name, node)
+    broadcast_shape = _broadcast_shape(
+        (condition_shape, x_shape, y_shape),
+        node,
+    )
+    if output_shape != broadcast_shape:
+        raise ShapeInferenceError(
+            f"Where output shape must be {broadcast_shape}, got {output_shape}"
+        )
+    return WhereOp(
+        condition=condition_name,
+        input_x=x_name,
+        input_y=y_name,
+        output=output_name,
+        condition_shape=condition_shape,
+        x_shape=x_shape,
+        y_shape=y_shape,
+        output_shape=output_shape,
+        dtype=output_dtype,
+    )

emx_onnx_cgen/onnx_import.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+from typing import Iterable, Mapping
+
+import onnx
+import numpy as np
+from onnx import helper, numpy_helper, shape_inference
+
+from shared.scalar_types import ScalarType
+
+from .dtypes import scalar_type_from_onnx
+from .errors import ShapeInferenceError, UnsupportedOpError
+from .ir.model import Graph, Initializer, Node, TensorType, Value
+
+
+def _normalize_initializer_data(dtype: ScalarType, data: object) -> np.ndarray:
+    if isinstance(data, (onnx.TensorProto, onnx.SparseTensorProto)):
+        array = numpy_helper.to_array(data)
+    elif isinstance(data, np.ndarray):
+        array = data
+    else:
+        array = np.array(data)
+    return array.astype(dtype.np_dtype, copy=False)
+
+
+def _format_elem_type(elem_type: int) -> str:
+    try:
+        name = onnx.TensorProto.DataType.Name(elem_type)
+    except ValueError:
+        name = "UNKNOWN"
+    return f"{elem_type} ({name})"
+
+
+def _unsupported_value_type(value_info: onnx.ValueInfoProto) -> UnsupportedOpError:
+    value_kind = value_info.type.WhichOneof("value")
+    if value_kind is None:
+        value_kind = "unknown"
+    return UnsupportedOpError(
+        f"Unsupported value type '{value_kind}' for '{value_info.name}'. "
+        "Hint: export the model with tensor inputs/outputs."
+    )
+
+
+def _tensor_type(
+    value_info: onnx.ValueInfoProto,
+    *,
+    dim_param_override: tuple[str | None, ...] | None = None,
+) -> TensorType:
+    if value_info.type.WhichOneof("value") != "tensor_type":
+        raise _unsupported_value_type(value_info)
+    tensor_type = value_info.type.tensor_type
+    if not tensor_type.HasField("elem_type"):
+        raise ShapeInferenceError(f"Missing elem_type for tensor '{value_info.name}'")
+    dtype = scalar_type_from_onnx(tensor_type.elem_type)
+    if dtype is None:
+        raise UnsupportedOpError(
+            "Unsupported elem_type "
+            f"{_format_elem_type(tensor_type.elem_type)} for tensor '{value_info.name}'."
+        )
+    shape = []
+    dim_params = []
+    for dim_index, dim in enumerate(tensor_type.shape.dim):
+        dim_param = dim.dim_param if dim.HasField("dim_param") else ""
+        if (
+            dim_param_override is not None
+            and dim_index < len(dim_param_override)
+            and dim_param_override[dim_index]
+        ):
+            dim_param = dim_param_override[dim_index] or ""
+        dim_params.append(dim_param or None)
+        if not dim.HasField("dim_value"):
+            if dim_param:
+                shape.append(1)
+                continue
+            raise ShapeInferenceError(f"Dynamic dim for tensor '{value_info.name}'")
+        shape.append(dim.dim_value)
+    return TensorType(
+        dtype=dtype,
+        shape=tuple(shape),
+        dim_params=tuple(dim_params),
+    )
+
+
+def _values(
+    value_infos: Iterable[onnx.ValueInfoProto],
+    *,
+    dim_param_by_name: Mapping[str, tuple[str | None, ...]] | None = None,
+) -> tuple[Value, ...]:
+    dim_param_by_name = dim_param_by_name or {}
+    return tuple(
+        Value(
+            name=vi.name,
+            type=_tensor_type(
+                vi, dim_param_override=dim_param_by_name.get(vi.name)
+            ),
+        )
+        for vi in value_infos
+    )
+
+
+def _collect_dim_params(
+    value_infos: Iterable[onnx.ValueInfoProto],
+) -> dict[str, tuple[str | None, ...]]:
+    dim_params: dict[str, tuple[str | None, ...]] = {}
+    for value_info in value_infos:
+        dims = []
+        for dim in value_info.type.tensor_type.shape.dim:
+            dim_param = dim.dim_param if dim.HasField("dim_param") else ""
+            dims.append(dim_param or None)
+        if any(dims):
+            dim_params[value_info.name] = tuple(dims)
+    return dim_params
+
+
+def _initializer(value: onnx.TensorProto) -> Initializer:
+    dtype = scalar_type_from_onnx(value.data_type)
+    if dtype is None:
+        raise UnsupportedOpError(
+            "Unsupported elem_type "
+            f"{_format_elem_type(value.data_type)} for initializer '{value.name}'. "
+            "Hint: export the model with float32 initializers."
+        )
+    data = _normalize_initializer_data(dtype, value)
+    return Initializer(
+        name=value.name,
+        type=TensorType(
+            dtype=dtype,
+            shape=tuple(data.shape),
+            dim_params=(None,) * len(data.shape),
+        ),
+        data=data,
+    )
+
+
+def _node_attrs(node: onnx.NodeProto) -> dict[str, object]:
+    return {attr.name: helper.get_attribute_value(attr) for attr in node.attribute}
+
+
+def _constant_initializer(node: onnx.NodeProto) -> Initializer:
+    if len(node.output) != 1:
+        raise UnsupportedOpError("Constant must have exactly one output")
+    attrs = _node_attrs(node)
+    output_name = node.output[0]
+    if "value" in attrs:
+        tensor = attrs["value"]
+        dtype = scalar_type_from_onnx(tensor.data_type)
+        if dtype is None:
+            raise UnsupportedOpError(
+                "Unsupported elem_type "
+                f"{_format_elem_type(tensor.data_type)} for Constant '{output_name}'."
+            )
+        data = _normalize_initializer_data(dtype, tensor)
+        return Initializer(
+            name=output_name,
+            type=TensorType(
+                dtype=dtype,
+                shape=tuple(data.shape),
+                dim_params=(None,) * len(data.shape),
+            ),
+            data=data,
+        )
+    if "sparse_value" in attrs:
+        tensor = attrs["sparse_value"]
+        dtype = scalar_type_from_onnx(tensor.values.data_type)
+        if dtype is None:
+            raise UnsupportedOpError(
+                "Unsupported elem_type "
+                f"{_format_elem_type(tensor.values.data_type)} for Constant '{output_name}'."
+            )
+        data = _normalize_initializer_data(dtype, tensor)
+        return Initializer(
+            name=output_name,
+            type=TensorType(
+                dtype=dtype,
+                shape=tuple(data.shape),
+                dim_params=(None,) * len(data.shape),
+            ),
+            data=data,
+        )
+    if "value_float" in attrs or "value_floats" in attrs:
+        values = attrs.get("value_floats", attrs.get("value_float"))
+        data = _normalize_initializer_data(ScalarType.F32, values)
+        return Initializer(
+            name=output_name,
+        type=TensorType(
+            dtype=ScalarType.F32,
+            shape=tuple(data.shape),
+            dim_params=(None,) * len(data.shape),
+        ),
+        data=data,
+    )
+    if "value_int" in attrs or "value_ints" in attrs:
+        values = attrs.get("value_ints", attrs.get("value_int"))
+        data = _normalize_initializer_data(ScalarType.I64, values)
+        return Initializer(
+            name=output_name,
+        type=TensorType(
+            dtype=ScalarType.I64,
+            shape=tuple(data.shape),
+            dim_params=(None,) * len(data.shape),
+        ),
+        data=data,
+    )
+    if "value_string" in attrs or "value_strings" in attrs:
+        raise UnsupportedOpError(
+            f"Constant '{output_name}' has unsupported string values"
+        )
+    raise UnsupportedOpError(f"Constant '{output_name}' requires a value attribute")
+
+
+def import_onnx(model: onnx.ModelProto) -> Graph:
+    dim_param_by_name = _collect_dim_params(
+        tuple(model.graph.input) + tuple(model.graph.output)
+    )
+    try:
+        model = shape_inference.infer_shapes(model, data_prop=True)
+    except Exception as exc:  # pragma: no cover - onnx inference errors
+        raise ShapeInferenceError("ONNX shape inference failed") from exc
+    graph = model.graph
+    base_initializers = [_initializer(value) for value in graph.initializer]
+    constant_initializers: list[Initializer] = []
+    input_names = {value_info.name for value_info in graph.input}
+    output_names = {value_info.name for value_info in graph.output}
+    nodes: list[Node] = []
+    for node in graph.node:
+        if node.op_type == "Constant":
+            constant_initializers.append(_constant_initializer(node))
+            continue
+        nodes.append(
+            Node(
+                op_type=node.op_type,
+                name=node.name or None,
+                inputs=tuple(node.input),
+                outputs=tuple(node.output),
+                attrs=_node_attrs(node),
+            )
+        )
+    initializers = tuple(base_initializers + constant_initializers)
+    initializer_names = {initializer.name for initializer in initializers}
+    inputs = _values(
+        (
+            value_info
+            for value_info in graph.input
+            if value_info.name not in initializer_names
+        ),
+        dim_param_by_name=dim_param_by_name,
+    )
+    outputs = _values(graph.output, dim_param_by_name=dim_param_by_name)
+    values = _values(
+        value_info
+        for value_info in graph.value_info
+        if value_info.name
+        not in initializer_names | input_names | output_names
+    )
+    return Graph(
+        inputs=inputs,
+        outputs=outputs,
+        nodes=nodes,
+        initializers=initializers,
+        values=values,
+    )