tinygrad codegen.linearizer

Note

You likely want the upstream tinygrad, not tinygrab. Tinygrab contains AI generated docstrings for a tinygrad snapshot. Upstream: https://tinygrad.org

class tinygrad.codegen.linearizer.Linearizer(ast: LazyOp, opts: LinearizerOptions | None = None)[source]

Bases: Kernel

The Linearizer class. Subclass of Kernel.

uop_alu_idx[source]

Takes a UOp, b, ops, ctx, op and dtype as input, returns the result of the ALU operation.

Type:: method

const[source]

Takes b, dtype and insert_before as input, returns a UOp with CONST operation.

Type:: method

cast[source]

Takes val and dtype as input, returns a UOp with CAST operation if val.dtype is not equal to dtype, else returns val.

Type:: method

get_reduce_acc[source]

Takes op and dtype as input, returns the initial accumulator value for the given reduce operation.

Type:: method

ast_parse(x: LazyOp, acc: List[UOp], offs: List[int] | None, loaded_buffers: Dict[MemBuffer | ConstBuffer | LocalBuffer, List[UOp]], do_reduce=False, loop_ctx=()) → List[UOp][source]

Parses abstract syntax tree (AST) for operations.

x

The operation to parse.

Type:: LazyOp

acc

The list of UOps accumulated so far.

Type:: List[UOp]

offs

List of offsets, if any.

Type:: Optional[List[int]]

loaded_buffers

A dictionary containing loaded buffers.

Type:: Dict[Union[MemBuffer, ConstBuffer, LocalBuffer], List[UOp]]

do_reduce

Flag indicating whether to perform a reduction operation. Default is False.

Type:: bool

loop_ctx

The loop context, if any.

Type:: tuple

Returns:: The list of UOps resulting from parsing the AST.
Return type:: List[UOp]

cast(val: UOp, dtype) → UOp[source]

If val.dtype is not equal to dtype, creates a CAST operation with the given parameters and returns the resultant UOp, else returns val.

Parameters:

val (UOp) – The input UOp.
dtype – The data type to cast to.

Returns:

The resultant UOp of the CAST operation, or val if val.dtype is equal to dtype.

Return type:

UOp

const(b: int | float, dtype=(5, 4, 'int', <class 'numpy.int32'>, 1), insert_before=None) → UOp[source]

Creates a CONST operation with the given parameters and returns the resultant UOp.

Parameters:

b (Union[int, float]) – The constant value.
dtype (DType, optional) – The data type of the result. Defaults to dtypes.int32.
insert_before (optional) – For insertion before a certain operation. Defaults to None.

Returns:

The resultant UOp of the CONST operation.

Return type:

UOp

get_reduce_acc(op, dtype: DType)[source]

Returns the initial accumulator value for the given reduce operation.

Parameters:

op – The reduce operation.
dtype (DType) – The data type of the result.

Returns:

0.0 if dtypes.is_float(dtype), 0 if integer, -math.inf for float and -(2**31) for int if ReduceOps.MAX.

Return type:

Depends on the operation

global_load(i: int, idxs: Sequence[Node], acc=None, barrier: UOp | None = None) → List[UOp][source]

Load data from a buffer into the GPU.

Parameters:

i (int) – The index of the buffer to load.
idxs (Sequence[Node]) – A sequence of nodes representing indices for each dimension.
acc (Optional[Any], optional) – An accumulator value, defaults to None.
barrier (Optional[UOp], optional) – A UOp barrier, defaults to None.

Returns:

A list of UOps representing the load operation.

Return type:

List[UOp]

Attributes:

buf (Buffer): The buffer from which data is loaded.
const (Union[ConstBuffer, Any]): The constant value if the buffer is a ConstBuffer, otherwise the accumulator.

global_store(i: int, idxs: List[Node], store: List[UOp]) → List[UOp][source]

Perform a global store operation on the buffer at index i.

Parameters:

i (int) – The index of the buffer to perform the store operation on.
idxs (List[Node]) – The list of Nodes representing indices for the store operation.
store (List[UOp]) – The list of UOps containing the values to be stored.

Returns:

A list of UOps representing the stored values.

Return type:

List[UOp]

kernel_cnt: Final[DefaultDict[str, int]] = {}

linearize()[source]

Linearize the computation graph.

This function checks if new options are applied and if it is necessary to relinearize. It also handles backups, global uop cache, limiting dimensions, uops, buffer uops, and loop uops.

self.applied_opts

The newly applied options.

Type:: Dict

self.applied_opts_cache

The previously applied options.

Type:: Dict

sts_backup

A backup of the current state.

Type:: List

gfr_backup

A backup of the group for reduce.

Type:: List

upc_backup

A backup of the upcasted value.

Type:: Any

self.saved_exprs

