tilelang.language.eager.ast

Attributes

類別

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

Abstract base class for generic types.

函式

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)

基底類別: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)

基底類別: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)

基底類別: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]

Abstract base class for generic types.

A generic type is typically declared by inheriting from this class parameterized with one or more type variables. For example, a generic mapping type might be defined as:

class Mapping(Generic[KT, VT]):
    def __getitem__(self, key: KT) -> VT:
        ...
    # Etc.

This class can then be used as follows:

def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
    try:
        return mapping[key]
    except KeyError:
        return default
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