gemini-code-assist[bot] commented on code in PR #18856: URL: https://github.com/apache/tvm/pull/18856#discussion_r2867986588
########## python/tvm/script/parser/ir/entry.py: ########## @@ -18,6 +18,7 @@ """The entry point of TVM parser for ir module.""" import inspect +from collections import ChainMap from collections.abc import Callable Review Comment:  Per PEP 8, standard library imports should be sorted alphabetically. `collections` should come before `inspect`. ```suggestion from collections import ChainMap from collections.abc import Callable import inspect ``` <details> <summary>References</summary> 1. PEP 8 suggests sorting standard library imports alphabetically. <sup>([link](http://console.cloud.google.com/gemini-code-assist/agents-tools))</sup> </details> ########## tests/python/tvmscript/test_tvmscript_pep563_closure.py: ########## @@ -0,0 +1,161 @@ +# 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. +"""Test TVMScript with PEP 563 (from __future__ import annotations). + +IMPORTANT: The `from __future__ import annotations` import below is the +test condition itself, because we need to test compatibility with it. +""" +from __future__ import annotations + +import tvm +import tvm.testing +from tvm.script import ir as I +from tvm.script import tir as T + + +def _normalize(func): + """Strip the global_symbol so function names do not affect structural equality.""" + return func.with_attr("global_symbol", "") + + +def test_prim_func_closure_shape(): + """Closure variable used in Buffer shape annotation.""" + + def f(M=16): + @T.prim_func + def func(A: T.Buffer((M,), "float32")): + T.evaluate(0) + + return func + + @T.prim_func + def expected_16(A: T.Buffer((16,), "float32")): + T.evaluate(0) + + @T.prim_func + def expected_32(A: T.Buffer((32,), "float32")): + T.evaluate(0) + + tvm.ir.assert_structural_equal(_normalize(f(16)), _normalize(expected_16)) + tvm.ir.assert_structural_equal(_normalize(f(32)), _normalize(expected_32)) + + +def test_prim_func_closure_dtype(): + """Closure variable used as Buffer dtype.""" + + def f(dtype="float32"): + @T.prim_func + def func(A: T.Buffer((16,), dtype)): + T.evaluate(0) + + return func + + @T.prim_func + def expected_f32(A: T.Buffer((16,), "float32")): + T.evaluate(0) + + @T.prim_func + def expected_f16(A: T.Buffer((16,), "float16")): + T.evaluate(0) + + tvm.ir.assert_structural_equal(_normalize(f("float32")), _normalize(expected_f32)) + tvm.ir.assert_structural_equal(_normalize(f("float16")), _normalize(expected_f16)) + + +def test_prim_func_nested_closure(): + """Variables from enclosing scope active on the call stack (grandparent frame fallback). + + With PEP 563, closure-only variables are missing from __closure__ unless they + appear in the function body. The ChainMap fallback walks the live call stack, + so this works when the enclosing frames are still active (outer calls middle + which applies the decorator, keeping outer's frame alive on the stack). + """ + + def outer(M=16): + def middle(N=8): + @T.prim_func + def func(A: T.Buffer((M, N), "float32")): + T.evaluate(0) + + return func + + return middle() + + @T.prim_func + def expected_16_8(A: T.Buffer((16, 8), "float32")): + T.evaluate(0) + + @T.prim_func + def expected_32_8(A: T.Buffer((32, 8), "float32")): + T.evaluate(0) + + tvm.ir.assert_structural_equal(_normalize(outer(16)), _normalize(expected_16_8)) + tvm.ir.assert_structural_equal(_normalize(outer(32)), _normalize(expected_32_8)) + + +def test_ir_module_closure(): + """Closure variable in @I.ir_module class method.""" + + def f(M=16): + @I.ir_module + class Mod: + @T.prim_func + def main(A: T.Buffer((M,), "float32")): + T.evaluate(0) + + return Mod + + @T.prim_func + def expected_16(A: T.Buffer((16,), "float32")): + T.evaluate(0) + + @T.prim_func + def expected_32(A: T.Buffer((32,), "float32")): + T.evaluate(0) + + tvm.ir.assert_structural_equal( + f(16)["main"], expected_16.with_attr("global_symbol", "main") + ) + tvm.ir.assert_structural_equal( + f(32)["main"], expected_32.with_attr("global_symbol", "main") + ) Review Comment:  For consistency with the other tests in this file, consider using the `_normalize` helper function to strip the `global_symbol` attribute before comparison. This makes the test's intent clearer and aligns with the pattern used in other test cases here. ```python tvm.ir.assert_structural_equal(_normalize(f(16)["main"]), _normalize(expected_16)) tvm.ir.assert_structural_equal(_normalize(f(32)["main"]), _normalize(expected_32)) ``` ########## python/tvm/script/parser/tir/entry.py: ########## @@ -17,6 +17,7 @@ """The entry point of TVM parser for tir.""" import inspect +from collections import ChainMap from collections.abc import Callable Review Comment:  Per PEP 8, standard library imports should be sorted alphabetically. `collections` should come before `inspect`. ```suggestion from collections import ChainMap from collections.abc import Callable import inspect ``` <details> <summary>References</summary> 1. PEP 8 suggests sorting standard library imports alphabetically. <sup>([link](http://console.cloud.google.com/gemini-code-assist/agents-tools))</sup> </details> ########## python/tvm/script/parser/ir/entry.py: ########## @@ -46,14 +47,24 @@ def ir_module(mod: type | None = None, check_well_formed: bool = True) -> IRModu The parsed ir module. """ + # need to capture this var outside the wrapper because the wrapper + # adds to the stack + outer_stack = inspect.stack() + def decorator_wrapper(mod): if not inspect.isclass(mod): raise TypeError(f"Expect a class, but got: {mod}") # Check BasePyModule inheritance base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__) - m = parse(mod, utils.inspect_class_capture(mod), check_well_formed=check_well_formed) + extra_vars = utils.inspect_class_capture(mod) + # Lazy fallback to caller frame locals for PEP 563 compatibility. + # With `from __future__ import annotations`, variables used only + # in annotations are not captured in __closure__. ChainMap defers + # lookup to caller locals only on cache miss. + extra_vars = ChainMap(extra_vars, *[f.frame.f_locals for f in outer_stack[1:]]) Review Comment:  This logic for creating a `ChainMap` with caller frame locals is duplicated in `python/tvm/script/parser/tir/entry.py`. To improve maintainability, consider extracting this logic into a common helper function in a shared utility module (e.g., `tvm.script.parser.utils`). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
