Skip to content

Rewrite Skip fusions with check functions #2259

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 2, 2025

Conversation

shubhambhokare1
Copy link
Contributor

@shubhambhokare1 shubhambhokare1 commented May 1, 2025

  • Rewrite SkipLayerNorm fusions and SkipRMSNorm fusions to match format of other ort-fusion patterns.
  • Added check functions for ensuring shapes are as expected.
  • Moving these fusions out of PR Add fusion rules (Whisper optimizations) #2221

Fusion support patterns with:

  • Add(input, skip) -> Norm
  • Add(input, skip) -> Add (result, bias) -> Norm
  • Add(input, bias) -> Add (result, skip) -> Norm

NOTE:
These fusions should support:

  • Planned whisper-related optimizations
  • Benchmark failures stemming from wrong bias shapes for SkipLayerNorm fusions

Copy link

codecov bot commented May 1, 2025

❌ 3 Tests Failed:

Tests completed Failed Passed Skipped
14458 3 14455 1881
View the top 3 failed test(s) by shortest run time
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0085_test_averagepool_2d_precomputed_strides
Stack Traces | 0.005s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_averagepool_2d_precomputed_strides'

The above exception was the direct cause of the following exception:
.nox\test_ort_nightly\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_averagepool_2d_precomputed_strides' (e=No module named 'tests.onnx_backend_test_code.test_averagepool_2d_precomputed_strides') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_averagepool_2d_precomputed_strides.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_averagepool_2d_precomputed_strides.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 opset22
E   
E   @script()
E   def bck_test_averagepool_2d_precomputed_strides(x: FLOAT[1,1,5,5]) -> (FLOAT[1,1,2,2]):
E       y = opset22.AveragePool(x, kernel_shape=[2, 2], strides=[2, 2])
E       return y
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0797_test_range_int32_type_negative_delta
Stack Traces | 0.005s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_range_int32_type_negative_delta'

The above exception was the direct cause of the following exception:
.nox\test_ort_nightly\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_range_int32_type_negative_delta' (e=No module named 'tests.onnx_backend_test_code.test_range_int32_type_negative_delta') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_range_int32_type_negative_delta.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_range_int32_type_negative_delta.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 INT32
E   from onnxscript.onnx_opset import opset11
E   
E   @script()
E   def bck_test_range_int32_type_negative_delta(start: INT32, limit: INT32, delta: INT32) -> (INT32[2]):
E       output = opset11.Range(start, limit, delta)
E       return output
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.006s 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:94: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:38: 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.

@github-project-automation github-project-automation bot moved this from Todo to Done in ONNX Script Review Board May 1, 2025
@shubhambhokare1 shubhambhokare1 force-pushed the sbhokare/skip-fusions-fix branch from d82adcb to b54e320 Compare May 1, 2025 21:16
@shubhambhokare1 shubhambhokare1 merged commit 34e7ba8 into main May 2, 2025
22 of 26 checks passed
@shubhambhokare1 shubhambhokare1 deleted the sbhokare/skip-fusions-fix branch May 2, 2025 17:38
shubhambhokare1 added a commit that referenced this pull request May 7, 2025
Add fusion rules to support the optimization of Whisper models.

Fusions added:

- Basic Fusions:
    * additional pattern for erfgelu [moved to #2222]
- SkipLayerNorm:
    * #2259 
    * Fusion patterns where skip_sum is also an output
    * Bias + SkipLayerNorm -> SkipLayerNorm (with bias) [moved to #2222]
- BiasGelu Fusion [moved to #2222]
- SDPA:
    * Support for pattern where only q is pre-scaled
- MHA:
    * Patterns with/without past/present keys/values
    * Patterns with non-rotary embeddings
    * Patterns with/without mask
    * Patterns with cross-attention (only for past key/value patterns)
 - MHA Bias Fusion:
* Bias was offloaded to Attention fusion previously, this fusion fixes
that
 - Attention:
    * Patterns where Q, K and V do not come from slicing

TODO:
- [x]  Fix SDPA singular prescale case, due to lost shape information
- [x] - Enable check conditions when #2210 is merged
- [x] - Improve/Rewrite whisper model test case to be similar to that of
smollm (for eg)
- [x] - Fix failing test cases to account for new patterns
- [x] - Add isolated test cases for new fusions like BiasGelu,
SkipLayerNorm etc
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
Development

Successfully merging this pull request may close these issues.

3 participants