// RTX 4090 (sm_89) float32 ReLU design-space sweep.
// Sweeps vector width x unroll x store-policy x load-policy x block size across
// tensor sizes, verifies correctness, reports best effective bandwidth per size.
#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <cmath>
#include <string>
#include <vector>
#include <functional>
#include <algorithm>
#include <cuda_runtime.h>

#define CK(x) do{ cudaError_t e=(x); if(e){ printf("CUDA err %s:%d: %s\n",__FILE__,__LINE__,cudaGetErrorString(e)); exit(1);} }while(0)

__device__ __forceinline__ float relu1(float x){ return x > 0.0f ? x : 0.0f; }

template<int STORE> __device__ __forceinline__ void st4(float4* p, float4 v){
  if (STORE) __stcs(p, v); else *p = v;
}
template<int LOAD> __device__ __forceinline__ float4 ld4(const float4* p){
  return LOAD ? __ldcs(p) : *p;
}
template<int STORE> __device__ __forceinline__ void st2(float2* p, float2 v){
  if (STORE) __stcs(p, v); else *p = v;
}
template<int LOAD> __device__ __forceinline__ float2 ld2(const float2* p){
  return LOAD ? __ldcs(p) : *p;
}

// ---- float4 kernel, templated unroll/store/load ----
template<int U,int S,int L>
__global__ void kv4(float4* __restrict__ out, const float4* __restrict__ in, long long n4){
  long long stride=(long long)gridDim.x*blockDim.x;
  for(long long b=(long long)blockIdx.x*blockDim.x+threadIdx.x; b<n4; b+=stride*U){
    float4 r[U];
    #pragma unroll
    for(int u=0;u<U;u++){ long long i=b+(long long)u*stride; if(i<n4) r[u]=ld4<L>(in+i); }
    #pragma unroll
    for(int u=0;u<U;u++){ r[u].x=relu1(r[u].x); r[u].y=relu1(r[u].y); r[u].z=relu1(r[u].z); r[u].w=relu1(r[u].w); }
    #pragma unroll
    for(int u=0;u<U;u++){ long long i=b+(long long)u*stride; if(i<n4) st4<S>(out+i, r[u]); }
  }
}
// ---- float2 kernel ----
template<int U,int S,int L>
__global__ void kv2(float2* __restrict__ out, const float2* __restrict__ in, long long n2){
  long long stride=(long long)gridDim.x*blockDim.x;
  for(long long b=(long long)blockIdx.x*blockDim.x+threadIdx.x; b<n2; b+=stride*U){
    float2 r[U];
    #pragma unroll
    for(int u=0;u<U;u++){ long long i=b+(long long)u*stride; if(i<n2) r[u]=ld2<L>(in+i); }
    #pragma unroll
    for(int u=0;u<U;u++){ r[u].x=relu1(r[u].x); r[u].y=relu1(r[u].y); }
    #pragma unroll
    for(int u=0;u<U;u++){ long long i=b+(long long)u*stride; if(i<n2) st2<S>(out+i, r[u]); }
  }
}
// ---- scalar grid-stride reference ----
__global__ void kscalar(float* __restrict__ out, const float* __restrict__ in, long long n){
  long long stride=(long long)gridDim.x*blockDim.x;
  for(long long i=(long long)blockIdx.x*blockDim.x+threadIdx.x; i<n; i+=stride) out[i]=relu1(in[i]);
}
__global__ void ktail(float* out,const float* in,long long start,long long n){
  long long i=start+(long long)blockIdx.x*blockDim.x+threadIdx.x; if(i<n) out[i]=relu1(in[i]);
}

static long long ceil_div(long long a,long long b){ return (a+b-1)/b; }

template<int U,int S,int L>
void launch_v4(float* out,const float* in,long long n,int threads){
  long long n4=n/4, tail=n4*4;
  if(n4>0){ long long blk=ceil_div(n4,(long long)threads*U);
    kv4<U,S,L><<<(unsigned)std::min<long long>(blk, 1u<<31 - 1), threads>>>((float4*)out,(const float4*)in,n4); }
  if(tail<n){ long long b=ceil_div(n-tail,threads); ktail<<<(unsigned)b,threads>>>(out,in,tail,n); }
}
template<int U,int S,int L>
void launch_v2(float* out,const float* in,long long n,int threads){
  long long n2=n/2, tail=n2*2;
  if(n2>0){ long long blk=ceil_div(n2,(long long)threads*U);
    kv2<U,S,L><<<(unsigned)std::min<long long>(blk, 1u<<31 - 1), threads>>>((float2*)out,(const float2*)in,n2); }
  if(tail<n){ long long b=ceil_div(n-tail,threads); ktail<<<(unsigned)b,threads>>>(out,in,tail,n); }
}
void launch_scalar(float* out,const float* in,long long n,int threads){
  // full-ish grid stride: cap blocks to a large number; one pass covers it
  long long blk=std::min<long long>(ceil_div(n,threads), 262144);
  kscalar<<<(unsigned)blk,threads>>>(out,in,n);
}

