tilelang.language.eager.ast

Attributes

Operator
BoolOp

Classes

QuoteVisitor	A NodeVisitor subclass that walks the abstract syntax tree and
BaseBuilder
DSLMutator	A NodeVisitor subclass that walks the abstract syntax tree and
SpanAttacher	A NodeVisitor subclass that walks the abstract syntax tree and
IRGenerator

Functions

ast_has_span(ast)
ast_get_span(ast)
ast_set_span(ast, span)
quote(expr, *[, passes, span])
quote1(expr, *[, passes, span])
quote_expr(expr, **kws)
get_operator_name(operator)
get_boolop_name(boolop)
eval_op(op, left, right)
eval_aug_assign(op, left, sl, right)
has_internal_prim_func(func)
mutate(func)	Transform a Python function into an IR (Intermediate Representation) generator.

Module Contents

tilelang.language.eager.ast.ast_has_span(ast)

参数:

ast (ast_has_span.ast)

返回类型:

bool

tilelang.language.eager.ast.ast_get_span(ast)

参数:

ast (ast_get_span.ast)

返回类型:

tuple[int, int, int, int]

tilelang.language.eager.ast.ast_set_span(ast, span)

参数:

• ast (ast_set_span.ast)

• span (tuple[int, int, int, int])

class tilelang.language.eager.ast.QuoteVisitor(names, passes=None, span=None)

Bases: ast.NodeTransformer

A NodeVisitor subclass that walks the abstract syntax tree and allows modification of nodes.

The NodeTransformer will walk the AST and use the return value of the visitor methods to replace or remove the old node. If the return value of the visitor method is None, the node will be removed from its location, otherwise it is replaced with the return value. The return value may be the original node in which case no replacement takes place.

Here is an example transformer that rewrites all occurrences of name lookups (foo) to data['foo']:

class RewriteName(NodeTransformer):

    def visit_Name(self, node):
        return Subscript(
            value=Name(id='data', ctx=Load()),
            slice=Constant(value=node.id),
            ctx=node.ctx
        )

Keep in mind that if the node you're operating on has child nodes you must either transform the child nodes yourself or call the generic_visit() method for the node first.

For nodes that were part of a collection of statements (that applies to all statement nodes), the visitor may also return a list of nodes rather than just a single node.

Usually you use the transformer like this:

node = YourTransformer().visit(node)

参数:

• names (dict[str, ast.AST])

• passes (list[Any] | None)

names

passes = []

span = None

generic_visit(node)

Called if no explicit visitor function exists for a node.

参数:

node (ast.AST)

visit_Name(node)

参数:

node (ast.Name)

返回类型:

Any

visit_Pass(node)

参数:

node (ast.Pass)

返回类型:

Any

tilelang.language.eager.ast.quote(expr, *, passes=None, span=None, **kws)

参数:

• expr (str)

• passes (list[Any] | None)

返回类型:

list[ast.AST]

tilelang.language.eager.ast.quote1(expr, *, passes=None, span=None, **kws)

参数:

• expr (str)

• passes (list[Any] | None)

返回类型:

ast.AST

tilelang.language.eager.ast.quote_expr(expr, **kws)

参数:

expr (str)

返回类型:

ast.expr

tilelang.language.eager.ast.Operator

tilelang.language.eager.ast.BoolOp

tilelang.language.eager.ast.get_operator_name(operator)

参数:

operator (ast.operator)

返回类型:

Operator

tilelang.language.eager.ast.get_boolop_name(boolop)

参数:

boolop (ast.boolop)

返回类型:

BoolOp

tilelang.language.eager.ast.eval_op(op, left, right)

参数:

• op (Operator)

• left (Any)

• right (Any)

返回类型:

Any

tilelang.language.eager.ast.eval_aug_assign(op, left, sl, right)

参数:

• op (Operator)

• left (Any)

• sl (slice)

• right (Any)

返回类型:

Any

class tilelang.language.eager.ast.BaseBuilder

empty

get_parent_locals()

ctx_if(cond)

返回类型:

collections.abc.Iterable[_T]

ctx_then(val)

参数:

val (_T)

返回类型:

collections.abc.Iterable[None]

ctx_else(val)

参数:

val (_T)

返回类型:

collections.abc.Iterable[None]

eval(val)

参数:

val (Any)

ctx_for(range)

参数:

range (collections.abc.Iterable[Any])

返回类型:

collections.abc.Iterable[Any]

ctx_continue()

返回类型:

bool

ctx_break()

返回类型:

bool

ctx_while(cond)

参数:

cond (Callable[[], Any])

返回类型:

collections.abc.Iterable[None]

bind(name, value, annot=empty)

参数:

• name (str)

• value (Any)

• annot (Any)

返回类型:

Any

unwrap_value(value)

assign_slice(lval, sl, value, annot=empty)

参数:

• lval (Any)

• sl (slice)

• value (Any)

• annot (Any)

aug_assign(op, target, aug_value, name=None)

参数:

• op (Operator)

• target (Any)

• aug_value (Any)

• name (str | None)

返回类型:

Any

aug_assign_slice(op, target, sl, aug_value)

参数:

• op (Operator)

• target (Any)

• sl (slice)

• aug_value (Any)

boolop(op, left, right=None)

参数:

• op (BoolOp)

• left (Any)

• right (Callable[[], Any] | None)

返回类型:

Any

ifexp(cond, then, otherwise)

参数:

• cond (Any)

• then (Callable[[], Any])

• otherwise (Callable[[], Any])

