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

emx_onnx_cgen/compiler.py

@@ -10,6 +10,7 @@ import onnx
 
 from shared.scalar_types import ScalarType
 
+from .onnxruntime_utils import make_deterministic_session_options
 from .codegen.c_emitter import (
     AttentionOp,
     AveragePoolOp,
@@ -27,18 +28,26 @@ from .codegen.c_emitter import (
     CEmitter,
     ConstTensor,
     ConvOp,
+    ConvTransposeOp,
     ConcatOp,
     ConstantOfShapeOp,
     CumSumOp,
     GemmOp,
     GatherOp,
     GatherElementsOp,
+    GatherNDOp,
+    ScatterNDOp,
     ExpandOp,
     RangeOp,
+    OneHotOp,
+    LpPoolOp,
+    QuantizeLinearOp,
     LrnOp,
     LstmOp,
     LogSoftmaxOp,
+    HardmaxOp,
     NegativeLogLikelihoodLossOp,
+    NonZeroOp,
     NodeInfo,
     PadOp,
     SplitOp,
@@ -52,6 +61,7 @@ from .codegen.c_emitter import (
     ReshapeOp,
     ResizeOp,
     GridSampleOp,
+    HardmaxOp,
     SoftmaxOp,
     ShapeOp,
     SliceOp,
@@ -61,12 +71,13 @@ from .codegen.c_emitter import (
 )
 from .dtypes import dtype_info
 from .errors import CodegenError, ShapeInferenceError, UnsupportedOpError
-from .ir.model import Graph, Value
+from .ir.model import Graph, TensorType, Value
 from .lowering.attention import AttentionSpec, resolve_attention_spec
 from .lowering.average_pool import (
     lower_average_pool,
     lower_global_average_pool,
 )
+from .lowering import global_max_pool as _global_max_pool  # noqa: F401
 from .lowering.batch_normalization import lower_batch_normalization
 from .lowering.cast import lower_cast
 from .lowering.concat import lower_concat
@@ -78,25 +89,33 @@ from .lowering.common import (
     value_shape,
 )
 from .lowering.conv import ConvSpec, resolve_conv_spec
+from .lowering import conv_transpose as _conv_transpose  # noqa: F401
 from .lowering.constant_of_shape import lower_constant_of_shape
 from .lowering.dropout import lower_dropout
 from .lowering import cumsum as _cumsum  # noqa: F401
+from .lowering import einsum as _einsum  # noqa: F401
 from .lowering.flatten import lower_flatten
 from .lowering.gather import lower_gather
 from .lowering.gather_elements import lower_gather_elements
+from .lowering.gather_nd import lower_gather_nd
+from .lowering import scatter_nd as _scatter_nd  # noqa: F401
 from .lowering.gemm import resolve_gemm_spec, validate_gemm_bias_shape
 from .lowering.lrn import LrnSpec, resolve_lrn_spec
 from .lowering.logsoftmax import lower_logsoftmax
+from .lowering import hardmax as _hardmax  # noqa: F401
 from .lowering import group_normalization as _group_normalization  # noqa: F401
 from .lowering import instance_normalization as _instance_normalization  # noqa: F401
 from .lowering import layer_normalization as _layer_normalization  # noqa: F401
 from .lowering import lp_normalization as _lp_normalization  # noqa: F401
+from .lowering import lp_pool as _lp_pool  # noqa: F401
 from .lowering import mean_variance_normalization as _mean_variance_normalization  # noqa: F401
 from .lowering.negative_log_likelihood_loss import (
     lower_negative_log_likelihood_loss,
 )
+from .lowering import nonzero as _nonzero  # noqa: F401
 from .lowering.expand import lower_expand
 from .lowering.range import lower_range
+from .lowering import one_hot as _one_hot  # noqa: F401
 from .lowering.split import lower_split
 from .lowering.softmax_cross_entropy_loss import (
     lower_softmax_cross_entropy_loss,
@@ -109,15 +128,18 @@ from .lowering.reduce import (
     REDUCE_OUTPUTS_FLOAT_ONLY,
 )
 from .lowering import arg_reduce as _arg_reduce  # noqa: F401
+from .lowering import topk as _topk  # noqa: F401
 from .lowering.reshape import lower_reshape
 from .lowering.resize import lower_resize
 from .lowering.grid_sample import lower_grid_sample
+from .lowering import quantize_linear as _quantize_linear  # noqa: F401
 from .lowering.slice import lower_slice
 from .lowering.squeeze import lower_squeeze
 from .lowering import depth_space as _depth_space  # noqa: F401
 from .lowering import eye_like as _eye_like  # noqa: F401
 from .lowering import identity as _identity  # noqa: F401
 from .lowering import tile as _tile  # noqa: F401
+from .lowering import trilu as _trilu  # noqa: F401
 from .lowering.shape import lower_shape
 from .lowering.size import lower_size
 from .lowering.softmax import lower_softmax
@@ -157,6 +179,16 @@ class CompilerOptions:
     model_checksum: str | None = None
     restrict_arrays: bool = True
     testbench_inputs: Mapping[str, np.ndarray] | None = None
+    truncate_weights_after: int | None = None
+    large_temp_threshold_bytes: int = 1024
+    large_weight_threshold: int = 1024
+
+
+def _onnx_elem_type(dtype: np.dtype) -> int:
+    for elem_type, info in onnx._mapping.TENSOR_TYPE_MAP.items():
+        if info.np_dtype == dtype:
+            return elem_type
+    raise UnsupportedOpError(f"Unsupported dtype {dtype} for ONNX output")
 
 
 class Compiler:
@@ -165,11 +197,16 @@ class Compiler:
             options = CompilerOptions(template_dir=Path("templates"))
         self._options = options
         self._emitter = CEmitter(
-            options.template_dir, restrict_arrays=options.restrict_arrays
+            options.template_dir,
+            restrict_arrays=options.restrict_arrays,
+            truncate_weights_after=options.truncate_weights_after,
+            large_temp_threshold_bytes=options.large_temp_threshold_bytes,
+            large_weight_threshold=options.large_weight_threshold,
         )
 
     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
@@ -185,6 +222,7 @@ class Compiler:
 
     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
@@ -198,6 +236,46 @@ class Compiler:
             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)
+        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)
+        return generated, data_source, weight_data
+
     @staticmethod
     def _collect_variable_dims(
         graph: Graph,
@@ -282,15 +360,93 @@ class Compiler:
             resolved[name] = tuple(array.ravel().tolist())
         return resolved
 
+    def _concretize_graph_shapes(
+        self, model: onnx.ModelProto, graph: Graph
+    ) -> Graph:
+        if not self._options.testbench_inputs:
+            return graph
+        if not any(value.type.dim_params for value in graph.values):
+            if not any(value.type.dim_params for value in graph.inputs):
+                if not any(value.type.dim_params for value in graph.outputs):
+                    return graph
+        try:
+            import onnxruntime as ort
+        except Exception:
+            return graph
+        try:
+            model_with_outputs = onnx.ModelProto()
+            model_with_outputs.CopyFrom(model)
+            existing_outputs = {
+                output.name for output in model_with_outputs.graph.output
+            }
+            value_info_by_name = {
+                value_info.name: value_info
+                for value_info in model_with_outputs.graph.value_info
+            }
+            for value in graph.values:
+                if value.name in existing_outputs:
+                    continue
+                value_info = value_info_by_name.get(value.name)
+                if value_info is None:
+                    dims: list[int | str | None] = []
+                    for index, dim in enumerate(value.type.shape):
+                        dim_param = None
+                        if index < len(value.type.dim_params):
+                            dim_param = value.type.dim_params[index]
+                        dims.append(dim_param if dim_param else None)
+                    elem_type = _onnx_elem_type(value.type.dtype.np_dtype)
+                    value_info = onnx.helper.make_tensor_value_info(
+                        value.name, elem_type, dims
+                    )
+                model_with_outputs.graph.output.append(value_info)
+                existing_outputs.add(value.name)
+            output_names = [output.name for output in model_with_outputs.graph.output]
+            sess_options = make_deterministic_session_options(ort)
+            sess = ort.InferenceSession(
+                model_with_outputs.SerializeToString(),
+                sess_options=sess_options,
+                providers=["CPUExecutionProvider"],
+            )
+            output_arrays = sess.run(None, self._options.testbench_inputs)
+        except Exception:
+            return graph
+
+        shapes_by_name: dict[str, tuple[int, ...]] = {
+            name: tuple(int(dim) for dim in array.shape)
+            for name, array in zip(output_names, output_arrays)
+        }
+        for name, array in self._options.testbench_inputs.items():
+            shapes_by_name[name] = tuple(int(dim) for dim in array.shape)
+
+        def concretize_value(value: Value) -> Value:
+            shape = shapes_by_name.get(value.name)
+            if shape is None:
+                return value
+            return Value(
+                name=value.name,
+                type=TensorType(
+                    dtype=value.type.dtype,
+                    shape=shape,
+                    dim_params=(None,) * len(shape),
+                ),
+            )
+
+        return Graph(
+            inputs=tuple(concretize_value(value) for value in graph.inputs),
+            outputs=tuple(concretize_value(value) for value in graph.outputs),
+            nodes=graph.nodes,
+            initializers=graph.initializers,
+            values=tuple(concretize_value(value) for value in graph.values),
+            opset_imports=graph.opset_imports,
+        )
+
     def _validate_graph(self, graph: Graph) -> None:
         if not graph.outputs:
             raise UnsupportedOpError("Graph must have at least one output")
         if not graph.nodes:
             raise UnsupportedOpError("Graph must contain at least one node")
         for value in graph.outputs:
-            element_count = shape_product(value.type.shape)
-            if element_count <= 0:
-                raise ShapeInferenceError("Output shape must be fully defined")
+            shape_product(value.type.shape)
 
     def _collect_io_specs(
         self, graph: Graph
@@ -330,11 +486,14 @@ class Compiler:
             | UnaryOp
             | ClipOp
             | CastOp
+            | QuantizeLinearOp
             | MatMulOp
             | GemmOp
             | AttentionOp
             | ConvOp
+            | ConvTransposeOp
             | AveragePoolOp
+            | LpPoolOp
             | BatchNormOp
             | LpNormalizationOp
             | InstanceNormalizationOp
@@ -346,12 +505,15 @@ class Compiler:
             | LstmOp
             | SoftmaxOp
             | LogSoftmaxOp
+            | HardmaxOp
             | NegativeLogLikelihoodLossOp
             | SoftmaxCrossEntropyLossOp
             | MaxPoolOp
             | ConcatOp
             | GatherElementsOp
             | GatherOp
+            | GatherNDOp
+            | ScatterNDOp
             | TransposeOp
             | ConstantOfShapeOp
             | ReshapeOp
@@ -362,9 +524,11 @@ class Compiler:
             | ArgReduceOp
             | ShapeOp
             | PadOp
+            | NonZeroOp
             | ExpandOp
             | CumSumOp
             | RangeOp
+            | OneHotOp
             | SplitOp
         ],
         list[NodeInfo],
@@ -375,11 +539,14 @@ class Compiler:
             | UnaryOp
             | ClipOp
             | CastOp
+            | QuantizeLinearOp
             | MatMulOp
             | GemmOp
             | AttentionOp
             | ConvOp
+            | ConvTransposeOp
             | AveragePoolOp
+            | LpPoolOp
             | BatchNormOp
             | LpNormalizationOp
             | InstanceNormalizationOp
@@ -391,12 +558,14 @@ class Compiler:
             | LstmOp
             | SoftmaxOp
             | LogSoftmaxOp
+            | HardmaxOp
             | NegativeLogLikelihoodLossOp
             | SoftmaxCrossEntropyLossOp
             | MaxPoolOp
             | ConcatOp
             | GatherElementsOp
             | GatherOp
+            | GatherNDOp
             | TransposeOp
             | ConstantOfShapeOp
             | ReshapeOp
@@ -406,9 +575,11 @@ class Compiler:
             | ArgReduceOp
             | ShapeOp
             | PadOp
+            | NonZeroOp
             | ExpandOp
             | CumSumOp
             | RangeOp
+            | OneHotOp
             | SplitOp
             | WhereOp
         ] = []
@@ -515,6 +686,8 @@ def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
         op_spec = binary_op_symbol(function, node.attrs, dtype=op_dtype)
         if op_spec is None:
             raise UnsupportedOpError("Unsupported op BitShift")
+        input0_shape = value_shape(graph, node.inputs[0], node)
+        input1_shape = value_shape(graph, node.inputs[1], node)
         output_shape = value_shape(graph, node.outputs[0], node)
         return BinaryOp(
             input0=node.inputs[0],
@@ -522,6 +695,8 @@ def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
             output=node.outputs[0],
             function=function,
             operator_kind=op_spec.kind,
+            input0_shape=input0_shape,
+            input1_shape=input1_shape,
             shape=output_shape,
             dtype=op_dtype,
             input_dtype=op_dtype,
@@ -555,6 +730,8 @@ def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
             raise UnsupportedOpError(
                 f"{node.op_type} expects bool output, got {output_dtype.onnx_name}"
             )
+        input0_shape = value_shape(graph, node.inputs[0], node)
+        input1_shape = value_shape(graph, node.inputs[1], node)
         output_shape = value_shape(graph, node.outputs[0], node)
         return BinaryOp(
             input0=node.inputs[0],
@@ -562,6 +739,8 @@ def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
             output=node.outputs[0],
             function=function,
             operator_kind=op_spec.kind,
+            input0_shape=input0_shape,
+            input1_shape=input1_shape,
             shape=output_shape,
             dtype=output_dtype,
             input_dtype=input_dtype,
@@ -576,6 +755,8 @@ def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
             raise UnsupportedOpError(
                 f"{node.op_type} must have 2 inputs and 1 output"
             )
+        input0_shape = value_shape(graph, node.inputs[0], node)
+        input1_shape = value_shape(graph, node.inputs[1], node)
         output_shape = value_shape(graph, node.outputs[0], node)
         return BinaryOp(
             input0=node.inputs[0],
@@ -583,6 +764,8 @@ def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
             output=node.outputs[0],
             function=function,
             operator_kind=op_spec.kind,
+            input0_shape=input0_shape,
+            input1_shape=input1_shape,
             shape=output_shape,
             dtype=op_dtype,
             input_dtype=op_dtype,

emx_onnx_cgen/ir/model.py

@@ -44,6 +44,7 @@ class Graph:
     nodes: tuple[Node, ...]
     initializers: tuple[Initializer, ...]
     values: tuple[Value, ...] = ()
+    opset_imports: tuple[tuple[str, int], ...] = ()
 
     def find_value(self, name: str) -> Value:
         for value in self.inputs + self.outputs + self.values: