Add RMS Normalization rule variant (#2638)

Add RMS Normalization rule variant to support different order of
multiplying by scale.

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/rewriter/ort_fusions/rms_normalization.py

@@ -31,6 +31,10 @@
 
 
 class RmsNormFusion(pattern.RewriteRuleClassBase):
+    def __init__(self, name: str, _mul_order: bool):
+        super().__init__(name)
+        self._mul_order = _mul_order
+
     def pattern(self, op, x, scale, epsilon, compute_dtype, target_dtype):
         x = pattern.OrValue([op.Cast(x, to=compute_dtype), x])
         x_square = op.Pow(x, 2.0)
@@ -42,7 +46,11 @@ def pattern(self, op, x, scale, epsilon, compute_dtype, target_dtype):
         normalized = pattern.OrValue([op.Cast(normalized, to=target_dtype), normalized])
         # To support float16, we need to ensure the scale is casted or not.
         scale = pattern.OrValue([op.Cast(scale, to=compute_dtype), scale])
-        return op.Mul(scale, normalized)
+        # Workaround: can't use OrValue for final (returned) value
+        if self._mul_order:
+            return op.Mul(normalized, scale)
+        else:
+            return op.Mul(scale, normalized)
 
     def check(
         self, op, x, scale, epsilon, compute_dtype, target_dtype, **_
@@ -77,8 +85,10 @@ def rewrite(self, op, x, scale, epsilon, **_):
         )
 
 
-_rule = RmsNormFusion.rule()
-rms_normalization_rules = [_rule]
+_rule1 = RmsNormFusion.rule("RmsNormFusion1", _mul_order=False)
+_rule2 = RmsNormFusion.rule("RmsNormFusion2", _mul_order=True)
+
+rms_normalization_rules = [_rule1, _rule2]
 rms_normalization_ruleset = pattern.RewriteRuleSet(rms_normalization_rules)