RMS Normalization and Skip RMS Normalization fusion optimizations

[gramalingam]

Implements RMS Normalization and Skip RMS Normalization fusion optimizations (for use of onnxruntime custom fused ops for these).

[codecov]

❌ 17 Tests Failed:

Tests completed	Failed	Passed	Skipped
16147	17	16130	3318

View the top 2 failed tests by shortest run time

::onnxscript.rewriter.onnxruntime.xformers._test_models

Stack Traces | 0s run time

No failure message available

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyEagerCPU::test_output_match_opinfo__clamp_cpu_float16

Stack Traces | 0.001s run time

No failure message available

View the full list of 1 ❄️ flaky tests

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime::test_function_input_and_attribute_by_kwargs_out_of_order

Flake rate in main: 39.38% (Passed 12384 times, Failed 8045 times)

Stack Traces | 0.003s run time

.nox\test_torch_nightly\Lib\site-packages\onnx\reference\ops\_op.py:91: in run
    res = self._run(x, y)
.nox\test_torch_nightly\Lib\site-packages\onnx\reference\ops\_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
.nox\test_torch_nightly\Lib\site-packages\onnx\reference\custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests\eager_mode_test.py:115: in test_function_input_and_attribute_by_kwargs_out_of_order
    self.assertEqual(add_with_alpha(alpha=3.0, other=2.0, this=1.0), 7.0)
onnxscript\values.py:576: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript\evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests\eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
onnxscript\onnx_opset\_impl\opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript\values.py:304: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript\evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript\evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
.nox\test_torch_nightly\Lib\site-packages\onnx\reference\reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
.nox\test_torch_nightly\Lib\site-packages\onnx\reference\ops\_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
.nox\test_torch_nightly\Lib\site-packages\onnx\reference\ops\_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types <class 'numpy.ndarray'>, <class 'numpy.ndarray'> (binary operator 'Mul').

To view more test analytics, go to the Test Analytics Dashboard
📢 Thoughts on this report? Let us know!

[xadupre]

It looks good. At some point, we have a page in the documentation which lists all the possible optimization and how to trigger them. Even for us, this page should prove useful.

[github-advanced-security]

lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.

[justinchuby]

It looks good. At some point, we have a page in the documentation which lists all the possible optimization and how to trigger them. Even for us, this page should prove useful.

Sounds good. I think we should autogen this documentation from source.