/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright © CNRS, INRIA, Université Bordeaux 1
 * See COPYING in top-level directory.
 */

#include <GTG.h>
#include <string.h>
#include <assert.h>
#include "ezt_demangle.h"
#include "eztrace_convert.h"
#include "cuda_ev_codes.h"
#include "eztrace_list.h"

static int recording_stats = 0;

struct cuda_memcpy_stat_t {
  unsigned nb_memcpy;
  ezt_cuda_size_t total_size;		/* total size memcpied */
  ezt_cuda_size_t min_size;		/* minimum size of memcpy */
  ezt_cuda_size_t max_size;		/* maximum size of memcpy */
};

struct cuda_malloc_stat_t {
  unsigned nb_malloc;
  ezt_cuda_size_t total_size;		/* total size allocated */
  ezt_cuda_size_t peak_mem;		/* peak memory consumption */
  ezt_cuda_size_t cur_mem;		/* current memory consumption */
};

struct cuda_device_info_t {
  struct ezt_list_token_t token;
  struct cuda_process_info_t *p_process;
  uint32_t device_id;
  char* gpu_id;
  char* gpu_name;
  int link_start, link_end;
  float start_time;
  int added;			/* set to 1 once the container is created in GTG */
  struct ezt_list_t pending_kernels;
  struct ezt_list_t finished_kernels;

  struct ezt_list_t pending_memcpy;
  struct ezt_list_t finished_memcpy;

  struct ezt_list_t pending_malloc; /* buffers that are currently in use (ie. not freed) */
  struct ezt_list_t finished_malloc;  /* buffers that have been freed */

  struct ezt_list_t kernel_stats;
  struct cuda_memcpy_stat_t memcpy_stats;
  struct cuda_malloc_stat_t malloc_stats;
};

struct cuda_process_info_t {
  struct process_info_t *p_process;
  float start_time;
  struct ezt_list_t kernel_names;
  struct ezt_list_t device_list;
};

struct cuda_kernel_name_t {
  kernel_id_t id;
  struct ezt_list_token_t token;
  char state_id[80];
  char state_desc[80];
  char name[80];
};

struct cuda_kernel_stat_t {
  struct cuda_kernel_name_t *p_kernel;
  struct cuda_device_info_t *p_device;
  unsigned nb_calls;
  double min_duration;
  double max_duration;
  double kernel_duration;	/* total duration of the kernel */
  struct ezt_list_token_t token;
};

struct cuda_kernel_info_t {
  /* time in seconds */
  float gpu_start_time;
  float gpu_stop_time;
  float cpu_start_time;
  float cpu_stop_time;
  unsigned kernel_id;
  struct cuda_kernel_name_t *name;
  char link_str[80];
  /* todo: store the parameters, etc. */
  struct ezt_list_token_t token;
};

struct cuda_memcpy_info_t {
  /* time in seconds */
  float gpu_start_time;
  float gpu_stop_time;
  float cpu_start_time;
  float cpu_stop_time;
  unsigned id;
  char link_str[80];
  char* src;
  char* dest;
  ezt_cuda_size_t size;
  enum ezt_cudaMemcpyKind kind;
  struct ezt_list_token_t token;
};

struct cuda_malloc_info_t {
  float cpu_start_time;		/* timestamp when cudaMalloc started */
  float cpu_stop_time;		/* timestamp when cudaMalloc stoped */
  ezt_cuda_size_t size;			/* size of the allocated buffer */
  app_ptr devPtr;		/* address of the buffer */
  struct ezt_list_token_t token;
};


static struct cuda_process_info_t *__register_process_hook(struct process_info_t *p_process)
{
  struct cuda_process_info_t *p_cuda = (struct cuda_process_info_t*) malloc(sizeof(struct cuda_process_info_t));
  p_cuda->p_process = p_process;
  ezt_hook_list_add(&p_cuda->p_process->hooks, p_cuda, (uint8_t)EZTRACE_CUDA_EVENTS_ID);
  p_cuda->start_time = 0; 

  ezt_list_new(&p_cuda->kernel_names);
  ezt_list_new(&p_cuda->device_list);

  return p_cuda;
}

#define  INIT_CUDA_PROCESS_INFO(p_process, var)				\
  struct cuda_process_info_t *var = (struct cuda_process_info_t*)	\
    ezt_hook_list_retrieve_data(&p_process->hooks, (uint8_t)EZTRACE_CUDA_EVENTS_ID); \
  if(!(var)) {								\
    var = __register_process_hook(p_process);				\
  }


#define CUDA_CHANGE() if(!recording_stats) CHANGE()

