tilelang.autodd

Attributes

ASTPatKind
JobBackend

Classes

ASTRewrite	Helper class that provides a standard way to create an ABC using
GeneralRemove	Helper class that provides a standard way to create an ABC using
CallFwdArg1	Helper class that provides a standard way to create an ABC using
AttachFullFuncArgs	Helper class that provides a standard way to create an ABC using
IntConstApply	Helper class that provides a standard way to create an ABC using
BinOpFwdArg	Helper class that provides a standard way to create an ABC using
ASTPat
ASTPatRewrite	Helper class that provides a standard way to create an ABC using
ASTMutator
LabeledRewrite
RewriteAttacher
RewriteApplier
Task
PDD
TaskManager	Helper class that provides a standard way to create an ABC using
ASTPDD	Helper class that provides a standard way to create an ABC using
LinePDD	Helper class that provides a standard way to create an ABC using
Ruff	Helper class that provides a standard way to create an ABC using
AsyncPythonRunner
SubProcRunner
ParTaskManager
Args

Functions

ast_replace(node, **changes)
parse_stmts(s)
parse_expr(s)
expr_to_zeros(target)
attach_rewrites(tree, rewrites)
apply_rewrites(tree, target_labels)
test_rewrite(rewrite, code)
ruff_fix_code(code_string[, fix_lint, format_code])
clean_empty_pass(code)
main(args)
cli_main([argv])

Module Contents

tilelang.autodd.ast_replace(node, **changes)

Parameters:

node (ast.AST)

Return type:

ast.AST

tilelang.autodd.parse_stmts(s)

Parameters:

s (str)

Return type:

list[ast.stmt]

tilelang.autodd.parse_expr(s)

Parameters:

s (str)

Return type:

ast.expr

class tilelang.autodd.ASTRewrite

Bases: abc.ABC

Helper class that provides a standard way to create an ABC using inheritance.

abstract get_name()

Return type:

str

abstract match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

abstract rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

ast.AST | list[ast.AST] | None

class tilelang.autodd.GeneralRemove

Bases: ASTRewrite

Helper class that provides a standard way to create an ABC using inheritance.

name: str

target_type: type[ast.AST]

inside_list: bool = True

replace_with: ast.AST | list[ast.AST] | None = None

get_name()

Return type:

str

match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

None

tilelang.autodd.expr_to_zeros(target)

Parameters:

target (ast.expr)

Return type:

ast.expr

class tilelang.autodd.CallFwdArg1

Bases: ASTRewrite

Helper class that provides a standard way to create an ABC using inheritance.

get_name()

Return type:

str

match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

ast.AST

class tilelang.autodd.AttachFullFuncArgs

Bases: ASTRewrite

Helper class that provides a standard way to create an ABC using inheritance.

get_name()

Return type:

str

match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

ast.AST

class tilelang.autodd.IntConstApply

Bases: ASTRewrite

Helper class that provides a standard way to create an ABC using inheritance.

matcher: Callable[[int], bool]

apply: Callable[[int], ast.AST]

name: str

get_name()

Return type:

str

match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

ast.AST

class tilelang.autodd.BinOpFwdArg

Bases: ASTRewrite

Helper class that provides a standard way to create an ABC using inheritance.

forward: Literal['left', 'right'] = 'left'

get_name()

Return type:

str

match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

ast.AST

tilelang.autodd.ASTPatKind

class tilelang.autodd.ASTPat

tree: ast.expr | list[ast.stmt]

placeholders: set[str]

classmethod from_code(kind, code, placeholders)

Parameters:

• kind (ASTPatKind)

• code (str)

• placeholders (set[str])

Return type:

ASTPat

match_placeholders(node)

Parameters:

node (ast.AST | list[ast.AST])

Return type:

dict[str, ast.AST] | bool

match(node)

Parameters:

node (ast.AST)

Return type:

bool

replace(repl)

Parameters:

repl (dict[str, ast.AST])

Return type:

ast.AST

class tilelang.autodd.ASTPatRewrite

Bases: ASTRewrite

Helper class that provides a standard way to create an ABC using inheritance.

name: str

match_pat: ASTPat

rewrite_pat: ASTPat

checker: Callable[[dict[str, ast.AST]], bool] | dict[str, Callable[[ast.AST], bool]] | None = None

derived: dict[str, Callable[[dict[str, ast.AST]], ast.AST]] | None = None

classmethod from_code(name, kind, match, rewrite, placeholders, checker=None, derived=None)

Parameters:

• name (str)

• kind (ASTPatKind)

• match (str)

• rewrite (str)

• placeholders (set[str])

• checker (Callable[[dict[str, ast.AST]], bool] | dict[str, Callable[[ast.AST], bool]] | None)

• derived (dict[str, Callable[[dict[str, ast.AST]], ast.AST]] | None)

Return type:

ASTPatRewrite

get_name()

Return type:

str

match_placeholders(node)

Parameters:

node (ast.AST)

