// Test SYCL on the Intel iGPU: naive (1 elem/work-item, what the xpu backend does)
// vs vectorized (16 bytes/work-item, coalesced) across fp32/fp16-bf16/int8.
#include <sycl/sycl.hpp>
#include <cstdio>
#include <cstdint>
#include <vector>
#include <chrono>
#include <algorithm>
using namespace sycl;

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

template<class T> T fillval(size_t i);
template<> float    fillval<float>(size_t i)   { return ((i&1)?-1.f:1.f)*float(i%97); }
template<> uint16_t fillval<uint16_t>(size_t i){ return (uint16_t)((i&1)?(0x8000u|(i%200)):(i%200)); }
template<> int8_t   fillval<int8_t>(size_t i)  { return (int8_t)((i&1)?-(int)(i%100):(int)(i%100)); }

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

template<class F> double best_of(sycl::queue&q, F f){
  for(int w=0;w<5;w++) f(); q.wait();
  double best=1e30;
  for(int r=0;r<12;r++){ auto a=std::chrono::high_resolution_clock::now(); f(); q.wait();
    auto b=std::chrono::high_resolution_clock::now();
    best=std::min(best,std::chrono::duration<double>(b-a).count()); }
  return best;
}

template<typename T>
void run(sycl::queue& q, const char* name){
  const size_t n=64ull*1024*1024; const int C=16/sizeof(T); const size_t nv=n/C;
  T* in=malloc_device<T>(n,q); T* out=malloc_device<T>(n,q);
  std::vector<T> h(n),o(n); for(size_t i=0;i<n;++i) h[i]=fillval<T>(i);
  q.memcpy(in,h.data(),n*sizeof(T)).wait();
  double gb=2.0*n*sizeof(T)/1e9;

  auto naive=[&](){ q.parallel_for(range<1>(n),[=](id<1> id){ size_t i=id[0]; out[i]=relu_dev(in[i]); }); };
  auto vec=[&](){
    auto* vin=reinterpret_cast<const V16*>(in); auto* vout=reinterpret_cast<V16*>(out);
    q.parallel_for(range<1>(nv),[=](id<1> id){
      size_t i=id[0]; V16 raw=vin[i]; T* e=reinterpret_cast<T*>(&raw);
      #pragma unroll
      for(int k=0;k<C;k++) e[k]=relu_dev(e[k]);
      vout[i]=raw;
    });
  };
  // correctness (vec path)
  vec(); q.wait(); q.memcpy(o.data(),out,n*sizeof(T)).wait();
  bool ok=true; for(size_t i=0;i<n&&ok;++i) if(o[i]!=relu_dev(h[i])) ok=false;

  double tn=best_of(q,naive), tv=best_of(q,vec);
  printf("  %-10s %dB | naive %5.1f GB/s | vec16 %5.1f GB/s (%.1fx) | %s\n",
         name,(int)sizeof(T), gb/tn, gb/tv, tn/tv, ok?"OK":"FAIL");
  free(in,q); free(out,q);
}

int main(){
  queue q{gpu_selector_v};
  printf("iGPU: %s (%u EUs)\n", q.get_device().get_info<info::device::name>().c_str(),
         q.get_device().get_info<info::device::max_compute_units>());
  run<float>(q,"fp32");
  run<uint16_t>(q,"fp16/bf16");
  run<int8_t>(q,"int8");
  return 0;
}