#define TOSTRING(x) #x
#define GENERATE_HANDLER(fname)						\
  void handle_cuda_##fname (int start)					\
  {									\
    FUNC_NAME;								\
    DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);		\
    if(start) {								\
      CUDA_CHANGE() pushState(CURRENT, "ST_Thread", thread_id, TOSTRING(cuda_##fname)); \
    }else{								\
      CUDA_CHANGE() popState(CURRENT, "ST_Thread", thread_id);		\
    }									\
  }

/* creates the GPU in the output GTG trace */
static void __create_device_container(struct cuda_device_info_t* p_cuda);

GENERATE_HANDLER(cudaThreadSynchronize)


/* a GPU timestamp is expressed in nanoseconds */
#define GPU_TIMESTAMP_TO_CPU(timestamp) ((timestamp)/1e6)



/* find the cuda_kernel_info_t that matches kernel_id
 * return NULL if not found
 */
struct cuda_kernel_stat_t* __cuda_find_kernel_stats(struct cuda_device_info_t *p_device,
						    struct cuda_kernel_name_t *p_kernel) {
  struct ezt_list_token_t *t = NULL;
  struct cuda_kernel_stat_t *p_stat;

  ezt_list_foreach(&p_device->kernel_stats, t){
    p_stat = (struct cuda_kernel_stat_t *) t->data;
    if(p_stat->p_kernel == p_kernel)
      return p_stat;
  }

  /* stat not found. Probably because it's the first time this kernel is called
   */
  p_stat = malloc(sizeof(struct cuda_kernel_stat_t));
  p_stat->p_kernel = p_kernel;
  p_stat->p_device = p_device;
  p_stat->nb_calls = 0;
  p_stat->min_duration = DBL_MAX;
  p_stat->max_duration = 0;
  p_stat->kernel_duration = 0;

  p_stat->token.data = p_stat;
  ezt_list_add(&p_device->kernel_stats, &p_stat->token);

  return p_stat;
}

/* find the cuda_kernel_info_t that matches kernel_id
 * return NULL if not found
 */
struct cuda_kernel_name_t* __cuda_find_kernel_name(struct ezt_list_t*list, kernel_id_t kernel_id) {
  struct ezt_list_token_t *t = NULL;
  struct cuda_kernel_name_t *p_name;

  ezt_list_foreach(list, t){
    p_name = (struct cuda_kernel_name_t *) t->data;
    if(p_name->id == kernel_id)
      return p_name;
  }
  return NULL;
}


/* create a struct cuda_kernel_info_t for kernel #kernel_id and add it to the list of pending kernels */
struct cuda_kernel_info_t* __cuda_new_kernel(unsigned kernel_id, struct cuda_device_info_t*p_device) {
  struct cuda_kernel_info_t* p_kernel = malloc(sizeof(struct cuda_kernel_info_t));
  p_kernel->kernel_id = kernel_id;
  p_kernel->gpu_start_time = -1;
  p_kernel->gpu_stop_time = -1;
  p_kernel->cpu_start_time = -1;
  p_kernel->cpu_stop_time = -1;

  p_kernel->name = NULL;
  snprintf(p_kernel->link_str, 80, "%s_cuda_%x", p_device->p_process->p_process->container->id, kernel_id);
  p_kernel->token.data = p_kernel;

  ezt_list_add(&p_device->pending_kernels, &p_kernel->token);

  return p_kernel;
}

/* find the cuda_kernel_info_t that matches kernel_id
 * return NULL if not found
 */
struct cuda_kernel_info_t* __cuda_find_kernel(struct ezt_list_t*list, unsigned kernel_id) {
  struct ezt_list_token_t *t = NULL;
  struct cuda_kernel_info_t *p_kernel;
  ezt_list_foreach(list, t){
    p_kernel = (struct cuda_kernel_info_t *) t->data;
    if(p_kernel->kernel_id == kernel_id)
      return p_kernel;
  }
  return NULL;
}

static struct cuda_device_info_t*
__get_device_info(struct cuda_process_info_t *p_cuda, uint32_t device_id) {
  struct ezt_list_token_t *t = NULL;
  struct cuda_device_info_t* p_device = NULL;

  ezt_list_foreach(&p_cuda->device_list, t){
    p_device = (struct cuda_device_info_t *) t->data;

    if(p_device->device_id == device_id)
      return p_device;
  }
  return NULL;

}

static void __ezt_cuda_init_memcpy_stats(struct cuda_memcpy_stat_t*p_stats) {
  p_stats->nb_memcpy = 0;
  p_stats->total_size = 0;
  p_stats->min_size = 0;
  p_stats->max_size = 0;
}

static void __ezt_cuda_init_malloc_stats(struct cuda_malloc_stat_t*p_malloc) {
  p_malloc->nb_malloc = 0;
  p_malloc->total_size = 0;
  p_malloc->peak_mem = 0;
  p_malloc->cur_mem = 0;
}

static struct cuda_device_info_t *
__new_device(struct cuda_process_info_t* p_process, uint32_t device_id) {

  struct cuda_device_info_t *p_device = malloc(sizeof(struct cuda_device_info_t));
  p_device->p_process = p_process;

  p_device->start_time = 0;
  p_device->device_id = device_id;

  if(p_process->start_time > 0)
    p_device->start_time = p_process->start_time;

  asprintf(&p_device->gpu_id, "%s_GPU_%d", p_process->p_process->container->id, device_id);
  asprintf(&p_device->gpu_name, "%s_GPU_%d", p_process->p_process->container->name, device_id);
  p_device->added = 0;
  p_device->link_start = -1;
  p_device->link_end = -1;
  /* add the hook in the thread info structure */

  ezt_list_new(&p_device->pending_kernels);
  ezt_list_new(&p_device->finished_kernels);

  ezt_list_new(&p_device->pending_memcpy);
  ezt_list_new(&p_device->finished_memcpy);

  ezt_list_new(&p_device->pending_malloc);
  ezt_list_new(&p_device->finished_malloc);

  ezt_list_new(&p_device->kernel_stats);

  __ezt_cuda_init_memcpy_stats(&p_device->memcpy_stats);
  __ezt_cuda_init_malloc_stats(&p_device->malloc_stats);

  p_device->token.data = p_device;
  ezt_list_add(&p_process->device_list, &p_device->token);

  return p_device;
}



/* create a struct cuda_memcpy_info_t and add it to the list of pending memcpy */
struct cuda_memcpy_info_t* __cuda_new_memcpy(enum ezt_cudaMemcpyKind kind,
					     ezt_cuda_size_t size,
					     struct cuda_device_info_t*p_device) {
  static unsigned memcpy_id = 0;

  struct cuda_memcpy_info_t* p_memcpy = malloc(sizeof(struct cuda_memcpy_info_t));
  p_memcpy->gpu_start_time = -1;
  p_memcpy->gpu_stop_time = -1;
  p_memcpy->cpu_start_time = -1;
  p_memcpy->cpu_stop_time = -1;

  p_memcpy->id = memcpy_id++;
  p_memcpy->size = size;
  p_memcpy->kind = kind;
  p_memcpy->src = NULL;
  p_memcpy->dest = NULL;

  snprintf(p_memcpy->link_str, 80, "%s_memcpy_%x", p_device->p_process->p_process->container->id, p_memcpy->id);
  p_memcpy->token.data = p_memcpy;

  ezt_list_add(&p_device->pending_memcpy, &p_memcpy->token);
  return p_memcpy;
}

/* find the cuda_memcpy_info_t
 * return NULL if not found
 */
struct cuda_memcpy_info_t* __cuda_find_memcpy(struct ezt_list_t*list,
					      enum ezt_cudaMemcpyKind kind,
					      ezt_cuda_size_t size) {
  struct ezt_list_token_t *t = NULL;
  struct cuda_memcpy_info_t *p_memcpy;
  ezt_list_foreach(list, t){
    p_memcpy = (struct cuda_memcpy_info_t *) t->data;
    if((p_memcpy->kind == kind) &&
       (p_memcpy->size == size))
      return p_memcpy;
  }
  return NULL;
}

/* create a struct cuda_malloc_info_t for buffer #devPtr_id and add it to the list of pending buffers */
struct cuda_malloc_info_t* __cuda_new_malloc(app_ptr devPtr, struct cuda_device_info_t*p_device) {
  struct cuda_malloc_info_t* p_malloc = malloc(sizeof(struct cuda_malloc_info_t));
  p_malloc->cpu_start_time = -1;
  p_malloc->cpu_stop_time = -1;
  p_malloc->token.data = p_malloc;

  p_malloc->devPtr = devPtr;
  ezt_list_add(&p_device->pending_malloc, &p_malloc->token);
  return p_malloc;
}

/* find the cuda_malloc_info_t that matches devPtr
 * return NULL if not found
 */
struct cuda_malloc_info_t* __cuda_find_malloc(struct ezt_list_t*list, app_ptr devPtr) {
  struct ezt_list_token_t *t = NULL;
  struct cuda_malloc_info_t *p_malloc;

  ezt_list_foreach(list, t){
    p_malloc = (struct cuda_malloc_info_t *) t->data;
    if(p_malloc->devPtr == devPtr)
      return p_malloc;
  }
  return NULL;
}


void handle_cuda_cuLaunchKernel (int start)				
{									
  FUNC_NAME;								
  DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);	

  uint32_t device_id;
  kernel_id_t kernel_id;
  GET_PARAM_PACKED_2(CUR_EV, device_id, kernel_id);

  struct cuda_device_info_t* p_device = __get_device_info(p_cuda, device_id);
  if(start) {		
    struct cuda_kernel_info_t *p_kernel = __cuda_new_kernel(kernel_id, p_device);
    p_kernel->cpu_start_time = CURRENT;

    CUDA_CHANGE() {
      __create_device_container(p_device);

      pushState(CURRENT, "ST_Thread", thread_id, "cuda_cuLaunchKernel");

      if(p_device->start_time) {
	startLink (CURRENT,
		   "hToD_kernel",
		   CUR_ID,
		   thread_id,
		   p_device->gpu_id,
		   p_kernel->link_str,
		   p_kernel->link_str);
      }
    }

  }else{								
    struct cuda_kernel_info_t *p_kernel = __cuda_find_kernel(&p_device->pending_kernels, kernel_id);
    p_kernel->cpu_stop_time = CURRENT;
    CUDA_CHANGE() popState(CURRENT, "ST_Thread", thread_id);		
  }									
}

