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

emx_onnx_cgen/codegen/emitter.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .c_emitter import CEmitter
+
+__all__ = ["CEmitter"]

emx_onnx_cgen/compiler.py

@@ -10,141 +10,24 @@ import onnx
 
 from shared.scalar_types import ScalarType
 
+from .onnxruntime_utils import make_deterministic_session_options
 from .codegen.c_emitter import (
-    AttentionOp,
-    AveragePoolOp,
-    BatchNormOp,
-    LpNormalizationOp,
-    InstanceNormalizationOp,
-    GroupNormalizationOp,
-    LayerNormalizationOp,
-    MeanVarianceNormalizationOp,
-    RMSNormalizationOp,
-    BinaryOp,
-    MultiInputBinaryOp,
-    CastOp,
-    ClipOp,
     CEmitter,
     ConstTensor,
-    ConvOp,
-    ConcatOp,
-    ConstantOfShapeOp,
-    CumSumOp,
-    GemmOp,
-    GatherOp,
-    GatherElementsOp,
-    ExpandOp,
-    RangeOp,
-    LrnOp,
-    LstmOp,
-    LogSoftmaxOp,
-    NegativeLogLikelihoodLossOp,
-    NodeInfo,
-    PadOp,
-    SplitOp,
-    SoftmaxCrossEntropyLossOp,
     LoweredModel,
     ModelHeader,
-    MatMulOp,
-    MaxPoolOp,
-    ReduceOp,
-    ArgReduceOp,
-    ReshapeOp,
-    ResizeOp,
-    GridSampleOp,
-    SoftmaxOp,
-    ShapeOp,
-    SliceOp,
-    TransposeOp,
-    UnaryOp,
-    WhereOp,
+    NodeInfo,
 )
 from .dtypes import dtype_info
 from .errors import CodegenError, ShapeInferenceError, UnsupportedOpError
-from .ir.model import Graph, Value
-from .lowering.attention import AttentionSpec, resolve_attention_spec
-from .lowering.average_pool import (
-    lower_average_pool,
-    lower_global_average_pool,
-)
-from .lowering.batch_normalization import lower_batch_normalization
-from .lowering.cast import lower_cast
-from .lowering.concat import lower_concat
-from .lowering.common import (
-    ensure_supported_dtype,
-    node_dtype,
-    shape_product,
-    value_dtype,
-    value_shape,
-)
-from .lowering.conv import ConvSpec, resolve_conv_spec
-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.flatten import lower_flatten
-from .lowering.gather import lower_gather
-from .lowering.gather_elements import lower_gather_elements
-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 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 mean_variance_normalization as _mean_variance_normalization  # noqa: F401
-from .lowering.negative_log_likelihood_loss import (
-    lower_negative_log_likelihood_loss,
-)
-from .lowering.expand import lower_expand
-from .lowering.range import lower_range
-from .lowering.split import lower_split
-from .lowering.softmax_cross_entropy_loss import (
-    lower_softmax_cross_entropy_loss,
-)
-from .lowering.matmul import lower_matmul
-from .lowering.maxpool import MaxPoolSpec, resolve_maxpool_spec
-from .lowering import pad as _pad  # noqa: F401
-from .lowering.reduce import (
-    REDUCE_KIND_BY_OP,
-    REDUCE_OUTPUTS_FLOAT_ONLY,
-)
-from .lowering import arg_reduce as _arg_reduce  # noqa: F401
-from .lowering.reshape import lower_reshape
-from .lowering.resize import lower_resize
-from .lowering.grid_sample import lower_grid_sample
-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.shape import lower_shape
-from .lowering.size import lower_size
-from .lowering.softmax import lower_softmax
-from .lowering.transpose import lower_transpose
-from .lowering.unsqueeze import lower_unsqueeze
-from .lowering.where import lower_where
-from .lowering.elementwise import (
-    lower_celu,
-    lower_clip,
-    lower_isinf,
-    lower_isnan,
-    lower_shrink,
-    lower_swish,
-)
-from .lowering import variadic as _variadic  # noqa: F401
-from .lowering import rms_normalization as _rms_normalization  # noqa: F401
-from .lowering.registry import get_lowering_registry, resolve_dispatch
+from .ir.context import GraphContext
+from .ir.model import Graph, TensorType, Value
+from .ir.op_base import OpBase
+from .ir.op_context import OpContext
+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 .ops import (
-    BINARY_OP_TYPES,
-    COMPARE_FUNCTIONS,
-    UNARY_OP_TYPES,
-    binary_op_symbol,
-    unary_op_symbol,
-    validate_unary_attrs,
-)
-from shared.scalar_functions import ScalarFunction, ScalarFunctionError
 from .runtime.evaluator import Evaluator
 
 