A dictionary storing expressions as keys and UOps as values.

Type:: Dict[Tuple, UOp]

self.opts.global_max

The maximum global dimension.

Type:: int

self.opts.local_max

The maximum local dimension.

Type:: int

self.uops

A list of UOp objects, representing the computation operations.

Type:: List[UOp]

self.buf_uops

A list of optional UOp objects for buffering.

Type:: List[Optional[UOp]]

self.loop_uops

A dictionary with loop names as keys and UOp objects as values.

Type:: Dict[str, UOp]

Returns:: self if relinearization is not needed, otherwise None.

render_ops: Any = {<class 'tinygrad.shape.symbolic.AndNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.DivNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.LtNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.ModNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.MulNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.NumNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.SumNode'>: <function Linearizer.<lambda>>, <class 'tinygrad.shape.symbolic.Variable'>: <function Linearizer.<lambda>>}

uop(uop: UOps, dtype: DType | None = None, vin: Tuple[UOp, ...] = (), arg: Any = None, cachable=True, insert_before=None, simplify=True) → UOp[source]

This function creates and manipulates unary operations (uops). It supports various uop types, data types, and operations.

self: The instance of the class that this method is called on

uop[source]: The type of unary operation to perform (e.g., UOps.PHI for phi function)

dtype: The desired output data type (optional)

vin: A tuple of input unary operations (default is an empty tuple)

arg: An optional argument that can be used in certain uops (default is None)

cachable: A boolean flag indicating whether the operation should be cached for faster access (default is True)

insert_before: An optional parameter to specify where this operation should be inserted (default is None)

simplify: A boolean flag indicating whether to attempt to simplify the uop before creating it (default is True)

Returns:: The created and manipulated UOp object based on the given parameters.

uop_alu_idx(a: ~tinygrad.codegen.linearizer.UOp, b, ops, ctx: ~tinygrad.codegen.linearizer.Linearizer, op, dtype=(5, 4, 'int', <class 'numpy.int32'>, 1))[source]

Renders the ALU operation with given parameters and returns the result.

Parameters:

a (UOp) – The first input UOp.
b – The second input.
ops – The operations.
ctx (Linearizer) – The linearizer context.
op – The operation to perform.
dtype (DType, optional) – The data type of the result. Defaults to dtypes.int32.

Returns:

The result of the ALU operation.

Return type:

UOp

class tinygrad.codegen.linearizer.UOp(uop: UOps, dtype: DType | None, vin: Tuple[UOp, ...], arg: Any)[source]

Bases: object

Data class for UOp.

uop

UOps - Operation to be performed.

Type:: tinygrad.codegen.linearizer.UOps

dtype

Optional[DType] - Datatype of the operation result.

Type:: tinygrad.helpers.DType | None

vin

Tuple[UOp, …] - Inputs for the operation.

Type:: Tuple[tinygrad.codegen.linearizer.UOp, …]

arg

Any - Additional argument for the operation.

Type:: Any

arg: Any

dtype: DType | None

uop: UOps

vin: Tuple[UOp, ...]

class tinygrad.codegen.linearizer.UOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

Enumeration for Unified Operations (UOps).

LOOP: Indicates a loop operation.

IF: Indicates an if operation.

END: Indicates the end of an operation block.

SPECIAL: Indicates a special operation, which can be global, local, or other type.

DEFINE_GLOBAL: Indicates defining a global variable.

DEFINE_LOCAL: Indicates defining a local variable.

DEFINE_ACC: Indicates defining a buffer.

LOAD: Indicates loading data.

STORE: Indicates storing data.

CONST: Indicates a constant operation.

BARRIER: Indicates a barrier operation.

PHI: Indicates a phi operation, which is used in SSA form for merging control flow.

ALU: Indicates an arithmetic logic unit operation.

WMMA: Indicates a warp matrix multiplication and accumulation operation.

CAST: Indicates a cast operation.

GEP: Indicates a get element pointer operation, which is used to compute an address of a data element in memory.

ALU = 13

BARRIER = 11

CAST = 15

CONST = 10

DEFINE_ACC = 7

DEFINE_GLOBAL = 5

DEFINE_LOCAL = 6

END = 3

GEP = 16

IF = 2

LOAD = 8

LOOP = 1

PHI = 12

SPECIAL = 4

STORE = 9

WMMA = 14

tinygrad.codegen.linearizer.get_grouped_dims(prefix, start_dim, local_dims, maxdim: int = 0)[source]

Generate grouped dimensions and loop local indices.

Parameters:

prefix – str - Prefix for variable names.
start_dim – int - Start dimension index.
local_dims – list - List of local dimensions.
maxdim – int, optional - Maximum dimension value, default is 0.

Returns:

tuple[list, list] - Grouped dimensions and loop local indices.