static void __ezt_cuda_update_memcpy_stats(struct cuda_device_info_t *p_device, ezt_cuda_size_t size) {
  p_device->memcpy_stats.nb_memcpy ++;
  p_device->memcpy_stats.total_size += size;

  if(size < p_device->memcpy_stats.min_size || p_device->memcpy_stats.nb_memcpy == 1) {
    p_device->memcpy_stats.min_size = size;
  }

  if(size > p_device->memcpy_stats.max_size || p_device->memcpy_stats.nb_memcpy == 1) {
    p_device->memcpy_stats.max_size = size;
  }
}

void handle_cuda_cuMemcpy (int start)				
{									
  FUNC_NAME;								
  DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);	

  int device_id;
  ezt_cuda_size_t  size;
  enum ezt_cudaMemcpyKind kind;

  GET_PARAM_PACKED_3(CUR_EV, device_id, size, kind);

  /* TODO: for now the memcpy record is stored in the list of device #0
   * in order to support multi-GPU, we should store it in the srcDevice list (since it is where
   * handle_cuda_gpu_memcpy_timestamps() searches for it)
   */
  struct cuda_device_info_t* p_device = __get_device_info(p_cuda, device_id);

  /* todo: add an event that displays the size of the data transfer */

  if(start) {
    struct cuda_memcpy_info_t *p_memcpy = __cuda_new_memcpy(kind, size, p_device);
    p_memcpy->cpu_start_time = CURRENT;

    __create_device_container(p_device);
    CUDA_CHANGE() pushState(CURRENT, "ST_Thread", thread_id, TOSTRING(cuda_cuMemcpy));

    if(IS_FROM_CPU(p_memcpy->kind))
      p_memcpy->src = thread_id;
    else
      p_memcpy->src = p_device->gpu_id;

    if(IS_TO_CPU(p_memcpy->kind))
      p_memcpy->dest = thread_id;
    else
      p_memcpy->dest = p_device->gpu_id;

    CUDA_CHANGE() {
      if(p_device->start_time && (IS_FROM_CPU(p_memcpy->kind))) {
	startLink (CURRENT,
		   MEMCPY_TYPE_STR(p_memcpy->kind),
		   CUR_ID,
		   p_memcpy->src,
		   p_memcpy->dest,
		   p_memcpy->link_str,
		   p_memcpy->link_str);
      }
    }

  }else{	
    struct cuda_memcpy_info_t *p_memcpy = __cuda_find_memcpy(&p_device->pending_memcpy, kind, size);
    p_memcpy->cpu_stop_time = CURRENT;

    __ezt_cuda_update_memcpy_stats(p_device, size);

    CUDA_CHANGE() popState(CURRENT, "ST_Thread", thread_id);		
  }
}

