Where can I find the source code of TVM's nested loop representation?


I’m new to TVM, and I’m reading TVM’s source code. Can anyone help me to find the source code of TVM’s internal nested loop representation? Also, the location of the scheduling (e.g., tiling, loop permutation) source code can be very helpful.

Thanks a lot!

I’m not quite sure about what you mean on the nested loop representation.

For the following code:

for i = 0 to 100
    for j = 0 to 100
        for k  = 0 to 100
            C[i,j] = A[i,k] * B[k,j]

This will be a tvm.tir AST.

These three for are ForNode of the AST.

And the location of scheduling should be easy to find. For example, you can have a look at python/tvm/te/schedule.py, include/tvm/te/schedule.h, src/te/schedule/schedule_lang.cc.

1 Like