Migrate mobius rewrites

docs/source/features/onnx-transformations.md

@@ -119,6 +119,117 @@ Olive provides ability to apply many graph `surgeries` on the ONNX model. In the
 }
 ```
 
+### Transformer fusions and compatibility lowerings
+
+The following surgeons accept ONNX graphs directly and do not depend on Mobius
+or another exporter. They are registered automatically when `GraphSurgeries` is
+loaded; no implementation-module imports are required.
+
+These are explicit transformations, not an automatic execution-provider profile.
+Choose the surgeries, their order, and the model component according to the
+target runtime's operator and dtype support. A fusion that emits an ORT contrib
+operator is not a portable standard-ONNX optimization.
+
+| Surgeon | Transformation |
+| --- | --- |
+| `FuseGelu` | Exact or tanh-approximate Gelu decomposition to standard ONNX `Gelu`; requires opset 20 or newer. |
+| `FuseBiasGelu` | Compatible 1-D bias Add followed by **exact** Gelu to `com.microsoft::BiasGelu`. Tanh Gelu is not equivalent and is left unchanged. |
+| `FuseLayerNormalization` | Last-axis ReduceMean-based normalization to standard `LayerNormalization`, with or without bias. |
+| `FuseSkipLayerNormalization` | Compatible residual Add and last-axis `LayerNormalization` to `com.microsoft::SkipLayerNormalization`. |
+| `FuseSkipRMSNormalization` | Compatible residual Add and last-axis `RMSNormalization` to `com.microsoft::SkipSimplifiedLayerNormalization`. |
+| `AttentionToGroupQueryAttention` | Recognized causal standard `Attention`, optionally with RoPE, to `com.microsoft::GroupQueryAttention`. |
+| `PackQKVForGroupQueryAttention` | Separate constant-weight Q/K/V projections to a packed GQA projection, including unequal Q/KV widths and bias. |
+| `SeparateGroupQueryAttentionRoPE` | Move supported GQA-integrated RoPE into separate standard `RotaryEmbedding` nodes. |
+| `UnpackGroupQueryAttentionQKV` | Split supported packed GQA projections into separate Q/K/V projections. |
+| `BlockDiagonalAttentionToPackedMHA` | Recognized block-diagonal mask and standard `Attention` to `com.microsoft::PackedMultiHeadAttention`. |
+| `ClipToMinMax` | BF16 `Clip` with both bounds, only a lower bound, or only an upper bound to `Max`/`Min`. |
+| `Rank4RMSNormToRank3` | Rank-4 last-axis RMSNorm with static head dimensions to rank-3 RMSNorm surrounded by reshapes. |
+| `DecomposeOnnxRotaryEmbedding` | Standard rank-3, full-width, non-interleaved `RotaryEmbedding` to primitive rotate-half operations. |
+| `TensorScatterToScatterND` | Linear rank-3 static-cache writes with batch size 1 and known cache capacity to `ScatterND`; an omitted write index starts at zero. |
+| `DecomposeAttention` | Supported standard rank-3 `Attention` to scaled dot-product primitives, including GQA, cache outputs, and causal/nonpadding masks. |
+| `StaticEmptyKV` | Recognized dynamic empty-KV construction to a static empty tensor for graph-capture compatibility. |
+| `FuseDenseMoEToQMoE` | Compatible `MatMulNBits` expert banks and routing to `com.microsoft::QMoE`. |
+| `FuseBlockQuantizedMoE` | Compatible native block-quantized expert banks; requires a runtime implementing that operator. |
+
+The pattern surgeons preserve graphs that do not match their supported forms.
+In particular, `DecomposeAttention` leaves a fourth QK output, an explicit
+`softmax_precision`, and masks whose full KV width cannot be established
+unchanged. It is not a general-purpose decomposition of every legal Attention
+configuration. Check the resulting graph for unsupported operators before
+deploying to a runtime without an Attention kernel.
+
+`DecomposeOnnxRotaryEmbedding` is distinct from the existing
+[`DecomposeRotaryEmbedding`](#decomposerotaryembedding), which accepts the
+Microsoft-domain four-input ABI. Likewise, `BlockDiagonalAttentionToPackedMHA`
+accepts a standard Attention subgraph, not the `custom::PackedAttention` input
+expected by `PackedAttentionToPackedMHA`.
+
+#### Ordering
+
+`GraphSurgeries` runs the list in order, once per surgeon. For example, on a
+runtime that supports GQA and packed QKV:
+
+```json
+{
+    "type": "GraphSurgeries",
+    "surgeries": [
+        {"surgeon": "AttentionToGroupQueryAttention"},
+        {"surgeon": "PackQKVForGroupQueryAttention"},
+        {"surgeon": "FuseSkipRMSNormalization"},
+        {"surgeon": "FuseSkipLayerNormalization"},
+        {"surgeon": "FuseGelu"},
+        {"surgeon": "FuseBiasGelu"}
+    ]
+}
+```
+
+`FuseGelu` must precede `FuseBiasGelu` when the source Gelu is decomposed.
+Packing must follow GQA fusion. For a target requiring separate RoPE and Q/K/V,
+apply `SeparateGroupQueryAttentionRoPE` followed by
+`UnpackGroupQueryAttentionQKV` instead of packing. These lists are not universal
+EP recipes: kernel availability also depends on the model dtype and component.
+The surgeries do not implicitly run shape inference or a whole-model optimizer.
+Projection packing may leave constant `Concat`/`Transpose` nodes, and traced
+subgraphs may need unused-node cleanup; schedule an appropriate optimization
+pass separately when required.
+
+#### Weight-aware MoE surgeries
+
+Run MoE surgeries **after** quantization or native-block import has supplied the
+required weights. Both support loaded external initializers. `GraphSurgeries`
+retains its normal external-data output options such as `save_as_external_data`.
+The surgeries validate representable expert layouts before emitting replacements;
+they do not silently drop an unrecognized expert.
+
+`FuseBlockQuantizedMoE` defaults to `allow_dense_moe=false`: an identified
+native-block MoE layer that cannot be fused raises `MoEGraphSurgeryError`, rather
+than quietly retaining an expensive all-expert path. Setting
+`allow_dense_moe=true` explicitly retains that path with a warning:
+
+```json
+{
+    "type": "GraphSurgeries",
+    "surgeries": [
+        {"surgeon": "FuseBlockQuantizedMoE", "allow_dense_moe": false}
+    ],
+    "save_as_external_data": true
+}
+```
+
+`FuseDenseMoEToQMoE` requires compatible 4-bit geometry, complete expert banks,
+and consistent zero-point presence within each FC bank. FC1 and FC2 may
+independently have zero points. Malformed recognized groups raise
+`MoEGraphSurgeryError`; unsupported QMoE ABI geometry is left unchanged with a
+warning. Neither surgeon chooses an execution provider or quantizes weights.
+
+#### Adding a surgeon
+
+Implement new transformations in a focused module under
+`olive/passes/onnx/graph_surgery/`, importing `Surgeon` or
+`RewriteRuleSurgeon` from `graph_surgery.base`. Export the class from the
+package's `__init__.py` to register it for normal `GraphSurgeries` use.
+Existing imports of the base classes from `graph_surgeries` remain supported.
+
 ### `RenameInputs`
 
 #### Description

olive/passes/onnx/graph_surgeries.py

@@ -34,6 +34,7 @@
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
 from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
+from olive.passes.onnx.graph_surgery import ProtoSurgeon, RewriteRuleSurgeon, Surgeon
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 if TYPE_CHECKING:
@@ -49,93 +50,6 @@
 # pylint: disable=W0621
 
 
-class Surgeon:
-    """Base class for surgeons that operate on the ONNX IR model."""
-
-    # Refer to https://microsoft.github.io/onnxscript/intermediate_representation/ir_api.html#onnxscript.ir.Model
-    # for the IR model API.
-
-    registry: ClassVar[dict[str, type[Surgeon]]] = {}
-
-    @classmethod
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        Surgeon.registry[cls.__name__.lower()] = cls
-
-    def __init__(self):
-        pass
-
-    def __call__(self, model: ModelProto) -> ModelProto:
-        return ir.to_proto(self.call_ir(ir.from_proto(model)))
-
-    def call_ir(self, model: ir.Model) -> ir.Model:
-        # Implement this method in subclasses to operate on the IR model.
-        raise NotImplementedError
-
-
-class ProtoSurgeon(Surgeon):
-    """Base class for surgeons that operate on the ONNX model proto directly."""
-
-    def __call__(self, model: ModelProto) -> ModelProto:
-        raise NotImplementedError
-
-    def call_ir(self, model: ir.Model) -> ir.Model:
-        raise RuntimeError("Implement __call__ method instead of operator on onnx.ModelProto directly.")
-
-    @staticmethod
-    def get_node_by_name(model, name: str, match_output: bool = False):
-        for node in model.graph.node:
-            if (match_output and node.output[0] == name) or (not match_output and node.name == name):
-                return node
-        return None
-
-    @staticmethod
-    def get_tensor_shapes(model) -> dict[str, list[int]]:
-        return {info.name: [x.dim_value for x in info.type.tensor_type.shape.dim] for info in model.graph.value_info}
-
-    @staticmethod
-    def get_tensor_types(model):
-        return {info.name: info.type.tensor_type.elem_type for info in model.graph.value_info}
-
-    @staticmethod
-    def get_initializer_types(model):
-        return {initializer.name: initializer.data_type for initializer in model.graph.initializer}
-
-    @staticmethod
-    def get_initializer_shapes(model) -> dict[str, list[int]]:
-        return {initializer.name: initializer.dims for initializer in model.graph.initializer}
-
-    @staticmethod
-    def get_initializer_by_name(model, name: str):
-        for initializer in model.graph.initializer:
-            if initializer.name == name:
-                return initializer
-        return None
-
-    @staticmethod
-    def create_new_name(name: str, old_op: str, new_op: str) -> str:
-        return name.replace(old_op, new_op) if old_op in name else f"{name}_{new_op}"
-
-
-class RewriteRuleSurgeon(Surgeon):
-    """Base class for surgeons implemented as onnxscript rewrite rules.
-
-    Subclasses implement :meth:`rules` to return an
-    :class:`onnxscript.rewriter.pattern.RewriteRuleSet`, expressing the match
-    pattern and its replacement with the ONNX IR op builder. ``call_ir`` applies
-    the rules to the IR model in place. Prefer this over manual proto/DAG
-    manipulation for local subgraph pattern replacements: the rewriter handles
-    operand commutativity, use-count bookkeeping, and dead-node cleanup.
-    """
-
-    def rules(self) -> pattern.RewriteRuleSet:
-        raise NotImplementedError
-
-    def call_ir(self, model: ir.Model) -> ir.Model:
-        self.rules().apply_to_model(model)
-        return model
-
-
 # TODO(anyone): This is incorrect, remove or fix
 class RenameInputs(Surgeon):
     def __init__(self, old_names: list[str], new_names: list[str]):

olive/passes/onnx/graph_surgery/__init__.py

@@ -0,0 +1,54 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Built-in surgeons, registered on import for the GraphSurgeries pass."""
+
+from olive.passes.onnx.graph_surgery.activations import FuseBiasGelu, FuseGelu
+from olive.passes.onnx.graph_surgery.attention import (
+    AttentionToGroupQueryAttention,
+    BlockDiagonalAttentionToPackedMHA,
+    PackQKVForGroupQueryAttention,
+    SeparateGroupQueryAttentionRoPE,
+    UnpackGroupQueryAttentionQKV,
+)
+from olive.passes.onnx.graph_surgery.base import ProtoSurgeon, RewriteRuleSurgeon, Surgeon
+from olive.passes.onnx.graph_surgery.lowering import (
+    ClipToMinMax,
+    DecomposeAttention,
+    DecomposeOnnxRotaryEmbedding,
+    Rank4RMSNormToRank3,
+    StaticEmptyKV,
+    TensorScatterToScatterND,
+)
+from olive.passes.onnx.graph_surgery.moe import FuseBlockQuantizedMoE, FuseDenseMoEToQMoE, MoEGraphSurgeryError
+from olive.passes.onnx.graph_surgery.normalization import (
+    FuseLayerNormalization,
+    FuseSkipLayerNormalization,
+    FuseSkipRMSNormalization,
+)
+
+__all__ = [
+    "AttentionToGroupQueryAttention",
+    "BlockDiagonalAttentionToPackedMHA",
+    "ClipToMinMax",
+    "DecomposeAttention",
+    "DecomposeOnnxRotaryEmbedding",
+    "FuseBiasGelu",
+    "FuseBlockQuantizedMoE",
+    "FuseDenseMoEToQMoE",
+    "FuseGelu",
+    "FuseLayerNormalization",
+    "FuseSkipLayerNormalization",
+    "FuseSkipRMSNormalization",
+    "MoEGraphSurgeryError",
+    "PackQKVForGroupQueryAttention",
+    "ProtoSurgeon",
+    "Rank4RMSNormToRank3",
+    "RewriteRuleSurgeon",
+    "SeparateGroupQueryAttentionRoPE",
+    "StaticEmptyKV",
+    "Surgeon",
+    "TensorScatterToScatterND",
+    "UnpackGroupQueryAttentionQKV",
+]

olive/passes/onnx/graph_surgery/_common.py

@@ -0,0 +1,44 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import onnx_ir as ir
+
+
+def check_scalar_constant(value: ir.Value, expected: float, name: str, *, rel_tol: float = 1e-4) -> str | None:
+    if value.const_value is None:
+        return f"{name} is not a constant"
+    array = value.const_value.numpy()
+    if array.size != 1:
+        return f"{name} must contain exactly one element"
+    actual = float(array.flat[0])
+    if not math.isclose(actual, expected, rel_tol=rel_tol):
+        return f"{name} is {actual}, expected approximately {expected}"
+    return None
+
+
+class ReplacedAddCleanupMixin:
+    """Clean up traced Adds outside the matched pattern after rewriting."""
+
+    _replaced_adds: list[ir.Node]
+
+    def setup(self):
+        self._replaced_adds = []
+
+    def _record_replaced_add(self, add):
+        self._replaced_adds.append(add)
+
+    def cleanup(self):
+        for add in self._replaced_adds:
+            if (
+                add.graph is not None
+                and all(output not in add.graph.outputs for output in add.outputs)
+                and all(not list(output.uses()) for output in add.outputs)
+            ):
+                add.graph.remove(add, safe=True)