#define _REENTRANT

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

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <semaphore.h>
#include <assert.h>

#include "ezt_cuda.h"
#include "cuda_ev_codes.h"
#include "eztrace.h"

extern "C" {

  // This CUDA module uses 2 different mechanisms for recording events:
  // - CPU events are recorded using wrappers (usual mechanism in EZTrace)
  // - GPU events are recording using the CUPTI activity mechanism.
  // 
  // For the CUPTI mechanism, we ask CUDA to record events in a buffer. When we reach a synchronization point,
  // eztrace reads the buffer and record the corresponding events.


#define DEFAULT_BUF_SIZE (2 * 1024 * 1024)
  static size_t  __cuda_buf_size = DEFAULT_BUF_SIZE;
  static int __cupti_enabled = 0;
#define ALIGN_SIZE (8)
#define ALIGN_BUFFER(buffer, align)					\
  (((uintptr_t) (buffer) &((align)-1)) ? ((buffer)+(align) - ((uintptr_t)(buffer)&((align)-1))) : (buffer))


#define CUDA_CHECK_RETVAL(__func__, __expected_retval__)		\
  do {								\
    if(__func__ != __expected_retval__) {			\
      fprintf(stderr, "EZTrace: CUDA returned an error\n");	\
      abort();							\
    }								\
  } while(0)


  /****
   * CUPTI-based functions
   ****/

  /**
   * Allocate a new BUF_SIZE buffer for CUPTI
   */
  static void bufferRequested(uint8_t **buffer, size_t *size, size_t *maxNumRecords)
  {
    uint8_t *bfr = (uint8_t *) malloc(__cuda_buf_size + ALIGN_SIZE);
    if (bfr == NULL) {
      fprintf(stderr, "Error: out of memory\n");
      exit(-1);
    }
  
    *size = __cuda_buf_size;
    *buffer = ALIGN_BUFFER(bfr, ALIGN_SIZE);
    *maxNumRecords = 0;
  }

  // todo: create one timestamp per GPU
  static uint64_t startTimestamp = 0;

  int __ezt_cuda_initialized = 0;
  void __ezt_cuda_initialize() {
    if(!__ezt_cuda_initialized) {
      __ezt_cuda_initialized = 1;

      if(!__cupti_enabled) {
	fprintf(stderr, "[EZTrace-CUDA] Only CPU events will be recorded. To enable the recording of GPU-events, set the EZTRACE_CUDA_CUPTI_ENABLED environment variable\n");
      }

      int num_gpus = 0;
      cudaGetDeviceCount(&num_gpus);
      CUDA_CHECK_RETVAL(cuptiGetTimestamp(&startTimestamp), CUPTI_SUCCESS);
      EZTRACE_EVENT_PACKED_0(EZTRACE_CUDA_CUINIT);
      EZTRACE_EVENT_PACKED_1(EZTRACE_CUDA_SETNBDEVICES, num_gpus);
    }
  }

  /* set to 1 when all the hooks are set....
   * This is usefull in order to avoid recursive calls
   */
  static int __cuda_initialized = 0;

  static const char *
  getMemcpyKindString(CUpti_ActivityMemcpyKind kind)
  {
    switch (kind) {
    case CUPTI_ACTIVITY_MEMCPY_KIND_HTOD:
      return "HtoD";
    case CUPTI_ACTIVITY_MEMCPY_KIND_DTOH:
      return "DtoH";
    case CUPTI_ACTIVITY_MEMCPY_KIND_HTOA:
      return "HtoA";
    case CUPTI_ACTIVITY_MEMCPY_KIND_ATOH:
      return "AtoH";
    case CUPTI_ACTIVITY_MEMCPY_KIND_ATOA:
      return "AtoA";
    case CUPTI_ACTIVITY_MEMCPY_KIND_ATOD:
      return "AtoD";
    case CUPTI_ACTIVITY_MEMCPY_KIND_DTOA:
      return "DtoA";
    case CUPTI_ACTIVITY_MEMCPY_KIND_DTOD:
      return "DtoD";
    case CUPTI_ACTIVITY_MEMCPY_KIND_HTOH:
      return "HtoH";
    default:
      break;
    }

    return "<unknown>";
  }

  static void __process_memcpy_record(CUpti_ActivityMemcpy2 *memcpy_info) {

    float start_date =  memcpy_info->start - startTimestamp;
    float stop_date = memcpy_info->end - startTimestamp;
    ezt_cuda_size_t cpy_size = memcpy_info->bytes;
    enum ezt_cudaMemcpyKind type = CUPTI_ACTIVITY_MEMCPY_KIND_TO_EZT((CUpti_ActivityMemcpyKind)memcpy_info->copyKind);
    int srcDeviceId = memcpy_info->deviceId;
    int destDeviceId = memcpy_info->dstDeviceId;

    EZTRACE_EVENT_PACKED_6(EZTRACE_CUDA_GPU_MEMCPY_TIMESTAMPS,
			   start_date,
			   stop_date,
			   cpy_size,
			   type,
			   srcDeviceId,
			   destDeviceId);
  }

  struct kernel_info_t{
    const char* name;
  };
  static uint32_t nb_kernels = 0; 
  static uint32_t nb_allocated_kernels = 0;
  static struct kernel_info_t *kernels;

  static kernel_id_t __register_kernel(const char* kernel_name) {
    uint32_t i;
    for(i=0; i<nb_kernels; i++) {
      if(!strcmp(kernels[i].name, kernel_name)) {
	// the kernel is already registered
	return i;
      }
    }

    // the kernel is not registered yet. register it
    nb_kernels++;
    if(nb_kernels > nb_allocated_kernels) {
      // not enough space in the array. expand the array
      nb_allocated_kernels *= 2;
      struct kernel_info_t *ptr = (struct kernel_info_t *)realloc(kernels, nb_allocated_kernels*sizeof(struct kernel_info_t));
      if(!ptr) {
	fprintf(stderr, "[EZTrace] Cannot allocate memory. Aborting\n");
	abort();
      }
      kernels = ptr;
    }

    // register the kernel
    kernels[nb_kernels-1].name = kernel_name;
    EZTRACE_EVENT_PACKED_1(EZTRACE_CUDA_REGISTER_KERNEL, nb_kernels-1);
    litl_write_probe_raw(__ezt_trace.litl_trace, EZTRACE_CUDA_KERNEL_NAME, strlen(kernels[nb_kernels-1].name), (litl_data_t*)kernels[nb_kernels-1].name);
    return nb_kernels-1;
  }

  //#define HAVE_CUPTI_ACTIVITY_KERNEL3 1
#if CUPTI_API_VERSION >= 7
  /* as of cupti 6.5, CUpti_ActivityKernel2 is deprecated, and CUpti_ActivityKernel3 should be used  */
  typedef CUpti_ActivityKernel3 __CUpti_ActivityKernel ;
#else
  typedef  CUpti_ActivityKernel2 __CUpti_ActivityKernel;
#endif
  
static void __process_kernel_record(__CUpti_ActivityKernel *kernel) {
    static unsigned next_kernel_id = 0; // todo: use the correlationId or something similar to identify the kernels
    unsigned current_id = next_kernel_id++;

    kernel_id_t kernel_id = __register_kernel(kernel->name);

    EZTRACE_EVENT_PACKED_5(EZTRACE_CUDA_GPU_KERNEL_TIMESTAMPS,
			   current_id,
			   (float)(kernel->start - startTimestamp),
			   (float)(kernel->end - startTimestamp),
			   kernel_id,
			   (uint32_t)kernel->deviceId);
  }


static sem_t __flush_sem;

void CUPTIAPI bufferCompleted(CUcontext ctx, uint32_t streamId, uint8_t *buffer, size_t size, size_t validSize)
{
  CUptiResult status;
  CUpti_Activity *record = NULL;

  if (validSize > 0) {

    __ezt_cuda_initialize();
    if(sem_wait(&__flush_sem) != 0) {
       abort();
    }
    do {
      status = cuptiActivityGetNextRecord(buffer, validSize, &record);
      if(status == CUPTI_SUCCESS) {
	switch(record->kind) {
	case CUPTI_ACTIVITY_KIND_MEMCPY:
	  {
	    __process_memcpy_record((CUpti_ActivityMemcpy2 *)record);
	    break;
	  }
	case CUPTI_ACTIVITY_KIND_KERNEL:
	case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL:
	  {
	    __process_kernel_record((__CUpti_ActivityKernel *)record);
	    break;
	  }
	default:
	  printf("  <unknown %d>\n", record->kind);
	  break;
	}
      }
      else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED) {
	break;
      }
      else {
	CUDA_CHECK_RETVAL(status, CUPTI_SUCCESS);
      }
    } while (1);

    if(sem_post(&__flush_sem) != 0) {
      abort();
    }

    // report any records dropped from the queue
    size_t dropped;
    CUDA_CHECK_RETVAL(cuptiActivityGetNumDroppedRecords(ctx, streamId, &dropped), CUPTI_SUCCESS);
    if (dropped != 0) {
      fprintf(stderr, "[EZTrace-CUDA] %u events were dropped because the buffer dedicated to CUDA events is too small (%u bytes)\n", (unsigned int)dropped, __cuda_buf_size);
      fprintf(stderr, "[EZTrace-CUDA]\tYou can change this buffer size by setting EZTRACE_CUDA_BUFFER_SIZE\n");
    }
  }

  free(buffer);
}



  /**** CUDA runtime interface ****/

  cudaError_t (*libcudaDeviceSynchronize)();
  cudaError_t (*libcudaThreadSynchronize)(); // note: deprecated

  /* Kernel management */
  cudaError_t (*libcudaLaunch)(const void *func);
  cudaError_t (*libcudaConfigureCall) ( dim3  gridDim, dim3  blockDim, size_t  sharedMem, cudaStream_t  stream );

  /* Memory management */
  cudaError_t  (*libcudaMalloc)(void **devPtr, size_t size);
  cudaError_t  (*libcudaMallocHost)(void **ptr, size_t size);
  cudaError_t  (*libcudaMallocPitch)(void **devPtr, size_t *pitch, size_t width, size_t height);
  cudaError_t  (*libcudaMallocArray)(struct cudaArray **array, const struct cudaChannelFormatDesc *desc, size_t width, size_t height, unsigned int flags);
  cudaError_t  (*libcudaMalloc3D)(struct cudaPitchedPtr* pitchedDevPtr, struct cudaExtent extent);
  cudaError_t  (*libcudaMalloc3DArray)(cudaArray_t *array, const struct cudaChannelFormatDesc* desc, struct cudaExtent extent, unsigned int flags);
  cudaError_t  (*libcudaMallocMipmappedArray)(cudaMipmappedArray_t *mipmappedArray, const struct cudaChannelFormatDesc* desc, struct cudaExtent extent, unsigned int numLevels, unsigned int flags);
  cudaError_t  (*libcudaFree)(void *devPtr);
  cudaError_t  (*libcudaFreeHost)(void *ptr);
  cudaError_t  (*libcudaFreeArray)(cudaArray_t array);
  cudaError_t  (*libcudaFreeMipmappedArray)(cudaMipmappedArray_t mipmappedArray);

  cudaError_t  (*libcudaHostAlloc)(void **pHost, size_t size, unsigned int flags);
  cudaError_t  (*libcudaHostRegister)(void *ptr, size_t size, unsigned int flags);
  cudaError_t  (*libcudaHostUnregister)(void *ptr);


  /* Memory Transfers */
  cudaError_t  (*libcudaMemcpy3D)(const struct cudaMemcpy3DParms *p);
  cudaError_t  (*libcudaMemcpy3DPeer)(const struct cudaMemcpy3DPeerParms *p);
  cudaError_t  (*libcudaMemcpy3DAsync)(const struct cudaMemcpy3DParms *p, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpy3DPeerAsync)(const struct cudaMemcpy3DPeerParms *p, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpy)(void *dst, const void *src, size_t count, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpyPeer)(void *dst, int dstDevice, const void *src, int srcDevice, size_t count);
  cudaError_t  (*libcudaMemcpyToArray)(cudaArray_t dst, size_t wOffset, size_t hOffset, const void *src, size_t count, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpyFromArray)(void *dst, cudaArray_const_t src, size_t wOffset, size_t hOffset, size_t count, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpyArrayToArray)(cudaArray_t dst, size_t wOffsetDst, size_t hOffsetDst, cudaArray_const_t src, size_t wOffsetSrc, size_t hOffsetSrc, size_t count, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpy2D)(void *dst, size_t dpitch, const void *src, size_t spitch, size_t width, size_t height, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpy2DToArray)(cudaArray_t dst, size_t wOffset, size_t hOffset, const void *src, size_t spitch, size_t width, size_t height, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpy2DFromArray)(void *dst, size_t dpitch, cudaArray_const_t src, size_t wOffset, size_t hOffset, size_t width, size_t height, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpy2DArrayToArray)(cudaArray_t dst, size_t wOffsetDst, size_t hOffsetDst, cudaArray_const_t src, size_t wOffsetSrc, size_t hOffsetSrc, size_t width, size_t height, enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpyToSymbol)(const void *symbol, const void *src, size_t count, size_t offset , enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpyFromSymbol)(void *dst, const void *symbol, size_t count, size_t offset , enum cudaMemcpyKind kind);
  cudaError_t  (*libcudaMemcpyAsync)(void *dst, const void *src, size_t count, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpyPeerAsync)(void *dst, int dstDevice, const void *src, int srcDevice, size_t count, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpyToArrayAsync)(cudaArray_t dst, size_t wOffset, size_t hOffset, const void *src, size_t count, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpyFromArrayAsync)(void *dst, cudaArray_const_t src, size_t wOffset, size_t hOffset, size_t count, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpy2DAsync)(void *dst, size_t dpitch, const void *src, size_t spitch, size_t width, size_t height, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpy2DToArrayAsync)(cudaArray_t dst, size_t wOffset, size_t hOffset, const void *src, size_t spitch, size_t width, size_t height, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpy2DFromArrayAsync)(void *dst, size_t dpitch, cudaArray_const_t src, size_t wOffset, size_t hOffset, size_t width, size_t height, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpyToSymbolAsync)(const void *symbol, const void *src, size_t count, size_t offset, enum cudaMemcpyKind kind, cudaStream_t stream );
  cudaError_t  (*libcudaMemcpyFromSymbolAsync)(void *dst, const void *symbol, size_t count, size_t offset, enum cudaMemcpyKind kind, cudaStream_t stream );

  cudaError_t  (*libcudaMemset)(void *devPtr, int value, size_t count);
  cudaError_t  (*libcudaMemset2D)(void *devPtr, size_t pitch, int value, size_t width, size_t height);
  cudaError_t  (*libcudaMemset3D)(struct cudaPitchedPtr pitchedDevPtr, int value, struct cudaExtent extent);
  cudaError_t  (*libcudaMemsetAsync)(void *devPtr, int value, size_t count, cudaStream_t stream );
  cudaError_t  (*libcudaMemset2DAsync)(void *devPtr, size_t pitch, int value, size_t width, size_t height, cudaStream_t stream );
  cudaError_t  (*libcudaMemset3DAsync)(struct cudaPitchedPtr pitchedDevPtr, int value, struct cudaExtent extent, cudaStream_t stream );



  /**** CUDA driver interface ****/

