// ============================================================
// bench_l2_cache.cu
//
// Usage:
//   ./bench_l2_cache [working_set_mib] [measured_iters] [warmup_iters]
//
// Examples:
//   ./bench_l2_cache
//   ./bench_l2_cache 32
//   ./bench_l2_cache 64 5000
//   ./bench_l2_cache 64 5000 200
//
// 目的：
//   测量 GPU L2 cache global load 命中带宽。
//
// 方法：
//   1. 使用 ld.global.cg.v4.f32。
//      .cg = cache globally，通常绕过 L1，主要通过 L2/global path。
//   2. working set 控制在 L2 容量附近或小于 L2。
//   3. kernel 内部分成两个阶段：
//        warmup 阶段：让数据尽量进入 L2。
//        measured 阶段：只测量该阶段的 clock cycles。
//   4. 使用 clock64() 在 device 端计时。
//   5. host 端使用 cudaOccupancyMaxActiveBlocksPerMultiprocessor()
//      过滤掉 GridBlk/SM 大于实际可驻留 block 数的配置，避免多 wave
//      导致带宽被高估。
//   6. 输出 GB/s 和 byte/clk。
//
// 注意：
//   L2 是整个 GPU 共享资源，所以 byte/clk 不除以 SM 数。
// ============================================================

#include <cuda_runtime.h>

#include <cstdio>
#include <cstdlib>
#include <vector>
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <nvml.h>

#define CHECK_CUDA(call)                                                   \
    do {                                                                   \
        cudaError_t err = call;                                            \
        if (err != cudaSuccess) {                                          \
            fprintf(stderr, "CUDA error %s:%d: %s\n",                      \
                    __FILE__, __LINE__, cudaGetErrorString(err));          \
            std::exit(EXIT_FAILURE);                                       \
        }                                                                  \
    } while (0)

#define NVML_CHECK(call)                                                        \
do {                                                                        \
    nvmlReturn_t _status = (call);                                          \
    if (_status != NVML_SUCCESS) {                                          \
        fprintf(stderr, "NVML error %s:%d: %s\n",                           \
                __FILE__, __LINE__, nvmlErrorString(_status));              \
        exit(EXIT_FAILURE);                                                 \
    }                                                                       \
} while (0)
// ------------------------------------------------------------
// .cg load: cache globally
// 通常绕过 L1，使用 L2/global path。
// ------------------------------------------------------------
__device__ __forceinline__ float4 load_cg_float4(const float4* ptr) {
    float4 v;

    asm volatile(
        "ld.global.cg.v4.f32 {%0, %1, %2, %3}, [%4];"
        : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
        : "l"(ptr)
        : "memory"
    );

    return v;
}

// ------------------------------------------------------------
// L2 read kernel with internal warmup + measured phase
//
// 所有 blocks/threads 共同访问同一个数组。
// 每个线程按照 grid-stride 方式访问。
// ------------------------------------------------------------
__global__ void l2_read_kernel_timed(
    const float4* __restrict__ data,
    float* __restrict__ sink,
    unsigned long long* __restrict__ block_cycles,
    size_t elems,
    int warmup_iters,
    int measured_iters
) {
    size_t tid =
        static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;

    size_t stride =
        static_cast<size_t>(gridDim.x) * blockDim.x;

    float acc = 0.0f;

    __shared__ unsigned long long start_cycle;

    // --------------------------------------------------------
    // Phase 1: warmup，不计时
    // --------------------------------------------------------
    for (int it = 0; it < warmup_iters; ++it) {
        for (size_t i = tid; i < elems; i += stride) {
            float4 v = load_cg_float4(data + i);
            acc += v.x + v.y + v.z + v.w;
        }
    }

    __syncthreads();

    if (threadIdx.x == 0) {
        start_cycle = clock64();
    }

    __syncthreads();

    // --------------------------------------------------------
    // Phase 2: measured，只统计这个阶段
    // --------------------------------------------------------
    for (int it = 0; it < measured_iters; ++it) {
        for (size_t i = tid; i < elems; i += stride) {
            float4 v = load_cg_float4(data + i);
            acc += v.x + v.y + v.z + v.w;
        }
    }

    __syncthreads();

    if (threadIdx.x == 0) {
        unsigned long long stop_cycle = clock64();
        block_cycles[blockIdx.x] = stop_cycle - start_cycle;
    }

    // 防止编译器优化掉 load。
    sink[tid] = acc;
}

