tvm_ffi.structural_map

Contents

tvm_ffi.structural_map#

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

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.

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, default: ()) –

    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, default: ()) – 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, default: "post") – "post"/WalkOrder.POSTORDER (the default) to invoke callbacks after children, or "pre"/WalkOrder.PREORDER to invoke callbacks before children.

Return type:

Any

Returns:

result – The mapped owning value.

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