diff --git a/src/winml/modelkit/pattern/__init__.py b/src/winml/modelkit/pattern/__init__.py index 38f0112c4..2bad9af6d 100644 --- a/src/winml/modelkit/pattern/__init__.py +++ b/src/winml/modelkit/pattern/__init__.py @@ -39,6 +39,12 @@ Conv2DInplaceLinear4DPatternInputGenerator, Conv2DInplaceLinearInputGeneratorBase, ) +from .conv_batchnorm_patterns import ( + CONV_ADD_BATCHNORM_SCHEMA, + AddConvBatchNormalizationPattern, + ConvAddBatchNormalizationPattern, + FoldedConvAddPattern, +) from .gelu_patterns import ( Gelu1Pattern, Gelu1PatternInputGenerator, @@ -84,7 +90,9 @@ __all__ = [ + "CONV_ADD_BATCHNORM_SCHEMA", "MATMUL_ADD_SCHEMA", + "AddConvBatchNormalizationPattern", "Conv2DInplaceLinear2DPattern", "Conv2DInplaceLinear2DPatternInputGenerator", "Conv2DInplaceLinear3DPattern", @@ -92,8 +100,10 @@ "Conv2DInplaceLinear4DPattern", "Conv2DInplaceLinear4DPatternInputGenerator", "Conv2DInplaceLinearInputGeneratorBase", + "ConvAddBatchNormalizationPattern", "ExpandedAttentionPattern", "ExpandedAttentionPatternInputGenerator", + "FoldedConvAddPattern", "Gelu1Pattern", "Gelu1PatternInputGenerator", "Gelu2Pattern", diff --git a/src/winml/modelkit/pattern/base.py b/src/winml/modelkit/pattern/base.py index 1ca29ab44..10cf2a007 100644 --- a/src/winml/modelkit/pattern/base.py +++ b/src/winml/modelkit/pattern/base.py @@ -1967,6 +1967,17 @@ def rewrite( ] used_graph_node_keys = set(graph_node_keys) generated_node_key_counter = 0 + used_graph_tensor_names = { + name + for name in ( + [value.name for value in graph.input] + + [value.name for value in graph.output] + + [value.name for value in graph.value_info] + + [initializer.name for initializer in graph.initializer] + + [name for node in graph.node for name in (*node.input, *node.output)] + ) + if name + } def _allocate_graph_node_key(node: Any) -> str: """Allocate a non-conflicting stable key for inserted nodes.""" @@ -1990,6 +2001,20 @@ def _allocate_graph_node_key(node: Any) -> str: used_graph_node_keys.add(key) return key + def _allocate_graph_tensor_name(name: str) -> str: + """Allocate a tensor name that does not collide with the parent graph.""" + if name not in used_graph_tensor_names: + used_graph_tensor_names.add(name) + return name + + suffix = 1 + candidate = f"{name}__{suffix}" + while candidate in used_graph_tensor_names: + suffix += 1 + candidate = f"{name}__{suffix}" + used_graph_tensor_names.add(candidate) + return candidate + # Track which nodes have been deleted to avoid double deletion deleted_node_names: set[str] = set() @@ -2097,6 +2122,34 @@ def _allocate_graph_node_key(node: Any) -> str: logger.debug("Skipping rewrite %s: %s", new_pattern_class.__name__, e) continue + # Remap target-local tensors that collide with any existing graph + # tensor. Schema inputs and outputs are graph boundaries and must + # retain their original names. + boundary_names = {*input_names, *output_names} + local_tensor_names = dict.fromkeys( + [initializer.name for initializer in new_subgraph_model.graph.initializer] + + [ + name + for node in new_subgraph_model.graph.node + for name in (*node.output, *node.input) + ] + ) + tensor_name_mapping = { + name: _allocate_graph_tensor_name(name) + for name in local_tensor_names + if name and name not in boundary_names + } + for node in new_subgraph_model.graph.node: + for index, name in enumerate(node.input): + node.input[index] = tensor_name_mapping.get(name, name) + for index, name in enumerate(node.output): + node.output[index] = tensor_name_mapping.get(name, name) + for initializer in new_subgraph_model.graph.initializer: + initializer.name = tensor_name_mapping.get( + initializer.name, + initializer.name, + ) + # Find insertion point: position of last matched node after deletions # Since original graph is topologically sorted, # last matched node is after all input producers @@ -2130,9 +2183,8 @@ def _allocate_graph_node_key(node: Any) -> str: # Append new initializers (constants) from the new subgraph for initializer in new_subgraph_model.graph.initializer: - # Check if initializer already exists (by name) - existing_names = {init.name for init in graph.initializer} - if initializer.name not in existing_names: + # Schema-input constants already exist in the parent graph. + if initializer.name not in boundary_names: graph.initializer.append(initializer) # Add any missing opset imports to the model diff --git a/src/winml/modelkit/pattern/conv_batchnorm_patterns.py b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py new file mode 100644 index 000000000..68c325f18 --- /dev/null +++ b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py @@ -0,0 +1,363 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Patterns for folding inference BatchNormalization after Conv and Add.""" + +from __future__ import annotations + +from math import prod +from typing import TYPE_CHECKING, Any + +import numpy as np +from onnx import ModelProto, helper, numpy_helper +from onnx.defs import OpSchema + +from ..onnx import EXTERNAL_DATA_THRESHOLD, ONNXDomain, SupportedONNXType +from .base import Pattern, PatternMismatchedError, PatternSchema, Skeleton + + +if TYPE_CHECKING: + from .match import PatternMatchResult, SkeletonMatchResult + + +CONV_ADD_BATCHNORM_SCHEMA = PatternSchema( + name="ConvAddBatchNormalizationPattern", + doc=( + "Inference BatchNormalization applied to the sum of a Conv output and " + "a static broadcast tensor." + ), + type_constraints=[ + OpSchema.TypeConstraintParam( + type_param_str="T", + allowed_type_strs=[ + "tensor(float16)", + "tensor(float)", + "tensor(double)", + ], + description="Constrain inputs and output to floating-point tensors.", + ) + ], + inputs=[ + OpSchema.FormalParameter(name, "T", description) + for name, description in ( + ("X", "Conv input."), + ("W", "Static Conv weights."), + ("A", "Static tensor added to the Conv output."), + ("scale", "BatchNormalization scale."), + ("B", "BatchNormalization bias."), + ("mean", "BatchNormalization input mean."), + ("var", "BatchNormalization input variance."), + ) + ], + outputs=[ + OpSchema.FormalParameter( + "Y", + "T", + "Folded Conv and Add output.", + ) + ], +) + + +def _scale_broadcast_tensor( + values: np.ndarray, + output_shape: tuple[int, ...], + gamma: np.ndarray, +) -> np.ndarray | None: + """Scale a broadcastable Add operand along the Conv channel axis.""" + if values.ndim > len(output_shape): + return None + padded_shape = (1,) * (len(output_shape) - values.ndim) + tuple(values.shape) + if any(dimension not in (1, output_shape[axis]) for axis, dimension in enumerate(padded_shape)): + return None + if padded_shape[1] not in (1, output_shape[1]): + return None + + expanded_shape = list(padded_shape) + expanded_shape[1] = output_shape[1] + if prod(expanded_shape) * values.dtype.itemsize >= EXTERNAL_DATA_THRESHOLD: + return None + try: + broadcast_values = np.broadcast_to(values.reshape(padded_shape), tuple(expanded_shape)) + except ValueError: + return None + factors = gamma.reshape((1, len(gamma)) + (1,) * (len(output_shape) - 2)) + with np.errstate(invalid="ignore", over="ignore"): + return np.asarray(broadcast_values * factors, dtype=values.dtype) + + +class _ConvAddBatchNormalizationPatternBase(Pattern): + """Shared validation for the two commutative Add input orders.""" + + static_input_index: int + + @property + def pattern_id(self) -> str: + """Return the source pattern ID used by rewrite registration.""" + return "SUBGRAPH/Conv-Add-Batch-NormalizationPattern" + + def get_schema(self) -> PatternSchema: + return CONV_ADD_BATCHNORM_SCHEMA + + def get_internal_constants_and_attributes( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + domain_versions: dict[ONNXDomain, int], + ) -> tuple[list[tuple[int, int, np.ndarray]], dict[tuple[int, str], Any]]: + return [], {} + + def get_skeleton(self) -> Skeleton: + conv_add_slot = 1 - self.static_input_index + return Skeleton( + node_op_types=["Conv", "Add", "BatchNormalization"], + node_domains=[ONNXDomain.AI_ONNX] * 3, + edges=[ + (-1, 0, 0, 0), + (-2, 0, 0, 1), + (0, 0, 1, conv_add_slot), + (-3, 0, 1, self.static_input_index), + (1, 0, 2, 0), + (-4, 0, 2, 1), + (-5, 0, 2, 2), + (-6, 0, 2, 3), + (-7, 0, 2, 4), + ], + exit_nodes=[2], + n_inputs=7, + ) + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + result = super().check_skeleton_result(skeleton_match_result) + if result is None: + return None + + conv, add, batch_norm = skeleton_match_result.matched_nodes + if ( + len(conv.input) not in (2, 3) + or len(conv.output) != 1 + or len(add.input) != 2 + or len(add.output) != 1 + or len(batch_norm.input) != 5 + or len(batch_norm.output) != 1 + ): + return None + + required_constants = ("W", "A", "scale", "B", "mean", "var") + if any( + name not in result.input_infos + or not result.input_infos[name].is_constant + or result.input_infos[name].value is None + for name in required_constants + ): + return None + + values = {name: np.asarray(result.input_infos[name].value) for name in required_constants} + weight = values["W"] + add_tensor = values["A"] + scale = values["scale"] + beta = values["B"] + mean = values["mean"] + variance = values["var"] + if weight.ndim < 1 or not np.issubdtype(weight.dtype, np.floating): + return None + + matcher = skeleton_match_result.matcher + output_shape_value = matcher.get_tensor_shape(conv.output[0]) + if ( + output_shape_value is None + or len(output_shape_value) < 2 + or any(not isinstance(dimension, int) for dimension in output_shape_value) + ): + return None + output_shape = tuple(int(dimension) for dimension in output_shape_value) + channels = output_shape[1] + if weight.shape[0] != channels: + return None + + parameters = (scale, beta, mean, variance) + if any(value.ndim != 1 or len(value) != channels for value in parameters): + return None + if any(value.dtype != weight.dtype for value in (*parameters, add_tensor)): + return None + + bias = np.zeros(channels, dtype=weight.dtype) + if len(conv.input) == 3 and conv.input[2]: + bias_value = matcher.tensor_values.get(conv.input[2]) + if bias_value is None or conv.input[2] not in matcher.constant_and_initializer_names: + return None + bias = np.asarray(bias_value) + if bias.ndim != 1 or len(bias) != channels or bias.dtype != weight.dtype: + return None + + conv_attributes = { + attribute.name: helper.get_attribute_value(attribute) for attribute in conv.attribute + } + batch_norm_attributes = { + attribute.name: helper.get_attribute_value(attribute) + for attribute in batch_norm.attribute + } + try: + if int(batch_norm_attributes.get("training_mode", 0)) != 0: + return None + epsilon = float(batch_norm_attributes.get("epsilon", 1e-5)) + except (TypeError, ValueError): + return None + + denominator = variance + epsilon + if epsilon < 0 or np.any(denominator <= 0) or not np.all(np.isfinite(denominator)): + return None + gamma = scale / np.sqrt(denominator) + if not np.all(np.isfinite(gamma)): + return None + + scaled_add = _scale_broadcast_tensor(add_tensor, output_shape, gamma) + if scaled_add is None: + return None + with np.errstate(invalid="ignore", over="ignore"): + folded_weight = np.asarray( + weight * gamma.reshape((channels,) + (1,) * (weight.ndim - 1)), + dtype=weight.dtype, + ) + folded_bias = np.asarray( + bias * gamma + beta - gamma * mean, + dtype=weight.dtype, + ) + if any( + not np.all(np.isfinite(value)) for value in (folded_weight, folded_bias, scaled_add) + ): + return None + + result.attributes.update( + { + "conv_attributes": conv_attributes, + "static_input_index": self.static_input_index, + "folded_weight": folded_weight, + "folded_bias": folded_bias, + "scaled_add": scaled_add, + } + ) + return result + + +class ConvAddBatchNormalizationPattern(_ConvAddBatchNormalizationPatternBase): + """Match ``Add(Conv(X, W[, bias]), A) -> BatchNormalization``.""" + + static_input_index = 1 + + +class AddConvBatchNormalizationPattern(_ConvAddBatchNormalizationPatternBase): + """Match ``Add(A, Conv(X, W[, bias])) -> BatchNormalization``.""" + + static_input_index = 0 + + +class FoldedConvAddPattern(Pattern): + """Generate a Conv and Add with BatchNormalization folded into constants.""" + + @property + def pattern_id(self) -> str: + """Return the folded target pattern ID used by rewrite registration.""" + return "SUBGRAPH/FoldedConvAddPattern" + + def get_schema(self) -> PatternSchema: + """Return the schema shared with the source patterns.""" + return CONV_ADD_BATCHNORM_SCHEMA + + def get_skeleton(self) -> Skeleton: + """Return the folded Conv-to-Add topology.""" + return Skeleton( + node_op_types=["Conv", "Add"], + node_domains=[ONNXDomain.AI_ONNX] * 2, + edges=[ + (-1, 0, 0, 0), + (-2, 0, 0, 1), + (0, 0, 1, 0), + (-3, 0, 1, 1), + ], + exit_nodes=[1], + n_inputs=7, + ) + + def get_internal_constants_and_attributes( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + domain_versions: dict[ONNXDomain, int], + ) -> tuple[list[tuple[int, int, np.ndarray]], dict[tuple[int, str], Any]]: + """Return no static constraints because generation is customized.""" + return [], {} + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + """Build the folded Conv and Add subgraph from captured constants.""" + required = ( + "conv_attributes", + "static_input_index", + "folded_weight", + "folded_bias", + "scaled_add", + ) + if any(name not in attributes for name in required): + raise PatternMismatchedError("Missing folded Conv/Add attributes") + if input_names is None or output_names is None: + raise PatternMismatchedError("Input and output names are required") + + weight_name = f"{prefix}weight" + bias_name = f"{prefix}bias" + add_name = f"{prefix}add" + conv_output = f"{prefix}conv_output" + conv = helper.make_node( + "Conv", + [input_names[0], weight_name, bias_name], + [conv_output], + name=f"{prefix}Conv", + **attributes["conv_attributes"], + ) + add_inputs = [conv_output, add_name] + if attributes["static_input_index"] == 0: + add_inputs.reverse() + add = helper.make_node( + "Add", + add_inputs, + [output_names[0]], + name=f"{prefix}Add", + ) + + input_type = SupportedONNXType.from_np_type(inputs["X"].dtype).tensor_proto_type + output_type = SupportedONNXType.from_onnx_type(output_dtypes[0]).tensor_proto_type + graph = helper.make_graph( + [conv, add], + f"{prefix}FoldedConvAdd", + [helper.make_tensor_value_info(input_names[0], input_type, inputs["X"].shape)], + [helper.make_tensor_value_info(output_names[0], output_type, None)], + initializer=[ + numpy_helper.from_array(attributes["folded_weight"], weight_name), + numpy_helper.from_array(attributes["folded_bias"], bias_name), + numpy_helper.from_array(attributes["scaled_add"], add_name), + ], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items() + ], + ) + model.ir_version = 11 + return model diff --git a/src/winml/modelkit/pattern/rules/default.json b/src/winml/modelkit/pattern/rules/default.json index a03df57fa..d97397308 100644 --- a/src/winml/modelkit/pattern/rules/default.json +++ b/src/winml/modelkit/pattern/rules/default.json @@ -177,6 +177,44 @@ } ] }, + { + "pattern_id": "SUBGRAPH/Conv-Add-Batch-NormalizationPattern", + "pattern_class": "ConvAddBatchNormalizationPattern", + "module": "winml.modelkit.pattern.conv_batchnorm_patterns", + "enabled": true, + "flag_name": "conv-add-batch-normalization", + "description": "Conv followed by static Add and inference BatchNormalization", + "explanation": "BatchNormalization can be folded into copied Conv and Add constants when the topology and parameters are static.", + "alternatives": [ + { + "pattern_to_id": "SUBGRAPH/FoldedConvAddPattern", + "pattern_class": "FoldedConvAddPattern", + "module": "winml.modelkit.pattern.conv_batchnorm_patterns", + "priority": 1, + "flag_name": "folding", + "reason": "Remove inference BatchNormalization by folding its affine transform into Conv and Add constants" + } + ] + }, + { + "pattern_id": "SUBGRAPH/Conv-Add-Batch-NormalizationPattern", + "pattern_class": "AddConvBatchNormalizationPattern", + "module": "winml.modelkit.pattern.conv_batchnorm_patterns", + "enabled": true, + "flag_name": "add-conv-batch-normalization", + "description": "Static Add with Conv followed by inference BatchNormalization", + "explanation": "BatchNormalization can be folded into copied Conv and Add constants when the topology and parameters are static.", + "alternatives": [ + { + "pattern_to_id": "SUBGRAPH/FoldedConvAddPattern", + "pattern_class": "FoldedConvAddPattern", + "module": "winml.modelkit.pattern.conv_batchnorm_patterns", + "priority": 1, + "flag_name": "folding", + "reason": "Remove inference BatchNormalization by folding its affine transform into Conv and Add constants" + } + ] + }, { "pattern_id": "SUBGRAPH/LayerNormalizationPattern", "pattern_class": "LayerNormalizationPowPattern", diff --git a/tests/unit/analyze/core/test_unified_pattern_config.py b/tests/unit/analyze/core/test_unified_pattern_config.py index 0de5dd5d5..1429539d0 100644 --- a/tests/unit/analyze/core/test_unified_pattern_config.py +++ b/tests/unit/analyze/core/test_unified_pattern_config.py @@ -19,9 +19,10 @@ def test_load_default_config(self): htp_patterns = config.get_htp_patterns() # Should load all patterns from default.json - # 8 patterns: Gelu1-4, MatMulAdd, LayerNormPow, LayerNormMul, ReshapeTransposeReshape - assert len(skeleton_patterns) == 8, ( - f"Expected 8 skeleton patterns, got {len(skeleton_patterns)}" + # 10 patterns: Gelu1-4, MatMulAdd, two Conv/Add/BatchNorm variants, + # LayerNormPow, LayerNormMul, and ReshapeTransposeReshape. + assert len(skeleton_patterns) == 10, ( + f"Expected 10 skeleton patterns, got {len(skeleton_patterns)}" ) assert len(htp_patterns) == 1, f"Expected 1 HTP pattern, got {len(htp_patterns)}" @@ -30,6 +31,7 @@ def test_load_default_config(self): expected_ids = { "SUBGRAPH/GeluPattern", # Shared by Gelu1-4 "SUBGRAPH/GemmPattern", # MatMulAdd + "SUBGRAPH/Conv-Add-Batch-NormalizationPattern", # Both Add input orders "SUBGRAPH/LayerNormalizationPattern", # Shared by Pow and Mul variants "SUBGRAPH/ReshapeTransposeReshapeOverlyHighDimPattern", } @@ -43,9 +45,9 @@ def test_load_qnn_config_with_inheritance(self): htp_patterns = config.get_htp_patterns() # Should load all patterns from default + QNN overrides - # 9 patterns: 8 from default + TransposeAttentionPattern from QNN - assert len(skeleton_patterns) == 9, ( - f"Expected 9 skeleton patterns, got {len(skeleton_patterns)}" + # 11 patterns: 10 from default + TransposeAttentionPattern from QNN + assert len(skeleton_patterns) == 11, ( + f"Expected 11 skeleton patterns, got {len(skeleton_patterns)}" ) # HTP patterns should be inherited from default assert len(htp_patterns) == 1, f"Expected 1 HTP pattern, got {len(htp_patterns)}" @@ -184,7 +186,7 @@ def test_missing_ihv_config_falls_back_to_default(self): # Should load default patterns with a warning skeleton_patterns = config.get_skeleton_patterns() - assert len(skeleton_patterns) == 8, "Should fall back to default patterns" + assert len(skeleton_patterns) == 10, "Should fall back to default patterns" def test_alternatives_with_pattern_class(self): """Test that alternatives with pattern_class field are loaded correctly.""" diff --git a/tests/unit/pattern/test_conv_batchnorm_patterns.py b/tests/unit/pattern/test_conv_batchnorm_patterns.py new file mode 100644 index 000000000..a884ef970 --- /dev/null +++ b/tests/unit/pattern/test_conv_batchnorm_patterns.py @@ -0,0 +1,208 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for Conv/Add/BatchNormalization folding patterns.""" + +from __future__ import annotations + +import numpy as np +from onnx import ModelProto, TensorProto, ValueInfoProto, checker, helper, numpy_helper +from onnx.reference import ReferenceEvaluator + +from winml.modelkit.optim.pipes.rewrite import RewritePipe +from winml.modelkit.pattern import ( + AddConvBatchNormalizationPattern, + ConvAddBatchNormalizationPattern, + FoldedConvAddPattern, +) +from winml.modelkit.pattern.conv_batchnorm_patterns import _scale_broadcast_tensor + + +def _value_info(name: str, shape: list[int]) -> ValueInfoProto: + return helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + + +def _build_model( + *, + static_first: bool = False, + with_bias: bool = False, + training_mode: int = 0, + dynamic_scale: bool = False, + public_conv_output: bool = False, + variance: np.ndarray | None = None, +) -> tuple[ModelProto, dict[str, np.ndarray]]: + rng = np.random.default_rng(29) + channels = 3 + initializers = { + "weight": rng.normal(size=(channels, 2, 1, 1)).astype(np.float32), + "static": rng.normal(size=(1, channels, 1, 1)).astype(np.float32), + "scale": rng.uniform(0.5, 1.5, size=channels).astype(np.float32), + "beta": rng.normal(size=channels).astype(np.float32), + "mean": rng.normal(size=channels).astype(np.float32), + "variance": ( + rng.uniform(0.5, 1.5, size=channels).astype(np.float32) + if variance is None + else variance + ), + } + conv_inputs = ["x", "weight"] + if with_bias: + initializers["conv_bias"] = rng.normal(size=channels).astype(np.float32) + conv_inputs.append("conv_bias") + + add_inputs = ["conv_out", "static"] + if static_first: + add_inputs.reverse() + nodes = [ + helper.make_node( + "Conv", + conv_inputs, + ["conv_out"], + pads=[0, 0, 0, 0], + strides=[1, 1], + ), + helper.make_node("Add", add_inputs, ["add_out"]), + helper.make_node( + "BatchNormalization", + ["add_out", "scale", "beta", "mean", "variance"], + ["y"], + epsilon=0.01, + training_mode=training_mode, + ), + ] + graph_inputs = [_value_info("x", [1, 2, 2, 2])] + if dynamic_scale: + initializers.pop("scale") + graph_inputs.append(_value_info("scale", [channels])) + + outputs = [_value_info("y", [1, channels, 2, 2])] + if public_conv_output: + outputs.append(_value_info("conv_out", [1, channels, 2, 2])) + graph = helper.make_graph( + nodes, + "conv_add_batch_norm", + graph_inputs, + outputs, + initializer=[numpy_helper.from_array(value, name) for name, value in initializers.items()], + value_info=[ + _value_info("conv_out", [1, channels, 2, 2]), + _value_info("add_out", [1, channels, 2, 2]), + ], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 10 + feeds = {"x": rng.normal(size=(1, 2, 2, 2)).astype(np.float32)} + if dynamic_scale: + feeds["scale"] = rng.uniform(0.5, 1.5, size=channels).astype(np.float32) + return model, feeds + + +def _fold(model: ModelProto) -> ModelProto: + config = RewritePipe.build_config(conv_add_batch_normalization_folding=True) + return RewritePipe().process(model, config) + + +def _run( + model: ModelProto, + feeds: dict[str, np.ndarray], +) -> list[np.ndarray]: + return ReferenceEvaluator(model).run(None, feeds) + + +def test_rewrite_capability_is_registered() -> None: + config = RewritePipe.build_config(conv_add_batch_normalization_folding=True) + assert len(config.rules) == 2 + assert {rule.source.pattern_id for rule in config.rules} == { + "SUBGRAPH/Conv-Add-Batch-NormalizationPattern" + } + assert {rule.target.pattern_id for rule in config.rules} == {"SUBGRAPH/FoldedConvAddPattern"} + assert isinstance(config.rules[0].source, ConvAddBatchNormalizationPattern) + assert isinstance(config.rules[1].source, AddConvBatchNormalizationPattern) + assert all(isinstance(rule.target, FoldedConvAddPattern) for rule in config.rules) + + +def test_both_add_orders_and_optional_bias_are_equivalent() -> None: + for static_first in (False, True): + for with_bias in (False, True): + model, feeds = _build_model( + static_first=static_first, + with_bias=with_bias, + ) + expected = _run(model, feeds) + transformed = _fold(model) + + checker.check_model(transformed) + assert [node.op_type for node in transformed.graph.node] == ["Conv", "Add"] + assert transformed.graph.node[-1].output == ["y"] + assert not any(node.op_type == "BatchNormalization" for node in transformed.graph.node) + actual = _run(transformed, feeds) + np.testing.assert_allclose(actual[0], expected[0], rtol=3e-5, atol=3e-5) + + +def test_training_dynamic_parameters_and_public_intermediate_are_rejected() -> None: + cases = ( + {"training_mode": 1}, + {"dynamic_scale": True}, + {"public_conv_output": True}, + ) + for kwargs in cases: + model, _ = _build_model(**kwargs) + transformed = _fold(model) + assert any(node.op_type == "BatchNormalization" for node in transformed.graph.node) + + +def test_invalid_variance_is_rejected() -> None: + model, _ = _build_model(variance=np.asarray([-0.02, 0.5, 1.0], dtype=np.float32)) + transformed = _fold(model) + assert any(node.op_type == "BatchNormalization" for node in transformed.graph.node) + + +def test_non_finite_folded_constants_are_rejected() -> None: + model, _ = _build_model() + replacements = { + "weight": np.full((3, 2, 1, 1), 40000, dtype=np.float16), + "static": np.zeros((1, 3, 1, 1), dtype=np.float16), + "scale": np.full(3, 2, dtype=np.float16), + "beta": np.zeros(3, dtype=np.float16), + "mean": np.zeros(3, dtype=np.float16), + "variance": np.ones(3, dtype=np.float16), + } + for initializer in model.graph.initializer: + if initializer.name in replacements: + initializer.CopyFrom( + numpy_helper.from_array(replacements[initializer.name], initializer.name) + ) + for value_info in (*model.graph.input, *model.graph.output, *model.graph.value_info): + value_info.type.tensor_type.elem_type = TensorProto.FLOAT16 + + transformed = _fold(model) + + assert any(node.op_type == "BatchNormalization" for node in transformed.graph.node) + + +def test_excessive_static_broadcast_expansion_is_rejected() -> None: + values = np.ones((1, 1, 1024, 1024), dtype=np.float32) + gamma = np.ones(512, dtype=np.float32) + + assert _scale_broadcast_tensor(values, (1, 512, 1024, 1024), gamma) is None + + +def test_generated_initializer_names_do_not_collide() -> None: + model, feeds = _build_model() + expected = _run(model, feeds)[0] + collision_name = "Rewrite_FoldedConvAddPattern_0_weight" + collision_value = np.full((3, 2, 1, 1), 17, dtype=np.float32) + model.graph.initializer.append(numpy_helper.from_array(collision_value, collision_name)) + model.graph.node.append(helper.make_node("Identity", [collision_name], ["collision_output"])) + model.graph.output.append(_value_info("collision_output", [3, 2, 1, 1])) + + transformed = _fold(model) + + checker.check_model(transformed) + conv = next(node for node in transformed.graph.node if node.op_type == "Conv") + assert conv.input[1] != collision_name + np.testing.assert_allclose(_run(transformed, feeds)[0], expected, rtol=3e-5, atol=3e-5)