// ============================================================
// bench_l1_cache.cu
//
// Usage:
//   ./bench_l1_cache [per_sm_working_set_kib] [measured_iters] [warmup_iters]
//
// Examples:
//   ./bench_l1_cache
//   ./bench_l1_cache 64
//   ./bench_l1_cache 64 20000
//   ./bench_l1_cache 64 20000 500
//
// 目的：
//   测量 GPU L1 cache global load 命中带宽。
//
// 方法：
//   1. 使用 ld.global.ca.v4.f32。
//      .ca = cache at all levels，使 global load 尽量经过 L1 + L2。
//   2. 每个 block 访问自己的独立 working set。
//   3. kernel 内部分成两个阶段：
//        warmup 阶段：让数据进入 L1。
//        measured 阶段：只测量该阶段的 clock cycles。
//   4. 使用 clock64() 在 device 端计时。
//   5. host 端使用 cudaOccupancyMaxActiveBlocksPerMultiprocessor()
//      过滤掉 GridBlk/SM 大于实际可驻留 block 数的配置，避免多 wave
//      导致带宽被高估。
//   6. 输出 GB/s 和 byte/clk/SM。
//
// 注意：
//   这个 benchmark 测的是 L1 global load hit throughput。
//   不测 global store bandwidth。
// ============================================================

#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)
// ------------------------------------------------------------
// .ca load: cache at all levels
// ------------------------------------------------------------
__device__ __forceinline__ float4 load_ca_float4(const float4* ptr) {
    float4 v;

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

    return v;
}

