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

emx_onnx_cgen/lowering/softmax_cross_entropy_loss.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import SoftmaxCrossEntropyLossOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import shape_product as _shape_product
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@register_lowering("SoftmaxCrossEntropyLoss")
+def lower_softmax_cross_entropy_loss(
+    graph: Graph, node: Node
+) -> SoftmaxCrossEntropyLossOp:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) not in {1, 2}:
+        raise UnsupportedOpError(
+            "SoftmaxCrossEntropyLoss must have 2 or 3 inputs and 1 or 2 outputs"
+        )
+    input_name = node.inputs[0]
+    target_name = node.inputs[1]
+    weight_name = node.inputs[2] if len(node.inputs) > 2 else None
+    input_dtype = _value_dtype(graph, input_name, node)
+    if not input_dtype.is_float:
+        raise UnsupportedOpError(
+            "SoftmaxCrossEntropyLoss supports float16, float, and double inputs only"
+        )
+    output_name = node.outputs[0]
+    output_dtype = _value_dtype(graph, output_name, node)
+    if output_dtype != input_dtype:
+        raise UnsupportedOpError(
+            "SoftmaxCrossEntropyLoss output dtype must match input dtype"
+        )
+    log_prob_name = node.outputs[1] if len(node.outputs) > 1 else None
+    if log_prob_name is not None:
+        log_prob_dtype = _value_dtype(graph, log_prob_name, node)
+        if log_prob_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "SoftmaxCrossEntropyLoss log_prob dtype must match input dtype"
+            )
+    target_dtype = _value_dtype(graph, target_name, node)
+    if target_dtype not in {ScalarType.I32, ScalarType.I64}:
+        raise UnsupportedOpError(
+            "SoftmaxCrossEntropyLoss target must be int32 or int64"
+        )
+    weight_dtype = None
+    weight_shape: tuple[int, ...] | None = None
+    if weight_name is not None:
+        weight_dtype = _value_dtype(graph, weight_name, node)
+        if weight_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "SoftmaxCrossEntropyLoss weight dtype must match input dtype"
+            )
+    input_shape = _value_shape(graph, input_name, node)
+    target_shape = _value_shape(graph, target_name, node)
+    output_shape = _value_shape(graph, output_name, node)
+    if len(input_shape) < 2:
+        raise ShapeInferenceError("SoftmaxCrossEntropyLoss input must be at least 2D")
+    if len(target_shape) != len(input_shape) - 1:
+        raise ShapeInferenceError(
+            "SoftmaxCrossEntropyLoss target rank must be input rank - 1"
+        )
+    if input_shape[0] != target_shape[0]:
+        raise ShapeInferenceError(
+            "SoftmaxCrossEntropyLoss target batch dimension must match input"
+        )
+    if input_shape[2:] != target_shape[1:]:
+        raise ShapeInferenceError(
+            "SoftmaxCrossEntropyLoss target spatial dimensions must match input"
+        )
+    if weight_name is not None:
+        weight_shape = _value_shape(graph, weight_name, node)
+        if len(weight_shape) != 1 or weight_shape[0] != input_shape[1]:
+            raise ShapeInferenceError(
+                "SoftmaxCrossEntropyLoss weight must have shape (C,)"
+            )
+    if log_prob_name is not None:
+        log_prob_shape = _value_shape(graph, log_prob_name, node)
+        if log_prob_shape != input_shape:
+            raise ShapeInferenceError(
+                "SoftmaxCrossEntropyLoss log_prob output must match input shape"
+            )
+    reduction = node.attrs.get("reduction", "mean")
+    if isinstance(reduction, bytes):
+        reduction = reduction.decode("utf-8")
+    if reduction not in {"none", "mean", "sum"}:
+        raise UnsupportedOpError(
+            "SoftmaxCrossEntropyLoss reduction must be none, mean, or sum"
+        )
+    if reduction == "none":
+        if output_shape != target_shape:
+            raise ShapeInferenceError(
+                "SoftmaxCrossEntropyLoss output must match target shape "
+                "when reduction is none"
+            )
+    else:
+        if output_shape not in {(), (1,)}:
+            raise ShapeInferenceError(
+                "SoftmaxCrossEntropyLoss output must be scalar when reduced"
+            )
+    n = input_shape[0]
+    c = input_shape[1]
+    d = _shape_product(input_shape[2:]) if len(input_shape) > 2 else 1
+    ignore_index = node.attrs.get("ignore_index")
+    if ignore_index is not None:
+        ignore_index = int(ignore_index)
+    return SoftmaxCrossEntropyLossOp(
+        input0=input_name,
+        target=target_name,
+        weight=weight_name,
+        output=output_name,
+        log_prob=log_prob_name,
+        input_shape=input_shape,
+        target_shape=target_shape,
+        output_shape=output_shape,
+        log_prob_shape=input_shape if log_prob_name is not None else None,
+        n=n,
+        c=c,
+        d=d,
+        reduction=reduction,
+        ignore_index=ignore_index,
+        input_dtype=input_dtype,
+        weight_dtype=weight_dtype,
+        weight_shape=weight_shape,
+        dtype=input_dtype,
+        target_dtype=target_dtype,
+    )

