From 4682ffd559ca11f65257f6f6b1d5ee6d41ea0d65 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 11 Aug 2026 08:47:35 +0800 Subject: [PATCH 01/15] feat(optim): extend channel affine route folding --- src/winml/modelkit/optim/pipes/algebraic.py | 314 +++--- tests/unit/optim/pipes/test_pipe_algebraic.py | 982 ++++++++++++++++++ 2 files changed, 1178 insertions(+), 118 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 29aac986e..1811b52b7 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -20,6 +20,7 @@ algebraic.STATIC_SPLIT_TO_SLICE, algebraic.CONV_CHANNEL_AFFINE_FOLDING, ) +MAX_AFFINE_ROUTE_DEPTH = 64 @dataclass @@ -38,6 +39,7 @@ class _GraphIndex: 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 @@ -73,6 +75,7 @@ def build(cls, model: onnx.ModelProto) -> _GraphIndex: 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}, ) @@ -140,16 +143,20 @@ 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 in ("", "ai.onnx") + + 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)) 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: @@ -319,7 +326,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) @@ -360,20 +376,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 @@ -458,6 +480,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"} @@ -477,11 +500,11 @@ def _channel_preserving_view_output( 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) @@ -498,8 +521,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,13 +535,16 @@ 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 @@ -527,7 +557,7 @@ def _collect_affine_chain( 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 +566,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 +612,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 +645,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 +658,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 +671,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 [] + 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 + candidates.extend(nested_candidates) + return 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) @@ -625,27 +731,31 @@ def _collect_routed_affine_candidates( 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, ) + if routed_candidates is None: + return None + candidates.extend(routed_candidates) return 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 len(conv.input) < 2 or conv.input[1] in index.graph_inputs: return False weight = next( ( @@ -664,6 +774,8 @@ def _copy_conv_parameters( return False if len(conv.input) > 2 and conv.input[2]: + if conv.input[2] in index.graph_inputs: + return False bias = next( ( initializer @@ -712,14 +824,15 @@ 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: @@ -733,72 +846,32 @@ def _fold_channel_affine( 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 +893,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} @@ -846,7 +919,12 @@ def _rewrite_static_splits( 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 diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index e9a977838..41f332d17 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -69,6 +69,22 @@ 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.""" @@ -194,6 +210,91 @@ 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"), + [("ai.onnx", True), ("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 + + @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( @@ -375,6 +476,137 @@ def test_shared_conv_output_is_ineligible( "Identity", ] + @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", "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) + 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( + "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( @@ -416,6 +648,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,6 +833,376 @@ def test_static_split_branches_fold_without_overlapping_ranges(self) -> None: atol=2e-5, ) + def test_nested_static_channel_splits_fold_affine_leaves(self) -> None: + rng = np.random.default_rng(22) + model = _model( + [ + 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])], + [ + _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, 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( + _run(model, values), + _run(transformed, values), + strict=True, + ): + np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) + + @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( + [ + 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]) @@ -679,3 +1419,245 @@ def test_constant_attribute_affine_is_folded_and_pruned(self) -> None: rtol=2e-5, atol=2e-5, ) + + @pytest.mark.parametrize( + ("domain", "should_fold"), + [("ai.onnx", True), ("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) From cc6947eec9b7cd6aeb986aae28735cc475f9643f Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 11 Aug 2026 09:36:39 +0800 Subject: [PATCH 02/15] feat(optim): fold positive Exp scales --- .../modelkit/optim/capabilities/algebraic.py | 10 + src/winml/modelkit/optim/pipes/algebraic.py | 370 +++++++++- tests/unit/optim/pipes/test_pipe_algebraic.py | 677 +++++++++++++++++- 3 files changed, 1052 insertions(+), 5 deletions(-) diff --git a/src/winml/modelkit/optim/capabilities/algebraic.py b/src/winml/modelkit/optim/capabilities/algebraic.py index a6ac37655..f3a6fd3a4 100644 --- a/src/winml/modelkit/optim/capabilities/algebraic.py +++ b/src/winml/modelkit/optim/capabilities/algebraic.py @@ -30,3 +30,13 @@ 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" + ), + category=CapabilityCategory.REWRITE, + default=False, +) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 1811b52b7..0e50d74b1 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -19,6 +19,7 @@ ALGEBRAIC_CAPABILITIES: dict[str, Any] = caps_dict( algebraic.STATIC_SPLIT_TO_SLICE, algebraic.CONV_CHANNEL_AFFINE_FOLDING, + algebraic.EXP_POSITIVE_SCALE_FOLDING, ) MAX_AFFINE_ROUTE_DEPTH = 64 @@ -29,6 +30,7 @@ class AlgebraicRewritePipeConfig(PipeConfig): static_split_to_slice: bool = False conv_channel_affine_folding: bool = False + exp_positive_scale_folding: bool = False @dataclass @@ -36,6 +38,8 @@ 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, ...]] @@ -46,10 +50,24 @@ class _GraphIndex: 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: @@ -61,6 +79,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): @@ -72,6 +112,8 @@ 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, @@ -94,6 +136,27 @@ class _AffineCandidate: offset: np.ndarray +@dataclass +class _ExpScaleCandidate: + """A positive post-Exp scale that can be merged into an existing bias.""" + + add: onnx.NodeProto + bias_input_index: int + output_node: onnx.NodeProto + mul: onnx.NodeProto + combined_bias: np.ndarray + + +@dataclass +class _ExpScaleInsertCandidate: + """A positive post-Exp scale that requires a new input Add.""" + + exp: onnx.NodeProto + output_node: onnx.NodeProto + mul: onnx.NodeProto + log_scale: np.ndarray + + class _NameAllocator: """Allocate names without relying on optional or duplicated node names.""" @@ -153,7 +216,12 @@ def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: 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 not _is_standard_onnx_node(producer) or producer.op_type != "Constant": @@ -164,10 +232,15 @@ def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: return np.asarray(onnx.numpy_helper.to_array(value)) except (TypeError, ValueError): 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 @@ -513,6 +586,285 @@ 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 {"Reshape", "Squeeze", "Unsqueeze"} + ): + return None + 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 np.prod(input_shape, dtype=np.int64) != np.prod(output_shape, dtype=np.int64) + ): + return 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 _post_exp_scale( + index: _GraphIndex, + exp: onnx.NodeProto, + visited: set[str] | None = None, +) -> tuple[onnx.NodeProto, 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 + scale_operand = _constant_input(index, next_node, current_name) + mul_output = _node_output(next_node) + output_shape = _static_shape(index, current_name) + if scale_operand is None or mul_output is None or output_shape is None: + 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) + 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] + ): + return None + post_exp = _post_exp_scale(index, next_node, visited) + if post_exp is None: + return None + output_node, mul, 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) + combined_bias = np.asarray(bias + np.log(scale), dtype=bias.dtype) + np.broadcast_to(combined_bias, add_shape) + except ValueError: + return None + if not np.isfinite(combined_bias).all(): + return None + return _ExpScaleCandidate( + add=add, + bias_input_index=bias_operand[0], + output_node=output_node, + mul=mul, + combined_bias=combined_bias, + ) + + +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, mul, 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 np.prod(input_shape, dtype=np.int64) != np.prod(output_shape, dtype=np.int64) + ): + return None + broadcast_scale = np.asarray(np.broadcast_to(scale, output_shape)) + log_scale = np.asarray(np.log(broadcast_scale).reshape(input_shape), dtype=scale.dtype) + if not np.isfinite(log_scale).all(): + return None + return _ExpScaleInsertCandidate( + exp=exp, + output_node=output_node, + mul=mul, + log_scale=log_scale, + ) + + +def _fold_exp_positive_scales( + model: onnx.ModelProto, + allocator: _NameAllocator, +) -> None: + """Fold eligible positive post-Exp constants into existing input biases.""" + index = _GraphIndex.build(model) + for add in list(model.graph.node): + candidate = _exp_scale_candidate(index, add) + if candidate is None: + continue + combined_name = _new_initializer( + model, + allocator, + candidate.combined_bias, + "algebraic_exp_log_bias", + ) + candidate.add.input[candidate.bias_input_index] = combined_name + candidate.output_node.output[0] = candidate.mul.output[0] + _remove_nodes(model, {id(candidate.mul)}) + index = _GraphIndex.build(model) + + for exp in list(model.graph.node): + candidate = _exp_scale_insert_candidate(index, exp) + if candidate is None: + continue + log_scale_name = _new_initializer( + model, + allocator, + candidate.log_scale, + "algebraic_exp_log_scale", + ) + adjusted_name = allocator.new("algebraic_exp_adjusted") + add = onnx.helper.make_node( + "Add", + [candidate.exp.input[0], log_scale_name], + [adjusted_name], + name=allocator.new("algebraic_exp_log_add"), + ) + candidate.exp.input[0] = adjusted_name + candidate.output_node.output[0] = candidate.mul.output[0] + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + if node is candidate.exp: + rewritten.append(add) + if node is not candidate.mul: + rewritten.append(node) + del model.graph.node[:] + model.graph.node.extend(rewritten) + index = _GraphIndex.build(model) + + def _collect_affine_chain( index: _GraphIndex, first: onnx.NodeProto, @@ -977,12 +1329,17 @@ def build_config(cls, **kwargs: Any) -> AlgebraicRewritePipeConfig: return AlgebraicRewritePipeConfig( static_split_to_slice=kwargs.get("static_split_to_slice", 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.conv_channel_affine_folding + or config.exp_positive_scale_folding + ) def process( self, @@ -995,10 +1352,15 @@ def process( result = onnx.ModelProto() result.CopyFrom(model) + 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: _fold_channel_affine(result, allocator) + if config.exp_positive_scale_folding: + _fold_exp_positive_scales(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 41f332d17..07928fc32 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: from collections.abc import Sequence + from pathlib import Path def _tensor(name: str, values: np.ndarray) -> onnx.TensorProto: @@ -90,7 +91,11 @@ class TestAlgebraicRegistration: 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", + "conv-channel-affine-folding", + "exp-positive-scale-folding", + } assert names <= capabilities.keys() assert all(capabilities[name].default is False for name in names) assert all( @@ -101,15 +106,19 @@ def test_capabilities_are_opt_in_and_independent(self) -> None: config = AlgebraicRewritePipe.build_config( static_split_to_slice=True, conv_channel_affine_folding=False, + exp_positive_scale_folding=True, ) assert config.static_split_to_slice 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-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] @@ -118,6 +127,115 @@ 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 TestStaticSplitToSlice: """Test static Split replacement using generated data.""" @@ -1661,3 +1779,560 @@ def test_custom_domain_shape_view_below_nested_split_is_unchanged( ) _assert_byte_identical(model, transformed) + + +class TestExpPositiveScaleFolding: + """Test conservative positive scale folding into an existing pre-Exp bias.""" + + @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]), + ], + ) + values = {"x": rng.normal(size=(1, 2, 2)).astype(np.float32)} + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + 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), + rtol=2e-6, + atol=2e-6, + ) + + 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]), + ], + ) + values = {"x": rng.normal(size=(1, 4)).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", "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, + ) + + 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(exp_positive_scale_folding=True), + ) + + 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) + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-6, + atol=2e-6, + ) + + def test_constant_attribute_operands_preserve_float32_dtype(self) -> None: + rng = np.random.default_rng(33) + model = _model( + [ + onnx.helper.make_node("Constant", [], ["bias"], value_float=-1.0), + onnx.helper.make_node("Constant", [], ["scale"], value_float=1.5), + 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, 4])], + [_info("y", [1, 4])], + [], + value_info=[_info("biased", [1, 4]), _info("exponential", [1, 4])], + ) + values = {"x": rng.normal(size=(1, 4)).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"] + 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-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) + + @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("node_index", [2, 3, 4]) + def test_custom_domain_interpreted_nodes_are_unchanged( + self, + exp_scale_model: onnx.ModelProto, + node_index: int, + ) -> None: + exp_scale_model.graph.node[node_index].domain = "com.example" + + 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) From 0b4d706c096c7c9fd6eddff4ddc1100337061043 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 11 Aug 2026 14:39:22 +0800 Subject: [PATCH 03/15] fix(optim): harden algebraic folding rewrites --- src/winml/modelkit/optim/pipes/algebraic.py | 198 +++++++++++++-- tests/unit/optim/pipes/test_pipe_algebraic.py | 237 ++++++++++++++++++ 2 files changed, 412 insertions(+), 23 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 0e50d74b1..48ab73ecc 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -7,6 +7,7 @@ from __future__ import annotations from dataclasses import dataclass +from itertools import pairwise from typing import Any, ClassVar, cast import numpy as np @@ -29,6 +30,7 @@ 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 @@ -157,6 +159,18 @@ class _ExpScaleInsertCandidate: 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.""" @@ -244,6 +258,12 @@ def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: 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): @@ -442,6 +462,148 @@ 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 = next( + (int(opset.version) for opset in model.opset_import if opset.domain in ("", "ai.onnx")), + 0, + ) + if 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, @@ -509,6 +671,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): @@ -1109,17 +1273,9 @@ def _copy_conv_parameters( ) -> bool: if len(conv.input) < 2 or conv.input[1] in index.graph_inputs: return False - weight = next( - ( - initializer - for initializer in model.graph.initializer - if initializer.name == conv.input[1] - ), - None, - ) - if weight is None: + 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): @@ -1128,17 +1284,9 @@ def _copy_conv_parameters( if len(conv.input) > 2 and conv.input[2]: if conv.input[2] in index.graph_inputs: return False - bias = next( - ( - initializer - for initializer in model.graph.initializer - if initializer.name == conv.input[2] - ), - None, - ) - if bias is None: + 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): @@ -1190,10 +1338,10 @@ def _fold_channel_affine( 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) @@ -1328,6 +1476,7 @@ 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), ) @@ -1337,6 +1486,7 @@ def should_process(cls, config: AlgebraicRewritePipeConfig) -> bool: """Return whether any algebraic rewrite is enabled.""" return ( config.static_split_to_slice + or config.sibling_slice_to_split or config.conv_channel_affine_folding or config.exp_positive_scale_folding ) @@ -1361,6 +1511,8 @@ def process( _fold_channel_affine(result, allocator) if config.exp_positive_scale_folding: _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 07928fc32..9d8bc997d 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -429,6 +429,143 @@ 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( + ("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( @@ -651,6 +788,67 @@ def test_overridable_conv_parameters_are_unchanged(self, parameter_name: str) -> 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", [ @@ -1784,6 +1982,45 @@ def test_custom_domain_shape_view_below_nested_split_is_unchanged( class TestExpPositiveScaleFolding: """Test conservative positive scale folding into an existing pre-Exp bias.""" + 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] From 3a665b08afac62e75259dca4223020d4abad6bf1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 11 Aug 2026 14:59:01 +0800 Subject: [PATCH 04/15] fix(optim): satisfy algebraic mypy checks --- src/winml/modelkit/optim/pipes/algebraic.py | 40 ++++++++++----------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 48ab73ecc..343504c4a 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -985,44 +985,44 @@ def _fold_exp_positive_scales( """Fold eligible positive post-Exp constants into existing input biases.""" index = _GraphIndex.build(model) for add in list(model.graph.node): - candidate = _exp_scale_candidate(index, add) - if candidate is None: + bias_candidate = _exp_scale_candidate(index, add) + if bias_candidate is None: continue combined_name = _new_initializer( model, allocator, - candidate.combined_bias, + bias_candidate.combined_bias, "algebraic_exp_log_bias", ) - candidate.add.input[candidate.bias_input_index] = combined_name - candidate.output_node.output[0] = candidate.mul.output[0] - _remove_nodes(model, {id(candidate.mul)}) + bias_candidate.add.input[bias_candidate.bias_input_index] = combined_name + bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] + _remove_nodes(model, {id(bias_candidate.mul)}) index = _GraphIndex.build(model) for exp in list(model.graph.node): - candidate = _exp_scale_insert_candidate(index, exp) - if candidate is None: + insert_candidate = _exp_scale_insert_candidate(index, exp) + if insert_candidate is None: continue log_scale_name = _new_initializer( model, allocator, - candidate.log_scale, + insert_candidate.log_scale, "algebraic_exp_log_scale", ) adjusted_name = allocator.new("algebraic_exp_adjusted") add = onnx.helper.make_node( "Add", - [candidate.exp.input[0], log_scale_name], + [insert_candidate.exp.input[0], log_scale_name], [adjusted_name], name=allocator.new("algebraic_exp_log_add"), ) - candidate.exp.input[0] = adjusted_name - candidate.output_node.output[0] = candidate.mul.output[0] + insert_candidate.exp.input[0] = adjusted_name + insert_candidate.output_node.output[0] = insert_candidate.mul.output[0] rewritten: list[onnx.NodeProto] = [] for node in model.graph.node: - if node is candidate.exp: + if node is insert_candidate.exp: rewritten.append(add) - if node is not candidate.mul: + if node is not insert_candidate.mul: rewritten.append(node) del model.graph.node[:] model.graph.node.extend(rewritten) @@ -1206,7 +1206,7 @@ def _collect_routed_affine_candidates( boundaries = nested_info[1] if len(boundaries) != len(nested_split.output): return [] - candidates: list[_AffineCandidate] = [] + nested_affine_candidates: list[_AffineCandidate] = [] for output_index, (local_start, local_end) in enumerate(boundaries): nested_candidates = _collect_routed_affine_candidates( index, @@ -1220,8 +1220,8 @@ def _collect_routed_affine_candidates( ) if nested_candidates is None: return None - candidates.extend(nested_candidates) - return candidates + 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 @@ -1245,7 +1245,7 @@ def _collect_routed_affine_candidates( ): return [] - candidates: list[_AffineCandidate] = [] + routed_affine_candidates: list[_AffineCandidate] = [] for routed_slice, local_start, local_end in routed_slices: routed_candidates = _collect_routed_affine_candidates( index, @@ -1259,8 +1259,8 @@ def _collect_routed_affine_candidates( ) if routed_candidates is None: return None - candidates.extend(routed_candidates) - return candidates + routed_affine_candidates.extend(routed_candidates) + return routed_affine_candidates def _copy_conv_parameters( From 562f96eb3d37c352bbea4594b2f93bf5e206ee4f Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 11:30:55 +0800 Subject: [PATCH 05/15] fix(optim): address exp scale review comments --- .../modelkit/optim/capabilities/algebraic.py | 5 +- src/winml/modelkit/optim/pipes/algebraic.py | 61 ++++++++++-- tests/unit/optim/pipes/test_pipe_algebraic.py | 99 ++++++++++++++++++- 3 files changed, 152 insertions(+), 13 deletions(-) diff --git a/src/winml/modelkit/optim/capabilities/algebraic.py b/src/winml/modelkit/optim/capabilities/algebraic.py index f3a6fd3a4..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 @@ -35,7 +35,8 @@ name="exp-positive-scale-folding", ort_name=None, description=( - "Fold a finite, strictly positive constant scale after Exp into the log-domain input bias" + "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/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 343504c4a..a40200e98 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -311,6 +311,16 @@ 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 _new_initializer( model: onnx.ModelProto, allocator: _NameAllocator, @@ -966,9 +976,8 @@ def _exp_scale_insert_candidate( or np.prod(input_shape, dtype=np.int64) != np.prod(output_shape, dtype=np.int64) ): return None - broadcast_scale = np.asarray(np.broadcast_to(scale, output_shape)) - log_scale = np.asarray(np.log(broadcast_scale).reshape(input_shape), dtype=scale.dtype) - if not np.isfinite(log_scale).all(): + log_scale = _compact_log_scale_for_input(scale, input_shape) + if log_scale is None: return None return _ExpScaleInsertCandidate( exp=exp, @@ -978,13 +987,26 @@ def _exp_scale_insert_candidate( ) -def _fold_exp_positive_scales( +def _compact_log_scale_for_input( + scale: np.ndarray, + input_shape: tuple[int, ...], +) -> np.ndarray | None: + """Return a log-scale initializer without expanding broadcast-only dimensions.""" + if _shape_broadcasts_to(scale.shape, input_shape): + log_scale = np.asarray(np.log(scale), dtype=scale.dtype) + elif scale.size == int(np.prod(input_shape, dtype=np.int64)): + log_scale = np.asarray(np.log(scale).reshape(input_shape), dtype=scale.dtype) + else: + return None + return log_scale if np.isfinite(log_scale).all() else None + + +def _fold_existing_exp_bias_scale( model: onnx.ModelProto, allocator: _NameAllocator, -) -> None: - """Fold eligible positive post-Exp constants into existing input biases.""" +) -> bool: index = _GraphIndex.build(model) - for add in list(model.graph.node): + for add in model.graph.node: bias_candidate = _exp_scale_candidate(index, add) if bias_candidate is None: continue @@ -997,9 +1019,16 @@ def _fold_exp_positive_scales( bias_candidate.add.input[bias_candidate.bias_input_index] = combined_name bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] _remove_nodes(model, {id(bias_candidate.mul)}) - index = _GraphIndex.build(model) + return True + return False + - for exp in list(model.graph.node): +def _fold_inserted_exp_scale( + model: onnx.ModelProto, + allocator: _NameAllocator, +) -> bool: + index = _GraphIndex.build(model) + for exp in model.graph.node: insert_candidate = _exp_scale_insert_candidate(index, exp) if insert_candidate is None: continue @@ -1026,7 +1055,19 @@ def _fold_exp_positive_scales( rewritten.append(node) del model.graph.node[:] model.graph.node.extend(rewritten) - index = _GraphIndex.build(model) + return True + return False + + +def _fold_exp_positive_scales( + model: onnx.ModelProto, + allocator: _NameAllocator, +) -> None: + """Fold eligible positive post-Exp constants into the Exp input.""" + while _fold_existing_exp_bias_scale(model, allocator): + pass + while _fold_inserted_exp_scale(model, allocator): + pass def _collect_affine_chain( diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 9d8bc997d..5de81241b 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -1980,7 +1980,104 @@ def test_custom_domain_shape_view_below_nested_split_is_unchanged( class TestExpPositiveScaleFolding: - """Test conservative positive scale folding into an existing pre-Exp bias.""" + """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_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_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]) def test_simple_numeric_example_matches_log_domain_identity(self) -> None: x = np.asarray([[0.0, 2.0]], dtype=np.float32) From 7a4434d388ab0b38ac92280373ccf541b514d0a1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 12:15:41 +0800 Subject: [PATCH 06/15] fix(optim): preserve compact exp scale shapes --- src/winml/modelkit/optim/pipes/algebraic.py | 100 +++++++++++++++--- tests/unit/optim/pipes/test_pipe_algebraic.py | 73 +++++++++++++ 2 files changed, 156 insertions(+), 17 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index a40200e98..d8ad789ed 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -140,13 +140,16 @@ class _AffineCandidate: @dataclass class _ExpScaleCandidate: - """A positive post-Exp scale that can be merged into an existing bias.""" + """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 mul: onnx.NodeProto - combined_bias: np.ndarray + add_output: str + route_consumer: onnx.NodeProto + combined_bias: np.ndarray | None + log_scale: np.ndarray | None @dataclass @@ -321,6 +324,10 @@ def _shape_broadcasts_to(shape: tuple[int, ...], target_shape: tuple[int, ...]) ) +def _shape_element_count(shape: tuple[int, ...]) -> int: + return int(np.prod(shape, dtype=np.int64)) + + def _new_initializer( model: onnx.ModelProto, allocator: _NameAllocator, @@ -893,6 +900,7 @@ def _exp_scale_candidate( 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: @@ -908,6 +916,7 @@ def _exp_scale_candidate( 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) @@ -929,18 +938,22 @@ def _exp_scale_candidate( try: np.broadcast_to(bias, add_shape) np.broadcast_to(scale, add_shape) - combined_bias = np.asarray(bias + np.log(scale), dtype=bias.dtype) - np.broadcast_to(combined_bias, 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(combined_bias).all(): + 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, mul=mul, + add_output=add_output, + route_consumer=route_consumer, combined_bias=combined_bias, + log_scale=None if combined_bias is not None else log_scale, ) @@ -976,7 +989,7 @@ def _exp_scale_insert_candidate( or np.prod(input_shape, dtype=np.int64) != np.prod(output_shape, dtype=np.int64) ): return None - log_scale = _compact_log_scale_for_input(scale, input_shape) + log_scale = _compact_log_scale_for_input(scale, input_shape, output_shape) if log_scale is None: return None return _ExpScaleInsertCandidate( @@ -990,17 +1003,37 @@ def _exp_scale_insert_candidate( 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 _shape_broadcasts_to(scale.shape, input_shape): + 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 == int(np.prod(input_shape, dtype=np.int64)): + elif scale.size == _shape_element_count(input_shape): log_scale = np.asarray(np.log(scale).reshape(input_shape), dtype=scale.dtype) 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 + if not _shape_broadcasts_to(combined_shape, target_shape) or _shape_element_count( + combined_shape + ) > 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 _fold_existing_exp_bias_scale( model: onnx.ModelProto, allocator: _NameAllocator, @@ -1010,15 +1043,48 @@ def _fold_existing_exp_bias_scale( bias_candidate = _exp_scale_candidate(index, add) if bias_candidate is None: continue - 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 - bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] - _remove_nodes(model, {id(bias_candidate.mul)}) + 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 + and bias_candidate.route_consumer.input + and bias_candidate.route_consumer.input[0] == bias_candidate.add_output + ): + 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 + bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + if node is bias_candidate.mul: + continue + rewritten.append(node) + if node is bias_candidate.add: + rewritten.append(log_add) + del model.graph.node[:] + model.graph.node.extend(rewritten) + else: + continue + if bias_candidate.combined_bias is not None: + bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] + _remove_nodes(model, {id(bias_candidate.mul)}) return True return False diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 5de81241b..2056604fc 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -2053,6 +2053,79 @@ def test_scalar_scale_without_bias_stays_compact(self) -> None: 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 From ca3a2d6172e2f225ab4064f531cceab455c5ef44 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 13:56:46 +0800 Subject: [PATCH 07/15] fix(optim): guard exp scale edge cases --- src/winml/modelkit/optim/pipes/algebraic.py | 44 ++++++++++--- tests/unit/optim/pipes/test_pipe_algebraic.py | 63 +++++++++++++++++++ 2 files changed, 98 insertions(+), 9 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index d8ad789ed..07af4bc0c 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -6,6 +6,7 @@ from __future__ import annotations +import math from dataclasses import dataclass from itertools import pairwise from typing import Any, ClassVar, cast @@ -23,6 +24,7 @@ algebraic.EXP_POSITIVE_SCALE_FOLDING, ) MAX_AFFINE_ROUTE_DEPTH = 64 +MAX_NUMPY_ELEMENTS = np.iinfo(np.intp).max @dataclass @@ -325,7 +327,23 @@ def _shape_broadcasts_to(shape: tuple[int, ...], target_shape: tuple[int, ...]) def _shape_element_count(shape: tuple[int, ...]) -> int: - return int(np.prod(shape, dtype=np.int64)) + 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 _standard_opset_version(model: onnx.ModelProto) -> int | None: + for opset in model.opset_import: + if opset.domain in ("", "ai.onnx"): + return int(opset.version) + return None def _new_initializer( @@ -749,7 +767,7 @@ 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 @@ -788,7 +806,7 @@ def _order_preserving_view_output( input_shape is None or output_shape is None or any(dimension <= 0 for dimension in (*input_shape, *output_shape)) - or np.prod(input_shape, dtype=np.int64) != np.prod(output_shape, dtype=np.int64) + or not _same_shape_element_count(input_shape, output_shape) ): return None @@ -986,7 +1004,7 @@ def _exp_scale_insert_candidate( if ( input_shape is None or any(dimension <= 0 for dimension in (*input_shape, *output_shape)) - or np.prod(input_shape, dtype=np.int64) != np.prod(output_shape, dtype=np.int64) + or not _same_shape_element_count(input_shape, output_shape) ): return None log_scale = _compact_log_scale_for_input(scale, input_shape, output_shape) @@ -1011,7 +1029,10 @@ def _compact_log_scale_for_input( ): log_scale = np.asarray(np.log(scale), dtype=scale.dtype) elif scale.size == _shape_element_count(input_shape): - log_scale = np.asarray(np.log(scale).reshape(input_shape), dtype=scale.dtype) + 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 @@ -1026,9 +1047,12 @@ def _compact_combined_bias( combined_shape = np.broadcast_shapes(bias.shape, log_scale.shape) except ValueError: return None - if not _shape_broadcasts_to(combined_shape, target_shape) or _shape_element_count( - combined_shape - ) > max(int(bias.size), int(log_scale.size)): + 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 @@ -1617,7 +1641,9 @@ def process( if config.conv_channel_affine_folding: _fold_channel_affine(result, allocator) if config.exp_positive_scale_folding: - _fold_exp_positive_scales(result, allocator) + standard_opset = _standard_opset_version(result) + if 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: diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 2056604fc..8ccccc9e5 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -2152,6 +2152,69 @@ def test_float32_boundary_behavior_is_relaxed_when_enabled(self) -> None: 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 + + 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) From 6fa5343aaf649b89b29597855508134ab52d64c1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 14:41:22 +0800 Subject: [PATCH 08/15] fix(optim): resolve standard opsets strictly --- src/winml/modelkit/optim/pipes/algebraic.py | 24 ++---- tests/unit/optim/pipes/test_pipe_algebraic.py | 83 +++++++++++++++++++ 2 files changed, 91 insertions(+), 16 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 07af4bc0c..819c3d677 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -339,11 +339,9 @@ def _same_shape_element_count( return left_count >= 0 and left_count == _shape_element_count(right) -def _standard_opset_version(model: onnx.ModelProto) -> int | None: - for opset in model.opset_import: - if opset.domain in ("", "ai.onnx"): - return int(opset.version) - return None +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( @@ -595,11 +593,8 @@ def _fold_sibling_slices_to_split( allocator: _NameAllocator, ) -> None: """Replace contiguous sibling Slice nodes with an equivalent Split.""" - opset = next( - (int(opset.version) for opset in model.opset_import if opset.domain in ("", "ai.onnx")), - 0, - ) - if opset < 13: + 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) @@ -1541,11 +1536,8 @@ 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]] = {} @@ -1641,7 +1633,7 @@ def process( if config.conv_channel_affine_folding: _fold_channel_affine(result, allocator) if config.exp_positive_scale_folding: - standard_opset = _standard_opset_version(result) + standard_opset = _strict_default_opset_version(result) if standard_opset is not None and standard_opset >= 7: _fold_exp_positive_scales(result, allocator) if config.sibling_slice_to_split: diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 8ccccc9e5..62dd58825 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -514,6 +514,52 @@ def test_sibling_static_slices_are_unchanged_before_split_input_opset(self) -> N 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"), [ @@ -2184,6 +2230,43 @@ def test_legacy_or_missing_standard_opset_is_unchanged( 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( From 786a0fd5a0ed7f09fc4ed2fcbf2b69cc57130a4f Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 16:13:02 +0800 Subject: [PATCH 09/15] fix(optim): address algebraic review compatibility --- src/winml/modelkit/optim/pipes/algebraic.py | 160 ++++++++++++------ tests/unit/optim/pipes/test_pipe_algebraic.py | 97 ++++++++++- 2 files changed, 204 insertions(+), 53 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 819c3d677..c8b6a3eab 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -9,7 +9,7 @@ import math from dataclasses import dataclass from itertools import pairwise -from typing import Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast import numpy as np import onnx @@ -18,6 +18,10 @@ 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, algebraic.CONV_CHANNEL_AFFINE_FOLDING, @@ -226,7 +230,7 @@ def _attribute(node: onnx.NodeProto, name: str, default: Any = None) -> Any: def _is_standard_onnx_node(node: onnx.NodeProto) -> bool: - return node.domain in ("", "ai.onnx") + return node.domain == "" def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: @@ -1053,15 +1057,66 @@ def _compact_combined_bias( return combined_bias if np.isfinite(combined_bias).all() else None -def _fold_existing_exp_bias_scale( +def _candidate_node_ids(nodes: Iterable[onnx.NodeProto]) -> set[int]: + return {id(node) for node in nodes} + + +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 = _candidate_node_ids( + [candidate.add, candidate.route_consumer, candidate.output_node, candidate.mul] + ) + 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, +) -> list[_ExpScaleInsertCandidate]: + selected: list[_ExpScaleInsertCandidate] = [] + reserved_nodes: set[int] = set() + for exp in model.graph.node: + candidate = _exp_scale_insert_candidate(index, exp) + if candidate is None: + continue + candidate_nodes = _candidate_node_ids([candidate.exp, candidate.output_node, candidate.mul]) + if candidate_nodes & reserved_nodes: + continue + reserved_nodes.update(candidate_nodes) + selected.append(candidate) + return selected + + +def _fold_existing_exp_bias_scales( model: onnx.ModelProto, allocator: _NameAllocator, ) -> bool: index = _GraphIndex.build(model) - for add in model.graph.node: - bias_candidate = _exp_scale_candidate(index, add) - if bias_candidate is None: - continue + bias_candidates = _select_exp_bias_scale_candidates(model, index) + if not bias_candidates: + return False + + removed: set[int] = set() + 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, @@ -1070,11 +1125,7 @@ def _fold_existing_exp_bias_scale( "algebraic_exp_log_bias", ) bias_candidate.add.input[bias_candidate.bias_input_index] = combined_name - elif ( - bias_candidate.log_scale is not None - and bias_candidate.route_consumer.input - and bias_candidate.route_consumer.input[0] == bias_candidate.add_output - ): + elif bias_candidate.log_scale is not None: log_scale_name = _new_initializer( model, allocator, @@ -1089,34 +1140,37 @@ def _fold_existing_exp_bias_scale( name=allocator.new("algebraic_exp_log_add"), ) bias_candidate.route_consumer.input[0] = adjusted_name - bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] - rewritten: list[onnx.NodeProto] = [] - for node in model.graph.node: - if node is bias_candidate.mul: - continue - rewritten.append(node) - if node is bias_candidate.add: - rewritten.append(log_add) - del model.graph.node[:] - model.graph.node.extend(rewritten) + insert_after[id(bias_candidate.add)] = log_add else: continue - if bias_candidate.combined_bias is not None: - bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] - _remove_nodes(model, {id(bias_candidate.mul)}) - return True - return False + bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] + removed.add(id(bias_candidate.mul)) + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + if id(node) in removed: + continue + rewritten.append(node) + inserted = insert_after.get(id(node)) + if inserted is not None: + rewritten.append(inserted) + del model.graph.node[:] + model.graph.node.extend(rewritten) + return True -def _fold_inserted_exp_scale( + +def _fold_inserted_exp_scales( model: onnx.ModelProto, allocator: _NameAllocator, ) -> bool: index = _GraphIndex.build(model) - for exp in model.graph.node: - insert_candidate = _exp_scale_insert_candidate(index, exp) - if insert_candidate is None: - continue + insert_candidates = _select_exp_scale_insert_candidates(model, index) + if not insert_candidates: + return False + + removed: set[int] = set() + insert_before: dict[int, onnx.NodeProto] = {} + for insert_candidate in insert_candidates: log_scale_name = _new_initializer( model, allocator, @@ -1132,16 +1186,19 @@ def _fold_inserted_exp_scale( ) insert_candidate.exp.input[0] = adjusted_name insert_candidate.output_node.output[0] = insert_candidate.mul.output[0] - rewritten: list[onnx.NodeProto] = [] - for node in model.graph.node: - if node is insert_candidate.exp: - rewritten.append(add) - if node is not insert_candidate.mul: - rewritten.append(node) - del model.graph.node[:] - model.graph.node.extend(rewritten) - return True - return False + insert_before[id(insert_candidate.exp)] = add + removed.add(id(insert_candidate.mul)) + + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + inserted = insert_before.get(id(node)) + if inserted is not None: + rewritten.append(inserted) + if id(node) not in removed: + rewritten.append(node) + del model.graph.node[:] + model.graph.node.extend(rewritten) + return True def _fold_exp_positive_scales( @@ -1149,9 +1206,9 @@ def _fold_exp_positive_scales( allocator: _NameAllocator, ) -> None: """Fold eligible positive post-Exp constants into the Exp input.""" - while _fold_existing_exp_bias_scale(model, allocator): + while _fold_existing_exp_bias_scales(model, allocator): pass - while _fold_inserted_exp_scale(model, allocator): + while _fold_inserted_exp_scales(model, allocator): pass @@ -1625,17 +1682,22 @@ 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: - standard_opset = _strict_default_opset_version(result) - if standard_opset is not None and standard_opset >= 7: - _fold_exp_positive_scales(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: diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 62dd58825..9ad07ccc8 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -21,6 +21,7 @@ AlgebraicRewritePipe, AlgebraicRewritePipeConfig, ) +from winml.modelkit.optim.pipes import algebraic as algebraic_pipe if TYPE_CHECKING: @@ -360,7 +361,7 @@ def test_overridable_split_sizes_are_unchanged(self) -> None: @pytest.mark.parametrize( ("domain", "should_rewrite"), - [("ai.onnx", True), ("com.example", False)], + [("", True), ("ai.onnx", False), ("com.example", False)], ) def test_only_standard_domain_split_is_rewritten( self, @@ -393,6 +394,23 @@ def test_only_standard_domain_split_is_rewritten( 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")], @@ -727,6 +745,36 @@ 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_float64_affine_values_preserve_weight_precision(self) -> None: shape = [1, 1, 1, 1] @@ -1784,7 +1832,7 @@ def test_constant_attribute_affine_is_folded_and_pruned(self) -> None: @pytest.mark.parametrize( ("domain", "should_fold"), - [("ai.onnx", True), ("com.example", False)], + [("", True), ("ai.onnx", False), ("com.example", False)], ) def test_only_standard_domain_constant_is_interpreted( self, @@ -2074,6 +2122,45 @@ def test_multiple_exp_chains_are_folded_from_live_graph_nodes(self) -> None: ): 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_scalar_scale_without_bias_stays_compact(self) -> None: model = _model( [ @@ -2770,13 +2857,15 @@ def test_malformed_exp_mul_cycle_is_unchanged(self) -> None: _assert_byte_identical(model, transformed) + @pytest.mark.parametrize("domain", ["ai.onnx", "com.example"]) @pytest.mark.parametrize("node_index", [2, 3, 4]) - def test_custom_domain_interpreted_nodes_are_unchanged( + 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 = "com.example" + exp_scale_model.graph.node[node_index].domain = domain transformed = AlgebraicRewritePipe().process( exp_scale_model, From 69599177595f6a26e1ce0276fab3e6147c354df3 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 16:56:22 +0800 Subject: [PATCH 10/15] fix(optim): complete exp scale batching --- src/winml/modelkit/optim/pipes/algebraic.py | 111 +++++++++++++----- tests/unit/optim/pipes/test_pipe_algebraic.py | 79 +++++++++++++ 2 files changed, 162 insertions(+), 28 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index c8b6a3eab..dc0c70a9d 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -151,7 +151,7 @@ class _ExpScaleCandidate: add: onnx.NodeProto bias_input_index: int output_node: onnx.NodeProto - mul: onnx.NodeProto + muls: list[onnx.NodeProto] add_output: str route_consumer: onnx.NodeProto combined_bias: np.ndarray | None @@ -164,7 +164,7 @@ class _ExpScaleInsertCandidate: exp: onnx.NodeProto output_node: onnx.NodeProto - mul: onnx.NodeProto + muls: list[onnx.NodeProto] log_scale: np.ndarray @@ -197,14 +197,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 @@ -851,11 +853,35 @@ def _constant_input( return constants[0] if len(constants) == 1 else None +def _compact_combined_scale( + left: np.ndarray, + right: np.ndarray, + target_shape: tuple[int, ...], +) -> np.ndarray | None: + if left.dtype != right.dtype: + return None + try: + combined_shape = np.broadcast_shapes(left.shape, right.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(left.size), int(right.size)) + ): + return None + combined_scale = np.asarray(left * right, dtype=left.dtype) + return ( + combined_scale if np.isfinite(combined_scale).all() and np.all(combined_scale > 0) else None + ) + + def _post_exp_scale( index: _GraphIndex, exp: onnx.NodeProto, visited: set[str] | None = None, -) -> tuple[onnx.NodeProto, onnx.NodeProto, np.ndarray, tuple[int, ...]] | 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) @@ -881,23 +907,45 @@ def _post_exp_scale( if next_node is None or not _is_standard_onnx_node(next_node) or next_node.op_type != "Mul": return None - scale_operand = _constant_input(index, next_node, current_name) - mul_output = _node_output(next_node) output_shape = _static_shape(index, current_name) - if scale_operand is None or mul_output is None or output_shape is None: + if output_shape is None: 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) + + muls: list[onnx.NodeProto] = [] + combined_scale: np.ndarray | None = None + while ( + next_node is not None and _is_standard_onnx_node(next_node) and next_node.op_type == "Mul" ): + 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: + break + scale = scale_operand[1] + if ( + not np.issubdtype(scale.dtype, np.floating) + or not np.isfinite(scale).all() + or not np.all(scale > 0) + ): + break + try: + np.broadcast_to(scale, output_shape) + except ValueError: + break + next_scale = ( + scale + if combined_scale is None + else _compact_combined_scale(combined_scale, scale, output_shape) + ) + if next_scale is None: + break + combined_scale = next_scale + muls.append(next_node) + current_name = mul_output + next_node = _single_unobserved_consumer(index, current_name) + + if combined_scale is None: return None - try: - np.broadcast_to(scale, output_shape) - except ValueError: - return None - return current_node, next_node, scale, output_shape + return current_node, muls, combined_scale, output_shape def _exp_scale_candidate( @@ -939,7 +987,7 @@ def _exp_scale_candidate( post_exp = _post_exp_scale(index, next_node, visited) if post_exp is None: return None - output_node, mul, scale, output_shape = post_exp + output_node, muls, scale, output_shape = post_exp if output_shape != add_shape: return None @@ -966,7 +1014,7 @@ def _exp_scale_candidate( add=add, bias_input_index=bias_operand[0], output_node=output_node, - mul=mul, + muls=muls, add_output=add_output, route_consumer=route_consumer, combined_bias=combined_bias, @@ -998,7 +1046,7 @@ def _exp_scale_insert_candidate( post_exp = _post_exp_scale(index, exp) if post_exp is None: return None - output_node, mul, scale, output_shape = post_exp + output_node, muls, scale, output_shape = post_exp input_shape = _static_shape(index, exp.input[0]) if ( input_shape is None @@ -1012,7 +1060,7 @@ def _exp_scale_insert_candidate( return _ExpScaleInsertCandidate( exp=exp, output_node=output_node, - mul=mul, + muls=muls, log_scale=log_scale, ) @@ -1078,7 +1126,12 @@ def _select_exp_bias_scale_candidates( ): continue candidate_nodes = _candidate_node_ids( - [candidate.add, candidate.route_consumer, candidate.output_node, candidate.mul] + [ + candidate.add, + candidate.route_consumer, + candidate.output_node, + *candidate.muls, + ] ) if candidate_nodes & reserved_nodes: continue @@ -1097,7 +1150,9 @@ def _select_exp_scale_insert_candidates( candidate = _exp_scale_insert_candidate(index, exp) if candidate is None: continue - candidate_nodes = _candidate_node_ids([candidate.exp, candidate.output_node, candidate.mul]) + candidate_nodes = _candidate_node_ids( + [candidate.exp, candidate.output_node, *candidate.muls] + ) if candidate_nodes & reserved_nodes: continue reserved_nodes.update(candidate_nodes) @@ -1143,8 +1198,8 @@ def _fold_existing_exp_bias_scales( insert_after[id(bias_candidate.add)] = log_add else: continue - bias_candidate.output_node.output[0] = bias_candidate.mul.output[0] - removed.add(id(bias_candidate.mul)) + bias_candidate.output_node.output[0] = bias_candidate.muls[-1].output[0] + removed.update(id(mul) for mul in bias_candidate.muls) rewritten: list[onnx.NodeProto] = [] for node in model.graph.node: @@ -1185,9 +1240,9 @@ def _fold_inserted_exp_scales( name=allocator.new("algebraic_exp_log_add"), ) insert_candidate.exp.input[0] = adjusted_name - insert_candidate.output_node.output[0] = insert_candidate.mul.output[0] + insert_candidate.output_node.output[0] = insert_candidate.muls[-1].output[0] insert_before[id(insert_candidate.exp)] = add - removed.add(id(insert_candidate.mul)) + removed.update(id(mul) for mul in insert_candidate.muls) rewritten: list[onnx.NodeProto] = [] for node in model.graph.node: diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 9ad07ccc8..58e202d76 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -238,6 +238,31 @@ def test_cli_combines_split_affine_and_exp_folding(self, tmp_path: Path) -> None 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.""" @@ -2161,6 +2186,60 @@ def counted_build(model: onnx.ModelProto) -> algebraic_pipe._GraphIndex: assert build_count <= 6 assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] * chain_count + def test_serial_exp_mul_scales_are_folded_without_per_scale_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.1 + index / 10, 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"] + _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( [ From 3267e4bf38ce787de0cd67d9f7cdad5503386152 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 18:20:49 +0800 Subject: [PATCH 11/15] fix(optim): preserve serial exp scale semantics --- src/winml/modelkit/optim/pipes/algebraic.py | 220 ++++++++---------- tests/unit/optim/pipes/test_pipe_algebraic.py | 152 +++++++++++- 2 files changed, 248 insertions(+), 124 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index dc0c70a9d..710529133 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -29,6 +29,12 @@ ) MAX_AFFINE_ROUTE_DEPTH = 64 MAX_NUMPY_ELEMENTS = np.iinfo(np.intp).max +EXP_FOLDING_GENERATED_PREFIXES = ( + "algebraic_exp_log_bias", + "algebraic_exp_log_scale", + "algebraic_exp_log_add", + "algebraic_exp_adjusted", +) @dataclass @@ -853,27 +859,13 @@ def _constant_input( return constants[0] if len(constants) == 1 else None -def _compact_combined_scale( - left: np.ndarray, - right: np.ndarray, - target_shape: tuple[int, ...], -) -> np.ndarray | None: - if left.dtype != right.dtype: - return None - try: - combined_shape = np.broadcast_shapes(left.shape, right.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(left.size), int(right.size)) - ): - return None - combined_scale = np.asarray(left * right, dtype=left.dtype) - return ( - combined_scale if np.isfinite(combined_scale).all() and np.all(combined_scale > 0) else None +def _is_generated_exp_folding_name(name: str) -> bool: + return any(name.startswith(prefix) for prefix in EXP_FOLDING_GENERATED_PREFIXES) + + +def _has_generated_exp_folding_marker(node: onnx.NodeProto) -> bool: + return bool(node.name and _is_generated_exp_folding_name(node.name)) or any( + _is_generated_exp_folding_name(input_name) for input_name in node.input if input_name ) @@ -911,41 +903,22 @@ def _post_exp_scale( if output_shape is None: return None - muls: list[onnx.NodeProto] = [] - combined_scale: np.ndarray | None = None - while ( - next_node is not None and _is_standard_onnx_node(next_node) and next_node.op_type == "Mul" + 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 + scale = scale_operand[1] + if ( + not np.issubdtype(scale.dtype, np.floating) + or not np.isfinite(scale).all() + or not np.all(scale > 0) ): - 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: - break - scale = scale_operand[1] - if ( - not np.issubdtype(scale.dtype, np.floating) - or not np.isfinite(scale).all() - or not np.all(scale > 0) - ): - break - try: - np.broadcast_to(scale, output_shape) - except ValueError: - break - next_scale = ( - scale - if combined_scale is None - else _compact_combined_scale(combined_scale, scale, output_shape) - ) - if next_scale is None: - break - combined_scale = next_scale - muls.append(next_node) - current_name = mul_output - next_node = _single_unobserved_consumer(index, current_name) - - if combined_scale is None: return None - return current_node, muls, combined_scale, output_shape + try: + np.broadcast_to(scale, output_shape) + except ValueError: + return None + return current_node, [next_node], scale, output_shape def _exp_scale_candidate( @@ -954,6 +927,8 @@ def _exp_scale_candidate( ) -> _ExpScaleCandidate | None: if not _is_standard_onnx_node(add) or add.op_type != "Add": return None + if _has_generated_exp_folding_marker(add): + return None add_output = _node_output(add) bias_operand = _constant_input(index, add) if add_output is None or bias_operand is None: @@ -1031,6 +1006,8 @@ def _exp_scale_insert_candidate( current_name = exp.input[0] visited = {current_name} producer = index.producers.get(current_name) + if producer is not None and _has_generated_exp_folding_marker(producer): + return None while producer is not None and producer.op_type in {"Reshape", "Squeeze", "Unsqueeze"}: if ( not producer.input @@ -1042,6 +1019,8 @@ def _exp_scale_insert_candidate( return None visited.add(current_name) producer = index.producers.get(current_name) + if producer is not None and _has_generated_exp_folding_marker(producer): + return None post_exp = _post_exp_scale(index, exp) if post_exp is None: @@ -1109,6 +1088,21 @@ 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, @@ -1125,14 +1119,7 @@ def _select_exp_bias_scale_candidates( or candidate.route_consumer.input[0] != candidate.add_output ): continue - candidate_nodes = _candidate_node_ids( - [ - candidate.add, - candidate.route_consumer, - candidate.output_node, - *candidate.muls, - ] - ) + candidate_nodes = _exp_bias_candidate_node_ids(candidate) if candidate_nodes & reserved_nodes: continue reserved_nodes.update(candidate_nodes) @@ -1143,33 +1130,38 @@ def _select_exp_bias_scale_candidates( def _select_exp_scale_insert_candidates( model: onnx.ModelProto, index: _GraphIndex, + reserved_nodes: set[int] | None = None, ) -> list[_ExpScaleInsertCandidate]: selected: list[_ExpScaleInsertCandidate] = [] - reserved_nodes: set[int] = set() + 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 = _candidate_node_ids( - [candidate.exp, candidate.output_node, *candidate.muls] - ) - if candidate_nodes & reserved_nodes: + candidate_nodes = _exp_insert_candidate_node_ids(candidate) + if candidate_nodes & selected_nodes: continue - reserved_nodes.update(candidate_nodes) + selected_nodes.update(candidate_nodes) selected.append(candidate) return selected -def _fold_existing_exp_bias_scales( +def _fold_exp_positive_scales( model: onnx.ModelProto, allocator: _NameAllocator, -) -> bool: +) -> None: + """Fold eligible positive post-Exp constants into the Exp input.""" index = _GraphIndex.build(model) bias_candidates = _select_exp_bias_scale_candidates(model, index) - if not bias_candidates: - return False + 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: @@ -1201,30 +1193,6 @@ def _fold_existing_exp_bias_scales( bias_candidate.output_node.output[0] = bias_candidate.muls[-1].output[0] removed.update(id(mul) for mul in bias_candidate.muls) - rewritten: list[onnx.NodeProto] = [] - for node in model.graph.node: - if id(node) in removed: - continue - rewritten.append(node) - inserted = insert_after.get(id(node)) - if inserted is not None: - rewritten.append(inserted) - del model.graph.node[:] - model.graph.node.extend(rewritten) - return True - - -def _fold_inserted_exp_scales( - model: onnx.ModelProto, - allocator: _NameAllocator, -) -> bool: - index = _GraphIndex.build(model) - insert_candidates = _select_exp_scale_insert_candidates(model, index) - if not insert_candidates: - return False - - removed: set[int] = set() - insert_before: dict[int, onnx.NodeProto] = {} for insert_candidate in insert_candidates: log_scale_name = _new_initializer( model, @@ -1246,25 +1214,16 @@ def _fold_inserted_exp_scales( rewritten: list[onnx.NodeProto] = [] for node in model.graph.node: - inserted = insert_before.get(id(node)) - if inserted is not None: - rewritten.append(inserted) + 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) - return True - - -def _fold_exp_positive_scales( - model: onnx.ModelProto, - allocator: _NameAllocator, -) -> None: - """Fold eligible positive post-Exp constants into the Exp input.""" - while _fold_existing_exp_bias_scales(model, allocator): - pass - while _fold_inserted_exp_scales(model, allocator): - pass def _collect_affine_chain( @@ -1301,10 +1260,19 @@ def _collect_affine_chain( 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, []) @@ -1509,6 +1477,8 @@ def _copy_conv_parameters( scale: np.ndarray, offset: np.ndarray, ) -> bool: + if not np.isfinite(scale).all() or not np.isfinite(offset).all(): + return False if len(conv.input) < 2 or conv.input[1] in index.graph_inputs: return False weights = _initializer_array(index, conv.input[1]) @@ -1518,6 +1488,8 @@ def _copy_conv_parameters( 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]: if conv.input[2] in index.graph_inputs: @@ -1529,24 +1501,34 @@ def _copy_conv_parameters( 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: diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index 58e202d76..f1a437725 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -800,6 +800,62 @@ def test_legacy_opset_conv_affine_broadcast_is_unchanged(self) -> None: 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] @@ -2186,7 +2242,95 @@ def counted_build(model: onnx.ModelProto) -> algebraic_pipe._GraphIndex: assert build_count <= 6 assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] * chain_count - def test_serial_exp_mul_scales_are_folded_without_per_scale_rebuilds( + 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)} + + transformed = AlgebraicRewritePipe().process( + model, + AlgebraicRewritePipeConfig(exp_positive_scale_folding=True), + ) + + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp", "Mul"] + _assert_valid_with_inferred_shapes(transformed) + np.testing.assert_allclose( + _run(model, values), + _run(transformed, values), + rtol=2e-6, + atol=2e-6, + ) + + 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] == ["Add", "Exp", "Mul"] + _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_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 any(node.op_type == "Mul" for node in transformed.graph.node) + _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: @@ -2201,9 +2345,7 @@ def test_serial_exp_mul_scales_are_folded_without_per_scale_rebuilds( 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.1 + index / 10, dtype=np.float32)) - ) + 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 @@ -2231,7 +2373,7 @@ def counted_build(model: onnx.ModelProto) -> algebraic_pipe._GraphIndex: ) assert build_count <= 6 - assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp"] + assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp", *["Mul"] * 5] _assert_valid_with_inferred_shapes(transformed) np.testing.assert_allclose( _run(model, values), From 954a8d5fc75fb7ca62d47e0c11a9cac4dcda4abe Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 13 Aug 2026 10:33:02 +0800 Subject: [PATCH 12/15] fix(optim): harden exp scale serial boundaries --- src/winml/modelkit/optim/pipes/algebraic.py | 66 +++-- tests/unit/optim/pipes/test_pipe_algebraic.py | 234 +++++++++++++++++- 2 files changed, 268 insertions(+), 32 deletions(-) diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py index 710529133..0a32a6af7 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -29,12 +29,7 @@ ) MAX_AFFINE_ROUTE_DEPTH = 64 MAX_NUMPY_ELEMENTS = np.iinfo(np.intp).max -EXP_FOLDING_GENERATED_PREFIXES = ( - "algebraic_exp_log_bias", - "algebraic_exp_log_scale", - "algebraic_exp_log_add", - "algebraic_exp_adjusted", -) +ORDER_PRESERVING_VIEW_OPS = frozenset({"Flatten", "Identity", "Reshape", "Squeeze", "Unsqueeze"}) @dataclass @@ -259,9 +254,11 @@ def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: 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, dtype in ( ("value_float", np.float32), @@ -804,9 +801,11 @@ def _order_preserving_view_output( or not _is_standard_onnx_node(node) or not node.input or node.input[0] != input_name - or node.op_type not in {"Reshape", "Squeeze", "Unsqueeze"} + 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 ( @@ -817,6 +816,18 @@ def _order_preserving_view_output( ): 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) @@ -859,14 +870,29 @@ def _constant_input( return constants[0] if len(constants) == 1 else None -def _is_generated_exp_folding_name(name: str) -> bool: - return any(name.startswith(prefix) for prefix in EXP_FOLDING_GENERATED_PREFIXES) - - -def _has_generated_exp_folding_marker(node: onnx.NodeProto) -> bool: - return bool(node.name and _is_generated_exp_folding_name(node.name)) or any( - _is_generated_exp_folding_name(input_name) for input_name in node.input if input_name - ) +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( @@ -907,6 +933,8 @@ def _post_exp_scale( 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) @@ -927,8 +955,6 @@ def _exp_scale_candidate( ) -> _ExpScaleCandidate | None: if not _is_standard_onnx_node(add) or add.op_type != "Add": return None - if _has_generated_exp_folding_marker(add): - return None add_output = _node_output(add) bias_operand = _constant_input(index, add) if add_output is None or bias_operand is None: @@ -1006,8 +1032,6 @@ def _exp_scale_insert_candidate( current_name = exp.input[0] visited = {current_name} producer = index.producers.get(current_name) - if producer is not None and _has_generated_exp_folding_marker(producer): - return None while producer is not None and producer.op_type in {"Reshape", "Squeeze", "Unsqueeze"}: if ( not producer.input @@ -1019,8 +1043,6 @@ def _exp_scale_insert_candidate( return None visited.add(current_name) producer = index.producers.get(current_name) - if producer is not None and _has_generated_exp_folding_marker(producer): - return None post_exp = _post_exp_scale(index, exp) if post_exp is None: diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index f1a437725..cdbd47760 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -2260,22 +2260,63 @@ def test_serial_exp_mul_keeps_later_scales_to_preserve_float_boundaries( 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 [node.op_type for node in transformed.graph.node] == ["Add", "Exp", "Mul"] + 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_allclose( - _run(model, values), - _run(transformed, values), - rtol=2e-6, - atol=2e-6, + 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])], ) - def test_public_optimizer_preserves_serial_exp_mul_float_boundaries(self) -> None: + 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"]), @@ -2291,13 +2332,160 @@ def test_public_optimizer_preserves_serial_exp_mul_float_boundaries(self) -> Non 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 [node.op_type for node in transformed.graph.node] == ["Add", "Exp", "Mul"] + 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( [ @@ -2326,7 +2514,7 @@ def test_public_optimizer_preserves_serial_exp_mul_after_marked_bias_view(self) transformed = optimize_onnx(model, exp_positive_scale_folding=True) - assert any(node.op_type == "Mul" for node in transformed.graph.node) + 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)) @@ -2373,7 +2561,7 @@ def counted_build(model: onnx.ModelProto) -> algebraic_pipe._GraphIndex: ) assert build_count <= 6 - assert [node.op_type for node in transformed.graph.node] == ["Add", "Exp", *["Mul"] * 5] + 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), @@ -2988,6 +3176,32 @@ def test_unloaded_external_scale_is_unchanged( _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, From 52d1a3056b5636fed0966b63c7945c9f1387066c Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 13 Aug 2026 11:03:10 +0800 Subject: [PATCH 13/15] fix(optim): report algebraic slice split analysis --- src/winml/modelkit/optim/capabilities/misc.py | 2 +- src/winml/modelkit/optim/pipes/algebraic.py | 3 +- tests/unit/optim/pipes/test_pipe_algebraic.py | 5 ++ tests/unit/optim/test_analysis.py | 49 +++++++++++++++++++ 4 files changed, 57 insertions(+), 2 deletions(-) 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 0a32a6af7..161d6a70c 100644 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ b/src/winml/modelkit/optim/pipes/algebraic.py @@ -14,7 +14,7 @@ import numpy as np import onnx -from ..capabilities import algebraic +from ..capabilities import algebraic, misc from .base import BasePipe, PipeConfig, caps_dict @@ -24,6 +24,7 @@ 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, ) diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py index cdbd47760..667e923bf 100644 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ b/tests/unit/optim/pipes/test_pipe_algebraic.py @@ -94,10 +94,12 @@ def test_capabilities_are_opt_in_and_independent(self) -> None: capabilities = get_all_capabilities() 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}") @@ -106,10 +108,12 @@ 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) @@ -118,6 +122,7 @@ 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 diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index 6c778ee4f..105028932 100644 --- a/tests/unit/optim/test_analysis.py +++ b/tests/unit/optim/test_analysis.py @@ -74,6 +74,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 # ============================================================================= @@ -457,3 +489,20 @@ def test_findings_match_analyze_model(self) -> None: iter_names = {finding.name for finding, _ in iter_optimization_outputs(model, caps)} 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" + and finding.pipe_name == "algebraic_rewrite" + ] + + assert len(matches) == 1 + finding, produced = matches[0] + + assert finding.enable_flag == "--enable-gather-slice-to-split-fusion" + 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"] From 00cfe33c522d051ebd8e05daf7ca4b84bf134f82 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 13 Aug 2026 13:50:45 +0800 Subject: [PATCH 14/15] fix(optim): aggregate shared capability analysis --- src/winml/modelkit/optim/analysis.py | 131 ++++++++++++++++++++++++++- tests/unit/optim/test_analysis.py | 2 +- 2 files changed, 127 insertions(+), 6 deletions(-) diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index af69d1e9e..32793bbf7 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: + 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(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(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/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index 525b2849e..18c78881e 100644 --- a/tests/unit/optim/test_analysis.py +++ b/tests/unit/optim/test_analysis.py @@ -724,13 +724,13 @@ def test_reports_algebraic_sibling_slice_to_split(self) -> None: (finding, produced) for finding, produced in pairs if finding.name == "gather-slice-to-split-fusion" - and finding.pipe_name == "algebraic_rewrite" ] 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"] From 8412d8dc4ad47536f289de560170c56046cc7739 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 13 Aug 2026 15:30:08 +0800 Subject: [PATCH 15/15] fix(optim): isolate shared analysis baseline --- src/winml/modelkit/optim/analysis.py | 6 +- tests/unit/optim/test_analysis.py | 94 ++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 3 deletions(-) diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index 32793bbf7..0af4c4014 100644 --- a/src/winml/modelkit/optim/analysis.py +++ b/src/winml/modelkit/optim/analysis.py @@ -482,7 +482,7 @@ def full_pipeline_baseline() -> ModelProto: probe_kwargs.update(optimizer_kwargs) try: - base_out = full_pipeline_baseline() + shared_base_out = full_pipeline_baseline() probe_out = _run_pipeline(PIPES, pipeline_input, probe_kwargs) except Exception as exc: logger.warning( @@ -494,14 +494,14 @@ def full_pipeline_baseline() -> ModelProto: shared_base_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {} shared_probe_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {} - _collect_nodes(base_out.graph, (), shared_base_nodes) + _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(base_out) + 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, diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index 18c78881e..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 @@ -735,6 +736,99 @@ def test_reports_algebraic_sibling_slice_to_split(self) -> None: 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: