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

emx_onnx_cgen/lowering/__init__.py

@@ -1,3 +1,81 @@
+from __future__ import annotations
+
+import importlib
+
 from .registry import get_lowering, register_lowering
 
-__all__ = ["get_lowering", "register_lowering"]
+_LOWERING_MODULES = [
+    "adagrad",
+    "arg_reduce",
+    "attention",
+    "average_pool",
+    "batch_normalization",
+    "cast",
+    "concat",
+    "constant_of_shape",
+    "conv",
+    "conv_transpose",
+    "cumsum",
+    "depth_space",
+    "dropout",
+    "einsum",
+    "elementwise",
+    "expand",
+    "eye_like",
+    "flatten",
+    "gather",
+    "gather_elements",
+    "gather_nd",
+    "gemm",
+    "global_max_pool",
+    "grid_sample",
+    "group_normalization",
+    "hardmax",
+    "identity",
+    "instance_normalization",
+    "layer_normalization",
+    "logsoftmax",
+    "lp_normalization",
+    "lp_pool",
+    "lrn",
+    "lstm",
+    "matmul",
+    "maxpool",
+    "mean_variance_normalization",
+    "negative_log_likelihood_loss",
+    "non_max_suppression",
+    "nonzero",
+    "one_hot",
+    "pad",
+    "qlinear_matmul",
+    "quantize_linear",
+    "range",
+    "reduce",
+    "reshape",
+    "resize",
+    "rms_normalization",
+    "rotary_embedding",
+    "scatter_nd",
+    "shape",
+    "size",
+    "slice",
+    "softmax",
+    "softmax_cross_entropy_loss",
+    "split",
+    "squeeze",
+    "tensor_scatter",
+    "tile",
+    "topk",
+    "transpose",
+    "trilu",
+    "unsqueeze",
+    "variadic",
+    "where",
+]
+
+
+def load_lowering_registry() -> None:
+    for module_name in _LOWERING_MODULES:
+        importlib.import_module(f"{__name__}.{module_name}")
+
+__all__ = ["get_lowering", "register_lowering", "load_lowering_registry"]

emx_onnx_cgen/lowering/adagrad.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import AdagradOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+@register_lowering("Adagrad")
+def lower_adagrad(graph: Graph, node: Node) -> AdagradOp:
+    if len(node.inputs) < 5:
+        raise UnsupportedOpError("Adagrad must have at least 5 inputs")
+    if len(node.outputs) < 2:
+        raise UnsupportedOpError("Adagrad must have at least 2 outputs")
+    if (len(node.inputs) - 2) % 3 != 0:
+        raise UnsupportedOpError(
+            "Adagrad inputs must be R, T, Xs, Gs, Hs with matching counts"
+        )
+    tensor_count = (len(node.inputs) - 2) // 3
+    if len(node.outputs) != tensor_count * 2:
+        raise UnsupportedOpError(
+            "Adagrad outputs must be X_news followed by H_news"
+        )
+    rate_name = node.inputs[0]
+    timestep_name = node.inputs[1]
+    rate_shape = value_shape(graph, rate_name, node)
+    timestep_shape = value_shape(graph, timestep_name, node)
+    if not _is_scalar_shape(rate_shape):
+        raise UnsupportedOpError("Adagrad R input must be a scalar")
+    if not _is_scalar_shape(timestep_shape):
+        raise UnsupportedOpError("Adagrad T input must be a scalar")
+    rate_dtype = value_dtype(graph, rate_name, node)
+    if rate_dtype not in {ScalarType.F32, ScalarType.F64}:
+        raise UnsupportedOpError(
+            "Adagrad R input must be float or double"
+        )
+    timestep_dtype = value_dtype(graph, timestep_name, node)
+    if timestep_dtype != ScalarType.I64:
+        raise UnsupportedOpError("Adagrad T input must be int64")
+
+    inputs = node.inputs[2 : 2 + tensor_count]
+    gradients = node.inputs[2 + tensor_count : 2 + tensor_count * 2]
+    accumulators = node.inputs[2 + tensor_count * 2 : 2 + tensor_count * 3]
+    outputs = node.outputs[:tensor_count]
+    accumulator_outputs = node.outputs[tensor_count:]
+    if not inputs or not gradients or not accumulators:
+        raise UnsupportedOpError("Adagrad requires X, G, H inputs")
+    dtype = value_dtype(graph, inputs[0], node)
+    if dtype not in {ScalarType.F32, ScalarType.F64}:
+        raise UnsupportedOpError("Adagrad supports float and double tensors only")
+    if rate_dtype != dtype:
+        raise UnsupportedOpError(
+            "Adagrad R input dtype must match tensor dtype"
+        )
+    input_shapes: list[tuple[int, ...]] = []
+    output_shapes: list[tuple[int, ...]] = []
+    for index, (x_name, g_name, h_name, out_name, h_out_name) in enumerate(
+        zip(inputs, gradients, accumulators, outputs, accumulator_outputs)
+    ):
+        x_dtype = value_dtype(graph, x_name, node)
+        g_dtype = value_dtype(graph, g_name, node)
+        h_dtype = value_dtype(graph, h_name, node)
+        out_dtype = value_dtype(graph, out_name, node)
+        h_out_dtype = value_dtype(graph, h_out_name, node)
+        if {x_dtype, g_dtype, h_dtype, out_dtype, h_out_dtype} != {dtype}:
+            raise UnsupportedOpError(
+                "Adagrad inputs and outputs must share the same dtype"
+            )
+        x_shape = value_shape(graph, x_name, node)
+        g_shape = value_shape(graph, g_name, node)
+        h_shape = value_shape(graph, h_name, node)
+        out_shape = value_shape(graph, out_name, node)
+        h_out_shape = value_shape(graph, h_out_name, node)
+        if x_shape != g_shape or x_shape != h_shape:
+            raise ShapeInferenceError(
+                f"Adagrad inputs X/G/H shapes must match for tensor {index}"
+            )
+        if out_shape != x_shape or h_out_shape != x_shape:
+            raise ShapeInferenceError(
+                f"Adagrad outputs must match X shape for tensor {index}"
+            )
+        input_shapes.append(x_shape)
+        output_shapes.append(out_shape)
+
+    norm_coefficient = float(node.attrs.get("norm_coefficient", 0.0))
+    epsilon = float(node.attrs.get("epsilon", 0.0))
+    decay_factor = float(node.attrs.get("decay_factor", 0.0))
+
+    return AdagradOp(
+        rate=rate_name,
+        timestep=timestep_name,
+        inputs=tuple(inputs),
+        gradients=tuple(gradients),
+        accumulators=tuple(accumulators),
+        outputs=tuple(outputs),
+        accumulator_outputs=tuple(accumulator_outputs),
+        rate_shape=rate_shape,
+        timestep_shape=timestep_shape,
+        tensor_shapes=tuple(input_shapes),
+        output_shapes=tuple(output_shapes),
+        dtype=dtype,
+        rate_dtype=rate_dtype,
+        timestep_dtype=timestep_dtype,
+        norm_coefficient=norm_coefficient,
+        epsilon=epsilon,
+        decay_factor=decay_factor,
+    )

