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
from dataclasses import dataclass

@dataclass
class GroupedTokens:
    tokens: torch.Tensor
    tokens_per_expert_range: torch.tensor
    tokens_per_expert_range_cpu: torch.tensor
    indices: torch.Tensor
    token_gather_indices: torch.Tensor

@triton.jit
def silu(tile):
    # tl.exp needs float32 or float64.
    return tile * (1.0 / (1.0 + tl.exp(-tile)))

def num_sms():
    # MI300 has 304 CUs
    return 2 * 304

configs = [
    # M, N, K, warps, stages
    # for AMD, num stages should be 1 or 2.
    (128, 128, 128, 4, 2),
    (128, 128, 64, 4, 2),
    (128, 128, 32, 4, 2),
    (256, 256, 64, 4, 2),
    (64, 128, 128, 4, 2),
    (128, 128, 128, 8, 2),
    (256, 256, 64, 8, 2),
    (64, 64, 128, 4, 2),
    (64, 64, 256, 4, 2),
    (128, 64, 256, 4, 2)
]

configs = [
    triton.Config({
        'BLOCK_SIZE_M': c[0],
        'BLOCK_SIZE_N': c[1],
        'BLOCK_SIZE_K': c[2],
        'num_warps': c[3],
        'num_stages': c[4],
        'NUM_SM': num_sms()
    }) for c in configs
]

@triton.autotune(
    configs=configs,
    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,
    ACTIVATION: tl.constexpr,
    # 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,
):
    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
            
            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])
                                
                a = tl.load(a_ptrs, mask=m_mask[:, None])
                b = tl.load(b_ptrs)
                accumulator = tl.dot(a, b, acc=accumulator)
                a_ptrs += BLOCK_SIZE_K
                b_ptrs += BLOCK_SIZE_K * ldb
                
            if ACTIVATION:
                accumulator = silu(accumulator)
                
            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, :]

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

            # go to the next tile by advancing NUM_SM
            tile_idx += tl.num_programs(axis=0)

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

def get_a_addrs(group_A: GroupedTokens, group_size: int):
    a_addrs = []
    a_shapes = []
    for i in range(group_size):
        end_idx = group_A.tokens_per_expert_range_cpu[i]
        start_idx = 0 if i == 0 else group_A.tokens_per_expert_range_cpu[i - 1]
        a_addrs.append(group_A.tokens[start_idx:-1].data_ptr())
        M = end_idx - start_idx
        K = group_A.tokens.size(-1)
        a_shapes.append((M, K))
    return a_addrs, a_shapes
        
def get_b_addrs(group_B: List[torch.Tensor], group_size: int):
    b_addrs = []
    b_shapes = []
    for i in range(group_size):
        b_addrs.append(group_B[i].data_ptr())
        K = group_B[i].size(0)
        N = group_B[i].size(1)
        b_shapes.append((K, N))
    return b_addrs, b_shapes

def group_gemm_fn(group_A: GroupedTokens, group_B: List[torch.Tensor], activation):
    DEVICE=group_B[0].device
    
    group_size = len(group_B)
    
    A_addrs, A_shapes = get_a_addrs(group_A, group_size)
    B_addrs, B_shapes = get_b_addrs(group_B, group_size)
    
    C_addrs = []
    g_sizes = []
    g_lds = []
        
    c = torch.empty((group_A.tokens.size(0), group_B[0].size(-1)), device=DEVICE, dtype=group_A.tokens.dtype)
    
    c_offset = 0
    for i in range(group_size):
        assert A_shapes[i][1] == B_shapes[i][0]
        M, K = A_shapes[i]
        K, N = B_shapes[i]
        
        C_addrs.append(c[c_offset:].data_ptr()) # using [c_offset:] indexing to get view
        g_sizes += [M, N, K]
        g_lds += [K, N, c.stride(0)]
        c_offset += M
                
    # 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)

    # invoke like normal kernel???    
    grid = lambda META: (triton.cdiv(c.size(0), META['BLOCK_SIZE_M']) * 
                         triton.cdiv(c.size(1), META['BLOCK_SIZE_N']), )
    grouped_matmul_kernel[grid](
        d_a_ptrs,
        d_b_ptrs,
        d_c_ptrs,
        d_g_sizes,
        d_g_lds,
        group_size,
        ACTIVATION=activation
    )
        
    return GroupedTokens(
        tokens=c,
        tokens_per_expert_range=group_A.tokens_per_expert_range,
        tokens_per_expert_range_cpu=group_A.tokens_per_expert_range_cpu,
        indices=group_A.indices,
        token_gather_indices=group_A.token_gather_indices,
    )

