From b61b8d27601990f054fb77b5d49d30f7cbe50986 Mon Sep 17 00:00:00 2001 From: Qiping Pan Date: Sat, 4 Jul 2026 15:00:48 -0700 Subject: [PATCH 1/5] fix: pyright --- tilelang/language/eager/ast.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tilelang/language/eager/ast.py b/tilelang/language/eager/ast.py index fd98f8c73b..b7c2c0d4ae 100644 --- a/tilelang/language/eager/ast.py +++ b/tilelang/language/eager/ast.py @@ -2,8 +2,8 @@ import ast from dataclasses import dataclass, field from pathlib import Path -from typing import Generic, Any, Literal, ParamSpec, TypeVar -from collections.abc import Callable +from typing import Generic, Any, Literal, ParamSpec, TypeVar, cast +from collections.abc import Callable, Sequence from contextlib import AbstractContextManager from collections.abc import Iterable @@ -307,7 +307,7 @@ def visit_For(self, node: ast.For): f"for {tmp} in __tb.ctx_for(range):\n pass\n", target=node.target, range=node.iter, - passes=[stmts + node.body], + passes=[*stmts, *node.body], span=node, ) @@ -319,7 +319,7 @@ def visit_Break(self, node: ast.Break): node = self.generic_visit(node) return quote("if __tb.ctx_break(): break", span=node) - def _emit_assign_target(self, target: ast.expr, rval: ast.expr, annot: ast.expr = None) -> list[ast.AST]: + def _emit_assign_target(self, target: ast.expr, rval: ast.expr, annot: ast.expr = None) -> Sequence[ast.AST]: if isinstance(target, ast.Name): if annot is None: return quote(f"name = __tb.bind('{target.id}', value)", name=target, value=rval, span=target) @@ -350,7 +350,7 @@ def _emit_assign_target(self, target: ast.expr, rval: ast.expr, annot: ast.expr # flatten nested tuple into a list of (tmp_name, target) unpacked = [] - def _visit_target(target: ast.expr) -> str: + def _visit_target(target: ast.expr) -> ast.expr: if isinstance(target, (ast.Name, ast.Subscript)): tmp = self.get_tmp() unpacked.append((tmp, target)) @@ -374,7 +374,7 @@ def _visit_target(target: ast.expr) -> str: def flush_binds(): if bind_lvals: - stmts.append(quote1(f"{', '.join(bind_lvals)}, = {', '.join(bind_rvals)},", span=target)) + stmts.append(cast(ast.Assign, quote1(f"{', '.join(bind_lvals)}, = {', '.join(bind_rvals)},", span=target))) bind_lvals.clear() bind_rvals.clear() @@ -415,14 +415,14 @@ def visit_Assign(self, node: ast.Assign) -> list[ast.AST]: node = self.generic_visit(node) rval = node.value if len(node.targets) == 1: - return self._emit_assign_target(node.targets[0], rval) + return list(self._emit_assign_target(node.targets[0], rval)) else: tmp_name = self.get_tmp() tmp_store = ast.Name(tmp_name, ctx=ast.Store()) tmp_load = ast.Name(tmp_name, ctx=ast.Load()) - ast_set_span(tmp_store, node.targets[0]) - ast_set_span(tmp_load, node.targets[0]) - stmt = self._emit_assign_target(tmp_store, rval) + ast_set_span(tmp_store, ast_get_span(node.targets[0])) + ast_set_span(tmp_load, ast_get_span(node.targets[0])) + stmt = list(self._emit_assign_target(tmp_store, rval)) for target in node.targets: stmt.extend(self._emit_assign_target(target, tmp_load)) return stmt From 6c6f07c4b9a4fd7af6779aca13e8a9f2f26ab03f Mon Sep 17 00:00:00 2001 From: Qiping Pan Date: Sat, 4 Jul 2026 15:06:39 -0700 Subject: [PATCH 2/5] fix: pyright wave 2 --- tilelang/language/eager/ast.py | 49 +++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/tilelang/language/eager/ast.py b/tilelang/language/eager/ast.py index b7c2c0d4ae..1ea56c4a6b 100644 --- a/tilelang/language/eager/ast.py +++ b/tilelang/language/eager/ast.py @@ -2,8 +2,8 @@ import ast from dataclasses import dataclass, field from pathlib import Path -from typing import Generic, Any, Literal, ParamSpec, TypeVar, cast -from collections.abc import Callable, Sequence +from typing import Generic, Any, Literal, ParamSpec, TypeVar, TypeAlias, cast +from collections.abc import Callable from contextlib import AbstractContextManager from collections.abc import Iterable @@ -14,32 +14,33 @@ from .. import dtypes _span_attrs = ["lineno", "col_offset", "end_lineno", "end_col_offset"] +Span: TypeAlias = tuple[int, int, int, int] def ast_has_span(ast: ast.AST) -> bool: return all(hasattr(ast, attr) for attr in _span_attrs) -def ast_get_span(ast: ast.AST) -> tuple[int, int, int, int]: +def ast_get_span(ast: ast.AST) -> Span | None: if not ast_has_span(ast): return None - return tuple(getattr(ast, attr) for attr in _span_attrs) + return cast(Span, tuple(getattr(ast, attr) for attr in _span_attrs)) -def ast_set_span(ast: ast.AST, span: tuple[int, int, int, int]): - if not ast_has_span(ast): +def ast_set_span(ast: ast.AST, span: Span | None): + if span is None or not ast_has_span(ast): return for attr, value in zip(_span_attrs, span): setattr(ast, attr, value) class QuoteVisitor(ast.NodeTransformer): - def __init__(self, names: dict[str, ast.AST], passes: list[Any] | None = None, span=None): + def __init__(self, names: dict[str, Any], passes: list[Any] | None = None, span: Span | None = None): self.names = names self.passes = passes or [] self.span = span - def generic_visit(self, node: ast.AST): + def generic_visit(self, node: ast.AST) -> Any: if self.span is not None: ast_set_span(node, self.span) return super().generic_visit(node) @@ -55,15 +56,15 @@ def visit_Pass(self, node: ast.Pass) -> Any: return item if item else node -def quote(expr: str, *, passes: list[Any] | None = None, span=None, **kws) -> list[ast.AST]: +def quote(expr: str, *, passes: list[Any] | None = None, span: ast.AST | Span | None = None, **kws: Any) -> list[ast.stmt]: tree = ast.parse(expr) if isinstance(span, ast.AST): span = ast_get_span(span) - tree = QuoteVisitor(kws, passes, span).visit(tree) + tree = cast(ast.Module, QuoteVisitor(kws, passes, span).visit(tree)) return tree.body -def quote1(expr: str, *, passes: list[Any] | None = None, span=None, **kws) -> ast.AST: +def quote1(expr: str, *, passes: list[Any] | None = None, span: ast.AST | Span | None = None, **kws: Any) -> ast.stmt: res = quote(expr, passes=passes, span=span, **kws) assert len(res) == 1 return res[0] @@ -80,11 +81,11 @@ def quote_expr(expr: str, **kws) -> ast.expr: def get_operator_name(operator: ast.operator) -> Operator: - return operator.__class__.__name__ + return cast(Operator, operator.__class__.__name__) def get_boolop_name(boolop: ast.boolop) -> BoolOp: - return boolop.__class__.__name__ + return cast(BoolOp, boolop.__class__.__name__) _T = TypeVar("_T") @@ -170,7 +171,9 @@ class BaseBuilder: empty = _empty def get_parent_locals(self): - return inspect.currentframe().f_back.f_back.f_locals + frame = inspect.currentframe() + assert frame is not None and frame.f_back is not None and frame.f_back.f_back is not None + return frame.f_back.f_back.f_locals def ctx_if(self, cond) -> Iterable[_T]: yield cond @@ -216,8 +219,10 @@ def aug_assign_slice(self, op: Operator, target: Any, sl: slice, aug_value: Any) def boolop(self, op: BoolOp, left: Any, right: Callable[[], Any] | None = None) -> Any: if op == "And": + assert right is not None return left and right() if op == "Or": + assert right is not None return left or right() if op == "Not": return not left @@ -319,7 +324,7 @@ def visit_Break(self, node: ast.Break): node = self.generic_visit(node) return quote("if __tb.ctx_break(): break", span=node) - def _emit_assign_target(self, target: ast.expr, rval: ast.expr, annot: ast.expr = None) -> Sequence[ast.AST]: + def _emit_assign_target(self, target: ast.expr, rval: ast.expr, annot: ast.expr | None = None) -> list[ast.stmt]: if isinstance(target, ast.Name): if annot is None: return quote(f"name = __tb.bind('{target.id}', value)", name=target, value=rval, span=target) @@ -348,7 +353,7 @@ def _emit_assign_target(self, target: ast.expr, rval: ast.expr, annot: ast.expr ) else: # flatten nested tuple into a list of (tmp_name, target) - unpacked = [] + unpacked: list[tuple[str, ast.expr]] = [] def _visit_target(target: ast.expr) -> ast.expr: if isinstance(target, (ast.Name, ast.Subscript)): @@ -368,9 +373,9 @@ def _visit_target(target: ast.expr) -> ast.expr: unpack_stmt = ast.Assign(targets=[_visit_target(target)], value=quote_expr("__tb.unwrap_value(rval)", rval=rval, span=rval)) ast_set_span(unpack_stmt, ast_get_span(target)) - stmts = [unpack_stmt] - bind_lvals = [] - bind_rvals = [] + stmts: list[ast.stmt] = [unpack_stmt] + bind_lvals: list[str] = [] + bind_rvals: list[str] = [] def flush_binds(): if bind_lvals: @@ -411,7 +416,7 @@ def flush_binds(): flush_binds() return stmts - def visit_Assign(self, node: ast.Assign) -> list[ast.AST]: + def visit_Assign(self, node: ast.Assign) -> list[ast.stmt]: node = self.generic_visit(node) rval = node.value if len(node.targets) == 1: @@ -427,7 +432,7 @@ def visit_Assign(self, node: ast.Assign) -> list[ast.AST]: stmt.extend(self._emit_assign_target(target, tmp_load)) return stmt - def visit_AugAssign(self, node: ast.AugAssign) -> list[ast.AST]: + def visit_AugAssign(self, node: ast.AugAssign) -> ast.AST | list[ast.stmt]: node = self.generic_visit(node) target, rval = node.target, node.value op = get_operator_name(node.op) @@ -573,7 +578,7 @@ def visit_Compare(self, node: ast.Compare) -> ast.expr: last = quote_expr("__tb.boolop('And', left, lambda: right)", left=split[i], right=last, span=node) return last - def visit_IfExp(self, node: ast.IfExp) -> ast.Expr: + def visit_IfExp(self, node: ast.IfExp) -> ast.expr: node = self.generic_visit(node) return quote_expr( "__tb.ifexp(cond, lambda: then, lambda: otherwise)", cond=node.test, then=node.body, otherwise=node.orelse, span=node From 8342379b97a6d49c0ec7a227cd6195cd5ab3a95d Mon Sep 17 00:00:00 2001 From: Qiping Pan Date: Sat, 4 Jul 2026 15:21:03 -0700 Subject: [PATCH 3/5] QuoteReplacement type alias --- tilelang/language/eager/ast.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/tilelang/language/eager/ast.py b/tilelang/language/eager/ast.py index 1ea56c4a6b..db970d6854 100644 --- a/tilelang/language/eager/ast.py +++ b/tilelang/language/eager/ast.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Generic, Any, Literal, ParamSpec, TypeVar, TypeAlias, cast -from collections.abc import Callable +from collections.abc import Callable, Sequence from contextlib import AbstractContextManager from collections.abc import Iterable @@ -15,6 +15,7 @@ _span_attrs = ["lineno", "col_offset", "end_lineno", "end_col_offset"] Span: TypeAlias = tuple[int, int, int, int] +QuoteReplacement: TypeAlias = ast.AST | list[ast.stmt] | None def ast_has_span(ast: ast.AST) -> bool: @@ -35,28 +36,30 @@ def ast_set_span(ast: ast.AST, span: Span | None): class QuoteVisitor(ast.NodeTransformer): - def __init__(self, names: dict[str, Any], passes: list[Any] | None = None, span: Span | None = None): + def __init__(self, names: dict[str, QuoteReplacement], passes: Sequence[QuoteReplacement] | None = None, span: Span | None = None): self.names = names - self.passes = passes or [] + self.passes = list(passes) if passes is not None else [] self.span = span - def generic_visit(self, node: ast.AST) -> Any: + def generic_visit(self, node: ast.AST) -> ast.AST | list[ast.AST] | None: if self.span is not None: ast_set_span(node, self.span) return super().generic_visit(node) - def visit_Name(self, node: ast.Name) -> Any: + def visit_Name(self, node: ast.Name) -> QuoteReplacement: if node.id in self.names: return self.names[node.id] else: return node - def visit_Pass(self, node: ast.Pass) -> Any: + def visit_Pass(self, node: ast.Pass) -> QuoteReplacement: item = self.passes.pop(0) return item if item else node -def quote(expr: str, *, passes: list[Any] | None = None, span: ast.AST | Span | None = None, **kws: Any) -> list[ast.stmt]: +def quote( + expr: str, *, passes: Sequence[QuoteReplacement] | None = None, span: ast.AST | Span | None = None, **kws: QuoteReplacement +) -> list[ast.stmt]: tree = ast.parse(expr) if isinstance(span, ast.AST): span = ast_get_span(span) @@ -64,13 +67,15 @@ def quote(expr: str, *, passes: list[Any] | None = None, span: ast.AST | Span | return tree.body -def quote1(expr: str, *, passes: list[Any] | None = None, span: ast.AST | Span | None = None, **kws: Any) -> ast.stmt: +def quote1( + expr: str, *, passes: Sequence[QuoteReplacement] | None = None, span: ast.AST | Span | None = None, **kws: QuoteReplacement +) -> ast.stmt: res = quote(expr, passes=passes, span=span, **kws) assert len(res) == 1 return res[0] -def quote_expr(expr: str, **kws) -> ast.expr: +def quote_expr(expr: str, **kws: QuoteReplacement) -> ast.expr: res = quote1(expr, **kws) assert isinstance(res, ast.Expr) return res.value From e81084610778bf8e8851611a3e1edf9736d3d14d Mon Sep 17 00:00:00 2001 From: Qiping Pan Date: Sat, 4 Jul 2026 16:24:05 -0700 Subject: [PATCH 4/5] simplify redundant branches --- tilelang/language/eager/ast.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tilelang/language/eager/ast.py b/tilelang/language/eager/ast.py index db970d6854..538bea47be 100644 --- a/tilelang/language/eager/ast.py +++ b/tilelang/language/eager/ast.py @@ -486,10 +486,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef): for arg in all_args: name = arg.arg arg_names.add(name) - if arg.annotation is not None: - arg_stmt = quote1(f'{name} = __tb.arg("{name}", {name})', span=arg) - else: - arg_stmt = quote1(f'{name} = __tb.arg("{name}", {name})', span=arg) + arg_stmt = quote1(f'{name} = __tb.arg("{name}", {name})', span=arg) arg.annotation = None stmts.append(arg_stmt) # trying to find `A: T.Tensor, b: T.float32` like type hints From 4807729ed9bd74c6e7ee2f41db2549698b268f3c Mon Sep 17 00:00:00 2001 From: Qiping Pan Date: Sat, 4 Jul 2026 17:08:17 -0700 Subject: [PATCH 5/5] early return on kernel call check --- tilelang/language/eager/ast.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tilelang/language/eager/ast.py b/tilelang/language/eager/ast.py index 538bea47be..c9c7c0d295 100644 --- a/tilelang/language/eager/ast.py +++ b/tilelang/language/eager/ast.py @@ -600,6 +600,7 @@ def visit_With(self, node: ast.With): if eval_res is Kernel: is_kernel_ctx = True + break node = self.generic_visit(node) for expr in node.items: expr.context_expr = quote_expr("__tb.ctx_with(e)", e=expr.context_expr, span=expr)