// ------------------------------------------------------------
// 初始化 float4 device array
// ------------------------------------------------------------
static void init_float4_pattern(float4* d, size_t elems) {
    const size_t chunk = 1 << 20;
    std::vector<float4> h(chunk);

    for (size_t i = 0; i < chunk; ++i) {
        float x = static_cast<float>((i & 1023) + 1);
        h[i] = make_float4(x, x + 1.0f, x + 2.0f, x + 3.0f);
    }

    size_t off = 0;

    while (off < elems) {
        size_t now = std::min(chunk, elems - off);

        CHECK_CUDA(cudaMemcpy(
            d + off,
            h.data(),
            now * sizeof(float4),
            cudaMemcpyHostToDevice
        ));

        off += now;
    }
}

static double gbps_from_cycles(
    double bytes,
    unsigned long long cycles,
    double clock_hz
) {
    double sec = static_cast<double>(cycles) / clock_hz;
    return bytes / sec / 1e9;
}

static void print_usage(const char* prog) {
    fprintf(stderr,
            "Usage:\n"
            "  %s [working_set_mib] [measured_iters] [warmup_iters]\n\n"
            "Examples:\n"
            "  %s\n"
            "  %s 32\n"
            "  %s 64 5000\n"
            "  %s 64 5000 200\n",
            prog, prog, prog, prog, prog);
}

