#include <torch/all.h>

#include <cstddef>
#include <cstdint>

#if defined(__x86_64__) || defined(_M_X64)
#include <immintrin.h>
#define RELU_X86 1
#endif

// Multi-dtype CPU ReLU: float32 / float16 / bfloat16 / int8.
// ReLU is memory-bound, so the kernel maximizes effective bandwidth:
//   * AVX2 (256-bit) vectorized, runtime-dispatched (falls back to a portable
//     auto-vectorized scalar loop if the CPU lacks AVX2).
//   * Non-temporal streaming stores when the output is 32B-aligned (skips
//     write-allocate/RFO traffic) -> ~2.4x the old SSE-1-thread backend.
//   * OpenMP across cores (active when the build links -fopenmp; correct serial
//     otherwise).
// fp16/bf16 relu = "zero the lane if the sign bit is set" (both are 16-bit
// sign-magnitude floats), done as branchless integer SIMD -- no fp conversion.
// int8 relu = signed byte max with 0 (_mm256_max_epi8).

namespace {

// ---------- scalar (any arch; auto-vectorizes with -O3) ----------
inline float    relu_s(float x)    { return x > 0.f ? x : 0.f; }
inline uint16_t relu_h(uint16_t x) { return (x & 0x8000u) ? uint16_t(0) : x; }  // fp16 & bf16
inline int8_t   relu_b(int8_t x)   { return x > 0 ? x : int8_t(0); }

void relu_f32_scalar(float* o, const float* in, int64_t n) {
  #pragma omp parallel for schedule(static)
  for (int64_t i = 0; i < n; ++i) o[i] = relu_s(in[i]);
}
void relu_h16_scalar(uint16_t* o, const uint16_t* in, int64_t n) {
  #pragma omp parallel for schedule(static)
  for (int64_t i = 0; i < n; ++i) o[i] = relu_h(in[i]);
}
void relu_i8_scalar(int8_t* o, const int8_t* in, int64_t n) {
  #pragma omp parallel for schedule(static)
  for (int64_t i = 0; i < n; ++i) o[i] = relu_b(in[i]);
}

#ifdef RELU_X86
// ---------- AVX2 paths (NT store when 32B-aligned) ----------
__attribute__((target("avx2")))
void relu_f32_avx2(float* o, const float* in, int64_t n) {
  const int W = 8; int64_t nv = n / W;
  const __m256 z = _mm256_setzero_ps();
  if ((reinterpret_cast<uintptr_t>(o) & 31u) == 0) {
    #pragma omp parallel for schedule(static)
    for (int64_t v = 0; v < nv; ++v)
      _mm256_stream_ps(o + v*W, _mm256_max_ps(_mm256_loadu_ps(in + v*W), z));
    _mm_sfence();
  } else {
    #pragma omp parallel for schedule(static)
    for (int64_t v = 0; v < nv; ++v)
      _mm256_storeu_ps(o + v*W, _mm256_max_ps(_mm256_loadu_ps(in + v*W), z));
  }
  for (int64_t i = nv*W; i < n; ++i) o[i] = relu_s(in[i]);
}
__attribute__((target("avx2")))
void relu_h16_avx2(uint16_t* o, const uint16_t* in, int64_t n) {
  const int W = 16; int64_t nv = n / W;
  if ((reinterpret_cast<uintptr_t>(o) & 31u) == 0) {
    #pragma omp parallel for schedule(static)
    for (int64_t v = 0; v < nv; ++v) {
      __m256i x = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(in + v*W));
      __m256i neg = _mm256_srai_epi16(x, 15);            // 0xFFFF where sign set
      _mm256_stream_si256(reinterpret_cast<__m256i*>(o + v*W),
                          _mm256_andnot_si256(neg, x));   // 0 if negative else x
    }
    _mm_sfence();
  } else {
    #pragma omp parallel for schedule(static)
    for (int64_t v = 0; v < nv; ++v) {
      __m256i x = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(in + v*W));
      __m256i neg = _mm256_srai_epi16(x, 15);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(o + v*W),
                          _mm256_andnot_si256(neg, x));
    }
  }
  for (int64_t i = nv*W; i < n; ++i) o[i] = relu_h(in[i]);
}
__attribute__((target("avx2")))
void relu_i8_avx2(int8_t* o, const int8_t* in, int64_t n) {
  const int W = 32; int64_t nv = n / W;
  const __m256i z = _mm256_setzero_si256();
  if ((reinterpret_cast<uintptr_t>(o) & 31u) == 0) {
    #pragma omp parallel for schedule(static)
    for (int64_t v = 0; v < nv; ++v)
      _mm256_stream_si256(reinterpret_cast<__m256i*>(o + v*W),
        _mm256_max_epi8(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(in + v*W)), z));
    _mm_sfence();
  } else {
    #pragma omp parallel for schedule(static)
    for (int64_t v = 0; v < nv; ++v)
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(o + v*W),
        _mm256_max_epi8(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(in + v*W)), z));
  }
  for (int64_t i = nv*W; i < n; ++i) o[i] = relu_b(in[i]);
}
#endif  // RELU_X86

inline bool has_avx2() {
#ifdef RELU_X86
  return __builtin_cpu_supports("avx2");
#else
  return false;
#endif
}

void relu_f32(float* o, const float* in, int64_t n) {
#ifdef RELU_X86
  if (has_avx2()) { relu_f32_avx2(o, in, n); return; }
#endif
  relu_f32_scalar(o, in, n);
}
void relu_h16(uint16_t* o, const uint16_t* in, int64_t n) {
#ifdef RELU_X86
  if (has_avx2()) { relu_h16_avx2(o, in, n); return; }
#endif
  relu_h16_scalar(o, in, n);
}
void relu_i8(int8_t* o, const int8_t* in, int64_t n) {
#ifdef RELU_X86
  if (has_avx2()) { relu_i8_avx2(o, in, n); return; }
#endif
  relu_i8_scalar(o, in, n);
}

}  // namespace

void relu(torch::Tensor &out, torch::Tensor const &input) {
  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
  TORCH_CHECK(out.is_contiguous(), "output must be contiguous");
  TORCH_CHECK(input.numel() == out.numel(),
              "Input and output tensors must have the same number of elements");
  TORCH_CHECK(input.scalar_type() == out.scalar_type(),
              "Input and output tensors must have the same dtype");
  const int64_t n = input.numel();
  switch (input.scalar_type()) {
    case torch::kFloat32:
      relu_f32(out.data_ptr<float>(), input.data_ptr<float>(), n);
      break;
    case torch::kHalf:
      relu_h16(reinterpret_cast<uint16_t*>(out.data_ptr<at::Half>()),
               reinterpret_cast<const uint16_t*>(input.data_ptr<at::Half>()), n);
      break;
    case torch::kBFloat16:
      relu_h16(reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>()),
               reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>()), n);
      break;
    case torch::kChar:
      relu_i8(out.data_ptr<int8_t>(), input.data_ptr<int8_t>(), n);
      break;
    default:
      TORCH_CHECK(false, "relu_cpu: unsupported dtype ", input.scalar_type(),
                  " (supported: float32, float16, bfloat16, int8)");
  }
}