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

# reference
# https://github.com/sgl-project/sglang/blob/v0.4.6/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py

import triton
import triton.language as tl
import torch
import torch.nn.functional as F
from torch import Tensor

from task import input_t, output_t

torch._dynamo.config.recompile_limit = 10000
MlpWeight = tuple[Tensor, Tensor, Tensor]


@triton.jit
def mlp_up_kernel(
    x_ptr,   # (M, K)
    w0_ptr,  # (K, N)
    w1_ptr,  # (K, N)
    o_ptr,   # (M, N)
    M,
    N,
    K,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # TODO: might be worth to concat w0 and w1, if M is large

    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)

    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    off_m = pid_m * BLOCK_M
    off_n = pid_n * BLOCK_N

    offs_m = off_m + tl.arange(0, BLOCK_M)
    offs_n = off_n + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    # if M % BLOCK_M == 0:
    #     offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_M), BLOCK_M)
    offs_m %= M
    if N % BLOCK_N == 0:
        offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_N), BLOCK_N)

    x_ptrs = x_ptr + (offs_m[:, None] * K + offs_k[None, :])
    w0_ptrs = w0_ptr + (offs_k[:, None] * N + offs_n[None, :])
    w1_ptrs = w1_ptr + (offs_k[:, None] * N + offs_n[None, :])

    acc0 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    acc1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    k_iters = K // BLOCK_K
    for _ in range(0, k_iters):
        x = tl.load(x_ptrs)

        w0 = tl.load(w0_ptrs)
        acc0 = tl.dot(x, w0, acc0)
        w1 = tl.load(w1_ptrs)
        acc1 = tl.dot(x, w1, acc1)

        x_ptrs += BLOCK_K
        w0_ptrs += BLOCK_K * N
        w1_ptrs += BLOCK_K * N

    acc = tl.sigmoid(acc0) * (acc0 * acc1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]
    mask = (offs_m < M) & (offs_n < N)
    o_ptrs = o_ptr + (offs_m * N + offs_n)
    tl.store(o_ptrs, acc, mask=mask)


seen_shapes = set()


def mlp_up(x: Tensor, w0: Tensor, w1: Tensor):
    M, K = x.shape
    _, N = w0.shape

    BLOCK_M = 64
    BLOCK_N = 32
    BLOCK_K = 64
    kwargs = dict(
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        waves_per_eu=2,
        kpack=1,
        matrix_instr_nonkdim=16,
        num_warps=8,
        num_stages=2,
    )
    num_blocks = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N)
    out = x.new_empty((M, N))
    mlp_up_kernel[(num_blocks,)](x, w0, w1, out, M, N, K, **kwargs)
    return out


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

    out = mlp_up(x, w[0], w[1])
    out = out @ w[2]
    return out


# 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

    # must make this CUDA-graph compatible
    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])

        # this can go to epilogue-fusion with down-projection
        expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])

        # can this also go with epilogue-fusion? requires atomic add -> just use a separate kernel
        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_flat, shared_experts).view(x.shape)

    return routed_output + shared_output


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"],
                weights[f"experts.{i}.1.weight"],
                weights[f"experts.{i}.2.weight"],
            )
            for i in range(config["n_routed_experts"])
        ],
        (
            weights["shared_experts.0.weight"],
            weights["shared_experts.1.weight"],
            weights["shared_experts.2.weight"],
        ),
        config["n_experts_per_token"],
    )

    return output

python eval.py benchmark /tmp/tmprkmw4gao