Core: fixed handle wrapping for CUcontext

include/isaac/driver/backend.h

@@ -48,6 +48,7 @@ RESTORE_MSVC_WARNING_C4251
       static void release();
   public:
       static Context const & get_default();
+      static Context const & import(CUcontext context);
       static Context const & import(cl_context context);
       static void get(std::list<Context const *> &);
   private:

include/isaac/driver/buffer.h

@@ -6,7 +6,7 @@
 #include "isaac/driver/common.h"
 #include "isaac/driver/context.h"
 #include "isaac/driver/handle.h"
-
+#include "isaac/driver/dispatch.h"
 namespace isaac
 {
 
@@ -18,6 +18,14 @@ class ISAACAPI Buffer
 {
   friend class CommandQueue;
   friend class Kernel;
+
+  static CUcontext context(CUdeviceptr h)
+  {
+      CUcontext res;
+      cuda::check(dispatch::cuPointerGetAttribute((void*)&res, CU_POINTER_ATTRIBUTE_CONTEXT, h));
+      return res;
+  }
+
 public:
   Buffer(CUdeviceptr h = 0, bool take_ownership = true);
   Buffer(cl_mem Buffer = 0, bool take_ownership = true);

include/isaac/driver/context.h

@@ -21,6 +21,7 @@ class ISAACAPI Context
   friend class Buffer;
 
 public:
+  explicit Context(CUcontext const & context, CUdevice const & device, bool take_ownership = true);
   explicit Context(cl_context const & context, bool take_ownership = true);
   explicit Context(Device const & device);
 

include/isaac/driver/device.h

@@ -47,8 +47,9 @@ private:
   int cuGetInfo() const;
 
 public:
-  Device(int ordinal);
-  Device(cl_device_id const & device, bool take_ownership = true);
+//  Device(int ordinal);
+  explicit Device(CUdevice const & device, bool take_ownership = true);
+  explicit Device(cl_device_id const & device, bool take_ownership = true);
 
   bool operator==(Device const &) const;
   bool operator<(Device const &) const;

include/isaac/driver/dispatch.h

@@ -113,6 +113,7 @@ public:
     static CUresult cuStreamDestroy_v2(CUstream hStream);
     static CUresult cuEventDestroy_v2(CUevent hEvent);
     static CUresult cuMemAlloc_v2(CUdeviceptr *dptr, size_t bytesize);
+    static CUresult cuPointerGetAttribute(void * data, CUpointer_attribute attribute, CUdeviceptr ptr);
 
     static nvrtcResult nvrtcCompileProgram(nvrtcProgram prog, int numOptions, const char **options);
     static nvrtcResult nvrtcGetProgramLogSize(nvrtcProgram prog, size_t *logSizeRet);
@@ -188,6 +189,7 @@ private:
     static void* cuStreamDestroy_v2_;
     static void* cuEventDestroy_v2_;
     static void* cuMemAlloc_v2_;
+    static void* cuPointerGetAttribute_;
 
     static void* nvrtcCompileProgram_;
     static void* nvrtcGetProgramLogSize_;

lib/driver/backend.cpp

@@ -93,6 +93,15 @@ void backend::contexts::release()
     cache_.clear();
 }
 
+Context const & backend::contexts::import(CUcontext context)
+{
+  for(driver::Context const * x: cache_)
+      if(x->handle().cu()==context)
+          return *x;
+  cache_.emplace_back(new Context(context, false));
+  return *cache_.back();
+}
+
 Context const & backend::contexts::import(cl_context context)
 {
   for(driver::Context const * x: cache_)

lib/driver/buffer.cpp

@@ -9,7 +9,7 @@ namespace isaac
 namespace driver
 {
 
-Buffer::Buffer(CUdeviceptr h, bool take_ownership) : backend_(CUDA), context_(backend::contexts::get_default()), h_(backend_, take_ownership)
+Buffer::Buffer(CUdeviceptr h, bool take_ownership) : backend_(CUDA), context_(backend::contexts::import(Buffer::context(h))), h_(backend_, take_ownership)
 {
   h_.cu() = h;
 }

lib/driver/command_queue.cpp

@@ -28,8 +28,10 @@ CommandQueue::CommandQueue(Context const & context, Device const & device, cl_co
   switch(backend_)
   {
     case CUDA:
+    {
       cuda::check(dispatch::cuStreamCreate(&h_.cu(), 0));
       break;
+    }
 
     case OPENCL:
     {
@@ -38,7 +40,9 @@ CommandQueue::CommandQueue(Context const & context, Device const & device, cl_co
       ocl::check(err);
       break;
     }
-    default: throw;
+
+    default:
+      throw;
   }
 }
 

lib/driver/context.cpp

@@ -12,6 +12,11 @@ namespace isaac
 namespace driver
 {
 
+Context::Context(CUcontext const & context, CUdevice const & device, bool take_ownership) : backend_(CUDA), device_(device, false), cache_path_(tools::getenv("ISAAC_CACHE_PATH")), h_(backend_, take_ownership)
+{
+    h_.cu() = context;
+}
+
 Context::Context(cl_context const & context, bool take_ownership) : backend_(OPENCL), device_(ocl::info<CL_CONTEXT_DEVICES>(context)[0], false), cache_path_(tools::getenv("ISAAC_CACHE_PATH")), h_(backend_, take_ownership)
 {
     h_.cl() = context;

lib/driver/device.cpp

@@ -21,9 +21,14 @@ int Device::cuGetInfo() const
   return res;
 }
 
-Device::Device(int ordinal): backend_(CUDA), h_(backend_, true)
+//Device::Device(int ordinal): backend_(CUDA), h_(backend_, true)
+//{
+//  cuda::check(dispatch::cuDeviceGet(&h_.cu(), ordinal));
+//}
+
+Device::Device(CUdevice const & device, bool take_ownership): backend_(CUDA), h_(backend_, take_ownership)
 {
-  cuda::check(dispatch::cuDeviceGet(&h_.cu(), ordinal));
+  h_.cu() = device;
 }
 
 Device::Device(cl_device_id const & device, bool take_ownership) : backend_(OPENCL), h_(backend_, take_ownership)

lib/driver/dispatch.cpp

@@ -165,6 +165,7 @@ CUDA_DEFINE1(CUresult, cuStreamSynchronize, CUstream)
 CUDA_DEFINE1(CUresult, cuStreamDestroy_v2, CUstream)
 CUDA_DEFINE1(CUresult, cuEventDestroy_v2, CUevent)
 CUDA_DEFINE2(CUresult, cuMemAlloc_v2, CUdeviceptr*, size_t)
+CUDA_DEFINE3(CUresult, cuPointerGetAttribute, void*, CUpointer_attribute, CUdeviceptr)
 
 NVRTC_DEFINE3(nvrtcResult, nvrtcCompileProgram, nvrtcProgram, int, const char **)
 NVRTC_DEFINE2(nvrtcResult, nvrtcGetProgramLogSize, nvrtcProgram, size_t *)
@@ -255,6 +256,7 @@ void* dispatch::cuStreamSynchronize_;
 void* dispatch::cuStreamDestroy_v2_;
 void* dispatch::cuEventDestroy_v2_;
 void* dispatch::cuMemAlloc_v2_;
+void* dispatch::cuPointerGetAttribute_;
 
 void* dispatch::nvrtcCompileProgram_;
 void* dispatch::nvrtcGetProgramLogSize_;

lib/driver/platform.cpp

@@ -58,8 +58,11 @@ void Platform::devices(std::vector<Device> & devices) const
     {
       int N;
       cuda::check(dispatch::cuDeviceGetCount(&N));
-      for(int i = 0 ; i < N ; ++i)
-        devices.push_back(Device(i));
+      for(int i = 0 ; i < N ; ++i){
+        CUdevice device;
+        dispatch::cuDeviceGet(&device, i);
+        devices.push_back(Device(device));
+      }
       break;
     }
     case OPENCL:

lib/driver/program.cpp

@@ -24,7 +24,6 @@ Program::Program(Context const & context, std::string const & source) : backend_
   {
     case CUDA:
     {
-
       std::string prefix = context_.device_.name() + "cuda";
       std::string sha1 = tools::sha1(prefix + source);
       std::string fname(cache_path + sha1);
@@ -71,8 +70,6 @@ Program::Program(Context const & context, std::string const & source) : backend_
         cached.write((char*)ptx.data(), std::streamsize(ptx_size));
       }
 
-
-
 //    std::ofstream oss(sha1 + ".cu", std::ofstream::out | std::ofstream::trunc);
 //    oss << source << std::endl;
 //    oss.close();

lib/wrap/cublas.cpp

@@ -7,6 +7,15 @@ namespace sc = isaac;
 extern "C"
 {
 
+//    sc::driver::Context current_context()
+//    {
+//        CUcontext ctx;
+//        CUdevice dev;
+//        cuCtxGetCurrent(&ctx);
+//        cuCtxGetDevice(&dev);
+//        return sc::driver::Context(ctx, dev, false);
+//    }
+
     struct cublasContext
     {