// ReLU throughput vs dtype on RTX 4090 (sm_89). ReLU is memory-bound, so at the
// ~byte-bandwidth wall, ELEMENT throughput scales ~1/sizeof(dtype). Each kernel
// does 128-bit (int4) vectorized loads/stores -> (16/sizeof) elements per thread,
// so every dtype gets the same memory pattern; only per-element relu + count vary.
// Covers fp64/fp32/fp16/bf16/fp8(e4m3,e5m2)/int64/int32/int16/int8/uint8.
#include <cstdio>
#include <cstdint>
#include <vector>
#include <string>
#include <algorithm>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#define CK(x) do{cudaError_t e=(x); if(e){printf("ERR %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e));exit(1);}}while(0)

// ---- device relu overloads ----
__device__ __forceinline__ double  drelu(double x){return x>0.0?x:0.0;}
__device__ __forceinline__ float   drelu(float x){return x>0.f?x:0.f;}
__device__ __forceinline__ __half  drelu(__half x){__half z=__float2half(0.f); return __hgt(x,z)?x:z;}
__device__ __forceinline__ __nv_bfloat16 drelu(__nv_bfloat16 x){__nv_bfloat16 z=__float2bfloat16(0.f); return __hgt(x,z)?x:z;}
__device__ __forceinline__ long long    drelu(long long x){return x>0?x:0;}
__device__ __forceinline__ int          drelu(int x){return x>0?x:0;}
__device__ __forceinline__ short        drelu(short x){return x>0?x:0;}
__device__ __forceinline__ signed char  drelu(signed char x){return x>0?x:0;}
__device__ __forceinline__ unsigned char drelu(unsigned char x){return x;}   // max(x,0)=x
__device__ __forceinline__ __nv_fp8_e4m3 drelu(__nv_fp8_e4m3 x){float f=__half2float(__half(x)); return __nv_fp8_e4m3(f>0.f?f:0.f);}
__device__ __forceinline__ __nv_fp8_e5m2 drelu(__nv_fp8_e5m2 x){float f=__half2float(__half(x)); return __nv_fp8_e5m2(f>0.f?f:0.f);}

template<typename T>
__global__ void reluK(int4* __restrict__ out, const int4* __restrict__ in, long long nvec){
  constexpr int VEC = 16/sizeof(T);
  long long stride=(long long)gridDim.x*blockDim.x;
  for(long long i=(long long)blockIdx.x*blockDim.x+threadIdx.x;i<nvec;i+=stride){
    int4 raw=in[i];
    T* e=reinterpret_cast<T*>(&raw);
    #pragma unroll
    for(int k=0;k<VEC;k++) e[k]=drelu(e[k]);
    out[i]=raw;
  }
}

// ---- host helpers ----
static float h_tof(double x){return (float)x;}
static float h_tof(float x){return x;}
static float h_tof(__half x){return __half2float(x);}
static float h_tof(__nv_bfloat16 x){return __bfloat162float(x);}
static float h_tof(long long x){return (float)x;}
static float h_tof(int x){return (float)x;}
static float h_tof(short x){return (float)x;}
static float h_tof(signed char x){return (float)x;}
static float h_tof(unsigned char x){return (float)x;}
static float h_tof(__nv_fp8_e4m3 x){return __half2float(__half(x));}
static float h_tof(__nv_fp8_e5m2 x){return __half2float(__half(x));}

template<class T> T h_fromf(float v);
template<> double h_fromf<double>(float v){return (double)v;}
template<> float  h_fromf<float>(float v){return v;}
template<> __half h_fromf<__half>(float v){return __float2half(v);}
template<> __nv_bfloat16 h_fromf<__nv_bfloat16>(float v){return __float2bfloat16(v);}
template<> long long h_fromf<long long>(float v){return (long long)v;}
template<> int   h_fromf<int>(float v){return (int)v;}
template<> short h_fromf<short>(float v){return (short)v;}
template<> signed char h_fromf<signed char>(float v){return (signed char)v;}
template<> unsigned char h_fromf<unsigned char>(float v){return (unsigned char)v;}
template<> __nv_fp8_e4m3 h_fromf<__nv_fp8_e4m3>(float v){return __nv_fp8_e4m3(v);}
template<> __nv_fp8_e5m2 h_fromf<__nv_fp8_e5m2>(float v){return __nv_fp8_e5m2(v);}

double g_peak;

template<typename T>
void run(const char* name){
  printf("%-9s %dB |", name, (int)sizeof(T));
  for(int reg=0; reg<2; reg++){
    long long N = reg==0 ? (2LL<<20) : (64LL<<20);   // L2-resident vs DRAM
    int threads = reg==0 ? 128 : 512;
    long long nbytes=N*sizeof(T), nvec=nbytes/16;
    int4 *din,*dout; CK(cudaMalloc(&din,nbytes)); CK(cudaMalloc(&dout,nbytes));
    std::vector<T> h(N);
    for(long long i=0;i<N;i++){ float v=((i&1)?-1.f:1.f)*float(i%37); h[i]=h_fromf<T>(v); }
    CK(cudaMemcpy(din,h.data(),nbytes,cudaMemcpyHostToDevice));
    long long blk=std::min<long long>((nvec+threads-1)/threads, 1<<20);
    // correctness
    CK(cudaMemset(dout,0x5a,nbytes));
    reluK<T><<<(unsigned)blk,threads>>>(dout,din,nvec); CK(cudaDeviceSynchronize());
    std::vector<T> ho(N); CK(cudaMemcpy(ho.data(),dout,nbytes,cudaMemcpyDeviceToHost));
    bool ok=true; for(long long i=0;i<N&&ok;i++){ float ex=h_tof(h[i]); ex=ex>0?ex:0; if(h_tof(ho[i])!=ex) ok=false; }
    // timing
    for(int w=0;w<50;w++) reluK<T><<<(unsigned)blk,threads>>>(dout,din,nvec);
    CK(cudaDeviceSynchronize());
    cudaEvent_t a,b; cudaEventCreate(&a); cudaEventCreate(&b); cudaEventRecord(a);
    const int IT=200; for(int it=0;it<IT;it++) reluK<T><<<(unsigned)blk,threads>>>(dout,din,nvec);
    cudaEventRecord(b); CK(cudaEventSynchronize(b));
    float ms; cudaEventElapsedTime(&ms,a,b); ms/=IT;
    double gbps=2.0*nbytes/(ms/1e3)/1e9, gelem=N/(ms/1e3)/1e9;
    printf("  %s %4.0f GB/s %6.1f Gel/s%s |", reg==0?"L2":"DRAM", gbps, gelem, ok?"":" FAIL");
    cudaEventDestroy(a); cudaEventDestroy(b); cudaFree(din); cudaFree(dout);
  }
  printf("\n");
}

int main(){
  cudaDeviceProp p; CK(cudaGetDeviceProperties(&p,0));
  g_peak=2.0*p.memoryClockRate*1e3*(p.memoryBusWidth/8)/1e9;
  printf("%s  theoretical peak ~%.0f GB/s\n", p.name, g_peak);
  printf("dtype     sz | regime  bandwidth  elem-rate     | regime  bandwidth  elem-rate\n");
  printf("------------------------------------------------------------------------------\n");
  run<double>("fp64");
  run<float>("fp32");
  run<__half>("fp16");
  run<__nv_bfloat16>("bf16");
  run<__nv_fp8_e4m3>("fp8e4m3");
  run<__nv_fp8_e5m2>("fp8e5m2");
  run<long long>("int64");
  run<int>("int32");
  run<short>("int16");
  run<signed char>("int8");
  run<unsigned char>("uint8");
  return 0;
}