Skip to content

Use ir methods to replace onnx helper #2091

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 22 commits into from
Apr 30, 2025
Merged

Conversation

justinchuby
Copy link
Collaborator

@justinchuby justinchuby commented Mar 7, 2025

Ban onnx.helper and onnx.numpy_helper because they can be slow. Selectively enable usages of some with noqa: TID251 and updated usages of the rest.

Fix ir.tensor to generate float32 tensors when a plain python float is provided.

Copy link

codecov bot commented Mar 7, 2025

❌ 11 Tests Failed:

Tests completed Failed Passed Skipped
14639 11 14628 1699
View the top 3 failed test(s) by shortest run time
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0032_test_and_bcast3v2d
Stack Traces | 0.003s 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_and_bcast3v2d'

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_and_bcast3v2d' (e=No module named 'tests.onnx_backend_test_code.test_and_bcast3v2d') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_and_bcast3v2d.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_and_bcast3v2d.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
E   from onnxscript.onnx_opset import opset7
E   
E   @script()
E   def bck_test_and_bcast3v2d(x: BOOL[3,4,5], y: BOOL[4,5]) -> (BOOL[3,4,5]):
E       r_and = opset7.And(x, y)
E       return r_and
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0243_test_clip_default_min
Stack Traces | 0.003s 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_clip_default_min'

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_clip_default_min' (e=No module named 'tests.onnx_backend_test_code.test_clip_default_min') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_clip_default_min.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_clip_default_min.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 opset13
E   
E   @script()
E   def bck_test_clip_default_min(x: FLOAT[3,4,5], min: FLOAT) -> (FLOAT[3,4,5]):
E       y = opset13.Clip(x, min)
E       return y
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0759_test_or_bcast4v4d
Stack Traces | 0.003s 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_or_bcast4v4d'

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_or_bcast4v4d' (e=No module named 'tests.onnx_backend_test_code.test_or_bcast4v4d') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_or_bcast4v4d.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_or_bcast4v4d.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
E   from onnxscript.onnx_opset import opset7
E   
E   @script()
E   def bck_test_or_bcast4v4d(x: BOOL[1,4,1,6], y: BOOL[3,1,5,6]) -> (BOOL[3,4,5,6]):
E       r_or = opset7.Or(x, y)
E       return r_or

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

@justinchuby justinchuby marked this pull request as draft March 7, 2025 20:27
@justinchuby justinchuby marked this pull request as ready for review March 24, 2025 23:20
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR aims to deprecate the use of onnx.helper and onnx.numpy_helper in favor of ir-based methods to improve performance. Key changes include replacing multiple onnx.helper calls with either their ir or ir.serde counterparts (and adding “# noqa: TID251” comments) and updating imports and utility functions accordingly.

Reviewed Changes

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

Show a summary per file
File Description
onnxscript/tensor.py Updated import and added noqa comment on onnx.helper usage.
onnxscript/_internal/utils.py Applied noqa comments to onnx.helper calls in utility functions.
onnxscript/optimizer/_constant_folding.py Replaced onnx.numpy_helper.from_array with ir.serde.serialize_tensor.
onnxscript/main.py Updated model export logic to use ir.Model and ir.save.
onnxscript/optimizer/_legacy/constant_folding.py, etc. Added ruff: noqa: TID251 markers and adjusted helper calls across files.
onnxscript/_internal/autocast.py Replaced tensor conversion via numpy_helper with ir.serde.serialize_tensor.
onnxscript/rewriter/llama_rule_sets.py Modified allowed tensor types check to use as_int() instead of value.
onnxscript/rewriter/cast_constant_of_shape.py and others Changed onnx.helper.make_tensor calls to use ir.tensor.
onnxscript/function_libs/torch_lib/ops/nn.py Updated empty-tensor creation to use ir.tensor.
Comments suppressed due to low confidence (2)

onnxscript/rewriter/llama_rule_sets.py:65

  • Ensure that converting the attribute value with as_int() returns the expected integer representation that matches the entries in _allowed_tensor_types. Verify that edge cases or unexpected types do not cause the check to fail.
to.as_int() in cls._allowed_tensor_types

onnxscript/_internal/autocast.py:39

  • Confirm that the new serialization using ir.serde.serialize_tensor produces outputs consistent with the previous behavior of numpy_helper.from_array, ensuring compatibility with downstream tensor utilities.
return ir.serde.serialize_tensor(ir.tensor(pyvalue, name=tensor_name))

titaiwangms
titaiwangms previously approved these changes Mar 24, 2025
@justinchuby justinchuby disabled auto-merge March 25, 2025 00:09
@justinchuby justinchuby enabled auto-merge (squash) March 25, 2025 00:11
@justinchuby justinchuby disabled auto-merge March 25, 2025 00:17
@justinchuby justinchuby enabled auto-merge (squash) March 25, 2025 00:28
@justinchuby justinchuby disabled auto-merge March 25, 2025 00:31
@justinchuby justinchuby enabled auto-merge (squash) March 25, 2025 00:36
@justinchuby justinchuby disabled auto-merge March 25, 2025 01:06
@justinchuby justinchuby marked this pull request as ready for review April 29, 2025 17:23
@justinchuby justinchuby enabled auto-merge (squash) April 29, 2025 23:48
@justinchuby justinchuby merged commit 9910215 into main Apr 30, 2025
24 of 27 checks passed
@justinchuby justinchuby deleted the justinchu/remove-onnx-helper branch April 30, 2025 06:23
elif isinstance(value, float):
# If the value is a single float, we use np.float32 as the default dtype
numpy_dtype = np.dtype(np.float32)
elif isinstance(value, Sequence) and all(
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Check that sequence is not empty and ask for a dtype if it is

justinchuby added a commit that referenced this pull request May 1, 2025
In #2091, the call to helper
was replaced with onnx ir attribute convertion. This was not properly
handling when the attribute is a subgraph and when it uses values from
the parent scope, which the ir doesn't have access to. The IR thus
raises warnings of it not being able to find those values.

This PR reverts the change.
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.

5 participants