With `from __future__ import annotations`, Python stores annotations as
strings
and does not capture annotation-only variables in `__closure__`. This
broke
TVMScript when buffer shapes/dtypes referenced closure variables.
Fix: wrap `extra_vars` in a `collections.ChainMap` with snapshots of all
live
caller-frame locals (from `inspect.stack()`) as fallback layers in both
`tir/entry.py` (`prim_func`) and `ir/entry.py` (`ir_module`). The
`ir_module`
function now also captures `outer_stack = inspect.stack()` at its entry
point,
mirroring the existing pattern in `prim_func`. Lookup falls back to
frame locals
only on cache miss, preserving existing behavior for non-PEP-563 code.
Add `tests/python/tvmscript/test_tvmscript_pep563_closure.py` (requires
`from __future__ import annotations` at the top) covering closure
variables in
buffer shapes, dtypes, nested scopes, ir_module, and mixed
annotation+body use.
This PR enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
* [TVMScript][Bugfix] Normalize relax::If with function's TIR var
Prior to this commit, the branches of `relax::If` were normalized
using `EraseToWellDefinedInScope`, using a fresh variable scope.
While this had the intended behavior of preventing variables defined
in a single branch from being usable outside of the conditional, it
also caused the conditional's branches to treat function-scope
symbolic variables as if they were undefined.
This commit updates the `tvm::relax::Normalizer` so that `relax::If`
is normalized within an inherited scope. This preserves the previous
behavior for symbolic variables defined within a branch, but allows
shapes within a branch to use symbolic variables defined outside of
the branch.
* [Relax] Canonicalize known symbolic shapes in Relax expressions
Prior to this commit, known constants in Relax functions would be
inlined by the `CanonicalizeBindings` pass, but only if they appeared as Relax
expressions (e.g. `R.const` or `R.prim_value`). Known constants that
appeared as TIR variables (e.g. symbolic shapes) would be kept as
dynamic parameters, even if they were known at compile time.
This commit updates the `CanonicalizeBindings` pass to identify known
values of symbolic shapes, and to use these known values in shape
expressions.
* [Relax][Refactor] Reorganize pattern-matching
A follow-up to https://github.com/apache/tvm/pull/16730. Now that the
implementations for `rewrite_call` and `rewrite_bindings` are in
separate classes, they can be further split out into separate files.
* [Relax][Refactor] Implement Rewriter class for pattern-rewrite
Prior to this commit, the pattern to be matched and the rewrite to be
performed were provided as separate arguments. This commit introduces
a new class `ExprRewriter`, which contains both parts.
This abstraction will make it easier to combine multiple different
rewrite rules, applying them in a single pass.
* lint fixes
* Remove unnecessary change which broke a unit test
* lint fix for import order
* Add docstrings
* lint fix
* Lint fix
* lint fixes
* lint fix
* Update based on review comments
* Add test case for matching against arbitrary dtype
* Fix breakage in unit tests
One unit test that had been relying on invalid shape propagation.
Another unit test that required constructed an ill-formed output to
test against.
* Updated base class name from ExprRewriter to PatternMatchingRewriter
* lint fix
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
This PR introduces a decorator `tvm.ir.base.deprecated`, which emits a
deprecation warning if an outdated API is used, but preserves backward
compatibility by still allowing the API to be used.
For example, currently the preferred way of TIR buffer declaration in
function signature is:
```python
def example(
A: T.Buffer(...), # legacy behavior is `T.Buffer[...]`
): ...
```
With this decorator, if a user writes `T.Buffer[...]`, the parser will
still function properly, but emits a warning that guides the user to
adopt `T.Buffer(...)` if possible.
While there is no breaking change at all in this PR, we believe this
is useful to help users upgrade before any breaking change eventually
takes place.
This PR try to fix the following bug:
```python
def test_var_capturing_order():
b = 2
@T.prim_func
def test_case():
k: T.int32 = b
if __name__ == "__main__":
b = 1
```
In the prim func `test_case`, the vaule of b should be 2, rather than 1. The parser wrongly uses global vars to shadow the value of nonlocal vars, which should be reversed.
Co-authored-by: lightzhan-intellif <zhan.liang@intellif.com>
This PR introduces some minor restructuring of the `python/tvm/script`
folder structure to make it more convenient for future upstreaming.
Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
Tianqi Chen