# SPDX-License-Identifier: Apache-2.0
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/hc_collapse.py,
# commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified.
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
"""`hc_collapse` — `HyperConnectCollapse`, the mHC `hc_pre`: the `hc` parallel copies of the
residual stream are mixed back down into the single hidden state a sublayer reads. One fp32
multiply-accumulate over the copies, no parameters. Whole module on one core; no collective.
"""

import nki
import nki.isa as nisa
import nki.language as nl

PMAX = 128

#: Free-axis chunk candidates, widest first. 1024 bf16 is 2 KiB per partition.
CHUNKS = (1024, 512, 256, 128, 64, 32)


def _chunk(dim):
    for cand in CHUNKS:
        if dim % cand == 0:
            return cand
    return dim


@nki.jit
def kernel(input_0, input_1):
    """`sum(pre_mix[..., c, None] * x[..., c, :] for c in range(hc))`."""
    batch, n_tok, hc, dim = input_0.shape
    rows = batch * n_tok
    chunk = _chunk(dim)

    y_out = nl.ndarray((batch, n_tok, dim), dtype=input_0.dtype, buffer=nl.shared_hbm)

    x3 = input_0.reshape((rows, hc, dim))
    pm2 = input_1.reshape((rows, hc))
    y2 = y_out.reshape((rows, dim))

    def body(row0, nr):
        pm = nl.load(pm2[nl.ds(row0, nr), 0:hc])

        for c in nl.static_range(dim // chunk):
            f0 = c * chunk

            acc = nl.ndarray((nr, chunk), dtype=nl.float32, buffer=nl.sbuf)
            xg = nl.ndarray((nr, hc, chunk), dtype=input_0.dtype, buffer=nl.sbuf)

            nisa.dma_copy(dst=xg, src=x3[nl.ds(row0, nr), 0:hc, nl.ds(f0, chunk)],
                          dge_mode=nisa.dge_mode.none)

            for j in nl.static_range(hc):
                if j == 0:
                    nisa.tensor_scalar(dst=acc, data=xg[:, j, :], op0=nl.multiply,
                                       operand0=pm[:, j:j + 1])
                else:
                    nisa.scalar_tensor_tensor(dst=acc, data=xg[:, j, :], op0=nl.multiply,
                                              operand0=pm[:, j:j + 1],
                                              op1=nl.add, operand1=acc)

            nisa.dma_copy(dst=y2[nl.ds(row0, nr), nl.ds(f0, chunk)], src=acc,
                          dge_mode=nisa.dge_mode.none)

    full = rows // PMAX
    tail = rows - full * PMAX
    for t in nl.affine_range(full):
        body(t * PMAX, PMAX)
    if tail:
        body(full * PMAX, tail)

    return y_out