#if 0
  CUresult (*libcuMemcpyAtoA_v2)(CUarray dstArray, size_t dstOffset, CUarray srcArray, size_t srcOffset, size_t ByteCount);
  CUresult (*libcuMemcpyAtoD_v2)(CUdeviceptr dstDevice, CUarray srcArray, size_t srcOffset, size_t ByteCount);
  CUresult (*libcuMemcpyAtoH_v2)(void *dstHost, CUarray srcArray, size_t srcOffset, size_t ByteCount);
  CUresult (*libcuMemcpyDtoA_v2)(CUarray dstArray, size_t dstOffset, CUdeviceptr srcDevice, size_t ByteCount);
  CUresult (*libcuMemcpyDtoD_v2)(CUdeviceptr dstDevice, CUdeviceptr srcDevice, size_t ByteCount);
  CUresult (*libcuMemcpyDtoH_v2)(void *dstHost, CUdeviceptr srcDevice, size_t ByteCount);
  CUresult (*libcuMemcpyHtoA_v2)(CUarray dstArray, size_t dstOffset, const void* srcHost, size_t ByteCount);
  CUresult (*libcuMemcpyHtoD_v2)(CUdeviceptr dstDevice, const void *srcHost, size_t ByteCount);
  CUresult (*libcuMemFree_v2)(CUdeviceptr dptr);
  CUresult (*libcuMemAlloc_v2)(CUdeviceptr *dptr, size_t bytesize);
  CUresult (*libcuMemcpy) (CUdeviceptr dst, CUdeviceptr src, size_t ByteCount);
  CUresult (*libcuLaunchKernel)(CUfunction  	f,
				unsigned int  	gridDimX,
				unsigned int  	gridDimY,
				unsigned int  	gridDimZ,
				unsigned int  	blockDimX,
				unsigned int  	blockDimY,
				unsigned int  	blockDimZ,
				unsigned int  	sharedMemBytes,
				CUstream  	hStream,
				void **  	kernelParams,
				void **  	extra);
