/*
 * Copyright 2025 NVIDIA Corporation. All rights reserved
 *
 * Sample CUPTI app to print a trace of CUDA API and GPU activity
 * with user-defined records enabled in CUPTI.
 *
 * Refer following functions to see how to setup CUPTI with
 * user-defined records:
 *  helper_cupti_activity.h: InitCuptiTrace()
 *    How to set the CUPTI_ACTIVITY_ATTR_USER_DEFINED_RECORDS using cuptiActivitySetAttribute()
 *    How to register buffer callbacks using cuptiActivityRegisterCallbacks_v2()
 *  cupti_user_defined_records.cu: SetupCupti()
 *    How to enable activity kinds with user-defined field selections using cuptiActivityEnable_v2()
 *  helper_cupti_activity.h: PrintActivityBuffer()
 *    How to parse the buffers and print the user-defined records.
 *  helper_cupti_activity_user_defined_records.h: PrintUserDefinedActivity()
 *    How to print user-defined records.
 *
 *  Refer these links for online documentation:
 *    CUPTI Activity User-Defined Records: https://docs.nvidia.com/cupti/main/main.html#cupti-user-defined-activity-records
 *    Tutorial: https://docs.nvidia.com/cupti/tutorial/tutorial.html#tutorial-activity-user-defined-records
 *
 * NOTE: The CUPTI user-defined activity records feature is currently in beta. APIs and behavior may change in future releases.
 */

// System headers
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unordered_map>
#include <vector>
#include <iostream>

// CUDA headers
#include <cuda.h>
#include <cuda_runtime.h>

// CUPTI headers
#include "helper_cupti_activity.h"

// Macros
#define COMPUTE_N 50000

// Kernels
__global__ void
VectorAdd(
    const int* pA,
    const int* pB,
    int *pC,
    int N)
{
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < N)
    {
        pC[i] = pA[i] + pB[i];
    }
}

__global__ void
VectorSubtract(
    const int* pA,
    const int* pB,
    int *pC,
    int N)
{
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < N)
    {
        pC[i] = pA[i] - pB[i];
    }
}

// Functions
static void
DoPass(
    cudaStream_t stream)
{
    int *pHostA, *pHostB, *pHostC;
    int *pDeviceA, *pDeviceB, *pDeviceC;
    size_t size = COMPUTE_N * sizeof(int);
    int threadsPerBlock = 256;
    int blocksPerGrid = 0;

    // Allocate input vectors pHostA and pHostB in host memory.
    // Don't bother to initialize.
    pHostA = (int *)malloc(size);
    MEMORY_ALLOCATION_CALL(pHostA);

    pHostB = (int *)malloc(size);
    MEMORY_ALLOCATION_CALL(pHostB);

    pHostC = (int *)malloc(size);
    MEMORY_ALLOCATION_CALL(pHostB);


    // Allocate vectors in device memory.
    RUNTIME_API_CALL(cudaMalloc((void **)&pDeviceA, size));
    RUNTIME_API_CALL(cudaMalloc((void **)&pDeviceB, size));
    RUNTIME_API_CALL(cudaMalloc((void **)&pDeviceC, size));

    RUNTIME_API_CALL(cudaMemcpyAsync(pDeviceA, pHostA, size, cudaMemcpyHostToDevice, stream));
    RUNTIME_API_CALL(cudaMemcpyAsync(pDeviceB, pHostB, size, cudaMemcpyHostToDevice, stream));

    blocksPerGrid = (COMPUTE_N + threadsPerBlock - 1) / threadsPerBlock;

    VectorAdd <<< blocksPerGrid, threadsPerBlock, 0, stream >>> (pDeviceA, pDeviceB, pDeviceC, COMPUTE_N);
    RUNTIME_API_CALL(cudaGetLastError());

    VectorSubtract <<< blocksPerGrid, threadsPerBlock, 0, stream >>> (pDeviceA, pDeviceB, pDeviceC, COMPUTE_N);
    RUNTIME_API_CALL(cudaGetLastError());

    RUNTIME_API_CALL(cudaMemcpyAsync(pHostC, pDeviceC, size, cudaMemcpyDeviceToHost, stream));

    if (stream == 0)
    {
        RUNTIME_API_CALL(cudaDeviceSynchronize());
    }
    else
    {
        RUNTIME_API_CALL(cudaStreamSynchronize(stream));
    }

    // Free host memory.
    if (pHostA)
    {
        free(pHostA);
    }
    if (pHostB)
    {
        free(pHostB);
    }
    if (pHostC)
    {
        free(pHostC);
    }

    // Free device memory.
    if (pDeviceA)
    {
        RUNTIME_API_CALL(cudaFree(pDeviceA));
    }
    if (pDeviceB)
    {
        RUNTIME_API_CALL(cudaFree(pDeviceB));
    }
    if (pDeviceC)
    {
        RUNTIME_API_CALL(cudaFree(pDeviceC));
    }
}

