Add Python functor support for TIR expressions and statements #18060
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR introduces Python functor support for TIR expressions and statements by adding new C++ classes and corresponding Python bindings for both visitor and mutator patterns. It also updates import paths and macro formatting to align with the new functor dispatch design.
- Introduces PyStmtExprVisitorNode and PyStmtExprMutatorNode classes in C++ and the necessary Python bindings.
- Updates module imports in the Python interface to reflect the new functor support.
- Adjusts macro formatting in the C++ header for clearer block structure.
Reviewed Changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| src/tir/analysis/buffer_access_lca_detector.cc | Removes an unused lambda capture parameter to simplify the lambda. |
| python/tvm/tir/init.py | Adds imports for the new functor classes. |
| python/tvm/runtime/support.py | Introduces the derived_object decorator for TVM objects. |
| python/tvm/relax/expr_functor.py | Changes the derived_object import and corrects the order of imports. |
| include/tvm/relax/expr_functor.h | Adjusts macro formatting for improved readability. |
| const PrimExpr& binding) -> const ScopeInfo* { | ||
| const ScopeInfo* highest_scope = nullptr; | ||
| PostOrderVisit(binding, [this, &itervar, &highest_scope](const ObjectRef& obj) { | ||
| PostOrderVisit(binding, [this, &highest_scope](const ObjectRef& obj) { |
There was a problem hiding this comment.
The lambda no longer uses the captured variable 'itervar'. Removing it from the capture list improves clarity and avoids potential confusion.
Member
Author
Member
|
cc @jinhongyii |
Hzfengsy
force-pushed
the
tir_py_functor
branch
2 times, most recently
from
40f6d9f to
14f0261
Compare
This commit introduces Python interfaces for TIR functors, enabling Python-side customization of expression and statement visiting and mutation operations. Key changes: - Add PyStmtExprVisitorNode and PyStmtExprMutatorNode classes in C++ - Implement Python bindings for all TIR expression and statement types - Support both visitor (read-only) and mutator (transforming) patterns - Add comprehensive dispatch mechanisms for Python callbacks - Include FFI registrations for Python-C++ interoperability This enables users to write custom TIR transformations and analysis passes directly in Python while maintaining performance through selective Python callback dispatch.
jinhongyii
approved these changes
Jun 16, 2025
ShiboXing
pushed a commit
to ShiboXing/tvm
that referenced
this pull request
Aug 10, 2025
…#18060) This commit introduces Python interfaces for TIR functors, enabling Python-side customization of expression and statement visiting and mutation operations. Key changes: - Add PyStmtExprVisitorNode and PyStmtExprMutatorNode classes in C++ - Implement Python bindings for all TIR expression and statement types - Support both visitor (read-only) and mutator (transforming) patterns - Add comprehensive dispatch mechanisms for Python callbacks - Include FFI registrations for Python-C++ interoperability This enables users to write custom TIR transformations and analysis passes directly in Python while maintaining performance through selective Python callback dispatch.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit introduces Python interfaces for TIR functors, enabling Python-side customization of expression and statement visiting and mutation operations.
Key changes:
This enables users to write custom TIR transformations and analysis passes directly in Python while maintaining performance through selective Python callback dispatch.