gpu: Actually fix PE to GPU mapping
[charm.git] / src / arch / cuda / hybridAPI / hapi_impl.cpp
1 #include "hapi.h"
2 #include "hapi_impl.h"
3 #include "converse.h"
4 #include "ckcallback.h"
5 #include <stdio.h>
6 #include <stdlib.h>
7 #include <algorithm>
8 #include <queue>
9 #include <atomic>
10 #include "cuda.h"
11
12 #ifdef HAPI_NVTX_PROFILE
13 #include "hapi_nvtx.h"
14 #endif
15
16 #if defined HAPI_MEMPOOL || defined HAPI_INSTRUMENT_WRS
17 #include "cklists.h"
18 #endif
19
20 #if defined HAPI_TRACE || defined HAPI_INSTRUMENT_WRS
21 extern "C" double CmiWallTimer();
22 #endif
23
24 #ifdef HAPI_TRACE
25 #define QUEUE_SIZE_INIT 128
26 extern "C" int traceRegisterUserEvent(const char* x, int e);
27 extern "C" void traceUserBracketEvent(int e, double beginT, double endT);
28
29 typedef struct gpuEventTimer {
30   int stage;
31   double cmi_start_time;
32   double cmi_end_time;
33   int event_type;
34   const char* trace_name;
35 } gpuEventTimer;
36 #endif
37
38 #ifdef HAPI_INSTRUMENT_WRS
39 static bool initializedInstrument();
40 #endif
41
42 // Event stages used for profiling.
43 enum WorkRequestStage{
44   DataSetup        = 1,
45   KernelExecution  = 2,
46   DataCleanup      = 3
47 };
48
49 enum ProfilingStage{
50   GpuMemSetup   = 8800,
51   GpuKernelExec = 8801,
52   GpuMemCleanup = 8802
53 };
54
55 #ifndef HAPI_CUDA_CALLBACK
56 typedef struct hapiEvent {
57   cudaEvent_t event;
58   void* cb;
59   void* cb_msg;
60   hapiWorkRequest* wr; // if this is not NULL, buffers and request itself are deallocated
61
62   hapiEvent(cudaEvent_t event_, void* cb_, void* cb_msg_, hapiWorkRequest* wr_ = NULL)
63             : event(event_), cb(cb_), cb_msg(cb_msg_), wr(wr_) {}
64 } hapiEvent;
65
66 CpvDeclare(std::queue<hapiEvent>, hapi_event_queue);
67 #endif
68 CpvDeclare(int, n_hapi_events);
69
70 void initEventQueues() {
71 #ifndef HAPI_CUDA_CALLBACK
72   CpvInitialize(std::queue<hapiEvent>, hapi_event_queue);
73 #endif
74   CpvInitialize(int, n_hapi_events);
75   CpvAccess(n_hapi_events) = 0;
76 }
77
78 // Returns the CUDA device associated with the given PE.
79 // TODO: should be updated to exploit the hardware topology instead of round robin
80 static inline int getMyCudaDevice(int my_node) {
81   int device_count;
82   hapiCheck(cudaGetDeviceCount(&device_count));
83   return my_node % device_count;
84 }
85
86 // A function in ck.C which casts the void* to a CkCallback object and invokes
87 // the Charm++ callback.
88 extern void CUDACallbackManager(void* fn);
89 extern int CmiMyPe();
90
91 // Functions used to support quiescence detection.
92 extern void QdCreate(int n);
93 extern void QdProcess(int n);
94
95 // Initial size of the user-addressed portion of host/device buffer arrays;
96 // the system-addressed portion of host/device buffer arrays (used when there
97 // is no need to share buffers between work requests) will be equivalant in size.
98 // FIXME hard-coded maximum
99 #if CMK_SMP
100 #define NUM_BUFFERS 4096
101 #else
102 #define NUM_BUFFERS 256
103 #endif
104
105 #define MAX_PINNED_REQ 64
106 #define MAX_DELAYED_FREE_REQS 64
107
108 // Contains data and methods needed by HAPI.
109 class GPUManager {
110
111 public:
112 #ifdef HAPI_MEMPOOL
113 // Update for new row, again this shouldn't be hard coded!
114 #define HAPI_MEMPOOL_NUM_SLOTS 20
115 // Pre-allocated buffers will be at least this big (in bytes).
116 #define HAPI_MEMPOOL_MIN_BUFFER_SIZE 256
117 // Scale the amount of memory each node pins.
118 #define HAPI_MEMPOOL_SCALE 1.0
119
120   CkVec<BufferPool> mempool_free_bufs_;
121   CkVec<size_t> mempool_boundaries_;
122 #endif // HAPI_MEMPOOL
123
124   // The runtime system keeps track of all allocated buffers on the GPU.
125   // The following arrays contain pointers to host (CPU) data and the
126   // corresponding data on the device (GPU).
127   void **host_buffers_;
128   void **device_buffers_;
129
130   // Used to assign buffer IDs automatically by the system if the user
131   // specifies an invalid buffer ID.
132   int next_buffer_;
133
134   cudaStream_t *streams_;
135   int n_streams_;
136   int last_stream_id_;
137
138 #ifdef HAPI_CUDA_CALLBACK
139   int host_to_device_cb_idx_;
140   int kernel_cb_idx_;
141   int device_to_host_cb_idx_;
142   int light_cb_idx_; // for lightweight version
143 #endif
144
145   int running_kernel_idx_;
146   int data_setup_idx_;
147   int data_cleanup_idx_;
148
149 #ifdef HAPI_TRACE
150   gpuEventTimer gpu_events_[QUEUE_SIZE_INIT * 3];
151   std::atomic<int> time_idx_;
152 #endif
153
154 #ifdef HAPI_INSTRUMENT_WRS
155   CkVec<CkVec<CkVec<hapiRequestTimeInfo> > > avg_times_;
156   bool init_instr_;
157 #endif
158
159 #if CMK_SMP || CMK_MULTICORE
160   CmiNodeLock buffer_lock_;
161   CmiNodeLock queue_lock_;
162   CmiNodeLock progress_lock_;
163   CmiNodeLock stream_lock_;
164 #endif
165
166   cudaDeviceProp device_prop_;
167 #ifdef HAPI_CUDA_CALLBACK
168   bool cb_support;
169 #endif
170
171   void init();
172   int createStreams();
173   void destroyStreams();
174   cudaStream_t getNextStream();
175   cudaStream_t getStream(int);
176   void allocateBuffers(hapiWorkRequest*);
177   void hostToDeviceTransfer(hapiWorkRequest*);
178   void deviceToHostTransfer(hapiWorkRequest*);
179   void freeBuffers(hapiWorkRequest*);
180   void runKernel(hapiWorkRequest*);
181 };
182
183 // Declare GPU Manager as a process-shared object.
184 CsvDeclare(GPUManager, gpu_manager);
185
186 void GPUManager::init() {
187   next_buffer_ = NUM_BUFFERS;
188   streams_ = NULL;
189   last_stream_id_ = -1;
190   running_kernel_idx_ = 0;
191   data_setup_idx_ = 0;
192   data_cleanup_idx_ = 0;
193
194 #if CMK_SMP || CMK_MULTICORE
195   // create mutex locks
196   buffer_lock_ = CmiCreateLock();
197   queue_lock_ = CmiCreateLock();
198   progress_lock_ = CmiCreateLock();
199   stream_lock_ = CmiCreateLock();
200 #endif
201
202 #ifdef HAPI_TRACE
203   time_idx_ = 0;
204 #endif
205
206   // store CUDA device properties
207   hapiCheck(cudaGetDeviceProperties(&device_prop_, getMyCudaDevice(CmiMyNode())));
208
209 #ifdef HAPI_CUDA_CALLBACK
210   // check if CUDA callback is supported
211   // CUDA 5.0 (compute capability 3.0) or newer
212   cb_support = (device_prop_.major >= 3);
213   if (!cb_support) {
214     CmiAbort("[HAPI] CUDA callback is not supported on this device");
215   }
216 #endif
217
218   // set which device to use
219   hapiCheck(cudaSetDevice(getMyCudaDevice(CmiMyNode())));
220
221   // allocate host/device buffers array (both user and system-addressed)
222   host_buffers_ = new void*[NUM_BUFFERS*2];
223   device_buffers_ = new void*[NUM_BUFFERS*2];
224
225   // initialize device array to NULL
226   for (int i = 0; i < NUM_BUFFERS*2; i++) {
227     device_buffers_[i] = NULL;
228   }
229
230 #ifdef HAPI_TRACE
231   traceRegisterUserEvent("GPU Memory Setup", GpuMemSetup);
232   traceRegisterUserEvent("GPU Kernel Execution", GpuKernelExec);
233   traceRegisterUserEvent("GPU Memory Cleanup", GpuMemCleanup);
234 #endif
235
236 #ifdef HAPI_MEMPOOL
237   mempool_boundaries_.reserve(HAPI_MEMPOOL_NUM_SLOTS);
238   mempool_boundaries_.length() = HAPI_MEMPOOL_NUM_SLOTS;
239
240   size_t buf_size = HAPI_MEMPOOL_MIN_BUFFER_SIZE;
241   for(int i = 0; i < HAPI_MEMPOOL_NUM_SLOTS; i++){
242     mempool_boundaries_[i] = buf_size;
243     buf_size = buf_size << 1;
244   }
245 #endif // HAPI_MEMPOOL
246
247 #ifdef HAPI_INSTRUMENT_WRS
248   init_instr_ = false;
249 #endif
250 }
251
252 // Creates streams equal to the maximum number of concurrent kernels,
253 // which depends on the compute capability of the device.
254 // Returns the number of created streams.
255 int GPUManager::createStreams() {
256   if (streams_)
257     return n_streams_;
258
259 #if CMK_SMP || CMK_MULTICORE
260   if (device_prop_.major == 3) {
261     if (device_prop_.minor == 0)
262       n_streams_ = 16;
263     else if (device_prop_.minor == 2)
264       n_streams_ = 4;
265     else // 3.5, 3.7 or unknown 3.x
266       n_streams_ = 32;
267   }
268   else if (device_prop_.major == 5) {
269     if (device_prop_.minor == 3)
270       n_streams_ = 16;
271     else // 5.0, 5.2 or unknown 5.x
272       n_streams_ = 32;
273   }
274   else if (device_prop_.major == 6) {
275     if (device_prop_.minor == 1)
276       n_streams_ = 32;
277     else if (device_prop_.minor == 2)
278       n_streams_ = 16;
279     else // 6.0 or unknown 6.x
280       n_streams_ = 128;
281   }
282   else // unknown (future) compute capability
283     n_streams_ = 128;
284 #else
285   n_streams_ = 4; // FIXME per PE in non-SMP mode
286 #endif
287
288   streams_ = new cudaStream_t[n_streams_];
289   for (int i = 0; i < n_streams_; i++) {
290     hapiCheck(cudaStreamCreate(&streams_[i]));
291   }
292
293   return n_streams_;
294 }
295
296 void GPUManager::destroyStreams() {
297   if (streams_) {
298     for (int i = 0; i < n_streams_; i++) {
299       hapiCheck(cudaStreamDestroy(streams_[i]));
300     }
301   }
302 }
303
304 cudaStream_t GPUManager::getNextStream() {
305   if (streams_ == NULL)
306     return NULL;
307
308   last_stream_id_ = (++last_stream_id_) % n_streams_;
309   return streams_[last_stream_id_];
310 }
311
312 cudaStream_t GPUManager::getStream(int i) {
313   if (streams_ == NULL)
314     return NULL;
315
316   if (i < 0 || i >= n_streams_)
317     CmiAbort("[HAPI] invalid stream ID");
318   return streams_[i];
319 }
320
321 // Allocates device buffers.
322 void GPUManager::allocateBuffers(hapiWorkRequest* wr) {
323   for (int i = 0; i < wr->getBufferCount(); i++) {
324     hapiBufferInfo& bi = wr->buffers[i];
325     int index = bi.id;
326     size_t size = bi.size;
327
328     // if index value is invalid, use an available ID
329     if (index < 0 || index >= NUM_BUFFERS) {
330       bool is_found = false;
331       for (int j = next_buffer_; j < NUM_BUFFERS*2; j++) {
332         if (device_buffers_[j] == NULL) {
333           index = j;
334           is_found = true;
335           break;
336         }
337       }
338
339       // if no index was found, try to search for a value at the
340       // beginning of the system addressed space
341       if (!is_found) {
342         for (int j = NUM_BUFFERS; j < next_buffer_; j++) {
343           if (device_buffers_[j] == NULL) {
344             index = j;
345             is_found = true;
346             break;
347           }
348         }
349       }
350
351       if (!is_found) {
352         CmiAbort("[HAPI] ran out of device buffer indices");
353       }
354
355       next_buffer_ = index + 1;
356       if (next_buffer_ == NUM_BUFFERS*2) {
357         next_buffer_ = NUM_BUFFERS;
358       }
359
360       bi.id = index;
361     }
362
363     if (device_buffers_[index] == NULL) {
364       // allocate device memory
365       hapiCheck(cudaMalloc((void **)&device_buffers_[index], size));
366
367 #ifdef HAPI_DEBUG
368       printf("[HAPI] allocated buffer %d at %p, time: %.2f, size: %zu\n",
369              index, device_buffers_[index], cutGetTimerValue(timerHandle),
370              size);
371 #endif
372     }
373   }
374 }
375
376 #ifndef HAPI_CUDA_CALLBACK
377 void recordEvent(cudaStream_t stream, void* cb, void* cb_msg, hapiWorkRequest* wr = NULL) {
378   // create CUDA event and insert into stream
379   cudaEvent_t ev;
380   cudaEventCreateWithFlags(&ev, cudaEventDisableTiming);
381   cudaEventRecord(ev, stream);
382
383   hapiEvent hev(ev, cb, cb_msg, wr);
384
385   // push event information in queue
386   CpvAccess(hapi_event_queue).push(hev);
387
388   // increase count so that scheduler can poll the queue
389   CpvAccess(n_hapi_events)++;
390 }
391 #endif
392
393 // Initiates host-to-device data transfer.
394 void GPUManager::hostToDeviceTransfer(hapiWorkRequest* wr) {
395   for (int i = 0; i < wr->getBufferCount(); i++) {
396     hapiBufferInfo& bi = wr->buffers[i];
397     int index = bi.id;
398     size_t size = bi.size;
399     host_buffers_[index] = bi.host_buffer;
400
401     if (bi.transfer_to_device) {
402       // TODO should be changed back to async for performance
403       hapiCheck(cudaMemcpy(device_buffers_[index], host_buffers_[index], size,
404                                 cudaMemcpyHostToDevice));
405
406 #ifdef HAPI_DEBUG
407       printf("[HAPI] transferring buffer %d from host to device, time: %.2f, "
408              "size: %zu\n", index, cutGetTimerValue(timerHandle), size);
409 #endif
410     }
411   }
412 }
413
414 // Initiates device-to-host data transfer.
415 void GPUManager::deviceToHostTransfer(hapiWorkRequest* wr) {
416   for (int i = 0; i < wr->getBufferCount(); i++) {
417     hapiBufferInfo& bi = wr->buffers[i];
418     int index = bi.id;
419     size_t size = bi.size;
420
421     if (bi.transfer_to_host) {
422       // TODO should be changed back to async for performance
423       hapiCheck(cudaMemcpy(host_buffers_[index], device_buffers_[index], size,
424                                 cudaMemcpyDeviceToHost));
425
426 #ifdef HAPI_DEBUG
427       printf("[HAPI] transferring buffer %d from device to host, time %.2f, "
428              "size: %zu\n", index, cutGetTimerValue(timerHandle), size);
429 #endif
430     }
431   }
432 }
433
434 // Frees device buffers.
435 void GPUManager::freeBuffers(hapiWorkRequest* wr) {
436   for (int i = 0; i < wr->getBufferCount(); i++) {
437     hapiBufferInfo& bi = wr->buffers[i];
438     int index = bi.id;
439
440     if (bi.need_free) {
441       hapiCheck(cudaFree(device_buffers_[index]));
442       device_buffers_[index] = NULL;
443
444 #ifdef HAPI_DEBUG
445       printf("[HAPI] freed buffer %d, time %.2f\n",
446              index, cutGetTimerValue(timerHandle));
447 #endif
448     }
449   }
450 }
451
452 inline static void hapiWorkRequestCleanup(hapiWorkRequest* wr) {
453 #if CMK_SMP || CMK_MULTICORE
454   CmiLock(CsvAccess(gpu_manager).progress_lock_);
455 #endif
456
457   // free device buffers
458   CsvAccess(gpu_manager).freeBuffers(wr);
459
460 #if CMK_SMP || CMK_MULTICORE
461   CmiUnlock(CsvAccess(gpu_manager).progress_lock_);
462 #endif
463
464   // free hapiWorkRequest
465   delete wr;
466 }
467
468 // Run the user's kernel for the given work request.
469 // This used to be a switch statement defined by the user to allow the runtime
470 // to execute the correct kernel.
471 void GPUManager::runKernel(hapiWorkRequest* wr) {
472         if (wr->runKernel) {
473                 wr->runKernel(wr, wr->stream, device_buffers_);
474         }
475         // else, might be only for data transfer (or might be a bug?)
476 }
477
478 #ifdef HAPI_CUDA_CALLBACK
479 // Invokes user's host-to-device callback.
480 static void* hostToDeviceCallback(void* arg) {
481 #ifdef HAPI_NVTX_PROFILE
482   NVTXTracer nvtx_range("hostToDeviceCallback", NVTXColor::Asbestos);
483 #endif
484   hapiWorkRequest* wr = *((hapiWorkRequest**)((char*)arg + CmiMsgHeaderSizeBytes + sizeof(int)));
485   CUDACallbackManager(wr->host_to_device_cb);
486
487   // inform QD that the host-to-device transfer is complete
488   QdProcess(1);
489
490   return NULL;
491 }
492
493 // Invokes user's kernel execution callback.
494 static void* kernelCallback(void* arg) {
495 #ifdef HAPI_NVTX_PROFILE
496   NVTXTracer nvtx_range("kernelCallback", NVTXColor::Asbestos);
497 #endif
498   hapiWorkRequest* wr = *((hapiWorkRequest**)((char*)arg + CmiMsgHeaderSizeBytes + sizeof(int)));
499   CUDACallbackManager(wr->kernel_cb);
500
501   // inform QD that the kernel is complete
502   QdProcess(1);
503
504   return NULL;
505 }
506
507 // Frees device buffers and invokes user's device-to-host callback.
508 // Invoked regardless of the availability of the user's callback.
509 static void* deviceToHostCallback(void* arg) {
510 #ifdef HAPI_NVTX_PROFILE
511   NVTXTracer nvtx_range("deviceToHostCallback", NVTXColor::Asbestos);
512 #endif
513   hapiWorkRequest* wr = *((hapiWorkRequest**)((char*)arg + CmiMsgHeaderSizeBytes + sizeof(int)));
514
515   // invoke user callback
516   if (wr->device_to_host_cb) {
517     CUDACallbackManager(wr->device_to_host_cb);
518   }
519
520   hapiWorkRequestCleanup(wr);
521
522   // inform QD that device-to-host transfer is complete
523   QdProcess(1);
524
525   return NULL;
526 }
527
528 // Used by lightweight HAPI.
529 static void* lightCallback(void *arg) {
530 #ifdef HAPI_NVTX_PROFILE
531   NVTXTracer nvtx_range("lightCallback", NVTXColor::Asbestos);
532 #endif
533
534   char* conv_msg_tmp = (char*)arg + CmiMsgHeaderSizeBytes + sizeof(int);
535   void* cb = *((void**)conv_msg_tmp);
536
537   // invoke user callback
538   if (cb != NULL) {
539     CUDACallbackManager(cb);
540   }
541
542   // notify process to QD
543   QdProcess(1);
544
545   return NULL;
546 }
547 #endif // HAPI_CUDA_CALLBACK
548
549 // Register callback functions. All PEs need to call this.
550 void hapiRegisterCallbacks() {
551 #ifdef HAPI_CUDA_CALLBACK
552   // FIXME: Potential race condition on assignments, but CmiAssignOnce
553   // causes a hang at startup.
554   CsvAccess(gpu_manager).host_to_device_cb_idx_
555     = CmiRegisterHandler((CmiHandler)hostToDeviceCallback);
556   CsvAccess(gpu_manager).kernel_cb_idx_
557     = CmiRegisterHandler((CmiHandler)kernelCallback);
558   CsvAccess(gpu_manager).device_to_host_cb_idx_
559     = CmiRegisterHandler((CmiHandler)deviceToHostCallback);
560   CsvAccess(gpu_manager).light_cb_idx_
561     = CmiRegisterHandler((CmiHandler)lightCallback);
562 #endif
563 }
564
565 #ifdef HAPI_CUDA_CALLBACK
566 // Callback function invoked by the CUDA runtime certain parts of GPU work are
567 // complete. It sends a converse message to the original PE to free the relevant
568 // device memory and invoke the user's callback. The reason for this method is
569 // that a thread created by the CUDA runtime does not have access to any of the
570 // CpvDeclare'd variables as it is not one of the threads created by the Charm++
571 // runtime.
572 static void CUDART_CB CUDACallback(cudaStream_t stream, cudaError_t status,
573                                    void *data) {
574 #ifdef HAPI_NVTX_PROFILE
575   NVTXTracer nvtx_range("CUDACallback", NVTXColor::Silver);
576 #endif
577
578   if (status == cudaSuccess) {
579     // send message to the original PE
580     char *conv_msg = (char*)data;
581     int dstRank = *((int *)(conv_msg + CmiMsgHeaderSizeBytes));
582     CmiPushPE(dstRank, conv_msg);
583   }
584   else {
585     CmiAbort("[HAPI] error before CUDACallback");
586   }
587 }
588
589 enum CallbackStage {
590   AfterHostToDevice,
591   AfterKernel,
592   AfterDeviceToHost
593 };
594
595 static void addCallback(hapiWorkRequest *wr, CallbackStage stage) {
596   // create converse message to be delivered to this PE after CUDA callback
597   char *conv_msg = (char *)CmiAlloc(CmiMsgHeaderSizeBytes + sizeof(int) +
598                                   sizeof(hapiWorkRequest *)); // FIXME memory leak?
599   *((int *)(conv_msg + CmiMsgHeaderSizeBytes)) = CmiMyRank();
600   *((hapiWorkRequest **)(conv_msg + CmiMsgHeaderSizeBytes + sizeof(int))) = wr;
601
602   int handlerIdx;
603   switch (stage) {
604     case AfterHostToDevice:
605       handlerIdx = CsvAccess(gpu_manager).host_to_device_cb_idx_;
606       break;
607     case AfterKernel:
608       handlerIdx = CsvAccess(gpu_manager).kernel_cb_idx_;
609       break;
610     case AfterDeviceToHost:
611       handlerIdx = CsvAccess(gpu_manager).device_to_host_cb_idx_;
612       break;
613     default: // wrong type
614       CmiFree(conv_msg);
615       return;
616   }
617   CmiSetHandler(conv_msg, handlerIdx);
618
619   // add callback into CUDA stream
620   hapiCheck(cudaStreamAddCallback(wr->stream, CUDACallback, (void*)conv_msg, 0));
621 }
622 #endif // HAPI_CUDA_CALLBACK
623
624 /******************** DEPRECATED ********************/
625 // User calls this function to offload work to the GPU.
626 void hapiEnqueue(hapiWorkRequest* wr) {
627 #ifdef HAPI_NVTX_PROFILE
628   NVTXTracer nvtx_range("enqueue", NVTXColor::Pomegranate);
629 #endif
630
631 #if CMK_SMP || CMK_MULTICORE
632   CmiLock(CsvAccess(gpu_manager).progress_lock_);
633 #endif
634
635   // allocate device memory
636   CsvAccess(gpu_manager).allocateBuffers(wr);
637
638   // transfer data to device
639   CsvAccess(gpu_manager).hostToDeviceTransfer(wr);
640
641   // add host-to-device transfer callback
642   if (wr->host_to_device_cb) {
643     // while there is an ongoing workrequest, quiescence should not be detected
644     // even if all PEs seem idle
645     QdCreate(1);
646
647 #ifdef HAPI_CUDA_CALLBACK
648     addCallback(wr, AfterHostToDevice);
649 #else
650     recordEvent(wr->stream, wr->host_to_device_cb, NULL);
651 #endif
652   }
653
654   // run kernel
655   CsvAccess(gpu_manager).runKernel(wr);
656
657   // add kernel callback
658   if (wr->kernel_cb) {
659     QdCreate(1);
660
661 #ifdef HAPI_CUDA_CALLBACK
662     addCallback(wr, AfterKernel);
663 #else
664     recordEvent(wr->stream, wr->kernel_cb, NULL);
665 #endif
666   }
667
668   // transfer data to host
669   CsvAccess(gpu_manager).deviceToHostTransfer(wr);
670
671   // add device-to-host transfer callback
672   QdCreate(1);
673 #ifdef HAPI_CUDA_CALLBACK
674   // always invoked to free memory
675   addCallback(wr, AfterDeviceToHost);
676 #else
677   if (wr->device_to_host_cb) {
678     recordEvent(wr->stream, wr->device_to_host_cb, NULL, wr);
679   }
680   else {
681     recordEvent(wr->stream, NULL, NULL, wr);
682   }
683 #endif
684
685 #if CMK_SMP || CMK_MULTICORE
686   CmiUnlock(CsvAccess(gpu_manager).progress_lock_);
687 #endif
688 }
689
690 /******************** DEPRECATED ********************/
691 // Creates a hapiWorkRequest object on the heap and returns it to the user.
692 hapiWorkRequest* hapiCreateWorkRequest() {
693   return (new hapiWorkRequest);
694 }
695
696 #ifdef HAPI_MEMPOOL
697 static void createPool(int *nbuffers, int n_slots, CkVec<BufferPool> &pools);
698 static void releasePool(CkVec<BufferPool> &pools);
699 #endif
700
701 // Initialization of HAPI functionalities.
702 void initHybridAPI() {
703   // create and initialize GPU Manager object
704   CsvInitialize(GPUManager, gpu_manager);
705   CsvAccess(gpu_manager).init();
706
707 #ifdef HAPI_MEMPOOL
708   // create pool of page-locked memory
709   int sizes[HAPI_MEMPOOL_NUM_SLOTS];
710         /*256*/ sizes[0]  =  4;
711         /*512*/ sizes[1]  =  2;
712        /*1024*/ sizes[2]  =  2;
713        /*2048*/ sizes[3]  =  4;
714        /*4096*/ sizes[4]  =  2;
715        /*8192*/ sizes[5]  =  6;
716       /*16384*/ sizes[6]  =  5;
717       /*32768*/ sizes[7]  =  2;
718       /*65536*/ sizes[8]  =  1;
719      /*131072*/ sizes[9]  =  1;
720      /*262144*/ sizes[10] =  1;
721      /*524288*/ sizes[11] =  1;
722     /*1048576*/ sizes[12] =  1;
723     /*2097152*/ sizes[13] =  2;
724     /*4194304*/ sizes[14] =  2;
725     /*8388608*/ sizes[15] =  2;
726    /*16777216*/ sizes[16] =  2;
727    /*33554432*/ sizes[17] =  1;
728    /*67108864*/ sizes[18] =  1;
729   /*134217728*/ sizes[19] =  7;
730   createPool(sizes, HAPI_MEMPOOL_NUM_SLOTS, CsvAccess(gpu_manager).mempool_free_bufs_);
731
732 #ifdef HAPI_MEMPOOL_DEBUG
733   printf("[HAPI (%d)] done creating buffer pool\n", CmiMyPe());
734 #endif
735
736 #endif // HAPI_MEMPOOL
737 }
738
739 // Set HAPI device for non-0 ranks
740 void setHybridAPIDevice() {
741   // set which device to use
742   hapiCheck(cudaSetDevice(getMyCudaDevice(CmiMyNode())));
743 }
744
745 // Clean up and delete memory used by HAPI.
746 void exitHybridAPI() {
747 #if CMK_SMP || CMK_MULTICORE
748   // destroy mutex locks
749   CmiDestroyLock(CsvAccess(gpu_manager).buffer_lock_);
750   CmiDestroyLock(CsvAccess(gpu_manager).queue_lock_);
751   CmiDestroyLock(CsvAccess(gpu_manager).progress_lock_);
752   CmiDestroyLock(CsvAccess(gpu_manager).stream_lock_);
753 #endif
754
755   // destroy streams (if they were created)
756   CsvAccess(gpu_manager).destroyStreams();
757
758 #ifdef HAPI_MEMPOOL
759   // release memory pool
760   releasePool(CsvAccess(gpu_manager).mempool_free_bufs_);
761 #endif // HAPI_MEMPOOL
762
763 #ifdef HAPI_TRACE
764   for (int i = 0; i < CsvAccess(gpu_manager).time_idx_; i++) {
765     switch (CsvAccess(gpu_manager).gpu_events_[i].event_type) {
766     case DataSetup:
767       printf("[HAPI] kernel %s data setup\n",
768              CsvAccess(gpu_manager).gpu_events_[i].trace_name);
769       break;
770     case DataCleanup:
771       printf("[HAPI] kernel %s data cleanup\n",
772              CsvAccess(gpu_manager).gpu_events_[i].trace_name);
773       break;
774     case KernelExecution:
775       printf("[HAPI] kernel %s execution\n",
776              CsvAccess(gpu_manager).gpu_events_[i].trace_name);
777       break;
778     default:
779       printf("[HAPI] invalid timer identifier\n");
780     }
781     printf("[HAPI] %.2f:%.2f\n",
782            CsvAccess(gpu_manager).gpu_events_[i].cmi_start_time -
783            CsvAccess(gpu_manager).gpu_events_[0].cmi_start_time,
784            CsvAccess(gpu_manager).gpu_events_[i].cmi_end_time -
785            CsvAccess(gpu_manager).gpu_events_[0].cmi_start_time);
786   }
787 #endif
788 }
789
790 /******************** DEPRECATED ********************/
791 // Need to be updated with the Tracing API.
792 static inline void gpuEventStart(hapiWorkRequest* wr, int* index,
793                                  WorkRequestStage event, ProfilingStage stage) {
794 #ifdef HAPI_TRACE
795   gpuEventTimer* shared_gpu_events_ = CsvAccess(gpu_manager).gpu_events_;
796   int shared_time_idx_ = CsvAccess(gpu_manager).time_idx_++;
797   shared_gpu_events_[shared_time_idx_].cmi_start_time = CmiWallTimer();
798   shared_gpu_events_[shared_time_idx_].event_type = event;
799   shared_gpu_events_[shared_time_idx_].trace_name = wr->trace_name;
800   *index = shared_time_idx_;
801   shared_gpu_events_[shared_time_idx_].stage = stage;
802 #ifdef HAPI_DEBUG
803   printf("[HAPI] start event %d of WR %s, profiling stage %d\n",
804          event, wr->trace_name, stage);
805 #endif
806 #endif // HAPI_TRACE
807 }
808
809 /******************** DEPRECATED ********************/
810 // Need to be updated with the Tracing API.
811 static inline void gpuEventEnd(int index) {
812 #ifdef HAPI_TRACE
813   CsvAccess(gpu_manager).gpu_events_[index].cmi_end_time = CmiWallTimer();
814   traceUserBracketEvent(CsvAccess(gpu_manager).gpu_events_[index].stage,
815                         CsvAccess(gpu_manager).gpu_events_[index].cmi_start_time,
816                         CsvAccess(gpu_manager).gpu_events_[index].cmi_end_time);
817 #ifdef HAPI_DEBUG
818   printf("[HAPI] end event %d of WR %s, profiling stage %d\n",
819           CsvAccess(gpu_manager).gpu_events_[index].event_type,
820           CsvAccess(gpu_manager).gpu_events_[index].trace_name,
821           CsvAccess(gpu_manager).gpu_events_[index].stage);
822 #endif
823 #endif // HAPI_TRACE
824 }
825
826 static inline void hapiWorkRequestStartTime(hapiWorkRequest* wr) {
827 #ifdef HAPI_INSTRUMENT_WRS
828   wr->phase_start_time = CmiWallTimer();
829 #endif
830 }
831
832 static inline void profileWorkRequestEvent(hapiWorkRequest* wr,
833                                            WorkRequestStage event) {
834 #ifdef HAPI_INSTRUMENT_WRS
835   if (initializedInstrument()) {
836     double tt = CmiWallTimer() - (wr->phase_start_time);
837     int index = wr->chare_index;
838     char type = wr->comp_type;
839     char phase = wr->comp_phase;
840
841     CkVec<hapiRequestTimeInfo> &vec = wr->avg_times_[index][type];
842     if (vec.length() <= phase){
843       vec.growAtLeast(phase);
844       vec.length() = phase+1;
845     }
846     switch (event) {
847       case DataSetup:
848         vec[phase].transfer_time += tt;
849         break;
850       case KernelExecution:
851         vec[phase].kernel_time += tt;
852         break;
853       case DataCleanup:
854         vec[phase].cleanup_time += tt;
855         vec[phase].n++;
856         break;
857       default:
858         printf("[HAPI] invalid event during profileWorkRequestEvent\n");
859     }
860   }
861 #endif
862 }
863
864 #ifdef HAPI_MEMPOOL
865 // Create a pool with n_slots slots.
866 // There are n_buffers[i] buffers for each buffer size corresponding to entry i.
867 // TODO list the alignment/fragmentation issues with either of two allocation schemes:
868 // if single, large buffer is allocated for each subpool
869 // if multiple, smaller buffers are allocated for each subpool
870 static void createPool(int *n_buffers, int n_slots, CkVec<BufferPool> &pools){
871   CkVec<size_t>& mempool_boundaries = CsvAccess(gpu_manager).mempool_boundaries_;
872
873   // initialize pools
874   pools.reserve(n_slots);
875   pools.length() = n_slots;
876   for (int i = 0; i < n_slots; i++) {
877     pools[i].size = mempool_boundaries[i];
878     pools[i].head = NULL;
879   }
880
881   // divide by # of PEs on physical node and multiply by # of PEs in logical node
882   size_t available_memory = CsvAccess(gpu_manager).device_prop_.totalGlobalMem /
883                            CmiNumPesOnPhysicalNode(CmiPhysicalNodeID(CmiMyPe()))
884                            * CmiMyNodeSize() * HAPI_MEMPOOL_SCALE;
885
886   // pre-calculate memory per size
887   int max_buffers = *std::max_element(n_buffers, n_buffers + n_slots);
888   int n_buffers_to_allocate[n_slots];
889   memset(n_buffers_to_allocate, 0, sizeof(n_buffers_to_allocate));
890   size_t buf_size;
891   while (available_memory >= mempool_boundaries[0] + sizeof(Header)) {
892     for (int i = 0; i < max_buffers; i++) {
893       for (int j = n_slots - 1; j >= 0; j--) {
894         buf_size = mempool_boundaries[j] + sizeof(Header);
895         if (i < n_buffers[j] && buf_size <= available_memory) {
896           n_buffers_to_allocate[j]++;
897           available_memory -= buf_size;
898         }
899       }
900     }
901   }
902
903   // pin the host memory
904   for (int i = 0; i < n_slots; i++) {
905     buf_size = mempool_boundaries[i] + sizeof(Header);
906     int num_buffers = n_buffers_to_allocate[i];
907
908     Header* hd;
909     Header* previous = NULL;
910
911     // pin host memory in a contiguous block for a slot
912     void* pinned_chunk;
913     hapiCheck(cudaMallocHost(&pinned_chunk, buf_size * num_buffers));
914
915     // initialize header structs
916     for (int j = num_buffers - 1; j >= 0; j--) {
917       hd = reinterpret_cast<Header*>(reinterpret_cast<unsigned char*>(pinned_chunk)
918                                      + buf_size * j);
919       hd->slot = i;
920       hd->next = previous;
921       previous = hd;
922     }
923
924     pools[i].head = previous;
925 #ifdef HAPI_MEMPOOL_DEBUG
926     pools[i].num = num_buffers;
927 #endif
928   }
929 }
930
931 static void releasePool(CkVec<BufferPool> &pools){
932   for (int i = 0; i < pools.length(); i++) {
933     Header* hdr = pools[i].head;
934     if (hdr != NULL) {
935       hapiCheck(cudaFreeHost((void*)hdr));
936     }
937   }
938   pools.free();
939 }
940
941 static int findPool(size_t size){
942   int boundary_array_len = CsvAccess(gpu_manager).mempool_boundaries_.length();
943   if (size <= CsvAccess(gpu_manager).mempool_boundaries_[0]) {
944     return 0;
945   }
946   else if (size > CsvAccess(gpu_manager).mempool_boundaries_[boundary_array_len-1]) {
947     // create new slot
948     CsvAccess(gpu_manager).mempool_boundaries_.push_back(size);
949
950     BufferPool newpool;
951     hapiCheck(cudaMallocHost((void**)&newpool.head, size + sizeof(Header)));
952     if (newpool.head == NULL) {
953       printf("[HAPI (%d)] findPool: failed to allocate newpool %d head, size %zu\n",
954              CmiMyPe(), boundary_array_len, size);
955       CmiAbort("[HAPI] failed newpool allocation");
956     }
957     newpool.size = size;
958 #ifdef HAPI_MEMPOOL_DEBUG
959     newpool.num = 1;
960 #endif
961     CsvAccess(gpu_manager).mempool_free_bufs_.push_back(newpool);
962
963     Header* hd = newpool.head;
964     hd->next = NULL;
965     hd->slot = boundary_array_len;
966
967     return boundary_array_len;
968   }
969   for (int i = 0; i < CsvAccess(gpu_manager).mempool_boundaries_.length()-1; i++) {
970     if (CsvAccess(gpu_manager).mempool_boundaries_[i] < size &&
971         size <= CsvAccess(gpu_manager).mempool_boundaries_[i+1]) {
972       return (i + 1);
973     }
974   }
975   return -1;
976 }
977
978 static void* getBufferFromPool(int pool, size_t size){
979   Header* ret;
980
981   if (pool < 0 || pool >= CsvAccess(gpu_manager).mempool_free_bufs_.length()) {
982     printf("[HAPI (%d)] getBufferFromPool, pool: %d, size: %zu invalid pool\n",
983            CmiMyPe(), pool, size);
984 #ifdef HAPI_MEMPOOL_DEBUG
985     printf("[HAPI (%d)] num: %d\n", CmiMyPe(),
986            CsvAccess(gpu_manager).mempool_free_bufs_[pool].num);
987 #endif
988     CmiAbort("[HAPI] exiting after invalid pool");
989   }
990   else if (CsvAccess(gpu_manager).mempool_free_bufs_[pool].head == NULL) {
991     Header* hd;
992     hapiCheck(cudaMallocHost((void**)&hd, sizeof(Header) +
993                              CsvAccess(gpu_manager).mempool_free_bufs_[pool].size));
994 #ifdef HAPI_MEMPOOL_DEBUG
995     printf("[HAPI (%d)] getBufferFromPool, pool: %d, size: %zu expand by 1\n",
996            CmiMyPe(), pool, size);
997 #endif
998     if (hd == NULL) {
999       CmiAbort("[HAPI] exiting after NULL hd from pool");
1000     }
1001     hd->slot = pool;
1002     return (void*)(hd + 1);
1003   }
1004   else {
1005     ret = CsvAccess(gpu_manager).mempool_free_bufs_[pool].head;
1006     CsvAccess(gpu_manager).mempool_free_bufs_[pool].head = ret->next;
1007 #ifdef HAPI_MEMPOOL_DEBUG
1008     ret->size = size;
1009     CsvAccess(gpu_manager).mempool_free_bufs_[pool].num--;
1010 #endif
1011     return (void*)(ret + 1);
1012   }
1013   return NULL;
1014 }
1015
1016 static void returnBufferToPool(int pool, Header* hd) {
1017   hd->next = CsvAccess(gpu_manager).mempool_free_bufs_[pool].head;
1018   CsvAccess(gpu_manager).mempool_free_bufs_[pool].head = hd;
1019 #ifdef HAPI_MEMPOOL_DEBUG
1020   CsvAccess(gpu_manager).mempool_free_bufs_[pool].num++;
1021 #endif
1022 }
1023
1024 void* hapiPoolMalloc(size_t size) {
1025 #if CMK_SMP || CMK_MULTICORE
1026   CmiLock(CsvAccess(gpu_manager).buffer_lock_);
1027 #endif
1028
1029   int pool = findPool(size);
1030   void* buf = getBufferFromPool(pool, size);
1031
1032 #ifdef HAPI_MEMPOOL_DEBUG
1033   printf("[HAPI (%d)] hapiPoolMalloc size %zu pool %d left %d\n",
1034          CmiMyPe(), size, pool,
1035          CsvAccess(gpu_manager).mempool_free_bufs_[pool].num);
1036 #endif
1037
1038 #if CMK_SMP || CMK_MULTICORE
1039   CmiUnlock(CsvAccess(gpu_manager).buffer_lock_);
1040 #endif
1041
1042   return buf;
1043 }
1044
1045 void hapiPoolFree(void* ptr) {
1046   Header* hd = ((Header*)ptr) - 1;
1047   int pool = hd->slot;
1048
1049 #ifdef HAPI_MEMPOOL_DEBUG
1050   size_t size = hd->size;
1051 #endif
1052
1053 #if CMK_SMP || CMK_MULTICORE
1054   CmiLock(CsvAccess(gpu_manager).buffer_lock_);
1055 #endif
1056
1057   returnBufferToPool(pool, hd);
1058
1059 #if CMK_SMP || CMK_MULTICORE
1060   CmiUnlock(CsvAccess(gpu_manager).buffer_lock_);
1061 #endif
1062
1063 #ifdef HAPI_MEMPOOL_DEBUG
1064   printf("[HAPI (%d)] hapiPoolFree size %zu pool %d left %d\n",
1065          CmiMyPe(), size, pool,
1066          CsvAccess(gpu_manager).mempool_free_bufs_[pool].num);
1067 #endif
1068 }
1069 #endif // HAPI_MEMPOOL
1070
1071 #ifdef HAPI_INSTRUMENT_WRS
1072 void hapiInitInstrument(int n_chares, char n_types) {
1073   avg_times_.reserve(n_chares);
1074   avg_times_.length() = n_chares;
1075   for (int i = 0; i < n_chares; i++) {
1076     avg_times_[i].reserve(n_types);
1077     avg_times_[i].length() = n_types;
1078   }
1079   init_instr_ = true;
1080 }
1081
1082 static bool initializedInstrument() {
1083   return init_instr_;
1084 }
1085
1086 hapiRequestTimeInfo* hapiQueryInstrument(int chare, char type, char phase) {
1087   if (phase < avg_times_[chare][type].length()) {
1088     return &avg_times_[chare][type][phase];
1089   }
1090   else {
1091     return NULL;
1092   }
1093 }
1094
1095 void hapiClearInstrument() {
1096   for (int chare = 0; chare < avg_times_.length(); chare++) {
1097     for (int type = 0; type < avg_times_[chare].length(); type++) {
1098       for (int phase = 0; phase < avg_times_[chare][type].length(); phase++) {
1099         avg_times_[chare][type][phase].transferTime = 0.0;
1100         avg_times_[chare][type][phase].kernelTime = 0.0;
1101         avg_times_[chare][type][phase].cleanupTime = 0.0;
1102         avg_times_[chare][type][phase].n = 0;
1103       }
1104       avg_times_[chare][type].length() = 0;
1105     }
1106     avg_times_[chare].length() = 0;
1107   }
1108   avg_times_.length() = 0;
1109   init_instr_ = false;
1110 }
1111 #endif // HAPI_INSTRUMENT_WRS
1112
1113 // Poll HAPI events stored in the PE's queue. Current strategy is to process
1114 // all successive completed events in the queue starting from the front.
1115 // TODO Maybe we should make one pass of all events in the queue instead,
1116 // since there might be completed events later in the queue.
1117 void hapiPollEvents() {
1118 #ifndef HAPI_CUDA_CALLBACK
1119   std::queue<hapiEvent>& queue = CpvAccess(hapi_event_queue);
1120   while (!queue.empty()) {
1121     hapiEvent hev = queue.front();
1122     if (cudaEventQuery(hev.event) == cudaSuccess) {
1123       // invoke Charm++ callback if one was given
1124       if (hev.cb) {
1125         ((CkCallback*)hev.cb)->send(hev.cb_msg);
1126       }
1127
1128       // clean up hapiWorkRequest
1129       if (hev.wr) {
1130         hapiWorkRequestCleanup(hev.wr);
1131       }
1132       cudaEventDestroy(hev.event);
1133       queue.pop();
1134       CpvAccess(n_hapi_events)--;
1135
1136       // inform QD that an event was processed
1137       QdProcess(1);
1138     }
1139     else {
1140       // stop going through the queue once we encounter a non-successful event
1141       break;
1142     }
1143   }
1144 #endif
1145 }
1146
1147 int hapiCreateStreams() {
1148 #if CMK_SMP || CMK_MULTICORE
1149   CmiLock(CsvAccess(gpu_manager).stream_lock_);
1150 #endif
1151
1152   int ret = CsvAccess(gpu_manager).createStreams();
1153
1154 #if CMK_SMP || CMK_MULTICORE
1155   CmiUnlock(CsvAccess(gpu_manager).stream_lock_);
1156 #endif
1157
1158   return ret;
1159 }
1160
1161 cudaStream_t hapiGetStream() {
1162 #if CMK_SMP || CMK_MULTICORE
1163   CmiLock(CsvAccess(gpu_manager).stream_lock_);
1164 #endif
1165
1166   cudaStream_t ret = CsvAccess(gpu_manager).getNextStream();
1167
1168 #if CMK_SMP || CMK_MULTICORE
1169   CmiUnlock(CsvAccess(gpu_manager).stream_lock_);
1170 #endif
1171
1172   return ret;
1173 }
1174
1175 // Lightweight HAPI, to be invoked after data transfer or kernel execution.
1176 void hapiAddCallback(cudaStream_t stream, void* cb, void* cb_msg) {
1177 #ifndef HAPI_CUDA_CALLBACK
1178   // record CUDA event
1179   recordEvent(stream, cb, cb_msg);
1180 #else
1181   /* FIXME works for now (faster too), but CmiAlloc might not be thread-safe
1182 #if CMK_SMP || CMK_MULTICORE
1183   CmiLock(CsvAccess(gpu_manager).queue_lock_);
1184 #endif
1185 */
1186
1187   // create converse message to be delivered to this PE after CUDA callback
1188   char* conv_msg = (char*)CmiAlloc(CmiMsgHeaderSizeBytes + sizeof(int) +
1189                                  sizeof(void*)); // FIXME memory leak?
1190   char* conv_msg_tmp = conv_msg + CmiMsgHeaderSizeBytes;
1191   *((int*)conv_msg_tmp) = CmiMyRank();
1192   conv_msg_tmp += sizeof(int);
1193   *((void**)conv_msg_tmp) = cb;
1194   CmiSetHandler(conv_msg, CsvAccess(gpu_manager).light_cb_idx_);
1195
1196   // push into CUDA stream
1197   hapiCheck(cudaStreamAddCallback(stream, CUDACallback, (void*)conv_msg, 0));
1198
1199   /*
1200 #if CMK_SMP || CMK_MULTICORE
1201   CmiUnlock(CsvAccess(gpu_manager).queue_lock_);
1202 #endif
1203 */
1204 #endif
1205
1206   // while there is an ongoing workrequest, quiescence should not be detected
1207   // even if all PEs seem idle
1208   QdCreate(1);
1209 }
1210
1211 cudaError_t hapiMalloc(void** devPtr, size_t size) {
1212   return cudaMalloc(devPtr, size);
1213 }
1214
1215 cudaError_t hapiFree(void* devPtr) {
1216   return cudaFree(devPtr);
1217 }
1218
1219 cudaError_t hapiMallocHost(void** ptr, size_t size) {
1220   return cudaMallocHost(ptr, size);
1221 }
1222
1223 cudaError_t hapiFreeHost(void* ptr) {
1224   return cudaFreeHost(ptr);
1225 }
1226
1227 cudaError_t hapiMemcpyAsync(void* dst, const void* src, size_t count, cudaMemcpyKind kind, cudaStream_t stream = 0) {
1228   return cudaMemcpyAsync(dst, src, count, kind, stream);
1229 }
1230
1231 void hapiErrorDie(cudaError_t retCode, const char* code, const char* file, int line) {
1232   if (retCode != cudaSuccess) {
1233     fprintf(stderr, "Fatal CUDA Error [%d] %s at %s:%d\n", retCode, cudaGetErrorString(retCode), file, line);
1234     CmiAbort("Exit due to CUDA error");
1235   }
1236 }