返回类型:

Any

ret(value)

参数:

value (Any)

返回类型:

Any

ctx_with(ctx)

参数:

ctx (contextlib.AbstractContextManager[Any])

返回类型:

contextlib.AbstractContextManager[Any]

assert_expr(cond, msg)

参数:

• cond (Any)

• msg (Any)

rval(name, value)

参数:

• name (str)

• value (Any)

arg(name, value)

参数:

• name (str)

• value (Any)

override(name)

参数:

name (str)

class tilelang.language.eager.ast.DSLMutator(nonlocals, globals, filename)

Bases: ast.NodeTransformer

A NodeVisitor subclass that walks the abstract syntax tree and allows modification of nodes.

The NodeTransformer will walk the AST and use the return value of the visitor methods to replace or remove the old node. If the return value of the visitor method is None, the node will be removed from its location, otherwise it is replaced with the return value. The return value may be the original node in which case no replacement takes place.

Here is an example transformer that rewrites all occurrences of name lookups (foo) to data['foo']:

class RewriteName(NodeTransformer):

    def visit_Name(self, node):
        return Subscript(
            value=Name(id='data', ctx=Load()),
            slice=Constant(value=node.id),
            ctx=node.ctx
        )

Keep in mind that if the node you're operating on has child nodes you must either transform the child nodes yourself or call the generic_visit() method for the node first.

For nodes that were part of a collection of statements (that applies to all statement nodes), the visitor may also return a list of nodes rather than just a single node.

Usually you use the transformer like this:

node = YourTransformer().visit(node)

参数:

• nonlocals (dict[str, Any])

• globals (dict[str, Any])

• filename (str)

tmp_counter = 0

nonlocals

globals

extra_type_hints: dict[str, Any]

filename

get_tmp()

返回类型:

str

visit_If(node)

参数:

node (ast.If)

visit_Expr(node)

参数:

node (ast.Expr)

visit_For(node)

参数:

node (ast.For)

visit_Continue(node)

参数:

node (ast.Continue)

visit_Break(node)

参数:

node (ast.Break)

visit_Assign(node)

参数:

node (ast.Assign)

返回类型:

list[ast.AST]

visit_AugAssign(node)

参数:

node (ast.AugAssign)

返回类型:

list[ast.AST]

visit_AnnAssign(node)

参数:

node (ast.AnnAssign)

visit_While(node)

visit_FunctionDef(node)

参数:

node (ast.FunctionDef)

visit_BoolOp(node)

参数:

node (ast.BoolOp)

visit_UnaryOp(node)

参数:

node (ast.UnaryOp)

visit_Compare(node)

参数:

node (ast.Compare)

返回类型:

ast.expr

visit_IfExp(node)

参数:

node (ast.IfExp)

返回类型:

ast.Expr

visit_Return(node)

参数:

node (ast.Return)

visit_With(node)

参数:

node (ast.With)

visit_Assert(node)

参数:

node (ast.Assert)

visit_Name(node)

参数:

node (ast.Name)

class tilelang.language.eager.ast.SpanAttacher(filename_var, func_name_var)

Bases: ast.NodeTransformer

A NodeVisitor subclass that walks the abstract syntax tree and allows modification of nodes.

The NodeTransformer will walk the AST and use the return value of the visitor methods to replace or remove the old node. If the return value of the visitor method is None, the node will be removed from its location, otherwise it is replaced with the return value. The return value may be the original node in which case no replacement takes place.

Here is an example transformer that rewrites all occurrences of name lookups (foo) to data['foo']:

class RewriteName(NodeTransformer):

    def visit_Name(self, node):
        return Subscript(
            value=Name(id='data', ctx=Load()),
            slice=Constant(value=node.id),
            ctx=node.ctx
        )

Keep in mind that if the node you're operating on has child nodes you must either transform the child nodes yourself or call the generic_visit() method for the node first.

For nodes that were part of a collection of statements (that applies to all statement nodes), the visitor may also return a list of nodes rather than just a single node.

Usually you use the transformer like this:

node = YourTransformer().visit(node)

参数:

• filename_var (str)

• func_name_var (str)

filename_var

func_name_var

visit(node)

Visit a node.

参数:

node (ast.AST)

class tilelang.language.eager.ast.IRGenerator

Bases: Generic[_P, _T]

gen: Callable[[BaseBuilder], Callable[_P, _T]]

source: str

extra_type_hints: dict[str, Any]

tilelang.language.eager.ast.has_internal_prim_func(func)

参数:

func (Callable[_P, _T])

返回类型:

bool

tilelang.language.eager.ast.mutate(func)

Transform a Python function into an IR (Intermediate Representation) generator. This function takes a regular Python function and performs AST (Abstract Syntax Tree) transformation to create an IRGenerator that can be used for code generation purposes. :param func: The Python function to be transformed. This should be a

callable that will be analyzed and mutated at the AST level. The function's signature is preserved through generic type parameters _P (parameters) and _T (return type).

返回:

An IRGenerator instance wrapping the transformed function.

The generator contains: - gen: The compiled and mutated version of the original function - source: The unparsed source code of the transformed AST as a string

返回类型:

IRGenerator[_P, _T]

参数:

func (Callable[_P, _T])

示例

>>> @mutate
... def my_function(x: int) -> int:
...     return x * 2
>>> # my_function is now an IRGenerator that can be used for code generation

备注

• The original function's closure variables and captured context are preserved

• The transformation is performed at compile-time through AST manipulation

• The returned IRGenerator maintains type information from the original function