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

import os
import sys
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

os.environ["CXX"] = "clang++"

cuda_src = r"""
#include <iostream>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/amd_detail/amd_hip_fp8.h>

__host__ __device__ __forceinline__ constexpr int cdiv(int a, int b) {
  return (a + b - 1) / b;
}

#define float16_t __fp16

__global__ void fp16_mm_ref_kernel(const float16_t *A, const float16_t *B,
                                   float16_t *C, uint32_t N, uint32_t K,
                                   uint32_t M) {
  uint32_t row = blockIdx.y * blockDim.y + threadIdx.y;
  uint32_t col = blockIdx.x * blockDim.x + threadIdx.x;

  if (row >= N || col >= M)
    return;

  float sum = 0.0f;
  for (uint32_t k = 0; k < K; ++k) {
    sum += (float)A[row * K + k] * (float)B[k * M + col];
  }
  C[row * M + col] = sum;
}

template <const uint32_t BN, const uint32_t BK, const uint32_t BM,
          const uint32_t WITERN, const uint32_t WITERM, const uint32_t SM_COUNT,
          bool COL_SCHEDULING>
__global__ void fp16_mm_kernel(const float16_t *A, const float16_t *B,
                               float16_t *C, uint32_t N, uint32_t K,
                               uint32_t M) {
  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
  using float16x4 =
      __attribute__((__vector_size__(4 * sizeof(float16_t)))) float16_t;

  static constexpr uint32_t VECTOR_SIZE = 4;
  static constexpr uint32_t WARPSIZE = 64;
  static constexpr uint32_t WN = 32 * WITERN;
  static constexpr uint32_t WM = 32 * WITERM;
  static constexpr uint32_t numThreads = (BN * BM) / (16 * WITERN * WITERM);
  static constexpr uint32_t strideA = numThreads / (BK / VECTOR_SIZE);
  static constexpr uint32_t strideB = (numThreads / BM) * VECTOR_SIZE;
  static constexpr uint32_t nstridesA = BN > strideA ? BN / strideA : 1;
  static constexpr uint32_t nstridesB = BK > strideA ? BK / strideB : 1;

  static_assert(numThreads % (BK / VECTOR_SIZE) == 0,
                "BK / VECTOR_SIZE should be a multiple of numThreads");
  static_assert(numThreads % BM == 0, "BM should be a multiple of numThreads");

  uint32_t numTiles = cdiv(N, BN) * cdiv(M, BM);

  for (uint32_t tileIdx = blockIdx.x; tileIdx < numTiles; tileIdx += SM_COUNT) {
    uint32_t rowOffsetC, colOffsetC;
    if constexpr (COL_SCHEDULING) {
      rowOffsetC = (tileIdx % cdiv(N, BN)) * BN;
      colOffsetC = (tileIdx / cdiv(N, BN)) * BM;
    } else {
      rowOffsetC = (tileIdx / cdiv(M, BM)) * BN;
      colOffsetC = (tileIdx % cdiv(M, BM)) * BM;
    }
    uint32_t rowOffsetA = rowOffsetC;
    uint32_t colOffsetB = colOffsetC;

    uint32_t innerColA = (threadIdx.x % (BK / VECTOR_SIZE)) * VECTOR_SIZE;
    uint32_t innerRowA = threadIdx.x / (BK / VECTOR_SIZE);
    uint32_t innerColB = threadIdx.x % BM;
    uint32_t innerRowB = (threadIdx.x / BM) * VECTOR_SIZE;

    uint32_t laneIdx = threadIdx.x % WARPSIZE;
    uint32_t warpIdx = threadIdx.x / WARPSIZE;
    uint32_t warpColOffset = (warpIdx % (BM / WM)) * WM;
    uint32_t warpRowOffset = (warpIdx / (BM / WM)) * WN;
    uint32_t warpX = laneIdx % 32;
    uint32_t warpY = laneIdx / 32;

    __shared__ float16_t As[2][BN][BK], Bs[2][BM][BK + 2];
    uint32_t curr = 0;

    for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BN;
         innerRowOffsetA += strideA) {
      if ((rowOffsetA + innerRowOffsetA + innerRowA) < N && (innerColA) < K &&
          (innerRowOffsetA + innerRowA) < BN) {
        *reinterpret_cast<double *>(
            &As[curr][innerRowOffsetA + innerRowA][innerColA]) =
            *reinterpret_cast<const double *>(
                &A[(rowOffsetA + innerRowOffsetA + innerRowA) * K +
                   (innerColA)]);
      } else if ((innerRowOffsetA + innerRowA) < BN) {
        *reinterpret_cast<double *>(
            &As[curr][innerRowOffsetA + innerRowA][innerColA]) = 0.0;
      }
    }
    for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
         innerRowOffsetB += strideB) {
      if ((innerRowOffsetB + innerRowB) < K && (colOffsetB + innerColB) < M &&
          (innerRowOffsetB + innerRowB) < BK) {
        uint64_t x = *reinterpret_cast<const uint16_t *>(
            &B[(innerRowOffsetB + innerRowB) * M + (colOffsetB + innerColB)]);
        uint64_t y = *reinterpret_cast<const uint16_t *>(
            &B[(innerRowOffsetB + innerRowB + 1) * M +
               (colOffsetB + innerColB)]);
        uint64_t z = *reinterpret_cast<const uint16_t *>(
            &B[(innerRowOffsetB + innerRowB + 2) * M +
               (colOffsetB + innerColB)]);
        uint64_t w = *reinterpret_cast<const uint16_t *>(
            &B[(innerRowOffsetB + innerRowB + 3) * M +
               (colOffsetB + innerColB)]);

        uint64_t pack = (w & 0xffff) << 48 | (z & 0xffff) << 32 |
                        (y & 0xffff) << 16 | (x & 0xffff);

        *reinterpret_cast<uint64_t *>(
            &Bs[curr][innerColB][innerRowOffsetB + innerRowB]) = pack;
      } else if ((innerRowOffsetB + innerRowB) < BK) {
        *reinterpret_cast<uint64_t *>(
            &Bs[curr][innerColB][innerRowOffsetB + innerRowB]) = 0;
      }
    }

    __syncthreads();

    floatx16 d[WITERN][WITERM] = {0};
    double A_tmp[nstridesA];
    uint16_t B_tmp1[nstridesB], B_tmp2[nstridesB], B_tmp3[nstridesB],
        B_tmp4[nstridesB];

    for (uint32_t tileOffset = BK; tileOffset < K + BK; tileOffset += BK) {
      // load from global to shared memory in coalesced manner
      if (tileOffset < K) {
        for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BN;
             innerRowOffsetA += strideA) {
          if ((rowOffsetA + innerRowOffsetA + innerRowA) < N &&
              (tileOffset + innerColA) < K &&
              (innerRowOffsetA + innerRowA) < BN) {
            A_tmp[innerRowOffsetA / strideA] =
                *reinterpret_cast<const double *>(
                    &A[(rowOffsetA + innerRowOffsetA + innerRowA) * K +
                       (tileOffset + innerColA)]);
          } else if ((innerRowOffsetA + innerRowA) < BN) {
            A_tmp[innerRowOffsetA / strideA] = 0.0;
          }
        }
        for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
             innerRowOffsetB += strideB) {
          if ((tileOffset + innerRowOffsetB + innerRowB) < K &&
              (colOffsetB + innerColB) < M &&
              (innerRowOffsetB + innerRowB) < BK) {
            B_tmp1[innerRowOffsetB / strideB] =
                *reinterpret_cast<const uint16_t *>(
                    &B[(tileOffset + innerRowOffsetB + innerRowB) * M +
                       (colOffsetB + innerColB)]);
            B_tmp2[innerRowOffsetB / strideB] =
                *reinterpret_cast<const uint16_t *>(
                    &B[(tileOffset + innerRowOffsetB + innerRowB + 1) * M +
                       (colOffsetB + innerColB)]);
            B_tmp3[innerRowOffsetB / strideB] =
                *reinterpret_cast<const uint16_t *>(
                    &B[(tileOffset + innerRowOffsetB + innerRowB + 2) * M +
                       (colOffsetB + innerColB)]);
            B_tmp4[innerRowOffsetB / strideB] =
                *reinterpret_cast<const uint16_t *>(
                    &B[(tileOffset + innerRowOffsetB + innerRowB + 3) * M +
                       (colOffsetB + innerColB)]);
          } else if ((innerRowOffsetB + innerRowB) < BK) {
            B_tmp1[innerRowOffsetB / strideB] = 0;
            B_tmp2[innerRowOffsetB / strideB] = 0;
            B_tmp3[innerRowOffsetB / strideB] = 0;
            B_tmp4[innerRowOffsetB / strideB] = 0;
          }
        }
      }

      float16x4 a[WITERN], b[WITERM];
      for (uint32_t BKOffset = 0; BKOffset < BK; BKOffset += 8) {
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          a[wn] = *reinterpret_cast<float16x4 *>(
              &As[curr][warpRowOffset + wn * 32 + warpX][BKOffset + warpY * 4]);
        }
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          b[wm] = *reinterpret_cast<float16x4 *>(
              &Bs[curr][warpColOffset + wm * 32 + warpX][BKOffset + warpY * 4]);
        }
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          for (uint32_t wm = 0; wm < WITERM; ++wm) {
            d[wn][wm] = __builtin_amdgcn_mfma_f32_32x32x8f16(
                a[wn], b[wm], d[wn][wm], 0, 0, 0);
          }
        }
      }

      if (tileOffset < K) {
        for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BN;
             innerRowOffsetA += strideA) {
          if ((innerRowOffsetA + innerRowA) < BN) {
            *reinterpret_cast<double *>(
                &As[1 - curr][innerRowOffsetA + innerRowA][innerColA]) =
                A_tmp[innerRowOffsetA / strideA];
          }
        }
        for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
             innerRowOffsetB += strideB) {
          if ((innerRowOffsetB + innerRowB) < BK) {
            uint64_t pack =
                ((uint64_t)B_tmp4[innerRowOffsetB / strideB] & 0xffff) << 48 |
                ((uint64_t)B_tmp3[innerRowOffsetB / strideB] & 0xffff) << 32 |
                ((uint64_t)B_tmp2[innerRowOffsetB / strideB] & 0xffff) << 16 |
                ((uint64_t)B_tmp1[innerRowOffsetB / strideB] & 0xffff);

            *reinterpret_cast<uint64_t *>(
                &Bs[1 - curr][innerColB][innerRowOffsetB + innerRowB]) = pack;
          }
        }
      }

      __syncthreads();

      curr = 1 - curr;
    }

    for (uint32_t wn = 0; wn < WITERN; ++wn) {
      for (uint32_t wm = 0; wm < WITERM; ++wm) {
        for (uint32_t j = 0; j < 4; ++j) {
          for (uint32_t i = 0; i < 4; ++i) {
            if ((rowOffsetC + warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i) <
                    N &&
                (colOffsetC + warpColOffset + wm * 32 + warpX) < M)
              C[(rowOffsetC + warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i) *
                    M +
                (colOffsetC + warpColOffset + wm * 32 + warpX)] =
                  d[wn][wm][i + j * 4];
          }
        }
      }
    }
  }
}

at::Tensor fp16_mm(at::Tensor A, at::Tensor B) {
  int N = A.size(0), K = A.size(1), M = B.size(1);
  at::Tensor C = at::empty({N, M}, A.options());

  const uint32_t BK = 16;
  const uint32_t BN = 128;
  const uint32_t BM = 128;
  const uint32_t WITERM = 1;
  const uint32_t WITERN = 1;
  const uint32_t SM_COUNT = 304;
  const bool COL_SCHEDULING = cdiv(N, BN) > cdiv(M, BM);

    dim3 numThreads(32, 32);
    dim3 numBlocks(cdiv(M, numThreads.x), cdiv(N, numThreads.y));
    fp16_mm_ref_kernel<<<numBlocks, numThreads>>>((float16_t *)A.data_ptr(), (float16_t *)B.data_ptr(),
        (float16_t *)C.data_ptr(), N, K, M);

  
  hipDeviceSynchronize();
  return C;
}
"""

