"""Mirrors ``tilelang_ops/asc/gather.py`` (``gather_tl``): bytewise row gather with zero-fill."""

from collections.abc import Callable

import torch


def _gather_rows(x: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
    """``x[index]`` for a byte table ``x [n, D]``.

    On the neuron device, index_select / advanced indexing of an int8 table allocates several times the table
    (+4.9 GiB for the 1.3 GB V4.1 embedding, which ran the full model out of HBM), while F.embedding through a
    4- or 2-byte view copies only the selected rows and is bit-exact (verified on arbitrary bytes, NaN patterns
    included).
    """
    if x.is_contiguous():
        for elem, dtype in ((4, torch.int32), (2, torch.float16)):
            if x.shape[1] % elem == 0 and (torch.compiler.is_compiling() or x.storage_offset() % elem == 0):
                return torch.nn.functional.embedding(index, x.view(dtype)).view(torch.int8)
    return torch.index_select(x, 0, index)


def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]:
    """Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous.

    ``y[i] = x[index[i]]`` if ``0 <= index[i] < n`` else 0; written in place.
    """

    def _gather(x: torch.Tensor, index: torch.Tensor, y: torch.Tensor) -> None:
        n = x.shape[0]
        m = index.shape[0]
        if m == 0:
            return
        if n == 0:
            y.zero_()
            return
        # No scalar tensor operands: on the neuron device a host scalar is a host->device copy, which drains the
        # stream.
        rows = _gather_rows(x[:, :D], index.clamp(0, n - 1))
        y.copy_(rows.masked_fill(((index < 0) | (index >= n)).view(m, 1), 0))

    return _gather