#endif

  /* Kernel management */
  CUresult  (*libcuLaunchKernel)(CUfunction f,
				 unsigned int gridDimX,
				 unsigned int gridDimY,
				 unsigned int gridDimZ,
				 unsigned int blockDimX,
				 unsigned int blockDimY,
				 unsigned int blockDimZ,
				 unsigned int sharedMemBytes,
				 CUstream hStream,
				 void **kernelParams,
				 void **extra);
  CUresult  (*libcuLaunch)(CUfunction f);
  CUresult  (*libcuLaunchGrid)(CUfunction f, int grid_width, int grid_height);
  CUresult  (*libcuLaunchGridAsync)(CUfunction f, int grid_width, int grid_height, CUstream hStream);

  /* Memory Management */
  CUresult  (*libcuMemGetInfo_v2)(size_t *free, size_t *total);
  CUresult  (*libcuMemAlloc_v2)(CUdeviceptr *dptr, size_t bytesize);
  CUresult  (*libcuMemAllocPitch_v2)(CUdeviceptr *dptr, size_t *pPitch, size_t WidthInBytes, size_t Height, unsigned int ElementSizeBytes);
  CUresult  (*libcuMemFree_v2)(CUdeviceptr dptr);
  CUresult  (*libcuMemAllocHost_v2)(void **pp, size_t bytesize);
  CUresult  (*libcuMemFreeHost)(void *p);
  CUresult  (*libcuMemHostAlloc)(void **pp, size_t bytesize, unsigned int Flags);
  CUresult  (*libcuMemHostRegister)(void *p, size_t bytesize, unsigned int Flags);
  CUresult  (*libcuMemHostUnregister)(void *p);

  /* Memory Transfers */
  CUresult  (*libcuMemcpy)(CUdeviceptr dst, CUdeviceptr src, size_t ByteCount);
  CUresult  (*libcuMemcpyPeer)(CUdeviceptr dstDevice, CUcontext dstContext, CUdeviceptr srcDevice, CUcontext srcContext, size_t ByteCount);
  CUresult  (*libcuMemcpyHtoD_v2)(CUdeviceptr dstDevice, const void *srcHost, size_t ByteCount);
  CUresult  (*libcuMemcpyDtoH_v2)(void *dstHost, CUdeviceptr srcDevice, size_t ByteCount);
  CUresult  (*libcuMemcpyDtoD_v2)(CUdeviceptr dstDevice, CUdeviceptr srcDevice, size_t ByteCount);
  CUresult  (*libcuMemcpyDtoA_v2)(CUarray dstArray, size_t dstOffset, CUdeviceptr srcDevice, size_t ByteCount);
  CUresult  (*libcuMemcpyAtoD_v2)(CUdeviceptr dstDevice, CUarray srcArray, size_t srcOffset, size_t ByteCount);
  CUresult  (*libcuMemcpyHtoA_v2)(CUarray dstArray, size_t dstOffset, const void *srcHost, size_t ByteCount);
  CUresult  (*libcuMemcpyAtoH_v2)(void *dstHost, CUarray srcArray, size_t srcOffset, size_t ByteCount);
  CUresult  (*libcuMemcpyAtoA_v2)(CUarray dstArray, size_t dstOffset, CUarray srcArray, size_t srcOffset, size_t ByteCount);
  CUresult  (*libcuMemcpy2D_v2)(const CUDA_MEMCPY2D *pCopy);
  CUresult  (*libcuMemcpy2DUnaligned_v2)(const CUDA_MEMCPY2D *pCopy);
  CUresult  (*libcuMemcpy3D_v2)(const CUDA_MEMCPY3D *pCopy);
  CUresult  (*libcuMemcpy3DPeer)(const CUDA_MEMCPY3D_PEER *pCopy);
  CUresult  (*libcuMemcpyAsync)(CUdeviceptr dst, CUdeviceptr src, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpyPeerAsync)(CUdeviceptr dstDevice, CUcontext dstContext, CUdeviceptr srcDevice, CUcontext srcContext, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpyHtoDAsync_v2)(CUdeviceptr dstDevice, const void *srcHost, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpyDtoHAsync_v2)(void *dstHost, CUdeviceptr srcDevice, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpyDtoDAsync_v2)(CUdeviceptr dstDevice, CUdeviceptr srcDevice, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpyHtoAAsync_v2)(CUarray dstArray, size_t dstOffset, const void *srcHost, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpyAtoHAsync_v2)(void *dstHost, CUarray srcArray, size_t srcOffset, size_t ByteCount, CUstream hStream);
  CUresult  (*libcuMemcpy2DAsync_v2)(const CUDA_MEMCPY2D *pCopy, CUstream hStream);
  CUresult  (*libcuMemcpy3DAsync_v2)(const CUDA_MEMCPY3D *pCopy, CUstream hStream);
  CUresult  (*libcuMemcpy3DPeerAsync)(const CUDA_MEMCPY3D_PEER *pCopy, CUstream hStream);



  CUresult  (*libcuMemsetD8_v2)(CUdeviceptr dstDevice, unsigned char uc, size_t N);
  CUresult  (*libcuMemsetD16_v2)(CUdeviceptr dstDevice, unsigned short us, size_t N);
  CUresult  (*libcuMemsetD32_v2)(CUdeviceptr dstDevice, unsigned int ui, size_t N);
  CUresult  (*libcuMemsetD2D8_v2)(CUdeviceptr dstDevice, size_t dstPitch, unsigned char uc, size_t Width, size_t Height);
  CUresult  (*libcuMemsetD2D16_v2)(CUdeviceptr dstDevice, size_t dstPitch, unsigned short us, size_t Width, size_t Height);
  CUresult  (*libcuMemsetD2D32_v2)(CUdeviceptr dstDevice, size_t dstPitch, unsigned int ui, size_t Width, size_t Height);
  CUresult  (*libcuMemsetD8Async)(CUdeviceptr dstDevice, unsigned char uc, size_t N, CUstream hStream);
  CUresult  (*libcuMemsetD16Async)(CUdeviceptr dstDevice, unsigned short us, size_t N, CUstream hStream);
  CUresult  (*libcuMemsetD32Async)(CUdeviceptr dstDevice, unsigned int ui, size_t N, CUstream hStream);
  CUresult  (*libcuMemsetD2D8Async)(CUdeviceptr dstDevice, size_t dstPitch, unsigned char uc, size_t Width, size_t Height, CUstream hStream);
  CUresult  (*libcuMemsetD2D16Async)(CUdeviceptr dstDevice, size_t dstPitch, unsigned short us, size_t Width, size_t Height, CUstream hStream);
  CUresult  (*libcuMemsetD2D32Async)(CUdeviceptr dstDevice, size_t dstPitch, unsigned int ui, size_t Width, size_t Height, CUstream hStream);

