# 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

```python
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

```python
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

```python
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

```python
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,
        )
```