r"""
Uses a group gemm, instead of a loop over experts.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, List, Optional
import math
import einops
import triton
import triton.language as tl

def is_cuda():
    return triton.runtime.driver.active.get_current_target().backend == "cuda"

def num_sms():
    if is_cuda():
        return torch.cuda.get_device_properties("cuda").multi_processor_count
    return 148

@triton.autotune(
    configs=[
        triton.Config({
            'BLOCK_SIZE_M': 128,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 84,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 128,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 128,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 64,
            'BLOCK_SIZE_N': 64,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 84,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 64,
            'BLOCK_SIZE_N': 64,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 128,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 128,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 64,
            'NUM_SM': num_sms(),
        }),
        triton.Config({
            'BLOCK_SIZE_M': 64,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 64,
            'NUM_SM': num_sms(),
        }),
    ],
    key=['group_size'],
)
@triton.jit
def grouped_matmul_kernel(
    # device tensor of matrices pointers
    group_a_ptrs,
    group_b_ptrs,
    group_c_ptrs,
    # device tensor of gemm sizes. its shape is [group_size, 3]
    # dim 0 is group_size, dim 1 is the values of <M, N, K> of each gemm
    group_gemm_sizes,
    # device tensor of leading dimension sizes. its shape is [group_size, 3]
    # dim 0 is group_size, dim 1 is the values of <lda, ldb, ldc> of each gemm
    g_lds,
    # number of gemms
    group_size,
    # number of virtual SM
    NUM_SM: tl.constexpr,
    # tile sizes
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
):
    r"""
    This is taken and modified from the triton tutorial. 
      - adds masking to get working with different shapes.
    """
    tile_idx = tl.program_id(0)
    last_problem_end = 0
    for g in range(group_size):
        # get the gemm size of the current problem
        gm = tl.load(group_gemm_sizes + g * 3)
        gn = tl.load(group_gemm_sizes + g * 3 + 1)
        gk = tl.load(group_gemm_sizes + g * 3 + 2)
        num_m_tiles = tl.cdiv(gm, BLOCK_SIZE_M)
        num_n_tiles = tl.cdiv(gn, BLOCK_SIZE_N)
        num_tiles = num_m_tiles * num_n_tiles
        # iterate through the tiles in the current gemm problem
        while (tile_idx >= last_problem_end and tile_idx < last_problem_end + num_tiles):
            # pick up a tile from the current gemm problem
            k = gk
            lda = tl.load(g_lds + g * 3)
            ldb = tl.load(g_lds + g * 3 + 1)
            ldc = tl.load(g_lds + g * 3 + 2)
            a_ptr = tl.load(group_a_ptrs + g).to(tl.pointer_type(tl.float16))
            b_ptr = tl.load(group_b_ptrs + g).to(tl.pointer_type(tl.float16))
            c_ptr = tl.load(group_c_ptrs + g).to(tl.pointer_type(tl.float16))
            # figure out tile coordinates
            tile_idx_in_gemm = tile_idx - last_problem_end
            tile_m_idx = tile_idx_in_gemm // num_n_tiles
            tile_n_idx = tile_idx_in_gemm % num_n_tiles

            # do regular gemm here
            offs_am = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
            offs_bn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
            offs_k = tl.arange(0, BLOCK_SIZE_K)
            a_ptrs = a_ptr + offs_am[:, None] * lda + offs_k[None, :]
            b_ptrs = b_ptr + offs_k[:, None] * ldb + offs_bn[None, :]
            
            m_mask = offs_am < gm
            n_mask = offs_bn < gn
            
            accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
            for kk in range(0, tl.cdiv(k, BLOCK_SIZE_K)):
                # hint to Triton compiler to do proper loop pipelining
                #tl.multiple_of(a_ptrs, [16, 16])
                #tl.multiple_of(b_ptrs, [16, 16])
                
                k_mask = kk + offs_k < gk
                a_mask = m_mask[:, None] and k_mask[None, :]
                b_mask = k_mask[:, None] and n_mask[None, :]
                
                a = tl.load(a_ptrs, mask=a_mask)
                b = tl.load(b_ptrs, mask=b_mask)
                accumulator += tl.dot(a, b)
                a_ptrs += BLOCK_SIZE_K
                b_ptrs += BLOCK_SIZE_K * ldb
            c = accumulator.to(tl.float16)

            offs_cm = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
            offs_cn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
            c_ptrs = c_ptr + ldc * offs_cm[:, None] + offs_cn[None, :]

            c_mask = m_mask[:, None] and n_mask[None, :]
            tl.store(c_ptrs, c, mask=c_mask)

            # go to the next tile by advancing NUM_SM
            tile_idx += NUM_SM

        # get ready to go to the next gemm problem
        last_problem_end = last_problem_end + num_tiles

def group_gemm_fn(group_A: List[torch.Tensor], group_B: List[torch.Tensor]):
    assert len(group_A) == len(group_B)
    DEVICE = 'cuda'
    group_size = len(group_A)

    A_addrs = []
    B_addrs = []
    C_addrs = []
    g_sizes = []
    g_lds = []
    group_C = []
    for i in range(group_size):        
        A = group_A[i]
        B = group_B[i]
        assert A.shape[1] == B.shape[0]
        M, K = A.shape
        K, N = B.shape
        
        C = torch.empty((M, N), device=DEVICE, dtype=A.dtype)
        group_C.append(C)
        A_addrs.append(A.data_ptr())
        B_addrs.append(B.data_ptr())
        C_addrs.append(C.data_ptr())
        g_sizes += [M, N, K]
        g_lds += [A.stride(0), B.stride(0), C.stride(0)]
        
    # note these are device tensors
    d_a_ptrs = torch.tensor(A_addrs, device=DEVICE)
    d_b_ptrs = torch.tensor(B_addrs, device=DEVICE)
    d_c_ptrs = torch.tensor(C_addrs, device=DEVICE)
    d_g_sizes = torch.tensor(g_sizes, dtype=torch.int32, device=DEVICE)
    d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=DEVICE)
    # we use a fixed number of CTA, and it's auto-tunable
    grid = lambda META: (META['NUM_SM'], )
    grouped_matmul_kernel[grid](
        d_a_ptrs,
        d_b_ptrs,
        d_c_ptrs,
        d_g_sizes,
        d_g_lds,
        group_size,
    )

    return group_C

def get_token_indices(expert_indices, num_experts_per_token):
    flat_expert_indices = expert_indices.view(-1)
    indices = flat_expert_indices.argsort()
    counts = flat_expert_indices.bincount().cpu().numpy() # TODO: why move to cpu?
    tokens_per_expert_range = counts.cumsum()
    token_indices = indices // num_experts_per_token
    return indices, tokens_per_expert_range, token_indices

def expert_input_permute(
    tokens: torch.Tensor, 
    expert_indices: torch.Tensor, 
    num_experts_per_token: int
) -> List[torch.Tensor]:
    _, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    
    expert_data = []
    for expert_id, end_idx in enumerate(tokens_per_expert_range):
        start_idx = 0 if expert_id == 0 else tokens_per_expert_range[expert_id - 1]
        if start_idx == end_idx:
            # triton group gemm needs len(group_A) == len(group_B),
            # adding an empty list is just a hack to ensure the group sizes are the
            # same. This could/should also be done by not passing in weights if the
            # corresponding expert gets no tokens.
            expert_data.append(torch.empty((0, tokens.size(-1))))
            continue
        
        exp_token_idxs = token_indices[start_idx:end_idx]
        expert_tokens = tokens[exp_token_idxs]
        expert_data.append(expert_tokens)
        
    return expert_data, tokens_per_expert_range

def expert_output_permute(
    expert_data: List[torch.Tensor], 
    expert_indices: torch.Tensor, 
    expert_weights: torch.Tensor, 
    num_experts_per_token: int,
    output_shape: Tuple[int]
) -> torch.Tensor:
    hidden_dim = expert_data[0].size(-1)
    flat_expert_weights = expert_weights.view(-1, 1)

    indices, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    
    output = torch.zeros(output_shape, device=expert_data[0].device, dtype=expert_data[0].dtype)
    for expert_id, end_idx in enumerate(tokens_per_expert_range):
        start_idx = 0 if expert_id == 0 else tokens_per_expert_range[expert_id - 1]
        if start_idx == end_idx:
            continue
        
        # multiply token by by expert weights and reduce
        expert_data[expert_id].mul_(flat_expert_weights[indices[start_idx:end_idx]])
        exp_token_idxs = token_indices[start_idx:end_idx]
        output.scatter_reduce_(
            0,
            exp_token_idxs.view(-1, 1).repeat(1, hidden_dim),
            expert_data[expert_id],
            reduce='sum'
        )
    
    return output

@torch.compile
class GRUSharedExpert(nn.Module):
    def __init__(self, d_hidden, d_expert):
        super().__init__()
        self.act_fn = nn.SiLU()
        self.W_gate = nn.Linear(d_hidden, d_expert, bias=False)
        self.W_up = nn.Linear(d_hidden, d_expert, bias=False)
        self.W_down = nn.Linear(d_expert, d_hidden, bias=False)
        
    def forward(self, x):
        gate = self.act_fn(self.W_gate(x))
        up = self.W_up(x)
        out = self.W_down(gate * up)
        return out
    
class GRURoutedExperts(nn.Module):
    def __init__(self, num_experts, d_hidden, d_expert):
        super().__init__()
        
        self.act_fn = nn.SiLU() # TODO: fuse SiLU with group gemm
        
        self.register_buffer('W_gate', torch.empty(num_experts, d_hidden, d_expert))
        self.register_buffer('W_up', torch.empty(num_experts, d_hidden, d_expert))
        self.register_buffer('W_down', torch.empty(num_experts, d_expert, d_hidden))

    def forward(self, group_token):
        # TODO: are these copies or a list of views?
        group_gate = [w for w in self.W_gate]
        group_up = [w for w in self.W_up]
        group_down = [w for w in self.W_down]
        
        gate = [self.act_fn(t) for t in group_gemm_fn(group_token, group_gate)]
        up = group_gemm_fn(group_token, group_up)
        lhs = [t_gate * t_up for (t_gate, t_up) in zip(gate, up)]
        out = group_gemm_fn(lhs, group_down)
        
        return out

@torch.compile
class MoEGate(nn.Module):
    def __init__(self, num_experts, d_hidden, topk):
        super().__init__()
        self.top_k = topk
        self.num_experts = num_experts
        self.d_hidden = d_hidden
        self.W_g = nn.Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.W_g(x)
        scores = logits.softmax(dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        return topk_scores, topk_indices

class MoEGroupGEMM(nn.Module):
    def __init__(self, num_experts, d_expert, d_hidden, topk):
        super().__init__()
        self.topk = topk
        self.shared_expert = GRUSharedExpert(d_hidden, d_expert)
        self.routed_experts = GRURoutedExperts(num_experts, d_hidden, d_expert)
        self.router = MoEGate(num_experts, d_hidden, self.topk)
        
    def forward(self, tokens):
        batch_size, seq_len, hidden_size = tokens.size()
        tokens_flat = tokens.view(-1, hidden_size)
        
        shared_output = self.shared_expert(tokens_flat)
        expert_scores, expert_indices = self.router(tokens_flat)
        
        # get tokens for each expert, but compute expert transforms in parallel with group gemms.
        group_token, _ = expert_input_permute(tokens_flat, expert_indices, self.topk)
        group_token = self.routed_experts(group_token)
        routed_output = expert_output_permute(group_token, expert_indices, expert_scores, self.topk, tokens_flat.shape)
        
        output = shared_output + routed_output
        return output.view(batch_size, seq_len, hidden_size)
        
def custom_kernel(data):
    """
    Reference implementation of DeepSeek-style Mixture of Experts using PyTorch.
    
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_dim]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    d_expert = config["d_expert"]
    d_hidden = config["d_hidden"]
    topk = config["n_experts_per_token"]
    moe = MoEGroupGEMM(num_experts, d_expert, d_hidden, topk).cuda().to(torch.float16)

    for i in range(num_experts):
        gate_proj_weight = weights[f'experts.{i}.0.weight']
        up_proj_weight = weights[f'experts.{i}.1.weight']
        down_proj_weight = weights[f'experts.{i}.2.weight']

        # weights not transposed, since normal tensor
        moe.routed_experts.W_gate[i] = gate_proj_weight
        moe.routed_experts.W_up[i] = up_proj_weight
        moe.routed_experts.W_down[i] = down_proj_weight
    
    # Transpose to match layout expected by linear
    moe.router.W_g.weight = nn.Parameter(weights['router.weight'])
    moe.shared_expert.W_gate.weight = nn.Parameter(weights['shared_experts.0.weight'].t())
    moe.shared_expert.W_up.weight = nn.Parameter(weights['shared_experts.1.weight'].t())
    moe.shared_expert.W_down.weight = nn.Parameter(weights['shared_experts.2.weight'].t())

    output = moe(input_tensor)

    return output