void handle_cuda_cuMemAlloc (int start)
{
  FUNC_NAME;
  DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);

  int device_id;
  ezt_cuda_size_t size;
  app_ptr devPtr;

  if(start) {
    GET_PARAM_PACKED_2(CUR_EV, device_id, size);
  } else {
    GET_PARAM_PACKED_3(CUR_EV, device_id, size, devPtr);
  }

  struct cuda_device_info_t* p_device = __get_device_info(p_cuda, device_id);
  assert(p_device);

  if(start) {
    __create_device_container(p_device);

    /* create a new malloc_info for this buffer.
     * For now, we don't know the address of the buffer, so let's say NULL
     */
    struct cuda_malloc_info_t * p_malloc = __cuda_new_malloc((app_ptr)NULL, p_device);
    p_malloc->size = size;

    CUDA_CHANGE() pushState(CURRENT, "ST_Thread", thread_id, "cuda_cuMemAlloc");
    CUDA_CHANGE() addVar(CURRENT, "V_GPU_Mem", p_device->gpu_id, size);
  }else{

    /* Get the malloc_info that was created at the entry of cudaMalloc.
     * We can now assign devPtr.
     * WARNING: This is not reentrant, so we may have a problem if several threads call cudaMalloc simultaneously
     */
    struct cuda_malloc_info_t * p_malloc = __cuda_find_malloc(&p_device->pending_malloc, (app_ptr)NULL);

    assert(p_malloc->size == size);
    p_malloc->devPtr = devPtr;

    /* update statistics */
    p_device->malloc_stats.nb_malloc ++;

    p_device->malloc_stats.total_size += size;
    p_device->malloc_stats.cur_mem += size;

    if(p_device->malloc_stats.cur_mem > p_device->malloc_stats.peak_mem) {
      p_device->malloc_stats.peak_mem = p_device->malloc_stats.cur_mem;
    }

    CUDA_CHANGE() popState(CURRENT, "ST_Thread", thread_id);
  }
}


