#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <torch/torch.h>

template<class T>
__global__ void packed_gemv(T* out, const uint8_t* x, const uint8_t* w,
    const float* xn, const c10::BFloat16* wn, const int8_t* ac, const int8_t* wc,
    const T* bias, int n, int packed_k, float as, float ws) {
  const int lane = threadIdx.x & 31;
  const int col = blockIdx.x * 4 + threadIdx.x / 32;
  const int row = blockIdx.y;
  __shared__ int a_codes[16], w_codes[16];
  if (threadIdx.x < 16) { a_codes[threadIdx.x] = ac[threadIdx.x]; w_codes[threadIdx.x] = wc[threadIdx.x]; }
  __syncthreads();
  if (col >= n) return;
  int sum = 0;
  if (packed_k % 4 == 0 && (reinterpret_cast<uintptr_t>(x) & 3) == 0 && (reinterpret_cast<uintptr_t>(w) & 3) == 0) {
    const uint32_t* xp = reinterpret_cast<const uint32_t*>(x + row * packed_k);
    const uint32_t* wp = reinterpret_cast<const uint32_t*>(w + col * packed_k);
    for (int k = lane; k < packed_k / 4; k += 32) {
      const uint32_t av = xp[k], bv = wp[k];
      uint32_t al = 0, ah = 0, bl = 0, bh = 0;
      #pragma unroll
      for (int j = 0; j < 4; ++j) {
        al |= uint32_t(uint8_t(a_codes[(av >> (j * 8)) & 15])) << (j * 8);
        ah |= uint32_t(uint8_t(a_codes[(av >> (j * 8 + 4)) & 15])) << (j * 8);
        bl |= uint32_t(uint8_t(w_codes[(bv >> (j * 8)) & 15])) << (j * 8);
        bh |= uint32_t(uint8_t(w_codes[(bv >> (j * 8 + 4)) & 15])) << (j * 8);
      }
      sum = __dp4a(static_cast<int>(al), static_cast<int>(bl), sum);
      sum = __dp4a(static_cast<int>(ah), static_cast<int>(bh), sum);
    }
  } else {
    for (int k = lane; k < packed_k; k += 32) {
      const uint8_t a = x[row * packed_k + k];
      const uint8_t b = w[col * packed_k + k];
      sum += a_codes[a & 15] * w_codes[b & 15];
      sum += a_codes[a >> 4] * w_codes[b >> 4];
    }
  }
  #pragma unroll
  for (int offset = 16; offset; offset >>= 1) sum += __shfl_down_sync(0xffffffff, sum, offset);
  if (lane == 0) {
    float value = static_cast<float>(sum);
    value *= xn[row] * static_cast<float>(wn[col]) * (as * ws);
    if (bias) value += static_cast<float>(bias[col]);
    out[row * n + col] = static_cast<T>(value);
  }
}

void orbitquant_gemv(torch::Tensor &out, torch::Tensor const &x, torch::Tensor const &w, torch::Tensor const &xn, torch::Tensor const &wn, torch::Tensor const &ac, torch::Tensor const &wc, torch::Tensor const &bias, bool has_bias, double activation_scale, double weight_scale) {
  TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kUInt8 && x.dim() == 2, "x must be rank-2 packed CUDA uint8");
  const torch::Tensor* tensors[] = {&out, &w, &xn, &wn, &ac, &wc};
  for (const auto* t : tensors) {
    TORCH_CHECK(t->device() == x.device() && t->is_contiguous(), "tensors must be contiguous on the input CUDA device");
  }
  TORCH_CHECK(x.is_contiguous() && w.scalar_type() == torch::kUInt8, "packed inputs must be contiguous uint8");
  TORCH_CHECK(x.size(0) >= 1 && x.size(0) <= 8 && x.size(1) > 0 && x.size(1) <= 8192, "supported rows 1..8 and packed K 1..8192");
  TORCH_CHECK(wn.scalar_type() == torch::kBFloat16 && xn.scalar_type() == torch::kFloat, "norm dtypes must be BF16 weights and FP32 activations");
  TORCH_CHECK(ac.scalar_type() == torch::kChar && wc.scalar_type() == torch::kChar && ac.numel() == 16 && wc.numel() == 16, "codebooks must contain 16 INT8 entries");
  TORCH_CHECK(out.dim() == 2 && out.size(0) == x.size(0) && out.size(1) == wn.numel() && xn.numel() == x.size(0), "output and norm shapes mismatch");
  TORCH_CHECK(w.numel() == wn.numel() * x.size(1), "packed weight size mismatch");
  TORCH_CHECK(out.scalar_type() == torch::kBFloat16 || out.scalar_type() == torch::kHalf, "output must be BF16 or FP16");
  if (has_bias) TORCH_CHECK(bias.device() == x.device() && bias.is_contiguous() && bias.scalar_type() == out.scalar_type() && bias.numel() == wn.numel(), "invalid bias");
  if (wn.numel() == 0) return;
  const at::cuda::OptionalCUDAGuard guard(device_of(x));
  const dim3 grid((wn.numel() + 3) / 4, x.size(0));
  const auto stream = at::cuda::getCurrentCUDAStream();
  #define LAUNCH(T) packed_gemv<T><<<grid, 128, 0, stream>>>(out.data_ptr<T>(), x.data_ptr<uint8_t>(), w.data_ptr<uint8_t>(), xn.data_ptr<float>(), wn.data_ptr<c10::BFloat16>(), ac.data_ptr<int8_t>(), wc.data_ptr<int8_t>(), has_bias ? bias.data_ptr<T>() : nullptr, wn.numel(), x.size(1), activation_scale, weight_scale)
  if (out.scalar_type() == torch::kBFloat16) { LAUNCH(c10::BFloat16); } else { LAUNCH(c10::Half); }
  #undef LAUNCH
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}