[DRAFT] First version of fusion optimizations for transformers

onnxscript/optimizer/__init__.py

@@ -4,13 +4,16 @@
 
 import onnx
 
+import onnxscript.optimizer._constant_folding as constant_folding
 import onnxscript.optimizer._legacy._optimizer as legacy_optimizer
+import onnxscript.optimizer._legacy.constant_folding as legacy_constant_folding
 from onnxscript import ir
-from onnxscript.optimizer._constant_folding import basic_constant_propagation
-from onnxscript.optimizer._legacy.constant_folding import fold_constants
 from onnxscript.optimizer._optimizer import optimize_ir
 from onnxscript.optimizer._remove_unused import remove_unused_nodes
 
+basic_constant_propagation = constant_folding.basic_constant_propagation
+fold_constants_ir = constant_folding.fold_constants
+
 
 def optimize(model: ir.Model | onnx.ModelProto, *args, **kwargs):
     if isinstance(model, ir.Model):
@@ -19,8 +22,16 @@ def optimize(model: ir.Model | onnx.ModelProto, *args, **kwargs):
         return legacy_optimizer.optimize(model, *args, **kwargs)
 
 
+def fold_constants(model: ir.Model | onnx.ModelProto, *args, **kwargs):
+    if isinstance(model, ir.Model):
+        return constant_folding.fold_constants(model, *args, **kwargs)
+    else:
+        return legacy_constant_folding.fold_constants(model, *args, **kwargs)
+
+
 __all__ = [
     "fold_constants",
+    "fold_constants_ir",
     "remove_unused_nodes",
     "optimize",
     "optimize_ir",

onnxscript/optimizer/_constant_folding.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import onnx
+import onnx.helper
 import onnx.reference.ops
 
 import onnxscript.ir as ir
@@ -434,25 +435,54 @@ def concat(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
 @register("Dropout", version=(12, None))
 def dropout(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     """Replace a Dropout by Identity when applicable."""
-    if len(node.outputs) != 1:
-        # If output mask is requested, optimization is more complex.
-        # TODO: handle this case. But unlikely to be needed in practice.
-        return None
+
+    def optimized_dropout():
+        input = node.inputs[0]
+        output = op.Identity(input)
+        if len(node.outputs) == 1:
+            return output
+        else:
+            true_tensor = onnx.helper.make_tensor("true", onnx.TensorProto.BOOL, [1], [True])
+            input_shape = op.Shape(input)
+            mask = op.ConstantOfShape(input_shape, value=true_tensor)
+            return output, mask
+
     inputs = node.inputs
     if (len(inputs) <= 2) or inputs[2] is None:
         # No training_mode specified:
-        return op.Identity(inputs[0])
+        return optimized_dropout()
     if _get_bool_value(inputs[2]) is False:
         # training_mode is False: dropout is not applied.
-        return op.Identity(inputs[0])
+        return optimized_dropout()
     ratio = _get_numpy_value(inputs[1])
     if ratio is None:
         return None
     if ratio.size != 1:  # Only scalar dropout ratio is supported.
         return None
     if ratio.item() == 0:
         # dropout ratio is 0: dropout is not applied.
-        return op.Identity(inputs[0])
+        return optimized_dropout()
+    return None
+
+
+@register("Expand")
+def expand(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
+    """Replace an Expand node by Identity when applicable."""
+    if len(node.inputs) != 2:
+        return None
+    if (input := node.inputs[0]) is None:
+        return None
+    if (input_shape := input.shape) is None:
+        # Input shape is not known.
+        return None
+    if (expanded_shape := _get_numpy_value(node.inputs[1])) is None:
+        # Target shape is not known.
+        return None
+    if expanded_shape.ndim != 1:
+        # Target shape must be a 1D tensor. Erroneous model.
+        return None
+    if input_shape.dims == tuple(expanded_shape.tolist()):
+        return op.Identity(input)
     return None
 
 

onnxscript/optimizer/_constant_folding_test.py

@@ -433,6 +433,22 @@ def test_concat_identity(self):
         self.assertEqual(len(optimized.graph.node), 1)
         self.assertEqual(optimized.graph.node[0].op_type, "Identity")
 
+    def test_expand_identity(self):
+        if not self.using_ir:
+            self.skipTest("New optimizations not supported for legacy optimizer")
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[128, 256] x) => (float[128, 256] z)
+            {
+                shape = Constant <value_ints=[128, 256]> ()
+                z = Expand (x, shape)
+            }
+        """
+        )
+        optimized = self._fold(model)
+        self.assertEqual(optimized.graph.node[-1].op_type, "Identity")
+
 
 if __name__ == "__main__":
     unittest.main()

onnxscript/optimizer/_inliner.py

@@ -305,6 +305,7 @@ def inline_calls_in(self, graph: ir.Graph) -> None:
 
 def inline(model: ir.Model) -> None:
     """Inline all function calls (recursively) in the model."""
-    inliner = _Inliner(model)
-    inliner.inline_calls_in(model.graph)
-    model.functions.clear()
+    if model.functions:
+        inliner = _Inliner(model)
+        inliner.inline_calls_in(model.graph)
+        model.functions.clear()

onnxscript/rewriter/_ir_utils.py

@@ -2,6 +2,8 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
+import numpy as np
+
 import onnxscript.ir as ir
 from onnxscript.optimizer import basic_constant_propagation
 
@@ -11,3 +13,24 @@ def get_const_value(value: ir.Value) -> ir.TensorProtocol | None:
     if node is not None:
         basic_constant_propagation([node])
     return value.const_value
+
+
+def get_numpy_value(val: ir.Value | None) -> np.ndarray | None:
+    if val is None:
+        return None
+    const_value = val.const_value
+    if const_value is not None:
+        try:
+            return const_value.numpy()
+        except FileNotFoundError:
+            # External data is not available.
+            return None
+    return None
+
+
+def get_singleton_value(val: ir.Value | None):
+    """Returns element of a single element tensor constant value, and None otherwise."""
+    np_val = get_numpy_value(val)
+    if np_val is not None and np_val.size == 1:
+        return np_val.item()
+    return None

onnxscript/rewriter/onnxruntime/xformers/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+from onnxscript.rewriter.onnxruntime.xformers.multi_head_attention import (
+    mha_rules as mha_rules,
+)
+from onnxscript.rewriter.onnxruntime.xformers.rms_normalization import (
+    rms_normalization_rules as rms_normalization_rules,
+)
+from onnxscript.rewriter.onnxruntime.xformers.rotary_embedding import (
+    rotary_embedding_rules as rotary_embedding_rules,
+)
+from onnxscript.rewriter.onnxruntime.xformers.sdpa import sdpa_rules as sdpa_rules
+from onnxscript.rewriter.onnxruntime.xformers.skip_normalization import (
+    skip_normalization_rules as skip_normalization_rules,
+)

onnxscript/rewriter/onnxruntime/xformers/_optimize_transformers.py

@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import onnxscript.ir as ir
+from onnxscript.optimizer import fold_constants_ir, remove_unused_nodes
+from onnxscript.rewriter.onnxruntime.xformers import (
+    mha_rules,
+    rms_normalization_rules,
+    rotary_embedding_rules,
+    sdpa_rules,
+    skip_normalization_rules,
+)
+
+
+def optimize(irmodel: ir.Model, verbose: int = 0) -> None:
+    def apply(rulename: str, rule):
+        count = rule.apply_to_model(irmodel, verbose=verbose)
+        print(f"{rulename} count: {count}")
+
+    fold_constants_ir(irmodel, input_size_limit=5120000 * 4, output_size_limit=5120000 * 4)
+    remove_unused_nodes(irmodel)
+
+    apply("RMS Normalization", rms_normalization_rules)
+    apply("Skip Normalization", skip_normalization_rules)
+
+    fold_constants_ir(irmodel)
+    remove_unused_nodes(irmodel)
+
+    apply("SDPA-Attention", sdpa_rules)
+    apply("RotaryEmbedding", rotary_embedding_rules)
+    apply("Multi-Head-Attention", mha_rules)

onnxscript/rewriter/onnxruntime/xformers/_optimize_transformers_test.py

@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import unittest
+
+import torch
+import transformers.models.llama.modeling_llama as modeling_llama
+from transformers import LlamaConfig
+
+import onnxscript.optimizer
+from onnxscript.rewriter.onnxruntime.xformers import (
+    _optimize_transformers as optimize_transformers,
+)
+
+# Create a LlamaConfig object with the desired parameters
+_config = LlamaConfig(
+    _name_or_path="HuggingFaceTB/SmolLM-1.7B",
+    architectures=["LlamaForCausalLM"],
+    attention_bias=False,
+    attention_dropout=0.0,
+    bos_token_id=0,
+    eos_token_id=0,
+    hidden_act="silu",
+    hidden_size=2048,
+    initializer_range=0.02,
+    intermediate_size=8192,
+    max_position_embeddings=2048,
+    model_type="llama",
+    num_attention_heads=32,
+    num_hidden_layers=24,
+    num_key_value_heads=32,
+    pretraining_tp=1,
+    rms_norm_eps=1e-05,
+    rope_scaling=None,
+    rope_theta=10000.0,
+    tie_word_embeddings=True,
+    torch_dtype="float32",
+    transformers_version="4.37.2",
+    use_cache=True,
+    vocab_size=49152,
+)
+
+# Create a LlamaAttention object with the desired parameters
+# model = modeling_llama.LlamaAttention(_config, 0)
+model = modeling_llama.LlamaSdpaAttention(_config, 0)
+
+# Dimensions for inputs:
+_batch_size = 1
+_seq_len = 10
+_hidden_size = _config.hidden_size
+_num_attention_heads = _config.num_attention_heads
+dim = _hidden_size // _num_attention_heads
+
+# Generate inputs:
+_hidden_states = torch.rand(_batch_size, _seq_len, _hidden_size, dtype=torch.float32)
+_attention_mask = torch.rand(_batch_size, 1, _seq_len, _seq_len, dtype=torch.float32)
+_position_ids = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64).reshape(1, 10)
+
+
+# Get model in ONNX format
+def _get_model(llama_attention_class, with_mask: bool):
+    model = llama_attention_class(_config, 0)
+    if with_mask:
+        inputs = (_hidden_states, _attention_mask, _position_ids)
+    else:
+        inputs = (_hidden_states, None, _position_ids)
+    exported = torch.onnx.export(model, inputs, dynamo=True)
+    onnxscript.optimizer.optimize(exported.model)
+    # optimize_transformers.basic_optimize(exported.model)
+    return exported.model
+
+
+class TestOptimizeTransformers(unittest.TestCase):
+    def test_attention(self):
+        model = _get_model(modeling_llama.LlamaAttention, with_mask=False)
+        optimize_transformers.optimize(model)
+        op_types = [n.op_type for n in model.graph]
+        self.assertIn("MultiHeadAttention", op_types)
+
+    def test_masked_attention(self):
+        model = _get_model(modeling_llama.LlamaAttention, with_mask=True)
+        optimize_transformers.optimize(model)
+        op_types = [n.op_type for n in model.graph]
+        self.assertIn("MultiHeadAttention", op_types)
+
+    def test_sdpa_attention(self):
+        model = _get_model(modeling_llama.LlamaSdpaAttention, with_mask=False)
+        optimize_transformers.optimize(model)
+        op_types = [n.op_type for n in model.graph]
+        self.assertIn("MultiHeadAttention", op_types)
+
+    def test_masked_sdpa_attention(self):
+        model = _get_model(modeling_llama.LlamaSdpaAttention, with_mask=True)
+        optimize_transformers.optimize(model)
+        op_types = [n.op_type for n in model.graph]
+        self.assertIn("MultiHeadAttention", op_types)
+
+
+if __name__ == "__main__":
+    unittest.main()