// Matrix-free tensor-product Laplacian apply, fp64, sum factorization.
// Sweeps polynomial degree p (n = p+1). One thread block per element,
// n^3 threads, all temporaries in shared memory.
//
// Global traffic per element: read u (n^3 doubles) + write v (n^3 doubles)
//   -> bytes/DOF = 16, independent of p.   (matches the model's g=0 case)
// Arithmetic: 18 one-dimensional contractions, each 2*n^4 flops per element
//   -> flops/DOF = 36*n = 36*(p+1).        (measured for real by ncu)
//
// Elements are independent (no gather/scatter), which isolates exactly the
// operator-apply trade the roofline model describes.

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>

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

// out[i,j,k] = sum_m A[row(m)] * in[...]   dir: 0=x,1=y,2=z ; TR: use A transposed
template<int N, int DIR, bool TR>
__device__ __forceinline__ void apply1d(const double* __restrict__ A,
                                        const double* __restrict__ in,
                                        double* __restrict__ out)
{
    const int t = threadIdx.x;
    const int i = t % N;
    const int j = (t / N) % N;
    const int k = t / (N * N);
    double s = 0.0;
#pragma unroll
    for (int m = 0; m < N; ++m) {
        int r = (DIR == 0) ? i : (DIR == 1) ? j : k;
        double a = TR ? A[m * N + r] : A[r * N + m];
        int idx = (DIR == 0) ? (m + N * (j + N * k))
                : (DIR == 1) ? (i + N * (m + N * k))
                             : (i + N * (j + N * m));
        s = fma(a, in[idx], s);          // 2 flops
    }
    out[t] = s;
}

template<int N>
__global__ __launch_bounds__(N*N*N)
void laplacian(const double* __restrict__ u, double* __restrict__ v,
               const double* __restrict__ B, const double* __restrict__ D,
               long nelem)
{
    constexpr int N3 = N * N * N;
    __shared__ double sB[N*N], sD[N*N];
    __shared__ double su[N3], t1[N3], t2[N3], sg[N3], sv[N3];

    const int t = threadIdx.x;
    if (t < N*N) { sB[t] = B[t]; sD[t] = D[t]; }

    for (long e = blockIdx.x; e < nelem; e += gridDim.x) {
        const double* ue = u + e * (long)N3;
        double*       ve = v + e * (long)N3;

        su[t] = ue[t];
        sv[t] = 0.0;
        __syncthreads();

        // --- component x: g = Bz( By( Dx u ) ) ; then v += Dx^T( By^T( Bz^T g ) )
        apply1d<N,0,false>(sD, su, t1); __syncthreads();
        apply1d<N,1,false>(sB, t1, t2); __syncthreads();
        apply1d<N,2,false>(sB, t2, sg); __syncthreads();
        apply1d<N,2,true >(sB, sg, t1); __syncthreads();
        apply1d<N,1,true >(sB, t1, t2); __syncthreads();
        apply1d<N,0,true >(sD, t2, t1); __syncthreads();
        sv[t] += t1[t]; __syncthreads();

        // --- component y: g = Bz( Dy( Bx u ) )
        apply1d<N,0,false>(sB, su, t1); __syncthreads();
        apply1d<N,1,false>(sD, t1, t2); __syncthreads();
        apply1d<N,2,false>(sB, t2, sg); __syncthreads();
        apply1d<N,2,true >(sB, sg, t1); __syncthreads();
        apply1d<N,1,true >(sD, t1, t2); __syncthreads();
        apply1d<N,0,true >(sB, t2, t1); __syncthreads();
        sv[t] += t1[t]; __syncthreads();

        // --- component z: g = Dz( By( Bx u ) )
        apply1d<N,0,false>(sB, su, t1); __syncthreads();
        apply1d<N,1,false>(sB, t1, t2); __syncthreads();
        apply1d<N,2,false>(sD, t2, sg); __syncthreads();
        apply1d<N,2,true >(sD, sg, t1); __syncthreads();
        apply1d<N,1,true >(sB, t1, t2); __syncthreads();
        apply1d<N,0,true >(sB, t2, t1); __syncthreads();
        sv[t] += t1[t]; __syncthreads();

        ve[t] = sv[t];
        __syncthreads();
    }
}

