[cuda] Clear cuda context after init

python/taichi/lang/_ndarray.py

@@ -307,11 +307,15 @@ def __init__(self, arr, indices_first, indices_second):
             self.indices = indices_first + indices_second
 
         def getter():
+            if impl.get_runtime().prog.config.arch == _ti_core.cuda:
+                _ti_core.push_cuda_context()
             self.ndarr._initialize_host_accessor()
             return self.ndarr.host_accessor.getter(
                 *self.ndarr._pad_key(self.indices))
 
         def setter(value):
+            if impl.get_runtime().prog.config.arch == _ti_core.cuda:
+                _ti_core.push_cuda_context()
             self.ndarr._initialize_host_accessor()
             self.ndarr.host_accessor.setter(value,
                                             *self.ndarr._pad_key(self.indices))

python/taichi/lang/misc.py

@@ -459,6 +459,9 @@ def init(arch=None,
 
     impl._root_fb = _snode.FieldsBuilder()
 
+    if get_runtime().prog.config.arch == cuda:
+        _ti_core.pop_cuda_context()
+
     if cfg.debug:
         impl.get_runtime()._register_signal_handlers()
 

taichi/python/export_lang.cpp

@@ -1100,6 +1100,14 @@ void export_lang(py::module &m) {
     }
   });
 
+#if defined(TI_WITH_CUDA)
+  m.def("pop_cuda_context",
+        []() { CUDADriver::get_instance().context_pop_current(NULL); });
+
+  m.def("push_cuda_context",
+        []() { CUDAContext::get_instance().make_current(); });
+#endif
+
   // Type system
 
   py::class_<Type>(m, "Type").def("to_string", &Type::to_string);

taichi/rhi/cuda/cuda_driver_functions.inc.h

@@ -14,6 +14,7 @@ PER_CUDA_FUNCTION(device_get_attribute, cuDeviceGetAttribute, int *, uint32, voi
 PER_CUDA_FUNCTION(context_create, cuCtxCreate_v2, void*, int, void *);
 PER_CUDA_FUNCTION(context_set_current, cuCtxSetCurrent, void *);
 PER_CUDA_FUNCTION(context_get_current, cuCtxGetCurrent, void **);
+PER_CUDA_FUNCTION(context_pop_current, cuCtxPopCurrent, void **);
 
 // Stream management
 PER_CUDA_FUNCTION(stream_create, cuStreamCreate, void **, uint32);

taichi/runtime/llvm/llvm_runtime_executor.cpp

@@ -176,6 +176,7 @@ void LlvmRuntimeExecutor::print_list_manager_info(void *list_manager,
 void LlvmRuntimeExecutor::synchronize() {
   if (config_->arch == Arch::cuda) {
 #if defined(TI_WITH_CUDA)
+    CUDAContext::get_instance().make_current();
     CUDADriver::get_instance().stream_synchronize(nullptr);
 #else
     TI_ERROR("No CUDA support");
@@ -190,6 +191,7 @@ uint64 LlvmRuntimeExecutor::fetch_result_uint64(int i, uint64 *result_buffer) {
   uint64 ret;
   if (config_->arch == Arch::cuda) {
 #if defined(TI_WITH_CUDA)
+    CUDAContext::get_instance().make_current();
     CUDADriver::get_instance().memcpy_device_to_host(&ret, result_buffer + i,
                                                      sizeof(uint64));
 #else