@@ -157,6 +40,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 * 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 +58,17 @@ 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,
         )
+        load_lowering_registry()
 
     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 +84,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 +98,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,
@@ -219,7 +159,8 @@ class Compiler:
         return collect(graph.inputs), collect(graph.outputs)
 
     def _lower_model(self, model: onnx.ModelProto, graph: Graph) -> LoweredModel:
-        constants = _lowered_constants(graph)
+        ctx = GraphContext(graph)
+        constants = _lowered_constants(ctx)
         self._validate_graph(graph)
         (
             input_names,
@@ -229,7 +170,14 @@ class Compiler:
             output_shapes,
             output_dtypes,
         ) = self._collect_io_specs(graph)
-        ops, node_infos = self._lower_nodes(graph)
+        ops, node_infos = self._lower_nodes(ctx)
+        op_ctx = OpContext(ctx)
+        for op in ops:
+            op.validate(op_ctx)
+        for op in ops:
+            op.infer_types(op_ctx)
+        for op in ops:
+            op.infer_shapes(op_ctx)
         header = self._build_header(model, graph)
         return LoweredModel(
             name=self._options.model_name,
@@ -243,6 +191,7 @@ class Compiler:
             ops=tuple(ops),
             node_infos=tuple(node_infos),
             header=header,
+            op_context=op_ctx,
         )
 
     def _resolve_testbench_inputs(
@@ -282,15 +231,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
@@ -322,107 +349,16 @@ class Compiler:
         )
 
     def _lower_nodes(
-        self, graph: Graph
-    ) -> tuple[
-        list[
-            BinaryOp
-            | MultiInputBinaryOp
-            | UnaryOp
-            | ClipOp
-            | CastOp
-            | MatMulOp
-            | GemmOp
-            | AttentionOp
-            | ConvOp
-            | AveragePoolOp
-            | BatchNormOp
-            | LpNormalizationOp
-            | InstanceNormalizationOp
-            | GroupNormalizationOp
-            | LayerNormalizationOp
-            | MeanVarianceNormalizationOp
-            | RMSNormalizationOp
-            | LrnOp
-            | LstmOp
-            | SoftmaxOp
-            | LogSoftmaxOp
-            | NegativeLogLikelihoodLossOp
-            | SoftmaxCrossEntropyLossOp
-            | MaxPoolOp
-            | ConcatOp
-            | GatherElementsOp
-            | GatherOp
-            | TransposeOp
-            | ConstantOfShapeOp
-            | ReshapeOp
-            | SliceOp
-            | ResizeOp
-            | GridSampleOp
-            | ReduceOp
-            | ArgReduceOp
-            | ShapeOp
-            | PadOp
-            | ExpandOp
-            | CumSumOp
-            | RangeOp
-            | SplitOp
-        ],
-        list[NodeInfo],
-    ]:
-        ops: list[
-            BinaryOp
-            | MultiInputBinaryOp
-            | UnaryOp
-            | ClipOp
-            | CastOp
-            | MatMulOp
-            | GemmOp
-            | AttentionOp
-            | ConvOp
-            | AveragePoolOp
-            | BatchNormOp
-            | LpNormalizationOp
-            | InstanceNormalizationOp
-            | GroupNormalizationOp
-            | LayerNormalizationOp
-            | MeanVarianceNormalizationOp
-            | RMSNormalizationOp
-            | LrnOp
-            | LstmOp
-            | SoftmaxOp
-            | LogSoftmaxOp
-            | NegativeLogLikelihoodLossOp
-            | SoftmaxCrossEntropyLossOp
-            | MaxPoolOp
-            | ConcatOp
-            | GatherElementsOp
-            | GatherOp
-            | TransposeOp
-            | ConstantOfShapeOp
-            | ReshapeOp
-            | SliceOp
-            | ResizeOp
-            | ReduceOp
-            | ArgReduceOp
-            | ShapeOp
-            | PadOp
-            | ExpandOp
-            | CumSumOp
-            | RangeOp
-            | SplitOp
-            | WhereOp
-        ] = []
+        self, ctx: GraphContext
+    ) -> tuple[list[OpBase], list[NodeInfo]]:
+        ops: list[OpBase] = []
         node_infos: list[NodeInfo] = []