template<int N>
void run(long target_dof, int reps, int smcount)
{
    constexpr int N3 = N * N * N;
    long nelem = target_dof / N3;
    if (nelem < 1) nelem = 1;
    long ndof = nelem * (long)N3;

    double *u, *v, *B, *D;
    CK(cudaMalloc(&u, ndof * sizeof(double)));
    CK(cudaMalloc(&v, ndof * sizeof(double)));
    CK(cudaMalloc(&B, N * N * sizeof(double)));
    CK(cudaMalloc(&D, N * N * sizeof(double)));

    double hB[N*N], hD[N*N], *hu = (double*)malloc(ndof * sizeof(double));
    for (int i = 0; i < N*N; ++i) { hB[i] = cos(0.7*i + 0.3); hD[i] = sin(0.4*i + 0.1); }
    for (long i = 0; i < ndof; ++i) hu[i] = 1.0 + 1e-3 * (double)(i % 17);
    CK(cudaMemcpy(u, hu, ndof*sizeof(double), cudaMemcpyHostToDevice));
    CK(cudaMemcpy(B, hB, N*N*sizeof(double), cudaMemcpyHostToDevice));
    CK(cudaMemcpy(D, hD, N*N*sizeof(double), cudaMemcpyHostToDevice));

    int grid = smcount * 2;
    if (grid > nelem) grid = (int)nelem;

    laplacian<N><<<grid, N3>>>(u, v, B, D, nelem);   // warm-up
    CK(cudaDeviceSynchronize());

    cudaEvent_t a, b; CK(cudaEventCreate(&a)); CK(cudaEventCreate(&b));
    CK(cudaEventRecord(a));
    for (int r = 0; r < reps; ++r) laplacian<N><<<grid, N3>>>(u, v, B, D, nelem);
    CK(cudaEventRecord(b));
    CK(cudaEventSynchronize(b));
    float ms = 0; CK(cudaEventElapsedTime(&ms, a, b));

    double sec_per_apply = (ms / 1e3) / reps;
    double ps_per_dof    = sec_per_apply / (double)ndof * 1e12;
    // model expectations (to be compared against ncu-measured values)
    double model_flops_dof = 36.0 * N;
    double model_bytes_dof = 16.0;
    double gflops = model_flops_dof * ndof / sec_per_apply / 1e9;
    double gbs    = model_bytes_dof * ndof / sec_per_apply / 1e9;

    printf("%3d %3d %12ld %10ld %10.4f %12.3f %12.2f %12.2f %10.2f\n",
           N - 1, N, nelem, ndof, ps_per_dof, sec_per_apply * 1e3,
           gflops, gbs, model_flops_dof / model_bytes_dof);
    fflush(stdout);

    free(hu);
    CK(cudaFree(u)); CK(cudaFree(v)); CK(cudaFree(B)); CK(cudaFree(D));
}

int main(int argc, char** argv)
{
    long target_dof = (argc > 1) ? atol(argv[1]) : 40000000L;
    int  reps       = (argc > 2) ? atoi(argv[2]) : 20;
    int  only       = (argc > 3) ? atoi(argv[3]) : 0;   // run a single degree if >0

    cudaDeviceProp prop; CK(cudaGetDeviceProperties(&prop, 0));
    if (!only) {
        printf("# device=%s  SMs=%d  clock=%.0fMHz\n", prop.name, prop.multiProcessorCount,
               prop.clockRate / 1e3);
        printf("#  p   n        nelem       ndof    ps/DOF     ms/apply   GFLOP/s(mdl)   GB/s(mdl)   AI(mdl)\n");
    }
    int sm = prop.multiProcessorCount;

#define RUN(NN) if (!only || only == NN) run<NN>(target_dof, reps, sm);
    RUN(2) RUN(3) RUN(4) RUN(5) RUN(6) RUN(7) RUN(8) RUN(9) RUN(10)
#undef RUN
    return 0;
}