# Input generation for the reference code

def generate_input(
    dhidden: int,
    dexpert: int,
    nroutedexperts: int,
    nsharedexperts: int,
    nexpertspertoken: int,
    bs: int,
    seqlen: int,
    seed: int
):

    # Really dumb but for now _ isn't parsing correctly.
    d_hidden = dhidden
    d_expert = dexpert
    n_routed_experts = nroutedexperts
    n_shared_experts = nsharedexperts
    n_experts_per_token = nexpertspertoken
    batch_size = bs
    seq_len = seqlen

    config = {
        "d_hidden": d_hidden,
        "d_expert": d_expert,
        "n_routed_experts": n_routed_experts,
        "n_shared_experts": n_shared_experts,
        "n_experts_per_token": n_experts_per_token,
        "batch_size": batch_size,
        "seq_len": seq_len,
    }

    gen = torch.Generator(device='cuda')
    gen.manual_seed(seed)

    num_experts = n_routed_experts
    expert_dim = d_expert
    weights = {}

    input_tensor = torch.randn(
        (batch_size, seq_len, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ).contiguous()

    # Initialize router weights
    weights['router.weight'] = torch.randn(
        (num_experts, d_hidden),
        device="cuda",
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    for i in range(num_experts):
        weights[f'experts.{i}.0.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.1.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.2.weight'] = torch.randn(
            (expert_dim, d_hidden),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(d_hidden)
    
    weights['shared_experts.0.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.1.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.2.weight'] = torch.randn(
        (expert_dim * n_shared_experts, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    return (input_tensor, weights, config)

if __name__ == '__main__':
    input = generate_input(128, 128, 4, 1, 2, 8, 8, 0)
    output = custom_kernel(input)
python eval.py benchmark /tmp/tmpho97ouiw