emx_onnx_cgen/lowering/split.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import SplitOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import optional_name, value_dtype, value_shape
+from ..validation import normalize_axis
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _read_split_sizes(graph: Graph, name: str, node: Node) -> 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(
+            f"{node.op_type} split input must be int64 or int32"
+        )
+    if len(initializer.type.shape) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} split input must be a 1D tensor"
+        )
+    values = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if values.size == 0:
+        raise ShapeInferenceError(
+            f"{node.op_type} split input cannot be empty"
+        )
+    return [int(value) for value in values]
+
+
+def _validate_static_dims(shape: tuple[int, ...], node: Node) -> None:
+    if any(dim < 0 for dim in shape):
+        raise ShapeInferenceError(
+            f"{node.op_type} does not support dynamic dims"
+        )
+
+
+def _normalize_num_outputs(node: Node, output_count: int) -> int:
+    num_outputs_attr = node.attrs.get("num_outputs")
+    if num_outputs_attr is None:
+        return output_count
+    num_outputs = int(num_outputs_attr)
+    if num_outputs <= 0:
+        raise UnsupportedOpError("Split num_outputs must be positive")
+    if output_count != num_outputs:
+        raise ShapeInferenceError(
+            f"Split expects {num_outputs} outputs, got {output_count}"
+        )
+    return num_outputs
+
+
+@register_lowering("Split")
+def lower_split(graph: Graph, node: Node) -> SplitOp:
+    if len(node.inputs) < 1 or len(node.outputs) < 1:
+        raise UnsupportedOpError("Split must have at least 1 input and 1 output")
+    if len(node.inputs) > 2:
+        raise UnsupportedOpError("Split supports up to 2 inputs")
+    input_name = node.inputs[0]
+    if not input_name:
+        raise UnsupportedOpError("Split input must be provided")
+    input_shape = value_shape(graph, input_name, node)
+    _validate_static_dims(input_shape, node)
+    axis = normalize_axis(int(node.attrs.get("axis", 0)), input_shape, node)
+    output_shapes = [
+        value_shape(graph, output, node) for output in node.outputs
+    ]
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtypes = {value_dtype(graph, output, node) for output in node.outputs}
+    if output_dtypes != {input_dtype}:
+        dtype_names = ", ".join(
+            dtype.onnx_name for dtype in sorted(output_dtypes, key=str)
+        )
+        raise UnsupportedOpError(
+            f"Split expects matching dtypes, got {dtype_names}"
+        )
+    split_name = optional_name(node.inputs, 1)
+    if split_name is not None and "num_outputs" in node.attrs:
+        raise UnsupportedOpError(
+            "Split cannot specify both split input and num_outputs"
+        )
+    if split_name is not None:
+        split_sizes = _read_split_sizes(graph, split_name, node)
+        if split_sizes is None:
+            split_shape, split_dtype = value_shape(
+                graph, split_name, node
+            ), value_dtype(graph, split_name, node)
+            if split_dtype not in {ScalarType.I64, ScalarType.I32}:
+                raise UnsupportedOpError(
+                    f"{node.op_type} split input must be int64 or int32"
+                )
+            if len(split_shape) != 1:
+                raise UnsupportedOpError(
+                    f"{node.op_type} split input must be a 1D tensor"
+                )
+            if split_shape[0] != len(node.outputs):
+                raise ShapeInferenceError(
+                    f"Split expects {len(node.outputs)} outputs, got {split_shape[0]}"
+                )
+            split_sizes = [shape[axis] for shape in output_shapes]
+        if len(split_sizes) != len(node.outputs):
+            raise ShapeInferenceError(
+                f"Split expects {len(split_sizes)} outputs, got {len(node.outputs)}"
+            )
+        if any(size < 0 for size in split_sizes):
+            raise ShapeInferenceError("Split sizes must be non-negative")
+        if sum(split_sizes) != input_shape[axis]:
+            raise ShapeInferenceError(
+                "Split sizes must sum to the axis dimension"
+            )
+    else:
+        num_outputs = _normalize_num_outputs(node, len(node.outputs))
+        axis_dim = input_shape[axis]
+        base = axis_dim // num_outputs
+        remainder = axis_dim % num_outputs
+        split_sizes = [base + 1] * remainder + [base] * (
+            num_outputs - remainder
+        )
+    computed_shapes: list[tuple[int, ...]] = []
+    for size, output_shape in zip(split_sizes, output_shapes):
+        if size < 0:
+            raise ShapeInferenceError("Split output size must be non-negative")
+        shape = list(input_shape)
+        shape[axis] = size
+        computed_shape = tuple(shape)
+        if output_shape != computed_shape:
+            raise ShapeInferenceError(
+                f"Split output shape must be {computed_shape}, got {output_shape}"
+            )
+        computed_shapes.append(computed_shape)
+    return SplitOp(
+        input0=input_name,
+        outputs=tuple(node.outputs),
+        input_shape=input_shape,
+        output_shapes=tuple(computed_shapes),
+        axis=axis,
+        split_sizes=tuple(split_sizes),
+        dtype=input_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/squeeze.py

@@ -0,0 +1,161 @@
+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], input_rank: int, node: Node
+) -> tuple[int, ...]:
+    normalized: list[int] = []
+    for axis in axes:
+        if axis < 0:
+            axis += input_rank
+        if axis < 0 or axis >= input_rank:
+            raise ShapeInferenceError(
+                f"{node.op_type} axis {axis} is out of range for rank {input_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(
+                    "Squeeze 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:
+        return None
+    return tuple(axes_values)
+
+
+def _expected_output_shape(
+    input_shape: tuple[int, ...], axes: tuple[int, ...]
+) -> tuple[int, ...]:
+    axis_set = set(axes)
+    return tuple(
+        dim for index, dim in enumerate(input_shape) if index not in axis_set
+    )
+
+
+def _validate_output_shape_for_unknown_axes(
+    input_shape: tuple[int, ...], output_shape: tuple[int, ...], node: Node
+) -> None:
+    output_index = 0
+    for dim in input_shape:
+        if output_index < len(output_shape) and dim == output_shape[output_index]:
+            output_index += 1
+            continue
+        if dim != 1:
+            raise ShapeInferenceError(
+                "Squeeze output shape must remove only dimensions of size 1"
+            )
+    if output_index != len(output_shape):
+        raise ShapeInferenceError(
+            "Squeeze output shape must preserve input order while removing size-1 axes"
+        )
+
+
+@register_lowering("Squeeze")
+def lower_squeeze(graph: Graph, node: Node) -> ReshapeOp:
+    if len(node.outputs) != 1 or len(node.inputs) not in {1, 2}:
+        raise UnsupportedOpError("Squeeze 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(
+            "Squeeze 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(
+                    "Squeeze axes input must be int64 or int32, "
+                    f"got {axes_dtype.onnx_name}"
+                )
+            _validate_output_shape_for_unknown_axes(input_shape, output_shape, node)
+        else:
+            expected_shape = tuple(dim for dim in input_shape if dim != 1)
+            if expected_shape != output_shape:
+                raise ShapeInferenceError(
+                    "Squeeze output shape must be "
+                    f"{expected_shape}, got {output_shape}"
+                )
+    else:
+        normalized_axes = _normalize_axes(list(axes), len(input_shape), node)
+        for axis in normalized_axes:
+            if input_shape[axis] != 1:
+                raise ShapeInferenceError(
+                    "Squeeze axes must target dimensions of size 1"
+                )
+        expected_shape = _expected_output_shape(input_shape, normalized_axes)
+        if expected_shape != output_shape:
+            raise ShapeInferenceError(
+                "Squeeze 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/tile.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import TileOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _read_repeats(graph: Graph, name: str, node: Node) -> tuple[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("Tile repeats input must be int64 or int32")
+    if len(initializer.type.shape) != 1:
+        raise UnsupportedOpError("Tile repeats input must be a 1D tensor")
+    values = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    return tuple(int(value) for value in values)
+
+
+def _compute_strides(shape: tuple[int, ...]) -> tuple[int, ...]:
+    strides: list[int] = []
+    stride = 1
+    for dim in reversed(shape):
+        strides.append(stride)
+        stride *= dim
+    return tuple(reversed(strides))
+
+
+@register_lowering("Tile")
+def lower_tile(graph: Graph, node: Node) -> TileOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Tile must have 2 inputs and 1 output")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    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(
+            "Tile expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    repeats = _read_repeats(graph, node.inputs[1], node)
+    if repeats is None:
+        raise UnsupportedOpError("Tile repeats input must be a constant initializer")
+    if len(repeats) != len(input_shape):
+        raise ShapeInferenceError(
+            "Tile repeats must have the same rank as input shape"
+        )
+    if any(value < 0 for value in repeats):
+        raise UnsupportedOpError("Tile repeats must be non-negative")
+    expected_shape = tuple(
+        int(dim) * int(repeat) for dim, repeat in zip(input_shape, repeats)
+    )
+    if output_shape and output_shape != expected_shape:
+        raise ShapeInferenceError(
+            "Tile output shape mismatch: "
+            f"expected {expected_shape}, got {output_shape}"
+        )
+    return TileOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=expected_shape,
+        repeats=repeats,
+        input_strides=_compute_strides(input_shape),
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/transpose.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import TransposeOp
+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
+
+
+@register_lowering("Transpose")
+def lower_transpose(graph: Graph, node: Node) -> TransposeOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Transpose must have 1 input and 1 output")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    perm = node.attrs.get("perm")
+    if perm is None:
+        perm = tuple(reversed(range(len(input_shape))))
+    else:
+        perm = tuple(int(axis) for axis in perm)
+    if len(perm) != len(input_shape):
+        raise ShapeInferenceError(
+            "Transpose perm must match input rank, "
+            f"got perm {perm} for shape {input_shape}"
+        )
+    if set(perm) != set(range(len(input_shape))):
+        raise UnsupportedOpError(
+            f"Transpose perm must be a permutation, got {perm}"
+        )
+    expected_shape = tuple(input_shape[axis] for axis in perm)
+    if output_shape != expected_shape:
+        raise ShapeInferenceError(
+            "Transpose output shape must match permuted input shape, "
+            f"expected {expected_shape}, got {output_shape}"
+        )
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    return TransposeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        perm=perm,
+        input_shape=input_shape,
+        output_shape=output_shape,
+        dtype=op_dtype,
+        input_dtype=op_dtype,
+    )