#if 0
  /* is it useful ? */
  CUresult  (*libcuArrayCreate)(CUarray *pHandle, const CUDA_ARRAY_DESCRIPTOR *pAllocateArray);
  CUresult  (*libcuArrayGetDescriptor)(CUDA_ARRAY_DESCRIPTOR *pArrayDescriptor, CUarray hArray);
  CUresult  (*libcuArrayDestroy)(CUarray hArray);
  CUresult  (*libcuArray3DCreate)(CUarray *pHandle, const CUDA_ARRAY3D_DESCRIPTOR *pAllocateArray);
  CUresult  (*libcuArray3DGetDescriptor)(CUDA_ARRAY3D_DESCRIPTOR *pArrayDescriptor, CUarray hArray);
  CUresult  (*libcuMipmappedArrayCreate)(CUmipmappedArray *pHandle, const CUDA_ARRAY3D_DESCRIPTOR *pMipmappedArrayDesc, unsigned int numMipmapLevels);
  CUresult  (*libcuMipmappedArrayGetLevel)(CUarray *pLevelArray, CUmipmappedArray hMipmappedArray, unsigned int level);
  CUresult  (*libcuMipmappedArrayDestroy)(CUmipmappedArray hMipmappedArray);
#endif

  // CUresult  (*libcuCtxSynchronize)(void);
  // CUresult  (*libcuDeviceTotalMem)(unsigned int *bytes, CUdevice dev);

