Add unit tests for rotary embedding, skip norm, rms norm, layer norm

22 new tests across 4 files:

rotary_embedding_unit_test.py (3 tests):
- Full rotary embedding pattern fusion
- Partial rotary embedding (adds rotary_embedding_dim attribute)
- 3D input rejection (negative)

skip_normalization_unit_test.py (8 tests):
- SkipRmsNorm: both Add orderings via OrValue (parameterized)
- SkipRmsNorm: post-add bias and pre-add bias variants
- SkipLayerNorm: no bias and post-add bias
- No skip Add (negative), rank-2 input (negative)

_rms_normalization_extended_test.py (5 tests):
- Both mul_order variants: scale*norm and norm*scale (parameterized)
- Mixed-precision: fp16 input with fp32 compute via Cast
- Double precision
- Integer input rejection (negative)

_layer_norm_extended_test.py (6 tests):
- OrValue: Pow(deviation,2) vs Mul(deviation,deviation)
- OrValue: Div(deviation,std_dev) vs Mul(deviation,Reciprocal)
- Both OrValue alternatives combined
- Div + bias fusion
- Double precision
- fp16 input rejection (negative)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: gramalingam

onnxscript/rewriter/ort_fusions/rotary_embedding_unit_test.py

@@ -0,0 +1,167 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Unit tests for RotaryEmbeddingFusion and PartialRotaryEmbeddingFusion rules.
+
+RotaryEmbeddingFusion matches: x * cos + rotate_half(x) * sin
+and rewrites to RotaryEmbedding(x, cos, sin, ...).
+
+PartialRotaryEmbeddingFusion matches: Concat(RotaryEmbedding(Slice(x)), Slice(x))
+and adds rotary_embedding_dim attribute to the RotaryEmbedding op.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import onnx_ir as ir
+
+from onnxscript import FLOAT, script, values
+from onnxscript import opset18 as op
+from onnxscript.optimizer import optimize
+from onnxscript.rewriter.ort_fusions.rotary_embedding import (
+    fuse_partial_rotary_embedding,
+    fuse_rotary_embedding,
+)
+
+fusion_op = values.Opset("ai.onnxruntime._fusion", 1)
+msft_op = values.Opset("com.microsoft", 1)
+
+_B, _H, _S, _Dh = 2, 4, 8, 8
+_HALF = _Dh // 2
+
+# Constants for slice boundaries
+_ZERO = ir.tensor(np.array([0], dtype=np.int64))
+_HALF_TENSOR = ir.tensor(np.array([_HALF], dtype=np.int64))
+_HEAD_SIZE_TENSOR = ir.tensor(np.array([_Dh], dtype=np.int64))
+_MAX_INT = ir.tensor(np.array([9223372036854775807], dtype=np.int64))
+
+
+# --- Full rotary embedding pattern ---
+
+
+@script()
+def _rotary_full(x, cos, sin):
+    """X * cos + rotate_half(x) * sin — standard non-interleaved pattern."""
+    start_0 = op.Constant(value=_ZERO)
+    end_half = op.Constant(value=_HALF_TENSOR)
+    start_half = op.Constant(value=_HALF_TENSOR)
+    end_full = op.Constant(value=_HEAD_SIZE_TENSOR)
+    # rotate_half: concat(-x2, x1)
+    x1 = op.Slice(x, start_0, end_half, [3], [1])
+    x2 = op.Slice(x, start_half, end_full, [3], [1])
+    neg_x2 = op.Neg(x2)
+    rotated = op.Concat(neg_x2, x1, axis=-1)
+    return op.Add(op.Mul(x, cos), op.Mul(rotated, sin))
+
+
+# --- Partial rotary embedding pattern ---
+
+
+@script()
+def _partial_rotary(x, cos, sin, position_ids):
+    """Slice → RotaryEmbedding on first half, concat with second half."""
+    end1 = op.Constant(value=_HALF_TENSOR)
+    start2 = op.Constant(value=_HALF_TENSOR)
+    max_end = op.Constant(value=_MAX_INT)
+    x_part1 = op.Slice(x, [0], end1, [3], [1])
+    x_part2 = op.Slice(x, start2, max_end, [3], [1])
+    x_part1_rope = msft_op.RotaryEmbedding(x_part1, position_ids, cos, sin, interleaved=0)
+    return op.Concat(x_part1_rope, x_part2, axis=-1)
+
+
+# --- Negative: 3D input instead of 4D ---
+
+
+@script()
+def _rotary_3d_input(x, cos, sin):
+    """3D input — should fail the 4D check."""
+    start_0 = op.Constant(value=_ZERO)
+    end_half = op.Constant(value=_HALF_TENSOR)
+    start_half = op.Constant(value=_HALF_TENSOR)
+    end_full = op.Constant(value=_HEAD_SIZE_TENSOR)
+    x1 = op.Slice(x, start_0, end_half, [3], [1])
+    x2 = op.Slice(x, start_half, end_full, [3], [1])
+    neg_x2 = op.Neg(x2)
+    rotated = op.Concat(neg_x2, x1, axis=-1)
+    return op.Add(op.Mul(x, cos), op.Mul(rotated, sin))
+
+
+class RotaryEmbeddingFusionTest(unittest.TestCase):
+    """Unit tests for RotaryEmbeddingFusion rule."""
+
+    def _build(self, script_fn, input_types, output_types) -> ir.Model:
+        model_proto = script_fn.to_model_proto(
+            input_types=input_types, output_types=output_types
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        optimize(model)
+        return model
+
+    def _count_op(self, model: ir.Model, op_type: str, domain: str = "") -> int:
+        return sum(1 for n in model.graph if n.op_type == op_type and n.domain == domain)
+
+    # --- Positive tests ---
+
+    def test_full_rotary_fuses(self):
+        """Standard full rotary embedding pattern → fuses to RotaryEmbedding op."""
+        model = self._build(
+            _rotary_full,
+            input_types=[
+                FLOAT[_B, _H, "S", _Dh],
+                FLOAT[_B, _H, "S", _Dh],
+                FLOAT[_B, _H, "S", _Dh],
+            ],
+            output_types=[FLOAT[_B, _H, "S", _Dh]],
+        )
+        count = fuse_rotary_embedding(model)
+        self.assertEqual(count, 1)
+        self.assertEqual(self._count_op(model, "RotaryEmbedding", "ai.onnxruntime._fusion"), 1)
+        # Pattern ops should be consumed
+        self.assertEqual(self._count_op(model, "Neg"), 0)
+
+    def test_partial_rotary_fuses(self):
+        """Partial rotary: Slice → RotaryEmbedding on first part, concat with rest."""
+        model = self._build(
+            _partial_rotary,
+            input_types=[
+                FLOAT[_B, _H, "S", _Dh],
+                FLOAT["S", _HALF],
+                FLOAT["S", _HALF],
+                FLOAT[_B, "S"],
+            ],
+            output_types=[FLOAT[_B, _H, "S", _Dh]],
+        )
+        count = fuse_partial_rotary_embedding(model)
+        self.assertEqual(count, 1)
+        # RotaryEmbedding should still exist but now with rotary_embedding_dim
+        rope_nodes = [
+            n
+            for n in model.graph
+            if n.op_type == "RotaryEmbedding" and n.domain == "com.microsoft"
+        ]
+        self.assertEqual(len(rope_nodes), 1)
+        self.assertIn("rotary_embedding_dim", rope_nodes[0].attributes)
+        self.assertEqual(rope_nodes[0].attributes["rotary_embedding_dim"].value, _HALF)
+
+    # --- Negative tests ---
+
+    def test_3d_input_no_fusion(self):
+        """3D input (missing batch or head dim) → check rejects."""
+        model_proto = _rotary_3d_input.to_model_proto(
+            input_types=[
+                FLOAT[_H, "S", _Dh],
+                FLOAT[_H, "S", _Dh],
+                FLOAT[_H, "S", _Dh],
+            ],
+            output_types=[FLOAT[_H, "S", _Dh]],
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        # Skip optimize — 3D shapes cause shape inference errors which is expected
+        count = fuse_rotary_embedding(model)
+        self.assertEqual(count, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

onnxscript/rewriter/ort_fusions/skip_normalization_unit_test.py

@@ -0,0 +1,217 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Unit tests for SkipRmsNormFusion and SkipLayerNormFusion rules.
+
+SkipRmsNormFusion: Add(input, skip) → SimplifiedLayerNormalization →
+    SkipSimplifiedLayerNormalization (com.microsoft).
+
+SkipLayerNormFusion: Add(input, skip) → LayerNormalization →
+    SkipLayerNormalization (com.microsoft).
+
+Covers: no bias, post-add bias, pre-add bias variants; negative tests.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import onnx_ir as ir
+from parameterized import parameterized
+
+from onnxscript import FLOAT, script
+from onnxscript import opset18 as op
+from onnxscript.optimizer import optimize
+from onnxscript.rewriter.ort_fusions.skip_normalization import (
+    fuse_skip_layer_normalization,
+    fuse_skip_rms_normalization,
+)
+
+_B, _S, _D = 2, 8, 16
+_EPS_F = ir.tensor(np.array([1e-6], dtype=np.float32))
+
+
+# ========== Skip RMS Norm patterns ==========
+
+
+@script()
+def _skip_rms_no_bias(input, skip, gamma):
+    skip_sum = op.Add(input, skip)
+    return op.SimplifiedLayerNormalization(
+        skip_sum, gamma, axis=-1, epsilon=1e-6, stash_type=1
+    )
+
+
+@script()
+def _skip_rms_no_bias_reversed(input, skip, gamma):
+    """Skip + input order reversed (OrValue alternative)."""
+    skip_sum = op.Add(skip, input)
+    return op.SimplifiedLayerNormalization(
+        skip_sum, gamma, axis=-1, epsilon=1e-6, stash_type=1
+    )
+
+
+@script()
+def _skip_rms_post_bias(input, skip, gamma, bias):
+    skip_sum = op.Add(input, skip)
+    skip_sum_biased = op.Add(skip_sum, bias)
+    return op.SimplifiedLayerNormalization(
+        skip_sum_biased, gamma, axis=-1, epsilon=1e-6, stash_type=1
+    )
+
+
+@script()
+def _skip_rms_pre_bias(input, skip, gamma, bias):
+    input_biased = op.Add(input, bias)
+    skip_sum = op.Add(input_biased, skip)
+    return op.SimplifiedLayerNormalization(
+        skip_sum, gamma, axis=-1, epsilon=1e-6, stash_type=1
+    )
+
+
+# ========== Skip Layer Norm patterns ==========
+
+
+@script()
+def _skip_ln_no_bias(input, skip, gamma, beta):
+    skip_sum = op.Add(input, skip)
+    return op.LayerNormalization(skip_sum, gamma, beta, axis=-1, epsilon=1e-6, stash_type=1)
+
+
+@script()
+def _skip_ln_post_bias(input, skip, gamma, beta, bias):
+    skip_sum = op.Add(input, skip)
+    skip_sum_biased = op.Add(skip_sum, bias)
+    return op.LayerNormalization(
+        skip_sum_biased, gamma, beta, axis=-1, epsilon=1e-6, stash_type=1
+    )
+
+
+# ========== Negative patterns ==========
+
+
+@script()
+def _skip_rms_no_add(input, gamma):
+    """No skip addition at all — just SimplifiedLayerNormalization."""
+    return op.SimplifiedLayerNormalization(input, gamma, axis=-1, epsilon=1e-6, stash_type=1)
+
+
+class SkipNormalizationTest(unittest.TestCase):
+    """Unit tests for SkipRmsNormFusion and SkipLayerNormFusion."""
+
+    def _build(self, script_fn, input_types, output_types) -> ir.Model:
+        model_proto = script_fn.to_model_proto(
+            input_types=input_types, output_types=output_types
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        optimize(model)
+        return model
+
+    def _count_op(self, model: ir.Model, op_type: str, domain: str = "") -> int:
+        return sum(1 for n in model.graph if n.op_type == op_type and n.domain == domain)
+
+    _3D = FLOAT["B", "S", _D]
+    _1D = FLOAT[_D]
+
+    # ---- Skip RMS Norm positive tests ----
+
+    @parameterized.expand(
+        [
+            ("input_plus_skip", _skip_rms_no_bias),
+            ("skip_plus_input", _skip_rms_no_bias_reversed),
+        ]
+    )
+    def test_skip_rms_no_bias(self, _name, script_fn):
+        """Skip + Input (both orderings) → SkipSimplifiedLayerNormalization."""
+        model = self._build(
+            script_fn,
+            input_types=[self._3D, self._3D, self._1D],
+            output_types=[self._3D],
+        )
+        count = fuse_skip_rms_normalization(model)
+        self.assertGreater(count, 0)
+        self.assertEqual(
+            self._count_op(model, "SkipSimplifiedLayerNormalization", "com.microsoft"), 1
+        )
+        self.assertEqual(self._count_op(model, "SimplifiedLayerNormalization"), 0)
+
+    def test_skip_rms_post_bias(self):
+        """(Input + Skip) + Bias → fuses with bias."""
+        model = self._build(
+            _skip_rms_post_bias,
+            input_types=[self._3D, self._3D, self._1D, self._1D],
+            output_types=[self._3D],
+        )
+        count = fuse_skip_rms_normalization(model)
+        self.assertGreater(count, 0)
+        self.assertEqual(
+            self._count_op(model, "SkipSimplifiedLayerNormalization", "com.microsoft"), 1
+        )
+
+    def test_skip_rms_pre_bias(self):
+        """(Input + Bias) + Skip → fuses with pre-add bias."""
+        model = self._build(
+            _skip_rms_pre_bias,
+            input_types=[self._3D, self._3D, self._1D, self._1D],
+            output_types=[self._3D],
+        )
+        count = fuse_skip_rms_normalization(model)
+        self.assertGreater(count, 0)
+        self.assertEqual(
+            self._count_op(model, "SkipSimplifiedLayerNormalization", "com.microsoft"), 1
+        )
+
+    # ---- Skip Layer Norm positive tests ----
+
+    def test_skip_ln_no_bias(self):
+        """Skip + Input → SkipLayerNormalization."""
+        model = self._build(
+            _skip_ln_no_bias,
+            input_types=[self._3D, self._3D, self._1D, self._1D],
+            output_types=[self._3D],
+        )
+        count = fuse_skip_layer_normalization(model)
+        self.assertGreater(count, 0)
+        self.assertEqual(self._count_op(model, "SkipLayerNormalization", "com.microsoft"), 1)
+        self.assertEqual(self._count_op(model, "LayerNormalization"), 0)
+
+    def test_skip_ln_post_bias(self):
+        """(Input + Skip) + Bias → fuses with bias."""
+        model = self._build(
+            _skip_ln_post_bias,
+            input_types=[self._3D, self._3D, self._1D, self._1D, self._1D],
+            output_types=[self._3D],
+        )
+        count = fuse_skip_layer_normalization(model)
+        self.assertGreater(count, 0)
+        self.assertEqual(self._count_op(model, "SkipLayerNormalization", "com.microsoft"), 1)
+
+    # ---- Negative tests ----
+
+    def test_no_skip_add_no_fusion(self):
+        """No Add before norm → rule should not match."""
+        model = self._build(
+            _skip_rms_no_add,
+            input_types=[self._3D, self._1D],
+            output_types=[self._3D],
+        )
+        count = fuse_skip_rms_normalization(model)
+        self.assertEqual(count, 0)
+        self.assertEqual(
+            self._count_op(model, "SkipSimplifiedLayerNormalization", "com.microsoft"), 0
+        )
+
+    def test_rank2_input_no_fusion(self):
+        """Rank-2 input [S, D] → shape check rejects (expects 3D)."""
+        model = self._build(
+            _skip_rms_no_bias,
+            input_types=[FLOAT["S", _D], FLOAT["S", _D], self._1D],
+            output_types=[FLOAT["S", _D]],
+        )
+        count = fuse_skip_rms_normalization(model)
+        self.assertEqual(count, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()