From d10efff3f4373faa4af9dc03f0123ab9cfb22828 Mon Sep 17 00:00:00 2001 From: Zhenchao Ni Date: Mon, 21 Sep 2026 15:27:16 +0800 Subject: [PATCH 1/2] Remove pad and identity rewrite (#1432) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Remove two CGC compatibility workarounds following upstream IX fixes already included in the current compiler version (`35c01c26`). | Removed workaround | Upstream issue | Fix | |---|---|---| | `eliminate-identity` | microsoft/ix#1198 — ONNX Identity fails to lower to Foundry and DXCGC | microsoft/ix#1201 adds native Identity lowering. | | Pad-specific constant folding (`fold_constant_pad_pads`) and the `fold-constant-pad-pads` compatibility alias | microsoft/ix#1199 — Support constant folding of ONNX Pad parameters before lowering | microsoft/ix#1206 improves Cast/Pad chain support, static shape inference, and lowering type compatibility. | ## Changes - Remove the Identity rewrite implementation, capability registration, and public export. - Remove the Pad-specific folding prepass and its dedicated dead-producer cleanup. - Remove the `fold-constant-pad-pads` alias and public function export. - Update documentation and tests to reflect the retired rules. - Preserve `cgc-constant-folding` and its shared integer/boolean expression evaluator. General folding remains necessary for static shape chains and continues to fold Pad parameters in graphs containing `Shape`. - Leave compiler options, including the default-enabled topological sorting, unchanged. ## Validation - 203 affected unit tests passed. - All five models that previously used the Pad-specific prepass retained identical Pad parameter values. - GPU outputs before and after Pad-prepass removal were elementwise identical across three input samples per model. - The three SOD models required topological sorting to be disabled for the GPU comparison because the current wheel does not yet include the fix for microsoft/ix#1207. Their default-path conversion failure remains unchanged. - End-to-end SOD validation through JSON config, build, perf, and eval is still in progress. (cherry picked from commit 855fe5c4dfaf2f501d4aacdfb0460939a813c26e) --- docs/commands/optimize.md | 3 +- .../optim/pipes/cgir_rewrite_rules.py | 26 +- src/winml/modelkit/pattern/cgc/__init__.py | 5 +- .../pattern/cgc/cgc_constant_folding.py | 153 +--------- .../modelkit/pattern/cgc/identity_rewrites.py | 206 ------------- .../optim/pipes/test_pipe_cgir_rewrite.py | 278 ++---------------- 6 files changed, 42 insertions(+), 629 deletions(-) delete mode 100644 src/winml/modelkit/pattern/cgc/identity_rewrites.py 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/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/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, From 3826ba5509e6952b0f5bcfeecb1f1a57d388543c Mon Sep 17 00:00:00 2001 From: Zhenchao Ni Date: Tue, 22 Sep 2026 13:53:55 +0800 Subject: [PATCH 2/2] Turn off toposort by default (#1434) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Disable topological sorting by default for CGIR conversion to work around microsoft/ix#1207, whose fix is not included in the compiler shipped with the current Windows ML wheel. This avoids conversion failures in the three P1 models and aligns with IX OnnxConverter’s default. Sorting remains available through `topo-sort-nodes=true`. (cherry picked from commit ff977dd98c68d6f1f520a69ea33e99cb4279cb86) --- src/winml/modelkit/export/cgc/exporter.py | 2 +- tests/unit/commands/test_export.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) 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/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"]