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.
