// Why is torch's int8 relu faster in the L2 regime? Hypothesis: per-thread byte
// granularity (more threads = more MLP). Sweep bytes/thread {4,8,16} x block, int8,
// at 16.7M (33.4MB ws, L2-resident). Uses __vmaxs4. Ranks configs in one process.
#include <cstdio>
#include <cstdint>
#include <vector>
#include <string>
#include <functional>
#include <algorithm>
#include <cuda_runtime.h>
#define CK(x) do{cudaError_t e=(x); if(e){printf("ERR %s\n",cudaGetErrorString(e));exit(1);}}while(0)

// 4 bytes/thread (like torch: 4 int8 per thread)
__global__ void k4(unsigned* __restrict__ o,const unsigned* __restrict__ in,int64_t nu){
  int64_t s=(int64_t)gridDim.x*blockDim.x;
  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<nu;i+=s) o[i]=__vmaxs4(__ldcs(in+i),0u);
}
// 8 bytes/thread (uint2 = 8 int8)
__global__ void k8(uint2* __restrict__ o,const uint2* __restrict__ in,int64_t nu){
  int64_t s=(int64_t)gridDim.x*blockDim.x;
  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<nu;i+=s){
    uint2 v=__ldcs(in+i); v.x=__vmaxs4(v.x,0u); v.y=__vmaxs4(v.y,0u); __stcs(o+i,v);
  }
}
// 16 bytes/thread (int4 = 16 int8, our current)
__global__ void k16(int4* __restrict__ o,const int4* __restrict__ in,int64_t nu){
  int64_t s=(int64_t)gridDim.x*blockDim.x;
  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<nu;i+=s){
    int4 v=__ldcs(in+i);
    v.x=__vmaxs4(v.x,0u);v.y=__vmaxs4(v.y,0u);v.z=__vmaxs4(v.z,0u);v.w=__vmaxs4(v.w,0u);
    __stcs(o+i,v);
  }
}

int main(){
  cudaDeviceProp p; CK(cudaGetDeviceProperties(&p,0));
  double peak=2.0*p.memoryClockRate*1e3*(p.memoryBusWidth/8)/1e9;
  const int64_t N=4096LL*4096; const size_t bytes=N; // int8
  signed char *di,*doo; CK(cudaMalloc(&di,bytes)); CK(cudaMalloc(&doo,bytes));
  std::vector<signed char> h(N); for(int64_t i=0;i<N;i++) h[i]=(signed char)((i%2?-1:1)*(i%100));
  CK(cudaMemcpy(di,h.data(),bytes,cudaMemcpyHostToDevice));
  double gb=2.0*bytes/1e9;
  printf("int8 relu N=%lld (ws %.0fMB, L2-resident).  torch ref ~3560 GB/s (9383ns)\n",N,gb*1e3/2);

  struct C{std::string name;int bpt;int block;std::function<void(int)> run;};
  std::vector<C> cs;
  for(int blk: {64,128,256,512}){
    cs.push_back({"4B/thread blk"+std::to_string(blk),4,blk,[=](int b){ int64_t nu=bytes/4,g=(nu+b-1)/b; k4<<<(unsigned)std::min<int64_t>(g,1<<22),b>>>((unsigned*)doo,(const unsigned*)di,nu);} });
    cs.push_back({"8B/thread blk"+std::to_string(blk),8,blk,[=](int b){ int64_t nu=bytes/8,g=(nu+b-1)/b; k8<<<(unsigned)std::min<int64_t>(g,1<<22),b>>>((uint2*)doo,(const uint2*)di,nu);} });
    cs.push_back({"16B/thread blk"+std::to_string(blk),16,blk,[=](int b){ int64_t nu=bytes/16,g=(nu+b-1)/b; k16<<<(unsigned)std::min<int64_t>(g,1<<22),b>>>((int4*)doo,(const int4*)di,nu);} });
  }
  std::string best;double bbw=0;
  for(auto&c:cs){
    for(int w=0;w<100;w++) c.run(c.block); CK(cudaDeviceSynchronize());
    // min over a few timed batches to get the least-throttled estimate
    double bestms=1e9;
    for(int rep=0;rep<5;rep++){
      cudaEvent_t a,b;cudaEventCreate(&a);cudaEventCreate(&b);cudaEventRecord(a);
      for(int it=0;it<200;it++) c.run(c.block);
      cudaEventRecord(b);CK(cudaEventSynchronize(b));
      float ms;cudaEventElapsedTime(&ms,a,b);ms/=200; bestms=std::min(bestms,(double)ms);
      cudaEventDestroy(a);cudaEventDestroy(b);
    }
    double bw=gb/(bestms/1e3);
    printf("  %-18s %.0f GB/s (%.1f%% peak)\n",c.name.c_str(),bw,100*bw/peak);
    if(bw>bbw){bbw=bw;best=c.name;}
  }
  printf(">> BEST: %s  %.0f GB/s\n",best.c_str(),bbw);
  return 0;
}