Siyuan Feng
121e1e7a03
The current cross-function calls in TVMScript will cause PyLint warnings,
since the GlobalVar will be marked as undefined vars, e.g.:
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
gv0 = R.call_tir(tir_func, x, R.Tensor((128,), dtype="float32")) # <= `tir_func` is not defined in Python syntax.
return gv0
```
This PR changes the behavior into `TestModule.tir_func` instead of direct `tir_func`
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
cls = TestModule # Use `cls` to refer the current Module
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128,), dtype="float32"))
return gv0
```
NOTE: It's a breaking change, the old style is deprecated.
Additionally, this PR contains the following minor fixes:
- mark `R.function` as staticmethod as what we do for `T.prim_func`
- make `I`, `R`, `T`, `cls` be the builtin keywords for the printer
- define names for functions, modules to prevent naming conflict
- checking the var names is valid via regex expression
- fix typos
119 lines
3.6 KiB
Python
119 lines
3.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""TVM Script Parser utils"""
|
|
import inspect
|
|
from types import FrameType
|
|
from typing import Any, Callable, Dict, List
|
|
|
|
from .diagnostics import findsource
|
|
|
|
|
|
def get_func_nonlocals(func):
|
|
"""A modified version of `inspect.getclosurevars`"""
|
|
|
|
if inspect.ismethod(func):
|
|
func = func.__func__
|
|
|
|
if not inspect.isfunction(func):
|
|
raise TypeError("{!r} is not a Python function".format(func))
|
|
|
|
code = func.__code__
|
|
# Nonlocal references are named in co_freevars and resolved
|
|
# by looking them up in __closure__ by positional index
|
|
nonlocal_vars = {}
|
|
if func.__closure__ is not None:
|
|
for var, cell in zip(code.co_freevars, func.__closure__):
|
|
try:
|
|
nonlocal_vars[var] = cell.cell_contents
|
|
except ValueError as err:
|
|
# cell_contents may raise ValueError if the cell is empty.
|
|
if "empty" not in str(err):
|
|
raise
|
|
return nonlocal_vars
|
|
|
|
|
|
def inspect_function_capture(func: Callable) -> Dict[str, Any]:
|
|
"""Capture function non-locals and global variables.
|
|
|
|
Parameters
|
|
----------
|
|
func : Callable
|
|
The function to inspect.
|
|
|
|
Returns
|
|
-------
|
|
res : Dict[str, Any]
|
|
The function variables map with non-local or global variables.
|
|
"""
|
|
captured = {
|
|
**func.__globals__, # type: ignore
|
|
**get_func_nonlocals(func),
|
|
}
|
|
return captured
|
|
|
|
|
|
def inspect_class_capture(cls: type) -> Dict[str, Any]:
|
|
"""Capture class non-locals and global variables.
|
|
|
|
Parameters
|
|
----------
|
|
cls : type
|
|
The class to inspect.
|
|
|
|
Returns
|
|
-------
|
|
res : Dict[str, Any]
|
|
The class variables map with non-local or global variables.
|
|
"""
|
|
result: Dict[str, Any] = {}
|
|
for _, v in cls.__dict__.items():
|
|
if inspect.isfunction(v):
|
|
func_vars = inspect_function_capture(v)
|
|
result.update(**func_vars)
|
|
return result
|
|
|
|
|
|
def is_defined_in_class(frames: List[FrameType], obj: Any) -> bool:
|
|
"""Check whether a object is defined in a class scope.
|
|
|
|
Parameters
|
|
----------
|
|
frames : List[FrameType]
|
|
The frame stack of the object, obtained by `inspect.stack()`.
|
|
|
|
Returns
|
|
-------
|
|
res : bool
|
|
The result if the object is defined in a class scope.
|
|
"""
|
|
if len(frames) > 2:
|
|
frame_info = frames[2]
|
|
code_context = frame_info.code_context
|
|
if code_context is None:
|
|
return False
|
|
line = code_context[0].strip()
|
|
if line.startswith("@") and "ir_module" in line:
|
|
return True
|
|
if line.startswith("class"):
|
|
lineno = frame_info.lineno
|
|
if lineno >= 2:
|
|
source, _ = findsource(obj)
|
|
line = source[lineno - 2].strip()
|
|
if line.startswith("@") and "ir_module" in line:
|
|
return True
|
|
return False
|