// Round-2: test the levers NOT covered in sweep.cu against the current champion
// (float4 + __ldcs + __stcs, adaptive block). Goal: find any real gain or prove
// we are at the wall. Sizes = the two benchmark points (1M base / 16.7M large).
#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:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e));exit(1);}}while(0)
__device__ __forceinline__ float r1(float x){return x>0.f?x:0.f;}
__device__ __forceinline__ float4 r4(float4 v){v.x=r1(v.x);v.y=r1(v.y);v.z=r1(v.z);v.w=r1(v.w);return v;}

// champion: streaming load + streaming store
__global__ void k_cs(float4* __restrict__ o,const float4* __restrict__ i,long long n){
  long long s=(long long)gridDim.x*blockDim.x;
  for(long long t=(long long)blockIdx.x*blockDim.x+threadIdx.x;t<n;t+=s) __stcs(o+t,r4(__ldcs(i+t)));
}
// write-through store
__global__ void k_wt(float4* __restrict__ o,const float4* __restrict__ i,long long n){
  long long s=(long long)gridDim.x*blockDim.x;
  for(long long t=(long long)blockIdx.x*blockDim.x+threadIdx.x;t<n;t+=s) __stwt(o+t,r4(__ldcs(i+t)));
}
// plain default (baseline)
__global__ void k_def(float4* __restrict__ o,const float4* __restrict__ i,long long n){
  long long s=(long long)gridDim.x*blockDim.x;
  for(long long t=(long long)blockIdx.x*blockDim.x+threadIdx.x;t<n;t+=s) o[t]=r4(i[t]);
}
// launch_bounds hinted (champion + bounds)
__global__ void __launch_bounds__(256,6) k_lb(float4* __restrict__ o,const float4* __restrict__ i,long long n){
  long long s=(long long)gridDim.x*blockDim.x;
  for(long long t=(long long)blockIdx.x*blockDim.x+threadIdx.x;t<n;t+=s) __stcs(o+t,r4(__ldcs(i+t)));
}

typedef void(*KF)(float4*,const float4*,long long);
struct C{std::string name; KF f; int carveoutL1;};

int main(){
  cudaDeviceProp p; CK(cudaGetDeviceProperties(&p,0));
  double peak=2.0*p.memoryClockRate*1e3*(p.memoryBusWidth/8)/1e9;
  std::vector<C> cs={
    {"champion ld.cs+st.cs", k_cs, 0},
    {"+launch_bounds",       k_lb, 0},
    {"+preferL1 carveout",   k_cs, 1},
    {"write-through st.wt",   k_wt, 0},
    {"default (baseline)",    k_def,0},
  };
  // match deployed adaptive block policy
  auto blockFor=[&](long long n){ return (n < (4LL<<20)) ? 128 : 512; };
  for(long long N : {1024LL*1024, 4096LL*4096}){
    int threads=blockFor(N);
    size_t bytes=N*4; long long n4=N/4;
    float *di,*doo; CK(cudaMalloc(&di,bytes)); CK(cudaMalloc(&doo,bytes));
    std::vector<float> h(N); for(long long i=0;i<N;i++) h[i]=((i&1)?-1.f:1.f)*float(i%101-50);
    CK(cudaMemcpy(di,h.data(),bytes,cudaMemcpyHostToDevice));
    double gb=2.0*bytes/1e9;
    printf("=== N=%lld  block=%d  (rw %.0f MB) ===\n",N,threads,gb*1e3);
    for(auto&c:cs){
      if(c.carveoutL1) cudaFuncSetAttribute(c.f,cudaFuncAttributePreferredSharedMemoryCarveout,0);
      else cudaFuncSetAttribute(c.f,cudaFuncAttributePreferredSharedMemoryCarveout,-1);
      long long blk=(n4+threads-1)/threads;
      CK(cudaMemset(doo,0xff,bytes));
      c.f<<<(unsigned)blk,threads>>>((float4*)doo,(const float4*)di,n4); CK(cudaDeviceSynchronize());
      std::vector<float> ho(N); CK(cudaMemcpy(ho.data(),doo,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;}
      for(int w=0;w<80;w++) c.f<<<(unsigned)blk,threads>>>((float4*)doo,(const float4*)di,n4);
      CK(cudaDeviceSynchronize());
      cudaEvent_t a,b;cudaEventCreate(&a);cudaEventCreate(&b);cudaEventRecord(a);
      const int IT=400; for(int it=0;it<IT;it++) c.f<<<(unsigned)blk,threads>>>((float4*)doo,(const float4*)di,n4);
      cudaEventRecord(b);CK(cudaEventSynchronize(b));
      float ms;cudaEventElapsedTime(&ms,a,b);ms/=IT;
      printf("  %-22s %.4f ms  %5.0f GB/s  (%.1f%%)  %s\n",c.name.c_str(),ms,gb/(ms/1e3),100*gb/(ms/1e3)/peak,ok?"ok":"FAIL");
      cudaEventDestroy(a);cudaEventDestroy(b);
    }
    printf("\n"); cudaFree(di);cudaFree(doo);
  }
  return 0;
}