TPU Dialect

Refer to the official documentation for more details.

Reactant.MLIR.Dialects.tpu.broadcast_in_sublanes Method

broadcast_in_sublanes

For each sublane i, broadcasts the value in lane lane + i along the entire sublane. If lane + i is not in [0, lane_count), then the value in sublane i is not defined (can be anything).

source

Reactant.MLIR.Dialects.tpu.create_subelement_mask Method

create_subelement_mask

The "half-sublanes", "quarter-sublanes", etc. (unit is determined by the type of output) of the mask are masked in the range specified by from and to.

• If from <= to, the range [from, to) is set and the rest is unset.

• If to <= from, the range [to, from) is unset and the rest is set.

All lanes are set identically.

Example

%msk = tpu.create_subelement_mask 3, 9 : vector<8x128x2xi1>

This creates a mask %msk where, for all lanes, %msk[*][lane][*] is:

[[0, 0], [0, 1], [1, 1], [1, 1], [1, 0], [0, 0], [0, 0], [0, 0]]

It is currently only supported:

• In TPU v4, for num_subelems of 1 and 2.

• In TPU v5, for num_subelems of 1, 2, and 4.

source

Reactant.MLIR.Dialects.tpu.dynamic_gather Method

dynamic_gather

Gathers elements from source using indices.

The specified dimensions of source are collapsed together and indexed by indices.

Given a shape N0 x N1 x ..., the output[i0, i1, ...] is given by collapsed_source[j0, j1, ..., indices[i0, i1, ...] mod M] where

• collapsed_source is the result of collapsing dimensions of source into a new trailing dimension of size M.

• jk is the subsequence of in for n not in dimensions.

When a single dimension is specified, this is similar to np.take_along_axis.

source

Reactant.MLIR.Dialects.tpu.iota Method

iota

Creates a vector that with values that start at 0 and increase along a dimension resulting from collapsing the given dimensions together in row-major order.

Example

tpu.iota {dimensions = array<i32: 2, 0>} : vector<4x3x2xi16>

This produces a vector with the following values:

[[[0, 4], [0, 4], [0, 4]]
 [[1, 5], [1, 5], [1, 5]]
 [[2, 6], [2, 6], [2, 6]]
 [[3, 7], [3, 7], [3, 7]]]

source

Reactant.MLIR.Dialects.tpu.rotate Method

rotate

Rotates the given vector by the given amount in the given dimension, i.e., for a 2D vector of shape (m, n), rotating dim 0 by amount will shift a row at index i to index (i + amount) % m

source