emx-onnx-cgen 0.3.8__py3-none-any.whl → 0.4.2.dev0__py3-none-any.whl

emx_onnx_cgen/compiler.py

@@ -3,7 +3,8 @@ from __future__ import annotations
 from dataclasses import dataclass, fields
 import hashlib
 from pathlib import Path
-from typing import Mapping
+import time
+from typing import Callable, Mapping, TypeVar
 
 import numpy as np
 import onnx
@@ -29,21 +30,27 @@ from .lowering import load_lowering_registry
 from .lowering.common import ensure_supported_dtype, shape_product, value_dtype
 from .lowering.registry import get_lowering_registry
 from .onnx_import import import_onnx
-from .runtime.evaluator import Evaluator
 
 
 @dataclass(frozen=True)
 class CompilerOptions:
-    template_dir: Path
+    template_dir: Path | None = None
     model_name: str = "model"
     emit_testbench: bool = False
     command_line: str | None = None
     model_checksum: str | None = None
     restrict_arrays: bool = True
+    fp32_accumulation_strategy: str = "fp64"
+    fp16_accumulation_strategy: str = "fp32"
     testbench_inputs: Mapping[str, np.ndarray] | None = None
+    testbench_optional_inputs: Mapping[str, bool] | None = None
     truncate_weights_after: int | None = None
     large_temp_threshold_bytes: int = 1024
-    large_weight_threshold: int = 1024 * 1024
+    large_weight_threshold: int = 100 * 1024
+    timings: dict[str, float] | None = None
+
+
+_T = TypeVar("_T")
 
 
 def _onnx_elem_type(dtype: np.dtype) -> int:
@@ -53,90 +60,155 @@ def _onnx_elem_type(dtype: np.dtype) -> int:
     raise UnsupportedOpError(f"Unsupported dtype {dtype} for ONNX output")
 
 
+def _optional_flag_name(name: str) -> str:
+    return f"{name}_present"
+
+
 class Compiler:
     def __init__(self, options: CompilerOptions | None = None) -> None:
         if options is None:
-            options = CompilerOptions(template_dir=Path("templates"))
+            options = CompilerOptions()
         self._options = options
         self._emitter = CEmitter(
             options.template_dir,
             restrict_arrays=options.restrict_arrays,
+            fp32_accumulation_strategy=options.fp32_accumulation_strategy,
+            fp16_accumulation_strategy=options.fp16_accumulation_strategy,
             truncate_weights_after=options.truncate_weights_after,
             large_temp_threshold_bytes=options.large_temp_threshold_bytes,
             large_weight_threshold=options.large_weight_threshold,
         )
         load_lowering_registry()
 
+    def _time_step(self, label: str, func: Callable[[], _T]) -> _T:
+        timings = self._options.timings
+        if timings is None:
+            return func()
+        started = time.perf_counter()
+        result = func()
+        timings[label] = time.perf_counter() - started
+        return result
+
     def compile(self, model: onnx.ModelProto) -> str:
