tqchen opened a new pull request, #741: URL: https://github.com/apache/tvm-ffi/pull/741
Reworks the `StructuralMap` mutation engine. Hooked `StructuralMap` was slower than the mutator it replaces, and the cause was structural rather than incidental: the engine moved every result through representations the compiler could not eliminate, and expressed the walk in a shape that could not inline. ## Result representation Hooks are C-ABI function pointers returning `TVMFFIAny`, a 16-byte POD that lives in registers. Wrapping a result in `Expected<Any>` forces it to memory, because `Expected<Any>` is not trivially destructible and so is classified MEMORY under the SysV ABI. The convention is now uniform, matching what the visit hooks already did: > **The ABI boundary is a raw `TVMFFIAny`. Hook and callback bodies work in `Expected<Any>` and move out at that boundary.** Descent through an unmatched node *is* that boundary, so it stays raw end to end. `TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN` still takes a typed `Expected<Any>` as input — callers keep writing `mutator->MutateExpected(...)` — and moves out only on return. It introduces no helper of its own, using the `AnyUnsafe` and `ExpectedUnsafe` entry points that already existed. ## Error context Hooks no longer attach it. The engine names a node where it dispatches into that node, giving **exactly one frame per node**; previously a hook named its own node and the engine named it again on some paths. Both `MAYBE_EARLY_RETURN` macros lose their node argument as a result, and the visit engine's own walk no longer routes through the hook-facing macro. ## Walk structure The callback chain was continuation-passing: a lambda nested per link, with the walk-order body at the bottom of the nest. The mutator now owns its callbacks and tests them with a `(... || ...)` fold, and the matched-node path — selection, identity remap, descent, invocation — is one straight-line function. ## Selection strategy Statically typed callbacks and runtime `ffi.Function` links differ only in how a link is chosen, but that difference ran through the middle of the walk. They are now separate classes — `StructuralMapMutatorObj` in the header, `StructuralMapDynMutatorObj` in the `.cc` — sharing only the identity remap in `StructuralMapMutatorBaseObj`. This removes a hazard rather than relocating it: the dynamic form recorded its selected `Function` in mutable state on one shared instance, where a post-order walk could descend into a matching child and overwrite the selection its parent was about to use. Selection is now a local that descent cannot reach. ## Removed `TVM_FFI_S_MUTATE_ASSIGN_FROM` existed only to skip a conversion that was expensive while the checked macro built its mismatch message eagerly; it and `SMutateResultUnchanged` are gone, along with `AssignOrReturnHelper` and the conversion helpers that existed to recover a type the preprocessor had swallowed. `MutateWithIdentityRemapExpected` had a single caller, two branches that were dead at it, and re-derived a check `IsRemappableIdentity` already owned. ## Behavior change When descent hands back a node whose type no longer matches the link that selected it, the engine now reports a `TypeError` instead of silently passing the node through. Default mutation is required to preserve a node's type, so this can only happen when a hook has broken that invariant. Everything else is unchanged: match-before-descent ordering, identity-remap entry and exit, and the pre-order in-place rules. ## Also `AnyUnsafe::CheckAnyStrict` was missing the `Any` special case that both of its siblings already had. ## Performance Measured on a TVM split/fuse expression, `-O3`, ns per node, medians of 11 pinned process runs. Both sides are self-consistent builds differing only in this change. | | before | after | change | |---|---:|---:|---:| | hooked `StructuralMap` (replace) | 39.35 | **28.58** | **−27.4%** | | hooked `StructuralMap` (no-op) | 38.26 | **28.90** | −24.5% | | minimal-vtable floor | 20.41 | 13.25 | −35.1% | | plain recursive descent (control) | 2.97 | 2.99 | +0.7% | A standalone harness mirroring the same node hierarchy, with no TVM dependency, reproduces the engine cost above its own floor to within 1% (+15.51 vs +15.33 ns/node), so the result is not specific to one downstream tree. The remaining ~+15 ns/node over the minimal-vtable floor is the matched-callback path — callback invocation plus identity remap on matched free variables. This PR does not address it. ## Review notes - **`TryLink` is 138 lines** in the header and 97 in the `.cc`. That is the cost of folding the matched-node path inline and duplicating the walk rather than sharing it. Both deliberate; nothing else in the file exceeds 18 lines. - **The two `TryLink` copies must stay in step.** They may differ only in how a link is found and called; their error-node sequences are identical today, and a `StructuralMapDyn` test covers the dynamic copy, which was previously reachable only from Python. - **`TVM_FFI_UNSAFE_S_MUTATE_ASSIGN_OR_RETURN_SKIP_CHECK`** is the one added public macro. It is undefined behavior in a release build on a wrong-typed result and has **zero call sites** here; it exists for hooks that have measured the check to matter. ## Verification C++ test suite: 477 pass. Downstream TVM builds clean against this. The error-context frame count was probed directly rather than reasoned about. https://claude.ai/code/session_018jQX1azGgoDx3zT6BwsGAk -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
