// ============================================================
// bench_l1_cache_float.cu
//
// Usage:
//   ./bench_l1_cache_float [per_sm_working_set_kib] [measured_iters] [warmup_iters]
//
// Examples:
//   ./bench_l1_cache_float
//   ./bench_l1_cache_float 64
//   ./bench_l1_cache_float 64 20000
//   ./bench_l1_cache_float 64 20000 500
//
// 目的：
//   测量 GPU L1 cache global scalar float load 命中带宽。
//
// 方法：
//   1. 使用 ld.global.ca.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 scalar float load hit throughput。
//   不测 global store bandwidth。
// ============================================================

#include <cuda_runtime.h>
#include <nvml.h>

#include <cstdio>
#include <cstdlib>
#include <vector>
#include <algorithm>
#include <cstdint>
#include <cstring>

#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));           \
            std::exit(EXIT_FAILURE);                                       \
        }                                                                   \
    } while (0)

// ------------------------------------------------------------
// .ca scalar float load: cache at all levels
// ------------------------------------------------------------
__device__ __forceinline__ float load_ca_float(const float* ptr) {
    float v;

    asm volatile(
        "ld.global.ca.f32 %0, [%1];"
        : "=f"(v)
        : "l"(ptr)
        : "memory"
    );

    return v;
}

// ------------------------------------------------------------
// L1 scalar float read kernel with internal warmup + measured phase
//
// 每个 block 访问自己独立的一段数据：
//   base = data + blockIdx.x * elems_per_block
//
// elems_per_block 单位是 float 元素数量。
// 计时只覆盖 measured phase。
// ------------------------------------------------------------
__global__ void l1_read_kernel_timed(
    const float* __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 float* 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)) {
            float v = load_ca_float(base + i);
            acc += v;
        }
    }

    __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)) {
            float v = load_ca_float(base + i);
            acc += v;
        }
    }

    __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;
}

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

    for (size_t i = 0; i < chunk; ++i) {
        h[i] = static_cast<float>((i & 1023) + 1);
    }

    size_t off = 0;

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

        CHECK_CUDA(cudaMemcpy(
            d + off,
            h.data(),
            now * sizeof(float),
            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());

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

    double clock_hz = static_cast<double>(graphics_clock) * 1e6;

    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.f32\n");
    printf("Data type: float scalar load, 4 bytes per load instruction\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("Clock used for GB/s:        %u MHz, source: NVML SM clock\n", sm_clock);
    printf("\n");

    // ------------------------------------------------------------
    // working set per SM, 单位 byte
    // ------------------------------------------------------------
    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,
    };

    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 = 100;

    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 = {
    //     1024
    // };      // 每个 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");
        NVML_CHECK(nvmlShutdown());
        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(float));

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

    size_t max_total_bytes =
        max_total_elems * sizeof(float);

    float* 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_float_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;

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

            for (int bpsm : block_per_sm_candidates) {
                // 只允许单 wave 配置，避免 per-block clock64 低估多 wave wall time。
                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(float),
                        target_per_sm_bytes / static_cast<size_t>(bpsm)
                    );

                size_t elems_per_block =
                    per_block_bytes / sizeof(float);

                per_block_bytes =
                    elems_per_block * sizeof(float);

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

                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(cudaGetLastError());
                CHECK_CUDA(cudaDeviceSynchronize());

                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>(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.f32\n");
    printf("  Element type: float\n");
    printf("  Bytes per scalar load instruction: 4\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");
    printf("  Clock source for GB/s: NVML SM clock, %u MHz\n", sm_clock);
    printf("  For stricter clock, use NCU metric: sm__cycles_elapsed.avg.per_second\n");

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

    NVML_CHECK(nvmlShutdown());

    return 0;
}