[refactor] Replace make_current with cuda context guard in jit cuda

session

Main goal is to get rid of the to-do in comment as it's fixed taichi-dev#5891

Author: Ailing Zhang

taichi/rhi/cuda/cuda_device.cpp

@@ -43,6 +43,7 @@ DeviceAllocation CudaDevice::allocate_memory_runtime(
       caching_allocator_ = std::make_unique<CudaCachingAllocator>(this);
     }
     info.ptr = caching_allocator_->allocate(params);
+    auto context_guard = CUDAContext::get_instance().get_guard();
     CUDADriver::get_instance().memset((void *)info.ptr, 0, info.size);
   } else {
     info.ptr = allocate_llvm_runtime_memory_jit(params);
@@ -117,6 +118,7 @@ DeviceAllocation CudaDevice::import_memory(void *ptr, size_t size) {
 }
 
 uint64 CudaDevice::fetch_result_uint64(int i, uint64 *result_buffer) {
+  auto context_guard = CUDAContext::get_instance().get_guard();
   CUDADriver::get_instance().stream_synchronize(nullptr);
   uint64 ret;
   CUDADriver::get_instance().memcpy_device_to_host(&ret, result_buffer + i,

taichi/runtime/cuda/jit_cuda.cpp

@@ -13,9 +13,7 @@ JITModule *JITSessionCUDA ::add_module(std::unique_ptr<llvm::Module> M,
                                      "module NVPTX");
     writer.write(ptx);
   }
-  // TODO: figure out why using the guard leads to wrong tests results
-  // auto context_guard = CUDAContext::get_instance().get_guard();
-  CUDAContext::get_instance().make_current();
+  auto context_guard = CUDAContext::get_instance().get_guard();
   // Create module for object
   void *cuda_module;
   TI_TRACE("PTX size: {:.2f}KB", ptx.size() / 1024.0);

taichi/runtime/cuda/jit_cuda.h

@@ -47,9 +47,7 @@ class JITModuleCUDA : public JITModule {
   }
 
   void *lookup_function(const std::string &name) override {
-    // TODO: figure out why using the guard leads to wrong tests results
-    // auto context_guard = CUDAContext::get_instance().get_guard();
-    CUDAContext::get_instance().make_current();
+    auto context_guard = CUDAContext::get_instance().get_guard();
     void *func = nullptr;
     auto t = Time::get_time();
     auto err = CUDADriver::get_instance().module_get_function.call_with_warning(

taichi/runtime/llvm/llvm_runtime_executor.cpp

@@ -476,6 +476,7 @@ void LlvmRuntimeExecutor::fill_ndarray(const DeviceAllocation &alloc,
   auto ptr = get_ndarray_alloc_info_ptr(alloc);
   if (config_->arch == Arch::cuda) {
 #if defined(TI_WITH_CUDA)
+    auto context_guard = CUDAContext::get_instance().get_guard();
     CUDADriver::get_instance().memsetd32((void *)ptr, data, size);
 #else
     TI_NOT_IMPLEMENTED