"""Task30 interleaved_rope: direct row-owned exact copy in one Triton launch.

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

改编说明:
- 改动了 docstring 与 author
- 保持 r16 全部 国产 NPU 套路不变
- 函数末尾保留 `reference = interleaved_rope`(平台 import 入口)
"""
from numbers import Integral

import torch
import triton
import triton.language as tl


@triton.jit
def _interleave_tiles(
    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,
    BLOCK: tl.constexpr,
):
    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))

    row = tl.program_id(1).to(tl.int64)
    col = (tl.program_id(0).to(tl.int64) * BLOCK
           + tl.arange(0, BLOCK).to(tl.int64))
    valid = col < D
    safe_col = tl.where(valid, col, 0)
    phase = safe_col % 3
    from_b = (phase == 1) & (safe_col < H_END)
    from_c = (phase == 2) & (safe_col < W_END)
    base = row * ROW_STRIDE + safe_col * COL_STRIDE

    a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0)
    b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0)
    c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0)
    value = tl.where(from_b, b, tl.where(from_c, c, a))
    tl.store(output_ptr + row * D + col, value, valid)


def interleaved_rope(x, mrope_section):
    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()))
    grid = (triton.cdiv(d, block), rows)
    with torch.get_device_module(x.device).device(x.device):
        _interleave_tiles[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),
            BLOCK=block, num_warps=4, enable_fp_fusion=False,
        )
    return output


reference = interleaved_rope