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