[RFC] TensorIR: A schedulable IR for TVM

[jcf94]

It’s really a great work! @Hzfengsy

1. Have you tried tensorize intrinsic(e.g. TensorCore schedule) using this new IR? Since I remember that to support tensorize is also one of your initial motivations.

I’m curious if this would requires a different programming logic writing such schedule. Before seeing any example, I feel hard to imagine how the schedule would be.

2. The if_then_else case is really a strong point!

The original implementation of te.compute makes spatial iterators to follow the shape of the output Tensor, which limit us to write more flexible computations using TVM.

But seems this would break the original design of Stage, or do we still have a concept of Stage in the new TensorIR? How to schedule them?

3. Looking forward to upgrade Ansor based on the new TensorIR.

Thanks!

[ds1231h]

Great work! I believe this will make “scheduling” more flexible and intuitive!

However, will this increase the coupling between the schedule and the lower pass, which may lead to an increase in the complexity of the lower pass?

By the way, I’m also looking forward to know how to auto-schedule based on the the TensorIR.

Thanks!

[spectrometerHBH]

Thanks for your reply! @jcf94

A1. We’ve tried to tensorize intrinsic using this new IR, and are working on the TensorCore demo. Our design is really close to the original tensorize programming logic, only differs in the declaration of description&implementation of HW intrinsic (we can use Hybrid Script to program them more conveniently)

A2. There is no concept of Stage in TensorIR. What you can see in Hybrid Script contains full IR information and our schedule transforms IR directly. The basic schedule unit is Block.

A3. Ansor + TensorIR is on our roadmap.

[spectrometerHBH]

@ds1231h

However, will this increase the coupling between the schedule and the lower pass, which may lead to an increase in the complexity of the lower pass?

Thanks for your reply! @ds1231h

At the moment, we at first transform TIR with block to TIR without block in one pass(the latter is current TIR in the main repo) and utilizes the current TIR lowering passes to do codegen. And we can gradually refactor this lowering process in the future.

[aca88]

Hi,

Even though I don’t think I understood everything, I like the idea of solving some of the limitations of te.compute. Since the te.compute is in a central part of the TVM stack changing it requires a lot of work and understanding. So thank you all for continuing such development.

Q1: I was wondering how this fits in the Relay ->Topi->TE->TIR flow. In a more specific case take FuseOps. AFAIK the FuseOps pass creates the multi-stage operator based on the te.computes. Since you mentioned that there would be no notion of a stage in the new TensorIR, how would FuseOps work? and more generally how TVM’s philosophy of “defining a compute rule and a separate schedule” be changed?

@Hzfengsy

TE has limited expressiveness since each stage is defined by Stage = te.compute(lambda expr) , while TensorIR is a full c+±like IR. We can write any program with TensorIR as you want. Although not all programs can be scheduled, there are still more workloads that can be optimized by TensorIR.

Q2: What exactly do you mean by “not all program can be scheduled”? maybe an example?

Q3: You mentioned “new scheduling primitives”, could you maybe give a list?

EDIT: Q4: Expected timeline of the steps?

[ds1231h]

Well-received with thanks!

[Hzfengsy]

Thank you for your interest.

A1: Current op fusing is based on stage but the critical point is fusing the injective computation. We can also inline injective computation by traverse_inline. So there is no doubt that FuseOps works. As for the philosophy, I think there are only few changes. TIR is not only an IR but only can be a computation declaration. We provide very user-friendly API (as easy as TE) to define compute rules.

A2: TIR is a general IR which can represent almost every program. It’s really hard to schedule a general program, but we promise TIR can schedule all programs which TE can schedule.

A3: Most of the primitives are similar to the TE ones. For now, only two new primitives are decompose_reduction and merge_reduction

A4: Upstream on Oct and Nov. Ansor supporting is WIP

I hope I can answer your question.

[kevinthesun]

Thank you for this proposal! This work does make scheduling much easier. I have a concern about using this way to write a tensor expression. It looks like more complicated than tvm.compute when defining matmul. We need to define some buffers and creating block with corresponding shape dimension. It would be helpful if you can add a conv2d example which can replace existing topi.nn.conv2d definition to better understand what developer would need to write.

Another question is about representing generic programming style ops such as shape functions. Since these programs don’t fit into tvm scheduling, I assume it would still be more convenient to use existing te hybrid script to create these ops?

[spectrometerHBH]

Thanks for your reply! @kevinthesun

@kevinthesun

Thank you for this proposal! This work does make scheduling much easier. I have a concern about using this way to write a tensor expression. It looks like more complicated than tvm.compute when defining matmul. We need to define some buffers and creating block with corresponding shape dimension. It would be helpful if you can add a conv2d example which can replace existing topi.nn.conv2d definition to better understand what developer would need to write.