// ------------------------------------------------------------
// L1 read kernel with internal warmup + measured phase
//
// 每个 block 访问自己独立的一段数据：
//   base = data + blockIdx.x * elems_per_block
//
// 计时只覆盖 measured phase。
// ------------------------------------------------------------
__global__ void l1_read_kernel_timed(
    const float4* __restrict__ data,
    float* __restrict__ sink,
    unsigned long long* __restrict__ block_cycles,
    size_t elems_per_block,
    int warmup_iters,
    int measured_iters
) {
    int bid = blockIdx.x;
    int tid = threadIdx.x;

    const float4* base =
        data + static_cast<size_t>(bid) * elems_per_block;

    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 = static_cast<size_t>(tid);
                i < elems_per_block;
                i += static_cast<size_t>(blockDim.x)) {
            float4 v = load_ca_float4(base + i);
            acc += v.x + v.y + v.z + v.w;
        }
    }

    __syncthreads();

    if (tid == 0) {
        start_cycle = clock64();
    }

    __syncthreads();

    // --------------------------------------------------------
    // Phase 2: measured，只统计这个阶段
    // --------------------------------------------------------
    for (int it = 0; it < measured_iters; ++it) {
        for (size_t i = static_cast<size_t>(tid);
                i < elems_per_block;
                i += static_cast<size_t>(blockDim.x)) {
            float4 v = load_ca_float4(base + i);
            acc += v.x + v.y + v.z + v.w;
        }
    }

    __syncthreads();

    if (tid == 0) {
        unsigned long long stop_cycle = clock64();
        block_cycles[bid] = stop_cycle - start_cycle;
    }

    // 防止编译器优化掉 load。
    sink[static_cast<size_t>(bid) * blockDim.x + 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 [per_sm_working_set_kib] [measured_iters] [warmup_iters]\n\n"
            "Examples:\n"
            "  %s\n"
            "  %s 64\n"
            "  %s 64 20000\n"
            "  %s 64 20000 500\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));

    // 尽量偏向 L1。
    CHECK_CUDA(cudaDeviceSetCacheConfig(cudaFuncCachePreferL1));

    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));


    printf("GPU: %s\n", prop.name);
    printf("SM count: %d\n", sms);
    printf("Max threads per block: %d\n", prop.maxThreadsPerBlock);
    printf("Benchmark: L1 cache read bandwidth using ld.global.ca.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 per SM
    // ------------------------------------------------------------
    std::vector<size_t> per_sm_sizes = {
        4ull   * 1024,
        8ull   * 1024,
        16ull  * 1024,
        17ull  * 1024,
        18ull  * 1024,
        19ull  * 1024,
        20ull  * 1024,
        21ull  * 1024,
        22ull  * 1024,
        23ull  * 1024,
        24ull  * 1024,
        32ull  * 1024,
        // 48ull  * 1024,
        // 64ull  * 1024,
        // 96ull  * 1024,
        // 128ull * 1024
    };

    if (argc >= 2) {
        per_sm_sizes.clear();

        double kib = std::atof(argv[1]);

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

        size_t bytes = static_cast<size_t>(kib * 1024.0);
        per_sm_sizes.push_back(bytes);
    }

    int measured_iters = 20000;
    int warmup_iters = 500;

    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 = {
    //     704
    // };      // 每个 block 的 threads 数量

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

    std::vector<int> thread_list;

    for (int t : thread_candidates) {
        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;

    printf("Scan settings:\n");

    printf("  Per-SM working set candidates KiB: ");
    for (size_t s : per_sm_sizes) {
        printf("%.1f ", s / 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);

    // ------------------------------------------------------------
    // 分配最大可能用到的内存
    //
    // 最大 per-block working set 出现在 bpsm = 1。
    // ------------------------------------------------------------
    size_t max_per_sm_bytes =
        *std::max_element(per_sm_sizes.begin(), per_sm_sizes.end());

    size_t max_elems_per_block =
        std::max<size_t>(1, max_per_sm_bytes / sizeof(float4));

    size_t max_total_elems =
        max_elems_per_block * static_cast<size_t>(max_blocks);

    size_t max_total_bytes =
        max_total_elems * sizeof(float4);

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

    CHECK_CUDA(cudaMalloc(&d_data, max_total_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_total_bytes / 1024.0 / 1024.0);

    init_float4_pattern(d_data, max_total_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 %10s %10s %8s %10s %10s %14s %14s %14s %16s %16s\n",
           "PerSM_KiB",
           "GridBlk/SM",
           "ActiveBlk",
           "Threads",
           "Blocks",
           "GridWarp/SM",
           "Warmup",
           "Measured",
           "Cycles",
           "GB/s",
           "byte/clk/SM");

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

    double best_gbps = 0.0;
    double best_bpc_sm = 0.0;
    size_t best_per_sm_bytes = 0;
    int best_threads = 0;
    int best_bpsm = 0;
    int best_active_blocks = 0;
    unsigned long long best_cycles = 0;
    
    for (size_t target_per_sm_bytes : per_sm_sizes) {
        for (int threads : thread_list) {
            int active_blocks_per_sm = 0;

            // 对于 l1_read_kernel_timed 这个 kernel，如果每个 block 有 threads 个
            // 线程，并且动态 shared memory 是 0，那么计算一个 SM 最多能同时驻留多少个 block
            CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &active_blocks_per_sm,
                l1_read_kernel_timed,
                threads,
                0
            ));

            for (int bpsm : block_per_sm_candidates) {
                // ------------------------------------------------
                // 只允许单 wave 配置：
                // 例如：
                // active_blocks_per_sm = 1
                // bpsm = 4
                // 那么每个 SM 实际执行过程是：
                // wave 0: 执行第 1 个 block
                // wave 1: 执行第 2 个 block
                // wave 2: 执行第 3 个 block
                // wave 3: 执行第 4 个 block
                // 而 clock64() 是测量 device 端每个 block 自己的计时
                // 如果 host 端只取 max(block_cycles)
                // 就只相当于测到了一个 wave 的时间，而不是四个 wave 的总时间。
                // 这会导致带宽会虚高大约 4 倍。
                // ------------------------------------------------
                if (bpsm > active_blocks_per_sm) {
                    continue;
                }

                int blocks = sms * bpsm;        

                // 保持每个 SM 的总 working set 约等于 target_per_sm_bytes
                size_t per_block_bytes =
                    std::max<size_t>(
                        sizeof(float4),
                        target_per_sm_bytes / static_cast<size_t>(bpsm)
                    );

                size_t elems_per_block =
                    per_block_bytes / sizeof(float4);

                per_block_bytes =
                    elems_per_block * sizeof(float4);

                if (elems_per_block == 0) {
                    continue;
                }

                size_t total_elems =
                    elems_per_block * static_cast<size_t>(blocks);

                size_t total_bytes =
                    total_elems * sizeof(float4);

                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)
                ));

                l1_read_kernel_timed<<<blocks, threads>>>(
                    d_data,
                    d_sink,
                    d_block_cycles,
                    elems_per_block,
                    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;
                
                // host 端找到一个最大的 block 计时
                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>(total_bytes) *
                    static_cast<double>(measured_iters);

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

                double byte_per_clk_per_sm =
                    logical_bytes /
                    static_cast<double>(max_cycles) /
                    static_cast<double>(sms);

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

                if (bandwidth > best_gbps) {
                    best_gbps = bandwidth;
                    best_bpc_sm = byte_per_clk_per_sm;
                    best_per_sm_bytes = target_per_sm_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("  Per-SM working set: %.1f KiB\n",
           best_per_sm_bytes / 1024.0);
    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/SM: %.2f\n", best_bpc_sm);

    printf("\nConfiguration summary:\n");
    printf("  Instruction: ld.global.ca.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;
}