/* Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
 * This is a simple test showing performance and usability
 * improvements with large kernel parameters introduced in CUDA 12.1
 */
#include <cassert>
#include <chrono>
#include <iostream>

// Utility includes
#include <helper_cuda.h>

using namespace std;
using namespace std::chrono;

#define TEST_ITERATIONS     (1000)
#define TOTAL_PARAMS        (8000) // ints
#define KERNEL_PARAM_LIMIT  (1024) // ints
#define CONST_COPIED_PARAMS (TOTAL_PARAMS - KERNEL_PARAM_LIMIT)

__constant__ int excess_params[CONST_COPIED_PARAMS];

typedef struct
{
    int param[KERNEL_PARAM_LIMIT];
} param_t;

typedef struct
{
    int param[TOTAL_PARAMS];
} param_large_t;

// Kernel with 4KB kernel parameter limit
__global__ void kernelDefault(__grid_constant__ const param_t p, int *result)
{
    int tmp = 0;

    // accumulate kernel parameters
    for (int i = 0; i < KERNEL_PARAM_LIMIT; ++i) {
        tmp += p.param[i];
    }

    // accumulate excess values passed via const memory
    for (int i = 0; i < CONST_COPIED_PARAMS; ++i) {
        tmp += excess_params[i];
    }

    *result = tmp;
}

// Kernel with 32,764 byte kernel parameter limit
__global__ void kernelLargeParam(__grid_constant__ const param_large_t p, int *result)
{
    int tmp = 0;

    // accumulate kernel parameters
    for (int i = 0; i < TOTAL_PARAMS; ++i) {
        tmp += p.param[i];
    }

    *result = tmp;
}

static void report_time(std::chrono::time_point<std::chrono::steady_clock> start,
                        std::chrono::time_point<std::chrono::steady_clock> end,
                        int                                                iters)
{
    auto usecs = duration_cast<duration<float, microseconds::period>>(end - start);
    cout << usecs.count() / iters << endl;
}

int main()
{
    int rc;
    cudaFree(0);

    param_t       p;
    param_large_t p_large;

    // pageable host memory that holds excess constants passed via constant memory
    int *copied_params = (int *)malloc(CONST_COPIED_PARAMS * sizeof(int));
    assert(copied_params);

    // storage for computed result
    int *d_result;
    int  h_result;
    checkCudaErrors(cudaMalloc(&d_result, sizeof(int)));

    int expected_result = 0;

    // fill in data for validation
    for (int i = 0; i < KERNEL_PARAM_LIMIT; ++i) {
        p.param[i] = (i & 0xFF);
    }
    for (int i = KERNEL_PARAM_LIMIT; i < TOTAL_PARAMS; ++i) {
        copied_params[i - KERNEL_PARAM_LIMIT] = (i & 0xFF);
    }
    for (int i = 0; i < TOTAL_PARAMS; ++i) {
        p_large.param[i] = (i & 0xFF);
        expected_result += (i & 0xFF);
    }

    // warmup, verify correctness
    checkCudaErrors(
        cudaMemcpyToSymbol(excess_params, copied_params, CONST_COPIED_PARAMS * sizeof(int), 0, cudaMemcpyHostToDevice));
    kernelDefault<<<1, 1>>>(p, d_result);
    checkCudaErrors(cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost));
    checkCudaErrors(cudaDeviceSynchronize());
    if (h_result != expected_result) {
        std::cout << "Test failed" << std::endl;
        rc = -1;
        goto Exit;
    }

    kernelLargeParam<<<1, 1>>>(p_large, d_result);
    checkCudaErrors(cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost));
    checkCudaErrors(cudaDeviceSynchronize());
    if (h_result != expected_result) {
        std::cout << "Test failed" << std::endl;
        rc = -1;
        goto Exit;
    }

    // benchmark default kernel parameter limit
    {
        auto start = steady_clock::now();
        for (int i = 0; i < TEST_ITERATIONS; ++i) {
            checkCudaErrors(cudaMemcpyToSymbol(
                excess_params, copied_params, CONST_COPIED_PARAMS * sizeof(int), 0, cudaMemcpyHostToDevice));
            kernelDefault<<<1, 1>>>(p, d_result);
        }
        checkCudaErrors(cudaDeviceSynchronize());
        auto end = steady_clock::now();
        std::cout << "Kernel 4KB parameter limit - time (us):";
        report_time(start, end, TEST_ITERATIONS);

        // benchmark large kernel parameter limit
        start = steady_clock::now();
        for (int i = 0; i < TEST_ITERATIONS; ++i) {
            kernelLargeParam<<<1, 1>>>(p_large, d_result);
        }
        checkCudaErrors(cudaDeviceSynchronize());
        end = steady_clock::now();
        std::cout << "Kernel 32,764 byte parameter limit - time (us):";
        report_time(start, end, TEST_ITERATIONS);
    }
    std::cout << "Test passed!" << std::endl;
    rc = 0;
Exit:
    // cleanup
    cudaFree(d_result);
    free(copied_params);
    return rc;
}
