Skip to content

tinygrad directory layout¤

This explains the flow of a big graph down to programs.

Directories are listed in order of how they are processed.

tinygrad/kernelize¤

Group UOps into kernels.

get_kernelize_map ¤

get_kernelize_map(sink: UOp) -> dict[UOp, UOp]

Function to transform the Tensor UOp graph into a version with Ops.KERNEL

Parameters:

  • sink (UOp) –

    The Ops.SINK rooting the Tensor graph.

Returns:

  • dict[UOp, UOp] –

    Map transforming each UOp in the sink to the Ops.KERNEL graph.

tinygrad/opt¤

Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.

get_optimized_ast ¤

get_optimized_ast(ast: UOp, renderer: Renderer) -> UOp

Optimize an AST based on heuristics or BEAM search.

Parameters:

  • ast (UOp) –

    The Ops.SINK rooted AST

  • renderer (Renderer) –

    The renderer used to generate the code

Returns:

  • UOp –

    The Ops.SINK rooted AST transformed to apply the opts and with a KernelInfo in the arg.

tinygrad/codegen¤

Transform the optimized ast into a linearized list of UOps.

full_rewrite ¤

full_rewrite(
    sink: UOp, opts: Renderer | None = None
) -> list[UOp]

Function to transform the Kernel UOp graph into a linearized program.

Parameters:

  • sink (UOp) –

    The Ops.SINK rooting the Kernel graph.

  • opts (Renderer | None, default: None ) –

    The Renderer (can change how things are processed, fix this).

Returns:

  • list[UOp] –

    Linear program in UOps.

tinygrad/renderer¤

Transform the linearized list of UOps into a program, represented as a string.

Renderer ¤

Methods:

render ¤

render(uops: list[UOp]) -> str

tinygrad/engine¤

Abstracted high level interface to the runtimes.

get_program ¤

get_program(ast: UOp, renderer: Renderer) -> ProgramSpec

Transform an AST into a ProgramSpec. May trigger BEAM search.

Parameters:

  • ast (UOp) –

    The Ops.SINK rooted AST

  • renderer (Renderer) –

    The renderer used to generate the code

Returns:

  • ProgramSpec –

    The ProgramSpec of the program.