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.load Method

load

Similar to vector::LoadOp but with sublane_mask and sublane_stride. When indices are negative, it means loading from negative offset of base address.

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