#if 0
  START_INTERCEPT
    INTERCEPT2("cudaLaunch", libcudaLaunch)
    INTERCEPT2("cudaMalloc", libcudaMalloc)
    INTERCEPT2("cudaFree", libcudaFree)
    INTERCEPT2("cudaConfigureCall", libcudaConfigureCall)
    INTERCEPT2("cudaMemcpy", libcudaMemcpy)
    INTERCEPT2("cudaMemcpyAsync", libcudaMemcpyAsync)
    INTERCEPT2("cudaDeviceSynchronize", libcudaDeviceSynchronize)
    INTERCEPT2("cudaThreadSynchronize", libcudaThreadSynchronize)
    INTERCEPT2("cuLaunchKernel", libcuLaunchKernel)
    INTERCEPT2("cuMemcpy", libcuMemcpy)
    INTERCEPT2("cuMemcpyAtoA_v2", libcuMemcpyAtoA_v2)
    INTERCEPT2("cuMemcpyAtoD_v2", libcuMemcpyAtoD_v2)
    INTERCEPT2("cuMemcpyAtoH_v2", libcuMemcpyAtoH_v2)
    INTERCEPT2("cuMemcpyDtoA_v2", libcuMemcpyDtoA_v2)
    INTERCEPT2("cuMemcpyDtoD_v2", libcuMemcpyDtoD_v2)
    INTERCEPT2("cuMemcpyDtoH_v2", libcuMemcpyDtoH_v2)
    INTERCEPT2("cuMemcpyHtoA_v2", libcuMemcpyHtoA_v2)
    INTERCEPT2("cuMemcpyHtoD_v2", libcuMemcpyHtoD_v2)
    INTERCEPT2("cuMemFree_v2",libcuMemFree_v2)
    INTERCEPT2("cuMemAlloc_v2", libcuMemAlloc_v2)
  END_INTERCEPT;
