[RFC] Better type hinting for TVM scripts

Current TVM’s scripting mode (TIR’s hybrid script/TVM’s hybrid script) uses sub-language embedded in python frontend, user cannot use type hinting and auto completion tools because the code is not parsed by python.

We can create a .pyi stub for these keywords which only annotates type and leave the implementations empty (actually we don’t need them). The stub could also help us generate language reference for user to grasp the syntax and semantics of TIR scripts.

For example, if I write a __init__.pyi under tvm.tir (from simplicity, it only cover a few keywords):

from typing import Tuple, Union, List
from .schedule import LoopRV, BlockRV
from ..ir import PrimExpr

class serial:
    def __init__(self, start: PrimExpr, end: PrimExpr) -> None: ...
    def __enter__(self) -> LoopRV: ...

class grid:
    def __init__(self, *args: Union[int, PrimExpr, ]) -> None: ...

    def __enter__(self) -> LoopRV: ...

class block:
    def __init__(self, *args: PrimExpr, name: str = '') -> None: ...

    def __enter__(self) -> List[LoopRV]: ...

The purpose of this stub is to create an environment where user can write TVM scripts as fluent as general python programs, with normal type hinting and auto completion when user write their TIR scripts.

To achieve this purpose, we want to retain the semantics of TIR keywords in Python. However, the API of tir functions nested in hybridscripts doesn’t necessarily align with that of functions defined under tvm.tir python module, the alignment requires some effort. We can also create another namespace (module) for keywords and import them via import tvm.tir.script.namespace as T.

Besides the alignment issue, I want to hear more from the community about the proposal, feel free to leave your suggestions and concerns here.


Thanks for the proposal! It is definitely extremely important for us to have better type hints and documents for TVM scripts.

One quick question I have been thinking of: is it possible to unify and prevent divergence between tir.XXX in the script and tvm.tir namespace? If so, there probably won’t be necessary to introduce a new namespace tvm.tir.script.namespace

1 Like

CC @tqchen @Hzfengsy @spectrometerHBH

I don’t think they can be merged into one. It has two reasons:

  1. There are some more things besides AST nodes in tvm.tir namespace, e.g., tir.transform
  2. The contractors for the same node are different for script and tir. script tries to provide the simplest format while the tir need to handle all kinds of attrs

Thanks for the proposal. Personally, I totally agree with that :+1: