#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpu MI300
from task import input_t, output_t
from reference import check_implementation, generate_input
import math
import torch
import torch.nn.functional as F
from typing import Dict, Tuple
import triton
import triton.language as tl
from collections import defaultdict

# Triton kernel for matrix multiplication (unchanged)
@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k_loop in range(0, K, BLOCK_SIZE_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    
    c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    tl.store(c_ptrs, acc.to(tl.float16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

def triton_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    M, K = a.shape
    _K, N = b.shape
    if K != _K:
        raise ValueError(f"Matrix dimensions mismatch for matmul: A.shape[1]={K}, B.shape[0]={_K}")
    c = torch.empty((M, N), device=a.device, dtype=torch.float16)
    
    BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 32
    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
    
    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K
    )
    return c

# Triton kernel for SiLU activation (unchanged)
@triton.jit
def silu_kernel(
    x_ptr, y_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    
    mask = offsets < n_elements
    x_val = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    y_val = x_val / (1.0 + tl.exp(-x_val))
    tl.store(y_ptr + offsets, y_val.to(tl.float16), mask=mask)

def triton_silu(x: torch.Tensor) -> torch.Tensor:
    n_elements = x.numel()
    output = torch.empty_like(x)
    
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
    
    silu_kernel[grid](
        x, output,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE
    )
    return output

def custom_kernel(data: input_t) -> output_t:
    input_tensor_arg, weights_dict, config_dict = data
    
    d_hidden = config_dict["d_hidden"]
    d_expert = config_dict["d_expert"]
    n_routed_experts = config_dict["n_routed_experts"]
    n_experts_per_token = config_dict["n_experts_per_token"]
    batch_size = config_dict["batch_size"]
    seq_len = config_dict["seq_len"]
    num_experts = n_routed_experts
    
    x_input = input_tensor_arg.view(-1, d_hidden)
    
    # Router weights are [out_features, in_features], so transpose for X @ W_g
    W_g = weights_dict['router.weight'].t()
    logits = triton_matmul(x_input, W_g)
    scores = F.softmax(logits, dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=n_experts_per_token, dim=-1)
    topk_scores_fp16 = topk_scores.to(torch.float16)

    flat_expert_indices = topk_indices.view(-1)
    flat_expert_weights = topk_scores_fp16.view(-1)
    
    idxs_argsort = flat_expert_indices.argsort()
    token_idxs_dispatch = idxs_argsort // n_experts_per_token

    counts = flat_expert_indices.bincount(minlength=num_experts)
    tokens_per_expert_cumsum = counts.cumsum(dim=0)
    
    total_tokens = batch_size * seq_len
    expert_outputs = torch.zeros((total_tokens, n_experts_per_token, d_hidden), device=x_input.device, dtype=x_input.dtype)
    
    for expert_id_loop in range(num_experts):
        start_idx = 0 if expert_id_loop == 0 else tokens_per_expert_cumsum[expert_id_loop - 1].item()
        end_idx = tokens_per_expert_cumsum[expert_id_loop].item()
        
        if start_idx == end_idx:
            continue
        
        original_idxs = idxs_argsort[start_idx:end_idx]
        k = original_idxs % n_experts_per_token
        exp_token_idxs = token_idxs_dispatch[start_idx:end_idx]
        
        expert_tokens = x_input[exp_token_idxs]
        
        W_gate = weights_dict[f'experts.{expert_id_loop}.0.weight']
        W_up = weights_dict[f'experts.{expert_id_loop}.1.weight']
        W_down = weights_dict[f'experts.{expert_id_loop}.2.weight']
        
        gate = triton_silu(triton_matmul(expert_tokens, W_gate))
        up = triton_matmul(expert_tokens, W_up)
        intermediate = gate * up
        expert_out_unweighted = triton_matmul(intermediate, W_down)
        
        expert_outputs[exp_token_idxs, k, :] = expert_out_unweighted
    
    routed_output = (expert_outputs * topk_scores_fp16.unsqueeze(-1)).sum(dim=1)
    
    # Shared Expert Computation
    W_gate_shared = weights_dict['shared_experts.0.weight']
    W_up_shared = weights_dict['shared_experts.1.weight']
    W_down_shared = weights_dict['shared_experts.2.weight']
    
    gate_shared = triton_silu(triton_matmul(x_input, W_gate_shared))
    up_shared = triton_matmul(x_input, W_up_shared)
    intermediate_shared = gate_shared * up_shared
    shared_output = triton_matmul(intermediate_shared, W_down_shared)
    
    if shared_output.shape[-1] != d_hidden:
        if shared_output.shape[-1] < d_hidden:
            padding_amount = d_hidden - shared_output.shape[-1]
            shared_output = F.pad(shared_output, (0, padding_amount), "constant", 0.0)
        else:
            shared_output = shared_output[..., :d_hidden]
            
    output_combined = routed_output + shared_output
    return output_combined.view(batch_size, seq_len, d_hidden)
python eval.py benchmark /tmp/tmpyeawqm79