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

using BF = c10::BFloat16;

// The positive RMS scale cancels in quantization except for quant_eps.
// Keep it in the returned token norm; deliberately omit BF16 intermediates.
__global__ void rmsquant_kernel(uint8_t* packed, float* norms, const BF* x,
    const BF* weight, const int32_t* permutation, const int8_t* signs,
    const float* boundaries, float rms_eps, float quant_eps) {
  constexpr int dim=2048;
  __shared__ float values[dim], raw_sum[256], weighted_sum[256], bounds[15];
  const int t=threadIdx.x, row=blockIdx.x;
  float sx=0, sy=0;
  for(int j=t;j<dim;j+=256){
    int source=permutation[j];
    float raw=float(x[row*dim+source]);
    float weighted=raw*float(weight[source]);
    values[j]=weighted*float(signs[j]);
    sx=fmaf(raw,raw,sx);sy=fmaf(weighted,weighted,sy);
  }
  raw_sum[t]=sx;weighted_sum[t]=sy;
  if(t<15)bounds[t]=boundaries[t];
  __syncthreads();
  for(int stride=128;stride>=32;stride>>=1){
    if(t<stride){raw_sum[t]+=raw_sum[t+stride];weighted_sum[t]+=weighted_sum[t+stride];}
    __syncthreads();
  }
  // Preserve the original reduction tree, but finish inside the first warp.
  if(t<32){
    float raw=raw_sum[t],weighted=weighted_sum[t];
    #pragma unroll
    for(int stride=16;stride;stride>>=1){
      raw+=__shfl_down_sync(0xffffffff,raw,stride);
      weighted+=__shfl_down_sync(0xffffffff,weighted,stride);
    }
    if(t==0){raw_sum[0]=raw;weighted_sum[0]=weighted;}
  }
  __syncthreads();
  float scale=rsqrtf(raw_sum[0]/float(dim)+rms_eps);
  float norm=sqrtf(weighted_sum[0]);
  if(t==0)norms[row]=norm*scale;
  float inv=1.f/(norm+quant_eps/scale);
  float local[8];
  #pragma unroll
  for(int i=0;i<8;i++)local[i]=values[t+256*i]*inv;
  // The first five Hadamard stages stay inside each warp.
  #pragma unroll
  for(int width=1;width<32;width<<=1){
    #pragma unroll
    for(int i=0;i<8;i++){
      float other=__shfl_xor_sync(0xffffffff,local[i],width);
      local[i]=(t&width)?other-local[i]:local[i]+other;
    }
  }
  #pragma unroll
  for(int i=0;i<8;i++)values[t+256*i]=local[i];
  __syncthreads();
  // Only three stages exchange values across warps.
  for(int width=32;width<=128;width<<=1){
    for(int b=t;b<dim/2;b+=256){
      int group=b/width,offset=b-group*width;
      int left=group*width*2+offset,right=left+width;
      float a=values[left],v=values[right];values[left]=a+v;values[right]=a-v;
    }
    __syncthreads();
  }
  #pragma unroll
  for(int i=0;i<8;i++)local[i]=values[t+256*i];
  // Dimensions256/512/1024 are eight register values owned by one thread.
  #pragma unroll
  for(int width=1;width<8;width<<=1){
    #pragma unroll
    for(int b=0;b<4;b++){
      int group=b/width,offset=b-group*width;
      int left=group*width*2+offset,right=left+width;
      float a=local[left],v=local[right];local[left]=a+v;local[right]=a-v;
    }
  }
  // All reads finished before shared storage is reused for packing.
  __syncthreads();
  #pragma unroll
  for(int i=0;i<8;i++)values[t+256*i]=local[i];
  __syncthreads();
  for(int j=t;j<dim/2;j+=256){
    uint8_t codes[2];
    #pragma unroll
    for(int k=0;k<2;k++){
      float v=values[2*j+k]*0.02209708691207961f;
      int index=v>bounds[7]?8:0;
      index+=v>bounds[index+3]?4:0;
      index+=v>bounds[index+1]?2:0;
      index+=v>bounds[index]?1:0;
      codes[k]=index;
    }
    packed[row*(dim/2)+j]=codes[0]|(codes[1]<<4);
  }
}

void rmsquant(torch::Tensor &packed, torch::Tensor &norms, torch::Tensor const &x,
    torch::Tensor const &weight, torch::Tensor const &permutation,
    torch::Tensor const &signs, torch::Tensor const &boundaries,
    double rms_eps, double quant_eps) {
  TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type()==torch::kBFloat16
    && x.dim()==2 && x.size(1)==2048,"requires contiguous CUDA BF16 [rows,2048]");
  for(auto* t:{&weight,&permutation,&signs,&boundaries})
    TORCH_CHECK(t->device()==x.device() && t->is_contiguous() && t->dim()==1,"invalid constant device/layout");
  TORCH_CHECK(weight.scalar_type()==torch::kBFloat16 && weight.numel()==2048,"invalid RMS weight");
  TORCH_CHECK(permutation.scalar_type()==torch::kInt && permutation.numel()==2048,"invalid permutation");
  TORCH_CHECK(signs.scalar_type()==torch::kChar && signs.numel()==2048,"invalid signs");
  TORCH_CHECK(boundaries.scalar_type()==torch::kFloat && boundaries.numel()==15,"invalid boundaries");
  TORCH_CHECK(packed.device()==x.device() && packed.is_contiguous() && packed.scalar_type()==torch::kUInt8
    && packed.dim()==2 && packed.size(0)==x.size(0) && packed.size(1)==1024,"invalid packed output");
  TORCH_CHECK(norms.device()==x.device() && norms.is_contiguous() && norms.scalar_type()==torch::kFloat
    && norms.dim()==1 && norms.numel()==x.size(0),"invalid norm output");
  TORCH_CHECK(std::isfinite(rms_eps) && rms_eps>0 && std::isfinite(quant_eps) && quant_eps>0,"epsilons must be finite and positive");
  const at::cuda::OptionalCUDAGuard guard(device_of(x));
  if(!x.size(0))return;
  rmsquant_kernel<<<x.size(0),256,0,at::cuda::getCurrentCUDAStream()>>>(packed.data_ptr<uint8_t>(),norms.data_ptr<float>(),
    x.data_ptr<BF>(),weight.data_ptr<BF>(),permutation.data_ptr<int32_t>(),signs.data_ptr<int8_t>(),
    boundaries.data_ptr<float>(),rms_eps,quant_eps);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}