#define _REENTRANT

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

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

#include <eztrace-core/eztrace_config.h>
#include <eztrace-core/eztrace_spinlock.h>
#include <eztrace-lib/eztrace_otf2.h>
#include <eztrace-lib/eztrace.h>
#include <eztrace-lib/eztrace_module.h>
#include <eztrace-core/ezt_demangle.h>

#define CURRENT_MODULE cuda
DECLARE_CURRENT_MODULE;

#include "ezt_cuda.h"

static int _cuda_initialized = 0;

struct ezt_cuda_gpu {
  int gpu_id;
  uint64_t device_startTimestamp; // timestamp (in ns) of the first event
  uint64_t host_startTimestamp; // timestamp (in ns) of the first event
  int otf2_gpu_id;
  OTF2_DefWriter *def_writer;
  OTF2_EvtWriter *evt_writer;
};

static struct ezt_cuda_gpu *ezt_gpus=NULL;
static int ezt_nb_gpus = 0;

static int otf2_cuda_unknown_kernel_id = -1;
static int otf2_cuda_memcpy_id = -1;

static CUpti_SubscriberHandle subscriber;

uint64_t ezt_get_gputimestamp(int gpu_id, uint64_t timestamp) {
  return (timestamp - ezt_gpus[gpu_id].device_startTimestamp) +  ezt_gpus[gpu_id].host_startTimestamp;
}

// 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__) {			\
      eztrace_error("EZTrace: CUDA returned an error\n");	\
    }								\
  } 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) {
    eztrace_abort("Error: out of memory\n");
    exit(-1);
  }
  
  *size = _cuda_buf_size;
  *buffer = ALIGN_BUFFER(bfr, ALIGN_SIZE);
  *maxNumRecords = 0;
}

/* process a memcpy event that happened on a GPU */
static void _process_memcpy_record(CUpti_ActivityMemcpy2 *memcpy_info) {
  int srcDeviceId = memcpy_info->deviceId;
  int destDeviceId = memcpy_info->dstDeviceId;

  size_t cpy_size = memcpy_info->bytes;
  char* copy_type = MEMCPY_TYPE_STR((CUpti_ActivityMemcpyKind)memcpy_info->copyKind);

  static OTF2_AttributeRef attr_id_type = OTF2_UNDEFINED_ATTRIBUTE;
  if(attr_id_type == OTF2_UNDEFINED_ATTRIBUTE) { attr_id_type = ezt_otf2_register_attribute("copyType", OTF2_TYPE_AUTO(copy_type)); }
  static OTF2_AttributeRef attr_id_dst = OTF2_UNDEFINED_ATTRIBUTE;
  if(attr_id_dst == OTF2_UNDEFINED_ATTRIBUTE) { attr_id_dst = ezt_otf2_register_attribute("destDevice", OTF2_TYPE_AUTO(destDeviceId)); }
  static OTF2_AttributeRef attr_id_bytes = OTF2_UNDEFINED_ATTRIBUTE;
  if(attr_id_bytes == OTF2_UNDEFINED_ATTRIBUTE) { attr_id_bytes = ezt_otf2_register_attribute("bytes", OTF2_TYPE_AUTO(cpy_size));}

  OTF2_AttributeList* attribute_list = OTF2_AttributeList_New();
  ADD_ATTRIBUTE(attribute_list, attr_id_type, copy_type);
  ADD_ATTRIBUTE(attribute_list, attr_id_dst, destDeviceId);
  ADD_ATTRIBUTE(attribute_list, attr_id_bytes, cpy_size);

  OTF2_EvtWriter_Enter( ezt_gpus[srcDeviceId].evt_writer,
			attribute_list,
			ezt_get_gputimestamp(srcDeviceId, memcpy_info->start),
			otf2_cuda_memcpy_id);

  OTF2_EvtWriter_Leave( ezt_gpus[srcDeviceId].evt_writer,
			attribute_list,
			ezt_get_gputimestamp(srcDeviceId, memcpy_info->end),
			otf2_cuda_memcpy_id);
}


struct kernel_info_t {
  const char* name;
  const char* demangled_name;
  int otf2_region_id;
};
static int nb_kernels = 0; 
static int nb_allocated_kernels = 0;
static struct kernel_info_t *kernels = NULL;

