Swizzle Helper Functions

gemm_swizzle2d_Nz

Converts a GEMM iteration index to a 2D tile coordinate in the output matrix.

Tiles are grouped by columns, with adjacent groups traversing rows in reverse order, forming a serpentine scheduling pattern. The last group with fewer than swizzle_offset columns is handled automatically.

Signature

gemm_swizzle2d_Nz(
    iter_id,
    data_row_shape,
    data_col_shape,
    tile_row_shape,
    tile_col_shape,
    swizzle_offset=7,
)

Parameters

Parameter	Description
iter_id	Current iteration index. Range: [0, ceil(data_row_shape / tile_row_shape) * ceil(data_col_shape / tile_col_shape)).
data_row_shape	Number of rows in the output matrix, typically M.
data_col_shape	Number of columns in the output matrix, typically N.
tile_row_shape	Number of rows per tile, typically BLOCK_SIZE_M.
tile_col_shape	Number of columns per tile, typically BLOCK_SIZE_N.
swizzle_offset	Number of column tiles per group, defaults to 7.

Return Value

Returns (data_row_idx, data_col_idx), the row and column indices of the tile.

Scheduling Order

For example, with a 3-row by 10-column tile grid and swizzle_offset=7: the first 7 columns traverse rows in the order 0 -> 1 -> 2, while the last 3 columns traverse in the order 2 -> 1 -> 0.

Use Cases

Used in the GEMM phase of distributed fused kernels.

• Grouping tiles from adjacent columns helps improve cache locality for matrix data.

• Serpentine traversal connects adjacent groups on the same row, reducing large jumps in tile coordinates.

• The returned (block_id_m, block_id_n) can be directly used to compute block addresses for A, B, and C.

Notes

• The shape, tile, and swizzle_offset parameters must be greater than 0.

• This function only adjusts the scheduling order and does not change the data layout.

Example

import triton
import triton.language as tl
from triton_dist.language.extra.ascend.algorithm import gemm_swizzle2d_Nz


@triton.jit
def gemm_tile_kernel(c_ptr, M, N, stride_cm, stride_cn,
                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid = tl.program_id(0)
    ncore = tl.num_programs(0)
    num_tiles = tl.cdiv(M, BLOCK_M) * tl.cdiv(N, BLOCK_N)

    for iter_id in range(pid, num_tiles, ncore):
        block_m, block_n = gemm_swizzle2d_Nz(
            iter_id, M, N, BLOCK_M, BLOCK_N,
        )
        # Use offs_m and offs_n to compute block addresses for A, B, and C.
        offs_m = block_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_n = block_n * BLOCK_N + tl.arange(0, BLOCK_N)

dist_swizzle2d_Nz

Converts a communication iteration index to a data tile, target rank, and tail block size.

Different tiles start communication from different ranks, avoiding multiple AI Cores concentrating access on the same rank.

Signature

dist_swizzle2d_Nz(
    iter_id,
    rank_size,
    data_row_shape,
    data_col_shape,
    tile_row_shape,
    tile_col_shape,
    comm_npu_split=1,
)

Parameters

Parameter	Description
iter_id	Current iteration index. Range: [0, rank_size * ceil(data_row_shape / tile_row_shape) * ceil(data_col_shape / tile_col_shape)).
rank_size	Number of ranks participating in communication.
data_row_shape	Number of data rows.
data_col_shape	Number of data columns.
tile_row_shape	Maximum number of rows per tile.
tile_col_shape	Maximum number of columns per tile.
comm_npu_split	Number of ranks per interleaved scheduling group, defaults to 1. Must be a divisor of rank_size.

Return Value

Returns (data_row_idx, data_col_idx, rank_idx, comm_row_size, comm_col_size):

• data_row_idx, data_col_idx: Row and column indices of the tile.

• rank_idx: Target rank.

• comm_row_size, comm_col_size: Actual number of rows and columns in the current tile.

Use Cases

Used in fused kernels such as AllGather-GEMM, GEMM-ReduceScatter, and GEMM-AllReduce.

• Rotating target ranks for different tiles helps reduce communication hotspots and access contention between ranks.

Notes

• comm_npu_split must satisfy rank_size % comm_npu_split == 0.

• This function only adjusts the scheduling order and does not change the data layout.

Example

import triton
import triton.language as tl
from triton_dist.language.extra.ascend.algorithm import dist_swizzle2d_Nz


@triton.jit
def communication_kernel(rank_size, rows, cols,
                         TILE_M: tl.constexpr, TILE_N: tl.constexpr):
    pid = tl.program_id(0)
    ncore = tl.num_programs(0)
    num_iters = rank_size * tl.cdiv(rows, TILE_M) * tl.cdiv(cols, TILE_N)

    for iter_id in range(pid, num_iters, ncore):
        # target_rank is used to select the remote rank; tile_rows and tile_cols are used to generate the tail block mask.
        tile_m, tile_n, target_rank, tile_rows, tile_cols = dist_swizzle2d_Nz(
            iter_id, rank_size, rows, cols, TILE_M, TILE_N,
        )