// Is 52 GB/s the wall? Compare our relu to a pure NT-copy (same 2N traffic =
// the achievable read+write ceiling), and test transparent huge pages. 24 threads.
#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <immintrin.h>
#include <omp.h>
#include <sys/mman.h>
#include <chrono>
#include <algorithm>

static void relu_nt(float* __restrict o, 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+=16){
    _mm256_stream_ps(o+j,   _mm256_max_ps(_mm256_loadu_ps(in+j),   z));
    _mm256_stream_ps(o+j+8, _mm256_max_ps(_mm256_loadu_ps(in+j+8), z));
  }
  _mm_sfence();
}
// pure NT copy: load + NT store, no compute -> the read+write BW ceiling (2N traffic)
static void copy_nt(float* __restrict o, const float* __restrict in, size_t n){
  #pragma omp parallel for schedule(static)
  for(size_t j=0;j<n;j+=16){
    _mm256_stream_ps(o+j,   _mm256_loadu_ps(in+j));
    _mm256_stream_ps(o+j+8, _mm256_loadu_ps(in+j+8));
  }
  _mm_sfence();
}

template<class F> static double bw(F f,float*o,const float*in,size_t n,double gb){
  for(int w=0;w<3;w++) f(o,in,n); double best=1e30;
  for(int r=0;r<10;r++){ auto a=std::chrono::high_resolution_clock::now(); f(o,in,n);
    auto b=std::chrono::high_resolution_clock::now();
    best=std::min(best,std::chrono::duration<double>(b-a).count()); }
  return gb/best;
}

int main(){
  omp_set_num_threads(24);
  const size_t n=256ull*1024*1024, bytes=n*4; double gb=2.0*bytes/1e9;
  // plain
  float* in =(float*)aligned_alloc(64,bytes); float* out=(float*)aligned_alloc(64,bytes);
  for(size_t i=0;i<n;++i) in[i]=((i&1)?-1.f:1.f)*float(i%97);
  printf("24 threads, AVX2 + NT-store:\n");
  printf("  relu                 %.0f GB/s\n", bw(relu_nt,out,in,n,gb));
  printf("  pure NT-copy (ceil)  %.0f GB/s   <- achievable read+write wall\n", bw(copy_nt,out,in,n,gb));
  // huge pages
  float* inh =(float*)aligned_alloc(2*1024*1024,bytes);
  float* outh=(float*)aligned_alloc(2*1024*1024,bytes);
  madvise(inh,bytes,MADV_HUGEPAGE); madvise(outh,bytes,MADV_HUGEPAGE);
  for(size_t i=0;i<n;++i) inh[i]=in[i];
  printf("  relu + hugepages     %.0f GB/s\n", bw(relu_nt,outh,inh,n,gb));
  return 0;
}