Tile API

Programming model

Tile

A tile is a multi-dimensional array of values in local memory (Unified Buffer, L1 Cache, etc.)

Tensor

A tensor is a contiguous ND-array of values in Global Memory.

tensor

Define a new tensor.

block_idx

block_num

num_tiles

Creation operations

full

full_like

zeros

zeros_like

Memory operations

load

store

Binary operations

add

Computes the element-wise addition of input and other.

div

Computes the element-wise division of input and other.

equal

Computes the element-wise 'equality' comparison of input and other.

greater

Computes the element-wise 'greater' comparison of input and other.

greater_equal

Computes the element-wise 'greater or equal' comparison of input and other.

left_shift

Computes the element-wise left shift (bitwise) of input and other.

less

Computes the element-wise 'less' comparison of input and other.

less_equal

Computes the element-wise 'less or equal' comparison of input and other.

matmul

Computes the matrix multiplication of input and other.

maximum

Computes the element-wise maximum of input and other.

minimum

Computes the element-wise minimum of input and other.

mul

Computes the element-wise multiplication of input and other.

not_equal

Computes the element-wise 'inequality' comparison of input and other.

right_shift

Computes the element-wise right shift (bitwise) of input and other.

sub

Computes the element-wise subtraction of input and other.

Unary operations

abs

Computes the element-wise absolute value of input.

ceil

Computes the element-wise ceil rounding of input.

cos

Computes the element-wise cosine of input.

cosh

Computes the element-wise hyperbolic cosine of input.

erf

Computes the element-wise error function of input.

exp

Computes the element-wise exponential of input.

exp2

Computes the element-wise exponential (base 2) of input.

floor

Computes the element-wise floor rounding of input.

log

Computes the element-wise natural logarithm of input.

log2

Computes the element-wise logarithm (base 2) of input.

negative

Computes the element-wise negation of input.

relu

Computes the element-wise ReLU value of input.

rsqrt

Computes the element-wise inverse square root of input.

sin

Computes the element-wise sine of input.

sinh

Computes the element-wise hyperbolic sine of input.

sqrt

Computes the element-wise square root of input.

softmax

Computes the element-wise softmax of input.

tan

Computes the element-wise tangent of input.

tanh

Computes the element-wise hyperbolic tangent of input.

Indexing operations

mask

where

Reduction operations

reduce_min

Returns the minimum value of each row of the input tile in the given dimensions dims.

reduce_max

Returns the maximum value of each row of the input tile in the given dimensions dims.

reduce_sum

Returns the sum of each row of the input tile in the given dimensions dims.

reduce_prod

Returns the product of each row of the input tile in the given dimensions dims.

Shape manipulation operations

broadcast_to

This function can also be called as a member function on Tile, as x.broadcast_to(...) instead of broadcast_to(x, ...).

expand_dims

This function can also be called as a member function on Tile, as x.expand_dims(...) instead of expand_dims(x, ...).

reshape

This function can also be called as a member function on Tile, as x.reshape(...) instead of reshape(x, ...).

squeeze

This function can also be called as a member function on Tile, as x.squeeze(...) instead of squeeze(x, ...).

Atomic operations

atomic_add

atomic_max

atomic_min