cpp_src = r"""
at::Tensor fp16_mm(at::Tensor A, at::Tensor B);
"""

if sys.stdout is None:
    sys.stdout = open("/dev/stdout", "w")
if sys.stderr is None:
    sys.stderr = open("/dev/stderr", "w")

class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = None
        self.W_up = None
        self.W_down = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn((x @ self.W_gate))
        out = (gate * (x @ self.W_up) @ self.W_down)
        return out


class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

        self.W_g = None

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

        return topk_indices, topk_scores # B, C, k


class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            Expert(config)
            for _ in range(config["n_routed_experts"])
        ])
        self.gating_network = MoEGate(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape
        x_flat = x.view(-1, hidden_dim) # B * C, d
        flat_expert_indices = expert_indices.view(-1) # B * C * k
        flat_expert_weights = expert_scores.view(-1, 1) # B * C * k
        routed_output_flat = self.moe_infer(x_flat,
                                            flat_expert_indices,
                                            flat_expert_weights)

        routed_output = routed_output_flat.view(*orig_shape)
        return routed_output + shared_output

    @torch.no_grad()
    def moe_infer(self,
                  x: torch.Tensor, # B * C, d
                  flat_expert_indices: torch.Tensor, # B * C * k
                  flat_expert_weights: torch.Tensor # B * C * k
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x) # B * C, d
        idxs = flat_expert_indices.argsort() # B * C * k
        counts = flat_expert_indices.bincount().cpu().numpy()
        tokens_per_expert = counts.cumsum() # 
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok
        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            expert = self.experts[expert_id]
            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idxs]
            expert_out    = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_reduce_(
                0,
                exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out,
                reduce='sum'
            )

        return expert_cache


def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for 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_size]
            - 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"]
    moe = MoE(config)

    # Fill in the given weights of the model
    moe.gating_network.W_g = weights['router.weight'].t()

    for i in range(num_experts):
        moe.experts[i].W_gate = weights[f'experts.{i}.0.weight']
        moe.experts[i].W_up = weights[f'experts.{i}.1.weight']
        moe.experts[i].W_down = weights[f'experts.{i}.2.weight']

    moe.shared_expert.W_gate = weights['shared_experts.0.weight']
    moe.shared_expert.W_up = weights['shared_experts.1.weight']
    moe.shared_expert.W_down = weights['shared_experts.2.weight']

    # Run the model
    output = moe(input_tensor)

    return output

python eval.py benchmark /tmp/tmpsao9oig7