// Push CPU relu past 50 GB/s on i9-13900K: sweep thread count x per-thread ILP,
// AVX2 + non-temporal stores. Find the bandwidth-optimal config.
#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <immintrin.h>
#include <omp.h>
#include <chrono>
#include <algorithm>

// AVX2 NT-store relu, U vectors (8 floats each) per thread step.
template<int U>
static void relu_nt(float* __restrict out, const float* __restrict in, size_t n){
  const __m256 z=_mm256_setzero_ps();
  #pragma omp parallel for schedule(static)
  for(size_t j=0;j<n;j+=(size_t)8*U){
    for(int u=0;u<U;++u){
      __m256 v=_mm256_loadu_ps(in+j+8*u);
      _mm256_stream_ps(out+j+8*u, _mm256_max_ps(v,z));
    }
  }
  _mm_sfence();
}

template<class F>
static 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<8;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(){
  const size_t n=256ull*1024*1024;            // 1 GB/array, divisible by 8*8
  float* in =(float*)aligned_alloc(64,n*4);
  float* out=(float*)aligned_alloc(64,n*4);
  for(size_t i=0;i<n;++i) in[i]=((i&1)?-1.f:1.f)*float(i%97);
  double gb=2.0*n*4/1e9;

  // correctness of the unrolled NT kernel
  for(size_t i=0;i<n;++i) out[i]=-7; relu_nt<4>(out,in,n);
  bool ok=true; for(size_t i=0;i<n&&ok;i++){ float e=in[i]>0?in[i]:0; if(out[i]!=e) ok=false; }
  printf("correctness: %s\n", ok?"PASS":"FAIL");

  printf("AVX2 + NT-store, sweep threads x ILP (GB/s).  DDR5 wall ~90 GB/s\n");
  printf("threads |  U=1   U=2   U=4   U=8\n");
  int ts[]={4,8,12,16,24,32};
  double bestbw=0; int bt=0,bu=0;
  for(int t: ts){
    omp_set_num_threads(t);
    double b1=gb/bench(relu_nt<1>,out,in,n);
    double b2=gb/bench(relu_nt<2>,out,in,n);
    double b4=gb/bench(relu_nt<4>,out,in,n);
    double b8=gb/bench(relu_nt<8>,out,in,n);
    printf("  %2d    | %5.0f %5.0f %5.0f %5.0f\n", t,b1,b2,b4,b8);
    double m[4]={b1,b2,b4,b8}; int uu[4]={1,2,4,8};
    for(int k=0;k<4;k++) if(m[k]>bestbw){bestbw=m[k];bt=t;bu=uu[k];}
  }
  printf(">> BEST: %.0f GB/s @ %d threads, U=%d  (%.0f%% of ~90 GB/s wall)\n", bestbw,bt,bu,100*bestbw/90.0);
  free(in); free(out);
  return 0;
}