@torch.compile
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()
    tokens_per_expert_range = counts.cumsum(dim=0)
    token_indices = indices // num_experts_per_token
    return indices, tokens_per_expert_range, token_indices

@torch.compile
def expert_input_permute(
    tokens: torch.Tensor, 
    expert_indices: torch.Tensor, 
    num_experts_per_token: int
) -> GroupedTokens:
    indices, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    token_dim = tokens.size(1)
    output = torch.gather(tokens, dim=0, index=token_indices.unsqueeze(1).expand(-1, token_dim))
    return GroupedTokens(
        tokens=output,
        tokens_per_expert_range=tokens_per_expert_range,
        tokens_per_expert_range_cpu=tokens_per_expert_range.detach().cpu(),
        indices=indices, 
        token_gather_indices=token_indices,
    )

@torch.compile
def expert_output_permute(
    grouped_tokens: GroupedTokens,
    expert_scores: torch.Tensor,
    num_experts_per_token: int,
    output: torch.Tensor
) -> torch.Tensor:
    grouped_tokens.tokens.mul_(expert_scores.view(-1, 1)[grouped_tokens.indices])
    output.scatter_reduce_(
        0,
        # index tensor doesn't broadcast, so repeat indices along hidden dim
        grouped_tokens.token_gather_indices.view(-1, 1).expand(-1, output.size(-1)),
        grouped_tokens.tokens,
        reduce='sum'
    )
    return output

@torch.compile
def compute_shared_expert(weights, tokens):
    w_gate = weights['shared_experts.0.weight']
    w_up = weights['shared_experts.1.weight']
    w_down = weights['shared_experts.2.weight']
    gate = F.silu(tokens @ w_gate)
    up = tokens @ w_up
    down = (gate * up) @ w_down
    return down

@torch.compile
def compute_routing(weights, tokens, topk):
    # For some reason, router weights are already transposed?? So transpose here...
    w_g = weights['router.weight'].t()
    logits = tokens @ w_g
    scores = logits.softmax(dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=topk, dim=-1, sorted=False)
    return topk_scores, topk_indices

def compute_routed_experts(
        weights, 
        tokens, 
        topk, 
        num_experts,
        expert_indices,
        expert_scores,
        output
    ):
    gate_proj_weights = [weights[f'experts.{i}.0.weight'] for i in range(num_experts)]
    up_proj_weights = [weights[f'experts.{i}.1.weight'] for i in range(num_experts)]
    down_proj_weights = [weights[f'experts.{i}.2.weight'] for i in range(num_experts)]
    
    # get tokens for each expert, but compute expert transforms in parallel with group gemms.
    group_tokens: GroupedTokens = expert_input_permute(tokens, expert_indices, topk)

    gate: GroupedTokens = group_gemm_fn(group_tokens, gate_proj_weights, activation=True)
    
    # this fuses gate * (token @ W_up)
    up: GroupedTokens = group_gemm_fn(group_tokens, up_proj_weights, activation=False)
    
    up_gated = GroupedTokens(
        tokens=gate.tokens * up.tokens,
        tokens_per_expert_range=gate.tokens_per_expert_range,
        tokens_per_expert_range_cpu=gate.tokens_per_expert_range_cpu,
        indices=gate.indices,
        token_gather_indices=gate.token_gather_indices,
    )
    
    group_token: GroupedTokens = group_gemm_fn(up_gated, down_proj_weights, activation=False)

    # routed experts get accumulated onto output
    expert_output_permute(
        group_token, 
        expert_scores, 
        topk, 
        output)

def custom_kernel(data):
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    topk = config["n_experts_per_token"]

    hidden_size = input_tensor.size(2)
    original_shape = input_tensor.size()
    
    flat_tokens = input_tensor.view(-1, hidden_size)

    shared_output = compute_shared_expert(weights, flat_tokens)
    scores, indices = compute_routing(weights, flat_tokens, topk)

    compute_routed_experts(
        weights,
        flat_tokens,
        topk,
        num_experts,
        indices,
        scores,
        shared_output
    )

    return shared_output.view(original_shape)

# 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 leaderboard /tmp/tmpkeslt1ii