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

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl
from typing import Dict, Tuple
from task import input_t, output_t

@triton.jit
def silu_mul_kernel(
    gate_ptr, up_ptr, weight_ptr, out_ptr,
    M, N,
    stride_gm, stride_gn, stride_um, stride_un, 
    stride_w, stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
):
    pid = tl.program_id(0)
    grid_n = tl.cdiv(N, BLOCK_N)
    pid_m, pid_n = pid // grid_n, pid % grid_n

    m_offs = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    
    mask = (m_offs[:, None] < M) & (n_offs[None, :] < N)

    # Prefetch weights - they're reused across all columns
    w = tl.load(weight_ptr + m_offs * stride_w, mask=m_offs < M, other=0.0)

    gate = tl.load(gate_ptr + m_offs[:, None] * stride_gm + n_offs[None, :] * stride_gn, mask=mask, other=0.0)
    up = tl.load(up_ptr + m_offs[:, None] * stride_um + n_offs[None, :] * stride_un, mask=mask, other=0.0)
    
    # Micro-optimized SiLU computation
    gate_clamped = tl.clamp(gate, -9.0, 9.0)
    gate_sigmoid = tl.sigmoid(gate_clamped)
    result = gate * gate_sigmoid * up * w[:, None]
    
    tl.store(out_ptr + m_offs[:, None] * stride_om + n_offs[None, :] * stride_on, result, mask=mask)


def silu_mul(gate_out, up_out, weight):
    M, N = gate_out.shape
    out = torch.empty_like(gate_out)
    
    # Micro-tuned for your dimensions (N=2048, varying M)
    if N >= 2048:
        BLOCK_M, BLOCK_N = 32, 128  # Wide blocks for N=2048
        num_warps = 8
    elif N >= 1024:
        BLOCK_M, BLOCK_N = 64, 64
        num_warps = 4
    else:
        BLOCK_M, BLOCK_N = 64, 32
        num_warps = 4
    
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
    
    silu_mul_kernel[grid](
        gate_out, up_out, weight, out, M, N,
        gate_out.stride(0), gate_out.stride(1),
        up_out.stride(0), up_out.stride(1),
        weight.stride(0), out.stride(0), out.stride(1),
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, num_warps=num_warps
    )
    return out


# Keep everything else EXACTLY as your working version
def fused_expert_forward(x, W_gate, W_up, W_down, weight):
    return silu_mul(x @ W_gate, x @ W_up, weight) @ W_down

@triton.jit
def softmax_only_kernel(
    logits_ptr, probs_ptr,
    N, D,
    stride_n, stride_d,
    BLOCK_SIZE: tl.constexpr,
):
    row_id = tl.program_id(0)
    if row_id >= N:
        return
    
    row_start = row_id * stride_n
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < D
    
    logits = tl.load(logits_ptr + row_start + cols * stride_d, mask=mask, other=0.0)
    
    # Exact PyTorch softmax
    row_max = tl.max(tl.where(mask, logits, -float('inf')), axis=0)
    shifted = logits - row_max
    exp_vals = tl.exp(shifted)
    exp_vals = tl.where(mask, exp_vals, 0.0)
    exp_sum = tl.sum(exp_vals, axis=0)
    probs = exp_vals / exp_sum
    
    tl.store(probs_ptr + row_start + cols * stride_d, probs, mask=mask)