void handle_cuda_cudaFree(int start)
{
  FUNC_NAME;
  DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);

  uint32_t device_id;
  app_ptr devPtr;
  GET_PARAM_PACKED_2(CUR_EV, device_id, devPtr);

  struct cuda_device_info_t* p_device = __get_device_info(p_cuda, device_id);
  assert(p_device);

  /* Get the malloc_info that was created in cudaMalloc.
   * This permits to know the buffer size.
   */
  struct cuda_malloc_info_t * p_malloc = __cuda_find_malloc(&p_device->pending_malloc, devPtr);

  if(start) {
    __create_device_container(p_device);
    CUDA_CHANGE() pushState(CURRENT, "ST_Thread", thread_id, "cuda_cudaFree");
    CUDA_CHANGE() subVar(CURRENT, "V_GPU_Mem", p_device->gpu_id, p_malloc->size);

  }else{
    CUDA_CHANGE() popState(CURRENT, "ST_Thread", thread_id);

    ezt_list_remove(&p_malloc->token);
    ezt_list_add(&p_device->finished_malloc, &p_malloc->token);
  }
}


void handle_cuda_gpu_kernel_timestamps()
{
  FUNC_NAME;
  DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);

  unsigned kernel_type;
  float start_time;
  float stop_time;
  kernel_id_t kernel_id;
  uint32_t device_id;
  GET_PARAM_PACKED_5(CUR_EV, kernel_type, start_time, stop_time, kernel_id, device_id);

  struct cuda_device_info_t* p_device = __get_device_info(p_cuda, device_id);

  struct cuda_kernel_info_t *p_kernel = __cuda_find_kernel(&p_device->pending_kernels, kernel_type);
  p_kernel->name = __cuda_find_kernel_name(&p_cuda->kernel_names, kernel_id);

  p_kernel->gpu_start_time = p_device->start_time + GPU_TIMESTAMP_TO_CPU(start_time);
  p_kernel->gpu_stop_time = p_device->start_time + GPU_TIMESTAMP_TO_CPU(stop_time);

  CUDA_CHANGE() pushState(p_kernel->gpu_start_time, "ST_GPU", p_device->gpu_id, p_kernel->name->state_id);

  /* todo: use the kernel id to identify the link */
  CUDA_CHANGE() {
    if(p_device->start_time) {
      endLink(p_kernel->gpu_start_time,
	      "hToD_kernel",
	      CUR_ID,
	      thread_id,
	      p_device->gpu_id,
	      p_kernel->link_str,
	      p_kernel->link_str);
    }
    popState(p_kernel->gpu_stop_time, "ST_GPU", p_device->gpu_id);
  }

  /* update the kernel statistics */
  struct cuda_kernel_stat_t*p_stats = __cuda_find_kernel_stats(p_device, p_kernel->name);
  p_stats->nb_calls++;
  double kernel_runtime = p_kernel->gpu_stop_time-p_kernel->gpu_start_time;
  p_stats->kernel_duration += kernel_runtime;

  if(kernel_runtime < p_stats->min_duration) {
    p_stats->min_duration = kernel_runtime;
  }

  if(kernel_runtime > p_stats->max_duration) {
    p_stats->max_duration = kernel_runtime;
  }


  /* remove the kernel from the list of pending kernels and add it to the list of finished kernels */
  ezt_list_remove(&p_kernel->token);
  ezt_list_add(&p_device->finished_kernels, &p_kernel->token);
}

