Skip to content

Serialize IRs alongside flatbuffers #947

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from

Conversation

sgligorijevicTT
Copy link
Contributor

@sgligorijevicTT sgligorijevicTT commented Aug 6, 2025

Ticket

closes #932

Problem description

Perf team needs TTIR and TTNN MLIR code to be serialized alongside flatbuffer binaries.

What's changed

The compile logic was changed to store MLIR code as necessary and pass it on, all the way to SerializedExecutableInstance class.
Serialization format was changed from just the fb to a custom format that incorporates the needed IRs alongside the fb.
Python side helper function was refactored to work with the new requirements.

Checklist

  • New/Existing tests provide coverage for changes

std::move(original_mlir_code), std::move(executable_name),
m_module_builder->getNumPartitions(),
std::move(original_mlir_code),
std::move(m_module_builder->getTTIRMlirCode()),
Copy link
Contributor Author

Choose a reason for hiding this comment

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

One questionable design decision I made: I made the string getters return a non const reference and move from that reference, leaving the ModuleBuilder with empty strings after constructing an ExecutableImage.
Arguably the interface of buildModule could be refactored to return something like std::expected or StatusOr from compile, instead of returning just the status code and storing results as it's state.

@sgligorijevicTT sgligorijevicTT changed the title Serialize IRs alongside flatbuffesr Serialize IRs alongside flatbuffers Aug 6, 2025
@codecov-commenter
Copy link

codecov-commenter commented Aug 6, 2025

Codecov Report

❌ Patch coverage is 34.54545% with 36 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.47%. Comparing base (8fbb97c) to head (d699b6a).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
...t_implementation/serialized_executable_instance.cc 0.00% 34 Missing ⚠️
.../common/pjrt_implementation/executable_instance.cc 0.00% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #947      +/-   ##
==========================================
- Coverage   68.41%   67.47%   -0.94%     
==========================================
  Files          25       25              
  Lines        1757     1808      +51     
==========================================
+ Hits         1202     1220      +18     
- Misses        555      588      +33     

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

@sgligorijevicTT sgligorijevicTT force-pushed the sgligorijevic/serialize-everything branch from 059032c to d699b6a Compare August 20, 2025 12:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add mlir modules serialization to tt-xla
2 participants