tilelang.language.eager.ast

Attributes

Operator
BoolOp

類別

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

函式

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