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

emx_onnx_cgen/ir/op_base.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Protocol
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from .op_context import OpContext
+
+
+class Emitter(Protocol):
+    def render_op(self, op: "OpBase", ctx: "EmitContext") -> str:
+        ...
+
+
+@dataclass(frozen=True)
+class EmitContext:
+    op_index: int
+
+
+class OpBase(ABC):
+    """Ops should not mutate themselves; store derived values in OpContext."""
+    inputs: tuple[str, ...]
+    outputs: tuple[str, ...]
+
+    def __getattr__(self, name: str) -> str:
+        if name == "kind":
+            return self.__class__.__name__
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object has no attribute '{name}'"
+        )
+
+    def validate(self, ctx: OpContext) -> None:
+        return None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        return None
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        return None
+
+    @abstractmethod
+    def emit(self, emitter: Emitter, ctx: EmitContext) -> str:
+        raise NotImplementedError
+
+
+class RenderableOpBase(OpBase):
+    def emit(self, emitter: Emitter, ctx: EmitContext) -> str:
+        return emitter.render_op(self, ctx)
+
+
+class ElementwiseOpBase(RenderableOpBase):
+    """Elementwise ops should validate against OpContext and store no derived state."""
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        raise NotImplementedError
+
+    def _elementwise_output(self) -> str:
+        raise NotImplementedError
+
+    def _elementwise_condition_inputs(self) -> tuple[str, ...]:
+        return ()
+
+    def _elementwise_compare(self) -> bool:
+        return False
+
+    def _elementwise_data_inputs(self) -> tuple[str, ...]:
+        inputs = self._elementwise_inputs()
+        condition_inputs = set(self._elementwise_condition_inputs())
+        return tuple(name for name in inputs if name not in condition_inputs)
+
+    def validate(self, ctx: OpContext) -> None:
+        condition_inputs = self._elementwise_condition_inputs()
+        for name in condition_inputs:
+            dtype = ctx.dtype(name)
+            if dtype != ScalarType.BOOL:
+                raise UnsupportedOpError(
+                    f"{self.kind} expects bool condition, got {dtype.onnx_name}"
+                )
+        data_inputs = self._elementwise_data_inputs()
+        if not data_inputs:
+            return None
+        data_dtypes = tuple(ctx.dtype(name) for name in data_inputs)
+        if any(dtype != data_dtypes[0] for dtype in data_dtypes[1:]):
+            dtype_names = ", ".join(dtype.onnx_name for dtype in data_dtypes)
+            raise UnsupportedOpError(
+                f"{self.kind} expects matching input dtypes, got {dtype_names}"
+            )
+        output_dtype = ctx.dtype(self._elementwise_output())
+        if self._elementwise_compare():
+            if output_dtype != ScalarType.BOOL:
+                raise UnsupportedOpError(
+                    f"{self.kind} expects bool output, got {output_dtype.onnx_name}"
+                )
+            return None
+        if output_dtype != data_dtypes[0]:
+            raise UnsupportedOpError(
+                f"{self.kind} expects output dtype {data_dtypes[0].onnx_name}, "
+                f"got {output_dtype.onnx_name}"
+            )
+        return None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        input_names = self._elementwise_inputs()
+        output_name = self._elementwise_output()
+        for name in input_names:
+            ctx.dtype(name)
+        ctx.dtype(output_name)
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_names = self._elementwise_inputs()
+        output_name = self._elementwise_output()
+        input_shapes = tuple(ctx.shape(name) for name in input_names)
+        if len(input_shapes) == 1:
+            output_shape = input_shapes[0]
+        else:
+            output_shape = BroadcastingOpBase.broadcast_shapes(*input_shapes)
+        ctx.set_shape(output_name, output_shape)
+        return None
+
+
+class ReduceOpBase(RenderableOpBase):
+    @staticmethod
+    def normalize_axes(
+        axes: tuple[int, ...] | None, rank: int
+    ) -> tuple[int, ...]:
+        if axes is None:
+            axes = tuple(range(rank))
+        normalized: list[int] = []
+        for axis in axes:
+            if axis < 0:
+                axis += rank
+            if axis < 0 or axis >= rank:
+                raise ShapeInferenceError(
+                    f"Reduce axis {axis} is out of bounds for rank {rank}"
+                )
+            normalized.append(axis)
+        return tuple(dict.fromkeys(normalized))
+
+    @staticmethod
+    def reduced_shape(
+        input_shape: tuple[int, ...],
+        axes: tuple[int, ...] | None,
+        *,
+        keepdims: bool,
+    ) -> tuple[int, ...]:
+        rank = len(input_shape)
+        normalized_axes = ReduceOpBase.normalize_axes(axes, rank)
+        if keepdims:
+            return tuple(
+                1 if axis in normalized_axes else dim
+                for axis, dim in enumerate(input_shape)
+            )
+        return tuple(
+            dim for axis, dim in enumerate(input_shape) if axis not in normalized_axes
+        )
+
+
+class BroadcastingOpBase(RenderableOpBase):
+    @staticmethod
+    def broadcast_shapes(
+        *shapes: tuple[int, ...],
+    ) -> tuple[int, ...]:
+        if not shapes:
+            return ()
+        max_rank = max(len(shape) for shape in shapes)
+        padded_shapes = [
+            (1,) * (max_rank - len(shape)) + shape for shape in shapes
+        ]
+        result: list[int] = []
+        for dims in zip(*padded_shapes):
+            dim = max(dims)
+            if any(d not in {1, dim} for d in dims):
+                raise ShapeInferenceError(
+                    "Broadcasting mismatch for shapes: "
+                    + ", ".join(str(shape) for shape in shapes)
+                )
+            result.append(dim)
+        return tuple(result)
+
+
+class MatMulLikeOpBase(RenderableOpBase):
+    pass
+
+
+class GemmLikeOpBase(RenderableOpBase):
+    pass
+
+
+class ConvLikeOpBase(RenderableOpBase):
+    pass

