Don't rely on co_varnames to get the list of locals

With Python 3.12, a list comprehension like [i*i for i in range(3)]
will add `i` to co_varnames, even though `i` is still not a local
variable.

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

src/cuda/tile/_passes/ast2hir.py

@@ -297,8 +297,9 @@ def _call_static_eval(expr: ast.expr, kind: hir.StaticEvalKind, ctx: _Context) -
     inner_lambda = outer_lambda(*tuple(None for _ in local_names))
 
     # Make sure the function doesn't store any locals, e.g. using the walrus operator
-    if len(inner_lambda.__code__.co_varnames) > 0:
-        name = inner_lambda.__code__.co_varnames[0]
+    stored_locals = ast_get_all_local_names(inner_lambda_ast).local_names
+    if len(stored_locals) > 0:
+        name = min(stored_locals)
         raise TileSyntaxError(f"static_eval() expression attempted"
                               f" to modify a local variable '{name}'")
 

test/test_ast_util.py

@@ -53,12 +53,14 @@ async def nested_async_def(): pass
     print(a, b, c, sneaky_1, sneaky_2, sneaky_3, sneaky_4, local_i)
 
 
-
 def test_get_all_local_names():
     parsed_ast = ast.parse(inspect.getsource(func))
     func_ast = parsed_ast.body[0]
     local_names, global_names, nonlocal_names = ast_get_all_local_names(func_ast)
-    expected = sorted(func.__code__.co_varnames + func.__code__.co_cellvars)
+    expected = ['a', 'b', 'c', 'cd', 'e', 'getcwd', 'local_i', 'match_alias', 'nested_async_def',
+                'nested_func', 'nonexistent', 'p1', 'p2', 'p3',
+                'pat1', 'pat2', 'pat3', 'pat4', 'pat5', 'pat6',
+                'sneaky_1', 'sneaky_2', 'sneaky_3', 'sneaky_4', 'walrus']
     assert sorted(local_names) == expected
     assert global_names == {"g"}
     assert nonlocal_names == set()

test/test_static_eval.py

@@ -25,6 +25,20 @@ def kernel(y):
     assert y.tolist() == [6, 6]
 
 
+def test_list_comprehension():
+    @ct.kernel
+    def kernel(y):
+        tup = (1, 2, 3)
+        s1 = ct.static_eval(sum([i*i for i in tup]))
+        s2 = cuda.tile.static_eval(sum([i*i for i in tup]))
+        ct.scatter(y, 0, s1)
+        ct.scatter(y, 1, s2)
+
+    y = torch.zeros((2,), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (y,))
+    assert y.tolist() == [1*1 + 2*2 + 3*3, 1*1 + 2*2 + 3*3]
+
+
 def test_mixed_tuple():
     @ct.kernel
     def kernel(y):