/* register a kernel name and return its index in the kernels array */
static int _register_kernel(const char* kernel_name) {
  uint32_t i;
  static ezt_spinlock lock = ezt_spin_unlocked;

  ezt_spin_lock(&lock);  
  for(i=0; i<nb_kernels; i++) {
    if(!strcmp(kernels[i].name, kernel_name)) {
      // the kernel is already registered
      ezt_spin_unlock(&lock);
      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) {
      eztrace_error("[EZTrace] Cannot allocate memory. Aborting\n");
    }
    kernels = ptr;
  }

  int kernel_id = nb_kernels-1;
  // register the kernel
  kernels[kernel_id].name = strdup(kernel_name);
  kernels[kernel_id].demangled_name = ezt_demangle(kernel_name);
  kernels[kernel_id].otf2_region_id = ezt_otf2_register_function(kernels[kernel_id].demangled_name);

  eztrace_log(dbg_lvl_debug,"[CUDA] Register kernel(%d, %s) -> %d\n", kernel_id,
	      kernels[kernel_id].demangled_name,
	      kernels[kernel_id].otf2_region_id);

  ezt_spin_unlock(&lock);  

  return kernel_id;
}

/* Cupti keeps changing its API so we need to know which data structure to use */
#if CUPTI_API_VERSION < 7
typedef  CUpti_ActivityKernel2 __CUpti_ActivityKernel;
/* as of cupti 6.5, CUpti_ActivityKernel2 is deprecated, and CUpti_ActivityKernel3 should be used  */
#elif CUPTI_API_VERSION < 9
typedef CUpti_ActivityKernel3 __CUpti_ActivityKernel ;
#elif CUPTI_API_VERSION < 11
typedef CUpti_ActivityKernel4 __CUpti_ActivityKernel ;
#elif CUPTI_API_VERSION <= 14
typedef CUpti_ActivityKernel6 __CUpti_ActivityKernel ;
#elif CUPTI_API_VERSION <= 18
typedef CUpti_ActivityKernel7 __CUpti_ActivityKernel ;
#else
typedef CUpti_ActivityKernel9 __CUpti_ActivityKernel ;
#endif
  
/* process a kernel event that happened on the GPU */
static void _process_kernel_record(__CUpti_ActivityKernel *kernel) {
  uint32_t kernel_id = _register_kernel(kernel->name);
  uint32_t device_id = kernel->deviceId;
  
  OTF2_EvtWriter_Enter( ezt_gpus[device_id].evt_writer,
			NULL,
			ezt_get_gputimestamp(device_id, kernel->start),
			kernels[kernel_id].otf2_region_id);

  OTF2_EvtWriter_Leave( ezt_gpus[device_id].evt_writer,
			NULL,
			ezt_get_gputimestamp(device_id, kernel->end),
			kernels[kernel_id].otf2_region_id);
}


static ezt_spinlock _flush_lock = ezt_spin_unlocked;


void CUPTIAPI eztrace_cuda_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) {
    eztrace_cuda_runtime_callback(userdata, domain, cbid, cbdata);
  } else if(domain == CUPTI_CB_DOMAIN_DRIVER_API) {
    eztrace_cuda_driver_callback(userdata, domain, cbid, cbdata);
  }

}

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_spin_lock(&_flush_lock);
    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:
	  eztrace_warn("unknown record type: %d\n", record->kind);
	  break;
	}
      }
      else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED) {
	break;
      }
      else {
	CUDA_CHECK_RETVAL(status, CUPTI_SUCCESS);
      }
    } while (1);

    ezt_spin_unlock(&_flush_lock);

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

  free(buffer);
}



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

PPTRACE_START_INTERCEPT_FUNCTIONS(cuda)
PPTRACE_END_INTERCEPT_FUNCTIONS(cuda)