#endif

  START_INTERCEPT_MODULE(cuda)
  INTERCEPT2("cudaDeviceSynchronize", libcudaDeviceSynchronize)
  INTERCEPT2("cudaThreadSynchronize", libcudaThreadSynchronize)
  INTERCEPT2("cudaLaunch", libcudaLaunch)
  INTERCEPT2("cudaConfigureCall", libcudaConfigureCall)
  INTERCEPT2("cudaMalloc", libcudaMalloc)
  INTERCEPT2("cudaMallocHost", libcudaMallocHost)
  INTERCEPT2("cudaMallocPitch", libcudaMallocPitch)
  INTERCEPT2("cudaMallocArray", libcudaMallocArray)
  INTERCEPT2("cudaMalloc3D", libcudaMalloc3D)
  INTERCEPT2("cudaMalloc3DArray", libcudaMalloc3DArray)
  INTERCEPT2("cudaMallocMipmappedArray", libcudaMallocMipmappedArray)
  INTERCEPT2("cudaFree", libcudaFree)
  INTERCEPT2("cudaFreeHost", libcudaFreeHost)
  INTERCEPT2("cudaFreeArray", libcudaFreeArray)
  INTERCEPT2("cudaFreeMipmappedArray", libcudaFreeMipmappedArray)
  INTERCEPT2("cudaHostAlloc", libcudaHostAlloc)
  INTERCEPT2("cudaHostRegister", libcudaHostRegister)
  INTERCEPT2("cudaHostUnregister", libcudaHostUnregister)
  INTERCEPT2("cudaMemcpy3D", libcudaMemcpy3D)
  INTERCEPT2("cudaMemcpy3DPeer", libcudaMemcpy3DPeer)
  INTERCEPT2("cudaMemcpy3DAsync", libcudaMemcpy3DAsync)
  INTERCEPT2("cudaMemcpy3DPeerAsync", libcudaMemcpy3DPeerAsync)
  INTERCEPT2("cudaMemcpy", libcudaMemcpy)
  INTERCEPT2("cudaMemcpyPeer", libcudaMemcpyPeer)
  INTERCEPT2("cudaMemcpyToArray", libcudaMemcpyToArray)
  INTERCEPT2("cudaMemcpyFromArray", libcudaMemcpyFromArray)
  INTERCEPT2("cudaMemcpyArrayToArray", libcudaMemcpyArrayToArray)
  INTERCEPT2("cudaMemcpy2D", libcudaMemcpy2D)
  INTERCEPT2("cudaMemcpy2DToArray", libcudaMemcpy2DToArray)
  INTERCEPT2("cudaMemcpy2DFromArray", libcudaMemcpy2DFromArray)
  INTERCEPT2("cudaMemcpy2DArrayToArray", libcudaMemcpy2DArrayToArray)
  INTERCEPT2("cudaMemcpyToSymbol", libcudaMemcpyToSymbol)
  INTERCEPT2("cudaMemcpyFromSymbol", libcudaMemcpyFromSymbol)
  INTERCEPT2("cudaMemcpyAsync", libcudaMemcpyAsync)
  INTERCEPT2("cudaMemcpyPeerAsync", libcudaMemcpyPeerAsync)
  INTERCEPT2("cudaMemcpyToArrayAsync", libcudaMemcpyToArrayAsync)
  INTERCEPT2("cudaMemcpyFromArrayAsync", libcudaMemcpyFromArrayAsync)
  INTERCEPT2("cudaMemcpy2DAsync", libcudaMemcpy2DAsync)
  INTERCEPT2("cudaMemcpy2DToArrayAsync", libcudaMemcpy2DToArrayAsync)
  INTERCEPT2("cudaMemcpy2DFromArrayAsync", libcudaMemcpy2DFromArrayAsync)
  INTERCEPT2("cudaMemcpyToSymbolAsync", libcudaMemcpyToSymbolAsync)
  INTERCEPT2("cudaMemcpyFromSymbolAsync", libcudaMemcpyFromSymbolAsync)
  INTERCEPT2("cudaMemset", libcudaMemset)
  INTERCEPT2("cudaMemset2D", libcudaMemset2D)
  INTERCEPT2("cudaMemset3D", libcudaMemset3D)
  INTERCEPT2("cudaMemsetAsync", libcudaMemsetAsync)
  INTERCEPT2("cudaMemset2DAsync", libcudaMemset2DAsync)
  INTERCEPT2("cudaMemset3DAsync", libcudaMemset3DAsync)
  INTERCEPT2("cuMemcpyAtoA_v2", libcuMemcpyAtoA_v2)
  INTERCEPT2("cuMemcpyAtoD_v2", libcuMemcpyAtoD_v2)
  INTERCEPT2("cuMemcpyAtoH_v2", libcuMemcpyAtoH_v2)
  INTERCEPT2("cuMemcpyDtoA_v2", libcuMemcpyDtoA_v2)
  INTERCEPT2("cuMemcpyDtoD_v2", libcuMemcpyDtoD_v2)
  INTERCEPT2("cuMemcpyDtoH_v2", libcuMemcpyDtoH_v2)
  INTERCEPT2("cuMemcpyHtoA_v2", libcuMemcpyHtoA_v2)
  INTERCEPT2("cuMemcpyHtoD_v2", libcuMemcpyHtoD_v2)
  INTERCEPT2("cuMemFree_v2", libcuMemFree_v2)
  INTERCEPT2("cuMemAlloc_v2", libcuMemAlloc_v2)
  INTERCEPT2("cuMemcpy", libcuMemcpy)
  INTERCEPT2("cuLaunchKernel", libcuLaunchKernel)
  INTERCEPT2("cuLaunchKernel", libcuLaunchKernel)
  INTERCEPT2("cuLaunch", libcuLaunch)
  INTERCEPT2("cuLaunchGrid", libcuLaunchGrid)
  INTERCEPT2("cuLaunchGridAsync", libcuLaunchGridAsync)
  INTERCEPT2("cuMemGetInfo_v2", libcuMemGetInfo_v2)
  INTERCEPT2("cuMemAlloc_v2", libcuMemAlloc_v2)
  INTERCEPT2("cuMemAllocPitch_v2", libcuMemAllocPitch_v2)
  INTERCEPT2("cuMemFree_v2", libcuMemFree_v2)
  INTERCEPT2("cuMemAllocHost_v2", libcuMemAllocHost_v2)
  INTERCEPT2("cuMemFreeHost", libcuMemFreeHost)
  INTERCEPT2("cuMemHostAlloc", libcuMemHostAlloc)
  INTERCEPT2("cuMemHostRegister", libcuMemHostRegister)
  INTERCEPT2("cuMemHostUnregister", libcuMemHostUnregister)
  INTERCEPT2("cuMemcpy", libcuMemcpy)
  INTERCEPT2("cuMemcpyPeer", libcuMemcpyPeer)
  INTERCEPT2("cuMemcpyHtoD_v2", libcuMemcpyHtoD_v2)
  INTERCEPT2("cuMemcpyDtoH_v2", libcuMemcpyDtoH_v2)
  INTERCEPT2("cuMemcpyDtoD_v2", libcuMemcpyDtoD_v2)
  INTERCEPT2("cuMemcpyDtoA_v2", libcuMemcpyDtoA_v2)
  INTERCEPT2("cuMemcpyAtoD_v2", libcuMemcpyAtoD_v2)
  INTERCEPT2("cuMemcpyHtoA_v2", libcuMemcpyHtoA_v2)
  INTERCEPT2("cuMemcpyAtoH_v2", libcuMemcpyAtoH_v2)
  INTERCEPT2("cuMemcpyAtoA_v2", libcuMemcpyAtoA_v2)
  INTERCEPT2("cuMemcpy2D_v2", libcuMemcpy2D_v2)
  INTERCEPT2("cuMemcpy2DUnaligned_v2", libcuMemcpy2DUnaligned_v2)
  INTERCEPT2("cuMemcpy3D_v2", libcuMemcpy3D_v2)
  INTERCEPT2("cuMemcpy3DPeer", libcuMemcpy3DPeer)
  INTERCEPT2("cuMemcpyAsync", libcuMemcpyAsync)
  INTERCEPT2("cuMemcpyPeerAsync", libcuMemcpyPeerAsync)
  INTERCEPT2("cuMemcpyHtoDAsync_v2", libcuMemcpyHtoDAsync_v2)
  INTERCEPT2("cuMemcpyDtoHAsync_v2", libcuMemcpyDtoHAsync_v2)
  INTERCEPT2("cuMemcpyDtoDAsync_v2", libcuMemcpyDtoDAsync_v2)
  INTERCEPT2("cuMemcpyHtoAAsync_v2", libcuMemcpyHtoAAsync_v2)
  INTERCEPT2("cuMemcpyAtoHAsync_v2", libcuMemcpyAtoHAsync_v2)
  INTERCEPT2("cuMemcpy2DAsync_v2", libcuMemcpy2DAsync_v2)
  INTERCEPT2("cuMemcpy3DAsync_v2", libcuMemcpy3DAsync_v2)
  INTERCEPT2("cuMemcpy3DPeerAsync", libcuMemcpy3DPeerAsync)
  INTERCEPT2("cuMemsetD8_v2", libcuMemsetD8_v2)
  INTERCEPT2("cuMemsetD16_v2", libcuMemsetD16_v2)
  INTERCEPT2("cuMemsetD32_v2", libcuMemsetD32_v2)
  INTERCEPT2("cuMemsetD2D8_v2", libcuMemsetD2D8_v2)
  INTERCEPT2("cuMemsetD2D16_v2", libcuMemsetD2D16_v2)
  INTERCEPT2("cuMemsetD2D32_v2", libcuMemsetD2D32_v2)
  INTERCEPT2("cuMemsetD8Async", libcuMemsetD8Async)
  INTERCEPT2("cuMemsetD16Async", libcuMemsetD16Async)
  INTERCEPT2("cuMemsetD32Async", libcuMemsetD32Async)
  INTERCEPT2("cuMemsetD2D8Async", libcuMemsetD2D8Async)
  INTERCEPT2("cuMemsetD2D16Async", libcuMemsetD2D16Async)
  INTERCEPT2("cuMemsetD2D32Async", libcuMemsetD2D32Async)
  END_INTERCEPT_MODULE(cuda);

