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

emx_onnx_cgen/lowering/common.py

@@ -5,7 +5,7 @@ from collections.abc import Sequence
 from shared.scalar_types import ScalarType
 
 from ..errors import ShapeInferenceError, UnsupportedOpError
-from ..ir.model import Graph, Node
+from ..ir.model import Graph, Initializer, Node
 
 
 def ensure_supported_dtype(dtype: ScalarType) -> ScalarType:
@@ -14,6 +14,17 @@ def ensure_supported_dtype(dtype: ScalarType) -> ScalarType:
     return dtype
 
 
+def onnx_opset_version(graph: Graph, domain: str = "") -> int | None:
+    if domain in {"", "ai.onnx"}:
+        domains = {"", "ai.onnx"}
+    else:
+        domains = {domain}
+    for opset_domain, version in graph.opset_imports:
+        if opset_domain in domains:
+            return int(version)
+    return None
+
+
 def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType:
     try:
         value = graph.find_value(name)
@@ -28,13 +39,379 @@ def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType
 
 def value_shape(graph: Graph, name: str, node: Node | None = None) -> tuple[int, ...]:
     try:
-        return graph.find_value(name).type.shape
+        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 shape for value '{name}' in op {op_type}. "
             "Hint: run ONNX shape inference or export with static shapes."
         ) from exc