static void
SetupCupti()
{
    UserData *pUserData = (UserData *)malloc(sizeof(UserData));
    MEMORY_ALLOCATION_CALL(pUserData);

    memset(pUserData, 0, sizeof(UserData));
    pUserData->pPostProcessActivityRecords = NULL;
    pUserData->printActivityRecords        = 1;
    // Check helper_cupti_activity.h for more details on how this flag
    // enables user defined activity structures in InitCuptiTrace()
    // Pasting snippet to ensure users enable the following CUPTI attribute
    /**
     *  size_t valueSize = sizeof(uint8_t);
     *  uint8_t value = 1;
     *  cuptiActivitySetAttribute(CUPTI_ACTIVITY_ATTR_USER_DEFINED_RECORDS, &valueSize, value);
    **/
    pUserData->userDefinedRecords = 1;

    // Common CUPTI Initialization
    InitCuptiTrace(pUserData, NULL, stdout);

    {
        CUpti_ActivityFieldSelection selection;
        CUpti_ActivityConfig activityConfig;
        std::vector<int> selectedFields;

        // Select fields
        selectedFields = {
            MEMCPY_FIELD_KIND,                      // CUpti_ActivityKind kind;
            MEMCPY_FIELD_DEVICE_ID,                 // uint32_t deviceId;
            MEMCPY_FIELD_CONTEXT_ID,                // uint32_t contextId;
            MEMCPY_FIELD_STREAM_ID,                 // uint32_t streamId;
            MEMCPY_FIELD_CORRELATION_ID,            // uint64_t correlationId;
            MEMCPY_FIELD_BYTES,                     // uint64_t bytes;
            MEMCPY_FIELD_START,                     // uint64_t start;
            MEMCPY_FIELD_END,                       // uint64_t end;
            MEMCPY_FIELD_COPY_KIND                  // uint8_t copyKind;
        };

        // Store the fields in CUpti_ActivityFieldSelection and the number of fields
        selection.pFieldIds = selectedFields.data();
        selection.numFields = static_cast<uint32_t>(selectedFields.size());

        // Store the field selection in CUpti_ActivityConfig
        activityConfig.fieldSelection = selection;

        // Enable activity kind with the specified field selection
        CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(globals.subscriberHandle, CUPTI_ACTIVITY_KIND_MEMCPY, &activityConfig));
    }
    {
        CUpti_ActivityFieldSelection selection;
        CUpti_ActivityConfig activityConfig;
        std::vector<int> selectedFields;

        // Use all fields by default
        selectedFields = {
            MEMCPY2_FIELD_KIND,             // CUpti_ActivityKind kind;
            MEMCPY2_FIELD_DEVICE_ID,        // uint32_t deviceId;
            MEMCPY2_FIELD_CONTEXT_ID,       // uint32_t contextId;
            MEMCPY2_FIELD_STREAM_ID,        // uint32_t streamId;
            MEMCPY2_FIELD_CORRELATION_ID,   // uint64_t correlationId;
            MEMCPY2_FIELD_BYTES,            // uint64_t bytes;
            MEMCPY2_FIELD_SRC_DEVICE_ID,    // uint32_t srcDeviceId;
            MEMCPY2_FIELD_SRC_CONTEXT_ID,   // uint64_t srcContextId;
            MEMCPY2_FIELD_DST_DEVICE_ID,    // uint32_t dstDeviceId;
            MEMCPY2_FIELD_DST_CONTEXT_ID,   // uint64_t dstContextId;
            MEMCPY2_FIELD_START,            // uint64_t start;
            MEMCPY2_FIELD_END,              // uint64_t end;
            MEMCPY2_FIELD_COPY_KIND,        // uint8_t copyKind;
        };

        selection.pFieldIds = selectedFields.data();
        selection.numFields = static_cast<uint32_t>(selectedFields.size());

        activityConfig.fieldSelection = selection;
        CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(globals.subscriberHandle, CUPTI_ACTIVITY_KIND_MEMCPY2, &activityConfig));
    }
    {
        CUpti_ActivityFieldSelection selection;
        CUpti_ActivityConfig activityConfig;
        std::vector<int> selectedFields;

        // Use all fields by default
        selectedFields = {
            MEMSET_FIELD_KIND,             // CUpti_ActivityKind kind;
            MEMSET_FIELD_DEVICE_ID,        // uint32_t deviceId;
            MEMSET_FIELD_CONTEXT_ID,       // uint32_t contextId;
            MEMSET_FIELD_STREAM_ID,        // uint32_t streamId;
            MEMSET_FIELD_CORRELATION_ID,   // uint64_t correlationId;
            MEMSET_FIELD_BYTES,            // uint64_t bytes;
            MEMSET_FIELD_START,            // uint64_t start;
            MEMSET_FIELD_END,              // uint64_t end;
            MEMSET_FIELD_VALUE             // uint32_t value;
        };

        selection.pFieldIds = selectedFields.data();
        selection.numFields = static_cast<uint32_t>(selectedFields.size());

        activityConfig.fieldSelection = selection;
        CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(globals.subscriberHandle, CUPTI_ACTIVITY_KIND_MEMSET, &activityConfig));
    }
    {
        CUpti_ActivityFieldSelection selection;
        CUpti_ActivityConfig activityConfig;
        std::vector<int> selectedFields;

        // Use all fields by default
        selectedFields = {
            KERNEL_FIELD_KIND,                              // CUpti_ActivityKind kind;
            KERNEL_FIELD_DEVICE_ID,                         // uint32_t deviceId;
            KERNEL_FIELD_CONTEXT_ID,                        // uint32_t contextId;
            KERNEL_FIELD_STREAM_ID,                         // uint32_t streamId;
            KERNEL_FIELD_CORRELATION_ID,                    // uint64_t correlationId;
            KERNEL_FIELD_NAME,                              // const char *name;
            KERNEL_FIELD_START,                             // uint64_t start;
            KERNEL_FIELD_END                                // uint64_t end;
        };

        selection.pFieldIds = selectedFields.data();
        selection.numFields = static_cast<uint32_t>(selectedFields.size());

        activityConfig.fieldSelection = selection;
        CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(globals.subscriberHandle, CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL, &activityConfig));
    }
    {
        CUpti_ActivityFieldSelection selection;
        CUpti_ActivityConfig activityConfig;
        std::vector<int> selectedFields;

        // Use all fields by default
        selectedFields = {
            API_FIELD_KIND,             // CUpti_ActivityKind kind;
            API_FIELD_CBID,             // CUpti_CallbackId cbid;
            API_FIELD_THREAD_ID,        // uint64_t threadId;
            API_FIELD_CORRELATION_ID,   // uint64_t correlationId;
            API_FIELD_START,            // uint64_t start;
            API_FIELD_END               // uint64_t end;
        };

        selection.pFieldIds = selectedFields.data();
        selection.numFields = static_cast<uint32_t>(selectedFields.size());

        activityConfig.fieldSelection = selection;
        CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(globals.subscriberHandle, CUPTI_ACTIVITY_KIND_DRIVER, &activityConfig));
    }
    {
        CUpti_ActivityFieldSelection selection;
        CUpti_ActivityConfig activityConfig;
        std::vector<int> selectedFields;

        // Use all fields by default
        selectedFields = {
            API_FIELD_KIND,             // CUpti_ActivityKind kind;
            API_FIELD_CBID,             // CUpti_CallbackId cbid;
            API_FIELD_THREAD_ID,        // uint64_t threadId;
            API_FIELD_CORRELATION_ID,   // uint64_t correlationId;
            API_FIELD_START,            // uint64_t start;
            API_FIELD_END               // uint64_t end;
        };

        selection.pFieldIds = selectedFields.data();
        selection.numFields = static_cast<uint32_t>(selectedFields.size());

        activityConfig.fieldSelection = selection;
        CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(globals.subscriberHandle, CUPTI_ACTIVITY_KIND_RUNTIME, &activityConfig));
    }
}

