#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpu MI300

import torch
import torch.nn.functional as F
from torch import Tensor

from task import input_t, output_t


MlpWeight = tuple[Tensor, Tensor, Tensor]


def mlp(x: Tensor, w: MlpWeight):
    return (F.silu(x @ w[0]) * (x @ w[1])) @ w[2]


# scores, indices
def moe_gate(x: Tensor, w: Tensor, k: int):
    return (x @ w.T).softmax(dim=-1).topk(k, dim=-1, sorted=False)


def moe_experts(
    x: torch.Tensor,
    flat_expert_indices: torch.Tensor,
    flat_expert_weights: torch.Tensor,
    experts: list[MlpWeight],
    n_experts_per_token: int,
):
    expert_cache = torch.zeros_like(x)
    idxs = flat_expert_indices.argsort()
    counts = flat_expert_indices.bincount().cpu().numpy()
    tokens_per_expert = counts.cumsum()
    token_idxs = idxs // n_experts_per_token

    for expert_id, end_idx in enumerate(tokens_per_expert):
        start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
        if start_idx == end_idx:
            continue

        exp_token_idxs = token_idxs[start_idx:end_idx]
        expert_out = mlp(x[exp_token_idxs], experts[expert_id])
        expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
        expert_cache.scatter_reduce_(0, exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]), expert_out, reduce="sum")

    return expert_cache


def moe(
    x: torch.Tensor,
    w_router: Tensor,
    experts: list[MlpWeight],
    shared_experts: MlpWeight,
    n_experts_per_token: int,
):
    expert_scores, expert_indices = moe_gate(x, w_router, n_experts_per_token)

    batch_size, seq_len, hidden_dim = x.shape
    orig_shape = x.shape
    x_flat = x.view(-1, hidden_dim)
    flat_expert_indices = expert_indices.view(-1)
    flat_expert_weights = expert_scores.view(-1, 1)
    routed_output_flat = moe_experts(
        x_flat,
        flat_expert_indices,
        flat_expert_weights,
        experts,
        n_experts_per_token,
    )
    routed_output = routed_output_flat.view(*orig_shape)

    shared_output = mlp(x, shared_experts)
    return routed_output + shared_output


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data

    output = moe(
        input_tensor,
        weights["router.weight"],
        [
            (
                weights[f"experts.{i}.0.weight"].T.contiguous().T,
                weights[f"experts.{i}.1.weight"].T.contiguous().T,
                weights[f"experts.{i}.2.weight"].T.contiguous().T,
            )
            for i in range(config["n_routed_experts"])
        ],
        (
            weights["shared_experts.0.weight"].T.contiguous().T,
            weights["shared_experts.1.weight"].T.contiguous().T,
            weights["shared_experts.2.weight"].T.contiguous().T,
        ),
        config["n_experts_per_token"],
    )

    return output

python eval.py benchmark /tmp/tmp0n0roqw_