[REFACTOR][IR] Add structural hooks for Expr and Stmt - #20275
Merged
tqchen merged 1 commit intoSep 6, 2026
Conversation
tlopex
approved these changes
Sep 5, 2026
tqchen
force-pushed
the
redo-expr-stmt-structural-hooks-current
branch
7 times, most recently
from
September 6, 2026 16:51
73402f2 to
e401670
Compare
Add compiled structural visit, mutate, and maybe-in-place-mutate hooks for every Expr and Stmt node family using the current tvm-ffi structural-hook ABI. The branch pins tvm-ffi at `9d784c4da74ff7360d76c79a89fe60a63516f880`, including the callback-owned `StructuralMutate` API from tvm-ffi apache#751 and the visit-hook success-tail macro from apache#753. Checked Map-valued hook fields use the current Expected-returning type-error path. `StmtNode` and `TypeNode` declare `_type_s_eq_hash_subclass_kind_fixed = true`. The current hierarchy census finds 18 final registered children under Stmt and 14 under Type, with no FreeVar or DAG kind override in either hierarchy. This lets typed structural-map callbacks over Stmt and Type fold out the identity-remap branch. `Expr`, `PrimExpr`, and `BaseFunc` retain the default because their descendants include FreeVar or DAG kinds. Definition-region annotations are aligned with reflected fallback behavior: direct definitions use non-recursive regions, container-owned definitions are established at their Var sites, and dynamic Var type descent remains a use. The hooks preserve the established visitor/mutator field set and same-object short circuits on unchanged mutations. Call hooks keep ordinary descent for args and nonempty type arguments, while skipping primitive result types, interned Op operators, and empty type-argument containers. Dynamic result types, function-valued operators, and populated type arguments retain ordinary descent. `PostOrderVisit` and `Substitute` remain unchanged. Direct `StructuralMap` roots in this PR are confined to focused tests: the intentionally shared lvalue exercises copy-on-write, while owned roots move into the by-value API to preserve eligible in-place reuse. All eight hook source files use one anonymous namespace and non-static file-local hook functions. Rebuilds box moved typed copies directly, unchanged paths box typed `self`, and constant leaves explicitly extract typed `self`. Every visit hook uses direct early-return macro calls per field and terminates with `TVM_FFI_S_VISIT_RETURN_NONE()`. `MaybeInplaceMutateSeqStmtRaw` separately verifies unique ownership of its `seq` field and borrows each element from raw array storage so the element remains eligible for in-place mutation. Mapped `Evaluate(0)` entries are omitted just as they are by `SeqStmt::Flatten`; zero or one remaining statement normalizes to `Evaluate(0)` or the sole statement without constructing a prohibited SeqStmt. This does not change the SeqStmt constructor or contract. Both SeqStmt mutation hooks scan their unchanged prefix with only one same-object test per element, then tail-return to changed-path helpers at the first divergence. Ordinary mutation allocates and initializes its unchanged prefix only there. In-place mutation begins its `total` cursor there: drops and splices compact into the existing array while writes remain behind unread input; a splice that would cross that boundary transfers once to an `InplaceSplice` output and continues the remaining callbacks exactly once. Differential tests cover first and last positions, multiple splices, growth, size preservation, `Evaluate(0)` shrink, zero/one-result normalization, array reuse, overflow replacement, and callback non-replay without constructing raw empty or size-one SeqStmt nodes. Owner measurements in the apache#373 harness on a 256-element SeqStmt record the accepted timing guardrail. Two replacements change from 23,511 to 24,444 ns (+4.0%), 16 from 24,430 to 24,768 ns (+1.4%), 64 from 25,464 to 25,132 ns (-1.3%), and 128 from 27,367 to 25,245 ns (-7.8%). Rebuilds fall from 3-5 to zero at every density, with crossover near one quarter. A follow-up L=8 allocation check finds that repaired shrink reuses the node array instead of allocating once, while growth remains one allocation with or without spare capacity. This repair lands for correct element-level in-place semantics and reduced splice allocation, not as a uniform speedup. Earlier general performance measurements for this work in apache#364/apache#365 were collected at the previous `556514c76225cb35a55e0649098b07be3fba2991` engine pin and predate this bump; they should not be interpreted as measurements of the newer engine. The SeqStmt-specific figures above are the later apache#373 measurement of the accepted repair tradeoff.
tqchen
force-pushed
the
redo-expr-stmt-structural-hooks-current
branch
from
September 6, 2026 18:09
e401670 to
a1031a2
Compare
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
Add compiled structural visit, mutate, and maybe-in-place-mutate
hooks for every Expr and Stmt node family using the current tvm-ffi
structural-hook ABI.
StmtNodeandTypeNodedeclare_type_s_eq_hash_subclass_kind_fixed = true. The current hierarchy censusfinds 18 final registered children under Stmt and 14 under Type, with no FreeVar
or DAG kind override in either hierarchy. This lets typed structural-map
callbacks over Stmt and Type fold out the identity-remap branch.
Expr,PrimExpr, andBaseFuncretain the default because their descendantsinclude FreeVar or DAG kinds.
Definition-region annotations are aligned with reflected fallback behavior:
direct definitions use non-recursive regions, container-owned definitions are
established at their Var sites, and dynamic Var type descent remains a use. The
hooks preserve the established visitor/mutator field set and same-object short
circuits on unchanged mutations.