In original te programming, we also have to declare buffers and create lambda expression with iter_vars having the correct shape dimension. If we take this additional info into account, TE programming is close to Hybrid Script’s programming in complexity.

Currently, we can not replace existing topi operators, since they are represented by Stage/Op and optimized by te schedule, while Hybrid Script will be parsed into TIR directly.

@kevinthesun

Another question is about representing generic programming style ops such as shape functions. Since these programs don’t fit into tvm scheduling, I assume it would still be more convenient to use existing te hybrid script to create these ops?

If we don’t have to schedule the PrimFunc, we don’t have to declare blocks in TIR. Actually, TE hybrid script is also a text representation of TIR to a large extent, with loop and condition statement directly representing the IR structure, using sugar like variable being translated to Array of size 1 to ease the format of Store&Load. At the moment, we haven’t introduced sugars to simplify such Load&Store, but the rest writing are largely simlilar.

[tqchen]

TIR and TE do not conflict with each other. TE is still a useful DSL to stitch fragments of TIR together to form a PrimFunc.

We could still define TE based DSL(backed by TIR) that enables primitives like compute and hybrid calls to stitch together a dataflow graph to form a PrimFunc And then use the TIR for scheduling.

[comaniac]

Thanks for the proposal! This definitely opens more opportunities for performance optimization. Two questions for clarification:

1. IIUC, based on the proposal and discussion, we will have both TE and TIR, but TE is more like a frontend wrapper of TIR to serve some users that prefer to write high-level DSL. Then, what will we do with the TE schedule primitives? Intuitively, we should still keep them; otherwise TE writers will have no way to schedule their computes, because they know nothing about TIR and blocks.

2. Does this proposal support dynamic shape (i.e., Any)? For example, can we have something like:

   @tvm.hybrid.script
   def matmul(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
       C = tir.match_buffer(c, (1024, 1024), "float32")
       A = tir.match_buffer(a, (1024, Any), "float32")
       B = tir.match_buffer(b, (Any, 1024), "float32")
       reducer = tir.comm_reducer(lambda x, y: x + y, tir.float32(0))
   
       with tir.block([1024, 1024, tir.reduce_axis(0, 1024)], "C") as [vi, vj, vk]:
           reducer.step(C[vi, vj], A[vi, vk] * B[vk, vj])
   
   s = tir.create_schedule(matmul)
   update = s.get_block("C")
   i, j, k = s.get_axes(update)
   i_o, i_i = s.split(i, bn)
   j_o, j_i = s.split(j, bn)
   k_o, k_i = s.split(k, 4)
   

   In this case, the length of vk (or k) is Any. Can we still apply split to it with a fixed factor

[kevinthesun]

Thanks for explanation. The relation between te and new tir is now more clear to me.

[kevinthesun]

Thanks for clarification. It would be nice if we can use various methods to create tensor programs and use new tir to schedule them.

[Hzfengsy]

Good questions!

1. As for as we know, we would like to let users use TensorIR schedule rather than TE schedule one we fully upstream the TensorIR. For three reasons:

   1. Just as you have mentioned, TE is a fronted wrapper, and it directly generates TIR with blocks. Somehow, TE is more like a sugar to define TIR.
   2. Most of the schedules and primitives in TensorIR are very similar to those in TE. The cost of learning TensorIR schedule is extremely low (maybe just one day).
   3. All primitives are based on the block (no stage concept in TensorIR schedule). It’s hard to keep the TE schedule with block

2. Dynamic shapes are not supported now. However, thanks to our de-coupled primitives, it’s ok to support later.

[kevinthesun]

Would love to see dynamic shape supported otherwise a large set of models can’t be backed by new TensorIR.

[comaniac]

So the scenario is like you can choose to use TE or TIR to write a compute, but if you choose TE, you have to first lower it to TIR and then add schedule primitives?

IIUC, it seems to me that this is nontrivial, because TIR was not written by human and you may need to first print it out to figure out how to schedule it. It sounds more straightforward to keep TE schedule as syntactic sugar. At least you can get the sense about how the schedule looks like by tracing the Python code.

[tqchen]

Because there is a 1-1 mapping between te.Stage and Block. It should actually not be hard to use tir schedule to schedule a te compute generated PrimFunc (either by getting block via name, or pragmatically traverse the blocks like we do pragmatically on stages). But i agree that we can keep te.schedule for a bit.

[comaniac]

Thanks for clarification. Make sense to me.

[MinminSun]

Thanks for the proposal! Just courious about the schuedule primitives like cache_write and cache_read, since there are no stages in TensorIR.

[spectrometerHBH]

Thanks for your reply! @MinminSun

The cache_read/cache_write API accepts a Buffer and new scope as input, do some checks to ensure it brings no problem to read/write the Buffer into cache, and create new blocks to do the cache transfer.