emx_onnx_cgen/ir/op_context.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from shared.scalar_types import ScalarType
+
+from .context import GraphContext
+
+
+_MISSING = object()
+
+
+@dataclass
+class OpContext:
+    graph: GraphContext
+    _dtype_overrides: dict[str, ScalarType] = field(default_factory=dict)
+    _shape_overrides: dict[str, tuple[int, ...]] = field(default_factory=dict)
+    _derived: dict[int, dict[str, object]] = field(default_factory=dict)
+
+    def dtype(self, name: str) -> ScalarType:
+        if name in self._dtype_overrides:
+            return self._dtype_overrides[name]
+        return self.graph.dtype(name)
+
+    def shape(self, name: str) -> tuple[int, ...]:
+        if name in self._shape_overrides:
+            return self._shape_overrides[name]
+        return self.graph.shape(name)
+
+    def set_dtype(self, name: str, dtype: ScalarType) -> None:
+        self._dtype_overrides[name] = dtype
+        self.graph.set_dtype(name, dtype)
+
+    def set_shape(self, name: str, shape: tuple[int, ...]) -> None:
+        self._shape_overrides[name] = shape
+        self.graph.set_shape(name, shape)
+
+    def set_derived(self, op: object, key: str, value: object) -> None:
+        self._derived.setdefault(id(op), {})[key] = value
+
+    def get_derived(
+        self, op: object, key: str, default: object = _MISSING
+    ) -> object:
+        derived = self._derived.get(id(op), {})
+        if key in derived:
+            return derived[key]
+        if default is _MISSING:
+            return _MISSING
+        return default
+
+    def require_derived(self, op: object, key: str) -> object:
+        derived = self._derived.get(id(op), {})
+        if key in derived:
+            return derived[key]
+        raise KeyError(
+            f"Missing derived value '{key}' for op {op.__class__.__name__}"
+        )
+
+    def copy_derived(self, source_op: object, target_op: object) -> None:
+        derived = self._derived.get(id(source_op))
+        if derived:
+            self._derived[id(target_op)] = dict(derived)
+
+    def __getattr__(self, name: str):
+        return getattr(self.graph, name)

