/*
 * Copyright 2025 NVIDIA Corporation. All rights reserved
 *
 * Sample CUPTI app to demonstrate the use of multiple subscribers doing CUDA tracing with CUPTI.
 *
 * Till 13.1, CUPTI has prohibited multiple CUPTI subscribers being active at the same time.
 * In 13.2, CUPTI provides limited support for multiple subscribers to trace CUDA activities at the same time.
 * 
 * More information about multiple subscribers can be found in the CUPTI documentation.
 * 
 * NOTE: The CUPTI multiple subscribers 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
#define SUBSCRIBER_COUNT 2

// Global variables
CUpti_SubscriberHandle g_subscriberHandle[SUBSCRIBER_COUNT];

// 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 CUPTIAPI
BufferRequested_subscriber0(
    uint8_t **ppBuffer,
    size_t *pSize,
    size_t *pMaxNumRecords,
    CUpti_BufferCallbackRequestInfo *pRequestInfo)
{
    uint8_t *pBuffer = (uint8_t *) malloc(globals.activityBufferSize + ALIGN_SIZE);
    MEMORY_ALLOCATION_CALL(pBuffer);

    *pSize = globals.activityBufferSize;
    *ppBuffer = ALIGN_BUFFER(pBuffer, ALIGN_SIZE);
    *pMaxNumRecords = 0;

    printf("BufferRequested_subscriber0\n");

}

static void CUPTIAPI
BufferCompleted_subscriber0(
    uint8_t *pBuffer,
    size_t size,
    size_t validSize,
    CUpti_BufferCallbackCompleteInfo *pCompleteInfo)
{
    printf("BufferCompleted_subscriber0\n");
    if (validSize > 0)
    {
        FILE *pOutputFile = globals.pOutputFile;
        if (!pOutputFile)
        {
            pOutputFile = stdout;
        }

        CUpti_Activity *pRecord = NULL;
        CUptiResult status = CUPTI_SUCCESS;

        do
        {
            status = cuptiActivityGetNextRecord(pBuffer, validSize, &pRecord);
            if (status == CUPTI_SUCCESS)
            {
                fprintf(pOutputFile, "[Subscriber 0] ");
                PrintActivity(pRecord, pOutputFile);
            }
            else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED)
            {
                break;
            }
            else if (status == CUPTI_ERROR_INVALID_KIND)
            {
                break;
            }
            else
            {
                CUPTI_API_CALL(status);
            }
        } while (1);
    }

    free(pBuffer);
}

static void CUPTIAPI
BufferRequested_subscriber1(
    uint8_t **ppBuffer,
    size_t *pSize,
    size_t *pMaxNumRecords,
    CUpti_BufferCallbackRequestInfo *pRequestInfo)
{
    uint8_t *pBuffer = (uint8_t *) malloc(globals.activityBufferSize + ALIGN_SIZE);
    MEMORY_ALLOCATION_CALL(pBuffer);

    printf("BufferRequested_subscriber1\n");

    *pSize = globals.activityBufferSize * 2;
    *ppBuffer = ALIGN_BUFFER(pBuffer, ALIGN_SIZE);
    *pMaxNumRecords = 0;
}

static void CUPTIAPI
BufferCompleted_subscriber1(
    uint8_t *pBuffer,
    size_t size,
    size_t validSize,
    CUpti_BufferCallbackCompleteInfo *pCompleteInfo)
{
    printf("BufferCompleted_subscriber1\n");
    
    if (validSize > 0)
    {
        FILE *pOutputFile = globals.pOutputFile;
        if (!pOutputFile)
        {
            pOutputFile = stdout;
        }

        CUpti_Activity *pRecord = NULL;
        CUptiResult status = CUPTI_SUCCESS;

        do
        {
            status = cuptiActivityGetNextRecord(pBuffer, validSize, &pRecord);
            if (status == CUPTI_SUCCESS)
            {
                fprintf(pOutputFile, "[Subscriber 1] ");
                PrintActivity(pRecord, pOutputFile);
            }
            else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED)
            {
                break;
            }
            else if (status == CUPTI_ERROR_INVALID_KIND)
            {
                break;
            }
            else
            {
                CUPTI_API_CALL(status);
            }
        } while (1);
    }

    free(pBuffer);
}

static void
SetupCupti()
{
    // Query if multiple subscribers are already enabled
    uint8_t value = 0;
    size_t valueSize = sizeof(uint8_t);
    CUPTI_API_CALL_VERBOSE(cuptiActivityGetAttribute_v2(NULL, CUPTI_ACTIVITY_ATTR_MULTIPLE_SUBSCRIBER_STATE, &valueSize, &value));
    if (value == 2)
    {
        std::cout << "Multiple subscribers are not set by any prior subscriber. Our first subscriber can enable it." << std::endl;
    }
    else
    {
        std::cout << "Multiple Subscribers option is previously set, we cannot modify it." << std::endl;
        exit(EXIT_FAILURE);
    }

    CUpti_SubscriberParams subscriberParams = { sizeof(CUpti_SubscriberParams), NULL, NULL, 0, 1}; // Allow multiple subscribers.
    CUPTI_API_CALL_VERBOSE(cuptiSubscribe_v2(&g_subscriberHandle[0], NULL, NULL, &subscriberParams));
    CUPTI_API_CALL_VERBOSE(cuptiSubscribe_v2(&g_subscriberHandle[1], NULL, NULL, &subscriberParams));

    // Set subscriber-specific attributes (must be done after subscription)
    CUpti_ActivityThreadIdType threadIdType0 = CUPTI_ACTIVITY_THREAD_ID_TYPE_DEFAULT;
    CUpti_ActivityThreadIdType threadIdType1 = CUPTI_ACTIVITY_THREAD_ID_TYPE_SYSTEM;
    size_t threadIdTypeSize = sizeof(CUpti_ActivityThreadIdType);
    CUPTI_API_CALL_VERBOSE(cuptiActivitySetAttribute_v2(g_subscriberHandle[0], CUPTI_ACTIVITY_ATTR_THREAD_ID_TYPE, &threadIdTypeSize, &threadIdType0));
    CUPTI_API_CALL_VERBOSE(cuptiActivitySetAttribute_v2(g_subscriberHandle[1], CUPTI_ACTIVITY_ATTR_THREAD_ID_TYPE, &threadIdTypeSize, &threadIdType1));

    // Enable different activities for different subscribers
    CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(g_subscriberHandle[0], CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL, NULL));
    CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(g_subscriberHandle[0], CUPTI_ACTIVITY_KIND_MEMCPY, NULL));

    CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(g_subscriberHandle[1], CUPTI_ACTIVITY_KIND_MEMSET, NULL));
    CUPTI_API_CALL_VERBOSE(cuptiActivityEnable_v2(g_subscriberHandle[1], CUPTI_ACTIVITY_KIND_DRIVER, NULL));

    // Register buffer callbacks for different subscribers
    CUPTI_API_CALL_VERBOSE(cuptiActivityRegisterCallbacks_v2(g_subscriberHandle[0], BufferRequested_subscriber0, BufferCompleted_subscriber0));
    CUPTI_API_CALL_VERBOSE(cuptiActivityRegisterCallbacks_v2(g_subscriberHandle[1], BufferRequested_subscriber1, BufferCompleted_subscriber1));

    // Print all activities enabled in total by all subscribers
    CUpti_ActivityKind *pActivityKinds = NULL;
    uint32_t activityKindsCount = 0;
    CUPTI_API_CALL_VERBOSE(cuptiActivityGetEnabledKinds(NULL, NULL, &activityKindsCount, &activityKindsCount));
    pActivityKinds = (CUpti_ActivityKind *)malloc(activityKindsCount * sizeof(CUpti_ActivityKind));
    MEMORY_ALLOCATION_CALL(pActivityKinds);
    CUPTI_API_CALL_VERBOSE(cuptiActivityGetEnabledKinds(NULL, pActivityKinds, &activityKindsCount, &activityKindsCount));
    std::cout << "Total activities enabled by all subscribers: " << activityKindsCount << std::endl;
    for (uint32_t i = 0; i < activityKindsCount; i++)
    {
        std::cout << "Activity kind " << i << ": " << pActivityKinds[i] << std::endl;
    }

    free(pActivityKinds);
}

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 multiple subscribers trace require CUPTI version 130200 or higher. Skipping running the sample." << std::endl;
        exit(EXIT_SUCCESS);
    }

    SetupCupti();

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

    return 0;
}