/******************************************************************************
 * Copyright (c) 2025, Haisheng Chen.
 ******************************************************************************/

#include "rotation.cuh"
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>

template <typename scalar_t, int CTA_M, int GROUP_SIZE, int KROT, bool USE_SCALE>
__global__ void rotate(const scalar_t *__restrict__ x, scalar_t *__restrict__ out,
                       const int16_t *__restrict__ idx_ij, const scalar_t *__restrict__ theta,
                       const scalar_t *__restrict__ scales, int s, int h) {
  // shared buffer
  static_assert(CTA_M % 2 == 0, "CTA_M must be even");
  static_assert(GROUP_SIZE % 2 == 0, "GROUP_SIZE must be even");
  __shared__ scalar_t x_grp[CTA_M * GROUP_SIZE];

  int j = blockIdx.x;  // row‐block
  int g = blockIdx.y;  // group
  int t = threadIdx.x; // pair within group

  // load & scale into shared memory
  RotateAccess<scalar_t>::template load_group<CTA_M, GROUP_SIZE, USE_SCALE>(x_grp, x, scales, s, h,
                                                                            j, g, t);

  // fetch sin/cos + idx into registers
  float reg_theta[KROT];
  int reg_idx[KROT];
  RotateAccess<scalar_t>::template load_coeffs<KROT, GROUP_SIZE>(reg_theta, reg_idx, idx_ij, theta,
                                                                 h, g, t);
  __syncthreads();

// apply all KROT rotations
#pragma unroll
  for (int r = 0; r < KROT; r++) {
    RotateAccess<scalar_t>::template apply_one<CTA_M>(x_grp, reg_idx[r], reg_theta[r]);
    __syncthreads();
  }

  // write back
  RotateAccess<scalar_t>::template store_group<CTA_M, GROUP_SIZE>(out, x_grp, s, h, j, g, t);
}

// Dispatch macro: cast pointers and launch the templated kernel for a given dtype.
#define LAUNCH_ROTATE(CUDA_T, TORCH_T)                                                             \
  {                                                                                                \
    auto *x_p = reinterpret_cast<CUDA_T *>(x.data_ptr<TORCH_T>());                                 \
    auto *o_p = reinterpret_cast<CUDA_T *>(out.data_ptr<TORCH_T>());                               \
    auto *t_p = reinterpret_cast<CUDA_T *>(theta_cast.data_ptr<TORCH_T>());                        \
    if (has_scale) {                                                                               \
      auto *s_p = reinterpret_cast<CUDA_T *>(scales_cast.data_ptr<TORCH_T>());                     \
      rotate<CUDA_T, CTA_M, GROUP_SIZE, KROT, true><<<grid, block, 0, stream>>>(                   \
          x_p, o_p, idx_ij.data_ptr<int16_t>(), t_p, s_p, seq_len, h);                             \
    } else {                                                                                       \
      rotate<CUDA_T, CTA_M, GROUP_SIZE, KROT, false><<<grid, block, 0, stream>>>(                  \
          x_p, o_p, idx_ij.data_ptr<int16_t>(), t_p, nullptr, seq_len, h);                         \
    }                                                                                              \
    break;                                                                                         \
  }

template <int KROT, int CTA_M, int GROUP_SIZE>
torch::Tensor rotate_launcher(at::Tensor x, at::Tensor idx_ij, at::Tensor theta,
                              at::Tensor scales) {
  int h = x.size(-1);
  TORCH_CHECK(h % GROUP_SIZE == 0, "h must be divisible by GROUP_SIZE");
  int groups_per_row = h / GROUP_SIZE;
  constexpr int pn = GROUP_SIZE / 2;
  int seq_len = x.numel() / x.size(-1);
  auto options = torch::TensorOptions().dtype(x.dtype()).device(x.device());
  at::Tensor out = torch::empty(x.sizes(), options);
  bool has_scale = scales.defined() && scales.numel() > 0;

  auto dtype = x.scalar_type();
  auto theta_cast = theta.scalar_type() == dtype ? theta : theta.to(x.dtype());
  auto scales_cast = !has_scale                      ? at::Tensor()
                     : scales.scalar_type() == dtype ? scales
                                                     : scales.to(x.dtype());

  dim3 grid((seq_len + CTA_M - 1) / CTA_M, groups_per_row);
  dim3 block(pn);
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  switch (dtype) {
  case at::kFloat:
    LAUNCH_ROTATE(float, float)
  case at::kHalf:
    LAUNCH_ROTATE(__half, c10::Half)
  case at::kBFloat16:
    LAUNCH_ROTATE(__nv_bfloat16, c10::BFloat16)
  default:
    TORCH_CHECK(false, "rotate supports Float, Half, and BFloat16, got ", x.scalar_type());
  }
  return out;
}

#undef LAUNCH_ROTATE

#define DISPATCH_KROT(GS)                                                                          \
  switch (krot) {                                                                                  \
  case 1:                                                                                          \
    return rotate_launcher<1, 4, GS>(x, idx, theta, scales);                                       \
  case 2:                                                                                          \
    return rotate_launcher<2, 4, GS>(x, idx, theta, scales);                                       \
  case 4:                                                                                          \
    return rotate_launcher<4, 4, GS>(x, idx, theta, scales);                                       \
  case 8:                                                                                          \
    return rotate_launcher<8, 4, GS>(x, idx, theta, scales);                                       \
  case 16:                                                                                         \
    return rotate_launcher<16, 4, GS>(x, idx, theta, scales);                                      \
  default:                                                                                         \
    TORCH_CHECK(false, "Unsupported KROT = ", krot, "; compiled variants: 1/2/4/8/16");            \
  }

torch::Tensor rotate_dynamic(at::Tensor x, at::Tensor idx, at::Tensor theta,
                             c10::optional<at::Tensor> scales_opt, int64_t group_size = 128) {
  int64_t krot = theta.size(0);
  TORCH_CHECK(krot == idx.size(0), "theta.size(0) must equal idx_ij.size(0)");
  at::Tensor scales = scales_opt.value_or(at::Tensor());

  if (group_size == 128) {
    DISPATCH_KROT(128)
  }
  if (group_size == 64) {
    DISPATCH_KROT(64)
  }
  TORCH_CHECK(false, "Unsupported group_size: ", group_size, "; expected 64 or 128");
}

#undef DISPATCH_KROT

TORCH_LIBRARY(rotation, m) {
  m.def("rotate(Tensor x, Tensor idx_ij, Tensor theta, Tensor? scales=None, "
        "int group_size=128) -> Tensor");
}

TORCH_LIBRARY_IMPL(rotation, CUDA, m) {
  m.impl("rotate", &rotate_dynamic);
}