+    if any(value.type.dim_params):
+        resolved = _resolve_value_shape(graph, name, node)
+        if resolved is not None:
+            return resolved
+        return value.type.shape
+    return value.type.shape
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _find_node_by_output(graph: Graph, name: str) -> Node | None:
+    for node in graph.nodes:
+        if name in node.outputs:
+            return node
+    return None
+
+
+def _shape_values_from_shape_node(
+    graph: Graph, shape_node: Node, node: Node | None
+) -> list[int]:
+    if len(shape_node.inputs) != 1 or len(shape_node.outputs) != 1:
+        raise UnsupportedOpError("Shape must have 1 input and 1 output")
+    source_shape = value_shape(graph, shape_node.inputs[0], node)
+    start = int(shape_node.attrs.get("start", 0))
+    end = int(shape_node.attrs.get("end", len(source_shape)))
+    if start < 0:
+        start += len(source_shape)
+    if end < 0:
+        end += len(source_shape)
+    start = max(start, 0)
+    end = min(end, len(source_shape))
+    if start > end:
+        return []
+    return list(source_shape[start:end])
+
+
+def _shape_values_from_initializer(
+    graph: Graph,
+    name: str,
+) -> list[int] | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    if initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            "Reshape expects int64 or int32 shape input, "
+            f"got {initializer.type.dtype.onnx_name}"
+        )
+    return [int(value) for value in initializer.data.reshape(-1)]
+
+
+def _shape_values_from_input(
+    graph: Graph,
+    name: str,
+    node: Node | None,
+    *,
+    _visited: set[str] | None = None,
+) -> list[int] | None:
+    if _visited is None:
+        _visited = set()
+    if name in _visited:
+        return None
+    _visited.add(name)
+    try:
+        shape_values = _shape_values_from_initializer(graph, name)
+        if shape_values is not None:
+            return shape_values
+        source_node = _find_node_by_output(graph, name)
+        if source_node is None:
+            return None
+        if source_node.op_type == "Shape":
+            return _shape_values_from_shape_node(graph, source_node, node)
+        if source_node.op_type == "Concat":
+            axis = int(source_node.attrs.get("axis", 0))
+            if axis != 0:
+                raise UnsupportedOpError("Reshape shape concat must use axis 0")
+            values: list[int] = []
+            for input_name in source_node.inputs:
+                input_values = _shape_values_from_input(
+                    graph,
+                    input_name,
+                    node,
+                    _visited=_visited,
+                )
+                if input_values is None:
+                    return None
+                values.extend(input_values)
+            return values
+        if source_node.op_type == "Cast":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Cast must have 1 input and 1 output")
+            return _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+        if source_node.op_type == "Unsqueeze":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Unsqueeze must have 1 input and 1 output")
+            return _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+        if source_node.op_type == "Identity":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Identity must have 1 input and 1 output")
+            return _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+        if source_node.op_type in {"Equal", "And", "Or", "Div", "Mod"}:
+            if len(source_node.inputs) != 2 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError(
+                    f"{source_node.op_type} must have 2 inputs and 1 output"
+                )
+            left = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            right = _shape_values_from_input(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            if left is None or right is None:
+                return None
+            if len(left) == 1 and len(right) != 1:
+                left = left * len(right)
+            if len(right) == 1 and len(left) != 1:
+                right = right * len(left)
+            if len(left) != len(right):
+                return None
+            if source_node.op_type == "Equal":
+                return [1 if l == r else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "And":
+                return [1 if (l and r) else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "Or":
+                return [1 if (l or r) else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "Div":
+                return [int(l / r) if r != 0 else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "Mod":
+                return [l % r if r != 0 else 0 for l, r in zip(left, right)]
+        if source_node.op_type == "Not":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Not must have 1 input and 1 output")
+            values = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            if values is None:
+                return None
+            return [0 if value else 1 for value in values]
+        if source_node.op_type == "Where":
+            if len(source_node.inputs) != 3 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Where must have 3 inputs and 1 output")
+            condition = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            if condition is None:
+                return None
+            on_true = _shape_values_from_input(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            on_false = _shape_values_from_input(
+                graph,
+                source_node.inputs[2],
+                node,
+                _visited=_visited,
+            )
+            if on_true is None or on_false is None:
+                return None
+            if len(condition) == 1:
+                condition = condition * max(len(on_true), len(on_false))
+            if len(on_true) == 1 and len(condition) != 1:
+                on_true = on_true * len(condition)
+            if len(on_false) == 1 and len(condition) != 1:
+                on_false = on_false * len(condition)
+            if not (len(condition) == len(on_true) == len(on_false)):
+                return None
+            return [
+                t if cond else f
+                for cond, t, f in zip(condition, on_true, on_false)
+            ]
+        return None
+    finally:
+        _visited.remove(name)
+
+
+def _broadcast_shapes(
+    left: tuple[int, ...],
+    right: tuple[int, ...],
+) -> tuple[int, ...] | None:
+    result = []
+    left_rev = list(reversed(left))
+    right_rev = list(reversed(right))
+    for index in range(max(len(left_rev), len(right_rev))):
+        left_dim = left_rev[index] if index < len(left_rev) else 1
+        right_dim = right_rev[index] if index < len(right_rev) else 1
+        if left_dim == right_dim:
+            result.append(left_dim)
+        elif left_dim == 1:
+            result.append(right_dim)
+        elif right_dim == 1:
+            result.append(left_dim)
+        else:
+            return None
+    return tuple(reversed(result))
+
+
+def _resolve_value_shape(
+    graph: Graph,
+    name: str,
+    node: Node | None,
+    *,
+    _visited: set[str] | None = None,
+) -> tuple[int, ...] | None:
+    if _visited is None:
+        _visited = set()
+    if name in _visited:
+        return None
+    _visited.add(name)
+    try:
+        value = graph.find_value(name)
+        shape = value.type.shape
+        if not any(value.type.dim_params):
+            return shape
+        source_node = _find_node_by_output(graph, name)
+        if source_node is None:
+            return None
+        if source_node.op_type == "Expand":
+            if len(source_node.inputs) != 2 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Expand must have 2 inputs and 1 output")
+            shape_values = _shape_values_from_input(
+                graph, source_node.inputs[1], node
+            )
+            if shape_values is not None and all(dim >= 0 for dim in shape_values):
+                return tuple(shape_values)
+            return None
+        if source_node.op_type == "Reshape":
+            if len(source_node.inputs) != 2 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Reshape must have 2 inputs and 1 output")
+            shape_values = _shape_values_from_input(
+                graph, source_node.inputs[1], node
+            )
+            if shape_values is None:
+                return None
+            allowzero = int(source_node.attrs.get("allowzero", 0))
+            input_shape = _resolve_value_shape(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            if input_shape is None:
+                return None
+            output_dims: list[int] = []
+            unknown_index: int | None = None
+            known_product = 1
+            contains_zero = False
+            for index, dim in enumerate(shape_values):
+                if dim == -1:
+                    if unknown_index is not None:
+                        return None
+                    unknown_index = len(output_dims)
+                    output_dims.append(-1)
+                else:
+                    if dim == 0:
+                        contains_zero = True
+                        if allowzero == 0:
+                            if index >= len(input_shape):
+                                return None
+                            dim = input_shape[index]
+                    if dim < 0:
+                        return None
+                    output_dims.append(dim)
+                    known_product *= dim
+            if allowzero == 1 and contains_zero and unknown_index is not None:
+                return None
+            input_product = shape_product(input_shape)
+            if unknown_index is not None:
+                if known_product == 0:
+                    if input_product != 0:
+                        return None
+                    output_dims[unknown_index] = 0
+                else:
+                    if input_product % known_product != 0:
+                        return None
+                    output_dims[unknown_index] = input_product // known_product
+            return tuple(output_dims)
+        if source_node.op_type in {
+            "Add",
+            "Sub",
+            "Mul",
+            "Div",
+            "Pow",
+            "Mod",
+            "And",
+            "Or",
+            "Xor",
+            "Equal",
+            "Greater",
+            "Less",
+            "GreaterOrEqual",
+            "LessOrEqual",
+        }:
+            if len(source_node.inputs) != 2 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError(
+                    f"{source_node.op_type} must have 2 inputs and 1 output"
+                )
+            left = _resolve_value_shape(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            right = _resolve_value_shape(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            if left is None or right is None:
+                return None
+            return _broadcast_shapes(left, right)
+        if source_node.op_type == "Where":
+            if len(source_node.inputs) != 3 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Where must have 3 inputs and 1 output")
+            on_true = _resolve_value_shape(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            on_false = _resolve_value_shape(
+                graph,
+                source_node.inputs[2],
+                node,
+                _visited=_visited,
+            )
+            if on_true is None or on_false is None:
+                return None
+            return _broadcast_shapes(on_true, on_false)
+        return None
+    finally:
+        _visited.remove(name)
 
 
 def node_dtype(graph: Graph, node: Node, *names: str) -> ScalarType:

emx_onnx_cgen/lowering/conv_transpose.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import ConvTransposeOp
+from ..errors import ShapeInferenceError, 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
+
+
+@dataclass(frozen=True)
+class ConvTransposeSpec:
+    batch: int
+    in_channels: int
+    out_channels: int
+    spatial_rank: int
+    in_spatial: tuple[int, ...]
+    out_spatial: tuple[int, ...]
+    kernel_shape: tuple[int, ...]
+    strides: tuple[int, ...]
+    pads: tuple[int, ...]
+    dilations: tuple[int, ...]
+    output_padding: tuple[int, ...]
+    group: int
+
+
+def _split_padding(
+    total_padding: int, auto_pad: str, *, dim: int
+) -> tuple[int, int]:
+    if total_padding < 0:
+        raise ShapeInferenceError(
+            "ConvTranspose output shape must be fully defined and non-negative"
+        )
+    pad_end = total_padding // 2
+    pad_begin = total_padding - pad_end
+    if auto_pad == "SAME_UPPER":
+        pad_begin, pad_end = pad_end, pad_begin
+    elif auto_pad not in {"SAME_LOWER", "NOTSET", ""}:
+        raise UnsupportedOpError(
+            f"ConvTranspose has unsupported auto_pad mode '{auto_pad}'"
+        )
+    if pad_begin < 0 or pad_end < 0:
+        raise ShapeInferenceError(
+            f"ConvTranspose pads must be non-negative for dim {dim}"
+        )
+    return pad_begin, pad_end
+
+
+def resolve_conv_transpose_spec(graph: Graph, node: Node) -> ConvTransposeSpec:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "ConvTranspose must have 2 or 3 inputs and 1 output"
+        )
+    supported_attrs = {
+        "auto_pad",
+        "dilations",
+        "group",
+        "kernel_shape",
+        "output_padding",
+        "output_shape",
+        "pads",
+        "strides",
+    }
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("ConvTranspose has unsupported attributes")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    weight_shape = _value_shape(graph, node.inputs[1], node)
+    if len(input_shape) < 3:
+        raise UnsupportedOpError("ConvTranspose expects NCHW inputs with spatial dims")
+    spatial_rank = len(input_shape) - 2
+    if spatial_rank not in {1, 2, 3}:
+        raise UnsupportedOpError("ConvTranspose supports 1D/2D/3D inputs only")
+    if len(weight_shape) != spatial_rank + 2:
+        raise UnsupportedOpError(
+            "ConvTranspose weight rank must match spatial rank"
+        )
+    batch, in_channels = input_shape[0], input_shape[1]
+    in_spatial = input_shape[2:]
+    weight_in_channels, weight_out_channels, *kernel_shape = weight_shape
+    kernel_attr = node.attrs.get("kernel_shape")
+    if kernel_attr is not None:
+        kernel_attr = tuple(int(value) for value in kernel_attr)
+        if len(kernel_attr) != spatial_rank:
+            raise UnsupportedOpError(
+                "ConvTranspose kernel_shape rank must match input spatial rank"
+            )
+        if kernel_attr != tuple(kernel_shape):
+            raise ShapeInferenceError(
+                "ConvTranspose kernel_shape must match weights, "
+                f"got {kernel_attr} and {tuple(kernel_shape)}"
+            )
+        kernel_shape = list(kernel_attr)
+    else:
+        kernel_shape = list(kernel_shape)
+    group = int(node.attrs.get("group", 1))
+    if group <= 0:
+        raise UnsupportedOpError("ConvTranspose expects group >= 1")
+    if in_channels % group != 0:
+        raise ShapeInferenceError(
+            "ConvTranspose expects group to evenly divide in channels, "
+            f"got group={group}, in_channels={in_channels}"
+        )
+    if weight_in_channels != in_channels:
+        raise ShapeInferenceError(
+            "ConvTranspose input channels must match weight channels, "
+            f"got {in_channels} and {weight_in_channels}"
+        )
+    out_channels = weight_out_channels * group
+    if out_channels % group != 0:
+        raise ShapeInferenceError(
+            "ConvTranspose expects group to evenly divide out channels, "
+            f"got group={group}, out_channels={out_channels}"
+        )
+    if len(node.inputs) == 3:
+        bias_shape = _value_shape(graph, node.inputs[2], node)
+        if bias_shape != (out_channels,):
+            raise ShapeInferenceError(
+                f"ConvTranspose bias shape must be {(out_channels,)}, got {bias_shape}"
+            )
+    strides = tuple(
+        int(value) for value in node.attrs.get("strides", (1,) * spatial_rank)
+    )
+    if len(strides) != spatial_rank:
+        raise UnsupportedOpError("ConvTranspose stride rank mismatch")
+    dilations = tuple(
+        int(value) for value in node.attrs.get("dilations", (1,) * spatial_rank)
+    )
+    if len(dilations) != spatial_rank:
+        raise UnsupportedOpError("ConvTranspose dilation rank mismatch")
+    output_padding = tuple(
+        int(value)
+        for value in node.attrs.get("output_padding", (0,) * spatial_rank)
+    )
+    if len(output_padding) != spatial_rank:
+        raise UnsupportedOpError("ConvTranspose output_padding rank mismatch")
+    for dim, (padding, stride) in enumerate(zip(output_padding, strides)):
+        if padding < 0:
+            raise UnsupportedOpError(
+                "ConvTranspose output_padding must be non-negative"
+            )
+        if padding >= stride:
+            raise UnsupportedOpError(
+                "ConvTranspose output_padding must be smaller than stride"
+            )
+    pads = tuple(
+        int(value)
+        for value in node.attrs.get("pads", (0,) * (2 * spatial_rank))
+    )
+    if len(pads) != 2 * spatial_rank:
+        raise UnsupportedOpError("ConvTranspose pads rank mismatch")
+    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 == "":
+        auto_pad = "NOTSET"
+    output_shape_attr = node.attrs.get("output_shape")
+    output_shape: list[int] | None = None
+    if output_shape_attr is not None:
+        output_shape = [int(value) for value in output_shape_attr]
+        if len(output_shape) != spatial_rank:
+            raise UnsupportedOpError("ConvTranspose output_shape rank mismatch")
+    if output_shape is not None:
+        if auto_pad == "VALID":
+            auto_pad = "NOTSET"
+        pad_begin = []
+        pad_end = []
+        for dim, (in_dim, stride, dilation, kernel, out_dim, out_pad) in enumerate(
+            zip(
+                in_spatial,
+                strides,
+                dilations,
+                kernel_shape,
+                output_shape,
+                output_padding,
+            )
+        ):
+            effective_kernel = dilation * (kernel - 1) + 1
+            total_padding = (
+                stride * (in_dim - 1)
+                + out_pad
+                + effective_kernel
+                - out_dim
+            )
+            pad_start, pad_finish = _split_padding(
+                total_padding, auto_pad, dim=dim
+            )
+            pad_begin.append(pad_start)
+            pad_end.append(pad_finish)
+        out_spatial = output_shape
+    else:
+        if auto_pad == "VALID":
+            pad_begin = [0] * spatial_rank
+            pad_end = [0] * spatial_rank
+        elif auto_pad in {"SAME_UPPER", "SAME_LOWER"}:
+            pad_begin = []
+            pad_end = []
+            for dim, (in_dim, stride, dilation, kernel, out_pad) in enumerate(
+                zip(in_spatial, strides, dilations, kernel_shape, output_padding)
+            ):
+                effective_kernel = dilation * (kernel - 1) + 1
+                out_dim = in_dim * stride
+                total_padding = (
+                    stride * (in_dim - 1)
+                    + out_pad
+                    + effective_kernel
+                    - out_dim
+                )
+                pad_start, pad_finish = _split_padding(
+                    total_padding, auto_pad, dim=dim
+                )
+                pad_begin.append(pad_start)
+                pad_end.append(pad_finish)
+        elif auto_pad in {"NOTSET"}:
+            pad_begin = list(pads[:spatial_rank])
+            pad_end = list(pads[spatial_rank:])
+        else:
+            raise UnsupportedOpError(
+                f"ConvTranspose has unsupported auto_pad mode '{auto_pad}'"
+            )
+        out_spatial = []
+        for dim, (in_dim, stride, dilation, kernel, pad_start, pad_finish, out_pad) in enumerate(
+            zip(
+                in_spatial,
+                strides,
+                dilations,
+                kernel_shape,
+                pad_begin,
+                pad_end,
+                output_padding,
+            )
+        ):
+            effective_kernel = dilation * (kernel - 1) + 1
+            out_dim = (
+                stride * (in_dim - 1)
+                + out_pad
+                + effective_kernel
+                - pad_start
+                - pad_finish
+            )
+            if out_dim < 0:
+                raise ShapeInferenceError(
+                    "ConvTranspose output shape must be non-negative"
+                )
+            out_spatial.append(out_dim)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    expected_output_shape = (batch, out_channels, *out_spatial)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "ConvTranspose output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    return ConvTransposeSpec(
+        batch=batch,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        spatial_rank=spatial_rank,
+        in_spatial=in_spatial,
+        out_spatial=tuple(out_spatial),
+        kernel_shape=tuple(kernel_shape),
+        strides=strides,
+        pads=(*pad_begin, *pad_end),
+        dilations=dilations,
+        output_padding=output_padding,
+        group=group,
+    )
+
+
+@register_lowering("ConvTranspose")
+def lower_conv_transpose(graph: Graph, node: Node) -> ConvTransposeOp:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "ConvTranspose must have 2 or 3 inputs and 1 output"
+        )
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "ConvTranspose supports float16, float, and double inputs only"
+        )
+    spec = resolve_conv_transpose_spec(graph, node)
+    return ConvTransposeOp(
+        input0=node.inputs[0],
+        weights=node.inputs[1],
+        bias=node.inputs[2] if len(node.inputs) == 3 else None,
+        output=node.outputs[0],
+        batch=spec.batch,
+        in_channels=spec.in_channels,
+        out_channels=spec.out_channels,
+        spatial_rank=spec.spatial_rank,
+        in_spatial=spec.in_spatial,
+        out_spatial=spec.out_spatial,
+        kernel_shape=spec.kernel_shape,
+        strides=spec.strides,
+        pads=spec.pads,
+        dilations=spec.dilations,
+        output_padding=spec.output_padding,
+        group=spec.group,
+        dtype=op_dtype,
+    )