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.