tilelang.language.v2.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

IRGenerator

Abstract base class for generic types.

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)

mutate(func)

Transform a Python function into an IR (Intermediate Representation) generator.

Module Contents¶

tilelang.language.v2.ast.ast_has_span(ast)¶
Parameters:

ast (ast_has_span.ast)

Return type:

bool

tilelang.language.v2.ast.ast_get_span(ast)¶
Parameters:

ast (ast_get_span.ast)

Return type:

tuple[int, int, int, int]

tilelang.language.v2.ast.ast_set_span(ast, span)¶
Parameters:

  • ast (ast_set_span.ast)

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

class tilelang.language.v2.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)
Parameters:

  • 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.

Parameters:

node (ast.AST)

visit_Name(node)¶
Parameters:

node (ast.Name)

Return type:

Any

visit_Pass(node)¶
Parameters:

node (ast.Pass)

Return type:

Any

tilelang.language.v2.ast.quote(expr, *, passes=None, span=None, **kws)¶
Parameters:

  • expr (str)

  • passes (list[Any] | None)

Return type:

list[ast.AST]

tilelang.language.v2.ast.quote1(expr, *, passes=None, span=None, **kws)¶
Parameters:

  • expr (str)

  • passes (list[Any] | None)

Return type:

ast.AST

tilelang.language.v2.ast.quote_expr(expr, **kws)¶
Parameters:

expr (str)

Return type:

ast.expr

tilelang.language.v2.ast.Operator¶
tilelang.language.v2.ast.BoolOp¶
tilelang.language.v2.ast.get_operator_name(operator)¶
Parameters:

operator (ast.operator)

Return type:

Operator

tilelang.language.v2.ast.get_boolop_name(boolop)¶
Parameters:

boolop (ast.boolop)

Return type:

BoolOp

tilelang.language.v2.ast.eval_op(op, left, right)¶
Parameters:

  • op (Operator)

  • left (Any)

  • right (Any)

Return type:

Any

tilelang.language.v2.ast.eval_aug_assign(op, left, sl, right)¶
Parameters:

  • op (Operator)

  • left (Any)

  • sl (slice)

  • right (Any)

Return type:

Any

class tilelang.language.v2.ast.BaseBuilder¶
empty¶
get_parent_locals()¶
ctx_if(cond)¶
Return type:

collections.abc.Iterable[_T]

ctx_then(val)¶
Parameters:

val (_T)

Return type:

collections.abc.Iterable[None]

ctx_else(val)¶
Parameters:

val (_T)

Return type:

collections.abc.Iterable[None]

eval(val)¶
Parameters:

val (Any)

ctx_for(range)¶
Parameters:

range (collections.abc.Iterable[Any])

Return type:

collections.abc.Iterable[Any]

ctx_continue()¶
Return type:

bool

ctx_break()¶
Return type:

bool

ctx_while(cond)¶
Parameters:

cond (Callable[[], Any])

Return type:

collections.abc.Iterable[None]

bind(name, value, annot=empty)¶
Parameters:

  • name (str)

  • value (Any)

  • annot (Any)

Return type:

Any

unwrap_value(value)¶
assign_slice(lval, sl, value, annot=empty)¶
Parameters:

  • lval (Any)

  • sl (slice)

  • value (Any)

  • annot (Any)

aug_assign(op, target, aug_value)¶
Parameters:

  • op (Operator)

  • target (Any)

  • aug_value (Any)

Return type:

Any

aug_assign_slice(op, target, sl, aug_value)¶
Parameters:

  • op (Operator)

  • target (Any)

  • sl (slice)

  • aug_value (Any)

boolop(op, left, right)¶
Parameters:

  • op (BoolOp)

  • left (Any)

  • right (Callable[[], Any])

Return type:

Any

ifexp(cond, then, otherwise)¶
Parameters:

  • cond (Any)

  • then (Callable[[], Any])

  • otherwise (Callable[[], Any])

Return type:

Any

ret(value)¶
Parameters:

value (Any)

Return type:

Any

ctx_with(ctx)¶
Parameters:

ctx (contextlib.AbstractContextManager[Any])

Return type:

contextlib.AbstractContextManager[Any]

assert_expr(cond, msg)¶
Parameters:

  • cond (Any)

  • msg (Any)

rval(name, value)¶
Parameters:

  • name (str)

  • value (Any)

arg(name, value)¶
Parameters:

  • name (str)

  • value (Any)

override(name)¶
Parameters:

name (str)

class tilelang.language.v2.ast.DSLMutator(closure_names)¶

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)
Parameters:

closure_names (list[str])

tmp_counter = 0¶
closure_names¶
get_tmp()¶
Return type:

str

visit_If(node)¶
Parameters:

node (ast.If)

visit_Expr(node)¶
Parameters:

node (ast.Expr)

visit_For(node)¶
Parameters:

node (ast.For)

visit_Continue(node)¶
Parameters:

node (ast.Continue)

visit_Break(node)¶
Parameters:

node (ast.Break)

visit_Assign(node)¶
Parameters:

node (ast.Assign)

Return type:

list[ast.AST]

visit_AugAssign(node)¶
Parameters:

node (ast.AugAssign)

Return type:

list[ast.AST]

visit_AnnAssign(node)¶
Parameters:

node (ast.AnnAssign)

visit_While(node)¶
visit_FunctionDef(node)¶
Parameters:

node (ast.FunctionDef)

visit_BoolOp(node)¶
Parameters:

node (ast.BoolOp)

visit_Compare(node)¶
Parameters:

node (ast.Compare)

Return type:

ast.expr

visit_IfExp(node)¶
Parameters:

node (ast.IfExp)

Return type:

ast.Expr

visit_Return(node)¶
Parameters:

node (ast.Return)

visit_With(node)¶
Parameters:

node (ast.With)

visit_Assert(node)¶
Parameters:

node (ast.Assert)

visit_Name(node)¶
Parameters:

node (ast.Name)

class tilelang.language.v2.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¶
tilelang.language.v2.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).

Returns:

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

Return type:

IRGenerator[_P, _T]

Parameters:

func (Callable[_P, _T])

Example

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

Note

  • 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