[aot] Add stream_ variable for CUDAContext to use a specific CUDA str…

…eam to launch CUDA kernel (#8579)

### Brief Summary

copilot:summary

### Walkthrough

copilot:walkthrough

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

c_api/include/taichi/taichi_cuda.h

@@ -20,10 +20,17 @@ ti_export_cuda_memory(TiRuntime runtime,
                       TiMemory memory,
                       TiCudaMemoryInteropInfo *interop_info);
 
+// Function `ti_import_cuda_memory`
 TI_DLL_EXPORT TiMemory TI_API_CALL ti_import_cuda_memory(TiRuntime runtime,
                                                          void *ptr,
                                                          size_t memory_size);
 
+// Function `ti_set_cuda_stream`
+TI_DLL_EXPORT void TI_API_CALL ti_set_cuda_stream(void *stream);
+
+// Function `ti_get_cuda_stream`
+TI_DLL_EXPORT void TI_API_CALL ti_get_cuda_stream(void **stream);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif  // __cplusplus

c_api/src/taichi_llvm_impl.cpp

@@ -14,6 +14,7 @@
 
 #ifdef TI_WITH_CUDA
 #include "taichi/rhi/cuda/cuda_device.h"
+#include "taichi/rhi/cuda/cuda_context.h"
 #include "taichi/runtime/cuda/kernel_launcher.h"
 #endif
 
@@ -242,4 +243,24 @@ TI_DLL_EXPORT TiMemory TI_API_CALL ti_import_cuda_memory(TiRuntime runtime,
 #endif
 }
 
+// function.set_cuda_stream
+TI_DLL_EXPORT void TI_API_CALL ti_set_cuda_stream(void *stream) {
+#ifdef TI_WITH_CUDA
+  taichi::lang::CUDAContext::get_instance().set_stream(stream);
+
+#else
+  TI_NOT_IMPLEMENTED;
+#endif
+}
+
+// function.get_cuda_stream
+TI_DLL_EXPORT void TI_API_CALL ti_get_cuda_stream(void **stream) {
+#ifdef TI_WITH_CUDA
+  *stream = taichi::lang::CUDAContext::get_instance().get_stream();
+#else
+  TI_NOT_IMPLEMENTED;
+
+#endif
+}
+
 #endif  // TI_WITH_LLVM

c_api/tests/c_api_interop_test.cpp

@@ -160,3 +160,23 @@ TEST_F(CapiTest, TestCUDAImport) {
   EXPECT_EQ(data_out[3], 4.0);
 }
 #endif  // TI_WITH_CUDA
+
+#ifdef TI_WITH_CUDA
+TEST_F(CapiTest, TestCUDAStreamSet) {
+  void *temp_stream = nullptr;
+
+  ti_get_cuda_stream(&temp_stream);
+  EXPECT_EQ(temp_stream, nullptr);
+
+  void *stream1 = reinterpret_cast<void *>(0x12345678);
+  void *stream2 = reinterpret_cast<void *>(0x87654321);
+
+  ti_set_cuda_stream(stream1);
+  ti_get_cuda_stream(&temp_stream);
+  EXPECT_EQ(temp_stream, stream1);
+
+  ti_set_cuda_stream(stream2);
+  ti_get_cuda_stream(&temp_stream);
+  EXPECT_EQ(temp_stream, stream2);
+}
+#endif

taichi/rhi/cuda/cuda_context.cpp

@@ -12,7 +12,9 @@
 namespace taichi::lang {
 
 CUDAContext::CUDAContext()
-    : profiler_(nullptr), driver_(CUDADriver::get_instance_without_context()) {
+    : profiler_(nullptr),
+      driver_(CUDADriver::get_instance_without_context()),
+      stream_(nullptr) {
   // CUDA initialization
   dev_count_ = 0;
   driver_.init(0);
@@ -156,14 +158,14 @@ void CUDAContext::launch(void *func,
           dynamic_shared_mem_bytes);
     }
     driver_.launch_kernel(func, grid_dim, 1, 1, block_dim, 1, 1,
-                          dynamic_shared_mem_bytes, nullptr,
+                          dynamic_shared_mem_bytes, stream_,
                           arg_pointers.data(), nullptr);
   }
   if (profiler_)
     profiler_->stop(task_handle);
 
   if (debug_) {
-    driver_.stream_synchronize(nullptr);
+    driver_.stream_synchronize(stream_);
   }
 }
 

taichi/rhi/cuda/cuda_context.h

@@ -29,6 +29,7 @@ class CUDAContext {
   int max_shared_memory_bytes_;
   bool debug_;
   bool supports_mem_pool_;
+  void *stream_;
 
  public:
   CUDAContext();
@@ -108,6 +109,14 @@ class CUDAContext {
   }
 
   static CUDAContext &get_instance();
+
+  void set_stream(void *stream) {
+    stream_ = stream;
+  }
+
+  void *get_stream() const {
+    return stream_;
+  }
 };
 
 }  // namespace taichi::lang