struct Cfg { std::string name; int threads; std::function<void(float*,const float*,long long,int)> fn; };

int main(){
  cudaDeviceProp p; CK(cudaGetDeviceProperties(&p,0));
  printf("GPU %s  SMs=%d  peakMemClk=%d kHz  busWidth=%d-bit\n", p.name, p.multiProcessorCount, p.memoryClockRate, p.memoryBusWidth);
  double peak_gbps = 2.0 * p.memoryClockRate * 1e3 * (p.memoryBusWidth/8) / 1e9; // 2x for DDR
  printf("theoretical DRAM peak ~%.0f GB/s\n\n", peak_gbps);

  std::vector<Cfg> cfgs;
  for(int t : {128,256,512}){
    char b[64];
    #define ADD4(U,S,L,tag) { snprintf(b,64,"v4 u%d %-8s blk%d",U,tag,t); cfgs.push_back({b,t,launch_v4<U,S,L>}); }
    ADD4(1,0,0,"def")  ADD4(2,0,0,"def")  ADD4(4,0,0,"def")
    ADD4(1,1,0,"st.cs")  ADD4(2,1,0,"st.cs")  ADD4(4,1,0,"st.cs")
    ADD4(1,0,1,"ld.cs")  ADD4(1,1,1,"ld+st.cs") ADD4(2,1,1,"ld+st.cs")
    #undef ADD4
  }
  // a couple of reference points
  for(int t : {256,512}){ char b[64]; snprintf(b,64,"v2 u1 def   blk%d",t); cfgs.push_back({b,t,launch_v2<1,0,0>}); }
  for(int t : {256}){ char b[64]; snprintf(b,64,"v2 u2 st.cs blk%d",t); cfgs.push_back({b,t,launch_v2<2,1,0>}); }
  cfgs.push_back({"scalar      blk256",256,launch_scalar});

  // sizes (in elements): include 1024^2 and 4096^2 plus a sweep
  std::vector<long long> sizes = { 1<<16, 1<<18, 1024LL*1024, 1<<22, 4096LL*4096, 1<<26, 1<<27 };

  const int WARM=50, ITER=200;
  for(long long n : sizes){
    size_t bytes=n*sizeof(float);
    float *din,*dout; CK(cudaMalloc(&din,bytes)); CK(cudaMalloc(&dout,bytes));
    // init host with mix of + and - (and a few exact patterns)
    std::vector<float> h(n);
    for(long long i=0;i<n;i++) h[i]= ((i*1103515245u+12345u)&1) ? -float((i%97)-48) : float((i%97)-48);
    CK(cudaMemcpy(din,h.data(),bytes,cudaMemcpyHostToDevice));
    double gb = 2.0*bytes/1e9;
    printf("=== N=%lld (%.1f MB/array, rw %.1f MB) ===\n", n, bytes/1e6, gb*1e3);

    std::string best; double bestbw=0;
    for(auto& c : cfgs){
      CK(cudaMemset(dout,0xff,bytes));
      c.fn(dout,din,n,c.threads); CK(cudaGetLastError()); CK(cudaDeviceSynchronize());
      // correctness
      std::vector<float> ho(n); CK(cudaMemcpy(ho.data(),dout,bytes,cudaMemcpyDeviceToHost));
      bool ok=true; for(long long i=0;i<n && ok;i++){ float e=h[i]>0?h[i]:0; if(ho[i]!=e) ok=false; }
      // timing
      for(int w=0;w<WARM;w++) c.fn(dout,din,n,c.threads);
      CK(cudaDeviceSynchronize());
      cudaEvent_t a,bb; cudaEventCreate(&a); cudaEventCreate(&bb);
      cudaEventRecord(a);
      for(int it=0;it<ITER;it++) c.fn(dout,din,n,c.threads);
      cudaEventRecord(bb); CK(cudaEventSynchronize(bb));
      float ms=0; cudaEventElapsedTime(&ms,a,bb); ms/=ITER;
      double bw=gb/(ms/1e3);
      cudaEventDestroy(a); cudaEventDestroy(bb);
      printf("  %-22s %7.4f ms  %6.0f GB/s  %s\n", c.name.c_str(), ms, bw, ok?"ok":"FAIL");
      if(ok && bw>bestbw){ bestbw=bw; best=c.name; }
    }
    printf("  >> BEST: %s  %.0f GB/s  (%.1f%% of peak)\n\n", best.c_str(), bestbw, 100*bestbw/peak_gbps);
    cudaFree(din); cudaFree(dout);
  }
  return 0;
}