diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index af69d1e9e..0af4c4014 100644 --- a/src/winml/modelkit/optim/analysis.py +++ b/src/winml/modelkit/optim/analysis.py @@ -15,7 +15,9 @@ only that capability enabled (plus auto-enabled dependencies), and the resulting graph is diffed against the baseline. A non-empty diff means the capability is applicable; the diff itself names the affected nodes and - constants. + constants. When one capability is intentionally owned by multiple pipes, it + is probed once across the full pipeline so the report matches the single + public ``--enable-*`` flag. No operator names, tensor names, or architectures are hardcoded — every result is derived from the concrete graph diff. @@ -233,9 +235,10 @@ def _initializers_equal(base: TensorProto, probe: TensorProto) -> bool: float_field_types = {"float_data": "f", "double_data": "d"} for field_name, typecode in float_field_types.items(): - if array(typecode, getattr(base, field_name)).tobytes() != array( - typecode, getattr(probe, field_name) - ).tobytes(): + if ( + array(typecode, getattr(base, field_name)).tobytes() + != array(typecode, getattr(probe, field_name)).tobytes() + ): return False if protobuf_equal: @@ -317,6 +320,19 @@ def _run_pipe(pipe: Any, model: ModelProto, config: Any) -> ModelProto: return result +def _run_pipeline( + pipe_classes: list[type[Any]], + model: ModelProto, + kwargs: dict[str, Any], +) -> ModelProto: + """Run the capability-driven pipe sequence from ``model`` with ``kwargs``.""" + current = model + for pipe_class in pipe_classes: + pipe = pipe_class() + current = _run_pipe(pipe, current, pipe.build_config(**kwargs)) + return current + + def _iter_findings( model: ModelProto, capabilities: dict[str, CapabilityDef], @@ -344,14 +360,30 @@ def _iter_findings( from .pipes import PIPES from .registry import BoolCapability, auto_enable_dependencies - remaining_probes = Counter( + pipe_probe_counts = Counter( name for pipe_class in PIPES for name, cap in pipe_class.capabilities.items() if isinstance(cap, BoolCapability) and not cap.default ) + remaining_probes = Counter(dict.fromkeys(pipe_probe_counts, 1)) + shared_cap_names = { + name for name, count in pipe_probe_counts.items() if count > 1 and name in capabilities + } + shared_cap_owners = { + name: [ + pipe_class.name + for pipe_class in PIPES + if name in pipe_class.capabilities + and isinstance(pipe_class.capabilities[name], BoolCapability) + and not pipe_class.capabilities[name].default + ] + for name in shared_cap_names + } def complete_probe(cap_name: str) -> None: + if remaining_probes[cap_name] <= 0: + return remaining_probes[cap_name] -= 1 if remaining_probes[cap_name] == 0 and on_probe_complete is not None: on_probe_complete(cap_name) @@ -366,6 +398,14 @@ def complete_probe(cap_name: str) -> None: # limit it round-trips through save_onnx with external data), and this # function must never modify the caller's input model. current = infer_shapes(_clone(model)) + pipeline_input = current + full_base_out: ModelProto | None = None + + def full_pipeline_baseline() -> ModelProto: + nonlocal full_base_out + if full_base_out is None: + full_base_out = _run_pipeline(PIPES, pipeline_input, default_kwargs) + return full_base_out for pipe_class in PIPES: pipe = pipe_class() @@ -408,6 +448,87 @@ def complete_probe(cap_name: str) -> None: ) for cap_name, cap in probe_caps: + if cap_name in shared_cap_names: + owners = shared_cap_owners[cap_name] + if owners and owners[0] == pipe.name: + if on_probe_start is not None: + on_probe_start(cap_name) + try: + ep_device = optimizer_kwargs.get("ep_device") + if ep_device is not None and cap.ep_constraint is not None: + from ..utils.constants import normalize_ep_name + + target_ep = normalize_ep_name(ep_device.device.ep_name) + if not any( + normalize_ep_name(name) == target_ep + for name in cap.ep_constraint + ): + logger.debug( + "Skipping capability '%s': target EP %s is not in %s", + cap.name, + target_ep, + cap.ep_constraint, + ) + continue + + kebab = dict(kebab_defaults) + kebab[cap_name] = True + kebab = auto_enable_dependencies(kebab, capabilities) + probe_kwargs = { + capabilities[name].python_name: value + for name, value in kebab.items() + if name in capabilities + } + probe_kwargs.update(optimizer_kwargs) + + try: + shared_base_out = full_pipeline_baseline() + probe_out = _run_pipeline(PIPES, pipeline_input, probe_kwargs) + except Exception as exc: + logger.warning( + "Could not evaluate shared capability '%s': %s", + cap_name, + exc, + ) + continue + + shared_base_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {} + shared_probe_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {} + _collect_nodes(shared_base_out.graph, (), shared_base_nodes) + _collect_nodes(probe_out.graph, (), shared_probe_nodes) + removed, added, modified = _diff_nodes( + shared_base_nodes, + shared_probe_nodes, + ) + + shared_base_inits = _collect_initializers(shared_base_out) + probe_inits = _collect_initializers(probe_out) + rem_init, add_init, mod_init = _diff_initializers( + shared_base_inits, + probe_inits, + ) + + finding = CapabilityFinding( + name=cap.name, + python_name=cap.python_name, + enable_flag=f"--enable-{cap.name}", + category=cap.category.value, + description=cap.description, + pipe_name="+".join(owners), + removed_nodes=removed, + added_nodes=added, + modified_nodes=modified, + removed_initializers=rem_init, + added_initializers=add_init, + modified_initializers=mod_init, + ) + + if finding.applicable: + yield finding, probe_out + finally: + complete_probe(cap_name) + continue + if on_probe_start is not None: on_probe_start(cap_name) try: diff --git a/src/winml/modelkit/optim/capabilities/algebraic.py b/src/winml/modelkit/optim/capabilities/algebraic.py index a6ac37655..a309d7a15 100644 --- a/src/winml/modelkit/optim/capabilities/algebraic.py +++ b/src/winml/modelkit/optim/capabilities/algebraic.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Opt-in, exact algebraic graph-rewrite capabilities.""" +"""Opt-in algebraic graph-rewrite capabilities.""" from __future__ import annotations @@ -30,3 +30,14 @@ category=CapabilityCategory.REWRITE, default=False, ) + +EXP_POSITIVE_SCALE_FOLDING = BoolCapability( + name="exp-positive-scale-folding", + ort_name=None, + description=( + "Fold a finite, strictly positive constant scale after Exp into the log-domain input " + "bias with relaxed floating-point overflow semantics" + ), + category=CapabilityCategory.REWRITE, + default=False, +) diff --git a/src/winml/modelkit/optim/capabilities/misc.py b/src/winml/modelkit/optim/capabilities/misc.py index cc54a9c46..3b2c99814 100644 --- a/src/winml/modelkit/optim/capabilities/misc.py +++ b/src/winml/modelkit/optim/capabilities/misc.py @@ -22,7 +22,7 @@ GATHER_SLICE_TO_SPLIT_FUSION = BoolCapability( name="gather-slice-to-split-fusion", ort_name="GatherSliceToSplitFusion", - description="Fuse Gather+Slice patterns to Split operation", + description="Fuse eligible Gather/Slice routes to Split operations", category=CapabilityCategory.MISC, default=False, ) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 29aac986e..161d6a70c 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -6,20 +6,31 @@ from __future__ import annotations +import math from dataclasses import dataclass -from typing import Any, ClassVar, cast +from itertools import pairwise +from typing import TYPE_CHECKING, Any, ClassVar, cast import numpy as np import onnx -from ..capabilities import algebraic +from ..capabilities import algebraic, misc from .base import BasePipe, PipeConfig, caps_dict +if TYPE_CHECKING: + from collections.abc import Iterable + + ALGEBRAIC_CAPABILITIES: dict[str, Any] = caps_dict( algebraic.STATIC_SPLIT_TO_SLICE, + misc.GATHER_SLICE_TO_SPLIT_FUSION, algebraic.CONV_CHANNEL_AFFINE_FOLDING, + algebraic.EXP_POSITIVE_SCALE_FOLDING, ) +MAX_AFFINE_ROUTE_DEPTH = 64 +MAX_NUMPY_ELEMENTS = np.iinfo(np.intp).max +ORDER_PRESERVING_VIEW_OPS = frozenset({"Flatten", "Identity", "Reshape", "Squeeze", "Unsqueeze"}) @dataclass @@ -27,7 +38,9 @@ class AlgebraicRewritePipeConfig(PipeConfig): """Configuration for exact algebraic rewrites.""" static_split_to_slice: bool = False + sibling_slice_to_split: bool = False conv_channel_affine_folding: bool = False + exp_positive_scale_folding: bool = False @dataclass @@ -35,19 +48,36 @@ class _GraphIndex: """Graph metadata required to identify statically bounded Split nodes.""" producers: dict[str, onnx.NodeProto] + definition_collisions: set[str] + has_cycle: bool consumers: dict[str, list[onnx.NodeProto]] initializers: dict[str, onnx.TensorProto] shapes: dict[str, tuple[int | None, ...]] + graph_inputs: set[str] graph_outputs: set[str] @classmethod def build(cls, model: onnx.ModelProto) -> _GraphIndex: graph = model.graph producers: dict[str, onnx.NodeProto] = {} + definition_collisions: set[str] = set() consumers: dict[str, list[onnx.NodeProto]] = {} + graph_input_names = [value.name for value in graph.input if value.name] + initializer_names = [ + initializer.name for initializer in graph.initializer if initializer.name + ] + for names in (graph_input_names, initializer_names): + seen: set[str] = set() + for name in names: + if name in seen: + definition_collisions.add(name) + seen.add(name) + protected_definitions = set(graph_input_names) | set(initializer_names) for node in graph.node: for output in node.output: if output: + if output in protected_definitions or output in producers: + definition_collisions.add(output) producers[output] = node consumed_names = {input_name for input_name in node.input if input_name} for attribute in node.attribute: @@ -59,6 +89,28 @@ def build(cls, model: onnx.ModelProto) -> _GraphIndex: for input_name in consumed_names: consumers.setdefault(input_name, []).append(node) + node_indexes = {id(node): index for index, node in enumerate(graph.node)} + successors: list[set[int]] = [set() for _ in graph.node] + indegrees = [0] * len(graph.node) + for consumer_index, node in enumerate(graph.node): + predecessor_indexes = { + node_indexes[id(producer)] + for input_name in node.input + if input_name and (producer := producers.get(input_name)) is not None + } + indegrees[consumer_index] = len(predecessor_indexes) + for predecessor_index in predecessor_indexes: + successors[predecessor_index].add(consumer_index) + ready = [index for index, indegree in enumerate(indegrees) if indegree == 0] + visited_count = 0 + while ready: + node_index = ready.pop() + visited_count += 1 + for successor_index in successors[node_index]: + indegrees[successor_index] -= 1 + if indegrees[successor_index] == 0: + ready.append(successor_index) + initializers = {initializer.name: initializer for initializer in graph.initializer} shapes: dict[str, tuple[int | None, ...]] = {} for value_info in (*graph.input, *graph.value_info, *graph.output): @@ -70,9 +122,12 @@ def build(cls, model: onnx.ModelProto) -> _GraphIndex: return cls( producers=producers, + definition_collisions=definition_collisions, + has_cycle=visited_count != len(graph.node), consumers=consumers, initializers=initializers, shapes=shapes, + graph_inputs={value.name for value in graph.input if value.name}, graph_outputs={output.name for output in graph.output if output.name}, ) @@ -91,6 +146,42 @@ class _AffineCandidate: offset: np.ndarray +@dataclass +class _ExpScaleCandidate: + """A positive post-Exp scale that can be moved before Exp with an existing bias.""" + + add: onnx.NodeProto + bias_input_index: int + output_node: onnx.NodeProto + muls: list[onnx.NodeProto] + add_output: str + route_consumer: onnx.NodeProto + combined_bias: np.ndarray | None + log_scale: np.ndarray | None + + +@dataclass +class _ExpScaleInsertCandidate: + """A positive post-Exp scale that requires a new input Add.""" + + exp: onnx.NodeProto + output_node: onnx.NodeProto + muls: list[onnx.NodeProto] + log_scale: np.ndarray + + +@dataclass +class _StaticSliceCandidate: + """A one-axis static Slice that can participate in a sibling Split.""" + + node: onnx.NodeProto + input_name: str + output_name: str + axis: int + start: int + end: int + + class _NameAllocator: """Allocate names without relying on optional or duplicated node names.""" @@ -108,14 +199,16 @@ def __init__(self, model: onnx.ModelProto) -> None: ) if name } + self._next_suffix: dict[str, int] = {} def new(self, prefix: str) -> str: - candidate = prefix - suffix = 0 + suffix = self._next_suffix.get(prefix, 0) + candidate = prefix if suffix == 0 else f"{prefix}_{suffix}" while candidate in self._used: suffix += 1 candidate = f"{prefix}_{suffix}" self._used.add(candidate) + self._next_suffix[prefix] = suffix + 1 return candidate @@ -140,30 +233,52 @@ def _attribute(node: onnx.NodeProto, name: str, default: Any = None) -> Any: return default +def _is_standard_onnx_node(node: onnx.NodeProto) -> bool: + return node.domain == "" + + def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: """Read an initializer or a regular ONNX Constant value.""" - if not name: + if not name or name in index.graph_inputs: return None initializer = index.initializers.get(name) if initializer is not None: - return np.asarray(onnx.numpy_helper.to_array(initializer)) + if initializer.data_location == onnx.TensorProto.EXTERNAL and not initializer.raw_data: + return None + try: + return np.asarray(onnx.numpy_helper.to_array(initializer)) + except (TypeError, ValueError, RuntimeError, onnx.checker.ValidationError): + return None producer = index.producers.get(name) - if producer is None or producer.op_type != "Constant": + if producer is None or not _is_standard_onnx_node(producer) or producer.op_type != "Constant": return None value = _attribute(producer, "value") if value is not None: + if value.data_location == onnx.TensorProto.EXTERNAL and not value.raw_data: + return None try: return np.asarray(onnx.numpy_helper.to_array(value)) - except (TypeError, ValueError): + except (TypeError, ValueError, RuntimeError, onnx.checker.ValidationError): return None - for attribute_name in ("value_float", "value_floats", "value_int", "value_ints"): + for attribute_name, dtype in ( + ("value_float", np.float32), + ("value_floats", np.float32), + ("value_int", np.int64), + ("value_ints", np.int64), + ): attribute_value = _attribute(producer, attribute_name) if attribute_value is not None: - return np.asarray(attribute_value) + return np.asarray(attribute_value, dtype=dtype) return None +def _initializer_array(index: _GraphIndex, name: str) -> np.ndarray | None: + if not name or name not in index.initializers: + return None + return _constant_array(index, name) + + def _constant_ints(index: _GraphIndex, name: str) -> list[int] | None: values = _constant_array(index, name) if values is None or not np.issubdtype(values.dtype, np.integer): @@ -211,6 +326,34 @@ def _static_shape(index: _GraphIndex, name: str) -> tuple[int, ...] | None: return cast("tuple[int, ...]", shape) +def _shape_broadcasts_to(shape: tuple[int, ...], target_shape: tuple[int, ...]) -> bool: + if len(shape) > len(target_shape): + return False + padded_shape = (1,) * (len(target_shape) - len(shape)) + shape + return all( + source_dimension in (1, target_dimension) + for source_dimension, target_dimension in zip(padded_shape, target_shape, strict=True) + ) + + +def _shape_element_count(shape: tuple[int, ...]) -> int: + count = math.prod(shape) + return count if count <= MAX_NUMPY_ELEMENTS else -1 + + +def _same_shape_element_count( + left: tuple[int, ...], + right: tuple[int, ...], +) -> bool: + left_count = _shape_element_count(left) + return left_count >= 0 and left_count == _shape_element_count(right) + + +def _strict_default_opset_version(model: onnx.ModelProto) -> int | None: + versions = [int(opset.version) for opset in model.opset_import if opset.domain == ""] + return versions[0] if len(versions) == 1 else None + + def _new_initializer( model: onnx.ModelProto, allocator: _NameAllocator, @@ -319,7 +462,16 @@ def _split_boundaries( ) -> tuple[int, list[tuple[int, int]]] | None: """Return a static Split axis and output boundaries.""" input_shape = index.shapes.get(input_name) - if input_shape is None or len(node.output) == 0: + outputs = list(node.output) + if ( + not _is_standard_onnx_node(node) + or node.op_type != "Split" + or input_shape is None + or not outputs + or any(not output for output in outputs) + or len(set(outputs)) != len(outputs) + or any(output in node.input for output in outputs) + ): return None axis_value = _attribute(node, "axis", 0) @@ -353,6 +505,145 @@ def _split_boundaries( return axis, boundaries +def _normalize_slice_bound(value: int, axis_size: int, *, is_end: bool) -> int: + if value < 0: + value += axis_size + if is_end and value > axis_size: + return axis_size + return max(0, min(value, axis_size)) + + +def _static_slice_candidate( + index: _GraphIndex, + node: onnx.NodeProto, +) -> _StaticSliceCandidate | None: + if ( + not _is_standard_onnx_node(node) + or node.op_type != "Slice" + or len(node.input) < 3 + or not node.input[0] + ): + return None + output_name = _node_output(node) + input_shape = _static_shape(index, node.input[0]) + if output_name is None or input_shape is None: + return None + starts = _constant_ints(index, node.input[1]) + ends = _constant_ints(index, node.input[2]) + if starts is None or ends is None or len(starts) != 1 or len(ends) != 1: + return None + if len(node.input) > 3 and node.input[3]: + axes = _constant_ints(index, node.input[3]) + if axes is None: + return None + else: + axes = [0] + if len(node.input) > 4 and node.input[4]: + steps = _constant_ints(index, node.input[4]) + if steps is None: + return None + else: + steps = [1] + if len(axes) != 1 or len(steps) != 1 or steps[0] != 1: + return None + axis = axes[0] + if axis < -len(input_shape) or axis >= len(input_shape): + return None + axis %= len(input_shape) + axis_size = input_shape[axis] + if axis_size <= 0: + return None + start = _normalize_slice_bound(starts[0], axis_size, is_end=False) + end = _normalize_slice_bound(ends[0], axis_size, is_end=True) + if end <= start: + return None + output_shape = _static_shape(index, output_name) + expected_shape = list(input_shape) + expected_shape[axis] = end - start + if output_shape is not None and output_shape != tuple(expected_shape): + return None + return _StaticSliceCandidate( + node=node, + input_name=node.input[0], + output_name=output_name, + axis=axis, + start=start, + end=end, + ) + + +def _sibling_slice_split_groups( + model: onnx.ModelProto, + index: _GraphIndex, +) -> list[list[_StaticSliceCandidate]]: + grouped: dict[tuple[str, int], list[_StaticSliceCandidate]] = {} + for node in model.graph.node: + candidate = _static_slice_candidate(index, node) + if candidate is not None: + grouped.setdefault((candidate.input_name, candidate.axis), []).append(candidate) + + groups: list[list[_StaticSliceCandidate]] = [] + for (input_name, axis), candidates in grouped.items(): + input_shape = _static_shape(index, input_name) + if input_shape is None or len(candidates) < 2: + continue + ordered = sorted(candidates, key=lambda candidate: candidate.start) + if len({candidate.output_name for candidate in ordered}) != len(ordered): + continue + if ordered[0].start != 0 or ordered[-1].end != input_shape[axis]: + continue + if any(left.end != right.start for left, right in pairwise(ordered)): + continue + groups.append(ordered) + return groups + + +def _fold_sibling_slices_to_split( + model: onnx.ModelProto, + allocator: _NameAllocator, +) -> None: + """Replace contiguous sibling Slice nodes with an equivalent Split.""" + opset = _strict_default_opset_version(model) + if opset is None or opset < 13: + return + index = _GraphIndex.build(model) + groups = _sibling_slice_split_groups(model, index) + if not groups: + return + + node_order = {id(node): position for position, node in enumerate(model.graph.node)} + replacements: dict[int, onnx.NodeProto] = {} + removed: set[int] = set() + for group in groups: + split_values = np.asarray( + [candidate.end - candidate.start for candidate in group], + dtype=np.int64, + ) + split_name = _new_initializer(model, allocator, split_values, "algebraic_slice_splits") + split = onnx.helper.make_node( + "Split", + [group[0].input_name, split_name], + [candidate.output_name for candidate in group], + name=allocator.new("algebraic_slice_split"), + axis=group[0].axis, + ) + first = min(group, key=lambda candidate: node_order[id(candidate.node)]) + replacements[id(first.node)] = split + removed.update( + id(candidate.node) for candidate in group if candidate.node is not first.node + ) + + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + replacement = replacements.get(id(node)) + if replacement is not None: + rewritten.append(replacement) + elif id(node) not in removed: + rewritten.append(node) + del model.graph.node[:] + model.graph.node.extend(rewritten) + + def _slice_channel_boundary( index: _GraphIndex, node: onnx.NodeProto, @@ -360,20 +651,26 @@ def _slice_channel_boundary( channel_axis: int, ) -> tuple[int, int] | None: """Read a Slice that selects a contiguous, full non-channel region.""" - if len(node.input) < 2: + if not _is_standard_onnx_node(node) or node.op_type != "Slice" or len(node.input) < 3: return None input_shape = _static_shape(index, input_name) if input_shape is None or channel_axis >= len(input_shape): return None starts = _constant_ints(index, node.input[1]) - ends = _constant_ints(index, node.input[2]) if len(node.input) > 2 else None - axes = _constant_ints(index, node.input[3]) if len(node.input) > 3 else None - steps = _constant_ints(index, node.input[4]) if len(node.input) > 4 else None + ends = _constant_ints(index, node.input[2]) if starts is None or ends is None: return None - if axes is None: + if len(node.input) > 3 and node.input[3]: + axes = _constant_ints(index, node.input[3]) + if axes is None: + return None + else: axes = list(range(len(starts))) - if steps is None: + if len(node.input) > 4 and node.input[4]: + steps = _constant_ints(index, node.input[4]) + if steps is None: + return None + else: steps = [1] * len(starts) if not (len(starts) == len(ends) == len(axes) == len(steps)): return None @@ -414,6 +711,8 @@ def _channel_affine_values( """Convert a scalar or a provably channel-only broadcast to ``[C]``.""" if not np.issubdtype(values.dtype, np.floating): return None + if not np.isfinite(values).all(): + return None if values.size == 1: return np.full(channels, values.reshape(-1)[0], dtype=values.dtype) if values.ndim > len(output_shape): @@ -458,6 +757,7 @@ def _channel_preserving_view_output( output_name = _node_output(node) if ( output_name is None + or not _is_standard_onnx_node(node) or len(node.input) == 0 or node.input[0] != input_name or node.op_type not in {"Reshape", "Squeeze", "Unsqueeze"} @@ -472,16 +772,16 @@ def _channel_preserving_view_output( or len(output_shape) < 2 or input_shape[:2] != output_shape[:2] or input_shape[1] != channels - or np.prod(input_shape[2:], dtype=np.int64) != np.prod(output_shape[2:], dtype=np.int64) + or not _same_shape_element_count(input_shape[2:], output_shape[2:]) ): return None if node.op_type == "Reshape": - if ( - len(node.input) < 2 - or _constant_ints(index, node.input[1]) is None - or _attribute(node, "allowzero", 0) != 0 - ): + target_shape = _constant_ints(index, node.input[1]) if len(node.input) >= 2 else None + allowzero = _attribute(node, "allowzero", 0) + if target_shape is None or allowzero not in (0, 1): + return None + if allowzero == 1 and 0 in target_shape: return None else: axes, conflict = _single_attribute_or_input_ints(index, node, "axes", 1) @@ -490,6 +790,465 @@ def _channel_preserving_view_output( return output_name +def _order_preserving_view_output( + index: _GraphIndex, + node: onnx.NodeProto, + input_name: str, +) -> str | None: + """Return a static shape-only view output that preserves element order.""" + output_name = _node_output(node) + if ( + output_name is None + or not _is_standard_onnx_node(node) + or not node.input + or node.input[0] != input_name + or node.op_type not in ORDER_PRESERVING_VIEW_OPS + ): + return None + if node.op_type == "Identity": + return output_name + input_shape = _static_shape(index, input_name) + output_shape = _static_shape(index, output_name) + if ( + input_shape is None + or output_shape is None + or any(dimension <= 0 for dimension in (*input_shape, *output_shape)) + or not _same_shape_element_count(input_shape, output_shape) + ): + return None + + if node.op_type == "Flatten": + axis = _attribute(node, "axis", 1) + if not isinstance(axis, int): + return None + rank = len(input_shape) + normalized_axis = axis + rank if axis < 0 else axis + if normalized_axis < 0 or normalized_axis > rank: + return None + outer_size = _shape_element_count(input_shape[:normalized_axis]) + inner_size = _shape_element_count(input_shape[normalized_axis:]) + return output_name if output_shape == (outer_size, inner_size) else None + + if node.op_type == "Reshape": + target_shape = _constant_ints(index, node.input[1]) if len(node.input) == 2 else None + allowzero = _attribute(node, "allowzero", 0) + if target_shape is None or allowzero not in (0, 1): + return None + if allowzero == 1 and 0 in target_shape: + return None + else: + axes, conflict = _single_attribute_or_input_ints(index, node, "axes", 1) + if conflict or axes is None: + return None + return output_name + + +def _single_unobserved_consumer( + index: _GraphIndex, + tensor_name: str, +) -> onnx.NodeProto | None: + if tensor_name in index.graph_outputs: + return None + consumers = index.consumers.get(tensor_name, []) + return consumers[0] if len(consumers) == 1 else None + + +def _constant_input( + index: _GraphIndex, + node: onnx.NodeProto, + data_name: str | None = None, +) -> tuple[int, np.ndarray] | None: + if len(node.input) != 2 or any(not input_name for input_name in node.input): + return None + if data_name is not None and sum(name == data_name for name in node.input) != 1: + return None + constants = [ + (position, values) + for position, input_name in enumerate(node.input) + if (data_name is None or input_name != data_name) + and (values := _constant_array(index, input_name)) is not None + ] + return constants[0] if len(constants) == 1 else None + + +def _feeds_standard_mul_through_order_preserving_views( + index: _GraphIndex, + tensor_name: str, +) -> bool: + pending = [(tensor_name, 0)] + visited = {tensor_name} + while pending: + current_name, depth = pending.pop() + consumers = index.consumers.get(current_name, []) + if depth >= MAX_AFFINE_ROUTE_DEPTH and consumers: + return True + for consumer in consumers: + if not _is_standard_onnx_node(consumer): + continue + if consumer.op_type == "Mul": + return True + if consumer.op_type in ORDER_PRESERVING_VIEW_OPS: + view_output = _order_preserving_view_output(index, consumer, current_name) + if view_output is None or view_output in visited: + return True + visited.add(view_output) + pending.append((view_output, depth + 1)) + return False + + +def _post_exp_scale( + index: _GraphIndex, + exp: onnx.NodeProto, + visited: set[str] | None = None, +) -> tuple[onnx.NodeProto, list[onnx.NodeProto], np.ndarray, tuple[int, ...]] | None: + if not _is_standard_onnx_node(exp) or exp.op_type != "Exp" or len(exp.input) != 1: + return None + exp_output = _node_output(exp) + if exp_output is None: + return None + route = set() if visited is None else set(visited) + if exp_output in route or len(route) >= MAX_AFFINE_ROUTE_DEPTH: + return None + route.add(exp_output) + current_name = exp_output + current_node = exp + next_node = _single_unobserved_consumer(index, current_name) + while next_node is not None: + view_output = _order_preserving_view_output(index, next_node, current_name) + if view_output is None: + break + if view_output in route or len(route) >= MAX_AFFINE_ROUTE_DEPTH: + return None + route.add(view_output) + current_name = view_output + current_node = next_node + next_node = _single_unobserved_consumer(index, current_name) + + if next_node is None or not _is_standard_onnx_node(next_node) or next_node.op_type != "Mul": + return None + output_shape = _static_shape(index, current_name) + if output_shape is None: + return None + + scale_operand = _constant_input(index, next_node, current_name) + mul_output = _node_output(next_node) + if scale_operand is None or mul_output is None: + return None + if _feeds_standard_mul_through_order_preserving_views(index, mul_output): + return None + scale = scale_operand[1] + if ( + not np.issubdtype(scale.dtype, np.floating) + or not np.isfinite(scale).all() + or not np.all(scale > 0) + ): + return None + try: + np.broadcast_to(scale, output_shape) + except ValueError: + return None + return current_node, [next_node], scale, output_shape + + +def _exp_scale_candidate( + index: _GraphIndex, + add: onnx.NodeProto, +) -> _ExpScaleCandidate | None: + if not _is_standard_onnx_node(add) or add.op_type != "Add": + return None + add_output = _node_output(add) + bias_operand = _constant_input(index, add) + if add_output is None or bias_operand is None: + return None + add_shape = _static_shape(index, add_output) + if add_shape is None: + return None + + current_name = add_output + visited = {add_output} + next_node = _single_unobserved_consumer(index, current_name) + route_consumer = next_node + while next_node is not None: + view_output = _order_preserving_view_output(index, next_node, current_name) + if view_output is None: + break + if view_output in visited or len(visited) >= MAX_AFFINE_ROUTE_DEPTH: + return None + visited.add(view_output) + current_name = view_output + next_node = _single_unobserved_consumer(index, current_name) + + if ( + next_node is None + or not _is_standard_onnx_node(next_node) + or next_node.op_type != "Exp" + or list(next_node.input) != [current_name] + or route_consumer is None + ): + return None + post_exp = _post_exp_scale(index, next_node, visited) + if post_exp is None: + return None + output_node, muls, scale, output_shape = post_exp + if output_shape != add_shape: + return None + + bias = bias_operand[1] + if ( + not np.issubdtype(bias.dtype, np.floating) + or bias.dtype != scale.dtype + or not np.isfinite(bias).all() + or not np.isfinite(scale).all() + or not np.all(scale > 0) + ): + return None + try: + np.broadcast_to(bias, add_shape) + np.broadcast_to(scale, add_shape) + log_scale = np.asarray(np.log(scale), dtype=scale.dtype) + np.broadcast_to(log_scale, add_shape) + except ValueError: + return None + if not np.isfinite(log_scale).all(): + return None + combined_bias = _compact_combined_bias(bias, log_scale, add_shape) + return _ExpScaleCandidate( + add=add, + bias_input_index=bias_operand[0], + output_node=output_node, + muls=muls, + add_output=add_output, + route_consumer=route_consumer, + combined_bias=combined_bias, + log_scale=None if combined_bias is not None else log_scale, + ) + + +def _exp_scale_insert_candidate( + index: _GraphIndex, + exp: onnx.NodeProto, +) -> _ExpScaleInsertCandidate | None: + if not _is_standard_onnx_node(exp) or exp.op_type != "Exp" or len(exp.input) != 1: + return None + current_name = exp.input[0] + visited = {current_name} + producer = index.producers.get(current_name) + while producer is not None and producer.op_type in {"Reshape", "Squeeze", "Unsqueeze"}: + if ( + not producer.input + or _order_preserving_view_output(index, producer, producer.input[0]) != current_name + ): + return None + current_name = producer.input[0] + if current_name in visited or len(visited) >= MAX_AFFINE_ROUTE_DEPTH: + return None + visited.add(current_name) + producer = index.producers.get(current_name) + + post_exp = _post_exp_scale(index, exp) + if post_exp is None: + return None + output_node, muls, scale, output_shape = post_exp + input_shape = _static_shape(index, exp.input[0]) + if ( + input_shape is None + or any(dimension <= 0 for dimension in (*input_shape, *output_shape)) + or not _same_shape_element_count(input_shape, output_shape) + ): + return None + log_scale = _compact_log_scale_for_input(scale, input_shape, output_shape) + if log_scale is None: + return None + return _ExpScaleInsertCandidate( + exp=exp, + output_node=output_node, + muls=muls, + log_scale=log_scale, + ) + + +def _compact_log_scale_for_input( + scale: np.ndarray, + input_shape: tuple[int, ...], + output_shape: tuple[int, ...], +) -> np.ndarray | None: + """Return a log-scale initializer without expanding broadcast-only dimensions.""" + if scale.size == 1 or ( + input_shape == output_shape and _shape_broadcasts_to(scale.shape, input_shape) + ): + log_scale = np.asarray(np.log(scale), dtype=scale.dtype) + elif scale.size == _shape_element_count(input_shape): + try: + log_scale = np.asarray(np.log(scale).reshape(input_shape), dtype=scale.dtype) + except (MemoryError, ValueError): + return None + else: + return None + return log_scale if np.isfinite(log_scale).all() else None + + +def _compact_combined_bias( + bias: np.ndarray, + log_scale: np.ndarray, + target_shape: tuple[int, ...], +) -> np.ndarray | None: + try: + combined_shape = np.broadcast_shapes(bias.shape, log_scale.shape) + except ValueError: + return None + combined_count = _shape_element_count(combined_shape) + if ( + combined_count < 0 + or not _shape_broadcasts_to(combined_shape, target_shape) + or combined_count > max(int(bias.size), int(log_scale.size)) + ): + return None + combined_bias = np.asarray(bias + log_scale, dtype=bias.dtype) + return combined_bias if np.isfinite(combined_bias).all() else None + + +def _candidate_node_ids(nodes: Iterable[onnx.NodeProto]) -> set[int]: + return {id(node) for node in nodes} + + +def _exp_bias_candidate_node_ids(candidate: _ExpScaleCandidate) -> set[int]: + return _candidate_node_ids( + [ + candidate.add, + candidate.route_consumer, + candidate.output_node, + *candidate.muls, + ] + ) + + +def _exp_insert_candidate_node_ids(candidate: _ExpScaleInsertCandidate) -> set[int]: + return _candidate_node_ids([candidate.exp, candidate.output_node, *candidate.muls]) + + +def _select_exp_bias_scale_candidates( + model: onnx.ModelProto, + index: _GraphIndex, +) -> list[_ExpScaleCandidate]: + selected: list[_ExpScaleCandidate] = [] + reserved_nodes: set[int] = set() + for add in model.graph.node: + candidate = _exp_scale_candidate(index, add) + if candidate is None: + continue + if candidate.combined_bias is None and ( + candidate.log_scale is None + or not candidate.route_consumer.input + or candidate.route_consumer.input[0] != candidate.add_output + ): + continue + candidate_nodes = _exp_bias_candidate_node_ids(candidate) + if candidate_nodes & reserved_nodes: + continue + reserved_nodes.update(candidate_nodes) + selected.append(candidate) + return selected + + +def _select_exp_scale_insert_candidates( + model: onnx.ModelProto, + index: _GraphIndex, + reserved_nodes: set[int] | None = None, +) -> list[_ExpScaleInsertCandidate]: + selected: list[_ExpScaleInsertCandidate] = [] + selected_nodes: set[int] = set() if reserved_nodes is None else set(reserved_nodes) + for exp in model.graph.node: + candidate = _exp_scale_insert_candidate(index, exp) + if candidate is None: + continue + candidate_nodes = _exp_insert_candidate_node_ids(candidate) + if candidate_nodes & selected_nodes: + continue + selected_nodes.update(candidate_nodes) + selected.append(candidate) + return selected + + +def _fold_exp_positive_scales( + model: onnx.ModelProto, + allocator: _NameAllocator, +) -> None: + """Fold eligible positive post-Exp constants into the Exp input.""" + index = _GraphIndex.build(model) + bias_candidates = _select_exp_bias_scale_candidates(model, index) + reserved_nodes = set().union( + *(_exp_bias_candidate_node_ids(candidate) for candidate in bias_candidates), + ) + insert_candidates = _select_exp_scale_insert_candidates(model, index, reserved_nodes) + if not bias_candidates and not insert_candidates: + return + + removed: set[int] = set() + insert_before: dict[int, onnx.NodeProto] = {} + insert_after: dict[int, onnx.NodeProto] = {} + for bias_candidate in bias_candidates: + if bias_candidate.combined_bias is not None: + combined_name = _new_initializer( + model, + allocator, + bias_candidate.combined_bias, + "algebraic_exp_log_bias", + ) + bias_candidate.add.input[bias_candidate.bias_input_index] = combined_name + elif bias_candidate.log_scale is not None: + log_scale_name = _new_initializer( + model, + allocator, + bias_candidate.log_scale, + "algebraic_exp_log_scale", + ) + adjusted_name = allocator.new("algebraic_exp_adjusted") + log_add = onnx.helper.make_node( + "Add", + [bias_candidate.add_output, log_scale_name], + [adjusted_name], + name=allocator.new("algebraic_exp_log_add"), + ) + bias_candidate.route_consumer.input[0] = adjusted_name + insert_after[id(bias_candidate.add)] = log_add + else: + continue + bias_candidate.output_node.output[0] = bias_candidate.muls[-1].output[0] + removed.update(id(mul) for mul in bias_candidate.muls) + + for insert_candidate in insert_candidates: + log_scale_name = _new_initializer( + model, + allocator, + insert_candidate.log_scale, + "algebraic_exp_log_scale", + ) + adjusted_name = allocator.new("algebraic_exp_adjusted") + add = onnx.helper.make_node( + "Add", + [insert_candidate.exp.input[0], log_scale_name], + [adjusted_name], + name=allocator.new("algebraic_exp_log_add"), + ) + insert_candidate.exp.input[0] = adjusted_name + insert_candidate.output_node.output[0] = insert_candidate.muls[-1].output[0] + insert_before[id(insert_candidate.exp)] = add + removed.update(id(mul) for mul in insert_candidate.muls) + + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + before = insert_before.get(id(node)) + if before is not None: + rewritten.append(before) + if id(node) not in removed: + rewritten.append(node) + after = insert_after.get(id(node)) + if after is not None: + rewritten.append(after) + del model.graph.node[:] + model.graph.node.extend(rewritten) + + def _collect_affine_chain( index: _GraphIndex, first: onnx.NodeProto, @@ -498,8 +1257,12 @@ def _collect_affine_chain( start: int, end: int, calculation_dtype: np.dtype[Any], -) -> _AffineCandidate | None: + visited_routes: set[tuple[int, int, str]], + depth: int, +) -> tuple[_AffineCandidate | None, bool]: """Collect a safe consecutive Mul/Add chain from one routed branch.""" + if not _is_standard_onnx_node(first) or first.op_type not in {"Mul", "Add"}: + return None, True current = first current_input = source_name scale = np.ones(end - start, dtype=calculation_dtype) @@ -508,26 +1271,38 @@ def _collect_affine_chain( while current.op_type in {"Mul", "Add"}: if len(current.input) != 2 or current_input not in current.input: - return None + return None, True current_output = _node_output(current) if current_output is None: - return None + return None, False + if not _visit_affine_route(current, 0, visited_routes, depth + 1): + return None, False + depth += 1 values = _affine_operand(index, current, current_input, output_shape, end - start) if values is None: - return None + return None, True values = values.astype(calculation_dtype, copy=False) if current.op_type == "Mul": - scale *= values - offset *= values + with np.errstate(over="ignore", invalid="ignore"): + next_scale = scale * values + next_offset = offset * values + if not np.isfinite(next_scale).all() or not np.isfinite(next_offset).all(): + return None, True + scale = next_scale + offset = next_offset else: - offset += values + with np.errstate(over="ignore", invalid="ignore"): + next_offset = offset + values + if not np.isfinite(next_offset).all(): + return None, True + offset = next_offset matched.append(current) consumers = index.consumers.get(current_output, []) if current_output in index.graph_outputs or len(consumers) != 1: break next_node = consumers[0] - if next_node.op_type not in {"Mul", "Add"}: + if not _is_standard_onnx_node(next_node) or next_node.op_type not in {"Mul", "Add"}: break current_input = current_output current = next_node @@ -536,19 +1311,45 @@ def _collect_affine_chain( if final_output is None or ( final_output not in index.graph_outputs and len(index.consumers.get(final_output, [])) == 0 ): - return None - return _AffineCandidate( - source_node=first, - source_output_index=0, - final_output=final_output, - nodes=matched, - start=start, - end=end, - scale=scale, - offset=offset, + return None, True + return ( + _AffineCandidate( + source_node=first, + source_output_index=0, + final_output=final_output, + nodes=matched, + start=start, + end=end, + scale=scale, + offset=offset, + ), + True, ) +def _visit_affine_route( + source_node: onnx.NodeProto, + source_output_index: int, + visited_routes: set[tuple[int, int, str]], + depth: int, +) -> bool: + """Record one unique, bounded source-slot and tensor route.""" + if ( + depth > MAX_AFFINE_ROUTE_DEPTH + or source_output_index < 0 + or source_output_index >= len(source_node.output) + ): + return False + source_name = source_node.output[source_output_index] + if not source_name: + return False + source_slot = (id(source_node), source_output_index) + if any(route[:2] == source_slot or route[2] == source_name for route in visited_routes): + return False + visited_routes.add((source_slot[0], source_slot[1], source_name)) + return True + + def _collect_routed_affine_candidates( index: _GraphIndex, source_node: onnx.NodeProto, @@ -556,12 +1357,19 @@ def _collect_routed_affine_candidates( start: int, end: int, calculation_dtype: np.dtype[Any], -) -> list[_AffineCandidate]: + visited_routes: set[tuple[int, int, str]], + depth: int, +) -> list[_AffineCandidate] | None: """Collect affine leaves below safe views and disjoint channel slices.""" - if source_output_index >= len(source_node.output): - return [] + if not _visit_affine_route( + source_node, + source_output_index, + visited_routes, + depth, + ): + return None source_name = source_node.output[source_output_index] - if not source_name or source_name in index.graph_outputs: + if source_name in index.graph_outputs: return [] current_node = source_node @@ -582,6 +1390,9 @@ def _collect_routed_affine_candidates( ) if view_output is None or current_name in index.graph_outputs: break + if not _visit_affine_route(view, 0, visited_routes, depth + 1): + return None + depth += 1 current_node = view current_output_index = 0 current_name = view_output @@ -592,8 +1403,12 @@ def _collect_routed_affine_candidates( if current_name in index.graph_outputs: return [] - if len(consumers) == 1 and consumers[0].op_type in {"Mul", "Add"}: - candidate = _collect_affine_chain( + if ( + len(consumers) == 1 + and _is_standard_onnx_node(consumers[0]) + and consumers[0].op_type in {"Mul", "Add"} + ): + candidate, route_is_valid = _collect_affine_chain( index, consumers[0], current_name, @@ -601,21 +1416,57 @@ def _collect_routed_affine_candidates( start, end, calculation_dtype, + visited_routes, + depth, ) + if not route_is_valid: + return None if candidate is None: return [] candidate.source_node = current_node candidate.source_output_index = current_output_index return [candidate] - if not consumers or any(node.op_type != "Slice" for node in consumers): + if len(consumers) == 1 and consumers[0].op_type == "Split": + nested_split = consumers[0] + nested_info = _split_boundaries(index, nested_split, current_name) + if nested_info is None or nested_info[0] != 1: + return None + boundaries = nested_info[1] + if len(boundaries) != len(nested_split.output): + return [] + nested_affine_candidates: list[_AffineCandidate] = [] + for output_index, (local_start, local_end) in enumerate(boundaries): + nested_candidates = _collect_routed_affine_candidates( + index, + nested_split, + output_index, + start + local_start, + start + local_end, + calculation_dtype, + visited_routes, + depth + 1, + ) + if nested_candidates is None: + return None + nested_affine_candidates.extend(nested_candidates) + return nested_affine_candidates + + if not consumers or any( + not _is_standard_onnx_node(node) or node.op_type != "Slice" for node in consumers + ): return [] routed_slices: list[tuple[onnx.NodeProto, int, int]] = [] + routed_outputs: list[str] = [] for routed_slice in consumers: boundary = _slice_channel_boundary(index, routed_slice, current_name, 1) - if boundary is None: - return [] + output_name = _node_output(routed_slice) + if boundary is None or output_name is None or output_name == current_name: + return None routed_slices.append((routed_slice, *boundary)) + routed_outputs.append(output_name) + if len(set(routed_outputs)) != len(routed_outputs): + return None if any( left_start < right_end and right_start < left_end for position, (_, left_start, left_end) in enumerate(routed_slices) @@ -623,80 +1474,84 @@ def _collect_routed_affine_candidates( ): return [] - candidates: list[_AffineCandidate] = [] + routed_affine_candidates: list[_AffineCandidate] = [] for routed_slice, local_start, local_end in routed_slices: - candidates.extend( - _collect_routed_affine_candidates( - index, - routed_slice, - 0, - start + local_start, - start + local_end, - calculation_dtype, - ) + routed_candidates = _collect_routed_affine_candidates( + index, + routed_slice, + 0, + start + local_start, + start + local_end, + calculation_dtype, + visited_routes, + depth + 1, ) - return candidates + if routed_candidates is None: + return None + routed_affine_candidates.extend(routed_candidates) + return routed_affine_candidates def _copy_conv_parameters( model: onnx.ModelProto, + index: _GraphIndex, allocator: _NameAllocator, conv: onnx.NodeProto, scale: np.ndarray, offset: np.ndarray, ) -> bool: - if len(conv.input) < 2: + if not np.isfinite(scale).all() or not np.isfinite(offset).all(): return False - weight = next( - ( - initializer - for initializer in model.graph.initializer - if initializer.name == conv.input[1] - ), - None, - ) - if weight is None: + if len(conv.input) < 2 or conv.input[1] in index.graph_inputs: + return False + weights = _initializer_array(index, conv.input[1]) + if weights is None: return False - weights = np.asarray(onnx.numpy_helper.to_array(weight)) if weights.ndim < 1 or weights.shape[0] != len(scale): return False if not np.issubdtype(weights.dtype, np.floating): return False + if not np.isfinite(weights).all(): + return False if len(conv.input) > 2 and conv.input[2]: - bias = next( - ( - initializer - for initializer in model.graph.initializer - if initializer.name == conv.input[2] - ), - None, - ) - if bias is None: + if conv.input[2] in index.graph_inputs: + return False + bias_values = _initializer_array(index, conv.input[2]) + if bias_values is None: return False - bias_values = np.asarray(onnx.numpy_helper.to_array(bias)) if bias_values.ndim != 1 or len(bias_values) != len(scale): return False if not np.issubdtype(bias_values.dtype, np.floating): return False + if not np.isfinite(bias_values).all(): + return False else: bias_values = np.zeros(len(scale), dtype=weights.dtype) - new_weights = weights * scale.reshape((len(scale),) + (1,) * (weights.ndim - 1)) + with np.errstate(over="ignore", invalid="ignore"): + new_weights = weights * scale.reshape((len(scale),) + (1,) * (weights.ndim - 1)) + folded_weights = np.asarray(new_weights, dtype=weights.dtype) + if not np.isfinite(folded_weights).all(): + return False + with np.errstate(over="ignore", invalid="ignore"): + new_bias = bias_values * scale + offset + folded_bias = np.asarray(new_bias, dtype=bias_values.dtype) + if not np.isfinite(folded_bias).all(): + return False weight_name = _new_initializer( model, allocator, - np.asarray(new_weights, dtype=weights.dtype), + folded_weights, "algebraic_conv_weight", ) - conv.input[1] = weight_name - new_bias = bias_values * scale + offset bias_name = _new_initializer( model, allocator, - np.asarray(new_bias, dtype=bias_values.dtype), + folded_bias, "algebraic_conv_bias", ) + conv.input[1] = weight_name if len(conv.input) > 2: conv.input[2] = bias_name else: @@ -712,93 +1567,54 @@ def _fold_channel_affine( index = _GraphIndex.build(model) for original_conv in list(model.graph.node): if ( - original_conv.op_type != "Conv" + not _is_standard_onnx_node(original_conv) + or original_conv.op_type != "Conv" or len(original_conv.output) != 1 or not original_conv.output[0] ): continue conv_output = original_conv.output[0] conv = index.producers.get(conv_output) - if conv is None or conv.op_type != "Conv": + if conv is None or not _is_standard_onnx_node(conv) or conv.op_type != "Conv": continue conv_shape = _static_shape(index, conv_output) if conv_shape is None or len(conv_shape) < 2: continue channels = conv_shape[1] - weight_initializer = index.initializers.get(conv.input[1]) if len(conv.input) > 1 else None - if weight_initializer is None: + weight_values = _initializer_array(index, conv.input[1]) if len(conv.input) > 1 else None + if weight_values is None: continue - weight_dtype = onnx.numpy_helper.to_array(weight_initializer).dtype + weight_dtype = weight_values.dtype if channels <= 0: continue calculation_dtype = np.result_type(weight_dtype, np.float32) - route_name = conv_output - route_shape = conv_shape - route_source_node = conv - route_source_output_index = 0 - direct_consumers = index.consumers.get(route_name, []) - while len(direct_consumers) == 1: - view = direct_consumers[0] - view_output = _channel_preserving_view_output( - index, - view, - route_name, - channels, - ) - if view_output is None or route_name in index.graph_outputs: - break - route_name = view_output - next_route_shape = _static_shape(index, route_name) - if next_route_shape is None: - break - route_shape = next_route_shape - route_source_node = view - route_source_output_index = 0 - direct_consumers = index.consumers.get(route_name, []) - - candidates: list[_AffineCandidate] = [] - if route_name not in index.graph_outputs and len(direct_consumers) == 1: - direct = _collect_affine_chain( - index, - direct_consumers[0], - route_name, - route_shape, - 0, - channels, - calculation_dtype, - ) - if direct is not None: - direct.source_node = route_source_node - direct.source_output_index = route_source_output_index - candidates.append(direct) - - if not candidates and route_name not in index.graph_outputs and len(direct_consumers) == 1: - router = direct_consumers[0] - boundaries: list[tuple[int, int]] | None = None - if router.op_type == "Split": - split_info = _split_boundaries(index, router, route_name) - if split_info is not None and split_info[0] == 1: - boundaries = split_info[1] - elif router.op_type == "Slice": - boundary = _slice_channel_boundary(index, router, route_name, 1) - if boundary is not None: - boundaries = [boundary] - - if boundaries is not None and len(boundaries) == len(router.output): - for output_index, (start, end) in enumerate(boundaries): - candidates.extend( - _collect_routed_affine_candidates( - index, - router, - output_index, - start, - end, - calculation_dtype, - ) - ) - - if not candidates: + collected_candidates = _collect_routed_affine_candidates( + index, + conv, + 0, + 0, + channels, + calculation_dtype, + set(), + 0, + ) + if not collected_candidates: + continue + candidates = collected_candidates + candidate_source_slots = [ + (id(candidate.source_node), candidate.source_output_index) for candidate in candidates + ] + candidate_source_tensors = [ + candidate.source_node.output[candidate.source_output_index] for candidate in candidates + ] + candidate_node_ids = [id(node) for candidate in candidates for node in candidate.nodes] + if ( + len({id(candidate) for candidate in candidates}) != len(candidates) + or len(set(candidate_source_slots)) != len(candidate_source_slots) + or len(set(candidate_source_tensors)) != len(candidate_source_tensors) + or len(set(candidate_node_ids)) != len(candidate_node_ids) + ): continue if any( left.start < right.end and right.start < left.end @@ -820,7 +1636,7 @@ def _fold_channel_affine( for candidate in candidates: scale[candidate.start : candidate.end] = candidate.scale offset[candidate.start : candidate.end] = candidate.offset - if not _copy_conv_parameters(model, allocator, conv, scale, offset): + if not _copy_conv_parameters(model, index, allocator, conv, scale, offset): continue removed = {id(node) for candidate in candidates for node in candidate.nodes} @@ -837,16 +1653,18 @@ def _rewrite_static_splits( ) -> None: """Replace statically bounded Split nodes with input-form Slice nodes.""" index = _GraphIndex.build(model) - opset = next( - (int(opset.version) for opset in model.opset_import if opset.domain in ("", "ai.onnx")), - 0, - ) - if opset and opset < 10: + opset = _strict_default_opset_version(model) + if opset is None or opset < 10: return replacements: dict[int, list[onnx.NodeProto]] = {} for split in list(model.graph.node): - if split.op_type != "Split" or len(split.input) < 1 or not split.input[0]: + if ( + not _is_standard_onnx_node(split) + or split.op_type != "Split" + or len(split.input) < 1 + or not split.input[0] + ): continue if any(not output for output in split.output): continue @@ -898,13 +1716,20 @@ def build_config(cls, **kwargs: Any) -> AlgebraicRewritePipeConfig: """Build the enabled algebraic rewrite configuration.""" return AlgebraicRewritePipeConfig( static_split_to_slice=kwargs.get("static_split_to_slice", False), + sibling_slice_to_split=kwargs.get("gather_slice_to_split_fusion", False), conv_channel_affine_folding=kwargs.get("conv_channel_affine_folding", False), + exp_positive_scale_folding=kwargs.get("exp_positive_scale_folding", False), ) @classmethod def should_process(cls, config: AlgebraicRewritePipeConfig) -> bool: """Return whether any algebraic rewrite is enabled.""" - return config.static_split_to_slice or config.conv_channel_affine_folding + return ( + config.static_split_to_slice + or config.sibling_slice_to_split + or config.conv_channel_affine_folding + or config.exp_positive_scale_folding + ) def process( self, @@ -917,10 +1742,24 @@ def process( result = onnx.ModelProto() result.CopyFrom(model) + if result.ir_version < 4: + return result + index = _GraphIndex.build(result) + if index.definition_collisions or index.has_cycle: + return result allocator = _NameAllocator(result) introduced_nodes: set[str] = set() - if config.conv_channel_affine_folding: + standard_opset = _strict_default_opset_version(result) + if ( + config.conv_channel_affine_folding + and standard_opset is not None + and standard_opset >= 7 + ): _fold_channel_affine(result, allocator) + if config.exp_positive_scale_folding and standard_opset is not None and standard_opset >= 7: + _fold_exp_positive_scales(result, allocator) + if config.sibling_slice_to_split: + _fold_sibling_slices_to_split(result, allocator) if config.static_split_to_slice: _rewrite_static_splits(result, allocator, introduced_nodes) _prune_generated_slices(result, introduced_nodes) diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index e9a977838..667e923bf 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -21,10 +21,12 @@ AlgebraicRewritePipe, AlgebraicRewritePipeConfig, ) +from winml.modelkit.optim.pipes import algebraic as algebraic_pipe if TYPE_CHECKING: from collections.abc import Sequence + from pathlib import Path def _tensor(name: str, values: np.ndarray) -> onnx.TensorProto: @@ -69,13 +71,35 @@ def _assert_valid_with_inferred_shapes(model: onnx.ModelProto) -> None: assert len(inferred.graph.output) == len(model.graph.output) +def _node_signatures( + model: onnx.ModelProto, +) -> list[tuple[str, str, str, tuple[str, ...], tuple[str, ...]]]: + return [ + (node.name, node.domain, node.op_type, tuple(node.input), tuple(node.output)) + for node in model.graph.node + ] + + +def _assert_byte_identical(original: onnx.ModelProto, transformed: onnx.ModelProto) -> None: + assert transformed.SerializeToString() == original.SerializeToString(), ( + f"graph mutated:\nbefore={_node_signatures(original)}\n" + f"after={_node_signatures(transformed)}" + ) + + class TestAlgebraicRegistration: """Verify capability registration, flags, and pipe ordering.""" def test_capabilities_are_opt_in_and_independent(self) -> None: capabilities = get_all_capabilities() - names = {"static-split-to-slice", "conv-channel-affine-folding"} + names = { + "static-split-to-slice", + "gather-slice-to-split-fusion", + "conv-channel-affine-folding", + "exp-positive-scale-folding", + } assert names <= capabilities.keys() + assert names <= AlgebraicRewritePipe.capabilities.keys() assert all(capabilities[name].default is False for name in names) assert all( capabilities[name].cli_flags() == (f"--enable-{name}", f"--disable-{name}") @@ -84,16 +108,23 @@ def test_capabilities_are_opt_in_and_independent(self) -> None: config = AlgebraicRewritePipe.build_config( static_split_to_slice=True, + gather_slice_to_split_fusion=True, conv_channel_affine_folding=False, + exp_positive_scale_folding=True, ) assert config.static_split_to_slice is True + assert config.sibling_slice_to_split is True assert config.conv_channel_affine_folding is False + assert config.exp_positive_scale_folding is True + assert AlgebraicRewritePipe.should_process(config) def test_cli_lists_algebraic_flag(self) -> None: result = CliRunner().invoke(optimize, ["--list-capabilities"]) assert result.exit_code == 0 assert "--enable-static-split-to-slice" in result.output + assert "--enable-gather-slice-to-split-fusion" in result.output assert "--enable-conv-channel-affine-folding" in result.output + assert "--enable-exp-positive-scale-folding" in result.output def test_pipe_is_after_ort_graph_and_before_cleanup(self) -> None: names = [pipe.name for pipe in PIPES] @@ -102,6 +133,140 @@ def test_pipe_is_after_ort_graph_and_before_cleanup(self) -> None: assert PIPES[names.index("algebraic_rewrite")] is AlgebraicRewritePipe assert not AlgebraicRewritePipe.should_process(AlgebraicRewritePipeConfig()) + def test_cli_combines_split_affine_and_exp_folding(self, tmp_path: Path) -> None: + rng = np.random.default_rng(40) + model = _model( + [ + onnx.helper.make_node("Conv", ["conv_input", "weight"], ["conv_out"]), + onnx.helper.make_node( + "Slice", + ["conv_out", "first_starts", "first_ends", "channel_axis"], + ["first"], + ), + onnx.helper.make_node( + "Slice", + ["conv_out", "second_starts", "second_ends", "channel_axis"], + ["second"], + ), + onnx.helper.make_node( + "Mul", + ["first", "first_scale"], + ["first_out"], + name="target_conv_mul", + ), + onnx.helper.make_node( + "Add", + ["second", "second_offset"], + ["second_out"], + name="target_conv_add", + ), + onnx.helper.make_node( + "Add", + ["exp_input", "exp_bias"], + ["biased"], + name="retained_exp_add", + ), + onnx.helper.make_node("Exp", ["biased"], ["exponential"]), + onnx.helper.make_node( + "Mul", + ["exponential", "exp_scale"], + ["exp_out"], + name="target_exp_mul", + ), + ], + [_info("conv_input", [1, 1, 2, 2]), _info("exp_input", [1, 2])], + [ + _info("first_out", [1, 2, 2, 2]), + _info("second_out", [1, 2, 2, 2]), + _info("exp_out", [1, 2]), + ], + [ + _tensor("weight", rng.normal(size=(4, 1, 1, 1)).astype(np.float32)), + _tensor("first_starts", np.asarray([0], dtype=np.int64)), + _tensor("first_ends", np.asarray([2], dtype=np.int64)), + _tensor("second_starts", np.asarray([2], dtype=np.int64)), + _tensor("second_ends", np.asarray([4], dtype=np.int64)), + _tensor("channel_axis", np.asarray([1], dtype=np.int64)), + _tensor( + "first_scale", np.asarray([1.25, 0.75], dtype=np.float32).reshape(1, 2, 1, 1) + ), + _tensor( + "second_offset", np.asarray([0.5, -0.5], dtype=np.float32).reshape(1, 2, 1, 1) + ), + _tensor("exp_bias", np.asarray(-1.0, dtype=np.float32)), + _tensor("exp_scale", np.asarray([1.5, 2.0], dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 4, 2, 2]), + _info("first", [1, 2, 2, 2]), + _info("second", [1, 2, 2, 2]), + _info("biased", [1, 2]), + _info("exponential", [1, 2]), + ], + ) + input_path = tmp_path / "input.onnx" + output_path = tmp_path / "output.onnx" + onnx.save_model(model, input_path) + + result = CliRunner().invoke( + optimize, + [ + "-m", + str(input_path), + "-o", + str(output_path), + "--enable-gather-slice-to-split-fusion", + "--enable-conv-channel-affine-folding", + "--enable-exp-positive-scale-folding", + "--no-color", + ], + ) + + assert result.exit_code == 0, result.output + transformed = onnx.load_model(output_path) + names = {node.name for node in transformed.graph.node} + assert not {"target_conv_mul", "target_conv_add", "target_exp_mul"} & names + assert any(node.op_type == "Split" for node in transformed.graph.node) + assert [output.SerializeToString() for output in transformed.graph.output] == [ + output.SerializeToString() for output in model.graph.output + ] + _assert_valid_with_inferred_shapes(transformed) + values = { + "conv_input": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), + "exp_input": rng.normal(size=(1, 2)).astype(np.float32), + } + for original, rewritten in zip( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + + +class TestNameAllocator: + """Test generated-name allocation behavior used by batched rewrites.""" + + def test_repeated_prefix_allocation_does_not_restart_suffix_probe(self) -> None: + class CountingNames(set[str]): + def __init__(self) -> None: + super().__init__() + self.contains_count = 0 + + def __contains__(self, value: object) -> bool: + self.contains_count += 1 + return super().__contains__(value) + + allocator = algebraic_pipe._NameAllocator(_model([], [], [], [])) + used = CountingNames() + allocator._used = used + + names = [allocator.new("algebraic_exp_log_scale") for _ in range(32)] + + assert names == ["algebraic_exp_log_scale"] + [ + f"algebraic_exp_log_scale_{suffix}" for suffix in range(1, 32) + ] + assert used.contains_count <= 40 + class TestStaticSplitToSlice: """Test static Split replacement using generated data.""" @@ -194,6 +359,108 @@ def test_dynamic_equal_split_and_malformed_split_are_unchanged(self) -> None: ) assert [node.op_type for node in transformed.graph.node] == ["Split", "Split"] + def test_overridable_split_sizes_are_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node( + "Split", + ["x", "split_sizes"], + ["left", "right"], + axis=1, + ) + ], + [ + _info("x", [1, 4, 2]), + onnx.helper.make_tensor_value_info( + "split_sizes", + onnx.TensorProto.INT64, + [2], + ), + ], + [_info("left", [1, 2, 2]), _info("right", [1, 2, 2])], + [_tensor("split_sizes", np.asarray([2, 2], dtype=np.int64))], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(static_split_to_slice=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize( + ("domain", "should_rewrite"), + [("", True), ("ai.onnx", False), ("com.example", False)], + ) + def test_only_standard_domain_split_is_rewritten( + self, + domain: str, + should_rewrite: bool, + ) -> None: + model = _model( + [ + onnx.helper.make_node( + "Split", + ["x", "split_sizes"], + ["left", "right"], + axis=1, + domain=domain, + ) + ], + [_info("x", [1, 4, 2])], + [_info("left", [1, 2, 2]), _info("right", [1, 2, 2])], + [_tensor("split_sizes", np.asarray([2, 2], dtype=np.int64))], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(static_split_to_slice=True), + ) + + if should_rewrite: + assert [node.op_type for node in transformed.graph.node] == ["Slice", "Slice"] + else: + assert transformed.SerializeToString() == original + + def test_legacy_ir_generated_initializer_rewrite_is_unchanged(self) -> None: + model = _model( + [onnx.helper.make_node("Split", ["x", "split_sizes"], ["left", "right"], axis=1)], + [_info("x", [1, 4, 2])], + [_info("left", [1, 2, 2]), _info("right", [1, 2, 2])], + [_tensor("split_sizes", np.asarray([2, 2], dtype=np.int64))], + ) + model.ir_version = 3 + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(static_split_to_slice=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize( + "outputs", + [("", "right"), ("part", "part"), ("x", "right")], + ) + def test_malformed_split_outputs_are_unchanged(self, outputs: tuple[str, str]) -> None: + model = _model( + [onnx.helper.make_node("Split", ["x"], list(outputs), axis=1)], + [_info("x", [1, 4, 2])], + [_info("y", [1, 2, 2])], + [], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(static_split_to_slice=True), + ) + + assert transformed.SerializeToString() == original + def test_dead_generated_slice_and_constants_are_pruned(self) -> None: x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 2]) model = _model( @@ -210,6 +477,189 @@ def test_dead_generated_slice_and_constants_are_pruned(self) -> None: assert transformed.graph.node[0].output[0] == "left" assert len(transformed.graph.initializer) == 4 + def test_sibling_static_slices_fold_to_split(self) -> None: + values = {"x": np.arange(12, dtype=np.float32).reshape(1, 6, 2)} + model = _model( + [ + onnx.helper.make_node( + "Slice", + ["x", "left_starts", "left_ends", "axis", "steps"], + ["left"], + name="left_slice", + ), + onnx.helper.make_node( + "Slice", + ["x", "right_starts", "right_ends", "axis", "steps"], + ["right"], + name="right_slice", + ), + onnx.helper.make_node("Relu", ["left"], ["left_out"]), + onnx.helper.make_node("Relu", ["right"], ["right_out"]), + ], + [_info("x", [1, 6, 2])], + [_info("left_out", [1, 2, 2]), _info("right_out", [1, 4, 2])], + [ + _tensor("left_starts", np.asarray([0], dtype=np.int64)), + _tensor("left_ends", np.asarray([2], dtype=np.int64)), + _tensor("right_starts", np.asarray([2], dtype=np.int64)), + _tensor("right_ends", np.asarray([6], dtype=np.int64)), + _tensor("axis", np.asarray([1], dtype=np.int64)), + _tensor("steps", np.asarray([1], dtype=np.int64)), + ], + value_info=[_info("left", [1, 2, 2]), _info("right", [1, 4, 2])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipe.build_config(gather_slice_to_split_fusion=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Split", "Relu", "Relu"] + split = transformed.graph.node[0] + assert list(split.input[:1]) == ["x"] + assert list(split.output) == ["left", "right"] + split_sizes_name = split.input[1] + split_sizes = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == split_sizes_name) + ) + np.testing.assert_array_equal(split_sizes, np.asarray([2, 4], dtype=np.int64)) + _assert_valid_with_inferred_shapes(transformed) + for original, rewritten in zip(_run(model, values), _run(transformed, values), strict=True): + np.testing.assert_array_equal(original, rewritten) + + def test_sibling_static_slices_are_unchanged_before_split_input_opset(self) -> None: + model = _model( + [ + onnx.helper.make_node( + "Slice", + ["x", "left_starts", "left_ends", "axis", "steps"], + ["left"], + ), + onnx.helper.make_node( + "Slice", + ["x", "right_starts", "right_ends", "axis", "steps"], + ["right"], + ), + ], + [_info("x", [1, 6, 2])], + [_info("left", [1, 2, 2]), _info("right", [1, 4, 2])], + [ + _tensor("left_starts", np.asarray([0], dtype=np.int64)), + _tensor("left_ends", np.asarray([2], dtype=np.int64)), + _tensor("right_starts", np.asarray([2], dtype=np.int64)), + _tensor("right_ends", np.asarray([6], dtype=np.int64)), + _tensor("axis", np.asarray([1], dtype=np.int64)), + _tensor("steps", np.asarray([1], dtype=np.int64)), + ], + ) + model.opset_import[0].version = 12 + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipe.build_config(gather_slice_to_split_fusion=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize( + "opset_imports", + [ + [onnx.helper.make_opsetid("ai.onnx", 17), onnx.helper.make_opsetid("", 12)], + [onnx.helper.make_opsetid("", 17), onnx.helper.make_opsetid("", 12)], + ], + ) + def test_sibling_static_slices_are_unchanged_for_ambiguous_standard_opset( + self, + opset_imports: list[onnx.OperatorSetIdProto], + ) -> None: + model = _model( + [ + onnx.helper.make_node( + "Slice", + ["x", "left_starts", "left_ends", "axis", "steps"], + ["left"], + ), + onnx.helper.make_node( + "Slice", + ["x", "right_starts", "right_ends", "axis", "steps"], + ["right"], + ), + ], + [_info("x", [1, 6, 2])], + [_info("left", [1, 2, 2]), _info("right", [1, 4, 2])], + [ + _tensor("left_starts", np.asarray([0], dtype=np.int64)), + _tensor("left_ends", np.asarray([2], dtype=np.int64)), + _tensor("right_starts", np.asarray([2], dtype=np.int64)), + _tensor("right_ends", np.asarray([6], dtype=np.int64)), + _tensor("axis", np.asarray([1], dtype=np.int64)), + _tensor("steps", np.asarray([1], dtype=np.int64)), + ], + ) + del model.opset_import[:] + model.opset_import.extend(opset_imports) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipe.build_config(gather_slice_to_split_fusion=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize( + ("left_starts", "left_ends", "right_starts", "right_ends", "axes", "steps"), + [ + ([0], [2], [3], [6], [1], [1]), + ([0], [3], [2], [6], [1], [1]), + ([1], [2], [2], [6], [1], [1]), + ([0], [2], [2], [6], [1], [2]), + ([0, 0], [1, 2], [1, 0], [6, 2], [0, 1], [1, 1]), + ], + ) + def test_ineligible_sibling_static_slices_are_unchanged( + self, + left_starts: list[int], + left_ends: list[int], + right_starts: list[int], + right_ends: list[int], + axes: list[int], + steps: list[int], + ) -> None: + model = _model( + [ + onnx.helper.make_node( + "Slice", + ["x", "left_starts", "left_ends", "axes", "steps"], + ["left"], + ), + onnx.helper.make_node( + "Slice", + ["x", "right_starts", "right_ends", "axes", "steps"], + ["right"], + ), + ], + [_info("x", [1, 6, 2])], + [_info("left", [1, 2, 2]), _info("right", [1, 4, 2])], + [ + _tensor("left_starts", np.asarray(left_starts, dtype=np.int64)), + _tensor("left_ends", np.asarray(left_ends, dtype=np.int64)), + _tensor("right_starts", np.asarray(right_starts, dtype=np.int64)), + _tensor("right_ends", np.asarray(right_ends, dtype=np.int64)), + _tensor("axes", np.asarray(axes, dtype=np.int64)), + _tensor("steps", np.asarray(steps, dtype=np.int64)), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipe.build_config(gather_slice_to_split_fusion=True), + ) + + assert transformed.SerializeToString() == original + def test_nested_subgraph_captures_keep_generated_slices_live(self) -> None: x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 2]) then_branch = onnx.helper.make_graph( @@ -325,6 +775,92 @@ def test_direct_affine_folding_is_exact_and_adds_optional_bias( atol=2e-5, ) + def test_legacy_opset_conv_affine_broadcast_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), + onnx.helper.make_node( + "Mul", + ["conv_out", "scale"], + ["y"], + broadcast=1, + axis=0, + ), + ], + [_info("x", [1, 2, 2, 2])], + [_info("y", [1, 3, 2, 2])], + [ + _tensor("weights", np.ones((3, 2, 1, 1), dtype=np.float32)), + _tensor("scale", np.ones((3, 1, 1), dtype=np.float32)), + ], + value_info=[_info("conv_out", [1, 3, 2, 2])], + ) + model.opset_import[0].version = 6 + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_nonfinite_synthesized_conv_parameters_are_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["branch"], axis=1), + onnx.helper.make_node("Split", ["branch"], ["leaf"], axis=1), + onnx.helper.make_node("Mul", ["leaf", "scale_a"], ["scaled_a"]), + onnx.helper.make_node("Mul", ["scaled_a", "scale_b"], ["y"]), + ], + [_info("x", [1, 1, 1, 1])], + [_info("y", [1, 1, 1, 1])], + [ + _tensor("weight", np.asarray([[[[1.0e-38]]]], dtype=np.float32)), + _tensor("scale_a", np.asarray(1.0e20, dtype=np.float32)), + _tensor("scale_b", np.asarray(1.0e20, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 1, 1, 1]), + _info("branch", [1, 1, 1, 1]), + _info("leaf", [1, 1, 1, 1]), + _info("scaled_a", [1, 1, 1, 1]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_nonfinite_synthesized_conv_bias_does_not_partially_mutate(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight", "bias"], ["conv_out"]), + onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), + ], + [_info("x", [1, 1, 1, 1])], + [_info("y", [1, 1, 1, 1])], + [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("bias", np.asarray([1.0e38], dtype=np.float32)), + _tensor("scale", np.asarray(10.0, dtype=np.float32)), + ], + value_info=[_info("conv_out", [1, 1, 1, 1])], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + def test_float64_affine_values_preserve_weight_precision(self) -> None: shape = [1, 1, 1, 1] @@ -375,21 +911,213 @@ def test_shared_conv_output_is_ineligible( "Identity", ] - def test_routed_view_graph_output_is_ineligible(self) -> None: - rng = np.random.default_rng(19) + @pytest.mark.parametrize("operand_name", ["scale", "offset"]) + def test_overridable_affine_operands_are_unchanged( + self, + affine_model: tuple[onnx.ModelProto, dict[str, np.ndarray]], + operand_name: str, + ) -> None: + model, _ = affine_model + initializer = next(value for value in model.graph.initializer if value.name == operand_name) + model.graph.input.append( + onnx.helper.make_tensor_value_info( + operand_name, + initializer.data_type, + list(initializer.dims), + ) + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize("parameter_name", ["weight", "bias"]) + def test_overridable_conv_parameters_are_unchanged(self, parameter_name: str) -> None: model = _model( [ - onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), - onnx.helper.make_node("Split", ["conv_out"], ["left", "right"], axis=1), - onnx.helper.make_node("Reshape", ["left", "view_shape"], ["left_view"]), - onnx.helper.make_node("Mul", ["left_view", "scale"], ["left_scaled"]), - onnx.helper.make_node("Identity", ["right"], ["right_out"]), + onnx.helper.make_node("Conv", ["x", "weight", "bias"], ["conv_out"]), + onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), ], [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], [ - _info("left_view", [1, 1, 1, 2, 2]), - _info("left_scaled", [1, 1, 1, 2, 2]), - _info("right_out", [1, 1, 2, 2]), + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("bias", np.ones(1, dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[_info("conv_out", [1, 1, 2, 2])], + ) + parameter = next(value for value in model.graph.initializer if value.name == parameter_name) + model.graph.input.append( + onnx.helper.make_tensor_value_info( + parameter_name, + parameter.data_type, + list(parameter.dims), + ) + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize("parameter_name", ["weight", "bias"]) + def test_unloaded_external_conv_parameter_is_unchanged( + self, + parameter_name: str, + ) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight", "bias"], ["conv_out"]), + onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("bias", np.ones(1, dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[_info("conv_out", [1, 1, 2, 2])], + ) + parameter = next(value for value in model.graph.initializer if value.name == parameter_name) + parameter.ClearField("raw_data") + parameter.data_location = onnx.TensorProto.EXTERNAL + location = parameter.external_data.add() + location.key = "location" + location.value = f"missing-{parameter_name}.bin" + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize("affine_op", ["Mul", "Add"]) + @pytest.mark.parametrize("affine_value", [np.nan, np.inf, -np.inf]) + def test_nonfinite_affine_operand_is_unchanged( + self, + affine_op: str, + affine_value: float, + ) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Constant", [], ["affine"], value_float=affine_value), + onnx.helper.make_node(affine_op, ["conv_out", "affine"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [_tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32))], + value_info=[_info("conv_out", [1, 1, 2, 2])], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize( + "constant_name", + [ + "view_shape", + "split_sizes", + "slice_starts", + "slice_ends", + "slice_axes", + "slice_steps", + ], + ) + def test_overridable_route_constants_are_unchanged(self, constant_name: str) -> None: + shape = [1, 2, 1, 2, 2] + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Reshape", ["conv_out", "view_shape"], ["viewed"]), + onnx.helper.make_node( + "Split", + ["viewed", "split_sizes"], + ["split_out"], + axis=1, + ), + onnx.helper.make_node( + "Slice", + [ + "split_out", + "slice_starts", + "slice_ends", + "slice_axes", + "slice_steps", + ], + ["sliced"], + ), + onnx.helper.make_node("Mul", ["sliced", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", shape)], + [ + _tensor("weight", np.ones((2, 1, 1, 1), dtype=np.float32)), + _tensor("view_shape", np.asarray(shape, dtype=np.int64)), + _tensor("split_sizes", np.asarray([2], dtype=np.int64)), + _tensor("slice_starts", np.asarray([0], dtype=np.int64)), + _tensor("slice_ends", np.asarray([2], dtype=np.int64)), + _tensor("slice_axes", np.asarray([1], dtype=np.int64)), + _tensor("slice_steps", np.asarray([1], dtype=np.int64)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 2, 2, 2]), + _info("viewed", shape), + _info("split_out", shape), + _info("sliced", shape), + ], + ) + initializer = next( + value for value in model.graph.initializer if value.name == constant_name + ) + model.graph.input.append( + onnx.helper.make_tensor_value_info( + constant_name, + initializer.data_type, + list(initializer.dims), + ) + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_routed_view_graph_output_is_ineligible(self) -> None: + rng = np.random.default_rng(19) + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["left", "right"], axis=1), + onnx.helper.make_node("Reshape", ["left", "view_shape"], ["left_view"]), + onnx.helper.make_node("Mul", ["left_view", "scale"], ["left_scaled"]), + onnx.helper.make_node("Identity", ["right"], ["right_out"]), + ], + [_info("x", [1, 1, 2, 2])], + [ + _info("left_view", [1, 1, 1, 2, 2]), + _info("left_scaled", [1, 1, 1, 2, 2]), + _info("right_out", [1, 1, 2, 2]), ], [ _tensor("weight", rng.normal(size=(2, 1, 1, 1)).astype(np.float32)), @@ -416,6 +1144,144 @@ def test_routed_view_graph_output_is_ineligible(self) -> None: ): np.testing.assert_array_equal(original, rewritten) + def test_reshape_allowzero_without_literal_zero_folds_affine_branches(self) -> None: + rng = np.random.default_rng(20) + branch_shape = [1, 1, 1, 2, 2] + model = _model( + [ + onnx.helper.make_node( + "Conv", + ["x", "weights"], + ["conv_out"], + name="conv", + ), + onnx.helper.make_node( + "Reshape", + ["conv_out", "view_shape"], + ["viewed"], + name="allowzero_view", + allowzero=1, + ), + onnx.helper.make_node( + "Split", + ["viewed", "split_sizes"], + ["scalar_branch", "affine_branch", "nonlinear_branch"], + name="channel_split", + axis=1, + ), + onnx.helper.make_node( + "Mul", + ["scalar_branch", "scalar_scale"], + ["scalar_out"], + name="scalar_mul", + ), + onnx.helper.make_node( + "Mul", + ["affine_branch", "affine_scale"], + ["affine_scaled"], + name="affine_mul", + ), + onnx.helper.make_node( + "Add", + ["affine_scaled", "affine_offset"], + ["affine_out"], + name="affine_add", + ), + onnx.helper.make_node( + "Relu", + ["nonlinear_branch"], + ["nonlinear_out"], + name="nonlinear_relu", + ), + ], + [_info("x", [1, 1, 2, 2])], + [ + _info("scalar_out", branch_shape), + _info("affine_out", branch_shape), + _info("nonlinear_out", branch_shape), + ], + [ + _tensor("weights", rng.normal(size=(3, 1, 1, 1)).astype(np.float32)), + _tensor("view_shape", np.asarray([1, -1, 1, 2, 2], dtype=np.int64)), + _tensor("split_sizes", np.asarray([1, 1, 1], dtype=np.int64)), + _tensor("scalar_scale", np.asarray(1.25, dtype=np.float32)), + _tensor("affine_scale", np.asarray([[[[[0.75]]]]], dtype=np.float32)), + _tensor("affine_offset", np.asarray([[[[[-0.5]]]]], dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 3, 2, 2]), + _info("viewed", [1, 3, 1, 2, 2]), + _info("scalar_branch", branch_shape), + _info("affine_branch", branch_shape), + _info("affine_scaled", branch_shape), + _info("nonlinear_branch", branch_shape), + ], + ) + values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} + config = AlgebraicRewritePipeConfig(conv_channel_affine_folding=True) + transformed = AlgebraicRewritePipe().process(model, config) + second = AlgebraicRewritePipe().process(transformed, config) + + remaining_names = {node.name for node in transformed.graph.node} + assert not {"scalar_mul", "affine_mul", "affine_add"} & remaining_names + assert "nonlinear_relu" in remaining_names + assert transformed.SerializeToString() == second.SerializeToString() + _assert_valid_with_inferred_shapes(transformed) + for original, rewritten in zip( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + + def test_reshape_allowzero_with_literal_zero_keeps_affine_nodes(self) -> None: + rng = np.random.default_rng(21) + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), + onnx.helper.make_node( + "Reshape", + ["conv_out", "view_shape"], + ["viewed"], + allowzero=1, + ), + onnx.helper.make_node( + "Mul", + ["viewed", "scale"], + ["scaled"], + name="zero_shape_mul", + ), + onnx.helper.make_node( + "Add", + ["scaled", "offset"], + ["y"], + name="zero_shape_add", + ), + ], + [_info("x", [0, 1, 2, 2])], + [_info("y", [0, 2, 2, 2])], + [ + _tensor("weights", rng.normal(size=(2, 1, 1, 1)).astype(np.float32)), + _tensor("view_shape", np.asarray([0, 2, 2, 2], dtype=np.int64)), + _tensor("scale", np.asarray(1.25, dtype=np.float32)), + _tensor("offset", np.asarray(-0.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [0, 2, 2, 2]), + _info("viewed", [0, 2, 2, 2]), + _info("scaled", [0, 2, 2, 2]), + ], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + remaining_names = {node.name for node in transformed.graph.node} + assert {"zero_shape_mul", "zero_shape_add"} <= remaining_names + _assert_valid_with_inferred_shapes(transformed) + def test_static_split_branches_fold_without_overlapping_ranges(self) -> None: rng = np.random.default_rng(12) x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 1, 2, 2]) @@ -463,78 +1329,73 @@ def test_static_split_branches_fold_without_overlapping_ranges(self) -> None: atol=2e-5, ) - def test_channel_preserving_views_and_nested_slices_fold(self) -> None: - rng = np.random.default_rng(15) - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 1, 2, 2]) - nodes = [ - onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), - onnx.helper.make_node("Reshape", ["conv_out", "view_shape"], ["viewed"]), - onnx.helper.make_node("Split", ["viewed", "split_sizes"], ["first", "second"], axis=1), - onnx.helper.make_node("Squeeze", ["first", "squeeze_axes"], ["first_view"]), - onnx.helper.make_node("Mul", ["first_view", "first_scale"], ["first_scaled"]), - onnx.helper.make_node("Relu", ["first_scaled"], ["first_out"]), - onnx.helper.make_node( - "Slice", - ["second", "slice_a_starts", "slice_a_ends"], - ["second_a"], - ), - onnx.helper.make_node( - "Slice", - [ - "second", - "slice_b_starts", - "slice_b_ends", - "slice_b_axes", - "slice_b_steps", - ], - ["second_b"], - ), - onnx.helper.make_node("Mul", ["second_a", "second_a_scale"], ["second_a_out"]), - onnx.helper.make_node("Add", ["second_b", "second_b_offset"], ["second_b_out"]), - ] + def test_nested_static_channel_splits_fold_affine_leaves(self) -> None: + rng = np.random.default_rng(22) model = _model( - nodes, - [x], [ - _info("first_out", [1, 2, 2, 2]), - _info("second_a_out", [1, 1, 1, 2, 2]), - _info("second_b_out", [1, 1, 1, 2, 2]), + onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), + onnx.helper.make_node( + "Split", + ["conv_out", "outer_sizes"], + ["depth", "colors", "keep"], + name="outer_split", + axis=1, + ), + onnx.helper.make_node("Mul", ["depth", "depth_scale"], ["depth_out"]), + onnx.helper.make_node( + "Split", + ["colors", "inner_sizes"], + ["rgb", "sh"], + name="inner_split", + axis=1, + ), + onnx.helper.make_node("Mul", ["rgb", "rgb_scale"], ["rgb_scaled"]), + onnx.helper.make_node("Add", ["rgb_scaled", "rgb_offset"], ["rgb_affine"]), + onnx.helper.make_node("Sigmoid", ["rgb_affine"], ["rgb_out"]), + onnx.helper.make_node("Mul", ["sh", "sh_scale"], ["sh_scaled"]), + onnx.helper.make_node("Add", ["sh_scaled", "sh_offset"], ["sh_out"]), + onnx.helper.make_node("Relu", ["keep"], ["keep_out"]), ], + [_info("x", [1, 1, 2, 2])], [ - _tensor("weights", rng.normal(size=(4, 1, 1, 1)).astype(np.float32)), - _tensor("view_shape", np.asarray([1, 4, 1, 2, 2], dtype=np.int64)), - _tensor("split_sizes", np.asarray([2, 2], dtype=np.int64)), - _tensor("squeeze_axes", np.asarray([2], dtype=np.int64)), - _tensor("slice_a_starts", np.asarray([0, 0, 0, 0, 0], dtype=np.int64)), - _tensor("slice_a_ends", np.asarray([1, 1, 1, 2, 2], dtype=np.int64)), - _tensor("slice_b_starts", np.asarray([0, 1, 0, 0, 0], dtype=np.int64)), - _tensor( - "slice_b_ends", - np.full(5, np.iinfo(np.int64).max, dtype=np.int64), - ), - _tensor("slice_b_axes", np.asarray([0, 1, 2, 3, 4], dtype=np.int64)), - _tensor("slice_b_steps", np.ones(5, dtype=np.int64)), - _tensor("first_scale", np.asarray(1.25, dtype=np.float32)), - _tensor("second_a_scale", np.asarray(0.75, dtype=np.float32)), - _tensor("second_b_offset", np.asarray(-0.5, dtype=np.float32)), + _info("depth_out", [1, 1, 2, 2]), + _info("rgb_out", [1, 1, 2, 2]), + _info("sh_out", [1, 3, 2, 2]), + _info("keep_out", [1, 1, 2, 2]), + ], + [ + _tensor("weights", rng.normal(size=(6, 1, 1, 1)).astype(np.float32)), + _tensor("outer_sizes", np.asarray([1, 4, 1], dtype=np.int64)), + _tensor("inner_sizes", np.asarray([1, 3], dtype=np.int64)), + _tensor("depth_scale", np.asarray(0.25, dtype=np.float32)), + _tensor("rgb_scale", np.asarray(0.75, dtype=np.float32)), + _tensor("rgb_offset", np.asarray(-0.125, dtype=np.float32)), + _tensor("sh_scale", np.asarray(1.5, dtype=np.float32)), + _tensor("sh_offset", np.asarray(0.25, dtype=np.float32)), ], value_info=[ - _info("conv_out", [1, 4, 2, 2]), - _info("viewed", [1, 4, 1, 2, 2]), - _info("first", [1, 2, 1, 2, 2]), - _info("second", [1, 2, 1, 2, 2]), - _info("first_view", [1, 2, 2, 2]), - _info("first_scaled", [1, 2, 2, 2]), - _info("second_a", [1, 1, 1, 2, 2]), - _info("second_b", [1, 1, 1, 2, 2]), + _info("conv_out", [1, 6, 2, 2]), + _info("depth", [1, 1, 2, 2]), + _info("colors", [1, 4, 2, 2]), + _info("keep", [1, 1, 2, 2]), + _info("rgb", [1, 1, 2, 2]), + _info("sh", [1, 3, 2, 2]), + _info("rgb_scaled", [1, 1, 2, 2]), + _info("rgb_affine", [1, 1, 2, 2]), + _info("sh_scaled", [1, 3, 2, 2]), ], ) values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} config = AlgebraicRewritePipeConfig(conv_channel_affine_folding=True) + transformed = AlgebraicRewritePipe().process(model, config) second = AlgebraicRewritePipe().process(transformed, config) assert not any(node.op_type in {"Mul", "Add"} for node in transformed.graph.node) + assert {node.name for node in transformed.graph.node if node.op_type == "Split"} == { + "outer_split", + "inner_split", + } assert transformed.SerializeToString() == second.SerializeToString() _assert_valid_with_inferred_shapes(transformed) for original, rewritten in zip( @@ -544,138 +1405,2016 @@ def test_channel_preserving_views_and_nested_slices_fold(self) -> None: ): np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) - public = optimize_onnx( - model, - conv_channel_affine_folding=True, - static_split_to_slice=True, - ) - second_public = optimize_onnx( - public, - conv_channel_affine_folding=True, - static_split_to_slice=True, - ) - assert not any(node.op_type in {"Mul", "Add", "Split"} for node in public.graph.node) - assert [ - (node.op_type, tuple(node.input), tuple(node.output), node.name) - for node in public.graph.node - ] == [ - (node.op_type, tuple(node.input), tuple(node.output), node.name) - for node in second_public.graph.node - ] - _assert_valid_with_inferred_shapes(second_public) - for original, rewritten in zip( - _run(model, values), - _run(second_public, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) - - def test_multiple_independent_affine_matches_fold(self) -> None: - rng = np.random.default_rng(16) - nodes = [ - onnx.helper.make_node("Conv", ["x1", "weight1"], ["conv1"]), - onnx.helper.make_node("Mul", ["conv1", "scale1"], ["y1"]), - onnx.helper.make_node("Conv", ["x2", "weight2"], ["conv2"]), - onnx.helper.make_node("Add", ["conv2", "offset2"], ["y2"]), + @pytest.mark.parametrize("case", ["custom_domain", "non_channel", "dynamic", "malformed"]) + def test_invalid_nested_split_is_unchanged(self, case: str) -> None: + nested_inputs = ["branch"] + nested_domain = "" + nested_axis = 1 + inputs = [_info("x", [1, 1, 2, 2])] + initializers = [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), ] + if case == "custom_domain": + nested_domain = "com.example" + elif case == "non_channel": + nested_axis = 2 + elif case == "dynamic": + nested_inputs.append("nested_sizes") + inputs.append( + onnx.helper.make_tensor_value_info( + "nested_sizes", + onnx.TensorProto.INT64, + [1], + ) + ) + else: + nested_inputs.append("nested_sizes") + initializers.append(_tensor("nested_sizes", np.asarray([2], dtype=np.int64))) model = _model( - nodes, + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["branch"], axis=1), + onnx.helper.make_node( + "Split", + nested_inputs, + ["leaf"], + axis=nested_axis, + domain=nested_domain, + ), + onnx.helper.make_node("Mul", ["leaf", "scale"], ["y"]), + ], + inputs, + [_info("y", [1, 1, 2, 2])], + initializers, + value_info=[ + _info("conv_out", [1, 1, 2, 2]), + _info("branch", [1, 1, 2, 2]), + _info("leaf", [1, 1, 2, 2]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_duplicate_nested_split_outputs_are_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["branch"], axis=1), + onnx.helper.make_node( + "Split", + ["branch"], + ["leaf", "leaf"], + axis=1, + ), + onnx.helper.make_node("Mul", ["leaf", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((2, 1, 1, 1), dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 2, 2, 2]), + _info("branch", [1, 2, 2, 2]), + _info("leaf", [1, 1, 2, 2]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_duplicate_routed_tensor_and_affine_node_are_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["branch"], axis=1), + onnx.helper.make_node( + "Slice", + ["branch", "left_starts", "left_ends", "axes"], + ["leaf"], + ), + onnx.helper.make_node( + "Slice", + ["branch", "right_starts", "right_ends", "axes"], + ["leaf"], + ), + onnx.helper.make_node("Mul", ["leaf", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((2, 1, 1, 1), dtype=np.float32)), + _tensor("left_starts", np.asarray([0], dtype=np.int64)), + _tensor("left_ends", np.asarray([1], dtype=np.int64)), + _tensor("right_starts", np.asarray([1], dtype=np.int64)), + _tensor("right_ends", np.asarray([2], dtype=np.int64)), + _tensor("axes", np.asarray([1], dtype=np.int64)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 2, 2, 2]), + _info("branch", [1, 2, 2, 2]), + _info("leaf", [1, 1, 2, 2]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_repeated_tensor_across_sibling_routes_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node( + "Split", + ["conv_out", "outer_sizes"], + ["left", "right"], + axis=1, + ), + onnx.helper.make_node("Split", ["left"], ["shared"], axis=1), + onnx.helper.make_node("Split", ["right"], ["shared"], axis=1), + onnx.helper.make_node("Mul", ["shared", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((3, 1, 1, 1), dtype=np.float32)), + _tensor("outer_sizes", np.asarray([1, 2], dtype=np.int64)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 3, 2, 2]), + _info("left", [1, 1, 2, 2]), + _info("right", [1, 2, 2, 2]), + _info("shared", [1, 1, 2, 2]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize("protection", ["graph_output", "shared", "captured"]) + def test_protected_nested_route_is_unchanged(self, protection: str) -> None: + nodes = [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["branch"], axis=1), + onnx.helper.make_node("Split", ["branch"], ["leaf"], axis=1), + onnx.helper.make_node("Mul", ["leaf", "scale"], ["y"]), + ] + outputs = [_info("y", [1, 1, 2, 2])] + initializers = [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ] + if protection == "graph_output": + outputs.append(_info("leaf", [1, 1, 2, 2])) + elif protection == "shared": + nodes.append(onnx.helper.make_node("Identity", ["leaf"], ["protected"])) + outputs.append(_info("protected", [1, 1, 2, 2])) + else: + branch = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["leaf"], ["branch_output"])], + "capturing_branch", + [], + [_info("branch_output", [1, 1, 2, 2])], + ) + initializers.append(_tensor("condition", np.asarray(True, dtype=np.bool_))) + nodes.append( + onnx.helper.make_node( + "If", + ["condition"], + ["protected"], + then_branch=branch, + else_branch=branch, + ) + ) + outputs.append(_info("protected", [1, 1, 2, 2])) + model = _model( + nodes, + [_info("x", [1, 1, 2, 2])], + outputs, + initializers, + value_info=[ + _info("conv_out", [1, 1, 2, 2]), + _info("branch", [1, 1, 2, 2]), + _info("leaf", [1, 1, 2, 2]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_nested_split_cycle_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["route_a"], axis=1), + onnx.helper.make_node("Split", ["route_a"], ["route_b"], axis=1), + onnx.helper.make_node("Split", ["route_b"], ["route_a"], axis=1), + onnx.helper.make_node("Identity", ["x"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [_tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32))], + value_info=[ + _info("conv_out", [1, 1, 2, 2]), + _info("route_a", [1, 1, 2, 2]), + _info("route_b", [1, 1, 2, 2]), + ], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_deep_nested_split_route_is_unchanged(self) -> None: + route_depth = 65 + nodes = [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Split", ["conv_out"], ["route_0"], axis=1), + ] + value_info = [ + _info("conv_out", [1, 1, 2, 2]), + _info("route_0", [1, 1, 2, 2]), + ] + for route_index in range(route_depth): + nodes.append( + onnx.helper.make_node( + "Split", + [f"route_{route_index}"], + [f"route_{route_index + 1}"], + axis=1, + ) + ) + value_info.append(_info(f"route_{route_index + 1}", [1, 1, 2, 2])) + nodes.append(onnx.helper.make_node("Mul", [f"route_{route_depth}", "scale"], ["y"])) + model = _model( + nodes, + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=value_info, + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_channel_preserving_views_and_nested_slices_fold(self) -> None: + rng = np.random.default_rng(15) + x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 1, 2, 2]) + nodes = [ + onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), + onnx.helper.make_node("Reshape", ["conv_out", "view_shape"], ["viewed"]), + onnx.helper.make_node("Split", ["viewed", "split_sizes"], ["first", "second"], axis=1), + onnx.helper.make_node("Squeeze", ["first", "squeeze_axes"], ["first_view"]), + onnx.helper.make_node("Mul", ["first_view", "first_scale"], ["first_scaled"]), + onnx.helper.make_node("Relu", ["first_scaled"], ["first_out"]), + onnx.helper.make_node( + "Slice", + ["second", "slice_a_starts", "slice_a_ends"], + ["second_a"], + ), + onnx.helper.make_node( + "Slice", + [ + "second", + "slice_b_starts", + "slice_b_ends", + "slice_b_axes", + "slice_b_steps", + ], + ["second_b"], + ), + onnx.helper.make_node("Mul", ["second_a", "second_a_scale"], ["second_a_out"]), + onnx.helper.make_node("Add", ["second_b", "second_b_offset"], ["second_b_out"]), + ] + model = _model( + nodes, + [x], + [ + _info("first_out", [1, 2, 2, 2]), + _info("second_a_out", [1, 1, 1, 2, 2]), + _info("second_b_out", [1, 1, 1, 2, 2]), + ], + [ + _tensor("weights", rng.normal(size=(4, 1, 1, 1)).astype(np.float32)), + _tensor("view_shape", np.asarray([1, 4, 1, 2, 2], dtype=np.int64)), + _tensor("split_sizes", np.asarray([2, 2], dtype=np.int64)), + _tensor("squeeze_axes", np.asarray([2], dtype=np.int64)), + _tensor("slice_a_starts", np.asarray([0, 0, 0, 0, 0], dtype=np.int64)), + _tensor("slice_a_ends", np.asarray([1, 1, 1, 2, 2], dtype=np.int64)), + _tensor("slice_b_starts", np.asarray([0, 1, 0, 0, 0], dtype=np.int64)), + _tensor( + "slice_b_ends", + np.full(5, np.iinfo(np.int64).max, dtype=np.int64), + ), + _tensor("slice_b_axes", np.asarray([0, 1, 2, 3, 4], dtype=np.int64)), + _tensor("slice_b_steps", np.ones(5, dtype=np.int64)), + _tensor("first_scale", np.asarray(1.25, dtype=np.float32)), + _tensor("second_a_scale", np.asarray(0.75, dtype=np.float32)), + _tensor("second_b_offset", np.asarray(-0.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", [1, 4, 2, 2]), + _info("viewed", [1, 4, 1, 2, 2]), + _info("first", [1, 2, 1, 2, 2]), + _info("second", [1, 2, 1, 2, 2]), + _info("first_view", [1, 2, 2, 2]), + _info("first_scaled", [1, 2, 2, 2]), + _info("second_a", [1, 1, 1, 2, 2]), + _info("second_b", [1, 1, 1, 2, 2]), + ], + ) + values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} + config = AlgebraicRewritePipeConfig(conv_channel_affine_folding=True) + transformed = AlgebraicRewritePipe().process(model, config) + second = AlgebraicRewritePipe().process(transformed, config) + + assert not any(node.op_type in {"Mul", "Add"} for node in transformed.graph.node) + assert transformed.SerializeToString() == second.SerializeToString() + _assert_valid_with_inferred_shapes(transformed) + for original, rewritten in zip( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + + public = optimize_onnx( + model, + conv_channel_affine_folding=True, + static_split_to_slice=True, + ) + second_public = optimize_onnx( + public, + conv_channel_affine_folding=True, + static_split_to_slice=True, + ) + assert not any(node.op_type in {"Mul", "Add", "Split"} for node in public.graph.node) + assert [ + (node.op_type, tuple(node.input), tuple(node.output), node.name) + for node in public.graph.node + ] == [ + (node.op_type, tuple(node.input), tuple(node.output), node.name) + for node in second_public.graph.node + ] + _assert_valid_with_inferred_shapes(second_public) + for original, rewritten in zip( + _run(model, values), + _run(second_public, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + + def test_multiple_independent_affine_matches_fold(self) -> None: + rng = np.random.default_rng(16) + nodes = [ + onnx.helper.make_node("Conv", ["x1", "weight1"], ["conv1"]), + onnx.helper.make_node("Mul", ["conv1", "scale1"], ["y1"]), + onnx.helper.make_node("Conv", ["x2", "weight2"], ["conv2"]), + onnx.helper.make_node("Add", ["conv2", "offset2"], ["y2"]), + ] + model = _model( + nodes, [_info("x1", [1, 1, 2, 2]), _info("x2", [1, 1, 2, 2])], [_info("y1", [1, 1, 2, 2]), _info("y2", [1, 1, 2, 2])], [ - _tensor("weight1", rng.normal(size=(1, 1, 1, 1)).astype(np.float32)), - _tensor("scale1", np.asarray(1.5, dtype=np.float32)), - _tensor("weight2", rng.normal(size=(1, 1, 1, 1)).astype(np.float32)), - _tensor("offset2", np.asarray(-0.25, dtype=np.float32)), + _tensor("weight1", rng.normal(size=(1, 1, 1, 1)).astype(np.float32)), + _tensor("scale1", np.asarray(1.5, dtype=np.float32)), + _tensor("weight2", rng.normal(size=(1, 1, 1, 1)).astype(np.float32)), + _tensor("offset2", np.asarray(-0.25, dtype=np.float32)), + ], + value_info=[_info("conv1", [1, 1, 2, 2]), _info("conv2", [1, 1, 2, 2])], + ) + values = { + "x1": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), + "x2": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), + } + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + assert [node.op_type for node in transformed.graph.node] == ["Conv", "Conv"] + for original, rewritten in zip( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + + def test_nested_subgraph_captures_make_affine_fold_ineligible( + self, + affine_model: tuple[onnx.ModelProto, dict[str, np.ndarray]], + ) -> None: + model, values = affine_model + branch_shape = [1, 3, 3, 3] + then_branch = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["mul_out"], ["then_output"])], + "then_branch", + [], + [_info("then_output", branch_shape)], + ) + else_branch = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["conv_out"], ["else_output"])], + "else_branch", + [], + [_info("else_output", branch_shape)], + ) + model.graph.initializer.append(_tensor("condition", np.asarray(True, dtype=np.bool_))) + model.graph.node.append( + onnx.helper.make_node( + "If", + ["condition"], + ["captured"], + then_branch=then_branch, + else_branch=else_branch, + ) + ) + model.graph.output.append(_info("captured", branch_shape)) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + assert [node.op_type for node in transformed.graph.node] == [ + "Conv", + "Mul", + "Add", + "If", + ] + _assert_valid_with_inferred_shapes(transformed) + for original, rewritten in zip( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=0, atol=0) + + def test_constant_attribute_affine_is_folded_and_pruned(self) -> None: + rng = np.random.default_rng(18) + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Constant", [], ["scale"], value_float=1.5), + onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [_tensor("weight", rng.normal(size=(1, 1, 1, 1)).astype(np.float32))], + value_info=[_info("conv_out", [1, 1, 2, 2])], + ) + values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + assert [node.op_type for node in transformed.graph.node] == ["Conv"] + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-5, + atol=2e-5, + ) + + @pytest.mark.parametrize( + ("domain", "should_fold"), + [("", True), ("ai.onnx", False), ("com.example", False)], + ) + def test_only_standard_domain_constant_is_interpreted( + self, + domain: str, + should_fold: bool, + ) -> None: + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node( + "Constant", + [], + ["scale"], + value_float=1.5, + domain=domain, + ), + onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [_tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32))], + value_info=[_info("conv_out", [1, 1, 2, 2])], + ) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + if should_fold: + assert [node.op_type for node in transformed.graph.node] == ["Conv"] + else: + assert transformed.SerializeToString() == original + + def test_custom_domain_conv_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node( + "Conv", + ["x", "weight"], + ["conv_out"], + name="custom_conv", + domain="com.example", + ), + onnx.helper.make_node( + "Mul", + ["conv_out", "scale"], + ["y"], + name="affine_mul", + ), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[_info("conv_out", [1, 1, 2, 2])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + _assert_byte_identical(model, transformed) + + @pytest.mark.parametrize("affine_op", ["Mul", "Add"]) + @pytest.mark.parametrize("route", ["direct", "nested_split"]) + def test_custom_domain_affine_node_is_unchanged( + self, + affine_op: str, + route: str, + ) -> None: + nodes = [onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"])] + value_info = [_info("conv_out", [1, 1, 2, 2])] + affine_input = "conv_out" + if route == "nested_split": + nodes.extend( + [ + onnx.helper.make_node( + "Split", + ["conv_out"], + ["outer_branch"], + axis=1, + ), + onnx.helper.make_node( + "Split", + ["outer_branch"], + ["affine_input"], + axis=1, + ), + ] + ) + value_info.extend( + [ + _info("outer_branch", [1, 1, 2, 2]), + _info("affine_input", [1, 1, 2, 2]), + ] + ) + affine_input = "affine_input" + nodes.append( + onnx.helper.make_node( + affine_op, + [affine_input, "affine_value"], + ["y"], + name="custom_affine", + domain="com.example", + ) + ) + model = _model( + nodes, + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((1, 1, 1, 1), dtype=np.float32)), + _tensor("affine_value", np.asarray(1.5, dtype=np.float32)), + ], + value_info=value_info, + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + _assert_byte_identical(model, transformed) + + def test_custom_domain_slice_below_nested_split_is_unchanged(self) -> None: + shape = [1, 2, 2, 2] + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node( + "Split", + ["conv_out"], + ["outer_branch"], + axis=1, + ), + onnx.helper.make_node( + "Split", + ["outer_branch"], + ["slice_input"], + axis=1, + ), + onnx.helper.make_node( + "Slice", + ["slice_input", "starts", "ends", "axes"], + ["sliced"], + name="custom_slice", + domain="com.example", + ), + onnx.helper.make_node("Mul", ["sliced", "scale"], ["y"]), + ], + [_info("x", [1, 1, 2, 2])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("weight", np.ones((2, 1, 1, 1), dtype=np.float32)), + _tensor("starts", np.asarray([0], dtype=np.int64)), + _tensor("ends", np.asarray([1], dtype=np.int64)), + _tensor("axes", np.asarray([1], dtype=np.int64)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", shape), + _info("outer_branch", shape), + _info("slice_input", shape), + _info("sliced", [1, 1, 2, 2]), + ], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + _assert_byte_identical(model, transformed) + + @pytest.mark.parametrize("view_op", ["Reshape", "Squeeze", "Unsqueeze"]) + def test_custom_domain_shape_view_below_nested_split_is_unchanged( + self, + view_op: str, + ) -> None: + source_shape = [1, 1, 1, 2, 2] if view_op == "Squeeze" else [1, 1, 2, 2] + output_shape = [1, 1, 2, 2] if view_op == "Squeeze" else [1, 1, 1, 2, 2] + weight_shape = (1, 1, 1, 1, 1) if view_op == "Squeeze" else (1, 1, 1, 1) + view_parameter = "view_shape" if view_op == "Reshape" else "view_axes" + view_parameter_values = ( + np.asarray(output_shape, dtype=np.int64) + if view_op == "Reshape" + else np.asarray([2], dtype=np.int64) + ) + model = _model( + [ + onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node( + "Split", + ["conv_out"], + ["outer_branch"], + axis=1, + ), + onnx.helper.make_node( + "Split", + ["outer_branch"], + ["view_input"], + axis=1, + ), + onnx.helper.make_node( + view_op, + ["view_input", view_parameter], + ["viewed"], + name="custom_view", + domain="com.example", + ), + onnx.helper.make_node("Mul", ["viewed", "scale"], ["y"]), + ], + [_info("x", source_shape)], + [_info("y", output_shape)], + [ + _tensor("weight", np.ones(weight_shape, dtype=np.float32)), + _tensor(view_parameter, view_parameter_values), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("conv_out", source_shape), + _info("outer_branch", source_shape), + _info("view_input", source_shape), + _info("viewed", output_shape), + ], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + ) + + _assert_byte_identical(model, transformed) + + +class TestExpPositiveScaleFolding: + """Test opt-in positive scale folding into the Exp input.""" + + def test_multiple_exp_chains_are_folded_from_live_graph_nodes(self) -> None: + rng = np.random.default_rng(35) + model = _model( + [ + onnx.helper.make_node("Add", ["x0", "bias0"], ["biased0"]), + onnx.helper.make_node("Exp", ["biased0"], ["exp0"]), + onnx.helper.make_node("Mul", ["exp0", "scale0"], ["y0"]), + onnx.helper.make_node("Add", ["x1", "bias1"], ["biased1"]), + onnx.helper.make_node("Exp", ["biased1"], ["exp1"]), + onnx.helper.make_node("Mul", ["exp1", "scale1"], ["y1"]), + ], + [_info("x0", [1, 2]), _info("x1", [1, 2])], + [_info("y0", [1, 2]), _info("y1", [1, 2])], + [ + _tensor("bias0", np.asarray([[0.25, -0.5]], dtype=np.float32)), + _tensor("scale0", np.asarray([[1.25, 0.75]], dtype=np.float32)), + _tensor("bias1", np.asarray([[1.0, -1.25]], dtype=np.float32)), + _tensor("scale1", np.asarray([[2.0, 1.5]], dtype=np.float32)), + ], + value_info=[ + _info("biased0", [1, 2]), + _info("exp0", [1, 2]), + _info("biased1", [1, 2]), + _info("exp1", [1, 2]), + ], + ) + values = { + "x0": rng.normal(size=(1, 2)).astype(np.float32), + "x1": rng.normal(size=(1, 2)).astype(np.float32), + } + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp", "Add", "Exp"] + assert [output.name for output in transformed.graph.output] == ["y0", "y1"] + _assert_valid_with_inferred_shapes(transformed) + for original, rewritten in zip( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-6, atol=2e-6) + + def test_independent_exp_scale_chains_are_batched(self, monkeypatch) -> None: + chain_count = 6 + nodes = [] + inputs = [] + outputs = [] + initializers = [] + value_info = [] + for index in range(chain_count): + inputs.append(_info(f"x{index}", [1, 2])) + outputs.append(_info(f"y{index}", [1, 2])) + initializers.append( + _tensor(f"scale{index}", np.asarray([1.25, 0.75], dtype=np.float32)) + ) + value_info.append(_info(f"exp{index}", [1, 2])) + nodes.extend( + [ + onnx.helper.make_node("Exp", [f"x{index}"], [f"exp{index}"]), + onnx.helper.make_node("Mul", [f"exp{index}", f"scale{index}"], [f"y{index}"]), + ] + ) + model = _model(nodes, inputs, outputs, initializers, value_info=value_info) + build_count = 0 + original_build = algebraic_pipe._GraphIndex.build + + def counted_build(model: onnx.ModelProto) -> algebraic_pipe._GraphIndex: + nonlocal build_count + build_count += 1 + return original_build(model) + + monkeypatch.setattr(algebraic_pipe._GraphIndex, "build", counted_build) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert build_count <= 6 + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] * chain_count + + def test_serial_exp_mul_keeps_later_scales_to_preserve_float_boundaries( + self, + ) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Mul", ["scaled", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [ + _tensor("scale_a", np.asarray(1.0e-30, dtype=np.float32)), + _tensor("scale_b", np.asarray(1.0e30, dtype=np.float32)), + ], + value_info=[_info("exp", [1]), _info("scaled", [1])], + ) + values = {"x": np.asarray([np.log(np.float32(1.0e-20))], dtype=np.float32)} + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert transformed.SerializeToString() == original + np.testing.assert_array_equal(_run(model, values), _run(transformed, values)) + + def test_public_optimizer_preserves_serial_exp_mul_float_boundaries(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Mul", ["scaled", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [ + _tensor("scale_a", np.asarray(1.0e-30, dtype=np.float32)), + _tensor("scale_b", np.asarray(1.0e30, dtype=np.float32)), + ], + value_info=[_info("exp", [1]), _info("scaled", [1])], + ) + values = {"x": np.asarray([np.log(np.float32(1.0e-20))], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert [node.op_type for node in transformed.graph.node] == ["Exp", "Mul", "Mul"] + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_array_equal(_run(model, values), _run(transformed, values)) + + def test_user_prefix_names_do_not_block_single_exp_scale_folding(self) -> None: + model = _model( + [ + onnx.helper.make_node("Add", ["x", "algebraic_exp_adjusted_user"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [ + _tensor("algebraic_exp_adjusted_user", np.asarray(0.5, dtype=np.float32)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[_info("biased", [1]), _info("exp", [1])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] + + def test_serial_exp_mul_boundary_survives_user_renaming_between_runs(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Mul", ["scaled", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [ + _tensor("scale_a", np.asarray(1.0e-30, dtype=np.float32)), + _tensor("scale_b", np.asarray(1.0e30, dtype=np.float32)), + ], + value_info=[_info("exp", [1]), _info("scaled", [1])], + ) + values = {"x": np.asarray([np.log(np.float32(1.0e-20))], dtype=np.float32)} + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + renamed_tensors: dict[str, str] = {} + for index, node in enumerate(transformed.graph.node): + node.name = f"user_node_{index}" + for output_index, output in enumerate(node.output): + if output.startswith("algebraic_"): + renamed_tensors[output] = f"user_tensor_{index}_{output_index}" + node.output[output_index] = renamed_tensors[output] + for input_index, input_name in enumerate(node.input): + if input_name.startswith("algebraic_"): + node.input[input_index] = renamed_tensors.setdefault( + input_name, + f"user_tensor_input_{index}_{input_index}", + ) + for initializer in transformed.graph.initializer: + if initializer.name.startswith("algebraic_"): + initializer.name = renamed_tensors.setdefault( + initializer.name, + f"user_initializer_{initializer.name}", + ) + for value in (*transformed.graph.value_info, *transformed.graph.output): + if value.name.startswith("algebraic_"): + value.name = renamed_tensors.setdefault(value.name, f"user_value_{value.name}") + renamed = optimize_onnx(transformed, exp_positive_scale_folding=True) + + assert [node.op_type for node in renamed.graph.node] == ["Exp", "Mul", "Mul"] + np.testing.assert_array_equal(_run(model, values), _run(renamed, values)) + + def test_public_optimizer_preserves_serial_exp_mul_boundary_through_view(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Reshape", ["scaled", "shape"], ["viewed"]), + onnx.helper.make_node("Mul", ["viewed", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [ + _tensor("shape", np.asarray([1], dtype=np.int64)), + _tensor("scale_a", np.asarray(1.0e-30, dtype=np.float32)), + _tensor("scale_b", np.asarray(1.0e30, dtype=np.float32)), + ], + value_info=[_info("exp", [1]), _info("scaled", [1]), _info("viewed", [1])], + ) + values = {"x": np.asarray([np.log(np.float32(1.0e-20))], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert sum(node.op_type == "Mul" for node in transformed.graph.node) == 2 + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_array_equal(_run(model, values), _run(transformed, values)) + + def test_public_optimizer_preserves_serial_exp_mul_boundary_through_flatten(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Flatten", ["scaled"], ["flattened"]), + onnx.helper.make_node("Mul", ["flattened", "scale_b"], ["y"]), + ], + [_info("x", [1, 1])], + [_info("flattened", [1, 1]), _info("y", [1, 1])], + [ + _tensor("scale_a", np.asarray(0.005727309, dtype=np.float32)), + _tensor("scale_b", np.asarray(3.3510647e28, dtype=np.float32)), + ], + value_info=[_info("exp", [1, 1]), _info("scaled", [1, 1])], + ) + values = {"x": np.asarray([[-96.411636]], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert sum(node.op_type == "Mul" for node in transformed.graph.node) == 2 + _assert_valid_with_inferred_shapes(transformed) + for expected, actual in zip(_run(model, values), _run(transformed, values), strict=True): + np.testing.assert_array_equal(actual, expected) + + def test_public_optimizer_preserves_serial_exp_mul_boundary_through_identity(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Identity", ["scaled"], ["passed"]), + onnx.helper.make_node("Mul", ["passed", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("passed", [1]), _info("y", [1])], + [ + _tensor("scale_a", np.asarray(0.005727309, dtype=np.float32)), + _tensor("scale_b", np.asarray(3.3510647e28, dtype=np.float32)), + ], + value_info=[_info("exp", [1]), _info("scaled", [1])], + ) + values = {"x": np.asarray([-96.411636], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert sum(node.op_type == "Mul" for node in transformed.graph.node) == 2 + _assert_valid_with_inferred_shapes(transformed) + for expected, actual in zip(_run(model, values), _run(transformed, values), strict=True): + np.testing.assert_array_equal(actual, expected) + + def test_public_optimizer_preserves_serial_exp_mul_boundary_when_intermediate_observed( + self, + ) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Mul", ["scaled", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("scaled", [1]), _info("y", [1])], + [ + _tensor("scale_a", np.asarray(0.005727309, dtype=np.float32)), + _tensor("scale_b", np.asarray(3.3510647e28, dtype=np.float32)), + ], + value_info=[_info("exp", [1])], + ) + values = {"x": np.asarray([-96.411636], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert sum(node.op_type == "Mul" for node in transformed.graph.node) == 2 + _assert_valid_with_inferred_shapes(transformed) + for expected, actual in zip(_run(model, values), _run(transformed, values), strict=True): + np.testing.assert_array_equal(actual, expected) + + def test_public_optimizer_preserves_serial_exp_mul_boundary_on_branch(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Mul", ["scaled", "scale_b"], ["y"]), + onnx.helper.make_node("Identity", ["scaled"], ["tap"]), + ], + [_info("x", [1])], + [_info("y", [1]), _info("tap", [1])], + [ + _tensor("scale_a", np.asarray(0.005727309, dtype=np.float32)), + _tensor("scale_b", np.asarray(3.3510647e28, dtype=np.float32)), + ], + value_info=[_info("exp", [1]), _info("scaled", [1])], + ) + values = {"x": np.asarray([-96.411636], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert sum(node.op_type == "Mul" for node in transformed.graph.node) == 2 + _assert_valid_with_inferred_shapes(transformed) + for expected, actual in zip(_run(model, values), _run(transformed, values), strict=True): + np.testing.assert_array_equal(actual, expected) + + def test_public_optimizer_preserves_serial_exp_mul_after_marked_bias_view(self) -> None: + model = _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Reshape", ["biased", "shape"], ["viewed"]), + onnx.helper.make_node("Exp", ["viewed"], ["exp"]), + onnx.helper.make_node("Mul", ["exp", "scale_a"], ["scaled"]), + onnx.helper.make_node("Mul", ["scaled", "scale_b"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [ + _tensor("bias", np.asarray(0.0, dtype=np.float32)), + _tensor("shape", np.asarray([1], dtype=np.int64)), + _tensor("scale_a", np.asarray(1.0e-30, dtype=np.float32)), + _tensor("scale_b", np.asarray(1.0e30, dtype=np.float32)), + ], + value_info=[ + _info("biased", [1]), + _info("viewed", [1]), + _info("exp", [1]), + _info("scaled", [1]), + ], + ) + values = {"x": np.asarray([np.log(np.float32(1.0e-20))], dtype=np.float32)} + + transformed = optimize_onnx(model, exp_positive_scale_folding=True) + + assert sum(node.op_type == "Mul" for node in transformed.graph.node) == 2 + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_array_equal(_run(model, values), _run(transformed, values)) + + def test_serial_exp_mul_stops_after_first_noncompact_scale_without_rebuilds( + self, + monkeypatch, + ) -> None: + scale_count = 6 + nodes = [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["exp"]), + ] + initializers = [_tensor("bias", np.asarray(-0.25, dtype=np.float32))] + value_info = [_info("biased", [1, 2]), _info("exp", [1, 2])] + current = "exp" + for index in range(scale_count): + output = "y" if index == scale_count - 1 else f"scaled{index}" + nodes.append(onnx.helper.make_node("Mul", [current, f"scale{index}"], [output])) + initializers.append(_tensor(f"scale{index}", np.asarray(1.0e30, dtype=np.float32))) + if output != "y": + value_info.append(_info(output, [1, 2])) + current = output + model = _model( + nodes, + [_info("x", [1, 2])], + [_info("y", [1, 2])], + initializers, + value_info=value_info, + ) + values = {"x": np.asarray([[0.5, -1.0]], dtype=np.float32)} + build_count = 0 + original_build = algebraic_pipe._GraphIndex.build + + def counted_build(model: onnx.ModelProto) -> algebraic_pipe._GraphIndex: + nonlocal build_count + build_count += 1 + return original_build(model) + + monkeypatch.setattr(algebraic_pipe._GraphIndex, "build", counted_build) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert build_count <= 6 + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp", *["Mul"] * 6] + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-6, + atol=2e-6, + ) + + def test_scalar_scale_without_bias_stays_compact(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [_info("x", [1, 2, 3])], + [_info("y", [1, 2, 3])], + [_tensor("scale", np.asarray(1.5, dtype=np.float32))], + value_info=[_info("exponential", [1, 2, 3])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] + log_scale_name = transformed.graph.node[0].input[1] + log_scale = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == log_scale_name) + ) + assert log_scale.shape == () + np.testing.assert_array_equal(log_scale, np.log(np.asarray(1.5, dtype=np.float32))) + + def test_post_view_scale_with_different_flattened_pattern_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node("Reshape", ["exponential", "output_shape"], ["reshaped"]), + onnx.helper.make_node("Mul", ["reshaped", "scale"], ["y"]), + ], + [_info("x", [2, 2, 3])], + [_info("y", [3, 2, 2])], + [ + _tensor("output_shape", np.asarray([3, 2, 2], dtype=np.int64)), + _tensor("scale", np.asarray([[[1.0], [2.0]]], dtype=np.float32)), + ], + value_info=[ + _info("exponential", [2, 2, 3]), + _info("reshaped", [3, 2, 2]), + ], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(model, transformed) + + def test_orthogonal_bias_and_scale_broadcast_uses_compact_log_scale_add(self) -> None: + rng = np.random.default_rng(36) + model = _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"], name="bias_add"), + onnx.helper.make_node("Exp", ["biased"], ["exponential"], name="exp"), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"], name="scale_mul"), + ], + [_info("x", [2, 2, 3])], + [_info("y", [2, 2, 3])], + [ + _tensor( + "bias", + rng.normal(size=(2, 1, 3)).astype(np.float32), + ), + _tensor("scale", np.asarray([[[1.25], [0.75]]], dtype=np.float32)), + ], + value_info=[ + _info("biased", [2, 2, 3]), + _info("exponential", [2, 2, 3]), + ], + ) + values = {"x": rng.normal(size=(2, 2, 3)).astype(np.float32)} + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Add", "Exp"] + assert transformed.graph.node[0].name == "bias_add" + assert transformed.graph.node[0].input[1] == "bias" + assert transformed.graph.node[-1].output[0] == "y" + log_scale_name = transformed.graph.node[1].input[1] + log_scale = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == log_scale_name) + ) + assert log_scale.shape == (1, 2, 1) + assert not any(tuple(value.dims) == (2, 2, 3) for value in transformed.graph.initializer) + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-6, + atol=2e-6, + ) + + def test_exp_scale_flag_discloses_relaxed_float_boundary_semantics(self) -> None: + description = get_all_capabilities()["exp-positive-scale-folding"].description + assert "relaxed floating-point overflow semantics" in description + + def test_float32_boundary_behavior_is_relaxed_when_enabled(self) -> None: + x = np.asarray([89.0], dtype=np.float32) + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [_tensor("scale", np.asarray(1.0e-10, dtype=np.float32))], + value_info=[_info("exponential", [1])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] + assert np.isinf(_run(model, {"x": x})[0][0]) + assert np.isfinite(_run(transformed, {"x": x})[0][0]) + + @pytest.mark.parametrize("opset_version", [None, 6]) + def test_legacy_or_missing_standard_opset_is_unchanged( + self, + opset_version: int | None, + ) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node( + "Mul", + ["exponential", "scale"], + ["y"], + broadcast=1, + ), + ], + [_info("x", [1, 2])], + [_info("y", [1, 2])], + [_tensor("scale", np.asarray([1.25, 0.75], dtype=np.float32))], + value_info=[_info("exponential", [1, 2])], + ) + del model.opset_import[:] + if opset_version is not None: + model.opset_import.append(onnx.helper.make_opsetid("", opset_version)) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert transformed.SerializeToString() == original + + @pytest.mark.parametrize( + "opset_imports", + [ + [onnx.helper.make_opsetid("ai.onnx", 17), onnx.helper.make_opsetid("", 6)], + [onnx.helper.make_opsetid("", 17), onnx.helper.make_opsetid("", 6)], + ], + ) + def test_ambiguous_standard_opset_is_unchanged( + self, + opset_imports: list[onnx.OperatorSetIdProto], + ) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node( + "Mul", + ["exponential", "scale"], + ["y"], + broadcast=1, + ), + ], + [_info("x", [1, 2])], + [_info("y", [1, 2])], + [_tensor("scale", np.asarray([1.25, 0.75], dtype=np.float32))], + value_info=[_info("exponential", [1, 2])], + ) + del model.opset_import[:] + model.opset_import.extend(opset_imports) + original = model.SerializeToString() + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert transformed.SerializeToString() == original + + def test_huge_shape_product_overflow_is_unchanged(self) -> None: + huge_dimension = 5_270_498_306_774_157_605 + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node("Reshape", ["exponential", "output_shape"], ["reshaped"]), + onnx.helper.make_node( + "Mul", + ["reshaped", "scale"], + ["y"], + ), + ], + [_info("x", [7, huge_dimension])], + [_info("y", [3])], + [ + _tensor("output_shape", np.asarray([3], dtype=np.int64)), + _tensor("scale", np.asarray([1.25, 0.75, 1.5], dtype=np.float32)), + ], + value_info=[ + _info("exponential", [7, huge_dimension]), + _info("reshaped", [3]), + ], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(model, transformed) + + def test_simple_numeric_example_matches_log_domain_identity(self) -> None: + x = np.asarray([[0.0, 2.0]], dtype=np.float32) + bias = np.asarray([[1.0, -1.0]], dtype=np.float32) + log_scale = np.asarray([[2.0, 0.5]], dtype=np.float32) + scale = np.exp(log_scale).astype(np.float32) + expected_folded_bias = np.asarray([[3.0, -0.5]], dtype=np.float32) + expected_output = np.exp(x + expected_folded_bias).astype(np.float32) + model = _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [_info("x", [1, 2])], + [_info("y", [1, 2])], + [_tensor("bias", bias), _tensor("scale", scale)], + value_info=[_info("biased", [1, 2]), _info("exponential", [1, 2])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] + combined_name = transformed.graph.node[0].input[1] + combined = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == combined_name) + ) + np.testing.assert_allclose(combined, expected_folded_bias, rtol=0, atol=2e-7) + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose(_run(model, {"x": x}), [expected_output], rtol=2e-6, atol=2e-6) + np.testing.assert_allclose( + _run(transformed, {"x": x}), + [expected_output], + rtol=2e-6, + atol=2e-6, + ) + + @pytest.fixture + def exp_scale_model(self) -> onnx.ModelProto: + tensor_shape = [1, 2, 1, 2, 2] + return _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Reshape", ["biased", "flat_shape"], ["flat"]), + onnx.helper.make_node("Exp", ["flat"], ["exponential"]), + onnx.helper.make_node( + "Reshape", + ["exponential", "tensor_shape"], + ["restored"], + ), + onnx.helper.make_node("Mul", ["restored", "scale"], ["y"]), + ], + [_info("x", tensor_shape)], + [_info("y", tensor_shape)], + [ + _tensor("bias", np.asarray(-2.0, dtype=np.float32)), + _tensor("flat_shape", np.asarray([1, 8], dtype=np.int64)), + _tensor("tensor_shape", np.asarray(tensor_shape, dtype=np.int64)), + _tensor( + "scale", + np.asarray([[[[[1.25, 1.5], [2.0, 0.75]]]]], dtype=np.float32), + ), + ], + value_info=[ + _info("biased", tensor_shape), + _info("flat", [1, 8]), + _info("exponential", [1, 8]), + _info("restored", tensor_shape), + ], + ) + + def test_broadcast_scale_folds_through_round_trip_reshapes( + self, + exp_scale_model: onnx.ModelProto, + ) -> None: + rng = np.random.default_rng(30) + model = exp_scale_model + tensor_shape = [1, 2, 1, 2, 2] + values = {"x": rng.normal(size=tensor_shape).astype(np.float32)} + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == [ + "Add", + "Reshape", + "Exp", + "Reshape", + ] + assert transformed.graph.node[-1].output[0] == "y" + combined_name = transformed.graph.node[0].input[1] + combined = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == combined_name) + ) + expected = np.asarray(-2.0 + np.log(onnx.numpy_helper.to_array(model.graph.initializer[3]))) + assert combined.shape == (1, 1, 1, 2, 2) + np.testing.assert_array_equal(combined, expected) + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-6, + atol=2e-6, + ) + + def test_squeeze_and_unsqueeze_views_fold(self) -> None: + rng = np.random.default_rng(31) + model = _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Unsqueeze", ["biased", "axes"], ["expanded"]), + onnx.helper.make_node("Exp", ["expanded"], ["exponential"]), + onnx.helper.make_node("Squeeze", ["exponential", "axes"], ["restored"]), + onnx.helper.make_node("Mul", ["restored", "scale"], ["y"]), + ], + [_info("x", [1, 2, 2])], + [_info("y", [1, 2, 2])], + [ + _tensor("bias", np.asarray(-1.0, dtype=np.float32)), + _tensor("axes", np.asarray([1], dtype=np.int64)), + _tensor("scale", np.asarray([[[1.0, 1.25], [1.5, 2.0]]], dtype=np.float32)), + ], + value_info=[ + _info("biased", [1, 2, 2]), + _info("expanded", [1, 1, 2, 2]), + _info("exponential", [1, 1, 2, 2]), + _info("restored", [1, 2, 2]), ], - value_info=[_info("conv1", [1, 1, 2, 2]), _info("conv2", [1, 1, 2, 2])], ) - values = { - "x1": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), - "x2": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), - } + values = {"x": rng.normal(size=(1, 2, 2)).astype(np.float32)} + transformed = AlgebraicRewritePipe().process( model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), ) - assert [node.op_type for node in transformed.graph.node] == ["Conv", "Conv"] - for original, rewritten in zip( + + assert not any(node.op_type == "Mul" for node in transformed.graph.node) + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose( _run(model, values), _run(transformed, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + rtol=2e-6, + atol=2e-6, + ) - def test_nested_subgraph_captures_make_affine_fold_ineligible( - self, - affine_model: tuple[onnx.ModelProto, dict[str, np.ndarray]], - ) -> None: - model, values = affine_model - branch_shape = [1, 3, 3, 3] - then_branch = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["mul_out"], ["then_output"])], - "then_branch", - [], - [_info("then_output", branch_shape)], + def test_scale_without_existing_bias_becomes_pre_exp_add(self) -> None: + rng = np.random.default_rng(32) + scale = np.asarray([[[[1.0, 1.25], [1.5, 2.0]]]], dtype=np.float32) + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node( + "Reshape", + ["exponential", "output_shape"], + ["restored"], + ), + onnx.helper.make_node("Mul", ["restored", "scale"], ["y"]), + ], + [_info("x", [1, 4])], + [_info("y", [1, 1, 2, 2])], + [ + _tensor("output_shape", np.asarray([1, 1, 2, 2], dtype=np.int64)), + _tensor("scale", scale), + ], + value_info=[ + _info("exponential", [1, 4]), + _info("restored", [1, 1, 2, 2]), + ], ) - else_branch = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["conv_out"], ["else_output"])], - "else_branch", - [], - [_info("else_output", branch_shape)], + values = {"x": rng.normal(size=(1, 4)).astype(np.float32)} + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), ) - model.graph.initializer.append(_tensor("condition", np.asarray(True, dtype=np.bool_))) - model.graph.node.append( - onnx.helper.make_node( - "If", - ["condition"], - ["captured"], - then_branch=then_branch, - else_branch=else_branch, - ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp", "Reshape"] + assert transformed.graph.node[-1].output[0] == "y" + log_scale_name = transformed.graph.node[0].input[1] + log_scale = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == log_scale_name) + ) + np.testing.assert_array_equal(log_scale, np.log(scale).reshape(1, 4)) + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-6, + atol=2e-6, ) - model.graph.output.append(_info("captured", branch_shape)) + + def test_runtime_bias_keeps_original_add_and_replaces_mul(self) -> None: + rng = np.random.default_rng(34) + model = _model( + [ + onnx.helper.make_node( + "Add", + ["x", "runtime_bias"], + ["biased"], + name="runtime_bias_add", + ), + onnx.helper.make_node("Exp", ["biased"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [_info("x", [1, 4]), _info("runtime_bias", [1, 4])], + [_info("y", [1, 4])], + [_tensor("scale", np.asarray([1.0, 1.25, 1.5, 2.0], dtype=np.float32))], + value_info=[_info("biased", [1, 4]), _info("exponential", [1, 4])], + ) + values = { + "x": rng.normal(size=(1, 4)).astype(np.float32), + "runtime_bias": rng.normal(size=(1, 4)).astype(np.float32), + } transformed = AlgebraicRewritePipe().process( model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), ) - assert [node.op_type for node in transformed.graph.node] == [ - "Conv", - "Mul", - "Add", - "If", - ] + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Add", "Exp"] + assert transformed.graph.node[0].name == "runtime_bias_add" + assert list(transformed.graph.node[0].input) == ["x", "runtime_bias"] _assert_valid_with_inferred_shapes(transformed) - for original, rewritten in zip( + np.testing.assert_allclose( _run(model, values), _run(transformed, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=0, atol=0) + rtol=2e-6, + atol=2e-6, + ) - def test_constant_attribute_affine_is_folded_and_pruned(self) -> None: - rng = np.random.default_rng(18) + def test_constant_attribute_operands_preserve_float32_dtype(self) -> None: + rng = np.random.default_rng(33) model = _model( [ - onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), + onnx.helper.make_node("Constant", [], ["bias"], value_float=-1.0), onnx.helper.make_node("Constant", [], ["scale"], value_float=1.5), - onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), ], - [_info("x", [1, 1, 2, 2])], - [_info("y", [1, 1, 2, 2])], - [_tensor("weight", rng.normal(size=(1, 1, 1, 1)).astype(np.float32))], - value_info=[_info("conv_out", [1, 1, 2, 2])], + [_info("x", [1, 4])], + [_info("y", [1, 4])], + [], + value_info=[_info("biased", [1, 4]), _info("exponential", [1, 4])], ) - values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} + values = {"x": rng.normal(size=(1, 4)).astype(np.float32)} + transformed = AlgebraicRewritePipe().process( model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), ) - assert [node.op_type for node in transformed.graph.node] == ["Conv"] + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] + combined_name = transformed.graph.node[0].input[1] + combined = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == combined_name) + ) + assert combined.dtype == np.float32 + _assert_valid_with_inferred_shapes(transformed) np.testing.assert_allclose( _run(model, values), _run(transformed, values), - rtol=2e-5, - atol=2e-5, + rtol=2e-6, + atol=2e-6, + ) + + def test_direct_float64_chain_preserves_initializer_precision(self) -> None: + shape = [1, 2] + + def double_info(name: str) -> onnx.ValueInfoProto: + return onnx.helper.make_tensor_value_info( + name, + onnx.TensorProto.DOUBLE, + shape, + ) + + bias = np.asarray(1.0 + 2**-30, dtype=np.float64) + scale = np.asarray([[1.0 + 2**-29, 1.0 + 2**-28]], dtype=np.float64) + model = _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [double_info("x")], + [double_info("y")], + [_tensor("bias", bias), _tensor("scale", scale)], + value_info=[double_info("biased"), double_info("exponential")], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + combined_name = transformed.graph.node[0].input[1] + combined = onnx.numpy_helper.to_array( + next(value for value in transformed.graph.initializer if value.name == combined_name) + ) + assert combined.dtype == np.float64 + np.testing.assert_array_equal(combined, bias + np.log(scale)) + _assert_valid_with_inferred_shapes(transformed) + + def test_invalid_scale_broadcast_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + ) -> None: + scale = next(value for value in exp_scale_model.graph.initializer if value.name == "scale") + scale.CopyFrom(_tensor("scale", np.ones((1, 3, 1, 2, 2), dtype=np.float32))) + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + @pytest.mark.parametrize("scale_value", [0.0, -1.0, np.nan, np.inf]) + def test_nonpositive_or_nonfinite_scale_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + scale_value: float, + ) -> None: + scale = next(value for value in exp_scale_model.graph.initializer if value.name == "scale") + scale.CopyFrom(_tensor("scale", np.full((1, 1, 1, 2, 2), scale_value, np.float32))) + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + @pytest.mark.parametrize( + "constant_name", + ["flat_shape", "tensor_shape", "scale"], + ) + def test_overridable_constants_are_unchanged( + self, + exp_scale_model: onnx.ModelProto, + constant_name: str, + ) -> None: + initializer = next( + value for value in exp_scale_model.graph.initializer if value.name == constant_name + ) + exp_scale_model.graph.input.append( + onnx.helper.make_tensor_value_info( + constant_name, + initializer.data_type, + list(initializer.dims), + ) + ) + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + def test_unloaded_external_scale_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + ) -> None: + scale = next(value for value in exp_scale_model.graph.initializer if value.name == "scale") + scale.ClearField("raw_data") + scale.data_location = onnx.TensorProto.EXTERNAL + location = scale.external_data.add() + location.key = "location" + location.value = "missing-scale.bin" + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + def test_unloaded_external_constant_value_is_unchanged(self) -> None: + external_scale = _tensor("external_scale_payload", np.asarray(1.5, dtype=np.float32)) + external_scale.ClearField("raw_data") + external_scale.data_location = onnx.TensorProto.EXTERNAL + location = external_scale.external_data.add() + location.key = "location" + location.value = "missing-constant-scale.bin" + model = _model( + [ + onnx.helper.make_node("Exp", ["x"], ["exponential"]), + onnx.helper.make_node("Constant", [], ["scale"], value=external_scale), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [_info("x", [1])], + [_info("y", [1])], + [], + value_info=[_info("exponential", [1])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(model, transformed) + + @pytest.mark.parametrize("duplicate_name", ["biased", "exponential", "y"]) + def test_duplicate_tensor_definition_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + duplicate_name: str, + ) -> None: + exp_scale_model.graph.node.append( + onnx.helper.make_node("Identity", ["x"], [duplicate_name]) + ) + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + @pytest.mark.parametrize("collision", ["initializer", "graph_input", "initializer_copy"]) + def test_cross_kind_definition_collision_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + collision: str, + ) -> None: + if collision == "initializer": + exp_scale_model.graph.node.append(onnx.helper.make_node("Identity", ["x"], ["scale"])) + elif collision == "graph_input": + exp_scale_model.graph.node.append(onnx.helper.make_node("Identity", ["x"], ["x"])) + else: + duplicate = onnx.TensorProto() + duplicate.CopyFrom( + next(value for value in exp_scale_model.graph.initializer if value.name == "scale") + ) + exp_scale_model.graph.initializer.append(duplicate) + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + def test_malformed_post_exp_cycle_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["route_0"]), + onnx.helper.make_node("Reshape", ["route_0", "shape"], ["route_1"]), + onnx.helper.make_node("Reshape", ["route_1", "shape"], ["route_0"]), + onnx.helper.make_node("Mul", ["route_1", "scale"], ["y"]), + ], + [_info("x", [1, 4])], + [_info("y", [1, 4])], + [ + _tensor("bias", np.asarray(-1.0, dtype=np.float32)), + _tensor("shape", np.asarray([1, 4], dtype=np.int64)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=[ + _info("biased", [1, 4]), + _info("route_0", [1, 4]), + _info("route_1", [1, 4]), + ], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(model, transformed) + + def test_malformed_exp_mul_cycle_is_unchanged(self) -> None: + model = _model( + [ + onnx.helper.make_node("Exp", ["y"], ["exponential"]), + onnx.helper.make_node("Mul", ["exponential", "scale"], ["y"]), + ], + [], + [_info("y", [1, 4])], + [_tensor("scale", np.asarray(1.5, dtype=np.float32))], + value_info=[_info("exponential", [1, 4])], + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(model, transformed) + + @pytest.mark.parametrize("domain", ["ai.onnx", "com.example"]) + @pytest.mark.parametrize("node_index", [2, 3, 4]) + def test_non_empty_domain_interpreted_nodes_are_unchanged( + self, + exp_scale_model: onnx.ModelProto, + node_index: int, + domain: str, + ) -> None: + exp_scale_model.graph.node[node_index].domain = domain + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + @pytest.mark.parametrize("protection", ["graph_output", "shared", "captured"]) + def test_observed_intermediate_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + protection: str, + ) -> None: + if protection == "graph_output": + exp_scale_model.graph.output.append(_info("restored", [1, 2, 1, 2, 2])) + elif protection == "shared": + exp_scale_model.graph.node.append( + onnx.helper.make_node("Identity", ["restored"], ["observed"]) + ) + exp_scale_model.graph.output.append(_info("observed", [1, 2, 1, 2, 2])) + else: + branch = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["restored"], ["branch_output"])], + "capturing_branch", + [], + [_info("branch_output", [1, 2, 1, 2, 2])], + ) + exp_scale_model.graph.initializer.append( + _tensor("condition", np.asarray(True, dtype=np.bool_)) + ) + exp_scale_model.graph.node.append( + onnx.helper.make_node( + "If", + ["condition"], + ["observed"], + then_branch=branch, + else_branch=branch, + ) + ) + exp_scale_model.graph.output.append(_info("observed", [1, 2, 1, 2, 2])) + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + def test_shape_domain_mismatch_is_unchanged( + self, + exp_scale_model: onnx.ModelProto, + ) -> None: + restored = next( + value for value in exp_scale_model.graph.value_info if value.name == "restored" + ) + restored.type.tensor_type.shape.dim[1].ClearField("dim_value") + restored.type.tensor_type.shape.dim[1].dim_param = "channels" + + transformed = AlgebraicRewritePipe().process( + exp_scale_model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + _assert_byte_identical(exp_scale_model, transformed) + + def test_deep_view_route_is_unchanged(self) -> None: + route_depth = 65 + nodes = [ + onnx.helper.make_node("Add", ["x", "bias"], ["biased"]), + onnx.helper.make_node("Exp", ["biased"], ["route_0"]), + ] + value_info = [_info("biased", [1, 4]), _info("route_0", [1, 4])] + for route_index in range(route_depth): + nodes.append( + onnx.helper.make_node( + "Reshape", + [f"route_{route_index}", "shape"], + [f"route_{route_index + 1}"], + ) + ) + value_info.append(_info(f"route_{route_index + 1}", [1, 4])) + nodes.append(onnx.helper.make_node("Mul", [f"route_{route_depth}", "scale"], ["y"])) + model = _model( + nodes, + [_info("x", [1, 4])], + [_info("y", [1, 4])], + [ + _tensor("bias", np.asarray(-1.0, dtype=np.float32)), + _tensor("shape", np.asarray([1, 4], dtype=np.int64)), + _tensor("scale", np.asarray(1.5, dtype=np.float32)), + ], + value_info=value_info, + ) + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), ) + + _assert_byte_identical(model, transformed) + + def test_public_optimize_path_is_idempotent( + self, + exp_scale_model: onnx.ModelProto, + ) -> None: + transformed = optimize_onnx(exp_scale_model, exp_positive_scale_folding=True) + second = optimize_onnx(transformed, exp_positive_scale_folding=True) + + assert transformed.SerializeToString() == second.SerializeToString() + assert not any(node.op_type == "Mul" for node in transformed.graph.node) + _assert_valid_with_inferred_shapes(second) diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index 82cb1a512..c061b4f99 100644 --- a/tests/unit/optim/test_analysis.py +++ b/tests/unit/optim/test_analysis.py @@ -16,6 +16,7 @@ import subprocess import sys from array import array +from typing import ClassVar from unittest.mock import MagicMock import numpy as np @@ -83,6 +84,38 @@ def _benign_model() -> ModelProto: return _finalize(helper.make_graph([node], "benign", [x], [z], initializer=[small])) +def _sibling_slice_model() -> ModelProto: + """Two contiguous sibling Slices that can be replaced by one Split.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 6, 2]) + left_out = helper.make_tensor_value_info("left_out", TensorProto.FLOAT, [1, 2, 2]) + right_out = helper.make_tensor_value_info("right_out", TensorProto.FLOAT, [1, 4, 2]) + left = helper.make_tensor_value_info("left", TensorProto.FLOAT, [1, 2, 2]) + right = helper.make_tensor_value_info("right", TensorProto.FLOAT, [1, 4, 2]) + nodes = [ + helper.make_node("Slice", ["x", "left_starts", "left_ends", "axis", "steps"], ["left"]), + helper.make_node("Slice", ["x", "right_starts", "right_ends", "axis", "steps"], ["right"]), + helper.make_node("Relu", ["left"], ["left_out"]), + helper.make_node("Relu", ["right"], ["right_out"]), + ] + initializers = [ + numpy_helper.from_array(np.asarray([0], dtype=np.int64), "left_starts"), + numpy_helper.from_array(np.asarray([2], dtype=np.int64), "left_ends"), + numpy_helper.from_array(np.asarray([2], dtype=np.int64), "right_starts"), + numpy_helper.from_array(np.asarray([6], dtype=np.int64), "right_ends"), + numpy_helper.from_array(np.asarray([1], dtype=np.int64), "axis"), + numpy_helper.from_array(np.asarray([1], dtype=np.int64), "steps"), + ] + graph = helper.make_graph( + nodes, + "sibling_slice", + [x], + [left_out, right_out], + initializer=initializers, + value_info=[left, right], + ) + return _finalize(graph) + + # ============================================================================= # NODE / INITIALIZER DIFF HELPERS # ============================================================================= @@ -298,9 +331,7 @@ def test_signed_zero_change_with_pure_python_protobuf(self) -> None: env=env, ) - def test_upb_equality_does_not_copy_float_fields( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_upb_equality_does_not_copy_float_fields(self, monkeypatch: pytest.MonkeyPatch) -> None: base_init = TensorProto( name="W", data_type=TensorProto.FLOAT, @@ -310,9 +341,7 @@ def test_upb_equality_does_not_copy_float_fields( probe_init = TensorProto() probe_init.CopyFrom(base_init) - monkeypatch.setattr( - "winml.modelkit.optim.analysis._PROTOBUF_IMPLEMENTATION", "upb" - ) + monkeypatch.setattr("winml.modelkit.optim.analysis._PROTOBUF_IMPLEMENTATION", "upb") def fail_array(*_args, **_kwargs): raise AssertionError("upb equality fast path copied typed fields") @@ -321,9 +350,7 @@ def fail_array(*_args, **_kwargs): assert _initializers_equal(base_init, probe_init) - def test_cpp_equality_still_compares_float_bits( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_cpp_equality_still_compares_float_bits(self, monkeypatch: pytest.MonkeyPatch) -> None: base_init = TensorProto( name="W", data_type=TensorProto.FLOAT, @@ -335,9 +362,7 @@ def test_cpp_equality_still_compares_float_bits( array_calls = 0 real_array = array - monkeypatch.setattr( - "winml.modelkit.optim.analysis._PROTOBUF_IMPLEMENTATION", "cpp" - ) + monkeypatch.setattr("winml.modelkit.optim.analysis._PROTOBUF_IMPLEMENTATION", "cpp") def track_array(*args, **kwargs): nonlocal array_calls @@ -694,6 +719,116 @@ def test_findings_match_analyze_model(self) -> None: analyze_names = {finding.name for finding in analyze_model(_matmul_add_model(), caps)} assert iter_names == analyze_names + def test_reports_algebraic_sibling_slice_to_split(self) -> None: + pairs = list(iter_optimization_outputs(_sibling_slice_model(), get_all_capabilities())) + matches = [ + (finding, produced) + for finding, produced in pairs + if finding.name == "gather-slice-to-split-fusion" + ] + + assert len(matches) == 1 + finding, produced = matches[0] + + assert finding.enable_flag == "--enable-gather-slice-to-split-fusion" + assert finding.pipe_name == "ort_graph+algebraic_rewrite" + assert any(ref.op_type == "Slice" for ref in finding.removed_nodes) + assert any(ref.op_type == "Split" for ref in finding.added_nodes) + assert [node.op_type for node in produced.graph.node] == ["Split", "Relu", "Relu"] + + def test_shared_capability_probe_does_not_advance_pipeline_cursor( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + shared_cap = BoolCapability( + name="shared-route", + ort_name=None, + description="Shared route", + category=CapabilityCategory.MISC, + ) + observer_cap = BoolCapability( + name="observer-probe", + ort_name=None, + description="Observer probe", + category=CapabilityCategory.MISC, + ) + cap_registry = { + shared_cap.name: shared_cap, + observer_cap.name: observer_cap, + } + + def append_identity(model: ModelProto, prefix: str) -> None: + suffix = sum( + output.startswith(prefix) for node in model.graph.node for output in node.output + ) + model.graph.node.append(helper.make_node("Identity", ["z"], [f"{prefix}{suffix}"])) + + class FirstSharedPipe: + name = "first-shared" + capabilities: ClassVar[dict[str, BoolCapability]] = {shared_cap.name: shared_cap} + + @classmethod + def build_config(cls, **kwargs): + return kwargs + + @staticmethod + def process(model, config): + if config.get("shared_route"): + append_identity(model, "first_shared_") + return model + + def prepare_analysis_model(self, model): + return model + + def process_analysis(self, model, config): + return self.process(model, config) + + @classmethod + def requires_analysis_clone(cls): + return True + + def finish_analysis(self): + pass + + class SecondSharedPipe(FirstSharedPipe): + name = "second-shared" + + @staticmethod + def process(model, config): + append_identity(model, "default_marker_") + if config.get("shared_route"): + append_identity(model, "second_shared_") + return model + + class ObserverPipe(FirstSharedPipe): + name = "observer" + capabilities: ClassVar[dict[str, BoolCapability]] = {observer_cap.name: observer_cap} + + @staticmethod + def process(model, config): + if config.get("observer_probe"): + append_identity(model, "observer_") + return model + + monkeypatch.setattr( + "winml.modelkit.optim.pipes.PIPES", + [FirstSharedPipe, SecondSharedPipe, ObserverPipe], + ) + monkeypatch.setattr("winml.modelkit.onnx.infer_shapes", lambda model: model) + + pairs = list(iter_optimization_outputs(_benign_model(), cap_registry)) + observer_produced = next( + produced for finding, produced in pairs if finding.name == "observer-probe" + ) + + default_markers = [ + output + for node in observer_produced.graph.node + for output in node.output + if output.startswith("default_marker_") + ] + assert default_markers == ["default_marker_0"] + def test_closing_iterator_cleans_up_prepared_pipe( self, monkeypatch: pytest.MonkeyPatch ) -> None: