import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Pack tokens into expert-major order for MoE execution.

    Each token is assigned to one expert and one slot within that expert's
    buffer. The kernel writes token activations into the packed dispatch layout.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        token_hidden: torch.Tensor,
        expert_idx: torch.Tensor,
        slot_idx: torch.Tensor,
        expert_offsets: torch.Tensor,
    ) -> torch.Tensor:
        num_rows = int(expert_offsets[-1].item())
        out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)
        for token in range(token_hidden.shape[0]):
            expert = int(expert_idx[token].item())
            row = int(expert_offsets[expert].item() + slot_idx[token].item())
            out[row] = token_hidden[token]
        return out


num_tokens = 2048
hidden_dim = 128
num_experts = 16


def get_inputs():
    return popcorn_pri.moe_dispatch_inputs(num_tokens, hidden_dim, num_experts)




def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = """
__global__ void scatter_tokens_kernel(
    const float* __restrict__ token_hidden,
    const int* __restrict__ expert_idx,
    const int* __restrict__ slot_idx,
    const int* __restrict__ expert_offsets,
    float* __restrict__ out,
    int num_tokens,
    int hidden_dim
) {
    int token = blockIdx.x;
    int col = threadIdx.x;
    
    if (token < num_tokens && col < hidden_dim) {
        int expert = expert_idx[token];
        int row = expert_offsets[expert] + slot_idx[token];
        out[row * hidden_dim + col] = token_hidden[token * hidden_dim + col];
    }
}

torch::Tensor scatter_tokens_cuda(
    torch::Tensor token_hidden,
    torch::Tensor expert_idx,
    torch::Tensor slot_idx,
    torch::Tensor expert_offsets
) {
    TORCH_CHECK(token_hidden.is_contiguous(), "token_hidden must be contiguous");
    TORCH_CHECK(expert_idx.is_contiguous(), "expert_idx must be contiguous");
    TORCH_CHECK(slot_idx.is_contiguous(), "slot_idx must be contiguous");
    TORCH_CHECK(expert_offsets.is_contiguous(), "expert_offsets must be contiguous");
    
    int num_tokens = token_hidden.size(0);
    int hidden_dim = token_hidden.size(1);
    int num_rows = expert_offsets[expert_offsets.size(0) - 1].item<int>();
    
    auto out = torch::zeros({num_rows, hidden_dim}, token_hidden.options());
    
    dim3 grid(num_tokens);
    dim3 block(hidden_dim);
    
    scatter_tokens_kernel<<<grid, block>>>(
        token_hidden.data_ptr<float>(),
        expert_idx.data_ptr<int>(),
        slot_idx.data_ptr<int>(),
        expert_offsets.data_ptr<int>(),
        out.data_ptr<float>(),
        num_tokens,
        hidden_dim
    );
    
    return out;
}
"""

cpp_source = """
torch::Tensor scatter_tokens_cuda(
    torch::Tensor token_hidden,
    torch::Tensor expert_idx,
    torch::Tensor slot_idx,
    torch::Tensor expert_offsets
);
"""

module = load_inline(
    name='scatter_tokens_ext5',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['scatter_tokens_cuda'],
    verbose=True,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = module
    
    def forward(
        self,
        token_hidden: torch.Tensor,
        expert_idx: torch.Tensor,
        slot_idx: torch.Tensor,
        expert_offsets: torch.Tensor,
    ) -> torch.Tensor:
        return self.module.scatter_tokens_cuda(
            token_hidden.contiguous(),
            expert_idx.contiguous(),
            slot_idx.contiguous(),
            expert_offsets.contiguous()
        )