-        for node in graph.nodes:
-            lowering = resolve_dispatch(
-                node.op_type,
-                get_lowering_registry(),
-                binary_types=BINARY_OP_TYPES,
-                unary_types=UNARY_OP_TYPES,
-                binary_fallback=lambda: _lower_binary_unary,
-                unary_fallback=lambda: _lower_binary_unary,
-            )
-            ops.append(lowering(graph, node))
+        registry = get_lowering_registry()
+        for node in ctx.nodes:
+            lowering = registry.get(node.op_type)
+            if lowering is None:
+                raise UnsupportedOpError(f"Unsupported op {node.op_type}")
+            ops.append(lowering(ctx, node))
             node_infos.append(
                 NodeInfo(
                     op_type=node.op_type,
@@ -473,7 +409,7 @@ class Compiler:
         return evaluator.run(feeds)
 
 
-def _lowered_constants(graph: Graph) -> tuple[ConstTensor, ...]:
+def _lowered_constants(graph: Graph | GraphContext) -> tuple[ConstTensor, ...]:
     constants: list[ConstTensor] = []
     for initializer in graph.initializers:
         dtype = ensure_supported_dtype(initializer.type.dtype)
@@ -489,113 +425,3 @@ def _lowered_constants(graph: Graph) -> tuple[ConstTensor, ...]:
             )
         )
     return tuple(constants)
-
-
-def _lower_binary_unary(graph: Graph, node: Node) -> BinaryOp | UnaryOp:
-    if node.op_type == "BitShift":
-        if len(node.inputs) != 2 or len(node.outputs) != 1:
-            raise UnsupportedOpError("BitShift must have 2 inputs and 1 output")
-        direction_attr = node.attrs.get("direction", "LEFT")
-        if isinstance(direction_attr, bytes):
-            direction = direction_attr.decode()
-        else:
-            direction = str(direction_attr)
-        if direction not in {"LEFT", "RIGHT"}:
-            raise UnsupportedOpError(
-                "BitShift direction must be LEFT or RIGHT"
-            )
-        op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
-        if not op_dtype.is_integer:
-            raise UnsupportedOpError("BitShift expects integer inputs")
-        function = (
-            ScalarFunction.BITWISE_LEFT_SHIFT
-            if direction == "LEFT"
-            else ScalarFunction.BITWISE_RIGHT_SHIFT
-        )
-        op_spec = binary_op_symbol(function, node.attrs, dtype=op_dtype)
-        if op_spec is None:
-            raise UnsupportedOpError("Unsupported op BitShift")
-        output_shape = value_shape(graph, node.outputs[0], node)
-        return BinaryOp(
-            input0=node.inputs[0],
-            input1=node.inputs[1],
-            output=node.outputs[0],
-            function=function,
-            operator_kind=op_spec.kind,
-            shape=output_shape,
-            dtype=op_dtype,
-            input_dtype=op_dtype,
-        )
-    if node.op_type == "Mod":
-        fmod = int(node.attrs.get("fmod", 0))
-        if fmod not in {0, 1}:
-            raise UnsupportedOpError("Mod only supports fmod=0 or fmod=1")
-        function = (
-            ScalarFunction.FMOD if fmod == 1 else ScalarFunction.REMAINDER
-        )
-    else:
-        try:
-            function = ScalarFunction.from_onnx_op(node.op_type)
-        except ScalarFunctionError as exc:
-            raise UnsupportedOpError(
-                f"Unsupported op {node.op_type}"
-            ) from exc
-    validate_unary_attrs(node.op_type, node.attrs)
-    if function in COMPARE_FUNCTIONS:
-        input_dtype = node_dtype(graph, node, *node.inputs)
-        output_dtype = value_dtype(graph, node.outputs[0], node)
-        op_spec = binary_op_symbol(function, node.attrs, dtype=input_dtype)
-        if op_spec is None:
-            raise UnsupportedOpError(f"Unsupported op {node.op_type}")
-        if len(node.inputs) != 2 or len(node.outputs) != 1:
-            raise UnsupportedOpError(
-                f"{node.op_type} must have 2 inputs and 1 output"
-            )
-        if output_dtype != ScalarType.BOOL:
-            raise UnsupportedOpError(
-                f"{node.op_type} expects bool output, got {output_dtype.onnx_name}"
-            )
-        output_shape = value_shape(graph, node.outputs[0], node)
-        return BinaryOp(
-            input0=node.inputs[0],
-            input1=node.inputs[1],
-            output=node.outputs[0],
-            function=function,
-            operator_kind=op_spec.kind,
-            shape=output_shape,
-            dtype=output_dtype,
-            input_dtype=input_dtype,
-        )
-    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
-    op_spec = binary_op_symbol(function, node.attrs, dtype=op_dtype)
-    unary_symbol = unary_op_symbol(function, dtype=op_dtype)
-    if op_spec is None and unary_symbol is None:
-        raise UnsupportedOpError(f"Unsupported op {node.op_type}")
-    if op_spec is not None:
-        if len(node.inputs) != 2 or len(node.outputs) != 1:
-            raise UnsupportedOpError(
-                f"{node.op_type} must have 2 inputs and 1 output"
-            )
-        output_shape = value_shape(graph, node.outputs[0], node)
-        return BinaryOp(
-            input0=node.inputs[0],
-            input1=node.inputs[1],
-            output=node.outputs[0],
-            function=function,
-            operator_kind=op_spec.kind,
-            shape=output_shape,
-            dtype=op_dtype,
-            input_dtype=op_dtype,
-        )
-    if len(node.inputs) != 1 or len(node.outputs) != 1:
-        raise UnsupportedOpError(f"{node.op_type} must have 1 input and 1 output")
-    output_shape = value_shape(graph, node.outputs[0], node)
-    return UnaryOp(
-        input0=node.inputs[0],
-        output=node.outputs[0],
-        function=function,
-        shape=output_shape,
-        dtype=op_dtype,
-        input_dtype=op_dtype,
-        params=(),
-    )

