From c8a587666a23e045f25dc871c3257364a5f6a7d5 Mon Sep 17 00:00:00 2001
From: David Medina <dmed256@gmail.com>
Date: Wed, 9 Sep 2020 23:35:45 -0400
Subject: [PATCH] [CUDA] Sets the CUDA context in more methods (#400)

---
 include/occa/modes/cuda/device.hpp |  2 ++
 src/modes/cuda/device.cpp          | 30 ++++++++++++++++++++----------
 src/modes/cuda/kernel.cpp          |  6 +++++-
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/include/occa/modes/cuda/device.hpp b/include/occa/modes/cuda/device.hpp
index bd7a929dd..9832a979f 100644
--- a/include/occa/modes/cuda/device.hpp
+++ b/include/occa/modes/cuda/device.hpp
@@ -42,6 +42,8 @@ namespace occa {
 
       void* getNullPtr();
 
+      void setCudaContext();
+
       //---[ Stream ]-------------------
       virtual modeStream_t* createStream(const occa::properties &props);
 
diff --git a/src/modes/cuda/device.cpp b/src/modes/cuda/device.cpp
index fbc79aa27..a8f06eddb 100644
--- a/src/modes/cuda/device.cpp
+++ b/src/modes/cuda/device.cpp
@@ -126,12 +126,17 @@ namespace occa {
       return (void*) &(nullPtr->cuPtr);
     }
 
+    void device::setCudaContext() {
+      OCCA_CUDA_ERROR("Device: Setting Context",
+                      cuCtxSetCurrent(cuContext));
+    }
+
     //---[ Stream ]---------------------
     modeStream_t* device::createStream(const occa::properties &props) {
       CUstream cuStream = NULL;
 
-      OCCA_CUDA_ERROR("Device: Setting Context",
-                      cuCtxSetCurrent(cuContext));
+      setCudaContext();
+
       OCCA_CUDA_ERROR("Device: createStream",
                       cuStreamCreate(&cuStream, CU_STREAM_DEFAULT));
 
@@ -141,8 +146,8 @@ namespace occa {
     occa::streamTag device::tagStream() {
       CUevent cuEvent = NULL;
 
-      OCCA_CUDA_ERROR("Device: Setting Context",
-                      cuCtxSetCurrent(cuContext));
+      setCudaContext();
+
       OCCA_CUDA_ERROR("Device: Tagging Stream (Creating Tag)",
                       cuEventCreate(&cuEvent,
                                     CU_EVENT_DEFAULT));
@@ -219,6 +224,8 @@ namespace occa {
       CUfunction cuFunction;
       CUresult error;
 
+      setCudaContext();
+
       error = cuModuleLoad(&cuModule, binaryFilename.c_str());
       if (error) {
         lock.release();
@@ -353,6 +360,8 @@ namespace occa {
       CUmodule cuModule;
       CUresult error;
 
+      setCudaContext();
+
       error = cuModuleLoad(&cuModule, binaryFilename.c_str());
       if (error) {
         lock.release();
@@ -410,6 +419,8 @@ namespace occa {
       CUmodule cuModule = NULL;
       CUfunction cuFunction = NULL;
 
+      setCudaContext();
+
       OCCA_CUDA_ERROR("Kernel [" + kernelName + "]: Loading Module",
                       cuModuleLoad(&cuModule, filename.c_str()));
 
@@ -438,8 +449,7 @@ namespace occa {
 
       cuda::memory &mem = *(new cuda::memory(this, bytes, props));
 
-      OCCA_CUDA_ERROR("Device: Setting Context",
-                      cuCtxSetCurrent(cuContext));
+      setCudaContext();
 
       OCCA_CUDA_ERROR("Device: malloc",
                       cuMemAlloc(&(mem.cuPtr), bytes));
@@ -456,8 +466,8 @@ namespace occa {
 
       cuda::memory &mem = *(new cuda::memory(this, bytes, props));
 
-      OCCA_CUDA_ERROR("Device: Setting Context",
-                      cuCtxSetCurrent(cuContext));
+      setCudaContext();
+
       OCCA_CUDA_ERROR("Device: malloc host",
                       cuMemAllocHost((void**) &(mem.mappedPtr), bytes));
       OCCA_CUDA_ERROR("Device: get device pointer from host",
@@ -481,8 +491,8 @@ namespace occa {
       const unsigned int flags = (props.get("attached_host", false) ?
                                   CU_MEM_ATTACH_HOST : CU_MEM_ATTACH_GLOBAL);
 
-      OCCA_CUDA_ERROR("Device: Setting Context",
-                      cuCtxSetCurrent(cuContext));
+      setCudaContext();
+
       OCCA_CUDA_ERROR("Device: Unified alloc",
                       cuMemAllocManaged(&(mem.cuPtr),
                                         bytes,
diff --git a/src/modes/cuda/kernel.cpp b/src/modes/cuda/kernel.cpp
index a65698de7..f63d90d2c 100644
--- a/src/modes/cuda/kernel.cpp
+++ b/src/modes/cuda/kernel.cpp
@@ -62,6 +62,8 @@ namespace occa {
     }
 
     void kernel::deviceRun() const {
+      device *devicePtr = (device*) modeDevice;
+
       const int args = (int) arguments.size();
       if (!args) {
         vArgs.resize(1);
@@ -74,10 +76,12 @@ namespace occa {
         vArgs[i] = arguments[i].ptr();
         // Set a proper NULL pointer
         if (!vArgs[i]) {
-          vArgs[i] = ((device*) modeDevice)->getNullPtr();
+          vArgs[i] = devicePtr->getNullPtr();
         }
       }
 
+      devicePtr->setCudaContext();
+
       OCCA_CUDA_ERROR("Launching Kernel",
                       cuLaunchKernel(cuFunction,
                                      outerDims.x, outerDims.y, outerDims.z,