emx_onnx_cgen/ir/ops/__init__.py

@@ -0,0 +1,130 @@
+from .elementwise import BinaryOp, ClipOp, IdentityOp, MultiInputBinaryOp, UnaryOp, WhereOp
+from .misc import (
+    CastOp,
+    ConcatOp,
+    ConstantOfShapeOp,
+    CumSumOp,
+    DepthToSpaceOp,
+    ExpandOp,
+    EyeLikeOp,
+    GatherElementsOp,
+    GatherNDOp,
+    GatherOp,
+    GridSampleOp,
+    NonMaxSuppressionOp,
+    NonZeroOp,
+    OneHotOp,
+    PadOp,
+    QuantizeLinearOp,
+    RangeOp,
+    ReshapeOp,
+    ResizeOp,
+    ScatterNDOp,
+    ShapeOp,
+    SizeOp,
+    SliceOp,
+    SpaceToDepthOp,
+    SplitOp,
+    TensorScatterOp,
+    TileOp,
+    TransposeOp,
+    TriluOp,
+)
+from .nn import (
+    AdagradOp,
+    AttentionOp,
+    AveragePoolOp,
+    BatchNormOp,
+    ConvOp,
+    ConvTransposeOp,
+    EinsumKind,
+    EinsumOp,
+    GemmOp,
+    GroupNormalizationOp,
+    HardmaxOp,
+    InstanceNormalizationOp,
+    LayerNormalizationOp,
+    LogSoftmaxOp,
+    LpNormalizationOp,
+    LpPoolOp,
+    LrnOp,
+    LstmOp,
+    MatMulOp,
+    MaxPoolOp,
+    MeanVarianceNormalizationOp,
+    NegativeLogLikelihoodLossOp,
+    QLinearMatMulOp,
+    RMSNormalizationOp,
+    RotaryEmbeddingOp,
+    SoftmaxCrossEntropyLossOp,
+    SoftmaxOp,
+)
+from .reduce import ArgReduceOp, ReduceOp, TopKOp
+
+__all__ = [
+    "AdagradOp",
+    "ArgReduceOp",
+    "AttentionOp",
+    "AveragePoolOp",
+    "BatchNormOp",
+    "BinaryOp",
+    "CastOp",
+    "ClipOp",
+    "ConcatOp",
+    "ConstantOfShapeOp",
+    "ConvOp",
+    "ConvTransposeOp",
+    "CumSumOp",
+    "DepthToSpaceOp",
+    "EinsumKind",
+    "EinsumOp",
+    "ExpandOp",
+    "EyeLikeOp",
+    "GatherElementsOp",
+    "GatherNDOp",
+    "GatherOp",
+    "GemmOp",
+    "GridSampleOp",
+    "GroupNormalizationOp",
+    "HardmaxOp",
+    "IdentityOp",
+    "InstanceNormalizationOp",
+    "LayerNormalizationOp",
+    "LogSoftmaxOp",
+    "LpNormalizationOp",
+    "LpPoolOp",
+    "LrnOp",
+    "LstmOp",
+    "MatMulOp",
+    "MaxPoolOp",
+    "MeanVarianceNormalizationOp",
+    "MultiInputBinaryOp",
+    "NegativeLogLikelihoodLossOp",
+    "NonMaxSuppressionOp",
+    "NonZeroOp",
+    "OneHotOp",
+    "PadOp",
+    "QuantizeLinearOp",
+    "QLinearMatMulOp",
+    "RangeOp",
+    "ReduceOp",
+    "ReshapeOp",
+    "ResizeOp",
+    "RMSNormalizationOp",
+    "RotaryEmbeddingOp",
+    "ScatterNDOp",
+    "ShapeOp",
+    "SizeOp",
+    "SliceOp",
+    "SoftmaxCrossEntropyLossOp",
+    "SoftmaxOp",
+    "SpaceToDepthOp",
+    "SplitOp",
+    "TensorScatterOp",
+    "TileOp",
+    "TopKOp",
+    "TransposeOp",
+    "TriluOp",
+    "UnaryOp",
+    "WhereOp",
+]

