From 55cc9b23edfa81f722f310657b5519348e9b13dd Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Thu, 23 Jul 2026 08:43:28 +0800 Subject: [PATCH 1/4] feat(optim): fold Conv Add BatchNormalization pattern --- src/winml/modelkit/pattern/__init__.py | 10 + .../pattern/conv_batchnorm_patterns.py | 349 ++++++++++++++++++ src/winml/modelkit/pattern/rules/default.json | 38 ++ .../pattern/test_conv_batchnorm_patterns.py | 161 ++++++++ 4 files changed, 558 insertions(+) create mode 100644 src/winml/modelkit/pattern/conv_batchnorm_patterns.py create mode 100644 tests/unit/pattern/test_conv_batchnorm_patterns.py 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/conv_batchnorm_patterns.py b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py new file mode 100644 index 000000000..f5a33ef59 --- /dev/null +++ b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py @@ -0,0 +1,349 @@ +# ------------------------------------------------------------------------- +# 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 typing import TYPE_CHECKING, Any + +import numpy as np +import onnx +from onnx import helper, numpy_helper +from onnx.defs import OpSchema + +from ..onnx import 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] + 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)) + 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 + scaled_weight = weight * gamma.reshape((channels,) + (1,) * (weight.ndim - 1)) + folded_bias = bias * gamma + beta - gamma * mean + + result.attributes.update( + { + "conv_attributes": conv_attributes, + "static_input_index": self.static_input_index, + "folded_weight": np.asarray(scaled_weight, dtype=weight.dtype), + "folded_bias": np.asarray(folded_bias, dtype=weight.dtype), + "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, + ) -> onnx.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/pattern/test_conv_batchnorm_patterns.py b/tests/unit/pattern/test_conv_batchnorm_patterns.py new file mode 100644 index 000000000..fefeab7bd --- /dev/null +++ b/tests/unit/pattern/test_conv_batchnorm_patterns.py @@ -0,0 +1,161 @@ +# ------------------------------------------------------------------------- +# 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 +import onnx +from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator + +from winml.modelkit.optim.pipes.rewrite import RewritePipe +from winml.modelkit.pattern import ( + AddConvBatchNormalizationPattern, + ConvAddBatchNormalizationPattern, + FoldedConvAddPattern, +) + + +def _value_info(name: str, shape: list[int]) -> onnx.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[onnx.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: onnx.ModelProto) -> onnx.ModelProto: + config = RewritePipe.build_config(conv_add_batch_normalization_folding=True) + return RewritePipe().process(model, config) + + +def _run( + model: onnx.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) + + onnx.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) From 5b6c6aa17754b8a01541b29760b6105e38667510 Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Thu, 23 Jul 2026 10:42:30 +0800 Subject: [PATCH 2/4] fix(pattern): use consistent ONNX imports --- .../modelkit/pattern/conv_batchnorm_patterns.py | 5 ++--- tests/unit/pattern/test_conv_batchnorm_patterns.py | 13 ++++++------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/winml/modelkit/pattern/conv_batchnorm_patterns.py b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py index f5a33ef59..4b93ab30c 100644 --- a/src/winml/modelkit/pattern/conv_batchnorm_patterns.py +++ b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py @@ -9,8 +9,7 @@ from typing import TYPE_CHECKING, Any import numpy as np -import onnx -from onnx import helper, numpy_helper +from onnx import ModelProto, helper, numpy_helper from onnx.defs import OpSchema from ..onnx import ONNXDomain, SupportedONNXType @@ -290,7 +289,7 @@ def get_onnx_model( prefix: str = "", input_names: list[str] | None = None, output_names: list[str] | None = None, - ) -> onnx.ModelProto: + ) -> ModelProto: """Build the folded Conv and Add subgraph from captured constants.""" required = ( "conv_attributes", diff --git a/tests/unit/pattern/test_conv_batchnorm_patterns.py b/tests/unit/pattern/test_conv_batchnorm_patterns.py index fefeab7bd..2971aa4ac 100644 --- a/tests/unit/pattern/test_conv_batchnorm_patterns.py +++ b/tests/unit/pattern/test_conv_batchnorm_patterns.py @@ -7,8 +7,7 @@ from __future__ import annotations import numpy as np -import onnx -from onnx import TensorProto, helper, numpy_helper +from onnx import ModelProto, TensorProto, ValueInfoProto, checker, helper, numpy_helper from onnx.reference import ReferenceEvaluator from winml.modelkit.optim.pipes.rewrite import RewritePipe @@ -19,7 +18,7 @@ ) -def _value_info(name: str, shape: list[int]) -> onnx.ValueInfoProto: +def _value_info(name: str, shape: list[int]) -> ValueInfoProto: return helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) @@ -31,7 +30,7 @@ def _build_model( dynamic_scale: bool = False, public_conv_output: bool = False, variance: np.ndarray | None = None, -) -> tuple[onnx.ModelProto, dict[str, np.ndarray]]: +) -> tuple[ModelProto, dict[str, np.ndarray]]: rng = np.random.default_rng(29) channels = 3 initializers = { @@ -101,13 +100,13 @@ def _build_model( return model, feeds -def _fold(model: onnx.ModelProto) -> onnx.ModelProto: +def _fold(model: ModelProto) -> ModelProto: config = RewritePipe.build_config(conv_add_batch_normalization_folding=True) return RewritePipe().process(model, config) def _run( - model: onnx.ModelProto, + model: ModelProto, feeds: dict[str, np.ndarray], ) -> list[np.ndarray]: return ReferenceEvaluator(model).run(None, feeds) @@ -135,7 +134,7 @@ def test_both_add_orders_and_optional_bias_are_equivalent() -> None: expected = _run(model, feeds) transformed = _fold(model) - onnx.checker.check_model(transformed) + 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) From c72ffb3b44345a01a991269018edb342169e0441 Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Thu, 23 Jul 2026 11:05:07 +0800 Subject: [PATCH 3/4] test(pattern): update unified config expectations --- .../analyze/core/test_unified_pattern_config.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) 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.""" From c753bbaaf8603f4188d0565b79fc0cbe5e769c17 Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Mon, 27 Jul 2026 10:55:43 +0800 Subject: [PATCH 4/4] fix(pattern): harden generated rewrites --- src/winml/modelkit/pattern/base.py | 58 ++++++++++++++++++- .../pattern/conv_batchnorm_patterns.py | 27 +++++++-- .../pattern/test_conv_batchnorm_patterns.py | 48 +++++++++++++++ 3 files changed, 124 insertions(+), 9 deletions(-) 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 index 4b93ab30c..68c325f18 100644 --- a/src/winml/modelkit/pattern/conv_batchnorm_patterns.py +++ b/src/winml/modelkit/pattern/conv_batchnorm_patterns.py @@ -6,13 +6,14 @@ 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 ONNXDomain, SupportedONNXType +from ..onnx import EXTERNAL_DATA_THRESHOLD, ONNXDomain, SupportedONNXType from .base import Pattern, PatternMismatchedError, PatternSchema, Skeleton @@ -75,12 +76,15 @@ def _scale_broadcast_tensor( 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)) - return np.asarray(broadcast_values * factors, dtype=values.dtype) + with np.errstate(invalid="ignore", over="ignore"): + return np.asarray(broadcast_values * factors, dtype=values.dtype) class _ConvAddBatchNormalizationPatternBase(Pattern): @@ -215,15 +219,26 @@ def check_skeleton_result( scaled_add = _scale_broadcast_tensor(add_tensor, output_shape, gamma) if scaled_add is None: return None - scaled_weight = weight * gamma.reshape((channels,) + (1,) * (weight.ndim - 1)) - folded_bias = bias * gamma + beta - gamma * mean + 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": np.asarray(scaled_weight, dtype=weight.dtype), - "folded_bias": np.asarray(folded_bias, dtype=weight.dtype), + "folded_weight": folded_weight, + "folded_bias": folded_bias, "scaled_add": scaled_add, } ) diff --git a/tests/unit/pattern/test_conv_batchnorm_patterns.py b/tests/unit/pattern/test_conv_batchnorm_patterns.py index 2971aa4ac..a884ef970 100644 --- a/tests/unit/pattern/test_conv_batchnorm_patterns.py +++ b/tests/unit/pattern/test_conv_batchnorm_patterns.py @@ -16,6 +16,7 @@ ConvAddBatchNormalizationPattern, FoldedConvAddPattern, ) +from winml.modelkit.pattern.conv_batchnorm_patterns import _scale_broadcast_tensor def _value_info(name: str, shape: list[int]) -> ValueInfoProto: @@ -158,3 +159,50 @@ 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)