# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Reusable tile scheduler helpers for TIR tests/kernels.

These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
instances are automatically treated as meta values inside @T.prim_func.
"""

from tvm.backend.cuda.lang.pipeline import Pipeline, PipelineState
from tvm.script import tirx as T


@T.meta_class
class BaseTileScheduler:
    """Base class for tile schedulers with common state and macros."""

    def __init__(self, prefix: str):
        self.m_idx = T.local_scalar("int32")
        self.n_idx = T.local_scalar("int32")
        self.linear_idx = T.local_scalar("int32")

    @T.inline
    def update_current_m_n_idx(self, linear_idx):
        # To be implemented by subclasses
        pass

    @T.inline
    def init(self, linear_init):
        self.linear_idx = linear_init
        self.update_current_m_n_idx(linear_init)

    @T.inline
    def next_tile(self, step):
        self.linear_idx = self.linear_idx + step
        self.update_current_m_n_idx(self.linear_idx)

    def valid(self, total_tiles):
        return self.linear_idx < total_tiles


class ClusterPersistentScheduler2D(BaseTileScheduler):
    """
    Tile scheduler for cluster-based persistent kernels.

    Distributes a 2D tile grid across persistent clusters using group-major ordering
    for L2 cache locality. Each cluster starts at its cluster_id and strides by
    num_clusters to process tiles.

    Tile Ordering (group-major for L2 locality):
    - Tiles are grouped into "L2 groups" of `l2_group_size` rows
    - Within a group, tiles are visited in column-major order within the group
    - Groups are processed in row-major order

    Example with 4x4 tiles, l2_group_size=2:
        Group 0 (rows 0-1):  0  2  4  6
                             1  3  5  7
        Group 1 (rows 2-3):  8 10 12 14
                             9 11 13 15

    Serpentine Mode (serpentine=True):
    - Uses CUTLASS-style 2D block swizzle with serpentine traversal
    - Grid is divided into swizzle_size x swizzle_size blocks
    - Within each block, tiles are visited in row-major order
    - Blocks are traversed in serpentine order (even block-rows forward, odd backward)
    - This provides better L2 locality by reusing both A and B tiles

    Example with 4x4 tiles, swizzle_size=2, serpentine=True:
        Block layout:
          Block(0,0)  Block(0,1)
          Block(1,0)  Block(1,1)

        Tile numbering with serpentine:
               n=0  n=1  n=2  n=3
          m=0   0    1   14   15
          m=1   2    3   12   13
          m=2   4    5   10   11
          m=3   6    7    8    9

        Traversal: Block(0,0) -> Block(1,0) -> Block(1,1) -> Block(0,1)
                   (serpentine: down in col 0, then up in col 1)

    Parameters
    ----------
    prefix : str
        Prefix for TIR variable names
    num_m_tiles : int | T.ExprLike
        Total number of tiles in M dimension (can be runtime expression)
    num_n_tiles : int
        Total number of tiles in N dimension
    num_clusters : int
        Number of persistent clusters (determines stride)
    l2_group_size : int
        Number of M-tile rows per L2 locality group (default: 8)
        When serpentine=True, this is used as swizzle_size for 2D blocks
    cluster_m : int
        Cluster dimension in M for hierarchical scheduling (default: 1)
    cluster_n : int
        Cluster dimension in N for hierarchical scheduling (default: 1)
    serpentine : bool
        If True, use CUTLASS-style 2D block swizzle with serpentine traversal (default: False)

    Attributes
    ----------
    m_idx : T.local_scalar
        Current M tile index (output)
    n_idx : T.local_scalar
        Current N tile index (output)
    work_idx : T.local_scalar
        Global work item index for this cluster
    tile_count : T.local_scalar
        Number of tiles processed by this cluster so far

    Usage
    -----
    ```python
    scheduler = ClusterPersistentScheduler2D(
        "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES,
        num_clusters=NUM_CLUSTERS, l2_group_size=8
    )
    scheduler.init(cluster_id)  # cluster_id = cta_idx // CLUSTER_SIZE

    while scheduler.valid():
        m = T.meta_var(scheduler.m_idx)  # current M tile
        n = T.meta_var(scheduler.n_idx)  # current N tile
        # ... process tile (m, n) ...
        scheduler.next_tile()
    ```

    Examples
    --------
    Example 1: Basic persistent kernel
    ```
    num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2
    cluster_m=1, cluster_n=1 (default, no tile subdivision)

    Group-major tile numbering (l2_group_size=2):
           n=0  n=1  n=2  n=3
      m=0   0    2    4    6   ┐ L2 group 0
      m=1   1    3    5    7   ┘
      m=2   8   10   12   14   ┐ L2 group 1
      m=3   9   11   13   15   ┘

    Work distribution (cluster starts at cluster_id, strides by num_clusters=3):
      cluster 0: work_idx 0,3,6,9,12,15  -> tiles 0,3,6,9,12,15
      cluster 1: work_idx 1,4,7,10,13    -> tiles 1,4,7,10,13
      cluster 2: work_idx 2,5,8,11,14    -> tiles 2,5,8,11,14

    Tile grid (which cluster handles each tile):
           n=0  n=1  n=2  n=3
      m=0   C0   C2   C1   C0   ┐ L2 group 0
      m=1   C1   C0   C2   C1   ┘
      m=2   C2   C1   C0   C2   ┐ L2 group 1
      m=3   C0   C2   C1   C0   ┘

    Tile sequence per cluster (in execution order):
      cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3)
      cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2)
      cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3)
    ```

    Example 2: 2SM GEMM (typical B200 config)
    ```
    M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1
    => M_TILES=8, N_TILES=4
    => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity)

    Scheduler params:
      num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8
      cluster_m=1, cluster_n=1

    Key: Scheduler outputs CLUSTER-level tiles.
         All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler.
         CTAs differentiate via cluster_rank (computed OUTSIDE scheduler):
           cluster_rank = cta_idx % CLUSTER_SIZE
           cb_m = cluster_rank % CLUSTER_M   # 0 or 1 for 2SM
           cb_n = cluster_rank // CLUSTER_M  # 0 for 2SM

    Final CTA tile:
      cta_m = m_idx * CLUSTER_M + cb_m
      cta_n = n_idx * CLUSTER_N + cb_n

    Example: cluster 5 gets scheduler tile (1,2)
      CTA rank=0 (cb_m=0): actual tile (2,2)
      CTA rank=1 (cb_m=1): actual tile (3,2)
    ```
    """

    def __init__(
        self,
        prefix: str,
        num_m_tiles,
        num_n_tiles: int,
        num_clusters: int,
        l2_group_size: int = 8,
        cluster_m: int = 1,
        cluster_n: int = 1,
        serpentine: bool = False,
    ):
        super().__init__(prefix)
        self._num_m_tiles = num_m_tiles
        self._num_n_tiles = num_n_tiles
        self._num_clusters = num_clusters
        self._l2_group_size = l2_group_size
        self._cluster_m = cluster_m
        self._cluster_n = cluster_n
        self._serpentine = serpentine

        # Rename internal state for clarity
        self.work_idx = self.linear_idx  # alias: global work item index
        self.tile_count = T.local_scalar("int32")
        self.tile_idx = self.tile_count  # alias for backward compatibility

        is_static_m = isinstance(num_m_tiles, int)

        # Number of tile columns after accounting for cluster_n
        n_tile_cols = (num_n_tiles + cluster_n - 1) // cluster_n
        self._N_TILE_COLS = n_tile_cols

        if is_static_m:
            self._M_TILE_ROWS = (num_m_tiles + cluster_m - 1) // cluster_m
            self._FULL_GROUPS = self._M_TILE_ROWS // l2_group_size
        else:
            # Dynamic expressions for runtime M
            self._M_TILE_ROWS = T.truncdiv(self._num_m_tiles + self._cluster_m - 1, self._cluster_m)
            self._FULL_GROUPS = T.truncdiv(self._M_TILE_ROWS, self._l2_group_size)

        self._TAIL_ROWS = self._M_TILE_ROWS - self._FULL_GROUPS * l2_group_size
        self._TOTAL_TILES = self._M_TILE_ROWS * n_tile_cols * cluster_m * cluster_n

        # For serpentine mode: precompute block counts
        if serpentine:
            self._N_BLOCKS = n_tile_cols // l2_group_size  # full blocks in N
            self._M_BLOCKS = (
                self._M_TILE_ROWS // l2_group_size
                if is_static_m
                else T.truncdiv(self._M_TILE_ROWS, l2_group_size)
            )
            self._BLOCK_SIZE = l2_group_size * l2_group_size  # tiles per block
            self._FULL_BLOCK_TILES = self._M_BLOCKS * self._N_BLOCKS * self._BLOCK_SIZE
            # Residual tiles (not covered by full blocks)
            self._RESIDUAL_N = n_tile_cols - self._N_BLOCKS * l2_group_size
            self._RESIDUAL_M = self._M_TILE_ROWS - self._M_BLOCKS * l2_group_size

    # fmt: off
    @T.inline
    def update_current_m_n_idx(self, work_idx):
        """Convert global work index to (m_idx, n_idx) tile coordinates."""
        CLUSTER_M = T.meta_var(self._cluster_m)
        CLUSTER_N = T.meta_var(self._cluster_n)

        # Extract hierarchical cluster-local offsets
        cluster_m_offset = T.meta_var(work_idx % CLUSTER_M)
        t = T.meta_var(work_idx // CLUSTER_M)
        cluster_n_offset = T.meta_var(t % CLUSTER_N)
        tile_linear = T.meta_var(t // CLUSTER_N)

        @T.inline
        def set_tile_coords(tile_row, tile_col):
            self.m_idx = tile_row * CLUSTER_M + cluster_m_offset
            self.n_idx = tile_col * CLUSTER_N + cluster_n_offset

        if self._serpentine:
            self._update_serpentine(tile_linear, set_tile_coords)
        else:
            self._update_group_major(tile_linear, set_tile_coords)

    def _update_group_major(self, tile_linear, set_tile_coords):
        """Group-major ordering with parse-time pruning of statically-dead branches.

        The TIR script parser does not constant-fold ``if False: ...``, so a
        Python-literal ``FULL_GROUPS == 0`` would otherwise produce
        ``T.bitwise_and(T.bool(False), tile_linear < 0)`` IR plus the dead
        then-leg.  Branch in plain Python here and only invoke the inline
        emitter that can actually fire.
        """
        full_zero = isinstance(self._FULL_GROUPS, int) and self._FULL_GROUPS == 0
        tail_zero = isinstance(self._TAIL_ROWS, int) and self._TAIL_ROWS == 0
        if full_zero and tail_zero:
            self._gm_emit_zero(set_tile_coords)
        elif full_zero:
            self._gm_emit_tail_only(tile_linear, set_tile_coords)
        elif tail_zero:
            self._gm_emit_full_only(tile_linear, set_tile_coords)
        else:
            self._gm_emit_full_and_tail(tile_linear, set_tile_coords)

    @T.inline
    def _gm_emit_zero(self, set_tile_coords):
        set_tile_coords(0, 0)

    @T.inline
    def _gm_emit_full_only(self, tile_linear, set_tile_coords):
        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
        GROUP_SIZE = T.meta_var(self._l2_group_size)
        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
        if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
            group_id: T.let = tile_linear // GROUP_SPAN
            within_group: T.let = tile_linear % GROUP_SPAN
            tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
            tile_col: T.let = within_group // GROUP_SIZE
            set_tile_coords(tile_row, tile_col)
        else:
            set_tile_coords(0, 0)

    @T.inline
    def _gm_emit_tail_only(self, tile_linear, set_tile_coords):
        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
        TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
        GROUP_SIZE = T.meta_var(self._l2_group_size)
        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
        if TAIL_ROWS > 0:
            rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
            tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
            tile_col: T.let = rem // TAIL_ROWS
            set_tile_coords(tile_row, tile_col)
        else:
            set_tile_coords(0, 0)

    @T.inline
    def _gm_emit_full_and_tail(self, tile_linear, set_tile_coords):
        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
        TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
        GROUP_SIZE = T.meta_var(self._l2_group_size)
        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
        if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
            group_id: T.let = tile_linear // GROUP_SPAN
            within_group: T.let = tile_linear % GROUP_SPAN
            tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
            tile_col: T.let = within_group // GROUP_SIZE
            set_tile_coords(tile_row, tile_col)
        elif TAIL_ROWS > 0:
            rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
            tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
            tile_col: T.let = rem // TAIL_ROWS
            set_tile_coords(tile_row, tile_col)
        else:
            set_tile_coords(0, 0)

    @T.inline
    def _update_serpentine(self, tile_linear, set_tile_coords):
        """CUTLASS-style 2D block swizzle with serpentine traversal.

        Algorithm:
        1. Divide grid into swizzle_size x swizzle_size blocks
        2. Within each block, visit tiles in row-major order
        3. Blocks are traversed column by column (along N)
        4. Within each column of blocks, use serpentine:
           - Even columns: top to bottom
           - Odd columns: bottom to top

        This maximizes L2 reuse for both A and B matrices.
        """
        S = T.meta_var(self._l2_group_size)  # swizzle_size
        M_BLOCKS = T.meta_var(self._M_BLOCKS)
        N_BLOCKS = T.meta_var(self._N_BLOCKS)
        BLOCK_SIZE = T.meta_var(self._BLOCK_SIZE)  # S * S
        FULL_BLOCK_TILES = T.meta_var(self._FULL_BLOCK_TILES)
        M_TILE_ROWS = T.meta_var(self._M_TILE_ROWS)
        T.meta_var(self._N_TILE_COLS)
        RESIDUAL_N = T.meta_var(self._RESIDUAL_N)
        RESIDUAL_M = T.meta_var(self._RESIDUAL_M)

        # Check if we're in the full block region
        if (M_BLOCKS > 0) & (N_BLOCKS > 0) & (tile_linear < FULL_BLOCK_TILES):
            # Which block (in linear order along columns of blocks)
            block_linear: T.let = tile_linear // BLOCK_SIZE
            within_block: T.let = tile_linear % BLOCK_SIZE

            # Block column and row
            block_col: T.let = block_linear // M_BLOCKS
            block_row_raw: T.let = block_linear % M_BLOCKS

            # Serpentine: odd columns go bottom-to-top
            block_row: T.let = T.Select(
                block_col % 2 == 0,
                block_row_raw,
                M_BLOCKS - 1 - block_row_raw
            )

            # Position within block (row-major within block)
            local_row: T.let = within_block // S
            local_col: T.let = within_block % S

            tile_row: T.let = block_row * S + local_row
            tile_col: T.let = block_col * S + local_col
            set_tile_coords(tile_row, tile_col)

        elif RESIDUAL_N > 0:
            # Residual tiles in the rightmost partial column of blocks
            # These are tiles where n >= N_BLOCKS * S
            rem: T.let = tile_linear - FULL_BLOCK_TILES

            # First handle the right residual strip (full M height, partial N width)
            right_strip_tiles: T.let = M_TILE_ROWS * RESIDUAL_N
            if rem < right_strip_tiles:
                # Row-major within the right strip
                tile_row: T.let = rem // RESIDUAL_N
                tile_col: T.let = N_BLOCKS * S + (rem % RESIDUAL_N)
                set_tile_coords(tile_row, tile_col)
            elif RESIDUAL_M > 0:
                # Bottom residual strip (already covered in right strip overlap)
                # This handles corner case - shouldn't normally reach here
                # as right strip already covers full M height
                set_tile_coords(0, 0)
            else:
                set_tile_coords(0, 0)

        elif RESIDUAL_M > 0:
            # Bottom residual strip only (no right residual)
            rem: T.let = tile_linear - FULL_BLOCK_TILES
            bottom_strip_tiles: T.let = RESIDUAL_M * (N_BLOCKS * S)
            if rem < bottom_strip_tiles:
                tile_row: T.let = M_BLOCKS * S + (rem % RESIDUAL_M)
                tile_col: T.let = rem // RESIDUAL_M
                set_tile_coords(tile_row, tile_col)
            else:
                set_tile_coords(0, 0)
        else:
            # Fallback
            set_tile_coords(0, 0)

    @T.inline
    def init(self, cluster_id):
        """Initialize scheduler for a given cluster.

        Parameters
        ----------
        cluster_id : int
            The cluster's index (typically cta_idx // CLUSTER_SIZE)
        """
        self.linear_idx = cluster_id
        self.tile_count = 0
        self.update_current_m_n_idx(cluster_id)

    @T.inline
    def next_tile(self):
        """Advance to the next tile for this cluster."""
        self.linear_idx = self.linear_idx + self._num_clusters
        self.tile_count = self.tile_count + 1
        self.update_current_m_n_idx(self.linear_idx)

    @T.inline
    def next_tile_stride(self, stride: int):
        """Advance by a custom stride (for non-standard scheduling)."""
        self.linear_idx = self.linear_idx + stride
        self.tile_count = self.tile_count + 1
        self.update_current_m_n_idx(self.linear_idx)
    # fmt: on

    def valid(self):
        """Check if this cluster has more tiles to process."""
        return self.linear_idx < self._TOTAL_TILES


class GroupMajor3D(BaseTileScheduler):
    """
    3D grouped-row scheduler (M,N,K) with tail handling on M.

    Args
    ----
    prefix: str
    m_tiles: int | T Expr   # tiles along M (static or runtime)
    n_tiles: int                # tiles along N (static)
    k_tiles: int                # tiles along K (static)
    group_rows: int             # rows per group along M
    step: int = 1               # default stride for next_tile()
    """

    def __init__(
        self, prefix: str, m_tiles, n_tiles: int, k_tiles: int, group_rows: int, step: int = 1
    ):
        super().__init__(prefix)
        self._step = step
        self.tile_idx = T.local_scalar("int32")
        self.k_idx = T.local_scalar("int32")

        # ---- constants / primexprs baked once ----
        self._G = group_rows
        self._N = n_tiles
        self._K = k_tiles

        if isinstance(m_tiles, int):
            self._GROUPS = m_tiles // group_rows
            self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
            self._SAFE_FINAL_ROWS = max(self._FINAL_ROWS, 1)
            self._GROUP_SIZE = group_rows * n_tiles * k_tiles
            self._TOTAL = m_tiles * n_tiles * k_tiles
        else:
            self._GROUPS = T.truncdiv(m_tiles, group_rows)
            self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
            self._SAFE_FINAL_ROWS = T.max(self._FINAL_ROWS, 1)
            self._GROUP_SIZE = self._G * self._N * self._K
            self._TOTAL = m_tiles * n_tiles * k_tiles

        # handy composites used in macro
        self._FULL_BOUND = self._GROUPS * self._GROUP_SIZE
        self._HAS_FULL = self._GROUPS > 0
        self._HAS_TAIL = self._FINAL_ROWS > 0

    # fmt: off
    @T.inline
    def update_current_m_n_idx(self, linear_idx):
        # full-group formulas
        full_m: T.let = T.floordiv(linear_idx, self._GROUP_SIZE) * self._G + T.floormod(
            linear_idx, self._G
        )
        full_n: T.let = T.floormod(T.floordiv(linear_idx, self._G), self._N)
        full_k: T.let = T.floordiv(T.floormod(linear_idx, self._GROUP_SIZE), self._G * self._N)

        # tail formulas (relative to FULL_BOUND)
        # Use _SAFE_FINAL_ROWS (max(FINAL_ROWS, 1)) to avoid divide-by-zero when there is no tail
        rem: T.let = linear_idx - self._FULL_BOUND
        tail_m: T.let = self._GROUPS * self._G + T.floormod(rem, self._SAFE_FINAL_ROWS)
        tail_n: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS) % self._N
        tail_k: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS * self._N)

        # choose phase
        if self._HAS_FULL & (linear_idx < self._FULL_BOUND):
            self.m_idx = full_m
            self.n_idx = full_n
            self.k_idx = full_k
        elif self._HAS_TAIL:
            self.m_idx = tail_m
            self.n_idx = tail_n
            self.k_idx = tail_k
        else:
            self.m_idx = 0
            self.n_idx = 0
            self.k_idx = 0

    @T.inline
    def init(self, linear_init):
        self.linear_idx = linear_init
        self.tile_idx = 0
        self.update_current_m_n_idx(linear_init)

    @T.inline
    def next_tile(self):
        self.linear_idx = self.linear_idx + self._step
        self.tile_idx = self.tile_idx + 1
        self.update_current_m_n_idx(self.linear_idx)

    @T.inline
    def next_tile_stride(self, stride: int):
        self.linear_idx = self.linear_idx + stride
        self.tile_idx = self.tile_idx + 1
        self.update_current_m_n_idx(self.linear_idx)
    # fmt: on

    def valid(self):
        return self.linear_idx < self._TOTAL


class RankAwareGroupMajorTileScheduler(BaseTileScheduler):
    """
    Group-major scheduler that applies a rank-aware remapping (remote rows first).
    Kept as a thin adapter because it depends on NVSHMEM rank at device-side.
    """

    def __init__(
        self, prefix: str, m_clusters: int, n_clusters: int, group_size: int, world_size: int
    ):
        super().__init__(prefix)
        self._m_clusters = m_clusters
        self._n_clusters = n_clusters
        self._group_size = group_size
        self._world_size = world_size

    @T.inline
    def update_current_m_n_idx(self, linear_idx):
        my_rank: T.let = T.nvshmem.my_pe()
        remote_m_clusters: T.let = self._m_clusters - self._m_clusters // self._world_size
        group_rows: T.let = (remote_m_clusters // self._group_size) * self._group_size
        final_rows: T.let = remote_m_clusters - group_rows
        group_repeat: T.let = self._group_size * self._n_clusters
        if linear_idx < group_rows * self._n_clusters and group_rows > 0:
            self.m_idx = (
                (linear_idx // group_repeat) * self._group_size
                + (linear_idx % self._group_size)
                + (my_rank + 1) * self._m_clusters // self._world_size
            ) % self._m_clusters
            self.n_idx = (linear_idx % group_repeat) // self._group_size
        elif linear_idx < remote_m_clusters * self._n_clusters:
            remainder_idx: T.let = linear_idx - group_rows * self._n_clusters
            self.m_idx = (
                group_rows
                + remainder_idx % final_rows
                + (my_rank + 1) * self._m_clusters // self._world_size
            ) % self._m_clusters
            self.n_idx = remainder_idx // final_rows
        else:
            remainder_idx: T.let = linear_idx - remote_m_clusters * self._n_clusters
            self.m_idx = (
                remote_m_clusters
                + remainder_idx % (self._m_clusters // self._world_size)
                + (my_rank + 1) * self._m_clusters // self._world_size
            ) % self._m_clusters
            self.n_idx = remainder_idx // (self._m_clusters // self._world_size)

    @T.inline
    def next_tile(self, stride: int):
        self.linear_idx = self.linear_idx + stride
        self.update_current_m_n_idx(self.linear_idx)

    def valid(self):
        return self.linear_idx < self._m_clusters * self._n_clusters


class IndexedTripleTileScheduler(BaseTileScheduler):
    """Scheduler that maps linear_idx to (b_idx, h_idx, q_idx) via index lists."""

    def __init__(self, prefix: str, b_indices, h_indices, q_indices, tiles_indptr):
        super().__init__(prefix)
        self.b_indices = b_indices
        self.h_indices = h_indices
        self.q_indices = q_indices
        self.tiles_indptr = tiles_indptr
        self.q_idx = T.local_scalar("int32")
        self.h_idx = T.local_scalar("int32")
        self.b_idx = T.local_scalar("int32")
        self.linear_lim = T.local_scalar("int32")

    @T.inline
    def _load(self):
        self.q_idx = self.q_indices[self.linear_idx]
        self.h_idx = self.h_indices[self.linear_idx]
        self.b_idx = self.b_indices[self.linear_idx]

    @T.inline
    def init(self, sm):
        self.linear_idx = self.tiles_indptr[sm]
        self.linear_lim = self.tiles_indptr[sm + 1]
        self._load()

    @T.inline
    def next_tile(self):
        self.linear_idx = self.linear_idx + 1
        self._load()

    def valid(self):
        return self.linear_idx < self.linear_lim


class FlashAttentionLinearScheduler(BaseTileScheduler):
    """Linear 3D scheduler for flash attention (batch, head, m_block).

    Used for non-causal attention with simple linear decomposition.
    Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:
        batch = linear_idx // (num_heads * num_m_blocks)
        head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks
        m_block = linear_idx % num_m_blocks

    Parameters
    ----------
    prefix : str
        Prefix for TIR variable names
    num_batches : int
        Number of batches
    num_heads : int
        Number of KV heads
    num_m_blocks : int
        Number of Q blocks (M dimension tiles)
    num_ctas : int
        Number of CTAs for persistent kernel stride
    """

    def __init__(
        self, prefix: str, num_batches: int, num_heads: int, num_m_blocks: int, num_ctas: int
    ):
        super().__init__(prefix)
        self._num_batches = num_batches
        self._num_heads = num_heads
        self._num_m_blocks = num_m_blocks
        self._num_ctas = num_ctas
        self._total_tasks = num_batches * num_heads * num_m_blocks

        # Output indices
        self.batch_idx = T.local_scalar("int32")
        self.head_idx = T.local_scalar("int32")
        self.m_block_idx = T.local_scalar("int32")

    # fmt: off
    @T.inline
    def update_current_m_n_idx(self, linear_idx):
        """Convert linear index to (batch, head, m_block) coordinates."""
        NUM_HEADS = T.meta_var(self._num_heads)
        NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
        HEAD_M_PRODUCT = T.meta_var(NUM_HEADS * NUM_M_BLOCKS)

        self.batch_idx = linear_idx // HEAD_M_PRODUCT
        self.head_idx = (linear_idx % HEAD_M_PRODUCT) // NUM_M_BLOCKS
        self.m_block_idx = linear_idx % NUM_M_BLOCKS

    @T.inline
    def init(self, cta_id):
        """Initialize scheduler with CTA ID."""
        self.linear_idx = cta_id
        self.update_current_m_n_idx(cta_id)

    @T.inline
    def next_tile(self):
        """Advance to next tile by striding by num_ctas."""
        self.linear_idx = self.linear_idx + self._num_ctas
        self.update_current_m_n_idx(self.linear_idx)
    # fmt: on

    def valid(self):
        """Check if there are more tiles to process."""
        return self.linear_idx < self._total_tasks


class FlashAttentionLPTScheduler(BaseTileScheduler):
    """LPT scheduler with L2 swizzle for causal flash attention.

    Processes high-work Q blocks (with more KV blocks to attend to) first using
    Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle
    for better cache locality across batch*head dimensions.

    The LPT aspect comes from reversing m_block order: lower Q blocks have more
    KV blocks to process due to causal masking, so processing them first balances load.

    The scheduler is only applied to non-persistent kernels.

    L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.

    Parameters
    ----------
    prefix : str
        Prefix for TIR variable names
    num_batches : int
        Number of batches
    num_heads : int
        Number of KV heads
    num_m_blocks : int
        Number of Q blocks (M dimension tiles)
    num_ctas : int
        Number of CTAs (should equal total_tasks for causal)
    l2_swizzle : int
        L2 swizzle factor for cache locality
    """

    def __init__(
        self,
        prefix: str,
        num_batches: int,
        num_heads: int,
        num_m_blocks: int,
        l2_swizzle: int,
        num_ctas: int | None = None,
    ):
        super().__init__(prefix)
        self._num_batches = num_batches
        self._num_heads = num_heads
        self._num_m_blocks = num_m_blocks
        self._l2_swizzle = l2_swizzle
        self._num_ctas = num_ctas
        self._total_tasks = num_batches * num_heads * num_m_blocks

        # Derived constants for L2 swizzle
        self._num_hb = num_batches * num_heads
        self._l2_major = l2_swizzle * num_m_blocks
        self._num_hb_quotient = self._num_hb // l2_swizzle

        # Output indices
        self.batch_idx = T.local_scalar("int32")
        self.head_idx = T.local_scalar("int32")
        self.m_block_idx = T.local_scalar("int32")

    # fmt: off
    @T.inline
    def update_current_m_n_idx(self, linear_idx):
        """Convert linear index to (batch, head, m_block) with LPT + L2 swizzle."""
        L2_SWIZZLE = T.meta_var(self._l2_swizzle)
        L2_MAJOR = T.meta_var(self._l2_major)
        NUM_HB_QUOTIENT = T.meta_var(self._num_hb_quotient)
        NUM_HB = T.meta_var(self._num_hb)
        NUM_HEADS = T.meta_var(self._num_heads)
        NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)

        # L2 swizzle decomposition
        bidhb: T.let = linear_idx // L2_MAJOR
        l2_mod: T.let = linear_idx % L2_MAJOR

        # Handle residual section (last partial swizzle group)
        num_hb_remainder: T.let = T.max(NUM_HB % L2_SWIZZLE, 1)
        m_block_raw: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod // L2_SWIZZLE, l2_mod // num_hb_remainder)  # noqa: E501
        bidhb_residual: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod % L2_SWIZZLE, l2_mod % num_hb_remainder)  # noqa: E501
        bidhb_actual: T.let = bidhb * L2_SWIZZLE + bidhb_residual

        self.batch_idx = bidhb_actual // NUM_HEADS
        self.head_idx = bidhb_actual % NUM_HEADS

        # LPT: Reverse block order so high-work blocks are processed first
        self.m_block_idx = (NUM_M_BLOCKS - 1) - m_block_raw

    @T.inline
    def init(self, cta_id):
        """Initialize scheduler with CTA ID."""
        self.linear_idx = cta_id
        self.update_current_m_n_idx(cta_id)

    @T.inline
    def next_tile(self):
        """Advance to the next tile.

        Single-tile mode (``num_ctas=None``, the default): each CTA owns one
        task; terminate. Persistent mode (``num_ctas=N``): stride by N, like
        :class:`FlashAttentionLinearScheduler`, while keeping the LPT + L2
        swizzle index mapping.
        """
        if self._num_ctas is None:
            self.linear_idx = self._total_tasks
        else:
            self.linear_idx = self.linear_idx + self._num_ctas
            self.update_current_m_n_idx(self.linear_idx)
    # fmt: on

    def valid(self):
        """Check if there are more tiles to process."""
        return self.linear_idx < self._total_tasks


class _CLCWorker(ClusterPersistentScheduler2D):
    """Per-role CLC handle: IS-A ClusterPersistentScheduler2D (so m_idx / n_idx work as
    usual) plus the role-local barrier phase and handshake. A coord-free role (e.g. an
    MMA warp consuming whatever a loader staged) arms the loop with reset() not init().
    """

    def __init__(self, clc, prefix):
        super().__init__(
            prefix,
            num_m_tiles=clc._num_m_tiles,
            num_n_tiles=clc._num_n_tiles,
            num_clusters=clc._num_m_tiles * clc._num_n_tiles,
            l2_group_size=clc._l2_group_size,
        )
        self._clc = clc
        self._sa = PipelineState(1, 0)
        self._done = T.local_scalar("int32")
        self._nxt = T.local_scalar("uint32")

    @T.inline
    def reset(self):
        self._done = 0

    @T.inline
    def init(self, cluster_id):
        # Explicit base call: TVMScript's parser has no zero-arg super().
        ClusterPersistentScheduler2D.init(self, cluster_id)
        self._done = 0

    def valid(self):
        return self._done == 0

    @T.inline
    def consume(self):
        # Single-elected-thread scope: wait for the handle, decode, release the slot.
        self._clc.sched_arr.full.wait(0, self._sa.phase)
        self._sa.advance()
        self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
        self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)

    @T.inline
    def consume_wg(self, wg_id, warp_id, lane_id):
        # Warpgroup scope: all threads decode; one elected lane releases the slot.
        self._clc.sched_arr.full.wait(0, self._sa.phase)
        self._sa.advance()
        self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
        T.cuda.warpgroup_sync(wg_id + 1)
        if (warp_id == 0) & (lane_id == 0):
            self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)

    @T.inline
    def advance_coords(self):
        if self._nxt != 0xFFFFFFFF:
            self.update_current_m_n_idx(self._nxt // self._clc._cta_group)

    @T.inline
    def mark_done_if_drained(self):
        if self._nxt == 0xFFFFFFFF:
            self._done = 1


@T.meta_class
class ClusterLaunchControlScheduler:
    """Blackwell Cluster Launch Control (CLC) tile scheduler.

    A scheduler warp runs ``run_scheduler`` (issues ``try_cancel`` to steal the next
    cluster); worker roles each take a ``worker()`` handle and pull the stolen tile
    through the shared smem handshake. Owns the CLC smem: the 16B response handle, the
    arrival barrier (handle ready), and the finished barrier (slot consumed;
    ``finish_arrivals`` arrivals per round). Tile-coord mapping is delegated to
    ``ClusterPersistentScheduler2D`` (group-major L2 ordering).
    """

    def __init__(self, pool, num_m_tiles, num_n_tiles, l2_group_size, cta_group, finish_arrivals):
        self._num_m_tiles = num_m_tiles
        self._num_n_tiles = num_n_tiles
        self._l2_group_size = l2_group_size
        self._cta_group = cta_group
        self.sched_arr = Pipeline(pool, 1, full="tma", empty="mbar", init_empty=1)
        self.sched_fin = Pipeline(pool, 1, full="mbar", empty="mbar", init_empty=finish_arrivals)
        self.clc_handle = pool.alloc((4,), "uint32", align=16)
        self._s_done = T.local_scalar("int32")
        self._s_nxt = T.local_scalar("uint32")

    def worker(self, prefix):
        return _CLCWorker(self, prefix)

    @T.inline
    def run_scheduler(self, cbx):
        # cta0 drives try_cancel; both CTAs expect_bytes + consume the handle so the
        # finished-barrier count is met and the slot can be reissued.
        if T.ptx.elect_sync():
            sa = PipelineState(1, 0)
            sf = PipelineState(1, 1)
            self._s_done = 0
            while self._s_done == 0:
                if cbx == 0:
                    self.sched_fin.empty.wait(0, sf.phase)
                    sf.advance()
                    T.ptx.clc_try_cancel(
                        T.address_of(self.clc_handle[0]), T.address_of(self.sched_arr.full.buf[0])
                    )
                self.sched_arr.full.arrive(0, 16)  # expect_bytes for the 16B handle
                self.sched_arr.full.wait(0, sa.phase)
                sa.advance()
                self._s_nxt = T.ptx.clc_query_cancel(T.address_of(self.clc_handle[0]))
                self.sched_fin.empty.arrive(0, cta_id=0, pred=True)
                if self._s_nxt == 0xFFFFFFFF:
                    self._s_done = 1
