"""Task35 rotary_embedding: fixed pair ownership and fused float32 rotation.

作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓
版权: 2026 GuanghuLab
基础参考: r6-四Head共享角度流水线(队友共享)

改编:
- 改 docstring 与 author
- 加 `enable_fp_fusion=False`(国产 NPU 套路,从 r16 M-RoPE 学)
- 函数末尾保留 `reference = rotary_embedding`(平台 import 入口)
"""
import torch
import triton
import triton.language as tl


@triton.jit
def _rotary_pairs(
    x_ptr, cos_ptr, sin_ptr, out_ptr,
    H: tl.constexpr, D: tl.constexpr, HALF: tl.constexpr,
    XTS: tl.constexpr, XHS: tl.constexpr, XDS: tl.constexpr,
    CTS: tl.constexpr, CDS: tl.constexpr,
    STS: tl.constexpr, SDS: tl.constexpr,
    OTS: tl.constexpr, OHS: tl.constexpr, ODS: tl.constexpr,
    BLOCK: tl.constexpr,
):
    row = tl.program_id(0)
    tile = tl.program_id(1)
    token = row // H
    head = row - token * H
    pair = tile * BLOCK + tl.arange(0, BLOCK)
    valid = pair < HALF

    x_base = token * XTS + head * XHS
    x_even = tl.load(x_ptr + x_base + (2 * pair) * XDS,
                     mask=valid, other=0.0).to(tl.float32)
    x_odd = tl.load(x_ptr + x_base + (2 * pair + 1) * XDS,
                    mask=valid, other=0.0).to(tl.float32)
    cosine = tl.load(cos_ptr + token * CTS + pair * CDS,
                     mask=valid, other=0.0).to(tl.float32)
    sine = tl.load(sin_ptr + token * STS + pair * SDS,
                   mask=valid, other=0.0).to(tl.float32)

    out_even = x_even * cosine - x_odd * sine
    out_odd = x_even * sine + x_odd * cosine
    out_base = token * OTS + head * OHS
    tl.store(out_ptr + out_base + (2 * pair) * ODS, out_even, mask=valid)
    tl.store(out_ptr + out_base + (2 * pair + 1) * ODS, out_odd, mask=valid)


def _block_size(half):
    if half <= 32:
        return 32
    if half <= 64:
        return 64
    if half <= 128:
        return 128
    return 256


def rotary_embedding(x, cos, sin, interleaved):
    if x.ndim != 3:
        raise ValueError('x must have shape [T, H, D]')
    if not isinstance(interleaved, bool) or interleaved:
        raise ValueError('Task35 requires interleaved=False')
    if x.device.type in ('cpu', 'meta', 'mps'):
        raise RuntimeError('a real Triton accelerator backend is required')
    tokens, heads, width = x.shape
    if width % 2:
        raise ValueError('the last dimension must be even')
    half = width // 2
    if cos.ndim != 2 or tuple(cos.shape) != (tokens, half):
        raise ValueError('cos must have shape [T, D//2]')
    if sin.ndim != 2 or tuple(sin.shape) != (tokens, half):
        raise ValueError('sin must have shape [T, D//2]')
    if cos.device != x.device or sin.device != x.device:
        raise ValueError('x, cos and sin must be on the same device')

    output = torch.empty(x.shape, dtype=x.dtype, device=x.device)
    if x.numel() == 0:
        return output
    block = _block_size(half)
    grid = (tokens * heads, triton.cdiv(half, block))
    module = torch.get_device_module(x.device)
    with module.device(x.device):
        _rotary_pairs[grid](
            x, cos, sin, output,
            H=heads, D=width, HALF=half,
            XTS=x.stride(0), XHS=x.stride(1), XDS=x.stride(2),
            CTS=cos.stride(0), CDS=cos.stride(1),
            STS=sin.stride(0), SDS=sin.stride(1),
            OTS=output.stride(0), OHS=output.stride(1), ODS=output.stride(2),
            BLOCK=block, num_warps=4, enable_fp_fusion=False,
        )
    return output


reference = rotary_embedding
