#include <sycl/sycl.hpp>
#include <torch/torch.h>

#include <cstdint>

using namespace sycl;

// Multi-dtype XPU (SYCL) ReLU: float32 / float16 / bfloat16 / int8.
// Vectorized: each work-item processes a 16-byte (128-bit) chunk -> coalesced,
// dword-aligned access. On the Intel iGPU this reaches the memory wall for every
// dtype; the naive 1-element/work-item kernel left sub-dword types issue-bound
// (int8 ~2.6x slower, fp16 ~1.2x). fp16 & bf16 are 16-bit sign-magnitude floats,
// so relu = "zero the lane if the sign bit is set" on the raw uint16.

namespace {

struct alignas(16) V16 { uint32_t w[4]; };  // 16-byte vector

template <typename T> inline T relu_dev(T x);
template <> inline float    relu_dev<float>(float x)       { return x > 0.f ? x : 0.f; }
template <> inline uint16_t relu_dev<uint16_t>(uint16_t x) { return (x & 0x8000u) ? uint16_t(0) : x; }
template <> inline int8_t   relu_dev<int8_t>(int8_t x)     { return x > 0 ? x : int8_t(0); }

template <typename T>
void launch(sycl::queue& q, T* o, const T* in, int64_t n) {
  constexpr int C = 16 / sizeof(T);           // elements per 16-byte chunk
  const int64_t nv = n / C;
  const int64_t tail_start = nv * C;

  if (nv > 0) {
    auto* vin  = reinterpret_cast<const V16*>(in);
    auto* vout = reinterpret_cast<V16*>(o);
    q.parallel_for(range<1>(nv), [=](id<1> idx) {
      int64_t i = idx[0];
      V16 raw = vin[i];
      T* e = reinterpret_cast<T*>(&raw);
#pragma unroll
      for (int k = 0; k < C; ++k) e[k] = relu_dev<T>(e[k]);
      vout[i] = raw;
    });
  }
  if (tail_start < n) {
    const int64_t rem = n - tail_start;
    q.parallel_for(range<1>(rem), [=](id<1> idx) {
      int64_t i = tail_start + idx[0];
      o[i] = relu_dev<T>(in[i]);
    });
  }
  q.wait();
}

}  // namespace

void relu(torch::Tensor& out, const torch::Tensor& input) {
  TORCH_CHECK(input.device().is_xpu(), "input must be a XPU tensor");
  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
  TORCH_CHECK(out.is_contiguous(), "output must be contiguous");
  TORCH_CHECK(input.sizes() == out.sizes(),
              "Tensors must have the same shape. Got input shape: ",
              input.sizes(), " and output shape: ", out.sizes());
  TORCH_CHECK(input.scalar_type() == out.scalar_type(),
              "Tensors must have the same data type. Got input dtype: ",
              input.scalar_type(), " and output dtype: ", out.scalar_type());
  TORCH_CHECK(input.device() == out.device(),
              "Tensors must be on the same device. Got input device: ",
              input.device(), " and output device: ", out.device());

  sycl::queue queue;
  const int64_t n = input.numel();
  switch (input.scalar_type()) {
    case torch::kFloat:
      launch<float>(queue, out.data_ptr<float>(), input.data_ptr<float>(), n);
      break;
    case torch::kHalf:
      launch<uint16_t>(queue, reinterpret_cast<uint16_t*>(out.data_ptr<at::Half>()),
                       reinterpret_cast<const uint16_t*>(input.data_ptr<at::Half>()), n);
      break;
    case torch::kBFloat16:
      launch<uint16_t>(queue, reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>()),
                       reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>()), n);
      break;
    case torch::kChar:
      launch<int8_t>(queue, out.data_ptr<int8_t>(), input.data_ptr<int8_t>(), n);
      break;
    default:
      TORCH_CHECK(false, "relu_xpu: unsupported dtype ", input.scalar_type(),
                  " (supported: float32, float16, bfloat16, int8)");
  }
}