-        graph = import_onnx(model)
-        graph = self._concretize_graph_shapes(model, graph)
-        testbench_inputs = self._resolve_testbench_inputs(graph)
-        variable_dim_inputs, variable_dim_outputs = self._collect_variable_dims(
-            graph
+        graph = self._time_step("import_onnx", lambda: import_onnx(model))
+        graph = self._time_step(
+            "concretize_shapes",
+            lambda: self._concretize_graph_shapes(model, graph),
+        )
+        testbench_inputs = self._time_step(
+            "resolve_testbench_inputs", lambda: self._resolve_testbench_inputs(graph)
+        )
+        variable_dim_inputs, variable_dim_outputs = self._time_step(
+            "collect_variable_dims", lambda: self._collect_variable_dims(graph)
         )
-        lowered = self._lower_model(model, graph)
-        return self._emitter.emit_model(
-            lowered,
-            emit_testbench=self._options.emit_testbench,
-            testbench_inputs=testbench_inputs,
-            variable_dim_inputs=variable_dim_inputs,
-            variable_dim_outputs=variable_dim_outputs,
+        lowered = self._time_step(
+            "lower_model", lambda: self._lower_model(model, graph)
+        )
+        return self._time_step(
+            "emit_model",
+            lambda: self._emitter.emit_model(
+                lowered,
+                emit_testbench=self._options.emit_testbench,
+                testbench_inputs=testbench_inputs,
+                testbench_optional_inputs=self._options.testbench_optional_inputs,
+                variable_dim_inputs=variable_dim_inputs,
+                variable_dim_outputs=variable_dim_outputs,
+            ),
         )
 
     def compile_with_data_file(self, model: onnx.ModelProto) -> tuple[str, str]:
-        graph = import_onnx(model)
-        graph = self._concretize_graph_shapes(model, graph)
-        testbench_inputs = self._resolve_testbench_inputs(graph)
-        variable_dim_inputs, variable_dim_outputs = self._collect_variable_dims(
-            graph
+        graph = self._time_step("import_onnx", lambda: import_onnx(model))
+        graph = self._time_step(
+            "concretize_shapes",
+            lambda: self._concretize_graph_shapes(model, graph),
+        )
+        testbench_inputs = self._time_step(
+            "resolve_testbench_inputs", lambda: self._resolve_testbench_inputs(graph)
+        )
+        variable_dim_inputs, variable_dim_outputs = self._time_step(
+            "collect_variable_dims", lambda: self._collect_variable_dims(graph)
+        )
+        lowered = self._time_step(
+            "lower_model", lambda: self._lower_model(model, graph)
         )
-        lowered = self._lower_model(model, graph)
-        return self._emitter.emit_model_with_data_file(
-            lowered,
-            emit_testbench=self._options.emit_testbench,
-            testbench_inputs=testbench_inputs,
-            variable_dim_inputs=variable_dim_inputs,
-            variable_dim_outputs=variable_dim_outputs,
+        return self._time_step(
+            "emit_model_with_data_file",
+            lambda: self._emitter.emit_model_with_data_file(
+                lowered,
+                emit_testbench=self._options.emit_testbench,
+                testbench_inputs=testbench_inputs,
+                testbench_optional_inputs=self._options.testbench_optional_inputs,
+                variable_dim_inputs=variable_dim_inputs,
+                variable_dim_outputs=variable_dim_outputs,
+            ),
         )
 
     def compile_with_weight_data(
         self, model: onnx.ModelProto
     ) -> tuple[str, bytes | None]:
-        graph = import_onnx(model)
-        graph = self._concretize_graph_shapes(model, graph)
-        testbench_inputs = self._resolve_testbench_inputs(graph)
-        variable_dim_inputs, variable_dim_outputs = self._collect_variable_dims(
-            graph
-        )
-        lowered = self._lower_model(model, graph)
-        generated = self._emitter.emit_model(
-            lowered,
-            emit_testbench=self._options.emit_testbench,
-            testbench_inputs=testbench_inputs,
-            variable_dim_inputs=variable_dim_inputs,
-            variable_dim_outputs=variable_dim_outputs,
-        )
-        weight_data = self._emitter.collect_weight_data(lowered.constants)
+        graph = self._time_step("import_onnx", lambda: import_onnx(model))
+        graph = self._time_step(
+            "concretize_shapes",
+            lambda: self._concretize_graph_shapes(model, graph),
+        )
+        testbench_inputs = self._time_step(
+            "resolve_testbench_inputs", lambda: self._resolve_testbench_inputs(graph)
+        )
+        variable_dim_inputs, variable_dim_outputs = self._time_step(
+            "collect_variable_dims", lambda: self._collect_variable_dims(graph)
+        )
+        lowered = self._time_step(
+            "lower_model", lambda: self._lower_model(model, graph)
+        )
+        generated = self._time_step(
+            "emit_model",
+            lambda: self._emitter.emit_model(
+                lowered,
+                emit_testbench=self._options.emit_testbench,
+                testbench_inputs=testbench_inputs,
+                testbench_optional_inputs=self._options.testbench_optional_inputs,
+                variable_dim_inputs=variable_dim_inputs,
+                variable_dim_outputs=variable_dim_outputs,
+            ),
+        )
+        weight_data = self._time_step(
+            "collect_weight_data",
+            lambda: self._emitter.collect_weight_data(lowered.constants),
+        )
         return generated, weight_data
 
     def compile_with_data_file_and_weight_data(
         self, model: onnx.ModelProto
     ) -> tuple[str, str, bytes | None]:
-        graph = import_onnx(model)
-        graph = self._concretize_graph_shapes(model, graph)
-        testbench_inputs = self._resolve_testbench_inputs(graph)
-        variable_dim_inputs, variable_dim_outputs = self._collect_variable_dims(
-            graph
-        )
-        lowered = self._lower_model(model, graph)
-        generated, data_source = self._emitter.emit_model_with_data_file(
-            lowered,
-            emit_testbench=self._options.emit_testbench,
-            testbench_inputs=testbench_inputs,
-            variable_dim_inputs=variable_dim_inputs,
-            variable_dim_outputs=variable_dim_outputs,
-        )
-        weight_data = self._emitter.collect_weight_data(lowered.constants)
+        graph = self._time_step("import_onnx", lambda: import_onnx(model))
+        graph = self._time_step(
+            "concretize_shapes",
+            lambda: self._concretize_graph_shapes(model, graph),
+        )
+        testbench_inputs = self._time_step(
+            "resolve_testbench_inputs", lambda: self._resolve_testbench_inputs(graph)
+        )
+        variable_dim_inputs, variable_dim_outputs = self._time_step(
+            "collect_variable_dims", lambda: self._collect_variable_dims(graph)
+        )
+        lowered = self._time_step(
+            "lower_model", lambda: self._lower_model(model, graph)
+        )
+        generated, data_source = self._time_step(
+            "emit_model_with_data_file",
+            lambda: self._emitter.emit_model_with_data_file(
+                lowered,
+                emit_testbench=self._options.emit_testbench,
+                testbench_inputs=testbench_inputs,
+                testbench_optional_inputs=self._options.testbench_optional_inputs,
+                variable_dim_inputs=variable_dim_inputs,
+                variable_dim_outputs=variable_dim_outputs,
+            ),
+        )
+        weight_data = self._time_step(
+            "collect_weight_data",
+            lambda: self._emitter.collect_weight_data(lowered.constants),
+        )
         return generated, data_source, weight_data
 
     @staticmethod
@@ -165,9 +237,11 @@ class Compiler:
         self._validate_graph(graph)
         (
             input_names,
+            input_optional_names,
             input_shapes,
             input_dtypes,
             output_names,
+            output_optional_names,
             output_shapes,
             output_dtypes,
         ) = self._collect_io_specs(graph)
@@ -220,9 +294,11 @@ class Compiler:
         return LoweredModel(
             name=self._options.model_name,
             input_names=input_names,
+            input_optional_names=input_optional_names,
             input_shapes=input_shapes,
             input_dtypes=input_dtypes,
             output_names=output_names,
+            output_optional_names=output_optional_names,
             output_shapes=output_shapes,
             output_dtypes=output_dtypes,
             constants=constants,
@@ -248,7 +324,6 @@ class Compiler:
                 "Testbench inputs include unknown inputs: "
                 + ", ".join(unknown_inputs)
             )
-        resolved: dict[str, tuple[float | int | bool, ...]] = {}
         for name, values in self._options.testbench_inputs.items():
             if not isinstance(values, np.ndarray):
                 raise CodegenError(
@@ -265,9 +340,7 @@ class Compiler:
                     "Testbench input "
                     f"{name} has {array.size} elements, expected {expected_count}"
                 )
-            array = array.reshape(expected_shape)
-            resolved[name] = tuple(array.ravel().tolist())
-        return resolved
+        return None
 
     def _concretize_graph_shapes(
         self, model: onnx.ModelProto, graph: Graph
@@ -337,6 +410,7 @@ class Compiler:
                     dtype=value.type.dtype,
                     shape=shape,
                     dim_params=(None,) * len(shape),
+                    is_optional=value.type.is_optional,
                 ),
             )
 
@@ -361,27 +435,39 @@ class Compiler:
         self, graph: Graph
     ) -> tuple[
         tuple[str, ...],
+        tuple[str | None, ...],
         tuple[tuple[int, ...], ...],
         tuple[ScalarType, ...],
         tuple[str, ...],
+        tuple[str | None, ...],
         tuple[tuple[int, ...], ...],
         tuple[ScalarType, ...],
     ]:
         input_names = tuple(value.name for value in graph.inputs)
+        input_optional_names = tuple(
+            _optional_flag_name(value.name) if value.type.is_optional else None
+            for value in graph.inputs
+        )
         input_shapes = tuple(value.type.shape for value in graph.inputs)
         input_dtypes = tuple(
             value_dtype(graph, value.name) for value in graph.inputs
         )
         output_names = tuple(value.name for value in graph.outputs)
+        output_optional_names = tuple(
+            _optional_flag_name(value.name) if value.type.is_optional else None
+            for value in graph.outputs
+        )
         output_shapes = tuple(value.type.shape for value in graph.outputs)
         output_dtypes = tuple(
             value_dtype(graph, value.name) for value in graph.outputs
         )
         return (
             input_names,
+            input_optional_names,
             input_shapes,
             input_dtypes,
             output_names,
+            output_optional_names,
             output_shapes,
             output_dtypes,
         )
@@ -439,26 +525,22 @@ class Compiler:
             initializer_count=len(graph.initializers),
         )
 
-    def run(
-        self, model: onnx.ModelProto, feeds: Mapping[str, np.ndarray]
-    ) -> dict[str, np.ndarray]:
-        graph = import_onnx(model)
-        evaluator = Evaluator(graph)
-        return evaluator.run(feeds)
-
-
 def _lowered_constants(graph: Graph | GraphContext) -> tuple[ConstTensor, ...]:
+    used_initializers = {value.name for value in graph.outputs}
+    for node in graph.nodes:
+        used_initializers.update(node.inputs)
     constants: list[ConstTensor] = []
     for initializer in graph.initializers:
+        if initializer.name not in used_initializers:
+            continue
         dtype = ensure_supported_dtype(initializer.type.dtype)
+        data_array = initializer.data.astype(dtype.np_dtype, copy=False)
+        data_tuple = tuple(data_array.ravel().tolist())
         constants.append(
             ConstTensor(
                 name=initializer.name,
                 shape=initializer.type.shape,
-                data=tuple(
-                    dtype.np_dtype.type(value)
-                    for value in initializer.data.ravel()
-                ),
+                data=data_tuple,
                 dtype=dtype,
             )
         )

emx_onnx_cgen/determinism.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+import os
+from typing import Iterator
+
+THREAD_ENV_VARS = (
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+    "NUMEXPR_NUM_THREADS",
+    "BLIS_NUM_THREADS",
+)
+
+
+@contextmanager
+def deterministic_reference_runtime() -> Iterator[None]:
+    previous = {name: os.environ.get(name) for name in THREAD_ENV_VARS}
+    for name in THREAD_ENV_VARS:
+        os.environ[name] = "1"
+    limits_context = None
+    try:
+        try:
+            from threadpoolctl import threadpool_limits
+        except Exception:
+            threadpool_limits = None
+        if threadpool_limits is not None:
+            limits_context = threadpool_limits(limits=1)
+            limits_context.__enter__()
+        yield
+    finally:
+        if limits_context is not None:
+            limits_context.__exit__(None, None, None)
+        for name, value in previous.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value

emx_onnx_cgen/ir/context.py

@@ -14,9 +14,28 @@ class GraphContext:
     _shape_cache: dict[str, tuple[int, ...]] = field(default_factory=dict)
     _initializer_cache: dict[str, Initializer] = field(default_factory=dict)
     _producer_cache: dict[str, Node] = field(default_factory=dict)
+    _value_cache: dict[str, Value] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        for value in self.graph.inputs + self.graph.outputs + self.graph.values:
+            self._value_cache[value.name] = value
+        for initializer in self.graph.initializers:
+            if initializer.name not in self._value_cache:
+                self._value_cache[initializer.name] = Value(
+                    name=initializer.name,
+                    type=initializer.type,
+                )
+            self._initializer_cache[initializer.name] = initializer
+        for node in self.graph.nodes:
+            for output in node.outputs:
+                if output and output not in self._producer_cache:
+                    self._producer_cache[output] = node
 
     def find_value(self, name: str) -> Value:
-        return self.graph.find_value(name)
+        value = self._value_cache.get(name)
+        if value is None:
+            raise KeyError(name)
+        return value
 
     def dtype(self, name: str, node: Node | None = None) -> ScalarType:
         if name in self._dtype_cache:
@@ -55,23 +74,14 @@ class GraphContext:
     def set_shape(self, name: str, shape: tuple[int, ...]) -> None:
         self._shape_cache[name] = shape
 
+    def has_shape(self, name: str) -> bool:
+        return name in self._shape_cache
+
     def initializer(self, name: str) -> Initializer | None:
-        if name in self._initializer_cache:
-            return self._initializer_cache[name]
-        for initializer in self.graph.initializers:
-            if initializer.name == name:
-                self._initializer_cache[name] = initializer
-                return initializer
-        return None
+        return self._initializer_cache.get(name)
 
     def producer(self, output_name: str) -> Node | None:
-        if output_name in self._producer_cache:
-            return self._producer_cache[output_name]
-        for node in self.graph.nodes:
-            if output_name in node.outputs:
-                self._producer_cache[output_name] = node
-                return node
-        return None
+        return self._producer_cache.get(output_name)
 
     def opset_version(self, domain: str = "") -> int | None:
         if domain in {"", "ai.onnx"}:

emx_onnx_cgen/ir/model.py

@@ -13,6 +13,7 @@ class TensorType:
     dtype: ScalarType
     shape: tuple[int, ...]
     dim_params: tuple[str | None, ...]
+    is_optional: bool = False
 
 
 @dataclass(frozen=True)

emx_onnx_cgen/ir/op_base.py

@@ -414,19 +414,20 @@ class VariadicLikeOpBase(RenderableOpBase):
 
     def infer_shapes(self, ctx: OpContext) -> None:
         input_shapes = tuple(ctx.shape(name) for name in self._variadic_inputs())
-        output_shape = BroadcastingOpBase.broadcast_shapes(*input_shapes)
-        for shape in input_shapes:
-            if shape != output_shape:
-                raise UnsupportedOpError(
-                    f"{self._variadic_kind()} expects identical input/output shapes"
-                )
+        try:
+            output_shape = BroadcastingOpBase.broadcast_shapes(*input_shapes)
+        except ShapeInferenceError as exc:
+            raise UnsupportedOpError(
+                f"{self._variadic_kind()} expects broadcastable input shapes"
+            ) from exc
         try:
             expected = ctx.shape(self._variadic_output())
         except ShapeInferenceError:
             expected = None
         if expected is not None and expected != output_shape:
             raise UnsupportedOpError(
-                f"{self._variadic_kind()} expects identical input/output shapes"
+                f"{self._variadic_kind()} output shape must be {output_shape}, "
+                f"got {expected}"
             )
         ctx.set_shape(self._variadic_output(), output_shape)
 
@@ -469,6 +470,30 @@ class ReduceOpBase(RenderableOpBase):
 
 
 class BroadcastingOpBase(RenderableOpBase):
+    @staticmethod
+    def unidirectional_broadcastable(
+        source: tuple[int, ...],
+        target: tuple[int, ...],
+    ) -> bool:
+        if len(source) > len(target):
+            return False
+        padded = (1,) * (len(target) - len(source)) + source
+        for source_dim, target_dim in zip(padded, target):
+            if source_dim not in {1, target_dim}:
+                return False
+        return True
+
+    @staticmethod
+    def prelu_channel_axis(
+        input_shape: tuple[int, ...],
+        slope_shape: tuple[int, ...],
+    ) -> int | None:
+        if len(input_shape) < 2 or len(slope_shape) != 1:
+            return None
+        if slope_shape[0] != input_shape[1]:
+            return None
+        return 1
+
     @staticmethod
     def broadcast_shapes(
         *shapes: tuple[int, ...],

emx_onnx_cgen/ir/ops/__init__.py

@@ -3,28 +3,35 @@ from .elementwise import (
     ClipOp,
     IdentityOp,
     MultiInputBinaryOp,
+    PowOp,
+    QLinearMulOp,
     UnaryOp,
     VariadicOp,
     WhereOp,
 )
 from .misc import (
+    BernoulliOp,
     CastOp,
     ConcatOp,
     ConstantOfShapeOp,
     CumSumOp,
     DepthToSpaceOp,
+    DequantizeLinearOp,
     ExpandOp,
     EyeLikeOp,
     GatherElementsOp,
     GatherNDOp,
     GatherOp,
     GridSampleOp,
+    HammingWindowOp,
     NonMaxSuppressionOp,
     NonZeroOp,
     OneHotOp,
+    OptionalHasElementOp,
     PadOp,
     QuantizeLinearOp,
     RangeOp,
+    HammingWindowOp,
     ReshapeOp,
     ResizeOp,
     ScatterNDOp,
@@ -34,6 +41,7 @@ from .misc import (
     SpaceToDepthOp,
     SplitOp,
     TensorScatterOp,
+    TfIdfVectorizerOp,
     TileOp,
     TransposeOp,
     TriluOp,
@@ -44,10 +52,12 @@ from .nn import (
     AveragePoolOp,
     BatchNormOp,
     ConvOp,
+    ConvIntegerOp,
     ConvTransposeOp,
     EinsumKind,
     EinsumOp,
     GemmOp,
+    GruOp,
     GroupNormalizationOp,
     HardmaxOp,
     InstanceNormalizationOp,
@@ -75,15 +85,18 @@ __all__ = [
     "AttentionOp",
     "AveragePoolOp",
     "BatchNormOp",
+    "BernoulliOp",
     "BinaryOp",
     "CastOp",
     "ClipOp",
     "ConcatOp",
     "ConstantOfShapeOp",
     "ConvOp",
+    "ConvIntegerOp",
     "ConvTransposeOp",
     "CumSumOp",
     "DepthToSpaceOp",
+    "DequantizeLinearOp",
     "EinsumKind",
     "EinsumOp",
     "ExpandOp",
@@ -93,6 +106,8 @@ __all__ = [
     "GatherOp",
     "GemmOp",
     "GridSampleOp",
+    "GruOp",
+    "HammingWindowOp",
     "GroupNormalizationOp",
     "HardmaxOp",
     "IdentityOp",
@@ -111,10 +126,14 @@ __all__ = [
     "NonMaxSuppressionOp",
     "NonZeroOp",
     "OneHotOp",
+    "OptionalHasElementOp",
     "PadOp",
+    "PowOp",
     "QuantizeLinearOp",
+    "QLinearMulOp",
     "QLinearMatMulOp",
     "RangeOp",
+    "HammingWindowOp",
     "ReduceOp",
     "ReshapeOp",
     "ResizeOp",
@@ -129,6 +148,7 @@ __all__ = [
     "SpaceToDepthOp",
     "SplitOp",
     "TensorScatterOp",
+    "TfIdfVectorizerOp",
     "TileOp",
     "TopKOp",
     "TransposeOp",