@triton.jit
def topk_only_kernel(
    probs_ptr, values_ptr, indices_ptr,
    N, D, top_k,
    stride_n, stride_d,
    BLOCK_SIZE: tl.constexpr,
):
    row_id = tl.program_id(0)
    if row_id >= N:
        return
    
    row_start = row_id * stride_n
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < D
    
    probs = tl.load(probs_ptr + row_start + cols * stride_d, mask=mask, other=0.0)
    
    # Create working copy for selection
    working_probs = tl.where(mask, probs, -1.0)
    
    # Vectorized topk selection - unrolled for top_k=4 (common in MoE)
    # K=0
    if top_k > 0:
        max_val = tl.max(working_probs, axis=0)
        is_max = (working_probs == max_val) & (working_probs >= 0.0)
        # Find minimum index among tied maximums (PyTorch behavior)
        selected_idx = tl.min(tl.where(is_max, cols, D), axis=0)
        
        tl.store(values_ptr + row_id * top_k + 0, max_val)
        tl.store(indices_ptr + row_id * top_k + 0, selected_idx)
        
        # Mark selected element
        working_probs = tl.where(cols == selected_idx, -1.0, working_probs)
    
    # K=1
    if top_k > 1:
        max_val = tl.max(working_probs, axis=0)
        is_max = (working_probs == max_val) & (working_probs >= 0.0)
        selected_idx = tl.min(tl.where(is_max, cols, D), axis=0)
        
        tl.store(values_ptr + row_id * top_k + 1, max_val)
        tl.store(indices_ptr + row_id * top_k + 1, selected_idx)
        
        working_probs = tl.where(cols == selected_idx, -1.0, working_probs)
    
    # K=2
    if top_k > 2:
        max_val = tl.max(working_probs, axis=0)
        is_max = (working_probs == max_val) & (working_probs >= 0.0)
        selected_idx = tl.min(tl.where(is_max, cols, D), axis=0)
        
        tl.store(values_ptr + row_id * top_k + 2, max_val)
        tl.store(indices_ptr + row_id * top_k + 2, selected_idx)
        
        working_probs = tl.where(cols == selected_idx, -1.0, working_probs)
    
    # K=3
    if top_k > 3:
        max_val = tl.max(working_probs, axis=0)
        is_max = (working_probs == max_val) & (working_probs >= 0.0)
        selected_idx = tl.min(tl.where(is_max, cols, D), axis=0)
        
        tl.store(values_ptr + row_id * top_k + 3, max_val)
        tl.store(indices_ptr + row_id * top_k + 3, selected_idx)


