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

#include <THC/THCAtomics.cuh>

#include "utils.h"
#include "vectorization.cuh"

template <typename scalar_t>
__global__ void sampling_scaling_penalties_kernel(const scalar_t* logits, const scalar_t* scaling_penalties,
                                                  scalar_t* output, const int32_t numel) {
  const int32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
  const int32_t stride = blockDim.x * gridDim.x;

  auto const* vectorized_logits = reinterpret_cast<vec4_t<scalar_t> const*>(logits);
  auto const* vectorized_penalties = reinterpret_cast<vec4_t<scalar_t> const*>(scaling_penalties);
  auto* vectorized_output = reinterpret_cast<vec4_t<scalar_t>*>(output);

  const int32_t num_vec_elems = numel >> 2;

#pragma unroll 4
  for (int32_t i = tid; i < num_vec_elems; i += stride) {
    vec4_t<scalar_t> logits_vec = vectorized_logits[i];
    vec4_t<scalar_t> penalties_vec = vectorized_penalties[i];
    vec4_t<scalar_t> out_vec;

    out_vec.x = logits_vec.x > 0 ? logits_vec.x / penalties_vec.x : logits_vec.x * penalties_vec.x;
    out_vec.y = logits_vec.y > 0 ? logits_vec.y / penalties_vec.y : logits_vec.y * penalties_vec.y;
    out_vec.z = logits_vec.z > 0 ? logits_vec.z / penalties_vec.z : logits_vec.z * penalties_vec.z;
    out_vec.w = logits_vec.w > 0 ? logits_vec.w / penalties_vec.w : logits_vec.w * penalties_vec.w;

    vectorized_output[i] = out_vec;
  }

  const int32_t start_idx = num_vec_elems * 4;
  for (int32_t i = start_idx + tid; i < numel; i += stride) {
    scalar_t logit = logits[i];
    scalar_t penalty = scaling_penalties[i];
    output[i] = logit > 0 ? logit / penalty : logit * penalty;
  }
}

torch::Tensor sampling_scaling_penalties(const torch::Tensor& logits, const torch::Tensor& scaling_penalties) {
  auto output = torch::empty_like(logits);
  const auto numel = logits.numel();
  const int threads = 512;

  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  AT_DISPATCH_FLOATING_TYPES_AND2(
      at::ScalarType::Half, at::ScalarType::BFloat16, logits.scalar_type(), "sampling_scaling_penalties_kernel", ([&] {
        const int blocks = (numel + threads * 4 - 1) / (threads * 4);
        sampling_scaling_penalties_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
            logits.data_ptr<scalar_t>(), scaling_penalties.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(), numel);
      }));

  return output;
}