#if 0
  INTERCEPT2("cuArrayCreate", libcuArrayCreate)
  INTERCEPT2("cuArrayGetDescriptor", libcuArrayGetDescriptor)
  INTERCEPT2("cuArrayDestroy", libcuArrayDestroy)
  INTERCEPT2("cuArray3DCreate", libcuArray3DCreate)
  INTERCEPT2("cuArray3DGetDescriptor", libcuArray3DGetDescriptor)
  INTERCEPT2("cuMipmappedArrayCreate", libcuMipmappedArrayCreate)
  INTERCEPT2("cuMipmappedArrayGetLevel", libcuMipmappedArrayGetLevel)
  INTERCEPT2("cuMipmappedArrayDestroy", libcuMipmappedArrayDestroy)

  INTERCEPT2("libcuCtxSynchronize", libcuCtxSynchronize)
  INTERCEPT2("libcuDeviceTotalMem", libcuDeviceTotalMem)
#endif

  void CUPTIAPI
  my_callback(void *userdata, CUpti_CallbackDomain domain,
	      CUpti_CallbackId cbid, const void *cbdata) {
    const CUpti_CallbackData *cbInfo = (CUpti_CallbackData *)cbdata;

    if(domain == CUPTI_CB_DOMAIN_RUNTIME_API) {
      switch(cbid) {
      case CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000:
      case CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_ptsz_v7000:
      case CUPTI_RUNTIME_TRACE_CBID_cudaLaunchCooperativeKernel_v9000:
      case CUPTI_RUNTIME_TRACE_CBID_cudaLaunchCooperativeKernel_ptsz_v9000:
      case CUPTI_RUNTIME_TRACE_CBID_cudaLaunchCooperativeKernelMultiDevice_v9000:
	{
	  cuLaunchKernel_params *params = (cuLaunchKernel_params*) (cbInfo->functionParams);
	  int deviceId;
	  int ret = cudaGetDevice(&deviceId);
	  assert(ret == cudaSuccess);
	  static __thread int current_id =0;
	  if (cbInfo->callbackSite == CUPTI_API_ENTER) {
	    static unsigned kernel_id = 0;
	    current_id = kernel_id++;
	    EZTRACE_EVENT_PACKED_2(EZTRACE_CUDA_CULAUNCHKERNEL_START, deviceId, current_id);
	  } else {
	    EZTRACE_EVENT_PACKED_2(EZTRACE_CUDA_CULAUNCHKERNEL_STOP, deviceId, current_id);
	  }
	}
	break;
      }
    }

  }