emx_onnx_cgen/lowering/arg_reduce.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ArgReduceOp
+from ..ir.ops import ArgReduceOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import shape_product, value_dtype, value_shape

emx_onnx_cgen/lowering/attention.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import AttentionOp
+from ..ir.ops import AttentionOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/average_pool.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from ..codegen.c_emitter import AveragePoolOp
+from ..ir.ops import AveragePoolOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/batch_normalization.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from ..codegen.c_emitter import BatchNormOp
+from ..ir.ops import BatchNormOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/cast.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import onnx
 
-from ..codegen.c_emitter import CastOp
+from ..ir.ops import CastOp
 from ..dtypes import scalar_type_from_onnx
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node

emx_onnx_cgen/lowering/common.py

@@ -5,7 +5,8 @@ from collections.abc import Sequence
 from shared.scalar_types import ScalarType
 
 from ..errors import ShapeInferenceError, UnsupportedOpError
-from ..ir.model import Graph, Node
+from ..ir.context import GraphContext
+from ..ir.model import Graph, Initializer, Node
 
 
 def ensure_supported_dtype(dtype: ScalarType) -> ScalarType:
@@ -14,7 +15,24 @@ def ensure_supported_dtype(dtype: ScalarType) -> ScalarType:
     return dtype
 
 
-def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType:
+def onnx_opset_version(graph: Graph | GraphContext, domain: str = "") -> int | None:
+    if isinstance(graph, GraphContext):
+        return graph.opset_version(domain)
+    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 | GraphContext, name: str, node: Node | None = None
+) -> ScalarType:
+    if isinstance(graph, GraphContext):
+        return graph.dtype(name, node)
     try:
         value = graph.find_value(name)
     except KeyError as exc:
@@ -26,18 +44,395 @@ def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType
     return ensure_supported_dtype(value.type.dtype)
 
 
-def value_shape(graph: Graph, name: str, node: Node | None = None) -> tuple[int, ...]:
+def value_shape(
+    graph: Graph | GraphContext, name: str, node: Node | None = None
+) -> tuple[int, ...]:
+    if isinstance(graph, GraphContext):
+        shape = graph.shape(name, node)
+        value = graph.find_value(name)
+    else:
+        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 shape for value '{name}' in op {op_type}. "
+                "Hint: run ONNX shape inference or export with static shapes."
+            ) from exc
+        shape = value.type.shape
+    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 shape
+
+
+def _find_initializer(graph: Graph | GraphContext, name: str) -> Initializer | None:
+    if isinstance(graph, GraphContext):
+        return graph.initializer(name)
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _find_node_by_output(graph: Graph | GraphContext, name: str) -> Node | None:
+    if isinstance(graph, GraphContext):
+        return graph.producer(name)
+    for node in graph.nodes:
+        if name in node.outputs:
+            return node
+    return None
+
+
+def _shape_values_from_shape_node(
+    graph: Graph | GraphContext, 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 | GraphContext,
+    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 | GraphContext,
+    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:
-        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
+        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 | GraphContext,
+    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:
+def node_dtype(graph: Graph | GraphContext, node: Node, *names: str) -> ScalarType:
     filtered = [name for name in names if name]
     if not filtered:
         raise UnsupportedOpError(

emx_onnx_cgen/lowering/concat.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import ConcatOp
+from ..ir.ops import ConcatOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype