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

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <stdint.h>

namespace {

constexpr int BM = 128;
constexpr int BK = 64;
constexpr int CPR = BK / 8;
constexpr int BN1 = 64;
constexpr int BN2 = 128;
constexpr int STAGES = 3;
constexpr int THREADS = 256;
constexpr int ROUTE_THREADS = 1024;
constexpr int ROW_BYTES = BK * 2;
constexpr int STAGE_BYTES = (BM + 2 * BN1) * ROW_BYTES;
constexpr int MAX_GROUPS = 320;

__device__ __forceinline__ uint32_t smem_addr(const void* p) {
  return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

__device__ __forceinline__ int swz(int row, int chunk) {
  return (chunk ^ (row & 7)) << 4;
}

__device__ __forceinline__ void cp_async16(void* smem, const void* gmem, bool valid) {
  const int src_bytes = valid ? 16 : 0;
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
               :: "r"(smem_addr(smem)), "l"(reinterpret_cast<uint64_t>(gmem)), "r"(src_bytes));
}

__device__ __forceinline__ void cp_async_commit() {
  asm volatile("cp.async.commit_group;\n" ::);
}

template <int N>
__device__ __forceinline__ void cp_async_wait() {
  asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

__device__ __forceinline__ void ldmatrix_x4(uint32_t* r, const void* smem) {
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
               : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
               : "r"(smem_addr(smem)));
}

__device__ __forceinline__ void mma_bf16(float* d, const uint32_t* a, uint32_t b0, uint32_t b1) {
  asm volatile(
      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
      : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
}

__global__ void routing_kernel(const int64_t* __restrict__ ids, int T, int topk, int E, int NS,
                               int* __restrict__ gcnt, int* __restrict__ goff,
                               int* __restrict__ tile_group, int* __restrict__ tile_row0,
                               int* __restrict__ ntiles, int* __restrict__ row_token,
                               int* __restrict__ row_slot) {
  __shared__ int s_cnt[MAX_GROUPS];
  __shared__ int s_cur[MAX_GROUPS];
  __shared__ int s_off[MAX_GROUPS + 1];
  __shared__ int s_toff[MAX_GROUPS + 1];
  const int tid = threadIdx.x;
  const int NG = E + NS;
  const int P = T * topk;
  for (int g = tid; g < NG; g += blockDim.x) {
    s_cnt[g] = 0;
    s_cur[g] = 0;
  }
  __syncthreads();
  for (int p = tid; p < P; p += blockDim.x) atomicAdd(&s_cnt[ids[p]], 1);
  __syncthreads();
  if (tid == 0) {
    for (int s = 0; s < NS; ++s) s_cnt[E + s] = T;
    int off = 0, toff = 0;
    for (int g = 0; g < NG; ++g) {
      s_off[g] = off;
      s_toff[g] = toff;
      off += s_cnt[g];
      toff += (s_cnt[g] + BM - 1) / BM;
    }
    ntiles[0] = toff;
  }
  __syncthreads();
  for (int g = tid; g < NG; g += blockDim.x) {
    gcnt[g] = s_cnt[g];
    goff[g] = s_off[g];
    const int nt = (s_cnt[g] + BM - 1) / BM;
    for (int j = 0; j < nt; ++j) {
      tile_group[s_toff[g] + j] = g;
      tile_row0[s_toff[g] + j] = j * BM;
    }
  }
  for (int p = tid; p < P; p += blockDim.x) {
    const int g = static_cast<int>(ids[p]);
    const int r = s_off[g] + atomicAdd(&s_cur[g], 1);
    row_token[r] = p / topk;
    row_slot[r] = p;
  }
  for (int i = tid; i < NS * T; i += blockDim.x) {
    const int s = i / T;
    const int t = i - s * T;
    const int r = s_off[E + s] + t;
    row_token[r] = t;
    row_slot[r] = P + i;
  }
}

__global__ void __launch_bounds__(THREADS, 1)
moe_gemm1_kernel(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ w1r,
                 const __nv_bfloat16* __restrict__ w1s, __nv_bfloat16* __restrict__ hbuf,
                 const int* __restrict__ gcnt, const int* __restrict__ goff,
                 const int* __restrict__ tile_group, const int* __restrict__ tile_row0,
                 const int* __restrict__ ntiles, const int* __restrict__ row_token,
                 int E, int H, int I, int ntn) {
  extern __shared__ __align__(16) unsigned char smem[];
  __shared__ int s_tok[BM];
  const int m = blockIdx.x / ntn;
  if (m >= ntiles[0]) return;
  const int nt = blockIdx.x - m * ntn;
  const int g = tile_group[m];
  const int row0 = tile_row0[m];
  const int cnt = gcnt[g];
  const int gbase = goff[g];
  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int wm = (tid >> 5) >> 1;
  const int wn = (tid >> 5) & 1;
  if (tid < BM) {
    const int r = row0 + tid;
    s_tok[tid] = r < cnt ? row_token[gbase + r] : -1;
  }
  __syncthreads();
  const size_t wstride = (size_t)2 * I * H;
  const __nv_bfloat16* W = g < E ? w1r + (size_t)g * wstride : w1s + (size_t)(g - E) * wstride;
  const int n0 = nt * BN1;

  auto load_stage = [&](int stage, int kt) {
    unsigned char* sb = smem + stage * STAGE_BYTES;
    const int k0 = kt * BK;
    for (int q = tid; q < BM * CPR; q += THREADS) {
      const int row = q / CPR, c = q % CPR;
      const int tok = s_tok[row];
      const __nv_bfloat16* src = tok >= 0 ? x + (size_t)tok * H + k0 + c * 8 : x;
      cp_async16(sb + row * ROW_BYTES + swz(row, c), src, tok >= 0);
    }
    for (int q = tid; q < BN1 * CPR; q += THREADS) {
      const int row = q / CPR, c = q % CPR;
      cp_async16(sb + BM * ROW_BYTES + row * ROW_BYTES + swz(row, c),
                 W + (size_t)(n0 + row) * H + k0 + c * 8, true);
      cp_async16(sb + (BM + BN1) * ROW_BYTES + row * ROW_BYTES + swz(row, c),
                 W + (size_t)(I + n0 + row) * H + k0 + c * 8, true);
    }
  };

  float accg[2][4][4], accu[2][4][4];
  for (int i = 0; i < 2; ++i)
    for (int j = 0; j < 4; ++j)
      for (int k = 0; k < 4; ++k) accg[i][j][k] = accu[i][j][k] = 0.f;

  const int KT = H / BK;
  for (int s = 0; s < STAGES - 1; ++s) {
    if (s < KT) load_stage(s, s);
    cp_async_commit();
  }
  for (int kt = 0; kt < KT; ++kt) {
    cp_async_wait<STAGES - 2>();
    __syncthreads();
    const int nk = kt + STAGES - 1;
    if (nk < KT) load_stage(nk % STAGES, nk);
    cp_async_commit();

    const unsigned char* sA = smem + (kt % STAGES) * STAGE_BYTES;
    const unsigned char* sG = sA + BM * ROW_BYTES;
    const unsigned char* sU = sA + (BM + BN1) * ROW_BYTES;
#pragma unroll
    for (int ks = 0; ks < BK / 16; ++ks) {
      uint32_t af[2][4];
#pragma unroll
      for (int mi = 0; mi < 2; ++mi) {
        const int row = wm * 32 + mi * 16 + (lane & 15);
        ldmatrix_x4(af[mi], sA + row * ROW_BYTES + swz(row, ks * 2 + (lane >> 4)));
      }
      uint32_t bg[4][2], bu[4][2];
#pragma unroll
      for (int nj = 0; nj < 2; ++nj) {
        const int n = wn * 32 + nj * 16 + (lane & 7) + ((lane >> 4) << 3);
        const int off = swz(n, ks * 2 + ((lane >> 3) & 1));
        uint32_t r[4];
        ldmatrix_x4(r, sG + n * ROW_BYTES + off);
        bg[2 * nj][0] = r[0]; bg[2 * nj][1] = r[1];
        bg[2 * nj + 1][0] = r[2]; bg[2 * nj + 1][1] = r[3];
        ldmatrix_x4(r, sU + n * ROW_BYTES + off);
        bu[2 * nj][0] = r[0]; bu[2 * nj][1] = r[1];
        bu[2 * nj + 1][0] = r[2]; bu[2 * nj + 1][1] = r[3];
      }
#pragma unroll
      for (int mi = 0; mi < 2; ++mi)
#pragma unroll
        for (int ni = 0; ni < 4; ++ni) {
          mma_bf16(accg[mi][ni], af[mi], bg[ni][0], bg[ni][1]);
          mma_bf16(accu[mi][ni], af[mi], bu[ni][0], bu[ni][1]);
        }
    }
  }

  const int gl = lane >> 2, tq = (lane & 3) * 2;
  for (int mi = 0; mi < 2; ++mi)
    for (int ni = 0; ni < 4; ++ni)
      for (int hf = 0; hf < 2; ++hf) {
        const int row = wm * 32 + mi * 16 + gl + hf * 8;
        if (row0 + row >= cnt) continue;
        const int col = n0 + wn * 32 + ni * 8 + tq;
        const float g0 = accg[mi][ni][2 * hf], g1 = accg[mi][ni][2 * hf + 1];
        const float u0 = accu[mi][ni][2 * hf], u1 = accu[mi][ni][2 * hf + 1];
        const float h0 = g0 / (1.f + expf(-g0)) * u0;
        const float h1 = g1 / (1.f + expf(-g1)) * u1;
        *reinterpret_cast<__nv_bfloat162*>(hbuf + (size_t)(gbase + row0 + row) * I + col) =
            __floats2bfloat162_rn(h0, h1);
      }
}

__global__ void __launch_bounds__(THREADS, 1)
moe_gemm2_kernel(const __nv_bfloat16* __restrict__ hbuf, const __nv_bfloat16* __restrict__ w2r,
                 const __nv_bfloat16* __restrict__ w2s, __nv_bfloat16* __restrict__ ybuf,
                 const int* __restrict__ gcnt, const int* __restrict__ goff,
                 const int* __restrict__ tile_group, const int* __restrict__ tile_row0,
                 const int* __restrict__ ntiles, const int* __restrict__ row_slot,
                 int E, int H, int I, int ntn) {
  extern __shared__ __align__(16) unsigned char smem[];
  const int m = blockIdx.x / ntn;
  if (m >= ntiles[0]) return;
  const int nt = blockIdx.x - m * ntn;
  const int g = tile_group[m];
  const int row0 = tile_row0[m];
  const int cnt = gcnt[g];
  const int gbase = goff[g];
  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int wm = (tid >> 5) >> 1;
  const int wn = (tid >> 5) & 1;
  const size_t wstride = (size_t)H * I;
  const __nv_bfloat16* W = g < E ? w2r + (size_t)g * wstride : w2s + (size_t)(g - E) * wstride;
  const int n0 = nt * BN2;

  auto load_stage = [&](int stage, int kt) {
    unsigned char* sb = smem + stage * STAGE_BYTES;
    const int k0 = kt * BK;
    for (int q = tid; q < BM * CPR; q += THREADS) {
      const int row = q / CPR, c = q % CPR;
      const bool valid = row0 + row < cnt;
      const __nv_bfloat16* src =
          valid ? hbuf + (size_t)(gbase + row0 + row) * I + k0 + c * 8 : hbuf;
      cp_async16(sb + row * ROW_BYTES + swz(row, c), src, valid);
    }
    for (int q = tid; q < BN2 * CPR; q += THREADS) {
      const int row = q / CPR, c = q % CPR;
      cp_async16(sb + BM * ROW_BYTES + row * ROW_BYTES + swz(row, c),
                 W + (size_t)(n0 + row) * I + k0 + c * 8, true);
    }
  };

  float acc[2][8][4];
  for (int i = 0; i < 2; ++i)
    for (int j = 0; j < 8; ++j)
      for (int k = 0; k < 4; ++k) acc[i][j][k] = 0.f;

  const int KT = I / BK;
  for (int s = 0; s < STAGES - 1; ++s) {
    if (s < KT) load_stage(s, s);
    cp_async_commit();
  }
  for (int kt = 0; kt < KT; ++kt) {
    cp_async_wait<STAGES - 2>();
    __syncthreads();
    const int nk = kt + STAGES - 1;
    if (nk < KT) load_stage(nk % STAGES, nk);
    cp_async_commit();

    const unsigned char* sA = smem + (kt % STAGES) * STAGE_BYTES;
    const unsigned char* sB = sA + BM * ROW_BYTES;
#pragma unroll
    for (int ks = 0; ks < BK / 16; ++ks) {
      uint32_t af[2][4];
#pragma unroll
      for (int mi = 0; mi < 2; ++mi) {
        const int row = wm * 32 + mi * 16 + (lane & 15);
        ldmatrix_x4(af[mi], sA + row * ROW_BYTES + swz(row, ks * 2 + (lane >> 4)));
      }
      uint32_t bf[8][2];
#pragma unroll
      for (int nj = 0; nj < 4; ++nj) {
        const int n = wn * 64 + nj * 16 + (lane & 7) + ((lane >> 4) << 3);
        uint32_t r[4];
        ldmatrix_x4(r, sB + n * ROW_BYTES + swz(n, ks * 2 + ((lane >> 3) & 1)));
        bf[2 * nj][0] = r[0]; bf[2 * nj][1] = r[1];
        bf[2 * nj + 1][0] = r[2]; bf[2 * nj + 1][1] = r[3];
      }
#pragma unroll
      for (int mi = 0; mi < 2; ++mi)
#pragma unroll
        for (int ni = 0; ni < 8; ++ni) mma_bf16(acc[mi][ni], af[mi], bf[ni][0], bf[ni][1]);
    }
  }

  const int gl = lane >> 2, tq = (lane & 3) * 2;
  for (int mi = 0; mi < 2; ++mi)
    for (int ni = 0; ni < 8; ++ni)
      for (int hf = 0; hf < 2; ++hf) {
        const int row = wm * 32 + mi * 16 + gl + hf * 8;
        if (row0 + row >= cnt) continue;
        const int slot = row_slot[gbase + row0 + row];
        const int col = n0 + wn * 64 + ni * 8 + tq;
        *reinterpret_cast<__nv_bfloat162*>(ybuf + (size_t)slot * H + col) =
            __floats2bfloat162_rn(acc[mi][ni][2 * hf], acc[mi][ni][2 * hf + 1]);
      }
}

__device__ __forceinline__ void accum8(float* acc, const uint4& v, float w) {
  const __nv_bfloat162* p = reinterpret_cast<const __nv_bfloat162*>(&v);
#pragma unroll
  for (int i = 0; i < 4; ++i) {
    const float2 f = __bfloat1622float2(p[i]);
    acc[2 * i] += w * f.x;
    acc[2 * i + 1] += w * f.y;
  }
}

__global__ void moe_reduce_kernel(const __nv_bfloat16* __restrict__ ybuf,
                                  const __nv_bfloat16* __restrict__ wts,
                                  __nv_bfloat16* __restrict__ out, int T, int H, int topk, int NS) {
  const int t = blockIdx.x;
  const int h = threadIdx.x * 8;
  float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
  for (int k = 0; k < topk; ++k) {
    const int slot = t * topk + k;
    const uint4 v = *reinterpret_cast<const uint4*>(ybuf + (size_t)slot * H + h);
    accum8(acc, v, __bfloat162float(wts[slot]));
  }
  for (int s = 0; s < NS; ++s) {
    const int slot = T * topk + s * T + t;
    const uint4 v = *reinterpret_cast<const uint4*>(ybuf + (size_t)slot * H + h);
    accum8(acc, v, 1.f);
  }
  uint4 o;
  __nv_bfloat162* po = reinterpret_cast<__nv_bfloat162*>(&o);
#pragma unroll
  for (int i = 0; i < 4; ++i) po[i] = __floats2bfloat162_rn(acc[2 * i], acc[2 * i + 1]);
  *reinterpret_cast<uint4*>(out + (size_t)t * H + h) = o;
}

}  // namespace

torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor wts,
                          torch::Tensor w1r, torch::Tensor w2r, torch::Tensor w1s,
                          torch::Tensor w2s) {
  const c10::cuda::CUDAGuard guard(x.device());
  const int T = x.size(0), H = x.size(1);
  const int topk = ids.size(1);
  const int E = w1r.size(0), I = w2r.size(2), NS = w1s.size(0);
  TORCH_CHECK(H % BN2 == 0 && H % BK == 0 && H / 8 <= 1024, "H must be a multiple of 128");
  TORCH_CHECK(I % BN1 == 0 && I % BK == 0, "I must be a multiple of 64");
  TORCH_CHECK(E + NS <= MAX_GROUPS, "too many expert groups");
  const int NG = E + NS;
  const int P = T * topk;
  const int R = P + NS * T;
  const int maxtiles = P / BM + E + NS * ((T + BM - 1) / BM);
  const int ntn1 = I / BN1, ntn2 = H / BN2;

  const auto i32 = x.options().dtype(torch::kInt32);
  auto gcnt = torch::empty({NG}, i32);
  auto goff = torch::empty({NG}, i32);
  auto tile_group = torch::empty({maxtiles}, i32);
  auto tile_row0 = torch::empty({maxtiles}, i32);
  auto ntiles = torch::empty({1}, i32);
  auto [REDACTED credential assignment]}, i32);
  auto row_slot = torch::empty({R}, i32);
  auto hbuf = torch::empty({R, I}, x.options());
  auto ybuf = torch::empty({R, H}, x.options());
  auto out = torch::empty({T, H}, x.options());
  auto stream = at::cuda::getCurrentCUDAStream().stream();

  routing_kernel<<<1, ROUTE_THREADS, 0, stream>>>(
      ids.data_ptr<int64_t>(), T, topk, E, NS, gcnt.data_ptr<int>(), goff.data_ptr<int>(),
      tile_group.data_ptr<int>(), tile_row0.data_ptr<int>(), ntiles.data_ptr<int>(),
      row_token.data_ptr<int>(), row_slot.data_ptr<int>());

  static bool smem_attr_set = false;
  if (!smem_attr_set) {
    cudaFuncSetAttribute(moe_gemm1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                         STAGES * STAGE_BYTES);
    cudaFuncSetAttribute(moe_gemm2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                         STAGES * STAGE_BYTES);
    smem_attr_set = true;
  }
  moe_gemm1_kernel<<<maxtiles * ntn1, THREADS, STAGES * STAGE_BYTES, stream>>>(
      reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
      reinterpret_cast<const __nv_bfloat16*>(w1r.data_ptr()),
      reinterpret_cast<const __nv_bfloat16*>(w1s.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(hbuf.data_ptr()), gcnt.data_ptr<int>(), goff.data_ptr<int>(),
      tile_group.data_ptr<int>(), tile_row0.data_ptr<int>(), ntiles.data_ptr<int>(),
      row_token.data_ptr<int>(), E, H, I, ntn1);

  moe_gemm2_kernel<<<maxtiles * ntn2, THREADS, STAGES * STAGE_BYTES, stream>>>(
      reinterpret_cast<const __nv_bfloat16*>(hbuf.data_ptr()),
      reinterpret_cast<const __nv_bfloat16*>(w2r.data_ptr()),
      reinterpret_cast<const __nv_bfloat16*>(w2s.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(ybuf.data_ptr()), gcnt.data_ptr<int>(), goff.data_ptr<int>(),
      tile_group.data_ptr<int>(), tile_row0.data_ptr<int>(), ntiles.data_ptr<int>(),
      row_slot.data_ptr<int>(), E, H, I, ntn2);

  moe_reduce_kernel<<<T, H / 8, 0, stream>>>(
      reinterpret_cast<const __nv_bfloat16*>(ybuf.data_ptr()),
      reinterpret_cast<const __nv_bfloat16*>(wts.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), T, H, topk, NS);
  return out;
}
"""

_CPP_SRC = (
    "torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor wts, "
    "torch::Tensor w1r, torch::Tensor w2r, torch::Tensor w1s, torch::Tensor w2s);"
)

_ext = load_inline(
    name="glm52_moe_bf16",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["moe_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "-arch=sm_120", "-lineinfo"],
    verbose=False,
)


class Model(nn.Module):
    def __init__(self, T, E, top_k, n_shared, H, I):
        super().__init__()
        self.T, self.E, self.top_k = T, E, top_k
        self.n_shared, self.H, self.I = n_shared, H, I
        self.w1_routed = nn.Parameter(torch.empty(E, 2 * I, H, dtype=torch.bfloat16, device="cuda"))
        self.w2_routed = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16, device="cuda"))
        self.w1_shared = nn.Parameter(
            torch.empty(n_shared, 2 * I, H, dtype=torch.bfloat16, device="cuda")
        )
        self.w2_shared = nn.Parameter(torch.empty(n_shared, H, I, dtype=torch.bfloat16, device="cuda"))
        for p in self.parameters():
            nn.init.normal_(p, std=0.02)

    def forward(
        self,
        x: torch.Tensor,
        expert_ids: torch.Tensor,
        expert_weights: torch.Tensor,
    ) -> torch.Tensor:
        return _ext.moe_forward(
            x.contiguous(),
            expert_ids.contiguous(),
            expert_weights.contiguous(),
            self.w1_routed,
            self.w2_routed,
            self.w1_shared,
            self.w2_shared,
        )