Visit AST in Python

I wanted to write my own visitor for the AST from an expression via the Python API.

For instance I create an IRModule like this:

import tvm
import tvm.testing
from tvm import te

def vec_add(N=32):
    A = te.placeholder((N,), name="A")
    B = te.placeholder((N,), name="B")

    C = te.compute((N,), lambda n: A[n] + B[n], name="C")

    operands = [A, B]

    return C, operands

C, operands = vec_add()

s: tvm.te.schedule.Schedule = te.create_schedule(C.op)
ir: tvm.ir.module.IRModule = tvm.lower(s, operands + [C], simple_mode=True)

I know I can use the tvm.tir.stmt_functor.post_order_visit to visit the Nodes in the AST, but I’d like to write my own visitor. e.g. something like extending the class NodeVisitor for the python AST (python - Simple example of how to use ast.NodeVisitor? - Stack Overflow).

Any hints?

Thanks in advance

I’m not much of an expert here, and I’m not really sure what the motivation is to implement a separate visitor instead of using the post_order_visit, but maybe you can take a look at the visitor support implemented for the Relay AST in expr_functor.py

It can be used as shown in the accompanying tests in tvm/test_expr_functor.py at main · apache/tvm · GitHub

Is that what you’re looking for?

1 Like

Thanks this is perfect!

1 Like