// CPU ReLU efficiency: current SSE-1-thread vs AVX2 auto-vec vs AVX2+OpenMP.
// Compile: icx -O3 -xHost -qopenmp cpu_relu.cpp -o cpu_relu   (i9-13900K = AVX2)
#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <immintrin.h>
#include <omp.h>
#include <chrono>
#include <algorithm>

// --- current backend: explicit SSE (128-bit = 4 floats), single thread ---
void relu_sse(float* out, const float* in, size_t n){
  size_t i=0;
  for(; i+4<=n; i+=4){
    __m128 v=_mm_loadu_ps(in+i);                 // (use unaligned: aligned would segfault)
    _mm_storeu_ps(out+i, _mm_max_ps(v, _mm_setzero_ps()));
  }
  for(; i<n; ++i) out[i]=in[i]>0.f?in[i]:0.f;
}
// --- clean loop, compiler auto-vectorizes to AVX2 (8 floats) with -xHost, 1 thread ---
void relu_autovec(float* __restrict out, const float* __restrict in, size_t n){
  for(size_t i=0;i<n;++i) out[i]=in[i]>0.f?in[i]:0.f;
}
// --- AVX2 + OpenMP across cores ---
void relu_omp(float* __restrict out, const float* __restrict in, size_t n){
  #pragma omp parallel for schedule(static)
  for(size_t i=0;i<n;++i) out[i]=in[i]>0.f?in[i]:0.f;
}
// --- AVX2 + OpenMP + non-temporal stores (skip write-allocate/RFO traffic) ---
void relu_omp_nt(float* __restrict out, const float* __restrict in, size_t n){
  #pragma omp parallel for schedule(static)
  for(size_t j=0;j<n;j+=8){
    __m256 v=_mm256_loadu_ps(in+j);                       // out is 64B-aligned, j*4 %32==0
    _mm256_stream_ps(out+j, _mm256_max_ps(v,_mm256_setzero_ps()));
  }
  _mm_sfence();
}

template<class F>
double bench(F f, float* out, const float* in, size_t n){
  for(int w=0;w<3;w++) f(out,in,n);
  double best=1e30;
  for(int r=0;r<10;r++){
    auto t0=std::chrono::high_resolution_clock::now();
    f(out,in,n);
    auto t1=std::chrono::high_resolution_clock::now();
    best=std::min(best,std::chrono::duration<double>(t1-t0).count());
  }
  return best;
}

int main(){
  printf("OpenMP max threads: %d\n", omp_get_max_threads());
  const size_t n=256ull*1024*1024;          // 1 GB/array -> DRAM-bound
  float* in =(float*)aligned_alloc(64, n*sizeof(float));
  float* out=(float*)aligned_alloc(64, n*sizeof(float));
  for(size_t i=0;i<n;++i) in[i]=((i&1)?-1.f:1.f)*float(i%97);
  double gb=2.0*n*sizeof(float)/1e9;

  auto report=[&](const char* name, double s){ printf("  %-22s %6.1f GB/s   (%.2f ms)\n", name, gb/s, s*1e3); };

  // verify correctness of each
  auto check=[&](void(*f)(float*,const float*,size_t),const char* nm){
    for(size_t i=0;i<n;++i) out[i]=-7;
    f(out,in,n);
    for(size_t i=0;i<n;++i){ float e=in[i]>0?in[i]:0; if(out[i]!=e){printf("  %s FAIL@%zu\n",nm,i);return;} }
  };
  check(relu_sse,"sse"); check(relu_autovec,"autovec"); check(relu_omp,"omp"); check(relu_omp_nt,"omp_nt");

  printf("CPU ReLU, 256M fp32 (1 GB/array, DRAM-bound):\n");
  report("SSE 4-wide 1-thread", bench(relu_sse, out, in, n));
  report("AVX2 auto-vec 1-thread", bench(relu_autovec, out, in, n));
  report("AVX2 + OpenMP", bench(relu_omp, out, in, n));
  report("AVX2 + OpenMP + NT-store", bench(relu_omp_nt, out, in, n));
  free(in); free(out);
  return 0;
}