For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo function
block_swizzle
def block_swizzle(block_idx: IndexList[Int(2), element_type=block_idx.element_type], grid_dim: IndexList[Int(2), element_type=block_idx.element_type]) -> IndexList[Int(2), element_type=block_idx.element_type]
Remaps a linear block index into a swizzled two-dimensional block coordinate.
Applies CUTLASS-style block swizzling along the N dimension to improve L2 cache locality for tall-and-narrow matmul grids.
Args:
- block_idx (
IndexList[Int(2), element_type=block_idx.element_type]): The original two-dimensional block coordinate. - grid_dim (
IndexList[Int(2), element_type=block_idx.element_type]): The grid dimensions matchingblock_idx.
Returns:
IndexList[Int(2), element_type=block_idx.element_type]: The swizzled two-dimensional block coordinate.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!