TVM's AST usage and documentation

Hello all,

I am wondering where I can get more information on the AST of TVM.

My current understanding is as follow:
From reading Dev’s API:
The AST can be composed of Expressions (tvm.expr) and Statements (tvm.stmt), which are both derived classes of the Node (tvm.node) class.
From reading Adding a Comp. Pass:
Also, there are AST traversers which are used for IR passes.
IR passes which only gather information embedded in the AST are called ExprVisitors and those which change the AST are called ExprMutators. These are derived classes fron the ExprFunctor class.

So technically, if I wanted to analyse the AST I would have to implement an ExprVisitor.
1. So far so good?
2. If there are tvm.expr and tvm.stmt, do they all get visited by ExprVisitor?
3. Are there any other kinds of AST nodes which I havent listed?

I think need to understand how the AST is built to know which subclasses of tvm.expr and tvm.stmt I need to visit (and implement the visitor)
4. Where can I find the description of the method to generate the AST?

Since what I plan to do would be right before code generation, I guess I need more info on how the AST is read
5. Where can I find the description of the method to read the AST? or is this not relevant since there is no need to reimplement any other tree traversal method?

The reason why I want to understand the AST in more detail is because I would like to generate something similar to a Control Flow Graph (CFG) of the code represented by the AST.
Assume I am compiling for the VTA, I would like to get the CFG of the calls to the VTA’s Runtime. Another analogy would be calls to linear algebra libraries.

I thought about using Clang for generating this, but I don’t see the C code generator being able to represent everything (one specific example is the tvm.parallel schedule primitive, which gets ignored in C). Also, the only way I could get C code was setting tvm.lower(...,simple_mode=false) which generates a lot of code for TVM Runtime API (which I don’t really want to have in my analysis).
7. I was wondering of the language used when I call tvm.lower(..., simple_mode=false), is this Rust, Relay, HalideIR or TVMIR? or is it just a pseudo-language only used for TVM printout?
8. Where can I find the Abstract Grammar used to build the trees? (see here for what I mean)

Thanks :slight_smile:

4 Likes