Skip to content

Conversation

justinchuby
Copy link
Member

@justinchuby justinchuby commented Oct 3, 2025

TODO: tests

This PR introduces the tofile method on tensors (similarly named as the one on numpy arrays), which allows for faster write and lower memory usage on external data by bypassing tobytes().

Compatibility with existing TensorProtocols is maintained in the external data module by using tofile only when it is available in the class. The TorchTensor class in PyTorch exporter should be updated accordingly to leverage the new logic when saving.

Note that io time to disk is reduced by 40% below.

Reference: https://github.com/microsoft/onnxscript/pull/2241/files/b2381658492510a9bcc8c0a8574db7368e33bceb

Before:

________________________________________________________
Executed in   48.08 secs    fish           external
   usr time   60.54 secs    0.00 millis   60.54 secs
   sys time   23.06 secs    1.22 millis   23.06 secs
image image

After:

________________________________________________________
Executed in   45.69 secs    fish           external
   usr time   60.68 secs  244.00 micros   60.68 secs
   sys time   22.22 secs  518.00 micros   22.22 secs
image image

Fix #207

Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Copy link

codecov bot commented Oct 3, 2025

Codecov Report

❌ Patch coverage is 55.10204% with 22 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.57%. Comparing base (feb51e5) to head (40cb60d).
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
src/onnx_ir/_core.py 48.57% 14 Missing and 4 partials ⚠️
src/onnx_ir/external_data.py 66.66% 1 Missing and 1 partial ⚠️
src/onnx_ir/tensor_adapters.py 75.00% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #210      +/-   ##
==========================================
- Coverage   76.83%   76.57%   -0.27%     
==========================================
  Files          40       40              
  Lines        4922     4965      +43     
  Branches      980      989       +9     
==========================================
+ Hits         3782     3802      +20     
- Misses        856      873      +17     
- Partials      284      290       +6     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
@justinchuby justinchuby marked this pull request as ready for review October 4, 2025 00:44
@justinchuby justinchuby requested review from titaiwangms and a team as code owners October 4, 2025 00:44
@justinchuby justinchuby added this to the 0.1.11 milestone Oct 4, 2025
@justinchuby
Copy link
Member Author

cc @iksnagreb

Copy link

sonarqubecloud bot commented Oct 4, 2025

@justinchuby justinchuby changed the title Implement tofile on tensors Implement tofile on tensors to reduce data write time by 40% Oct 6, 2025
# Implement tobytes to support native PyTorch types so we can use types like bloat16
# Reading from memory directly is also more efficient because
# it avoids copying to a NumPy array
_, address = self._get_data_chunk()
Copy link
Member Author

Choose a reason for hiding this comment

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

Rename variable to be more accurate. This should be a list of some py objects.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Create a write() method on Tensors that will avoid potential tobytes() call when serializing
1 participant