tvm_tirx.tirx.ComposeLayout

tvm_tirx.tirx.ComposeLayout#

class tvm_tirx.tirx.ComposeLayout(per_element, swizzle_len, atom_len, tile_layout, swizzle_inner=True)#

Bases: Layout

Apply an XOR swizzle to the memory coordinate of a tile layout.

All three bit counts must be nonnegative.

Parameters:

  • per_element (int) – Number of low address bits left unchanged by the swizzle.

  • swizzle_len (int) – Number of bits in each XOR group; zero disables swizzling.

  • atom_len (int) – Separation between the bit groups; must be at least swizzle_len.

  • tile_layout (TileLayout) – The underlying mapping to memory coordinates.

  • swizzle_inner (bool, optional) – If true, XOR the higher group into the lower group; otherwise XOR the lower group into the higher group. Defaults to true.

Examples

from tvm_tirx.script import tirx as T

tile = T.TileLayout(T.S[8, 8])
layout = T.ComposeLayout(0, 2, 3, tile)
# Memory offset 10 becomes 10 ^ ((10 >> 3) & 3) = 11.
assert int(layout.apply(1, 2)["m"]) == 11
property per_element#
property swizzle_len#
property atom_len#
property swizzle_inner#
property tile_layout#
property inner_mask#
property outer_mask#