/*
 * Copyright 2010-2017 NVIDIA Corporation. All rights reserved
 *
 * Sample app to demonstrate use of CUPTI library to obtain profiler event values
 * using callbacks for CUDA runtime APIs
 *
 */

#include <stdio.h>
#include <cuda.h>
#include <cupti.h>

#define EVENT_NAME "inst_executed"

#define CHECK_CU_ERROR(err, cufunc)                                     \
  if (err != CUDA_SUCCESS)                                              \
    {                                                                   \
      printf ("%s:%d: error %d for CUDA Driver API function '%s'\n",    \
              __FILE__, __LINE__, err, cufunc);                         \
      exit(-1);                                                         \
    }

#define CHECK_CUPTI_ERROR(err, cuptifunc)                               \
  if (err != CUPTI_SUCCESS)                                             \
    {                                                                   \
      const char *errstr;                                               \
      cuptiGetResultString(err, &errstr);                               \
      printf ("%s:%d:Error %s for CUPTI API function '%s'.\n",          \
              __FILE__, __LINE__, errstr, cuptifunc);                   \
      if(err == CUPTI_ERROR_LEGACY_PROFILER_NOT_SUPPORTED)              \
          exit(0);                                                      \
      else                                                              \
          exit(-1);                                                     \
    }

typedef struct cupti_eventData_st {
  CUpti_EventGroup eventGroup;
  CUpti_EventID eventId;
} cupti_eventData;

// Structure to hold data collected by callback
typedef struct RuntimeApiTrace_st {
  cupti_eventData *eventData;
  uint64_t eventVal;
} RuntimeApiTrace_t;

  CUptiResult cuptiErr;
  CUpti_SubscriberHandle subscriber;
  cupti_eventData cuptiEvent;
  RuntimeApiTrace_t trace;

void CUPTIAPI
getEventValueCallback(void *userdata, CUpti_CallbackDomain domain,
                      CUpti_CallbackId cbid, const CUpti_CallbackData *cbInfo)
{
  CUptiResult cuptiErr;
  RuntimeApiTrace_t *traceData = (RuntimeApiTrace_t*)userdata;
  size_t bytesRead; 
     
  // This callback is enabled only for launch so we shouldn't see anything else.
  if ((cbid != CUPTI_RUNTIME_TRACE_CBID_cudaLaunch_v3020) &&
      (cbid != CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000))
  {
    printf("%s:%d: unexpected cbid %d\n", __FILE__, __LINE__, cbid);
    exit(-1);
  }

  if (cbInfo->callbackSite == CUPTI_API_ENTER) {
    printf("callback enter\n");
    cudaDeviceSynchronize();
    cuptiErr = cuptiSetEventCollectionMode(cbInfo->context, 
                                           CUPTI_EVENT_COLLECTION_MODE_KERNEL);
    CHECK_CUPTI_ERROR(cuptiErr, "cuptiSetEventCollectionMode");

    cuptiErr = cuptiEventGroupEnable(traceData->eventData->eventGroup);
    CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupEnable");
    printf("enter: event group %d \n", traceData->eventData->eventGroup);
  }
    
  if (cbInfo->callbackSite == CUPTI_API_EXIT) {
    uint32_t numInstances = 0, i;
    uint64_t *values = NULL;
    size_t valueSize = sizeof(numInstances);
    printf("callback exit\n");
    printf("eixt: event group %d \n", traceData->eventData->eventGroup);

    cuptiErr = cuptiEventGroupGetAttribute(traceData->eventData->eventGroup, 
                                           CUPTI_EVENT_GROUP_ATTR_INSTANCE_COUNT, 
                                           &valueSize, &numInstances);
    CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupGetAttribute");

    bytesRead = sizeof (uint64_t) * numInstances;
    values = (uint64_t *) malloc(bytesRead);
    if (values == NULL) {
        printf("%s:%d: Out of memory\n", __FILE__, __LINE__);
        exit(-1);
    }
    cudaDeviceSynchronize();
    cuptiErr = cuptiEventGroupReadEvent(traceData->eventData->eventGroup, 
                                        CUPTI_EVENT_READ_FLAG_NONE, 
                                        traceData->eventData->eventId, 
                                        &bytesRead, values);
    CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupReadEvent");
    printf("Instance %d \n", numInstances);

    printf("eixt: event %d : %d \n", traceData->eventData->eventId, bytesRead);

    traceData->eventVal = 0;
    for (i=0; i<numInstances; i++) {
        traceData->eventVal += values[i];
	printf(" %d ", values[i]);
    }
    free(values);
    printf("\n event value %d \n", traceData->eventVal);

    cuptiErr = cuptiEventGroupDisable(traceData->eventData->eventGroup);
    CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupDisable");
  }
}