void handle_cuda_gpu_memcpy_timestamps()
{
  FUNC_NAME;
  DECLARE_THREAD_ID_STR(thread_id, CUR_INDEX, CUR_THREAD_ID);
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);

  float start_time;
  float stop_time;
  ezt_cuda_size_t size;
  enum ezt_cudaMemcpyKind kind;
  int srcDeviceId;
  int destDeviceId;

  GET_PARAM_PACKED_6(CUR_EV, start_time, stop_time, size, kind, srcDeviceId, destDeviceId);

  struct cuda_device_info_t* p_srcDevice = __get_device_info(p_cuda, srcDeviceId);
  assert(p_srcDevice);
  struct cuda_device_info_t* p_destDevice = __get_device_info(p_cuda, destDeviceId);
  assert(p_destDevice);

  struct cuda_memcpy_info_t *p_memcpy = __cuda_find_memcpy(&p_srcDevice->pending_memcpy, kind, size);
  
  assert(p_memcpy);
  p_memcpy->gpu_start_time = p_srcDevice->start_time + GPU_TIMESTAMP_TO_CPU(start_time);
  p_memcpy->gpu_stop_time  = p_srcDevice->start_time + GPU_TIMESTAMP_TO_CPU(stop_time);

  CUDA_CHANGE() {
    pushState(p_memcpy->gpu_start_time, "ST_GPU", p_srcDevice->gpu_id, "cuda_gpu_memcpy"); 

    if(p_srcDevice->start_time) {
      if(IS_FROM_GPU(p_memcpy->kind)) {
	startLink (p_memcpy->gpu_start_time,
		   MEMCPY_TYPE_STR(p_memcpy->kind),
		   CUR_ID,
		   p_memcpy->src,
		   p_memcpy->dest,
		   p_memcpy->link_str,
		   p_memcpy->link_str);
      }

      endLink(p_memcpy->gpu_stop_time,
	      MEMCPY_TYPE_STR(p_memcpy->kind),
	      CUR_ID,
	      p_memcpy->src,
	      p_memcpy->dest,
	      p_memcpy->link_str,
	      p_memcpy->link_str);
    }

    popState(p_memcpy->gpu_stop_time, "ST_GPU", p_srcDevice->gpu_id);
  }

  /* remove the memcpy from the list of pending memcpy and add it to the list of finished memcpy */
  ezt_list_remove(&p_memcpy->token);
  ezt_list_add(&p_srcDevice->finished_memcpy, &p_memcpy->token);
}

int eztrace_convert_cuda_init()
{
  if(get_mode() == EZTRACE_CONVERT) {
    addContType("CT_GPU", "0", "GPU");
    addStateType("ST_GPU", "CT_GPU", "GPU");
    addLinkType("hToD_kernel", "kernel", "CT_Process", "CT_Thread", "CT_GPU");
    addLinkType("HToD_memcpy", "memcpy", "CT_Process", "CT_Thread", "CT_GPU");
    addLinkType("DToH_memcpy", "memcpy", "CT_Process", "CT_GPU", "CT_Thread");
    addLinkType("HToH_memcpy", "memcpy", "CT_Process", "CT_Thread", "CT_Thread");
    addLinkType("DToD_memcpy", "memcpy", "CT_Process", "CT_GPU", "CT_GPU");
    addEntityValue("cuda_cuMemAlloc", "ST_Thread", "cuda_cuMemAlloc", GTG_YELLOW);
    addEntityValue("cuda_cuLaunchKernel", "ST_Thread", "cuda_cuLaunchKernel", GTG_GREEN);
    addEntityValue("cuda_cudaThreadSynchronize", "ST_Thread", "cuda_cudaThreadSynchronize", GTG_RED);
    addEntityValue("cuda_cuMemcpy", "ST_Thread", "cuda_cuMemcpy", GTG_BLACK);

    addEntityValue("cuda_gpu_kernel", "ST_GPU", "cuda_gpu_kernel", GTG_GREEN);
    addEntityValue("cuda_gpu_memcpy", "ST_GPU", "cuda_gpu_memcpy", GTG_BLACK);
    addEntityValue("cuda_gpu_idle", "ST_GPU", "cuda_gpu_idle", GTG_RED);

    addVarType ("V_GPU_Mem", "Memory used", "CT_GPU");
  } 
  return 0;
}

static void __create_device_container(struct cuda_device_info_t* p_device)
{
  CUDA_CHANGE() {
    if(! p_device->added) {
      addContainer(CURRENT, p_device->gpu_id, "CT_GPU", p_device->p_process->p_process->container->id, p_device->gpu_name, "0");
      pushState(CURRENT, "ST_GPU", p_device->gpu_id, "cuda_gpu_idle");

      setVar(CURRENT, "V_GPU_Mem", p_device->gpu_id, 0);

      p_device->added = 1;
    }
  }
}


