cusolver parallel stream launch macro

pytorch · xwang233 · Nov 25, 2020 · Nov 25, 2020 · Nov 25, 2020 · Nov 25, 2020
commit 625a32fb4eb8dc23bdf1950c41d3a5edcdd6425a
diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cu
@@ -48,26 +48,14 @@ static void apply_batched_inverse_lib(Tensor& self, Tensor& self_inv, Tensor& in
 
   if (use_loop_launch(batch_size, n)) {
     int* p_infos = infos.data_ptr<int>();
-    auto main_stream = at::cuda::getCurrentCUDAStream();
-
-    at::cuda::CUDAEvent main_event;
-    main_event.record(main_stream);
-
-    for (int64_t i = 0; i < batch_size; i++) {
-      auto stream = at::cuda::getStreamFromPool();
-      at::cuda::CUDAStreamGuard guard(stream);
-
-      main_event.block(stream);
 
+    CUDA_PARALLEL_STREAM_LAUNCH(i, batch_size, [&] {
       auto dataPtr = allocator.allocate(sizeof(int) * n);
       int* pivot = reinterpret_cast<int*>(dataPtr.get());
       _apply_single_inverse_helper<scalar_t>(
         &self_data[i * self_mat_stride], &self_inv_data[i * self_inv_mat_stride], pivot, p_infos + i * 2, n);
+    });
 
-      at::cuda::CUDAEvent finished;
-      finished.record(stream);
-      finished.block(main_stream);
-    }
   } else {
     // cublas batched kernels require input be "device array of device pointers"
     Tensor self_array = at::arange(

diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.h b/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.h
@@ -12,6 +12,41 @@
 #define USE_CUSOLVER
 #endif
 
+// comment out this line for the emergency break if there is something wrong with parallel launch
+// e.g. if you have a script failure with linear algebra and set environment variable CUDA_LAUNCH_BLOCKING=1
+// solves the problem, then you can try removing the cuda parallel stream launch
+#define _USE_PARALLEL_LAUNCH
+
+
+#ifdef _USE_PARALLEL_LAUNCH
+
+#define CUDA_PARALLEL_STREAM_LAUNCH(i, _batch_size, _stmt)  \
+do {                                                        \
+  auto _main_stream = at::cuda::getCurrentCUDAStream();     \
+  at::cuda::CUDAEvent _main_event;                          \
+  _main_event.record(_main_stream);                         \
+  for (int i = 0; i < _batch_size; i++) {                   \
+    auto _stream = at::cuda::getStreamFromPool();           \
+    at::cuda::CUDAStreamGuard guard(_stream);               \
+    _main_event.block(_stream);                             \
+    _stmt();                                                \
+    at::cuda::CUDAEvent _finished;                          \
+    _finished.record(_stream);                              \
+    _finished.block(_main_stream);                          \
+  }                                                         \
+} while(0)
+
+#else // _USE_PARALLEL_LAUNCH is not defined
+
+#define CUDA_PARALLEL_STREAM_LAUNCH(i, _batch_size, _stmt)  \
+do {                                                        \
+  for (int i = 0; i < _batch_size; i++) {                   \
+    _stmt();                                                \
+  }                                                         \
+} while (0)
+
+#endif  // ifdef _USE_PARALLEL_LAUNCH
+
 #ifdef USE_CUSOLVER
 
 namespace at {