[RFC] Hybrid Script Support for TIR

[spectrometerHBH]

Thank you for your comment, @merrymercy!

We call this “hybrid script” because we intend to substitute the old hybrid script with current one. I think maybe we can consider another name if appropriate @tqchen

It is indeed a text format, but we also hope to enhance the syntax to make it as easy to use as old hybrid script and te.compute in the future. Hence, we choose Python which is popular in ML community.

I agree that Python’s AST is not stable, and we should keep it in mind in later development. Now the syntax we use is a really simple subset of Python, which in my opinion may not change abruptly.

[tqchen]

I do believe using python AST is useful, as we can view it(and build features around) in several ways:

• V0: Faithful text format for the TIR. In this case, there is some value to embed things in python:
  • User can directly write AST rather than strings in python to construct these IRs.
  • The code works with the python syntax highlighter.
• V1: Enhanced syntax sugars: After we achieved the goal of V0, we can then bring some enhanced syntax sugars that are more concise. The code below shows an hypothetical example:

@tvm.hybrid.script
def myfunc(a, b):
    A = tir.buffer_bind(a, (100, 100), "float32")
    B = tir.buffer_bind(a, (100, 100), "float32")

    for i, j in grid(100, 100):
        B[i, j] = A[i, j] + 1

Notably, the grid construct is a sugar that expands to two for loops, and we might omit other things like realize etc.

The syntax of V1 is intended to make things easy for developers to concisely construct TIRs. There is a fixed rule to lower these sugars to a well-defined TIR program. Importantly, V1 needs to be consistent with V0 – any V0 program is also a valid V1 program. We can also developer smart printers to potentially print out sugared versions which are more concise.

• V2: potential meta programming. We might want to develop meta-programming models to construct a collection of hybrid scripts, and use of some the python features for macro expansion.

So to sum up, while the text repr is one of the goals of the hybrid script and serves as an anchor for the design. Our ultimate goal should be building a single, unified infrastructure that serves as:

• An text repr for the IR.
• Sugars(DSLs) to construct the IR and compute themselves
• Potential user interface for specifying the program.

[junrushao]

From my past experience, relying on full feature of Python AST is not sustainable approach. For example, AST in python 3.8 differs from 3.7 for no reason, which crashes all my scripts…It is definitely important to find some stable subset.

[tqchen]

One way to achieve this is to define a stable subset of AST object, and then build canonicalization parser that canonicalize the AST to the stable one, and then implement the parse on top of the stable ast.

[ziheng]

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?

[spectrometerHBH]

@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")

[hongh]

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

[tqchen]

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

[hongh]

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

[sanirudh]

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?

[Hzfengsy]

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.

[sanirudh]

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.

[Hzfengsy]

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.

[sanirudh]

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.

[sanirudh]

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.

[Hzfengsy]

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.

[SYangDong]

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!

[masahi]

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

[sanirudh]

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.

[sanirudh]

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.