void handle_cuda_register_kernel() {
  FUNC_NAME;
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);

  kernel_id_t kernel_id;
  GET_PARAM_PACKED_1(CUR_EV, kernel_id);

  wait_for_an_event(CUR_INDEX, EZTRACE_CUDA_KERNEL_NAME);

  struct cuda_kernel_name_t* p_kernel = malloc(sizeof(struct cuda_kernel_name_t));

  p_kernel->id = kernel_id;

  assert(LITL_READ_GET_TYPE(CUR_EV) == LITL_TYPE_RAW);
  const char* mangled_name = (char*)LITL_READ_RAW(CUR_EV)->data;
  const char* demangled_name = ezt_demangle(mangled_name);

  strncpy(p_kernel->name, demangled_name, 80);
  p_kernel->token.data = p_kernel;
  sprintf(p_kernel->state_id, "cuda_kernel_%x", p_kernel->id);
  sprintf(p_kernel->state_desc, "kernel %s", p_kernel->name);

  CUDA_CHANGE() addEntityValue(p_kernel->state_id, "ST_GPU", p_kernel->state_desc, GTG_GREEN);

  ezt_list_add(&p_cuda->kernel_names, &p_kernel->token);
}

void handle_cuda_cuInit() {
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);
  FUNC_NAME;

  p_cuda->start_time = (CURRENT);
}

void handle_cuda_setnbdevices() {
  FUNC_NAME;
  DECLARE_CUR_PROCESS(p_process);
  INIT_CUDA_PROCESS_INFO(p_process, p_cuda);

  int num_gpus;
  GET_PARAM_PACKED_1(CUR_EV, num_gpus);

  int i;
  for(i=0; i<num_gpus; i++) {
    struct cuda_device_info_t* p_device = __new_device(p_cuda, i);
    assert(p_device);
    p_device->start_time = (CURRENT);
    __create_device_container(p_device);
  }
}

/* return 1 if the event was handled */
int handle_cuda_events(eztrace_event_t *ev)
{
  if(!STARTED)
    return 0;
  switch (LITL_READ_GET_CODE(ev)) {
    case EZTRACE_CUDA_CUMEMALLOC_START:    handle_cuda_cuMemAlloc(1); break;
    case EZTRACE_CUDA_CUMEMALLOC_STOP:    handle_cuda_cuMemAlloc(0); break;
    case EZTRACE_CUDA_CUDAFREE_START: handle_cuda_cudaFree(1); break;
    case EZTRACE_CUDA_CUDAFREE_STOP: handle_cuda_cudaFree(0); break;
    case EZTRACE_CUDA_CULAUNCHKERNEL_START:   handle_cuda_cuLaunchKernel(1); break;
    case EZTRACE_CUDA_CULAUNCHKERNEL_STOP:    handle_cuda_cuLaunchKernel(0); break;
    case EZTRACE_CUDA_CUMEMCPY_START:   handle_cuda_cuMemcpy(1); break;
    case EZTRACE_CUDA_CUMEMCPY_STOP:    handle_cuda_cuMemcpy(0); break;
    case EZTRACE_CUDA_CUDATHREADSYNCHRONIZE_START:    handle_cuda_cudaThreadSynchronize(1); break;
    case EZTRACE_CUDA_CUDATHREADSYNCHRONIZE_STOP:    handle_cuda_cudaThreadSynchronize(0); break;

    case EZTRACE_CUDA_CUDADEVICESYNCHRONIZE_START:    handle_cuda_cudaThreadSynchronize(1); break;
    case EZTRACE_CUDA_CUDADEVICESYNCHRONIZE_STOP:    handle_cuda_cudaThreadSynchronize(0); break;

    case EZTRACE_CUDA_CUINIT:
      handle_cuda_cuInit();
      break;
    case EZTRACE_CUDA_SETNBDEVICES:
      handle_cuda_setnbdevices();
      break;

    case EZTRACE_CUDA_REGISTER_KERNEL:
      handle_cuda_register_kernel();
      break;

    case EZTRACE_CUDA_GPU_KERNEL_TIMESTAMPS:    handle_cuda_gpu_kernel_timestamps(); break;
    case EZTRACE_CUDA_GPU_MEMCPY_TIMESTAMPS:    handle_cuda_gpu_memcpy_timestamps(); break;
    default:
      /* this is not a cuda event */
      return 0;
    }
  return 1;
}

int handle_cuda_stats(eztrace_event_t *ev)
{
  recording_stats = 1;
  return handle_cuda_events(ev);
}

#define FORMAT_BYTES(nb_bytes)						\
  ((uint64_t)nb_bytes<1024?"B":						\
   ((uint64_t)nb_bytes<1024*1024?"KB":					\
    ((uint64_t)nb_bytes<1024*1024*1024?"MB":				\
     ((uint64_t)nb_bytes<(uint64_t)1024*1024*1024*1024?"GB":		\
      ((uint64_t)nb_bytes<(uint64_t)1024*1024*1024*1024*1024?"TB":	\
       "PB")))))