void register_callback()
{
  CUcontext context = 0;
  CUdevice dev = 0;
  CUresult err;
  int i;
  int computeCapabilityMajor=0;
  int computeCapabilityMinor=0;
  int deviceNum;
  int deviceCount;
  char deviceName[32];
  const char *eventName;
  uint32_t profile_all = 1;

  err = cuInit(0);
  CHECK_CU_ERROR(err, "cuInit");

  err = cuDeviceGetCount(&deviceCount);
  CHECK_CU_ERROR(err, "cuDeviceGetCount");

  if (deviceCount == 0) {
    printf("There is no device supporting CUDA.\n");
    exit(-1);
  }

  deviceNum = 0;
  printf("CUDA Device Number: %d\n", deviceNum);

  err = cuDeviceGet(&dev, deviceNum);
  CHECK_CU_ERROR(err, "cuDeviceGet");

  err = cuDeviceGetName(deviceName, 32, dev);
  CHECK_CU_ERROR(err, "cuDeviceGetName");

  printf("CUDA Device Name: %s\n", deviceName);

  err = cuDeviceGetAttribute(&computeCapabilityMajor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev);
  CHECK_CU_ERROR(err, "cuDeviceGetAttribute");
  
  err = cuDeviceGetAttribute(&computeCapabilityMinor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev);
  CHECK_CU_ERROR(err, "cuDeviceGetAttribute");

  err = cuCtxCreate(&context, 0, dev);
  CHECK_CU_ERROR(err, "cuCtxCreate");

    
  // Creating event group for profiling
  cuptiErr = cuptiEventGroupCreate(context, &cuptiEvent.eventGroup, 0);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupCreate");
    printf("eixt: event group %d \n", cuptiEvent.eventGroup);

  eventName = EVENT_NAME;

  cuptiErr = cuptiEventGetIdFromName(dev, eventName, &cuptiEvent.eventId);
  if (cuptiErr != CUPTI_SUCCESS) 
    { 
      printf("Invalid eventName: %s\n", eventName);
      exit(-1);
    }

  cuptiErr = cuptiEventGroupAddEvent(cuptiEvent.eventGroup, cuptiEvent.eventId);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupAddEvent");

  cuptiErr = cuptiEventGroupSetAttribute(cuptiEvent.eventGroup,
                                         CUPTI_EVENT_GROUP_ATTR_PROFILE_ALL_DOMAIN_INSTANCES,
                                         sizeof(profile_all), &profile_all);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupSetAttribute");

  trace.eventData = &cuptiEvent;

  cuptiErr = cuptiSubscribe(&subscriber, (CUpti_CallbackFunc)getEventValueCallback , &trace);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiSubscribe");

  cuptiErr = cuptiEnableCallback(1, subscriber, CUPTI_CB_DOMAIN_RUNTIME_API, 
                                 CUPTI_RUNTIME_TRACE_CBID_cudaLaunch_v3020);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEnableCallback");
  cuptiErr = cuptiEnableCallback(1, subscriber, CUPTI_CB_DOMAIN_RUNTIME_API, 
                                 CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEnableCallback");

  printf("register_callback \n");
 }

void remove_callback() {
  cuptiErr = cuptiEventGroupRemoveEvent(cuptiEvent.eventGroup, cuptiEvent.eventId);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupRemoveEvent");

  cuptiErr = cuptiEventGroupDestroy(cuptiEvent.eventGroup);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiEventGroupDestroy");

  cuptiErr = cuptiUnsubscribe(subscriber);
  CHECK_CUPTI_ERROR(cuptiErr, "cuptiUnsubscribe");

  cudaDeviceSynchronize();

  cudaDeviceSynchronize();
}