emx_onnx_cgen/ir/ops/elementwise.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_functions import ScalarFunction
+from shared.scalar_types import ScalarType
+
+from ...ops import COMPARE_FUNCTIONS, OperatorKind
+from ..op_base import ElementwiseOpBase
+from ..op_context import OpContext
+
+
+@dataclass(frozen=True)
+class BinaryOp(ElementwiseOpBase):
+    input0: str
+    input1: str
+    output: str
+    function: ScalarFunction
+    operator_kind: OperatorKind
+    input0_shape: tuple[int, ...]
+    input1_shape: tuple[int, ...]
+    shape: tuple[int, ...]
+    dtype: ScalarType
+    input_dtype: ScalarType
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        return (self.input0, self.input1)
+
+    def _elementwise_output(self) -> str:
+        return self.output
+
+    def _elementwise_compare(self) -> bool:
+        return self.function in COMPARE_FUNCTIONS
+
+
+@dataclass(frozen=True)
+class MultiInputBinaryOp(ElementwiseOpBase):
+    inputs: tuple[str, ...]
+    output: str
+    function: ScalarFunction
+    operator_kind: OperatorKind
+    shape: tuple[int, ...]
+    dtype: ScalarType
+    input_dtype: ScalarType
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        return self.inputs
+
+    def _elementwise_output(self) -> str:
+        return self.output
+
+    def _elementwise_compare(self) -> bool:
+        return self.function in COMPARE_FUNCTIONS
+
+
+@dataclass(frozen=True)
+class WhereOp(ElementwiseOpBase):
+    condition: str
+    input_x: str
+    input_y: str
+    output: str
+    condition_shape: tuple[int, ...]
+    x_shape: tuple[int, ...]
+    y_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    dtype: ScalarType
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        return (self.condition, self.input_x, self.input_y)
+
+    def _elementwise_output(self) -> str:
+        return self.output
+
+    def _elementwise_condition_inputs(self) -> tuple[str, ...]:
+        return (self.condition,)
+
+
+@dataclass(frozen=True)
+class UnaryOp(ElementwiseOpBase):
+    input0: str
+    output: str
+    function: ScalarFunction
+    shape: tuple[int, ...]
+    dtype: ScalarType
+    input_dtype: ScalarType
+    params: tuple[float, ...] = ()
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        return (self.input0,)
+
+    def _elementwise_output(self) -> str:
+        return self.output
+
+    def validate(self, ctx: OpContext) -> None:
+        super().validate(ctx)
+        return None
+
+    def _elementwise_compare(self) -> bool:
+        return self.function in {ScalarFunction.ISINF, ScalarFunction.ISNAN}
+
+
+@dataclass(frozen=True)
+class ClipOp(ElementwiseOpBase):
+    input0: str
+    input_min: str | None
+    input_max: str | None
+    output: str
+    input_shape: tuple[int, ...]
+    min_shape: tuple[int, ...] | None
+    max_shape: tuple[int, ...] | None
+    output_shape: tuple[int, ...]
+    dtype: ScalarType
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        inputs = [self.input0]
+        if self.input_min is not None:
+            inputs.append(self.input_min)
+        if self.input_max is not None:
+            inputs.append(self.input_max)
+        return tuple(inputs)
+
+    def _elementwise_output(self) -> str:
+        return self.output
+
+    def validate(self, ctx: OpContext) -> None:
+        super().validate(ctx)
+        return None
+
+
+@dataclass(frozen=True)
+class IdentityOp(ElementwiseOpBase):
+    input0: str
+    output: str
+    shape: tuple[int, ...]
+    dtype: ScalarType
+    input_dtype: ScalarType
+
+    def _elementwise_inputs(self) -> tuple[str, ...]:
+        return (self.input0,)
+
+    def _elementwise_output(self) -> str:
+        return self.output
+
+    def validate(self, ctx: OpContext) -> None:
+        super().validate(ctx)
+        return None