match(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

bool

rewrite(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST)

• field (str)

• inside_list (bool)

Return type:

ast.AST

class tilelang.autodd.ASTMutator

generic_visit(node)

visit(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST | None)

• field (str | None)

• inside_list (bool)

class tilelang.autodd.LabeledRewrite

label: int

rewrite: ASTRewrite

class tilelang.autodd.RewriteAttacher(rewrites)

Bases: ASTMutator

Parameters:

rewrites (list[ASTRewrite])

rewrites

uid_counter = 0

rewrite_counter = 0

rewrite_names

visit(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST | None)

• field (str | None)

• inside_list (bool)

tilelang.autodd.attach_rewrites(tree, rewrites)

Parameters:

• tree (ast.AST)

• rewrites (list[ASTRewrite])

Return type:

tuple[ast.AST, int, int]

class tilelang.autodd.RewriteApplier(target_labels)

Bases: ASTMutator

Parameters:

target_labels (set[int])

target_labels

applied_rewrites: set[int]

visited: set[int]

visit(node, parent, field, inside_list)

Parameters:

• node (ast.AST)

• parent (ast.AST | None)

• field (str | None)

• inside_list (bool)

tilelang.autodd.apply_rewrites(tree, target_labels)

Parameters:

• tree (ast.AST)

• target_labels (set[int])

Return type:

tuple[ast.AST, set[int]]

tilelang.autodd.test_rewrite(rewrite, code)

Parameters:

• rewrite (ASTRewrite)

• code (str)

class tilelang.autodd.Task

source: str

applied: list[int]

masked: list[int]

with_source(source)

Parameters:

source (str)

Return type:

Task

class tilelang.autodd.PDD(all_labels, init_proba=0.93)

Parameters:

• all_labels (list[int])

• init_proba (float)

all_labels

probas

apply(target_labels)

Parameters:

target_labels (set[int])

Return type:

set[int]

generator()

Return type:

collections.abc.Iterable[Task]

update(task, is_interesting)

Parameters:

• task (Task)

• is_interesting (bool)

class tilelang.autodd.TaskManager

Bases: abc.ABC

Helper class that provides a standard way to create an ABC using inheritance.

abstract task_generator()

Return type:

collections.abc.Iterable[Task]

abstract task_update(task, is_interesting)

Parameters:

• task (Task)

• is_interesting (bool)

classmethod from_source(source, *args, **kwargs)

Abstractmethod:

Parameters:

source (str)

Return type:

TaskManager

class tilelang.autodd.ASTPDD(tree, rewrites, init_proba=0.93)

Bases: TaskManager, PDD

Helper class that provides a standard way to create an ABC using inheritance.

Parameters:

• tree (ast.AST)

• rewrites (list[ASTRewrite])

• init_proba (float)

classmethod from_source(source, *args, **kwargs)

apply(target_labels)

Parameters:

target_labels (set[int])

Return type:

set[int]

task_generator()

Return type:

collections.abc.Iterable[Task]

task_update(task, is_interesting)

Parameters:

• task (Task)

• is_interesting (bool)

tilelang.autodd.ruff_fix_code(code_string, fix_lint=True, format_code=True)

Parameters:

• code_string (str)

• fix_lint (bool)

• format_code (bool)

Return type:

str

class tilelang.autodd.LinePDD(source, init_proba=0.93)

Bases: TaskManager, PDD

Helper class that provides a standard way to create an ABC using inheritance.

Parameters:

• source (str)

• init_proba (float)

lines

classmethod from_source(source, *args, **kwargs)

task_generator()

Return type:

collections.abc.Iterable[Task]

task_update(task, is_interesting)

Parameters:

• task (Task)

• is_interesting (bool)

class tilelang.autodd.Ruff(source, fix_lint=True, format_code=True)

Bases: TaskManager

Helper class that provides a standard way to create an ABC using inheritance.

Parameters:

• source (str)

• fix_lint (bool)

• format_code (bool)

source

fix_lint = True

format_code = True

finished = False

classmethod from_source(source, *args, **kwargs)

Parameters:

source (str)

Return type:

Ruff

task_generator()

task_update(task, is_interesting)

Parameters:

• task (Task)

• is_interesting (bool)

class tilelang.autodd.AsyncPythonRunner

process = None

input_queue = None

output_queue = None

lock

start_proc()

stop_proc()

__enter__()

__exit__(exc_type, exc_value, traceback)

async run(code, timeout=5.0)

Parameters:

• code (str)

• timeout (float)

class tilelang.autodd.SubProcRunner

__enter__()

__exit__(exc_type, exc_value, traceback)

async run(code, timeout=5.0)

Parameters:

• code (str)

• timeout (float)

tilelang.autodd.clean_empty_pass(code)

Parameters:

code (str)

Return type:

str

tilelang.autodd.JobBackend

class tilelang.autodd.ParTaskManager

err_msg: str

text: str

output_file: pathlib.Path

timeout: int = 60

num_workers: int = 1

backend: JobBackend = 'runner'

allow_larger: bool = False

__post_init__()

property text_len

reset(task_manager)

Parameters:

task_manager (TaskManager)

async get_next_task()

Return type:

Task | None

async submit_result(task, is_interested)

Parameters:

• task (Task)

• is_interested (bool)

post_proc(text)

async worker(wid)

Parameters:

wid (int)

async start_workers()

async stop_workers()

async run_async(task_manager)

Parameters:

task_manager (TaskManager)

async run_with(cls, *args, **kwargs)

Parameters:

cls (type[TaskManager])

class tilelang.autodd.Args

Bases: NamedTuple

source: pathlib.Path

err_msg: str

output: pathlib.Path

backend: JobBackend

timeout: int

jobs: int

async tilelang.autodd.main(args)

Parameters:

args (Args)

tilelang.autodd.cli_main(argv=None)

Parameters:

argv (Sequence[str] | None)

Return type:

None