Skip to content

[TIRx] Generalize expression functor signatures#19931

Merged
tqchen merged 1 commit into
apache:mainfrom
tqchen:tvm-generalize-tirx-exprvisitor-mutator
Jul 3, 2026
Merged

[TIRx] Generalize expression functor signatures#19931
tqchen merged 1 commit into
apache:mainfrom
tqchen:tvm-generalize-tirx-exprvisitor-mutator

Conversation

@tqchen

@tqchen tqchen commented Jul 3, 2026

Copy link
Copy Markdown
Member

Expression unification gives TIRX a shared Expr surface, but its visitor and mutator APIs still expose primitive-only signatures. That mismatch prevents general expressions from flowing through the existing traversal structure and leaves statement traversal with overlapping customization hooks.

This refactor generalizes the existing ExprFunctor, ExprVisitor, and ExprMutator signatures in place to accept and return Expr. Statement visitors and mutators expose a single virtual VisitExpr(const Expr&) hook, while primitive statement reconstruction uses a non-virtual checked VisitPrimExpr helper so invalid narrowing fails at the boundary. Public pre-order and post-order traversal entry points accept general Expr roots.

The existing specialization, vtable, dispatch registration, and class structure remain intact; the change adds no parallel functor, fallback dispatcher, or alternate implementation path.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request removes the SizeVar and SizeVarNode classes, unifying symbolic variables under Var and VarNode. It also updates ExprFunctor and its subclasses to operate on Expr instead of PrimExpr to allow visiting and mutating general expressions. Additionally, BlockBuilder in Relax is updated to track active scopes and blocks to safely unwind them on exceptional exits. The reviewer identified a potential issue in _exit_function_scope where exceptions during block or scope finalization could skip subsequent cleanup steps, and suggested using a try...finally block to ensure robust cleanup.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread python/tvm/relax/block_builder.py
@tqchen tqchen force-pushed the tvm-generalize-tirx-exprvisitor-mutator branch 4 times, most recently from a4ba091 to 9d02d61 Compare July 3, 2026 03:03
@tqchen tqchen force-pushed the tvm-generalize-tirx-exprvisitor-mutator branch 4 times, most recently from 5c57f32 to c232afb Compare July 3, 2026 15:43
Unify TIRX expression traversal on the shared Expr surface. Expose one generalized statement hook while keeping primitive reconstruction behind checked VisitPrimExpr boundaries.
@tqchen tqchen force-pushed the tvm-generalize-tirx-exprvisitor-mutator branch from c232afb to b8e34e0 Compare July 3, 2026 16:18
@tqchen tqchen merged commit d27ed72 into apache:main Jul 3, 2026
9 of 10 checks passed
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.

2 participants