The GitHub Actions job "mainline-only" on tvm-ffi.git/main has succeeded. Run started by GitHub user tqchen (triggered by tqchen).
Head commit for run: ca8803c9f34d26cdc665f5dc5adfc0bc816771eb / Kathryn (Jinqi) Chen <[email protected]> [FEAT] Add structural map, maybe_inplace_mutate, and var_remap APIs (#649) ## 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: ```py 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: ```py 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. Report URL: https://github.com/apache/tvm-ffi/actions/runs/31765623403 With regards, GitHub Actions via GitBox --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