CUpti_SubscriberHandle subscriber;

  static void __init_buffer_size() {
    char* str = getenv("EZTRACE_CUDA_BUFFER_SIZE");
    if(str) {
      __cuda_buf_size = atoi(str);
      fprintf(stderr, "Setting CUDA buffer size to %d bytes\n", __cuda_buf_size);
    }
  }

  static void __init_cupti() {
    char* str = getenv("EZTRACE_CUDA_CUPTI_DISABLED");
    if(str) {
      __cupti_enabled = 0;
    } else {
      __cupti_enabled = 1;
    }

    if(__cupti_enabled) {

      fprintf(stderr, "[EZTrace][CUDA] CUPTI is enabled\n");
      __cupti_enabled = 1;

      kernels = (struct kernel_info_t*) malloc(sizeof(struct kernel_info_t)*1024);
      nb_allocated_kernels = 1024;

      // device activity record is created when CUDA initializes, so we
      // want to enable it before cuInit() or any CUDA runtime call
      CUDA_CHECK_RETVAL(cuptiActivityEnable(CUPTI_ACTIVITY_KIND_KERNEL), CUPTI_SUCCESS);
      CUDA_CHECK_RETVAL(cuptiActivityEnable(CUPTI_ACTIVITY_KIND_MEMCPY), CUPTI_SUCCESS);
      CUDA_CHECK_RETVAL(cuptiActivityEnable(CUPTI_ACTIVITY_KIND_MEMSET), CUPTI_SUCCESS);

      // Register callbacks for buffer requests and for buffers completed by CUPTI.
      CUDA_CHECK_RETVAL(cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted), CUPTI_SUCCESS);

      if(sem_init(&__flush_sem, 0, 1) != 0 ){
        abort();
      }


  cuptiSubscribe(&subscriber, 
                 (CUpti_CallbackFunc)my_callback , NULL);
  cuptiEnableDomain(1, subscriber, 
                    CUPTI_CB_DOMAIN_RUNTIME_API);
    }
  }

  // flush all pending buffers
  static void __ezt_cuda_flush(void*param) {
      if(__cupti_enabled) {
	CUDA_CHECK_RETVAL(cuptiActivityFlushAll(CUPTI_ACTIVITY_FLAG_FORCE_INT), CUPTI_SUCCESS);

	// make sure there isn't another thread that is flushing the cuda buffer
	if(sem_wait(&__flush_sem) != 0) {
	  abort();
	}
      }
  }

  static void __cuda_init (void) __attribute__ ((constructor));
  static void 
    __cuda_init (void)
  {
    DYNAMIC_INTERCEPT_ALL_MODULE(cuda);

    __init_buffer_size();

    __init_cupti();

    // we need to flush all the pending buffers before closing the trace so that we don't lose
    // any event from the GPU
    eztrace_atexit(__ezt_cuda_flush, NULL);

#ifdef EZTRACE_AUTOSTART
    eztrace_start();
#endif
    __cuda_initialized = 1;
  }

  static void __cuda_conclude (void) __attribute__ ((destructor));
  static void __cuda_conclude (void)
  {

    __cuda_initialized = 0;
    if(__cupti_enabled) {
      cuptiUnsubscribe(subscriber);
    }

    eztrace_stop();
  }
}
