Skip to content

Optimize away zero-length concat operands #2150

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
Mar 31, 2025
Merged

Conversation

gramalingam
Copy link
Collaborator

@gramalingam gramalingam commented Mar 31, 2025

We optimize Concat (x1, x2, x3) if one or more the concat operands has zero length along the concatenated axis-dimension.

This pattern shows up, for example, in Phi models. See this line in the implementation of partial-rotary-embedding:

q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1)

In the special case of total-rotary-embedding, the second operand q_pass of the concat is empty. This also interferes with the pattern-matching for GQA in the generated graph. Optimizing the redundant Concat away will help with GQA fusion as well.

Handle the edge case when all operands have zero size.

Copy link

codecov bot commented Mar 31, 2025

❌ 12 Tests Failed:

Tests completed Failed Passed Skipped
14813 12 14801 2225
View the top 2 failed test(s) by shortest run time
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0530_test_layer_normalization_3d_axis_negative_3_epsilon
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.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_layer_normalization_3d_axis_negative_3_epsilon'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\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_layer_normalization_3d_axis_negative_3_epsilon' (e=No module named 'tests.onnx_backend_test_code.test_layer_normalization_3d_axis_negative_3_epsilon') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_3d_axis_negative_3_epsilon.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_3d_axis_negative_3_epsilon.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 opset17
E   
E   @script()
E   def bck_test_layer_normalization_3d_axis_negative_3_epsilon(X: FLOAT[2,3,5], W: FLOAT[2,3,5], B: FLOAT[2,3,5]) -> (FLOAT[2,3,5], FLOAT[1,1,1], FLOAT[1,1,1]):
E       Y, Mean, InvStdDev = opset17.LayerNormalization(X, W, B, axis=-3, epsilon=0.10000000149011612)
E       return Y, Mean, InvStdDev
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0747_test_onehot_with_axis
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.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_onehot_with_axis'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\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_onehot_with_axis' (e=No module named 'tests.onnx_backend_test_code.test_onehot_with_axis') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_onehot_with_axis.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_onehot_with_axis.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 opset11
E   
E   @script()
E   def bck_test_onehot_with_axis(indices: FLOAT[2,2], depth: FLOAT, values: FLOAT[2]) -> (FLOAT[2,10,2]):
E       y = opset11.OneHot(indices, depth, values, axis=1)
E       return y
View the full list of 1 ❄️ flaky tests
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0818_test_reduce_l1_empty_set

Flake rate in main: 7.14% (Passed 39 times, Failed 3 times)

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.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_reduce_l1_empty_set'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\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_reduce_l1_empty_set' (e=No module named 'tests.onnx_backend_test_code.test_reduce_l1_empty_set') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l1_empty_set.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l1_empty_set.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, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_reduce_l1_empty_set(data: FLOAT[2,0,4], axes: INT64[1]) -> (FLOAT[2,1,4]):
E       reduced = opset18.ReduceL1(data, axes, keepdims=1)
E       return reduced

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

@gramalingam gramalingam enabled auto-merge (squash) March 31, 2025 18:03
@justinchuby
Copy link
Collaborator

Some tests are failing

@gramalingam gramalingam merged commit 5d969c4 into main Mar 31, 2025
23 of 27 checks passed
@gramalingam gramalingam deleted the rama/concat-zero branch March 31, 2025 18:44
bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025
We optimize `Concat (x1, x2, x3)` if one or more the concat operands has
zero length along the concatenated axis-dimension.

This pattern shows up, for example, in Phi models. See [this
line](https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/models/phi3/modeling_phi3.py#L152)
in the implementation of partial-rotary-embedding:
```py
q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1)
```
In the special case of total-rotary-embedding, the second operand
`q_pass` of the concat is empty. This also interferes with the
pattern-matching for GQA in the generated graph. Optimizing the
redundant Concat away will help with GQA fusion as well.

Handle the edge case when all operands have zero size.
bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025
We optimize `Concat (x1, x2, x3)` if one or more the concat operands has
zero length along the concatenated axis-dimension.

This pattern shows up, for example, in Phi models. See [this
line](https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/models/phi3/modeling_phi3.py#L152)
in the implementation of partial-rotary-embedding:
```py
q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1)
```
In the special case of total-rotary-embedding, the second operand
`q_pass` of the concat is empty. This also interferes with the
pattern-matching for GQA in the generated graph. Optimizing the
redundant Concat away will help with GQA fusion as well.

Handle the edge case when all operands have zero size.
bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025
We optimize `Concat (x1, x2, x3)` if one or more the concat operands has
zero length along the concatenated axis-dimension.

This pattern shows up, for example, in Phi models. See [this
line](https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/models/phi3/modeling_phi3.py#L152)
in the implementation of partial-rotary-embedding:
```py
q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1)
```
In the special case of total-rotary-embedding, the second operand
`q_pass` of the concat is empty. This also interferes with the
pattern-matching for GQA in the generated graph. Optimizing the
redundant Concat away will help with GQA fusion as well.

Handle the edge case when all operands have zero size.
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.

2 participants