def triton_softmax_topk(logits: torch.Tensor, top_k: int) -> Tuple[torch.Tensor, torch.Tensor]:
    if top_k > 4:
        raise ValueError("This kernel supports top_k <= 4")
        
    N, D = logits.shape
    
    # Intermediate probabilities tensor
    probs = torch.empty_like(logits)
    values = torch.empty((N, top_k), dtype=logits.dtype, device=logits.device)
    indices = torch.empty((N, top_k), dtype=torch.long, device=logits.device)
    
    BLOCK_SIZE = triton.next_power_of_2(D)
    grid = (N,)
    
    # Step 1: Softmax
    softmax_only_kernel[grid](
        logits, probs,
        N, D,
        logits.stride(0), logits.stride(1),
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    # Step 2: TopK
    topk_only_kernel[grid](
        probs, values, indices,
        N, D, top_k,
        probs.stride(0), probs.stride(1),
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    return values, indices


def fused_gating_topk(x: torch.Tensor, W_g: torch.Tensor, top_k: int) -> Tuple[torch.Tensor, torch.Tensor]:
    logits = x @ W_g
    return triton_softmax_topk(logits, top_k)

@triton.jit
def count_experts_kernel(
    topk_indices_ptr, counts_ptr,
    N, top_k, num_experts,
    BLOCK_SIZE: tl.constexpr,
):
    """First pass: count occurrences of each expert"""
    pid = tl.program_id(0)
    offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    total_tokens = N * top_k
    mask = offset < total_tokens
    
    token_id = offset // top_k
    k_id = offset % top_k
    
    expert_idx = tl.load(topk_indices_ptr + token_id * top_k + k_id, mask=mask, other=num_experts)
    valid_mask = mask & (expert_idx < num_experts)
    
    tl.atomic_add(counts_ptr + expert_idx, 1, mask=valid_mask)


@triton.jit
def scatter_sorted_kernel(
    topk_scores_ptr, topk_indices_ptr,
    sorted_expert_indices_ptr, sorted_token_indices_ptr, sorted_expert_scores_ptr,
    counts_ptr, offsets_ptr,
    N, top_k, num_experts,
    BLOCK_SIZE: tl.constexpr,
):
    """Second pass: scatter to sorted positions"""
    pid = tl.program_id(0)
    offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    total_tokens = N * top_k
    mask = offset < total_tokens
    
    token_id = offset // top_k
    k_id = offset % top_k
    
    expert_idx = tl.load(topk_indices_ptr + token_id * top_k + k_id, mask=mask, other=0)
    scores = tl.load(topk_scores_ptr + token_id * top_k + k_id, mask=mask, other=0.0)
    
    valid_mask = mask & (expert_idx < num_experts)
    
    # Get next available position for this expert
    local_pos = tl.atomic_add(counts_ptr + expert_idx, 1, mask=valid_mask)
    expert_offset = tl.load(offsets_ptr + expert_idx, mask=valid_mask, other=0)
    final_pos = expert_offset + local_pos
    
    # Write to sorted positions
    tl.store(sorted_expert_indices_ptr + final_pos, expert_idx, mask=valid_mask)
    tl.store(sorted_token_indices_ptr + final_pos, token_id, mask=valid_mask)
    tl.store(sorted_expert_scores_ptr + final_pos, scores, mask=valid_mask)


def triton_fused_routing_prep(topk_scores: torch.Tensor, topk_indices: torch.Tensor, 
                             N: int, top_k: int, num_experts: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Ultra-optimized two-pass implementation"""
    total_tokens = N * top_k
    
    # Output tensors
    sorted_expert_indices = torch.empty(total_tokens, dtype=torch.long, device=device)
    sorted_token_indices = torch.empty(total_tokens, dtype=torch.long, device=device)
    sorted_expert_scores = torch.empty(total_tokens, dtype=topk_scores.dtype, device=device)
    counts = torch.zeros(num_experts, dtype=torch.long, device=device)
    
    # Optimal block size for your hardware
    BLOCK_SIZE = 256
    grid_size = triton.cdiv(total_tokens, BLOCK_SIZE)
    
    # Pass 1: Count experts
    count_experts_kernel[(grid_size,)](
        topk_indices, counts,
        N, top_k, num_experts,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=8,
    )
    
    # Compute offsets (CPU is faster for small arrays)
    offsets = torch.zeros_like(counts)
    offsets[1:] = counts[:-1].cumsum(0)
    
    # Reset counts for second pass
    temp_counts = torch.zeros_like(counts)
    
    # Pass 2: Scatter to sorted positions
    scatter_sorted_kernel[(grid_size,)](
        topk_scores, topk_indices,
        sorted_expert_indices, sorted_token_indices, sorted_expert_scores,
        temp_counts, offsets,
        N, top_k, num_experts,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=8,
    )
    
    return sorted_expert_indices, sorted_token_indices, sorted_expert_scores, counts


def fused_routing_prep(topk_scores: torch.Tensor, topk_indices: torch.Tensor, 
                      N: int, top_k: int, num_experts: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    return triton_fused_routing_prep(topk_scores, topk_indices, N, top_k, num_experts, device)


@torch.jit.script
def fused_output_aggregation(output: torch.Tensor, expert_output: torch.Tensor, token_indices: torch.Tensor, D: int) -> torch.Tensor:
    output.scatter_add_(0, token_indices.unsqueeze(1).expand(-1, D), expert_output)
    return output


class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.num_experts = config["n_routed_experts"]
        self.top_k = config["n_experts_per_token"]
        
        self.router = None
        self.shared_gate = None
        self.shared_up = None
        self.shared_down = None
        self.gate = [None] * self.num_experts
        self.up = [None] * self.num_experts
        self.down = [None] * self.num_experts

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        N = B * T
        x_flat = x.view(-1, D)

        # Keep optimized BLAS operations
        shared_weights = torch.ones(N, device=x.device, dtype=x.dtype)
        shared_out = silu_mul(x_flat @ self.shared_gate, x_flat @ self.shared_up, shared_weights) @ self.shared_down

        # Routing
        topk_scores, topk_indices = fused_gating_topk(x_flat, self.router, self.top_k)
        expert_indices, token_indices, expert_scores, counts = fused_routing_prep(
            topk_scores, topk_indices, N, self.top_k, self.num_experts, x.device
        )
        
        cumsum = counts.cumsum(0).cpu().numpy()
        output = torch.zeros_like(x_flat)

        # Expert processing
        for i in range(self.num_experts):
            start = 0 if i == 0 else cumsum[i - 1]
            end = cumsum[i]
            if start == end:
                continue

            tokens = token_indices[start:end]
            weights = expert_scores[start:end]
            inputs = x_flat[tokens]

            out = fused_expert_forward(inputs, self.gate[i], self.up[i], self.down[i], weights)
            output = fused_output_aggregation(output, out, tokens, D)

        return (shared_out + output).view(B, T, D)


def custom_kernel(data: input_t) -> output_t:
    x, weights, config = data
    moe = MoE(config)
    moe.router = weights["router.weight"].t()
    moe.shared_gate = weights["shared_experts.0.weight"]
    moe.shared_up = weights["shared_experts.1.weight"]
    moe.shared_down = weights["shared_experts.2.weight"]
    for i in range(config["n_routed_experts"]):
        moe.gate[i] = weights[f"experts.{i}.0.weight"]
        moe.up[i] = weights[f"experts.{i}.1.weight"]
        moe.down[i] = weights[f"experts.{i}.2.weight"]
    return moe(x)
python eval.py benchmark /tmp/tmpjxxj8yhy