[FEAT] Add structural map, maybe_inplace_mutate, and var_remap APIs

[Kathryn-cat]

Summary

This PR adds structural transformation support to TVM FFI, complementing the existing structural equality, hashing, and walking APIs.

The new StructuralMutator recursively transforms reflected object graphs, while structural_map provides a convenient callback-based interface for compiler passes.

Structural mapping

structural_map follows the same typed callback model as structural_walk. Callbacks are matched by runtime type and return either the unchanged value or a replacement.

Mapping defaults to post-order, so callbacks observe values whose children have already been transformed. This makes bottom-up rewrites such as constant folding straightforward:

  def fold_add(expr):
      if isinstance(expr.lhs, IntImm) and isinstance(expr.rhs, IntImm):
          return IntImm(expr.lhs.value + expr.rhs.value)
      return expr

  optimized = tvm_ffi.structural_map(function, (Add, fold_add))

The PR exposes:

• structural_map in Python.
• StructuralMap and StructuralMapExpected in C++.
• Ordered typed callbacks and grouped callback types.
• Pre-order and post-order transformation.
• Callbacks receiving either value or (value, def_region_kind).
• Structural error context when transformation fails.

StructuralMutator

StructuralMutator provides the low-level transformation engine through two operations:

• Mutate transforms a value without intentionally modifying the input.
• MaybeInplaceMutate permits implementations to reuse an object when doing so is safe.

Custom mutation behavior

Object types can customize mutation using:

• __s_mutate__ for canonical non-in-place transformation.
• __s_maybe_inplace_mutate__ for an optional type-specific in-place optimization.

The maybe-in-place hook owns its safety policy and may reuse the input, delegate to normal mutation, or return another value. A type providing it must also provide __s_mutate__.

Built-in hooks are registered for Array, List, Map, and Dict.

Variable identity remapping

The mutator maintains an identity-substitution environment for FreeVar objects. Once an identity is mapped, later occurrences reuse the same result.

For example:

  def specialize_shape_var(var):
      if var.name == "n":
          return IntImm(10)
      return var

  specialized = tvm_ffi.structural_map(
      function,
      (Var, specialize_shape_var),
  )

If the same Var occurs in a function parameter and its body, both occurrences resolve to the same mapped result.

The mutator also exposes get_var_remap and set_var_remap, allowing custom DAG-style variable wrappers to use their underlying identity object as the remapping key.

[gemini-code-assist]

Code Review

This pull request introduces the StructuralMapper and StructuralMapperObj APIs in include/tvm/ffi/extra/structural_map.h to support structural mapping and in-place mutation of object-backed values. It also adds corresponding custom hooks (kStructuralMap and kStructuralInplaceMutate), registers test leaf objects, and includes comprehensive unit tests. The review feedback suggests two minor optimizations in structural_map.h: avoiding a const_cast when obtaining the field address, and reusing the calculated field address instead of recalculating it for the field setter.

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.

[tqchen]

main issue to to clarify the var handling in callback and testcase simplification and coverage

[tqchen]

these are too basic test, we just need to integrate test StructuralMap with the new instances

[tqchen]

Need test case coverage of inplace mutate case

[tqchen]

One case we need to fix is when value is a freevar, we need a special path for it. Since callback may again update the Var mapped value, in which case we need to override by setting free var map

[tlopex]

agreed, just tested and now preorder and postorder would both fail when the callback directly replaces the same freevar used multiple times. I think we need to have tests to cover them

[tqchen]

Let us have a few compact cases covering high-level usage of StructurualMap, need to be able to test

• inplace happens correctly for Map/Array when certain values are unique, while none implace happens when certain values are shared copies in map
• the hook works e2e as part of StructurualMap, no need to have shallow UT

[tqchen]

One thing that worth double click on is var remapping handling when callback also presence. In an ideal case, the logic should be like follows:

• In MutateExpectedImpl and MaybeInplaceMutateExpectedImpl
• First run a quick check if value is free_var (through TypeInfo and var , two path:
  • not a free var: apply the original old logic
  • value is a free var, look up var remap
    • if exist, return original value(not invoking callback)
    • if missing, run callback + default, or default + callback, in both case, after execution finish, need to override the final value to var remap(since callback may further change the var remapped value)
    • future invocations will hit the var remap and not recurse into callback.

To keep things consistent with structural hash, we can turn on the remapping for both

structural_eq_hash_kind == kTVMFFISEqHashKindDAGNode ||
structural_eq_hash_kind == kTVMFFISEqHashKindFreeVar

This allows node being declared as DAG node also being mapped once on first occurance.

[Seven-Streams]

seems that it does not support classes like ffi::string.

import tvm_ffi

identity = (int, lambda v: v)

tvm_ffi.structural_walk(tvm_ffi.Array([123, "12345678"]), (int, lambda v: None))
tvm_ffi.structural_map(tvm_ffi.Array([123, "12345678"]), (int, lambda v: v))

it cannot work, not sure if expected

[tqchen]

seems something to address, might post the full error message?