static float VALUE_BYTES(uint64_t nb_bytes) {
  int i;
  uint64_t div=1;
  for(i=0; i<6; i++) {
    if((nb_bytes/div)<1024)
      return (nb_bytes/(double)div);
    div*=1024;
  }
  return (float)nb_bytes;
}

void __ezt_print_device_stats(struct cuda_device_info_t *p_device) {

  /* 
     For each device:
        total allocated memory
  */

  printf("%s\n", p_device->gpu_name);
  struct cuda_kernel_stat_t *p_stat;
  struct ezt_list_token_t *t = NULL;

  /* print the list of executed kernels */
  ezt_list_foreach(&p_device->kernel_stats, t) {
    p_stat = (struct cuda_kernel_stat_t *) t->data;
    if(p_stat->nb_calls) {
      double avg_runtime = p_stat->kernel_duration/p_stat->nb_calls;
      printf("\tkernel %s was called %d times. Total run time: %lf ns, min: %lf ns, max: %lf ns, avg: %lf ns\n",
	     p_stat->p_kernel->name, p_stat->nb_calls, p_stat->kernel_duration,
	     p_stat->min_duration, p_stat->max_duration, avg_runtime);
    }
  }

  /* print the list of executed memcpy */
  struct cuda_memcpy_stat_t *p_memcpy = &p_device->memcpy_stats;
  if(p_memcpy->nb_memcpy>0) {
    printf("\tNumber of calls to cudaMemcpy: %d. Total size copied: %f %s (min: %f %s, max: %f %s, avg: %f %s)\n",
	   p_memcpy->nb_memcpy,
	   VALUE_BYTES(p_memcpy->total_size), FORMAT_BYTES(p_memcpy->total_size),
	   VALUE_BYTES(p_memcpy->min_size), FORMAT_BYTES(p_memcpy->min_size),
	   VALUE_BYTES(p_memcpy->max_size), FORMAT_BYTES(p_memcpy->max_size),
	   VALUE_BYTES(p_memcpy->total_size/p_memcpy->nb_memcpy),
	   FORMAT_BYTES(p_memcpy->total_size/p_memcpy->nb_memcpy));
  }

  /* print the memory information */
  struct cuda_malloc_stat_t *p_malloc = &p_device->malloc_stats;
  if(p_malloc->nb_malloc>0) {
    printf("\tNumber of calls to cudaMalloc: %d. Total size allocated: %f %s (peak: %f %s)\n",
	   p_malloc->nb_malloc,
	   VALUE_BYTES(p_malloc->total_size), FORMAT_BYTES(p_malloc->total_size),
	   VALUE_BYTES(p_malloc->peak_mem),   FORMAT_BYTES(p_malloc->peak_mem));
  }

}

void print_cuda_stats() {
  printf("\nCUDA:\n");
  printf("----------\n");

  /* browse the processes finished_section_list and extract statistics */
  int i;
  for (i = 0; i < NB_TRACES; i++) {
    struct process_info_t *p_process = GET_PROCESS_INFO(i);
    struct cuda_process_info_t *p_info =
      (struct cuda_process_info_t*) ezt_hook_list_retrieve_data(
          &p_process->hooks, (uint8_t) EZTRACE_CUDA_EVENTS_ID);

    if (!p_info)
      continue; /* No cuda process info attached, skip this process */

    struct ezt_list_token_t *token;
    ezt_list_foreach(&p_info->device_list, token) {
      struct cuda_device_info_t * p_device = (struct cuda_device_info_t *) token->data;
      assert(p_device);
      __ezt_print_device_stats(p_device);

    }
  }
}


struct eztrace_convert_module cuda_module;

void libinit(void) __attribute__ ((constructor));
void libinit(void)
{
  cuda_module.api_version = EZTRACE_API_VERSION;
  cuda_module.init = eztrace_convert_cuda_init;
  cuda_module.handle = handle_cuda_events;
  cuda_module.handle_stats = handle_cuda_stats;
  cuda_module.print_stats = print_cuda_stats;

  cuda_module.module_prefix = EZTRACE_CUDA_EVENTS_ID;
  asprintf(&cuda_module.name, "cuda");
  asprintf(&cuda_module.description, "Module for cuda functions (cuMemAlloc, cuMemcopy, etc.)");

  cuda_module.token.data = &cuda_module;
  eztrace_convert_register_module(&cuda_module);
  if(ezt_demangle_init() == 0) {
    fprintf(stderr, "Warning: C++ demangling is not available. Only the mangled names of CUDA kernels will be displayed.\n");
  }
}

void libfinalize(void) __attribute__ ((destructor));
void libfinalize(void)
{
}