emx_onnx_cgen/ir/context.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from .model import Graph, Initializer, Node, Value
+from shared.scalar_types import ScalarType
+
+
+@dataclass
+class GraphContext:
+    graph: Graph
+    _dtype_cache: dict[str, ScalarType] = field(default_factory=dict)
+    _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)
+
+    def find_value(self, name: str) -> Value:
+        return self.graph.find_value(name)
+
+    def dtype(self, name: str, node: Node | None = None) -> ScalarType:
+        if name in self._dtype_cache:
+            return self._dtype_cache[name]
+        try:
+            value = self.graph.find_value(name)
+        except KeyError as exc:
+            op_type = node.op_type if node is not None else "unknown"
+            raise ShapeInferenceError(
+                f"Missing dtype for value '{name}' in op {op_type}. "
+                "Hint: run ONNX shape inference or export with static shapes."
+            ) from exc
+        dtype = value.type.dtype
+        if not isinstance(dtype, ScalarType):
+            raise UnsupportedOpError(f"Unsupported dtype {dtype}")
+        self._dtype_cache[name] = dtype
+        return dtype
+
+    def set_dtype(self, name: str, dtype: ScalarType) -> None:
+        self._dtype_cache[name] = dtype
+
+    def shape(self, name: str, node: Node | None = None) -> tuple[int, ...]:
+        if name in self._shape_cache:
+            return self._shape_cache[name]
+        try:
+            value = self.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
+        self._shape_cache[name] = value.type.shape
+        return value.type.shape
+
+    def set_shape(self, name: str, shape: tuple[int, ...]) -> None:
+        self._shape_cache[name] = shape
+
+    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
+
+    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
+
+    def opset_version(self, domain: str = "") -> int | None:
+        if domain in {"", "ai.onnx"}:
+            domains = {"", "ai.onnx"}
+        else:
+            domains = {domain}
+        for opset_domain, version in self.graph.opset_imports:
+            if opset_domain in domains:
+                return int(version)
+        return None
+
+    def __getattr__(self, name: str):
+        return getattr(self.graph, name)