Add BiasGelu, Erfgelu and SkipLayerNormalization fusions

[shubhambhokare1]

This pull request introduces new fusion patterns and enhancements to the ONNXScript rewriter module, focusing on optimization and test coverage improvements. The key changes include adding support for BiasGelu and additional ErfGelu patterns, extending SkipLayerNormalization to handle bias addition, and updating test utilities for better accuracy validation.

New fusion patterns:

• BiasGelu Fusion: Added a new fusion pattern for BiasGelu operations, including its implementation in onnxscript/rewriter/ort_fusions/bias_gelu.py and integration into the fuse_xformers pipeline. A corresponding unit test was added to validate the functionality. [1] [2] [3] [4]

• ErfGelu Enhancements: Introduced a second pattern for ErfGelu fusion and refactored the corresponding implementation. The file was renamed from erfgelu.py to ort_fusions/erfgelu.py for consistency. [1] [2] [3] [4]

Enhancements to existing fusions:

• SkipLayerNormalization with Bias: Extended the SkipLayerNormalization fusion to support an additional bias term. This includes new patterns and rewrite rules in onnxscript/rewriter/ort_fusions/skip_normalization.py.

Test utility updates:

• Tolerance Adjustment: Increased the relative and absolute tolerances in assert_allclose to 1e-3 for better handling of numerical discrepancies in tests.

[codecov]

❌ 4 Tests Failed:

Tests completed	Failed	Passed	Skipped
14440	4	14436	1879

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0335_test_div_example

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_div_example'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_div_example' (e=No module named 'tests.onnx_backend_test_code.test_div_example') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_div_example.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_div_example.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset14
E   
E   @script()
E   def bck_test_div_example(x: FLOAT[2], y: FLOAT[2]) -> (FLOAT[2]):
E       z = opset14.Div(x, y)
E       return z

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0418_test_greater_bcast

Stack Traces | 0.006s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_greater_bcast'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_greater_bcast' (e=No module named 'tests.onnx_backend_test_code.test_greater_bcast') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_greater_bcast.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_greater_bcast.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import BOOL, FLOAT
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_greater_bcast(x: FLOAT[3,4,5], y: FLOAT[5]) -> (BOOL[3,4,5]):
E       greater = opset13.Greater(x, y)
E       return greater

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0125_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.007s run time

onnxscript/converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript/backend/onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript/backend/onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:95: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:39: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:825: in _translate_call_expr
    attrs = [
onnxscript/converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

Pull Request Overview

This PR adds new fusion rules for BiasGelu, Erfgelu, and updates SkipLayerNormalization fusion to support an additional output. Key changes include:

• Updating SkipLayerNormalization to return a new output (skip_sum) and adding a corresponding fusion rule for Add+SkipLayerNormalization.
• Implementing a new BiasGelu fusion along with its test.
• Updating Erfgelu rewriting with two pattern functions for better matching capability.

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxscript/rewriter/ort_fusions/skip_normalization.py	Modified SkipLayerNormalization to return an extra output and added support for bias fusion.
onnxscript/rewriter/ort_fusions/bias_gelu_test.py	Added a unit test to verify the BiasGelu fusion.
onnxscript/rewriter/ort_fusions/bias_gelu.py	Implemented the BiasGelu fusion rule.
onnxscript/rewriter/ort_fusions/_core.py	Updated the fusion count to include BiasGelu fusion.
onnxscript/rewriter/erfgelu.py	Introduced two Erfgelu pattern rules to enhance rewrite flexibility.
onnxscript/rewriter/init.py	Updated to include the new Erfgelu fusion rule.

Comments suppressed due to low confidence (1)

onnxscript/rewriter/erfgelu.py:9

• [nitpick] The current names 'erf_gelu_pattern_1' and 'erf_gelu_pattern_2' are not very descriptive. Consider renaming them to indicate the specific matching or transformation behavior they implement.

def erf_gelu_pattern_1(op, x):

[gramalingam]

Just a minor suggestion to use onnxscript for test-case model-proto