Attention fusion (part 2) (#2013)

Continuation of attention fusion.

* Adds a version of GroupQueryAttention
* Adds support in Cos-Sin cache fusion for constant-folded position-ids
* Restructure MHA fusion into a class-based rewrite rule

Also restructure the folder structure.
* Eventually eliminate folders called "onnxruntime" and "transfomers",
which hinder importing the original packages with those names. For now
moving just the relevant new files. (Will restructure older files
later.)
* ORT-specific fusions go into the ort_fusions folder.

Author: gramalingam

.lintrunner.toml

@@ -50,7 +50,7 @@ exclude_patterns = [
     'onnxscript/optimizer/_legacy/constant_folding.py',  # FIXME
     'onnxscript/rewriter/onnxruntime/transformers/fastgelu.py',  # FIXME
     'onnxscript/rewriter/onnxruntime/instance_to_group_normalization.py',  # FIXME
-    'onnxscript/rewriter/onnxruntime/xformers/_smollm_*.py',  # onnxscript code
+    'onnxscript/rewriter/ort_fusions/_smollm_*.py',  # onnxscript code
     'onnxscript/_legacy_ir/irbuilder.py',  # FIXME
     'onnxscript/rewriter/onnxruntime/transformers/multihead_attention.py',  # FIXME
     'onnxscript/tools/function_unittest_producer.py',  # FIXME

onnxscript/rewriter/generic_pattern.py

@@ -549,8 +549,10 @@ def match(
         model: ir.Model,
         graph_or_function: ir.Graph | ir.Function,
         node: ir.Node,
+        *,
         verbose: int = 0,
         remove_nodes: bool = True,
+        tracer: orp.MatchingTracer | None = None,
     ) -> orp.MatchResult | None:
         if not remove_nodes:
             raise NotImplementedError(

onnxscript/rewriter/onnxruntime/xformers/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Fusion optimizations for ORT backend."""
+
+__all__ = [
+    "optimize_for_ort",
+]
+
+from onnxscript.rewriter.ort_fusions._core import optimize_for_ort

onnxscript/rewriter/onnxruntime/xformers/fuse_xformers.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import onnxscript.ir as ir
+from onnxscript.optimizer import optimize, remove_unused_nodes
+from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
+from onnxscript.rewriter.ort_fusions.mha import fuse_mha
+from onnxscript.rewriter.ort_fusions.rms_normalization import fuse_rms_normalization
+from onnxscript.rewriter.ort_fusions.rotary_embedding import fuse_rotary_embedding
+from onnxscript.rewriter.ort_fusions.sdpa import fuse_sdpa
+from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_normalization
+
+
+def fuse_xformers(model: ir.Model) -> None:
+    optimize(model)
+    fuse_rms_normalization(model)
+    fuse_normalization(model)
+    fuse_rotary_embedding(model)
+    fuse_cos_sin_cache(model)
+    fuse_sdpa(model)
+    fuse_mha(model)
+    remove_unused_nodes(model)
+
+
+def optimize_for_ort(model: ir.Model) -> None:
+    # TODO(rama): Include the other optimizations
+    fuse_xformers(model)