[RFC] Hybrid Script Support for TIR

Development pre-RFC
#21

Nice work @spectrometerHBH!

I would like to give it a try! One question, how can I build the tir.function after I define it with the script? Is the path has been pushed on master branch now?

1 Like
#22

@ziheng Sorry for my late reply. The code of tvm.build in build_module.py on master branch is attached below.

After we define the TIR PrimFunc with the script, we should be able to put it inside an IRModule and using tvm.build as normal.

if isinstance(inputs, schedule.Schedule):
    if args is None:
        raise ValueError("args must be given for build from schedule")
    input_mod = lower(inputs, args, name=name, binds=binds)
elif isinstance(inputs, (list, tuple, container.Array)):
    merged_mod = tvm.IRModule({})
    for x in inputs:
        merged_mod.update(x)
    input_mod = merged_mod
elif isinstance(inputs, tvm.IRModule):
    input_mod = inputs
elif not isinstance(inputs, (dict, container.Map)):
    raise ValueError("inputs must be Schedule, IRModule or dict of target to IRModule")
1 Like
#23

According to the discussion above, @tvm.script.tir mainly aims to extend te + schedule. Now all relay ops are implemented through te.compute + schedule to lower to TIR on master branch. If want modify one op to the new script method and then i realize a tir.function by tvm.script.tir decorated, how should I replace the current te.compute + schedule process

In addition, All @tvm.script.tir code on the master existed in tvm/tests/python/unittest/. I donā€™t find this script used in our core function-code. When will it be officially launched in the core code? Can you give an example about how to use it?

Take the get_valid_counts op as an example:

@override_native_generic_func("get_valid_counts_strategy")
def get_valid_counts_strategy(attrs, inputs, out_type, target):
    """get_valid_counts generic strategy"""
    strategy = _op.OpStrategy()
    strategy.add_implementation(
        wrap_compute_get_valid_counts(topi.vision.get_valid_counts),
        wrap_topi_schedule(topi.generic.schedule_get_valid_counts),
        name="get_valid_counts.generic",
    )
    return strategy

wrap_compute_get_valid_counts(topi.vision.get_valid_counts), wrap_topi_schedule(topi.generic.schedule_get_valid_counts), this part how to realize by using @tvm.script.tir?@spectrometerHBH @tqchen

#24

the tir support is still being upstreamed atm. The main has not yet switch any of the codegen path to use the tir.

#25

About this part function, has upstream plan? Need extend the relay op strategy to support this tvm.script.tir landing?

#26

I see that the TIR tracking issue has been closed and looks like all the support has been upstreamed. Is tvm.script.tir usable to write ops now, and are there any plans for a documentation/user guide on how to use tvm.script to write ops?

#27

Thanks @sanirudh for your interests. TVMScipt can write and schedule ops now. We are preparing a blitz course to show how to write and schedule a IRModule by TVMScript and TensorIR. It will be upstreamed soon (before v0.8 release). Also more detailed tutorial and dev docs are also in our plan.

#28

Great, thanks a lot for the quick reply. I would love to see blitz course. Iā€™ve been trying to translate a complex cpp custom op into a TVM op, and using tensor expressions turned out to be too complicated with unnecessary if_then_else statements.

So, after I saw that the TIR tracking has been closed, I started trying to implement using TVMScript, but understanding what to do by going through the source and unittests was getting tedious and in some cases not enough to understand properly, which is why I asked that question.

Thanks again for the amazing work, this might make it really easy to implement custom ops in TVM without having to resort to TE tricks and hacks.

#29

We try to move fast. Also we have updated the script namespace(link) and will update block syntax (link). You can take a look at them if you are interested.

1 Like
#30

Great, thanks a lot. I actually did get confused as to when tir.block vs nested for loops should be used, so yeah this change makes a lot of sense.

Thanks a lot for the info.

#31

Is there a way to call a topi compute from within TVM Script prim_func. For example, if we would like to implement an op like NMS(Non-Max Suppression) using TVM Script, is there a way to call the topi.argsort instead of defining our own sorting function.

Thanks.

#32

It is not support to call topi compute inside TVMScript. However, you can convert TE compute to prim_func by calling create_prim_func. Then copy-paste may work for you.

#33

Is there a way to call these functions rather than inline all functions? Thatā€™s useful when implementing ops like nms, proposal, etc. Thank you!

#34

See the discussion in [RFC] TensorIR: A schedulable IR for TVM - #60 by masahi. Currently we cannot meta-program tvmscript, so we can only write a monolithic piece of code. There is an example of NMS written in hybrid script, but we shouldnā€™t encourage this style of programming (hard to read and maintain etc). https://github.com/apache/tvm/blob/main/python/tvm/topi/vision/nms.py#L244

I think meta programming support for tvmscript (by a macro system, for example) is a big open question. Iā€™m very interested in this topic too, because often I need to write a complicated op like NMS, FFT etc and the only possibility is to use the IR builder. cc @junrushao @vinx13

2 Likes
[RFC] TensorIR: A schedulable IR for TVM
#35

Thanks for the suggestion. That works great. Just out of curiosity, is there a plan to support some way to define and call functions while writing TVMScript. The way I can understand it, one of the biggest selling points for TVMScript would be to enable writing complex ops, ones that are too complicated to be written in TE.

But on the other hand, if we thereā€™s no easy way to define functions, it might push people away from adopting TVMScript, and go back to defining all computes using te.

Another question I had in mind is, if we define an IR_Module with multiple prim_func's would it be possible to support calling one prim_func's from another? That might make writing IR Modules much much easier I think.

I understand that this might not be the goal of TVMScript, and Iā€™m probably too late to the discussion, but just wanted to understand your thoughts on this.

Thanks again for the neat suggestion on calling topi scripts.

#36

Right, I did look at that discussion and I agree that it might make more sense to be able to modularize our code. Thanks for the answer.

One other method of thinking I had was to maybe allow calling other TVM prim_funcs and then write a pass in TVM if needed to either inline those prim_funcs or flatten it the way AOT executor is planning to do (either automatically or when manually defined). Maybe this way, even operators can be modularized into multiple prim_funcs, which would make long term maintenance much easier, especially when working with operators like NMS.

Thanks.

#37

Thanks Masa! I believe meta programming is definitely something we should look at after Siyuanā€™s current major refactoring is done.

The first step is to ping down the syntax (maybe TE-flavored?), then enhance the parser. Would be nice if you are interested to bring up proposals :slight_smile:

CC: @Hzfengsy

2 Likes
#38

The tutorial PR is on: https://github.com/apache/tvm/pull/9315

Comments and suggestions are welcomed

3 Likes