int
main(
    int argc,
    char *argv[])
{
    CUdevice device;
    char deviceName[256];
    int deviceId = 0, deviceCount = 0;

    uint32_t version;
    CUPTI_API_CALL_VERBOSE(cuptiGetVersion(&version));
    std::cout << "Current CUPTI Version: " << version << std::endl;
    if (version < 130200)
    {
        std::cerr << "CUPTI user-defined activity records require CUPTI version 130200 or higher. Skipping running the sample." << std::endl;
        exit(EXIT_SUCCESS);
    }

    SetupCupti();

    // Intialize CUDA
    DRIVER_API_CALL(cuInit(0));

    RUNTIME_API_CALL(cudaGetDeviceCount(&deviceCount));

    for (deviceId = 0; deviceId < deviceCount; deviceId++)
    {
        DRIVER_API_CALL(cuDeviceGet(&device, deviceId));
        DRIVER_API_CALL(cuDeviceGetName(deviceName, 256, device));
        std::cout << "Device Name: " << deviceName << std::endl;

        RUNTIME_API_CALL(cudaSetDevice(deviceId));

        // DoPass with user stream
        cudaStream_t stream;
        RUNTIME_API_CALL(cudaStreamCreate(&stream));
        DoPass(stream);

        RUNTIME_API_CALL(cudaDeviceSynchronize());

        RUNTIME_API_CALL(cudaDeviceReset());
    }

    DeInitCuptiTrace();

    exit(EXIT_SUCCESS);
}
