diff --git a/docs/commands/optimize.md b/docs/commands/optimize.md index 5c7134451..950fa8ec3 100644 --- a/docs/commands/optimize.md +++ b/docs/commands/optimize.md @@ -135,10 +135,9 @@ the shared pattern package does not re-export these backend-specific patterns. |------------|--------------------------| | `normalize-int32-dq` | Normalize initializer-backed INT32 `DequantizeLinear` in the standard domain (opset >= 10) and `com.microsoft` (opset 1): omit immutable all-zero scalar/singleton zero points and clone singleton scales as scalars. Preserve shared initializers and domains; skip overridable parameters, per-axis vectors, unsupported attributes and nonlocal inputs. Handles nested graphs, not local functions. Enabled by CGC build configuration, disabled by default elsewhere. | | `deduplicate-opset-imports` | Remove repeated model-level opset declarations with identical domain and version, retaining the first declaration and domain order. Reject conflicting versions for the same domain. Run before operator rewrites and opset upgrades; do not alter graph content, local functions, or the retained versions. | -| `eliminate-identity` | Remove safe internal tensor Identity aliases. Additionally replace top-level standard-domain FP32 graph-output Identities with same-shape Reshape when input/output types match exactly, all dimensions are positive static integers, opset >= 5, and the model has no subgraphs. Preserve output names/order, annotated aliases, unknown or conflicting types, scalar/dynamic/zero-size outputs and protected captures. Does not rewrite local functions. Workaround for [microsoft/ix#1198](https://github.com/microsoft/ix/issues/1198). | | `gridsample-to-gather` | Decompose 2D `GridSample` with linear interpolation and zero padding into four `GatherND` reads, bounds masks and weighted sums. Supports both `align_corners` settings, FP16/FP32 IO and dynamic batch, with known positive channel, input spatial and grid spatial dimensions. Rank-3 indices contain explicit batch and spatial coordinates; `batch_dims=0` avoids the ORT symbolic shape inference defect tracked in [onnxruntime#24206](https://github.com/microsoft/onnxruntime/pull/24206). Batch coordinates are generated dynamically and shared across the four reads; sampled values are reshaped back to the grid layout. FP16 interpolation is computed in FP32 and cast back. Requires opset >= 16 (`bilinear` before opset 20); other modes are unchanged. Enabled by CGC builds, disabled in ordinary optimization. Floating-point rounding may differ from native sampling. | | `omit-empty-resize-inputs` | Replace statically empty Resize ROI/scales with omitted inputs. Do not rely on graph-input defaults or rewrite crop-and-resize semantics. Requires opset 13; upgrade older matching models using ONNX version conversion. | -| `cgc-constant-folding` | Fill FoundryToolbox constant-folding gaps without an ORT Session. Fold standard `Pad.pads` constant integer chains; in graphs containing `Shape`, also fold statically known selected dimensions and bounded constant integer/boolean expressions (`Gather`, `Concat`, `Reshape`, `Slice`, `Transpose`, `Squeeze`, `Unsqueeze`, integer `Cast`, `ConstantOfShape`, arithmetic, `Equal`, `Where`). Iterate with shape inference, up to 32 rounds. Requires opset >= 11. Only the main graph is rewritten; preserve tensor names for shared uses and subgraph captures. Runtime floating-point computations and unresolved dimensions remain unchanged. This rule does not freeze inputs: specialize dimensions before optimization when needed; later Foundry `freeze-dims` does not retroactively affect this rule. Limits: 128 dependency values per traversal, 65,536 elements per operation and 1,048,576 cached elements per round. Enabled by CGC builds; disabled in ordinary optimization. `fold-constant-pad-pads` remains a compatibility alias. | +| `cgc-constant-folding` | Fill FoundryToolbox constant-folding gaps without an ORT Session. In graphs containing `Shape`, fold statically known selected dimensions and bounded constant integer/boolean expressions (`Gather`, `Concat`, `Reshape`, `Slice`, `Transpose`, `Squeeze`, `Unsqueeze`, integer `Cast`, `ConstantOfShape`, arithmetic, `Equal`, `Where`). Iterate with shape inference, up to 32 rounds. Requires opset >= 11. Only the main graph is rewritten; preserve tensor names for shared uses and subgraph captures. Graphs without `Shape`, runtime floating-point computations and unresolved dimensions remain unchanged. This rule does not freeze inputs: specialize dimensions before optimization when needed; later Foundry `freeze-dims` does not retroactively affect this rule. Limits: 128 dependency values per traversal, 65,536 elements per operation and 1,048,576 cached elements per round. Enabled by CGC builds; disabled in ordinary optimization. | | `resize-tf-half-pixel-for-nn-to-asymmetric` | Change only the coordinate mode for nearest/floor Resize with static, non-overridable, positive integer scales. Dynamic/fractional scales and sizes-based inference are outside this rule. | | `approximate-cubic-resize-with-linear` | **Lossy**, explicit cubic-to-linear approximation. Excludes antialiasing, outside exclusion, and crop-and-resize semantics. Prints a warning when applied. | | `gathernd-to-reshape` | Replace GatherND only when data/indices/output ranks are not all equal and static, non-overridable int64 indices visit every input slice exactly once in storage order. Require positive static data dimensions; support batch dimensions, multi-coordinate indices, and equivalent negative indices. Dynamic data shapes or indices, overridable defaults, empty tensors, partial selection, repetition, and reordering are outside this rule. | diff --git a/src/winml/modelkit/export/cgc/exporter.py b/src/winml/modelkit/export/cgc/exporter.py index 6518517f9..55d8e4ceb 100644 --- a/src/winml/modelkit/export/cgc/exporter.py +++ b/src/winml/modelkit/export/cgc/exporter.py @@ -31,7 +31,7 @@ class CGCOptions: """Configuration shared by CGC export steps.""" external_weights: bool = False - topo_sort_nodes: bool = True + topo_sort_nodes: bool = False update_opset: bool = True freeze_dims: str = "" diff --git a/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py b/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py index 6e5e951ef..475478e14 100644 --- a/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py +++ b/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py @@ -29,7 +29,6 @@ ResizeWithTfHalfPixelForNNPattern, cgc_constant_folding, deduplicate_opset_imports, - eliminate_identity, normalize_int32_dq, ) from ..registry import BoolCapability, CapabilityCategory @@ -73,26 +72,10 @@ class CGIRModelRewriteRule: default=False, ) -ELIMINATE_IDENTITY = BoolCapability( - name="eliminate-identity", - ort_name=None, - description="Eliminate internal tensor Identity aliases without changing graph IO for CGIR", - category=CapabilityCategory.REWRITE, - default=False, -) - -FOLD_CONSTANT_PAD_PADS = BoolCapability( - name="fold-constant-pad-pads", - ort_name=None, - description="Alias for cgc-constant-folding", - category=CapabilityCategory.REWRITE, - default=False, -) - CGC_CONSTANT_FOLDING = BoolCapability( name="cgc-constant-folding", ort_name=None, - description="Fill FoundryToolbox folding gaps for Pad parameters and static shape subgraphs", + description="Fill FoundryToolbox folding gaps for static shape subgraphs", category=CapabilityCategory.REWRITE, default=False, ) @@ -169,7 +152,6 @@ class CGIRModelRewriteRule: CGIRModelRewriteRule( capability=CGC_CONSTANT_FOLDING, transform=cgc_constant_folding, - aliases=(FOLD_CONSTANT_PAD_PADS,), ), CGIRModelRewriteRule( capability=DEDUPLICATE_OPSET_IMPORTS, @@ -214,10 +196,6 @@ class CGIRModelRewriteRule: target=MatMulDFTPattern, minimum_opset=17, ), - CGIRModelRewriteRule( - capability=ELIMINATE_IDENTITY, - transform=eliminate_identity, - ), CGIRRewriteRule( capability=GRIDSAMPLE_TO_GATHER, source=LinearGridSamplePattern, @@ -240,8 +218,6 @@ class CGIRModelRewriteRule: "CGIR_REWRITE_RULES", "DEDUPLICATE_OPSET_IMPORTS", "DFT_TO_MATMUL", - "ELIMINATE_IDENTITY", - "FOLD_CONSTANT_PAD_PADS", "GATHERND_TO_RESHAPE", "GRIDSAMPLE_TO_GATHER", "OMIT_EMPTY_RESIZE_INPUTS", diff --git a/src/winml/modelkit/pattern/cgc/__init__.py b/src/winml/modelkit/pattern/cgc/__init__.py index 82eda35ad..6aba322a5 100644 --- a/src/winml/modelkit/pattern/cgc/__init__.py +++ b/src/winml/modelkit/pattern/cgc/__init__.py @@ -4,12 +4,11 @@ # -------------------------------------------------------------------------- """Opt-in CGC compatibility patterns and model metadata rewrites.""" -from .cgc_constant_folding import cgc_constant_folding, fold_constant_pad_pads +from .cgc_constant_folding import cgc_constant_folding from .dft_patterns import DFTWithStaticParametersPattern, MatMulDFTPattern from .dq_rewrites import normalize_int32_dq from .gathernd_patterns import GatherNDWithIdentityIndicesPattern, ReshapedGatherNDPattern from .gridsample_patterns import GatherLinearGridSamplePattern, LinearGridSamplePattern -from .identity_rewrites import eliminate_identity from .opset_rewrites import deduplicate_opset_imports from .prelu_patterns import ExpandedPReluPattern, PReluWithFiniteSlopePattern from .resize_patterns import ( @@ -39,7 +38,5 @@ "ResizeWithTfHalfPixelForNNPattern", "cgc_constant_folding", "deduplicate_opset_imports", - "eliminate_identity", - "fold_constant_pad_pads", "normalize_int32_dq", ] diff --git a/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py b/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py index f4cf2c49a..021cea034 100644 --- a/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py +++ b/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py @@ -3,7 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Bounded constant folding for CGC Pad parameters and static shape subgraphs. +"""Bounded constant folding for CGC static shape subgraphs. FoundryToolbox currently lacks some constant folding needed by ONNX lowering. These rewrites fill that gap without an ORT Session or execution-provider graph @@ -14,13 +14,10 @@ import logging import math -from collections import Counter, deque from typing import cast import numpy as np from onnx import ( - AttributeProto, - GraphProto, ModelProto, TensorProto, ValueInfoProto, @@ -62,7 +59,6 @@ def __init__(self, model: ModelProto, *, static_shapes: bool = False) -> None: self.initializers = {value.name: value for value in model.graph.initializer} self.inputs = {value.name for value in model.graph.input} self.values: dict[str, np.ndarray] = {} - self.evaluated: set[str] = set() self.cached_elements = 0 def tensor(self, value: TensorProto) -> np.ndarray: @@ -153,7 +149,7 @@ def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray raise ValueError("Only integer Cast targets are supported") fragment = helper.make_model( helper.make_graph( - [node], "constant_pad_parameter", [], + [node], "constant_parameter", [], [ValueInfoProto(name=name)], [numpy_helper.from_array(value, key) for key, value in inputs.items()], ), @@ -161,7 +157,6 @@ def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray ir_version=self.model.ir_version, ) result = cast("list[np.ndarray]", ReferenceEvaluator(fragment).run(None, {}))[0] - self.evaluated.add(name) if result.dtype.kind not in "iub" or result.size > _MAX_ELEMENTS: raise ValueError("Constant result exceeds supported type or size") if self.cached_elements + result.size > _MAX_CACHED_ELEMENTS: @@ -173,150 +168,22 @@ def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray active.remove(name) -def _referenced_names(graph: GraphProto) -> list[str]: - names = [value.name for value in graph.output] - for annotation in graph.quantization_annotation: - names.append(annotation.tensor_name) - names.extend(item.value for item in annotation.quant_parameter_tensor_names) - for node in graph.node: - names.extend(name for name in node.input if name) - for attribute in node.attribute: - if attribute.type == AttributeProto.GRAPH: - names.extend(_referenced_names(attribute.g)) - elif attribute.type == AttributeProto.GRAPHS: - for child in attribute.graphs: - names.extend(_referenced_names(child)) - return names - - -def fold_constant_pad_pads(model: ModelProto) -> ModelProto: - """Fold constant integer Pad widths without specializing runtime input shapes. - - Only the main graph is rewritten. Nested graph captures and quantization - annotations conservatively protect shared producers from removal. - """ - versions = [item.version for item in model.opset_import if item.domain in {"", "ai.onnx"}] - if not versions or len(set(versions)) != 1 or versions[0] < 11: - return model - candidates = [ - (index, node) for index, node in enumerate(model.graph.node) - if node.domain in {"", "ai.onnx"} and node.op_type == "Pad" - and len(node.input) >= 2 and node.input[1] - ] - if not candidates: - return model - evaluator = _ConstantParameters(model) - types = { - value.name: value.type for value in - [*model.graph.input, *model.graph.value_info, *model.graph.output] - } - replacements: dict[int, np.ndarray] = {} - selected_dependencies: set[str] = set() - for index, node in candidates: - producer = evaluator.producers.get(node.input[1]) - if producer is None or producer.op_type == "Constant": - continue - try: - visited: set[str] = set() - pads = evaluator.evaluate(node.input[1], visited, set()) - if pads.dtype != np.int64 or pads.ndim != 1 or pads.size % 2: - continue - tensor_type = types.get(node.input[0]) - rank = ( - len(tensor_type.tensor_type.shape.dim) - if tensor_type is not None and tensor_type.tensor_type.HasField("shape") else None - ) - if len(node.input) > 3 and node.input[3]: - if versions[0] < 18: - continue - axes = evaluator.evaluate(node.input[3], visited, set()) - if axes.ndim != 1 or axes.dtype not in {np.dtype("int32"), np.dtype("int64")}: - continue - if pads.size != 2 * axes.size: - continue - if rank is not None: - if np.any(axes < -rank) or np.any(axes >= rank): - continue - if len({int(axis) % rank for axis in axes}) != axes.size: - continue - elif rank is not None and pads.size != 2 * rank: - continue - replacements[index] = pads - selected_dependencies.update(visited) - except ( - ValueError, KeyError, TypeError, IndexError, StopIteration, - NotImplementedError, OverflowError, - ): - logger.debug("Pad constant parameter is not foldable: %s", node.name, exc_info=True) - if not replacements: - return model - - rewritten = ModelProto() - rewritten.CopyFrom(model) - used_names = set(evaluator.producers) | set(evaluator.initializers) | evaluator.inputs - used_names.update(_referenced_names(model.graph)) - used_names.update(value.name for value in model.graph.value_info) - constants = [] - folded_names: dict[str, str] = {} - for index, value in replacements.items(): - node = rewritten.graph.node[index] - source = node.input[1] - if source not in folded_names: - name = source + "_folded_pads" - while name in used_names: - name += "_" - used_names.add(name) - folded_names[source] = name - constants.append(helper.make_node( - "Constant", [], [name], value=numpy_helper.from_array(value), - )) - node.input[1] = folded_names[source] - - references = Counter(_referenced_names(rewritten.graph)) - removable = { - node.output[0]: node for node in rewritten.graph.node - if len(node.output) == 1 - and node.output[0] in evaluator.evaluated & selected_dependencies - } - pending = deque(name for name in removable if not references[name]) - removed: set[str] = set() - while pending: - name = pending.popleft() - if name in removed: - continue - removed.add(name) - for source in removable[name].input: - references[source] -= 1 - if source in removable and not references[source]: - pending.append(source) - nodes = [node for node in rewritten.graph.node if not any(n in removed for n in node.output)] - del rewritten.graph.node[:] - rewritten.graph.node.extend([*constants, *nodes]) - infos = [value for value in rewritten.graph.value_info if value.name not in removed] - del rewritten.graph.value_info[:] - rewritten.graph.value_info.extend(infos) - logger.info("Folded constant pads for %d Pad node(s)", len(replacements)) - return rewritten - - def cgc_constant_folding(model: ModelProto) -> ModelProto: - """Fill FoundryToolbox constant-folding gaps for Pad and static Shape chains. + """Fill FoundryToolbox constant-folding gaps for static Shape chains. - Only the main graph is changed. Pad widths use the existing bounded folder; - graphs containing Shape also fold bounded integer/boolean constant chains + Only the main graph is changed. Graphs containing Shape fold bounded integer/boolean chains to a fixed point with shape inference. This never freezes symbolic input dimensions: callers must specialize inputs explicitly before this rule when required. Casts of runtime data, including floating-point outputs, remain. """ - prepared = fold_constant_pad_pads(model) if not any(node.op_type == "Shape" and node.domain in {"", "ai.onnx"} - for node in prepared.graph.node): - return prepared - versions = {item.version for item in prepared.opset_import if item.domain in {"", "ai.onnx"}} + for node in model.graph.node): + return model + versions = {item.version for item in model.opset_import if item.domain in {"", "ai.onnx"}} if len(versions) != 1 or next(iter(versions)) < 11: - return prepared + return model rewritten = ModelProto() - rewritten.CopyFrom(prepared) + rewritten.CopyFrom(model) changed = False for _iteration in range(32): for value_info in rewritten.graph.value_info: @@ -339,7 +206,7 @@ def cgc_constant_folding(model: ModelProto) -> ModelProto: )) folded += 1 if not folded: - return rewritten if changed else prepared + return rewritten if changed else model changed = True logger.info("CGC constant folding: folded %d shape/integer node(s)", folded) logger.warning("CGC constant folding reached the 32-round limit; retaining partial folding") diff --git a/src/winml/modelkit/pattern/cgc/identity_rewrites.py b/src/winml/modelkit/pattern/cgc/identity_rewrites.py deleted file mode 100644 index 665f949ab..000000000 --- a/src/winml/modelkit/pattern/cgc/identity_rewrites.py +++ /dev/null @@ -1,206 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""Eliminate internal tensor Identity aliases for IX compatibility.""" - -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -from onnx import AttributeProto, GraphProto, ModelProto, NodeProto, TensorProto, TypeProto, helper - -from ...onnx import ONNXDomain - - -if TYPE_CHECKING: - from collections.abc import Iterator - - -logger = logging.getLogger(__name__) - - -def _subgraphs(node: NodeProto) -> Iterator[GraphProto]: - for attribute in node.attribute: - if attribute.type == AttributeProto.GRAPH: - yield attribute.g - elif attribute.type == AttributeProto.GRAPHS: - yield from attribute.graphs - - -def _local_names(graph: GraphProto) -> set[str]: - return ( - {value.name for value in graph.input} - | {value.name for value in graph.initializer} - | {value.values.name for value in graph.sparse_initializer} - | {name for node in graph.node for name in node.output if name} - ) - - -def _types(graph: GraphProto, outer: dict[str, TypeProto]) -> dict[str, TypeProto]: - local = _local_names(graph) - types = {name: value for name, value in outer.items() if name not in local} - for value in (*graph.value_info, *graph.input, *graph.output): - types[value.name] = value.type - for tensor in graph.initializer: - types.setdefault(tensor.name, helper.make_tensor_type_proto(tensor.data_type, tensor.dims)) - return types - - -def _compatible_tensor(source: TypeProto | None, output: TypeProto | None) -> bool: - if source is None or not source.HasField("tensor_type"): - return False - if output is None: - return True - if not output.HasField("tensor_type"): - return False - left, right = source.tensor_type, output.tensor_type - if left.elem_type and right.elem_type and left.elem_type != right.elem_type: - return False - if left.HasField("shape") and right.HasField("shape"): - if len(left.shape.dim) != len(right.shape.dim): - return False - for a, b in zip(left.shape.dim, right.shape.dim, strict=True): - if a.HasField("dim_value") and b.HasField("dim_value") and a.dim_value != b.dim_value: - return False - if a.dim_param and b.dim_param and a.dim_param != b.dim_param: - return False - return True - - -def _redirect_uses( - graph: GraphProto, old: str, new: str, *, apply: bool, shadowed: bool = False -) -> bool: - if any(value.name == old for value in graph.output): - return False - # Quantization annotations can name aliases independently of node inputs. - if any( - annotation.tensor_name == old - or any(parameter.value == old for parameter in annotation.quant_parameter_tensor_names) - for annotation in graph.quantization_annotation - ): - return False - for node in graph.node: - for index, name in enumerate(node.input): - if name == old: - if shadowed: - return False - if apply: - node.input[index] = new - for child in _subgraphs(node): - local = _local_names(child) - if old in local: - continue - if not _redirect_uses( - child, old, new, apply=apply, shadowed=shadowed or new in local - ): - return False - if apply: - retained = [value for value in graph.value_info if value.name != old] - del graph.value_info[:] - graph.value_info.extend(retained) - return True - - -def _eliminate(graph: GraphProto, outer: dict[str, TypeProto]) -> int: - types = _types(graph, outer) - removed = 0 - for node in list(graph.node): - if node.domain not in {"", ONNXDomain.AI_ONNX.value} or node.op_type != "Identity": - continue - if len(node.input) != 1 or len(node.output) != 1 or node.attribute: - continue - source, output = node.input[0], node.output[0] - if not source or not output or source == output: - continue - if not _compatible_tensor(types.get(source), types.get(output)): - logger.debug( - "Retaining Identity %r: tensor types are unknown or incompatible", node.name, - ) - continue - if not _redirect_uses(graph, output, source, apply=False): - logger.debug( - "Retaining Identity %r: output interface or scoped alias is protected", node.name - ) - continue - _redirect_uses(graph, output, source, apply=True) - graph.node.remove(node) - types.pop(output, None) - removed += 1 - for node in graph.node: - for child in _subgraphs(node): - removed += _eliminate(child, types) - return removed - - -def _reshape_output_aliases(model: ModelProto) -> int: - versions = [entry.version for entry in model.opset_import if entry.domain == ""] - graph = model.graph - if len(versions) != 1 or versions[0] < 5 or any(list(_subgraphs(node)) for node in graph.node): - return 0 - types: dict[str, TypeProto] = {} - for value in (*graph.input, *graph.output, *graph.value_info): - if value.name in types and types[value.name] != value.type: - return 0 - types[value.name] = value.type - outputs = {value.name for value in graph.output} - inputs = {value.name for value in graph.input} - protected = set() - for annotation in graph.quantization_annotation: - protected.add(annotation.tensor_name) - protected.update(parameter.value for parameter in annotation.quant_parameter_tensor_names) - names = _local_names(graph) | set(types) | protected - names.update(name for node in graph.node for name in node.input) - rewritten = 0 - for node in graph.node: - if node.domain or node.op_type != "Identity" or node.attribute: - continue - if len(node.input) != 1 or len(node.output) != 1: - continue - source, output = node.input[0], node.output[0] - if (not source or source == output or output not in outputs or output in inputs - or source in protected or output in protected): - continue - source_type, output_type = types.get(source), types.get(output) - if (source_type is None or source_type != output_type - or not source_type.HasField("tensor_type")): - continue - tensor = source_type.tensor_type - if tensor.elem_type != TensorProto.FLOAT or not tensor.HasField("shape"): - continue - dimensions = tensor.shape.dim - if not dimensions or any(not dim.HasField("dim_value") or dim.dim_value <= 0 - for dim in dimensions): - continue - shape_name = output + "_identity_shape" - while shape_name in names: - shape_name += "_" - names.add(shape_name) - graph.initializer.append(helper.make_tensor( - shape_name, TensorProto.INT64, [len(dimensions)], - [dim.dim_value for dim in dimensions], - )) - node.op_type = "Reshape" - node.input.append(shape_name) - rewritten += 1 - return rewritten - - -def eliminate_identity(model: ModelProto) -> ModelProto: - """Remove safe internal tensor Identities, preserving graph IO and lexical bindings. - - Top-level static positive-shape FP32 output aliases use Reshape in models - without subgraphs. Other protected aliases and lexical bindings are retained. - """ - result = ModelProto() - result.CopyFrom(model) - reshaped = _reshape_output_aliases(result) - removed = _eliminate(result.graph, {}) - if not removed and not reshaped: - return model - logger.info( - "CGIR compatibility: eliminate-identity removed %d node(s), reshaped %d output alias(es)", - removed, reshaped, - ) - return result diff --git a/tests/unit/commands/test_export.py b/tests/unit/commands/test_export.py index a9eda7b73..6a614e283 100644 --- a/tests/unit/commands/test_export.py +++ b/tests/unit/commands/test_export.py @@ -151,7 +151,7 @@ def test_target_config_and_cli_option_precedence( elif config_flag == "both": build_path = tmp_path / "build.json" build_path.write_text(json.dumps({"export": settings})) - config_path.write_text(json.dumps({"options": {"topo_sort_nodes": False}})) + config_path.write_text(json.dumps({"options": {"topo_sort_nodes": True}})) args += ["-c", str(build_path), "--export-config", str(config_path)] else: args += [config_flag, str(config_path)] @@ -178,7 +178,7 @@ def test_target_config_and_cli_option_precedence( assert backend.call_args.args[0].options == CGCOptions( external_weights=not override and config_flag != "both", update_opset=override or config_flag == "both", - topo_sort_nodes=override or config_flag != "both", + topo_sort_nodes=not override and config_flag == "both", ) if not onnx_input: config = backend.call_args.kwargs["export_config"] diff --git a/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py b/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py index 6d9a578c2..60c0212c6 100644 --- a/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py +++ b/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py @@ -129,19 +129,26 @@ def test_normalize_int32_dq_preserves_unsupported(guard): @pytest.mark.parametrize("opset", [11, 17, 18]) @pytest.mark.parametrize("shared", [False, True]) -def test_fold_constant_pad_pads(opset, shared): - from winml.modelkit.pattern.cgc import fold_constant_pad_pads +@pytest.mark.parametrize("has_shape", [False, True]) +def test_cgc_constant_folding_pad_parameters(opset, shared, has_shape): + from winml.modelkit.pattern.cgc import cgc_constant_folding seed = np.random.default_rng(42) widths = seed.integers(0, 3, size=4, dtype=np.int32) + nodes = [ + helper.make_node("Cast", ["widths"], ["pads"], to=TensorProto.INT64), + helper.make_node("Pad", ["source", "pads"], ["result"]), + ] outputs = [helper.make_tensor_value_info("result", TensorProto.FLOAT, [None, None])] + if has_shape: + nodes.append(helper.make_node("Shape", ["result"], ["shape"])) + outputs.append(helper.make_tensor_value_info("shape", TensorProto.INT64, [2])) if shared: outputs.append(helper.make_tensor_value_info("pads", TensorProto.INT64, [4])) model = helper.make_model( helper.make_graph( - [helper.make_node("Cast", ["widths"], ["pads"], to=TensorProto.INT64), - helper.make_node("Pad", ["source", "pads"], ["result"])], - "constant_pad", [helper.make_tensor_value_info("source", TensorProto.FLOAT, [2, 3])], + nodes, "constant_pad", + [helper.make_tensor_value_info("source", TensorProto.FLOAT, [2, 3])], outputs, [numpy_helper.from_array(widths, "widths")], ), opset_imports=[helper.make_opsetid("", opset)], ir_version=10, @@ -150,14 +157,22 @@ def test_fold_constant_pad_pads(opset, shared): feeds = {"source": seed.normal(size=(2, 3)).astype(np.float32)} expected = ReferenceEvaluator(model).run(None, feeds) result = CGIRRewritePipe().process( - model, CGIRRewritePipe.build_config(fold_constant_pad_pads=True), + model, CGIRRewritePipe.build_config(cgc_constant_folding=True), ) checker.check_model(result) assert model.SerializeToString() == original assert list(result.graph.input) == list(model.graph.input) - assert list(result.graph.output) == list(model.graph.output) - assert sum(node.op_type == "Cast" for node in result.graph.node) == int(shared) - assert fold_constant_pad_pads(result) is result + assert [(value.name, value.type.tensor_type.elem_type) for value in result.graph.output] == [ + (value.name, value.type.tensor_type.elem_type) for value in model.graph.output + ] + if has_shape: + for output, reference in zip(result.graph.output, expected, strict=True): + assert tuple(dim.dim_value for dim in output.type.tensor_type.shape.dim) == ( + reference.shape + ) + assert (result is model) == (not has_shape) + assert sum(node.op_type == "Cast" for node in result.graph.node) == int(not has_shape) + assert cgc_constant_folding(result) is result for actual, reference in zip( ReferenceEvaluator(result).run(None, feeds), expected, strict=True, ): @@ -234,141 +249,12 @@ def test_gridsample_to_gather_preserves_unsupported_modes(mode, padding): assert result is model -def _constant_pad_chain(*, axes=False): - rank = 2 - nodes = [] - constants = { - "count": np.asarray([rank], dtype=np.int64), - "widths": np.random.default_rng(7).integers(0, 3, rank, dtype=np.int64), - "matrix": np.asarray([-1, 2], dtype=np.int64), - "start": np.asarray([-1], dtype=np.int64), - "end": np.asarray([np.iinfo(np.int64).min], dtype=np.int64), - "axis": np.asarray([0], dtype=np.int64), - "step": np.asarray([-1], dtype=np.int64), - "vector": np.asarray([-1], dtype=np.int64), - } - for name, value in constants.items(): - nodes.append(helper.make_node("Constant", [], [name], value=numpy_helper.from_array(value))) - nodes.extend([ - helper.make_node( - "ConstantOfShape", ["count"], ["zeros"], - value=numpy_helper.from_array(np.zeros(1, dtype=np.int64)), - ), - helper.make_node("Concat", ["widths", "zeros"], ["joined"], axis=0), - helper.make_node("Reshape", ["joined", "matrix"], ["pairs"]), - helper.make_node("Slice", ["pairs", "start", "end", "axis", "step"], ["reversed"]), - helper.make_node("Transpose", ["reversed"], ["transposed"], perm=[1, 0]), - helper.make_node("Reshape", ["transposed", "vector"], ["flattened"]), - helper.make_node("Cast", ["flattened"], ["pads"], to=TensorProto.INT64), - ]) - pad_inputs = ["source", "pads"] - if axes: - nodes.append(helper.make_node( - "Constant", [], ["pad_axes"], - value=numpy_helper.from_array(np.arange(rank, dtype=np.int64)), - )) - pad_inputs.extend(["", "pad_axes"]) - nodes.append(helper.make_node("Pad", pad_inputs, ["result"])) - return helper.make_model( - helper.make_graph( - nodes, "pad_chain", - [helper.make_tensor_value_info("source", TensorProto.FLOAT, [2, 3])], - [helper.make_tensor_value_info("result", TensorProto.FLOAT, [None, None])], - ), opset_imports=[helper.make_opsetid("", 18 if axes else 17)], ir_version=10, - ) - - -@pytest.mark.parametrize("axes", [False, True]) -@pytest.mark.parametrize("protected", ["none", "capture", "annotation"]) -def test_pad_constant_chain_preserves_values_and_references(axes, protected): - from winml.modelkit.pattern.cgc import fold_constant_pad_pads - - model = _constant_pad_chain(axes=axes) - if protected == "capture": - model.graph.input.append(helper.make_tensor_value_info("condition", TensorProto.BOOL, [])) - branch = helper.make_graph( - [helper.make_node("Identity", ["pads"], ["captured"])], "capture", [], - [helper.make_tensor_value_info("captured", TensorProto.INT64, [4])], - ) - model.graph.node.append(helper.make_node( - "If", ["condition"], ["observed"], then_branch=branch, else_branch=branch, - )) - model.graph.output.append(helper.make_tensor_value_info("observed", TensorProto.INT64, [4])) - if protected == "annotation": - annotation = model.graph.quantization_annotation.add(tensor_name="source") - annotation.quant_parameter_tensor_names.add(key="SCALE_TENSOR", value="pads") - feeds = {"source": np.random.default_rng(21).normal(size=(2, 3)).astype(np.float32)} - if protected == "capture": - feeds["condition"] = np.asarray(True) - expected = ReferenceEvaluator(model).run(None, feeds) - original = model.SerializeToString() - result = fold_constant_pad_pads(model) - checker.check_model(result) - assert result is not model - assert original == model.SerializeToString() - assert any("pads" in node.output for node in result.graph.node) == (protected != "none") - assert fold_constant_pad_pads(result) is result - for actual, reference in zip( - ReferenceEvaluator(result).run(None, feeds), expected, strict=True, - ): - np.testing.assert_array_equal(actual, reference) - - -@pytest.mark.parametrize("reason", [ - "runtime", "overridable", "unsupported", "float", "length", "domain", - "allocation", "nodes", "cache", "invalid_axes", "direct", -]) -def test_pad_constant_folding_rejects_unsafe_candidates(reason, monkeypatch): - from winml.modelkit.pattern.cgc import fold_constant_pad_pads - - folding_module = import_module("winml.modelkit.pattern.cgc.cgc_constant_folding") - - model = _constant_pad_chain(axes=reason == "invalid_axes") - producers = {name: node for node in model.graph.node for name in node.output} - if reason in {"runtime", "overridable"}: - model.graph.node.remove(producers["widths"]) - model.graph.input.append(helper.make_tensor_value_info("widths", TensorProto.INT64, [2])) - if reason == "overridable": - model.graph.initializer.append(numpy_helper.from_array(np.zeros(2, np.int64), "widths")) - elif reason == "unsupported": - producers["pads"].CopyFrom(helper.make_node("Identity", ["flattened"], ["pads"])) - elif reason == "float": - producers["pads"].attribute[0].i = TensorProto.FLOAT - elif reason == "length": - model.graph.input[0].type.tensor_type.shape.dim.add(dim_value=2) - elif reason == "domain": - model.graph.node[-1].domain = "custom" - elif reason == "allocation": - producers["count"].attribute[0].t.CopyFrom( - numpy_helper.from_array(np.asarray([folding_module._MAX_ELEMENTS + 1], np.int64)), - ) - elif reason == "nodes": - monkeypatch.setattr(folding_module, "_MAX_NODES", 2) - elif reason == "cache": - monkeypatch.setattr(folding_module, "_MAX_CACHED_ELEMENTS", 1) - elif reason == "invalid_axes": - producers["pad_axes"].attribute[0].t.CopyFrom( - numpy_helper.from_array(np.zeros(2, np.int64)), - ) - elif reason == "direct": - producers["pads"].CopyFrom(helper.make_node( - "Constant", [], ["pads"], value=numpy_helper.from_array(np.zeros(4, np.int64)), - )) - original = model.SerializeToString() - assert fold_constant_pad_pads(model) is model - assert original == model.SerializeToString() - - -def test_pad_folding_is_enabled_only_by_cgc_defaults(): +def test_cgc_constant_folding_is_enabled_only_by_cgc_defaults(): from winml.modelkit.optim import WinMLOptimizationConfig assert not CGIRRewritePipe.build_config().rules assert WinMLOptimizationConfig.for_cgc()["cgc_constant_folding"] is True - canonical = CGIRRewritePipe.build_config(cgc_constant_folding=True) - assert canonical == CGIRRewritePipe.build_config(fold_constant_pad_pads=True) - assert canonical == CGIRRewritePipe.build_config( - cgc_constant_folding=True, fold_constant_pad_pads=True, - ) + assert len(CGIRRewritePipe.build_config(cgc_constant_folding=True).rules) == 1 @pytest.mark.parametrize("dynamic", [False, True]) @@ -480,7 +366,7 @@ def match(self): prelu_rule = CGIRRewritePipe.build_config(prelu_to_relu=True).rules[0] rules = [resize_rule] if model_barrier: - rules.extend(CGIRRewritePipe.build_config(eliminate_identity=True).rules) + rules.extend(CGIRRewritePipe.build_config(deduplicate_opset_imports=True).rules) rules.extend([prelu_rule, resize_rule]) original = model.SerializeToString() feeds = {"source": np.exp(np.random.default_rng(42).normal(size=(2, 3))).astype(np.float32)} @@ -506,6 +392,8 @@ def match(self): "log-to-reduce-log-sum", "materialize-initializer-parameters", "fold-scalar-initializer-casts", + "eliminate-identity", + "fold-constant-pad-pads", ]) def test_retired_cgir_rules_are_not_registered(capability): from click.testing import CliRunner @@ -522,114 +410,6 @@ def test_retired_cgir_rules_are_not_registered(capability): assert f"--disable-{capability}" not in result.output -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.parametrize("protected_output", [False, True]) -def test_eliminate_identity_preserves_results_and_model( - enabled: bool, protected_output: bool, -) -> None: - shape = [2, 3] - nodes = [ - helper.make_node("Identity", ["source"], ["alias"]), - helper.make_node("Identity", ["alias"], ["second_alias"]), - helper.make_node("Add", ["second_alias", "source"], ["result"]), - ] - outputs = [helper.make_tensor_value_info("result", TensorProto.FLOAT, shape)] - if protected_output: - outputs.append(helper.make_tensor_value_info("alias", TensorProto.FLOAT, shape)) - model = helper.make_model( - helper.make_graph( - nodes, "aliases", - [helper.make_tensor_value_info("source", TensorProto.FLOAT, shape)], - outputs, - value_info=[ - helper.make_tensor_value_info("alias", TensorProto.FLOAT, shape), - helper.make_tensor_value_info("second_alias", TensorProto.FLOAT, shape), - ], - ), - opset_imports=[helper.make_opsetid("", 18)], - ir_version=11, - ) - original = model.SerializeToString() - feeds = {"source": np.random.default_rng(42).standard_normal(shape).astype(np.float32)} - expected = ReferenceEvaluator(model).run(None, feeds) - - result = CGIRRewritePipe().process( - model, CGIRRewritePipe.build_config(eliminate_identity=enabled), - ) - - checker.check_model(result) - assert model.SerializeToString() == original - assert result.graph.input == model.graph.input - assert result.graph.output == model.graph.output - retained = [node for node in result.graph.node if node.op_type == "Identity"] - assert len(retained) == (0 if enabled else 2) - assert sum(node.op_type == "Reshape" for node in result.graph.node) == int( - enabled and protected_output, - ) - for actual, reference in zip( - ReferenceEvaluator(result).run(None, feeds), expected, strict=True, - ): - np.testing.assert_array_equal(actual, reference) - repeated = CGIRRewritePipe().process( - result, CGIRRewritePipe.build_config(eliminate_identity=enabled), - ) - assert repeated.SerializeToString() == result.SerializeToString() - - -@pytest.mark.parametrize("guard", ["none", "dynamic", "zero", "scalar", "fp16", "annotation", - "subgraph", "mismatch", "unknown", "opset"]) -def test_identity_output_scope_and_bits(guard): - from winml.modelkit.pattern.cgc import eliminate_identity - - shape = {"dynamic": ["batch"], "zero": [0], "scalar": []}.get(guard, [8]) - dtype = TensorProto.FLOAT16 if guard == "fp16" else TensorProto.FLOAT - model = helper.make_model(helper.make_graph( - [helper.make_node("Identity", ["source"], ["result"])], "output_alias", - [helper.make_tensor_value_info("source", dtype, shape)], - [helper.make_tensor_value_info("result", dtype, shape)], - ), opset_imports=[helper.make_opsetid("", 4 if guard == "opset" else 18)], ir_version=10) - if guard == "annotation": - model.graph.quantization_annotation.add(tensor_name="result") - elif guard == "subgraph": - branch = helper.make_graph([], "branch", [], []) - model.graph.node.append(helper.make_node("If", ["cond"], ["other"], then_branch=branch)) - elif guard == "mismatch": - model.graph.value_info.append(helper.make_tensor_value_info("source", dtype, [9])) - elif guard == "unknown": - model.graph.input[0].type.tensor_type.ClearField("shape") - original = model.SerializeToString() - result = eliminate_identity(model) - assert model.SerializeToString() == original - if guard != "none": - assert result is model - return - checker.check_model(result, full_check=True) - assert result.graph.output == model.graph.output - assert result.graph.input == model.graph.input - assert result.graph.node[0].op_type == "Reshape" - assert eliminate_identity(result) is result - random = np.random.default_rng(42) - values = [random.normal(size=shape).astype(np.float32)] - values.extend(np.full(shape, special, dtype=np.float32) for special in ( - 0.0, -0.0, np.inf, -np.inf, np.nan, - np.nextafter(np.float32(0), np.float32(1)), - )) - for data in values: - expected = ReferenceEvaluator(model).run(None, {"source": data})[0] - actual = ReferenceEvaluator(result).run(None, {"source": data})[0] - assert actual.tobytes() == expected.tobytes() - - - - - - - - - - - - def _make_resize_model( *, opset: int,