tvm_ffi.structural_map

Contents

tvm_ffi.structural_map#

tvm_ffi.structural_map(root, callbacks=(), with_def_region_kind=(), order='post')#

Structurally map a value using typed replacement callbacks.

Each callback must follow map semantics: it returns the unchanged input or a replacement value and must not mutate its input in place. The policy is mutate(x) = post(D(pre(x))): callbacks run at every occurrence, while only default descent D reads or writes the variable-remap cache. A graph rewrite that must preserve sharing keeps its own callback memo.

At the outer error boundary, after recursive traversal collects visit context, errors are offered once to the original input root type’s __ffi_enrich_visit_error__(error, root) hook. Only the exact root dynamic type’s attribute is used; parent types are not searched. The hook receives the original input root and returns the original or enriched FFI error. An absent or failing hook leaves the original error unchanged. Recursive visitor/mutator methods only collect context.

Parameters:

  • root (Any) – Root value to map. Passing a regular Python reference preserves it through copy-on-write; passing root._move() transfers ownership and permits in-place mutation along unique paths.

  • callbacks (tuple | Sequence | Callable) –

    Normal callbacks. These callbacks receive one argument, value, and return its mapped value. Callback entries are tried in order.

    May be one of:

    • A single callback, used as a typing.Any catch-all.

    • A (type, callback) entry.

    • A grouped ((type1, type2, ...), callback) entry.

    • A sequence of entries.

    Types may be builtins, registered FFI object classes, or typing.Any/object as a catch-all.

  • with_def_region_kind (tuple | Sequence | Callable) – Def-region-aware callbacks. These callbacks receive (value, def_region_kind) and return the mapped value. They accept the same callback entry forms as callbacks.

  • order (str | WalkOrder) – "post"/WalkOrder.POSTORDER (the default) to invoke callbacks after children, or "pre"/WalkOrder.PREORDER to invoke callbacks before children.

Returns:

result – The mapped owning value.

Return type:

Any

Examples

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),
)
Contents