For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

Swizzle

struct Swizzle

Swizzle functor for memory access pattern optimization.

Implements a swizzling pattern to reduce bank conflicts in shared memory accesses. It XORs specific bits of memory indices based on configurable parameters.

Swizzle operation: Given index i, and Swizzle[bits, base, shift]:

1. Extract bits number of bits from i starting from position base + max(shift, 0). Let's call this YYY.
2. Extract bits number of bits from i starting from position base - min(shift, 0). Let's call this ZZZ.
3. Result is i ^ (YYY shifted by 'shift' positions).

Example (Swizzle[2, 0, 3]): Input index bits: xxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxxx Output index bits: xxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxxx where AA = ZZ ^ YY.

Attributes: bits (Int): Number of bits in the mask (YYY). base (Int): Number of least significant bits to keep constant. shift (Int): Shift distance for the mask (positive: right, negative: left). yyy_mask (Int): Mask for the bits to be shifted (YYY). zzz_mask (Int): Mask for the target bits (ZZZ).

Fields

• ​bits (Int): Number of bits in the mask.
• ​base (Int): Number of least significant bits to keep constant.
• ​shift (Int): Distance to shift the mask (pos right, neg left).
• ​yyy_mask (Int): Mask for the bits to be shifted.
• ​zzz_mask (Int): Mask for the target bits.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable, Writable

comptime members

has_shape

comptime has_shape = False

Indicates if layout has shape. Swizzle always False.

Methods

__init__

def __init__(bits: Int, base: Int, shift: Int) -> Self

Initialize a Swizzle object.

Configures the swizzle operation based on bits, base, and shift parameters.

Args:

• ​bits (Int): Number of bits in the mask.
• ​base (Int): Least significant bits to keep constant.
• ​shift (Int): Distance to shift the mask.

__call__

def __call__(self, index: IntTuple) -> Int

Apply swizzle to an IntTuple index.

Unwraps the IntTuple and applies the swizzle to the integer value.

Args:

• ​index (IntTuple): The IntTuple index to swizzle.

Returns:

Int: The swizzled index value.

def __call__(self, offset: Int) -> Int

Apply swizzle to an integer offset.

Performs the swizzle operation on an integer offset to rearrange memory access patterns.

Args:

• ​offset (Int): The integer offset to swizzle.

Returns:

Int: The swizzled offset value.

def __call__(self, offset: Scalar) -> Scalar[offset.dtype]

Apply swizzle to a scalar offset.

Scalar version of the swizzle operation. Applies swizzle to a scalar offset.

Args:

• ​offset (Scalar): The scalar offset to swizzle.

Returns:

Scalar[offset.dtype]: The swizzled scalar value.

size

def size(self) -> Int

Get the size of the swizzle pattern.

Calculates the size of the memory region affected by the swizzle pattern.

Returns:

Int: The size of the swizzle pattern.

cosize

def cosize(self) -> Int

Get the cosize of the swizzle pattern.

Cosize is the same as size for swizzle layouts, representing the output size.

Returns:

Int: The cosize of the swizzle pattern (same as size).

write_to

def write_to(self, mut writer: T)

Write the swizzle parameters to a writer.

Outputs the swizzle parameters (bits, base, shift) in a tuple format.

Args:

• ​writer (T): The writer to write to.