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