int main(int argc, char** argv) {
    if (argc > 4) {
        print_usage(argv[0]);
        return 1;
    }

    int dev = 0;
    CHECK_CUDA(cudaSetDevice(dev));

    cudaDeviceProp prop{};
    CHECK_CUDA(cudaGetDeviceProperties(&prop, dev));

    int sms = prop.multiProcessorCount;

    // 使用 NVML 获取当前实时频率
    NVML_CHECK(nvmlInit());
    // 通过 PCI Bus ID 找到对应 NVML 设备
    char pci_bus_id[32] = {};
    CHECK_CUDA(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), dev));
    nvmlDevice_t nvml_dev;
    NVML_CHECK(nvmlDeviceGetHandleByPciBusId(pci_bus_id, &nvml_dev));
    unsigned int graphics_clock = 0;
    unsigned int sm_clock = 0;
    unsigned int mem_clock = 0;

    NVML_CHECK(nvmlDeviceGetClockInfo(nvml_dev, NVML_CLOCK_GRAPHICS, &graphics_clock));
    NVML_CHECK(nvmlDeviceGetClockInfo(nvml_dev, NVML_CLOCK_SM, &sm_clock));
    NVML_CHECK(nvmlDeviceGetClockInfo(nvml_dev, NVML_CLOCK_MEM, &mem_clock));

    size_t l2_bytes = prop.l2CacheSize;
    size_t global_mem = prop.totalGlobalMem;


    printf("GPU: %s\n", prop.name);
    printf("SM count: %d\n", sms);
    printf("Max threads per block: %d\n", prop.maxThreadsPerBlock);
    printf("L2 cache size: %.2f MiB\n", l2_bytes / 1024.0 / 1024.0);
    printf("Global memory: %.2f GiB\n", global_mem / 1024.0 / 1024.0 / 1024.0);
    printf("Benchmark: L2 cache read bandwidth using ld.global.cg.v4.f32\n");
    printf("Timing: device-side clock64, measured phase only\n");
    printf("Occupancy rule: skip GridBlk/SM > active_blocks_per_SM\n");
    printf("Current GPU graphics clock: %u MHz\n", graphics_clock);
    printf("Current SM clock:           %u MHz\n", sm_clock);
    printf("Current memory clock:       %u MHz\n", mem_clock);
    printf("\n");

    double clock_hz = static_cast<double>(graphics_clock) * 1e6;
    // ------------------------------------------------------------
    // working set
    // ------------------------------------------------------------
    std::vector<size_t> working_set_list;

    if (argc >= 2) {
        double mib = std::atof(argv[1]);

        if (mib <= 0.0) {
            fprintf(stderr, "Invalid working_set_mib: %s\n", argv[1]);
            return 1;
        }

        size_t bytes =
            static_cast<size_t>(mib * 1024.0 * 1024.0);

        working_set_list.push_back(bytes);
    } else {
        // 默认扫描 L2 容量比例。
        // ratio > 1.0 的结果不再是纯 L2 hit，可用于观察容量拐点。
        std::vector<double> ratios = {
            0.05,
            0.1,
            0.15,
            0.2,
            0.25,
            0.3,
            0.35,
            0.4,
            0.45,
            0.5,
            0.55,
            0.6,
            0.65,
            0.7,
            0.75,
            0.8
        };

        for (double r : ratios) {
            size_t bytes =
                static_cast<size_t>(
                    static_cast<double>(l2_bytes) * r
                );

            working_set_list.push_back(bytes);
        }
    }

    int measured_iters = 5000;
    int warmup_iters = 200;

    if (argc >= 3) {
        measured_iters = std::atoi(argv[2]);

        if (measured_iters <= 0) {
            fprintf(stderr, "Invalid measured_iters: %s\n", argv[2]);
            return 1;
        }
    }

    if (argc >= 4) {
        warmup_iters = std::atoi(argv[3]);

        if (warmup_iters < 0) {
            fprintf(stderr, "Invalid warmup_iters: %s\n", argv[3]);
            return 1;
        }
    }

    // ------------------------------------------------------------
    // 扫描配置
    // ------------------------------------------------------------
    std::vector<int> thread_candidates = {};
    for (int i = 32; i <= 1536; i += 32)
    {   
        thread_candidates.push_back(i);
    }
    std::vector<int> block_per_sm_candidates = {};
    for (int i = 1; i <= 24; ++i)
    {
        block_per_sm_candidates.push_back(i);
    }
    // std::vector<int> thread_candidates = {
    //     512,
    // };

    // std::vector<int> block_per_sm_candidates = {
    //     2,
    // };

    std::vector<int> thread_list;

    for (int t : thread_candidates) {
        // printf("maxThreadsPerBlock: %d\n", prop.maxThreadsPerBlock);
        if (t <= prop.maxThreadsPerBlock) {
            thread_list.push_back(t);
        }
    }

    if (thread_list.empty()) {
        fprintf(stderr, "No valid thread block size found.\n");
        return 1;
    }

    int max_threads =
        *std::max_element(thread_list.begin(), thread_list.end());

    int max_bpsm =
        *std::max_element(
            block_per_sm_candidates.begin(),
            block_per_sm_candidates.end()
        );

    int max_blocks = sms * max_bpsm;

    size_t max_working_bytes =
        *std::max_element(
            working_set_list.begin(),
            working_set_list.end()
        );

    size_t max_elems =
        max_working_bytes / sizeof(float4);

    max_working_bytes =
        max_elems * sizeof(float4);

    if (max_elems == 0) {
        fprintf(stderr, "Working set too small.\n");
        return 1;
    }

    if (max_working_bytes > global_mem / 2) {
        fprintf(stderr,
                "Requested working set too large: %.2f MiB\n",
                max_working_bytes / 1024.0 / 1024.0);
        return 1;
    }

    printf("Scan settings:\n");

    printf("  Working set candidates MiB: ");
    for (size_t s : working_set_list) {
        printf("%.2f ", s / 1024.0 / 1024.0);
    }
    printf("\n");

    printf("  Thread candidates: ");
    for (int t : thread_list) {
        printf("%d ", t);
    }
    printf("\n");

    printf("  GridBlk/SM candidates: ");
    for (int b : block_per_sm_candidates) {
        printf("%d ", b);
    }
    printf("\n");

    printf("  measured_iters = %d\n", measured_iters);
    printf("  warmup_iters   = %d\n\n", warmup_iters);

    float4* d_data = nullptr;
    float* d_sink = nullptr;
    unsigned long long* d_block_cycles = nullptr;

    CHECK_CUDA(cudaMalloc(&d_data, max_working_bytes));

    CHECK_CUDA(cudaMalloc(
        &d_sink,
        static_cast<size_t>(max_blocks) *
        static_cast<size_t>(max_threads) *
        sizeof(float)
    ));

    CHECK_CUDA(cudaMalloc(
        &d_block_cycles,
        static_cast<size_t>(max_blocks) *
        sizeof(unsigned long long)
    ));

    printf("Initializing data: %.3f MiB\n",
           max_working_bytes / 1024.0 / 1024.0);

    init_float4_pattern(d_data, max_elems);

    CHECK_CUDA(cudaMemset(
        d_sink,
        0,
        static_cast<size_t>(max_blocks) *
        static_cast<size_t>(max_threads) *
        sizeof(float)
    ));

    CHECK_CUDA(cudaMemset(
        d_block_cycles,
        0,
        static_cast<size_t>(max_blocks) *
        sizeof(unsigned long long)
    ));

    printf("Initialization done.\n\n");

    std::vector<unsigned long long> h_block_cycles(max_blocks);

    printf("%12s %12s %10s %10s %8s %10s %10s %14s %14s %14s %16s %16s\n",
           "L2_Ratio",
           "Size_MiB",
           "GridBlk/SM",
           "ActiveBlk",
           "Threads",
           "Blocks",
           "GridWarp/SM",
           "Warmup",
           "Measured",
           "Cycles",
           "GB/s",
           "byte/clk");

    printf("------------------------------------------------------------------------------------------------------------------------------------------------------\n");

    double best_gbps = 0.0;
    double best_bpc = 0.0;
    size_t best_bytes = 0;
    int best_threads = 0;
    int best_bpsm = 0;
    int best_active_blocks = 0;
    unsigned long long best_cycles = 0;

    for (size_t raw_bytes : working_set_list) {
        size_t elems =
            raw_bytes / sizeof(float4);

        size_t bytes =
            elems * sizeof(float4);

        if (elems == 0) {
            continue;
        }

        for (int threads : thread_list) {
            int active_blocks_per_sm = 0;

            CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &active_blocks_per_sm,
                l2_read_kernel_timed,
                threads,
                0
            ));

            for (int bpsm : block_per_sm_candidates) {
                // ------------------------------------------------
                // 关键修正：
                // 只允许单 wave 配置。
                // 如果 bpsm 大于 active blocks per SM，
                // 则 block 会分多 wave 执行，per-block clock64()
                // 无法代表整个 grid measured phase 的 wall time。
                // ------------------------------------------------
                if (bpsm > active_blocks_per_sm) {
                    continue;
                }

                int blocks = sms * bpsm;

                double grid_warps_per_sm =
                    static_cast<double>(bpsm) *
                    static_cast<double>(threads) / 32.0;

                CHECK_CUDA(cudaMemset(
                    d_sink,
                    0,
                    static_cast<size_t>(blocks) *
                    static_cast<size_t>(threads) *
                    sizeof(float)
                ));

                CHECK_CUDA(cudaMemset(
                    d_block_cycles,
                    0,
                    static_cast<size_t>(blocks) *
                    sizeof(unsigned long long)
                ));

                l2_read_kernel_timed<<<blocks, threads>>>(
                    d_data,
                    d_sink,
                    d_block_cycles,
                    elems,
                    warmup_iters,
                    measured_iters
                );

                CHECK_CUDA(cudaDeviceSynchronize());
                CHECK_CUDA(cudaGetLastError());

                CHECK_CUDA(cudaMemcpy(
                    h_block_cycles.data(),
                    d_block_cycles,
                    static_cast<size_t>(blocks) *
                    sizeof(unsigned long long),
                    cudaMemcpyDeviceToHost
                ));

                unsigned long long max_cycles = 0;

                for (int i = 0; i < blocks; ++i) {
                    max_cycles = std::max(max_cycles, h_block_cycles[i]);
                }

                if (max_cycles == 0) {
                    continue;
                }

                double logical_bytes =
                    static_cast<double>(bytes) *
                    static_cast<double>(measured_iters);

                double bandwidth =
                    gbps_from_cycles(logical_bytes, max_cycles, clock_hz);

                // L2 是全 GPU 共享资源，所以 byte/clk 不除以 SM 数。
                double byte_per_clk =
                    logical_bytes /
                    static_cast<double>(max_cycles);

                double ratio =
                    static_cast<double>(bytes) /
                    static_cast<double>(l2_bytes);

                printf("%12.3f %12.3f %10d %10d %8d %10d %10.1f %14d %14d %14llu %16.2f %16.2f\n",
                       ratio,
                       bytes / 1024.0 / 1024.0,
                       bpsm,
                       active_blocks_per_sm,
                       threads,
                       blocks,
                       grid_warps_per_sm,
                       warmup_iters,
                       measured_iters,
                       max_cycles,
                       bandwidth,
                       byte_per_clk);

                if (bandwidth > best_gbps) {
                    best_gbps = bandwidth;
                    best_bpc = byte_per_clk;
                    best_bytes = bytes;
                    best_threads = threads;
                    best_bpsm = bpsm;
                    best_active_blocks = active_blocks_per_sm;
                    best_cycles = max_cycles;
                }
            }

            printf("------------------------------------------------------------------------------------------------------------------------------------------------------\n");
        }
    }

    printf("\nBest result:\n");
    printf("  Working set: %.3f MiB\n",
           best_bytes / 1024.0 / 1024.0);
    printf("  L2 ratio: %.3f\n",
           static_cast<double>(best_bytes) /
           static_cast<double>(l2_bytes));
    printf("  GridBlk/SM: %d\n", best_bpsm);
    printf("  Active blocks/SM for this threads/block: %d\n",
           best_active_blocks);
    printf("  Threads/block: %d\n", best_threads);
    printf("  Cycles: %llu\n", best_cycles);
    printf("  Best GB/s: %.2f\n", best_gbps);
    printf("  Best byte/clk: %.2f\n", best_bpc);

    printf("\nConfiguration summary:\n");
    printf("  Instruction: ld.global.cg.v4.f32\n");
    printf("  Timing method: internal warmup + internal measured phase\n");
    printf("  Warmup phase is excluded from measured cycles.\n");
    printf("  Occupancy filter: GridBlk/SM <= active_blocks_per_SM\n");

    CHECK_CUDA(cudaFree(d_data));
    CHECK_CUDA(cudaFree(d_sink));
    CHECK_CUDA(cudaFree(d_block_cycles));

    return 0;
}