"""Task30 interleaved_rope v2: 多行合一 + 简化 mask 优化

v1 (commit d1eeaf3) → v2 改动:
- ROWS_PER_PROGRAM = 2: 每个 program 处理 2 行, 减少 launch 次数
- 2D tile [ROWS_PER_PROGRAM, BLOCK]: 提升 cache locality
- a + b + c 替代 3 路 where: 利用 mask 互斥, 减少 ALU
- from_a 显式计算: 编译器更好优化

不破坏 v1 的所有正确性套路(指针 cast / num_warps / enable_fp_fusion)

作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓
版权: 2026 GuanghuLab
基础参考: r16-整条取料流水线-待平台验证(队友共享实现)
"""
from numbers import Integral

import torch
import triton
import triton.language as tl


@triton.jit
def _interleave_tiles_v2(
    input_ptr, output_ptr,
    N: tl.constexpr, D: tl.constexpr,
    PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr, COL_STRIDE: tl.constexpr,
    H_END: tl.constexpr, W_END: tl.constexpr,
    ROWS_PER_PROGRAM: tl.constexpr,  # 新增
    BLOCK: tl.constexpr,
):
    # 指针 cast 防 NaN
    if input_ptr.dtype.element_ty.primitive_bitwidth == 16:
        input_ptr = input_ptr.to(tl.pointer_type(tl.int16))
        output_ptr = output_ptr.to(tl.pointer_type(tl.int16))
    elif input_ptr.dtype.element_ty.primitive_bitwidth == 32:
        input_ptr = input_ptr.to(tl.pointer_type(tl.int32))
        output_ptr = output_ptr.to(tl.pointer_type(tl.int32))

    # 1D grid: (col_blocks, rows / ROWS_PER_PROGRAM)
    pid_col = tl.program_id(0).to(tl.int64)
    pid_row_grp = tl.program_id(1).to(tl.int64)

    # row 维度: ROWS_PER_PROGRAM 行
    rows_in_grp = tl.arange(0, ROWS_PER_PROGRAM)  # [RPP]
    row = pid_row_grp * ROWS_PER_PROGRAM + rows_in_grp  # [RPP]
    num_rows = N // D
    row_mask = row < num_rows  # [RPP]

    # col 维度: BLOCK 个
    col = pid_col * BLOCK + tl.arange(0, BLOCK).to(tl.int64)  # [BLOCK]
    col_mask = col < D  # [BLOCK]

    # 2D base [RPP, BLOCK]
    base = row[:, None] * ROW_STRIDE + col[None, :] * COL_STRIDE

    # phase + mask 显式互斥
    safe_col = tl.where(col_mask, col, 0)
    phase = safe_col % 3
    from_b = (phase == 1) & col_mask & (safe_col < H_END)  # [BLOCK]
    from_c = (phase == 2) & col_mask & (safe_col < W_END)  # [BLOCK]
    # from_a = ~from_b & ~from_c & col_mask (phase 0 或超界)
    from_a = col_mask & ~from_b & ~from_c  # [BLOCK]

    # 3 路 load(互斥 mask)
    a = tl.load(input_ptr + base, mask=row_mask[:, None] & from_a[None, :], other=0)
    b = tl.load(input_ptr + PLANE_STRIDE + base, mask=row_mask[:, None] & from_b[None, :], other=0)
    c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, mask=row_mask[:, None] & from_c[None, :], other=0)

    # 利用 mask 互斥: a + b + c 在每个 (row, col) 上只有 1 个非 0
    # 比 3 路 where 少一次 where 选择
    value = a + b + c

    # 写
    out_ptr = row[:, None] * D + col[None, :]
    tl.store(output_ptr + out_ptr, value, mask=row_mask[:, None] & col_mask[None, :])


def interleaved_rope(x, mrope_section, ROWS_PER_PROGRAM=2):
    if x.ndim != 3 or x.shape[0] != 3:
        raise ValueError('x must have shape [3, S, D]')
    if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3
            or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)):
        raise ValueError('mrope_section must contain three nonnegative integers')
    if x.device.type in ('cpu', 'meta', 'mps'):
        raise RuntimeError('a real Triton accelerator backend is required')
    _, rows, d = x.shape
    output = torch.empty((rows, d), dtype=x.dtype, device=x.device)
    if not rows or not d:
        return output

    block = min(1024, 1 << max(5, (d - 1).bit_length()))
    num_rows_grp = (rows + ROWS_PER_PROGRAM - 1) // ROWS_PER_PROGRAM
    grid = (triton.cdiv(d, block), num_rows_grp)
    with torch.get_device_module(x.device).device(x.device):
        _interleave_tiles_v2[grid](
            x, output, N=rows * d, D=d,
            PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1),
            COL_STRIDE=x.stride(2),
            H_END=min(d, int(mrope_section[1]) * 3),
            W_END=min(d, int(mrope_section[2]) * 3),
            ROWS_PER_PROGRAM=ROWS_PER_PROGRAM,
            BLOCK=block, num_warps=4, enable_fp_fusion=False,
        )
    return output


reference = interleaved_rope
