[RFC] Rename Hybrid Script

Development pre-RFC
#1

Current issue

TVM current has two different hybrid scripts: te.hybrid.script and tvm.hybrid.script. This leads to confusion as both scripts are similar but share different use cases and properties. This is especially confusing for new users as hybrid script can refer to either of these names. (I have had discussions with people where we were both referring to different hybrid scripts.)

I propose we rename tvm.hybrid.script as it is going to be used a lot in the future and is newer than te.hybrid.script. I have a couple of possible new names.

  • tir.script: this name reflects the current use of tvm.hybrid.script as a sugar for TIR.
  • tvm.script: this is a similar name to torch’s torchscript. Although tvm.hybrid.script supports only TIR right now, in the future we would like it to support relay. This generic name covers this use case.
  • tvm.script.tir and tvm.script.relay: like tvm.script, but we make the difference between script supporting TIR and script supporting Relay. In the future this may cause an issue if we have a unified script for TIR and Relay (which is a goal of tvm.hybrid.script).

Personally I am leaning towards tvm.script as I think it will make the most sense to new users and communicates the end goals we have.

Which options does everyone like? If you don’t like any of these, could you propose a new name?

Disadvantages

Name changing is churn, and can make old docs/discuss posts out of date. But right now there is not that much usage of tvm.hybrid.script. Renaming it earlier instead of later will cause fewer issues. We can also mark the tvm.hybrid.script identifier as deprecated to allow users time to switch over to the new naming scheme.

The initial tvm.hybrid.script PR: [RFC] Hybrid Script Support for TIR

4 Likes
#2

CC: @spectrometerHBH @Hzfengsy @were

#3

:heart_eyes: :star_struck: Thank you for proposing this! Completely agreed that we should disambiguate.

Having tvm.script and te.hybrid.script is an improvement, but I still think there is far more overlap than ideal between the two.

What about pytir? Then it opens the door to other embedded TIR DSLs in other languages being succeessfully disambiguated as well.

#4

I like the idea of tvm.script as the intention is to have the script to be able to represent both relay and tir module collectively.

2 Likes
#5

tvm.script looks good to me.

#6

tvm.script would be a great name

#7

I don’t think having a single decorator makes sense given that the two scripts will require disambiguation with separate decorators. My recent syntactic rewriting work on Relay shares very little concrete syntax with TIR. Either way we will probably need tvm.script.tir and tvm.script.relay even if we use them to mark function parsed by a single big parser.

1 Like
#8

One thing that is worth considering is how are we going to parse a module with mixed relay and tir functions (or those with external ones).

@tvm.script
class MixedModule:
     def relay_func(x: ty.Tensor):
         return relay.call_tir_dest_passing(tir_func, x)
   
     def tir_func(x: ty.handle)
          ...

Having such mixed module(and unified design) is a motivation for a common decorator. Although I agree we need to think about ways to distinguish the two.

#9

How about this for mixed TIR and Relay:

class MixedModule:
     @relay.script
     def relay_func(x: ty.Tensor):
          return relay.call_tir_dest_passing(tir_func, x)

     @tir.script
     def tir_func(x: ty.handle)
          ...
#10

There are some tradeoffs, for example, given the type system is shared across relay and tir. We might in the future have things like

@tvm.script
class MixedModule:
    MyStruct = ADT[Integer, Integer]    

    def relay_func(x: ty.Tensor):
         return relay.call_tir_dest_passing(tir_func, x)
   
     def tir_func(x: ty.handle)

So parsing might need to go beyond the function level. Additionally, we might also need to serialize the meta-data.

Conceptually, the basic unit in the transformation is an IRModule, and it would be nice to have a decorator on a module to support that(although it is also important to decorate a function as well)

#11

I’ve put up an initial PR here: https://github.com/apache/incubator-tvm/pull/6522.

An issue has come up, what do we name the python module?

Option 1

We name the module tvm.tvmscript. Example usage:

import tvm

# Can still use this though
@tvm.script # or tvm.script.tir
def my_func():
  pass

@tvm.script # or tvm.script.module
class Mod:
  def my_func():
    pass

string = my_func.asscript()
assert(string == tvm.tvmscript.parse(string))

# can also do
from tvm import tvmscript

assert(string == tvmscript.parse(string))

The disadvantage here is that tvm.tvmscript repeats tvm twice. But it does make it explicit that the script we are using is tvm script (as opposed to hybrid script).

Option 2

We name the module tvm.script. We still refer to this as “TVM Script” in all documentation, etc.

import tvm

# Can't use tvm.script as it is a namespace
@tvm.script.tvm # or tvm.script.tir (see option 2a)
def my_func():
  pass

@tvm.script.tvm # or tvm.script.module
class Mod:
  def my_func():
    pass

string = my_func.asscript()
assert(string == tvm.script.parse(string))

# can also do
from tvm import script

assert(string == script.parse(string))

If we use tvm.script as the module name we cannot use the @tvm.script decorator. We have two options for the decorator. Option 2a: use @tvm.script.tvm. Option 2b: use @tvm.script.tir for functions and @tvm.script.module for modules.

The disadvantage here is that the name script can be confusing when used unqualified (when using from imports). Pytorch uses this approach, but they only have a single script in their package.

Let me know which you like best. (Hopefully this isn’t too much bike shedding).

#12

No matter which option we take, do we have to discriminate between function and class when annotating with decorator?

#13

Yes and no. Right now we do not need to differentiate. But in the future, functions in a module may either use be for TIR or for relay.

2 Likes
#14

To move forward, how about we go with option 2b for now.