"""Task29 gelu_and_mul: 跨芯片通用版(v8.1 国际卡写法 + 国产 NPU 套路).

作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓
版权: 2026 GuanghuLab
基础参考: v8.1(国际 A/B) + 国产 NPU 套路(r16 M-RoPE 学的)

设计选择:
- 用 `tl.math.erf`(国际卡 SFU 加速: NVIDIA __nv_erf / AMD __ocml_erf 硬件实现)
- 2D grid + BLOCK=1024 + num_warps=4(国际卡最优)
- `enable_fp_fusion=False`(国产 NPU 必需)
- `with torch.get_device_module(x.device).device(x.device):`(国产 NPU 必需)
- 函数末尾加 `reference = gelu_and_mul`(平台 import 入口)
"""
import torch
import triton
import triton.language as tl


@triton.jit
def _gelu_and_mul_kernel(x_ptr, out_ptr, d, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    col_block = tl.program_id(1)
    d2 = 2 * d
    col = col_block * BLOCK + tl.arange(0, BLOCK)
    mask = col < d

    gate = tl.load(x_ptr + row * d2 + col, mask=mask)
    up = tl.load(x_ptr + row * d2 + col + d, mask=mask)
    gate_f = gate.to(tl.float32)
    up_f = up.to(tl.float32)
    gelu = gate_f * 0.5 * (1.0 + tl.math.erf(gate_f * 0.7071067811865476))
    out = gelu * up_f
    tl.store(out_ptr + row * d + col, out.to(x_ptr.dtype.element_ty), mask=mask)


def gelu_and_mul(hidden_states):
    if hidden_states.device.type in ('cpu', 'meta', 'mps'):
        raise RuntimeError('a real Triton accelerator backend is required')
    d = hidden_states.shape[-1] // 2
    num_rows = hidden_states.numel() // (2 * d)
    out = torch.empty_like(hidden_states[..., :d])

    BLOCK = 1024
    grid = (num_rows, triton.cdiv(d, BLOCK))

    module = torch.get_device_module(hidden_states.device)
    with module.device(hidden_states.device):
        _gelu_and_mul_kernel[grid](
            hidden_states, out, d,
            BLOCK=BLOCK, num_warps=4, enable_fp_fusion=False,
        )
    return out


reference = gelu_and_mul
