[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]

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.

[tlopex]

var_remap_ should use object identity, but ffi::Map may treat different freevars as the same key, likely the second freevar may reuse the first's mapping

[tqchen]

This is fine given the assumption that var is defined once

[tqchen]

Comment on simplifying the str/bytes handling