Cherry pick cgir fixes to main

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. |

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 = ""
 

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",

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",
 ]

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,15 +149,14 @@ 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()],
                         ),
                         opset_imports=list(self.model.opset_import),
                         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")