Swizzle Helper API¶

`gemm_swizzle2d`¶

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

By default, tiles are grouped by columns, with adjacent groups traversing rows in reverse order to form a serpentine scheduling pattern. swizzle_direction can instead group rows and traverse columns. A partial final group is handled automatically in either direction.

gemm_swizzle2d_Nz remains supported with its original signature and defaults. It is a special case of gemm_swizzle2d with swizzle_direction=1. We recommend gemm_swizzle2d for new code.

Signature¶

gemm_swizzle2d(
    iter_id,
    data_row_shape,
    data_col_shape,
    tile_row_shape,
    tile_col_shape,
    swizzle_offset=7,
    swizzle_direction: tl.constexpr = 1,
)

Parameters¶

Parameter	Data type	Description
`iter_id`	Integer	Current iteration index. Range: `[0, ceil(data_row_shape / tile_row_shape) * ceil(data_col_shape / tile_col_shape))`.
`data_row_shape`	Integer	Number of rows in the output matrix, typically `M`.
`data_col_shape`	Integer	Number of columns in the output matrix, typically `N`.
`tile_row_shape`	Integer	Number of rows per tile, typically `BLOCK_SIZE_M`.
`tile_col_shape`	Integer	Number of columns per tile, typically `BLOCK_SIZE_N`.
`swizzle_offset`	Integer	Maximum number of tiles in the grouped dimension. Valid range: `[1, 7]`; defaults to `7`.
`swizzle_direction`	`tl.constexpr`	Grouped dimension: `0` groups rows and traverses columns; `1` groups columns and traverses rows. Defaults to `1` and must be compile-time constant.

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, swizzle_offset=7, and swizzle_direction=1: 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¶

iter_id and all four shape parameters must be greater than 0; swizzle_offset must be in [1, 7]. A value of 0 is invalid because it makes the group size zero.
swizzle_direction must be 0 or 1.
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


@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(
            iter_id, M, N, BLOCK_M, BLOCK_N,
        )
        # Use offs_m and offs_n to compute the 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)
        c_offsets = (
            offs_m[:, None] * stride_cm
            + offs_n[None, :] * stride_cn
        )
        c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

        # Omit the A and B tile loads and the accumulation over the K dimension.
        accumulator = ...
        tl.store(c_ptr + c_offsets, accumulator, mask=c_mask)

Compatibility API¶

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

`dist_swizzle2d`¶

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.

dist_swizzle2d_Nz remains supported with its original signature and defaults. It is a special case of dist_swizzle2d with comm_data_split=1 and swizzle_direction=1. We recommend dist_swizzle2d for new code.

Signature¶

dist_swizzle2d(
    iter_id,
    rank_size,
    data_row_shape,
    data_col_shape,
    tile_row_shape,
    tile_col_shape,
    comm_npu_split=1,
    comm_data_split=1,
    swizzle_direction: tl.constexpr = 1,
)

Parameters¶

Parameter	Data type	Description
`iter_id`	Integer	Current iteration index. Range: `[0, rank_size * ceil(data_row_shape / tile_row_shape) * ceil(data_col_shape / tile_col_shape))`.
`rank_size`	Integer	Number of ranks participating in communication.
`data_row_shape`	Integer	Number of data rows.
`data_col_shape`	Integer	Number of data columns.
`tile_row_shape`	Integer	Maximum number of rows per tile.
`tile_col_shape`	Integer	Maximum number of columns per tile.
`comm_npu_split`	Integer	Number of ranks per interleaved scheduling group, defaults to `1`. Must be a divisor of `rank_size`.
`comm_data_split`	Integer	Maximum number of data tiles per scheduling group when `swizzle_direction=0`, defaults to `1`.
`swizzle_direction`	`tl.constexpr`	Grouped dimension: `0` groups flattened data tiles and traverses ranks; `1` groups ranks and traverses data tiles. Defaults to `1` and must be compile-time constant.

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 current 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.
comm_data_split must be greater than 0.
swizzle_direction must be 0 or 1.
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


@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(
            iter_id, rank_size, rows, cols, TILE_M, TILE_N,
        )

Compatibility API¶

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

Swizzle Helper API¶

gemm_swizzle2d¶

Signature¶

Parameters¶

Return Value¶

Scheduling Order¶

Use Cases¶

Notes¶

Example¶

Compatibility API¶

dist_swizzle2d¶

Signature¶

Parameters¶

Return Value¶

Use Cases¶

Notes¶

Example¶

Compatibility API¶

`gemm_swizzle2d`¶

`dist_swizzle2d`¶