feat(optim): add algebraic graph rewrites

src/winml/modelkit/optim/capabilities/__init__.py

@@ -14,6 +14,7 @@
 
 from . import (
     activation,
+    algebraic,
     attention,
     conv,
     elimination,
@@ -30,6 +31,7 @@
 
 __all__ = [
     "activation",
+    "algebraic",
     "attention",
     "conv",
     "elimination",

src/winml/modelkit/optim/capabilities/algebraic.py

@@ -0,0 +1,43 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Opt-in, exact algebraic graph-rewrite capabilities."""
+
+from __future__ import annotations
+
+from ..registry import BoolCapability, CapabilityCategory
+
+
+STATIC_SPLIT_TO_SLICE = BoolCapability(
+    name="static-split-to-slice",
+    ort_name=None,
+    description=(
+        "Replace statically bounded Split operations with standard Slice operations "
+        "while preserving output tensors"
+    ),
+    category=CapabilityCategory.REWRITE,
+    default=False,
+)
+
+CONV_CHANNEL_AFFINE_FOLDING = BoolCapability(
+    name="conv-channel-affine-folding",
+    ort_name=None,
+    description=(
+        "Fold safe scalar or channel-wise constant Mul/Add branches into Conv "
+        "weights and bias; direct and static Split/Slice routing only"
+    ),
+    category=CapabilityCategory.REWRITE,
+    default=False,
+)
+
+CONV_ADD_BATCH_NORMALIZATION_FOLDING = BoolCapability(
+    name="conv-add-batch-normalization-folding",
+    ort_name=None,
+    description=(
+        "Fold inference-mode BatchNormalization after Conv plus a static Add "
+        "tensor into Conv parameters"
+    ),
+    category=CapabilityCategory.REWRITE,
+    default=False,
+)

src/winml/modelkit/optim/pipes/__init__.py

@@ -11,6 +11,11 @@
 
 from typing import Any
 
+from .algebraic import (
+    ALGEBRAIC_CAPABILITIES,
+    AlgebraicRewritePipe,
+    AlgebraicRewritePipeConfig,
+)
 from .base import BasePipe, OptimizationError, PipeConfig, caps_dict
 from .fusion import ORTFusionPipe, ORTFusionPipeConfig
 from .graph import GRAPH_CAPABILITIES, ORTGraphPipe, ORTGraphPipeConfig
@@ -22,12 +27,19 @@
 # - ORTGraphPipe: ORT graph-level optimizations (C++ optimizer), including constant folding.
 #   Runs first so downstream pipes see a constant-folded graph (e.g. Reshape shape inputs
 #   become literal constants, enabling skeleton-based pattern matching).
+# - AlgebraicRewritePipe: Exact topology-based algebraic rewrites (after ORT folding).
 # - RewritePipe: Pattern-based subgraph rewriting (runs after ORT constant folding so that
 #   shape constants are visible, but before ORTFusionPipe so normalised patterns are
 #   available for transformer fusions).
 # - ORTFusionPipe: ORT transformer fusions (Python optimizer)
 # - SurgeryPipe: Post-optimization model surgery (runs last to clamp constants after folding)
-PIPES: list[type[BasePipe]] = [ORTGraphPipe, RewritePipe, ORTFusionPipe, SurgeryPipe]
+PIPES: list[type[BasePipe]] = [
+    ORTGraphPipe,
+    AlgebraicRewritePipe,
+    RewritePipe,
+    ORTFusionPipe,
+    SurgeryPipe,
+]
 
 
 def get_all_capabilities() -> dict[str, Any]:
@@ -43,9 +55,12 @@ def get_all_capabilities() -> dict[str, Any]:
 
 
 __all__ = [
+    "ALGEBRAIC_CAPABILITIES",
     "GRAPH_CAPABILITIES",
     "PIPES",
     "SURGERY_CAPABILITIES",
+    "AlgebraicRewritePipe",
+    "AlgebraicRewritePipeConfig",
     "BasePipe",
     "ORTFusionPipe",
     "ORTFusionPipeConfig",