Mojo struct
Swizzle
@register_passable(trivial)
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]:
- Extract
bitsnumber of bits fromistarting from positionbase + max(0, shift). Let's call thisYYY. - Extract
bitsnumber of bits fromistarting from positionbase - min(0, shift). Let's call thisZZZ. - 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).
Aliases
has_shape = False: Indicates if layout has shape. Swizzle always False.
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,
CollectionElement,
Copyable,
ExplicitlyCopyable,
LayoutTrait,
Movable,
Stringable,
UnknownDestructibility,
Writable
Methods
__init__
__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__
__call__(self, index: IntTuple[origin]) -> Int
Apply swizzle to an IntTuple index.
Unwraps the IntTuple and applies the swizzle to the integer value.
Args:
- index (
IntTuple[origin]): The IntTuple index to swizzle.
Returns:
The swizzled index value.
__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:
The swizzled offset value.
__call__(self, offset: SIMD[dtype, 1]) -> SIMD[dtype, 1]
Apply swizzle to a scalar offset.
Scalar version of the swizzle operation. Applies swizzle to a scalar offset.
Args:
- offset (
SIMD[dtype, 1]): The scalar offset to swizzle.
Returns:
The swizzled scalar value.
size
size(self) -> Int
Get the size of the swizzle pattern.
Calculates the size of the memory region affected by the swizzle pattern.
Returns:
The size of the swizzle pattern.
cosize
cosize(self) -> Int
Get the cosize of the swizzle pattern.
Cosize is the same as size for swizzle layouts, representing the output size.
Returns:
The cosize of the swizzle pattern (same as size).
write_to
write_to[W: Writer](self, mut writer: W)
Write the swizzle parameters to a writer.
Outputs the swizzle parameters (bits, base, shift) in a tuple format.
Parameters:
- W (
Writer): The writer type that implements the Writer trait.
Args:
- writer (
W): The writer to write to.
__str__
__str__(self) -> String
Convert the swizzle to a string representation.
Returns:
String representation of the swizzle parameters.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!