static void _ezt_init_gpu(int gpu_id) {
  ezt_gpus[gpu_id].gpu_id = gpu_id;
  CUDA_CHECK_RETVAL(cuptiGetTimestamp(&ezt_gpus[gpu_id].device_startTimestamp), CUPTI_SUCCESS);
  ezt_gpus[gpu_id].host_startTimestamp = ezt_get_timestamp();

  ezt_gpus[gpu_id].otf2_gpu_id = ezt_otf2_register_gpu(gpu_id);
  ezt_gpus[gpu_id].def_writer = OTF2_Archive_GetDefWriter( _ezt_trace.archive,
							   ezt_gpus[gpu_id].otf2_gpu_id );
  ezt_gpus[gpu_id].evt_writer = OTF2_Archive_GetEvtWriter( _ezt_trace.archive,
							   ezt_gpus[gpu_id].otf2_gpu_id );
}

static void _ezt_finalize_gpu(int gpu_id){
  if(eztrace_should_trace && ezt_gpus[gpu_id].def_writer) {
    static ezt_spinlock archive_lock;
    ezt_spin_lock(&archive_lock);
    OTF2_Archive_CloseEvtWriter( _ezt_trace.archive, ezt_gpus[gpu_id].evt_writer );
    OTF2_Archive_CloseDefWriter( _ezt_trace.archive, ezt_gpus[gpu_id].def_writer );
    ezt_spin_unlock(&archive_lock);
  }

}

static void ezt_cuda_init_otf2() {
  if(todo_get_status("ezt_cuda_init_otf2") == init_complete)
    return;

  todo_set_status("ezt_cuda_init_otf2", init_started);
  cudaGetDeviceCount(&ezt_nb_gpus);
  printf("ezt_cuda_init_otf2(): there are %d gpus\n", ezt_nb_gpus);
  ezt_gpus = malloc(sizeof(struct ezt_cuda_gpu) * ezt_nb_gpus);
  for(int i=0; i< ezt_nb_gpus; i++) {
    _ezt_init_gpu(i);
  }

  otf2_cuda_unknown_kernel_id = ezt_otf2_register_function("CUDA kernel");
  otf2_cuda_memcpy_id = ezt_otf2_register_function("CUDA Memcpy");

  
  cuptiSubscribe(&subscriber, 
		 (CUpti_CallbackFunc)eztrace_cuda_callback, NULL);
  cuptiEnableDomain(1, subscriber, CUPTI_CB_DOMAIN_RUNTIME_API);
  char* str = getenv("EZTRACE_CUDA_DRIVER");
  if(str) {
    eztrace_log(dbg_lvl_debug, "CUDA: record CUDA driver API events\n");
    cuptiEnableDomain(1, subscriber, CUPTI_CB_DOMAIN_DRIVER_API);
  }

  todo_set_status("ezt_cuda_init_otf2", init_complete);
}

static void _init_buffer_size() {
  char* str = getenv("EZTRACE_CUDA_BUFFER_SIZE");
  if(str) {
    _cuda_buf_size = atoi(str);
    eztrace_log(dbg_lvl_debug, "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) {

    eztrace_log(dbg_lvl_debug, "[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);

    ezt_spin_init(&_flush_lock);
  }

  enqueue_todo("ezt_cuda", ezt_cuda_init_otf2, "ezt_otf2", status_invalid);
}

// 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
    ezt_spin_lock(&_flush_lock);
    ezt_spin_unlock(&_flush_lock);
  }
}

static void  init_cuda() {   
  INSTRUMENT_FUNCTIONS(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);

  if (eztrace_autostart_enabled())
    eztrace_start();

  _cuda_initialized = 1;
}

static void finalize_cuda() {
  _cuda_initialized = 0;
  _ezt_cuda_flush(NULL);
  for(int i=0; i< ezt_nb_gpus; i++) {
    _ezt_finalize_gpu(i);
  }

  if(_cupti_enabled) {
    cuptiUnsubscribe(subscriber);
  }

  eztrace_stop();
}

static void _cuda_init(void) __attribute__((constructor));
static void _cuda_init(void) {
  eztrace_log(dbg_lvl_debug, "eztrace_cuda constructor starts\n");
  EZT_REGISTER_MODULE(cuda, "Module for CUDA",
		      init_cuda, finalize_cuda);
  eztrace_log(dbg_lvl_debug, "eztrace_cuda constructor ends\n");
}
