use tbb::concurrent_vector for multithreaded use of matrix_cl types

SteveBronder · SteveBronder · commit 400c94de6700 · 2023-05-16T17:37:29.000-04:00
diff --git a/stan/math/opencl/copy.hpp b/stan/math/opencl/copy.hpp
@@ -98,9 +98,10 @@ inline auto from_matrix_cl(const T& src) {
     try {
       cl::Event copy_event;
       const cl::CommandQueue queue = opencl_context.queue();
+      std::vector<cl::Event> copy_write_events(src.write_events().begin(), src.write_events().end());
       queue.enqueueReadBuffer(src.buffer(), opencl_context.in_order(), 0,
                               sizeof(T_val) * dst.size(), dst.data(),
-                              &src.write_events(), &copy_event);
+                              &copy_write_events, &copy_event);
       copy_event.wait();
       src.clear_write_events();
     } catch (const cl::Error& e) {
@@ -151,8 +152,9 @@ inline T_dst from_matrix_cl(const matrix_cl<T>& src) {
   try {
     cl::Event copy_event;
     const cl::CommandQueue queue = opencl_context.queue();
+      std::vector<cl::Event> copy_write_events(src.write_events().begin(), src.write_events().end());
     queue.enqueueReadBuffer(src.buffer(), opencl_context.in_order(), 0,
-                            sizeof(T), &dst, &src.write_events(), &copy_event);
+                            sizeof(T), &dst, &copy_write_events, &copy_event);
     copy_event.wait();
     src.clear_write_events();
   } catch (const cl::Error& e) {
@@ -183,9 +185,10 @@ inline T_dst from_matrix_cl(const matrix_cl<T>& src) {
   try {
     cl::Event copy_event;
     const cl::CommandQueue queue = opencl_context.queue();
+    std::vector<cl::Event> copy_write_events(src.write_events().begin(), src.write_events().end());
     queue.enqueueReadBuffer(src.buffer(), opencl_context.in_order(), 0,
                             sizeof(T) * src.rows(), dst.data(),
-                            &src.write_events(), &copy_event);
+                            &copy_write_events, &copy_event);
     copy_event.wait();
     src.clear_write_events();
   } catch (const cl::Error& e) {
@@ -257,7 +260,7 @@ inline auto packed_copy(const T& src) {
                                      packed, src, src.rows(), src.rows(),
                                      src.view());
     const std::vector<cl::Event> mat_events
-        = vec_concat(packed.read_write_events(), src.write_events());
+        = vec_concat(std::vector<cl::Event>{}, packed.read_write_events(), src.write_events());
     cl::Event copy_event;
     queue.enqueueReadBuffer(packed.buffer(), opencl_context.in_order(), 0,
                             sizeof(T_val) * packed_size, dst.data(),
diff --git a/stan/math/opencl/kernel_cl.hpp b/stan/math/opencl/kernel_cl.hpp
@@ -109,17 +109,17 @@ inline void assign_events(const cl::Event& new_event, CallArg& m,
  * @return A vector of OpenCL events.
  */
 template <typename T, require_not_matrix_cl_t<T>* = nullptr>
-inline std::vector<cl::Event> select_events(const T& m) {
-  return {};
+inline tbb::concurrent_vector<cl::Event> select_events(const T& m) {
+  return tbb::concurrent_vector<cl::Event>{};
 }
 template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
           require_same_t<T, in_buffer>* = nullptr>
-inline const std::vector<cl::Event>& select_events(const K& m) {
+inline const tbb::concurrent_vector<cl::Event>& select_events(const K& m) {
   return m.write_events();
 }
 template <typename T, typename K, require_matrix_cl_t<K>* = nullptr,
           require_any_same_t<T, out_buffer, in_out_buffer>* = nullptr>
-inline std::vector<cl::Event> select_events(K& m) {
+inline tbb::concurrent_vector<cl::Event> select_events(K& m) {
   static_assert(!std::is_const<K>::value, "Can not write to const matrix_cl!");
   return m.read_write_events();
 }
@@ -205,7 +205,7 @@ struct kernel_cl {
       opencl_context.register_kernel_cache(&kernel_);
     }
     cl::EnqueueArgs eargs(opencl_context.queue(),
-                          vec_concat(internal::select_events<Args>(args)...),
+                          vec_concat(std::vector<cl::Event>{}, internal::select_events<Args>(args)...),
                           global_thread_size);
     cl::KernelFunctor<internal::to_const_buffer_t<Args>&...> kernel_functor(
         kernel_);
@@ -232,7 +232,7 @@ struct kernel_cl {
       opencl_context.register_kernel_cache(&kernel_);
     }
     cl::EnqueueArgs eargs(opencl_context.queue(),
-                          vec_concat(internal::select_events<Args>(args)...),
+                          vec_concat(std::vector<cl::Event>{}, internal::select_events<Args>(args)...),
                           global_thread_size, thread_block_size);
     cl::KernelFunctor<internal::to_const_buffer_t<Args>&...> kernel_functor(
         kernel_);
diff --git a/stan/math/opencl/matrix_cl.hpp b/stan/math/opencl/matrix_cl.hpp
@@ -12,6 +12,7 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/vec_concat.hpp>
 #include <CL/opencl.hpp>
+#include <tbb/concurrent_vector.h>
 #include <algorithm>
 #include <iostream>
 #include <string>
@@ -50,8 +51,8 @@ class matrix_cl : public matrix_cl_base {
   int cols_{0};           // Number of columns.
   // Holds info on if matrix is a special type
   matrix_cl_view view_{matrix_cl_view::Entire};
-  mutable std::vector<cl::Event> write_events_;  // Tracks write jobs
-  mutable std::vector<cl::Event> read_events_;   // Tracks reads
+  mutable tbb::concurrent_vector<cl::Event> write_events_;  // Tracks write jobs
+  mutable tbb::concurrent_vector<cl::Event> read_events_;   // Tracks reads
 
  public:
   using Scalar = T;  // Underlying type of the matrix
@@ -99,23 +100,23 @@ class matrix_cl : public matrix_cl_base {
    * Get the events from the event stacks.
    * @return The write event stack.
    */
-  inline const std::vector<cl::Event>& write_events() const {
+  inline const tbb::concurrent_vector<cl::Event>& write_events() const {
     return write_events_;
   }
 
   /**
    * Get the events from the event stacks.
    * @return The read/write event stack.
    */
-  inline const std::vector<cl::Event>& read_events() const {
+  inline const tbb::concurrent_vector<cl::Event>& read_events() const {
     return read_events_;
   }
 
   /**
    * Get the events from the event stacks.
    * @return The read/write event stack.
    */
-  inline const std::vector<cl::Event> read_write_events() const {
+  inline const tbb::concurrent_vector<cl::Event> read_write_events() const {
     return vec_concat(this->read_events(), this->write_events());
   }
 
@@ -615,15 +616,29 @@ class matrix_cl : public matrix_cl_base {
    * @param A matrix_cl
    */
   void initialize_buffer_cl(const matrix_cl<T>& A) {
+    cl::Event cstr_event;
+    std::vector<cl::Event>* dep_events = 
+      new std::vector<cl::Event>(A.write_events().begin(), 
+      A.write_events().end()); 
     try {
-      cl::Event cstr_event;
       opencl_context.queue().enqueueCopyBuffer(A.buffer(), this->buffer(), 0, 0,
                                                A.size() * sizeof(T),
-                                               &A.write_events(), &cstr_event);
+                                               dep_events, &cstr_event);
+      if (opencl_context.device()[0].getInfo<CL_DEVICE_HOST_UNIFIED_MEMORY>()) {
+        buffer_cl_.setDestructorCallback(
+          &delete_it_destructor<std::vector<cl::Event>>, dep_events);
+      } else {
+        cstr_event.setCallback(CL_COMPLETE, 
+          &delete_it_event<std::vector<cl::Event>>, dep_events);
+      }
       this->add_write_event(cstr_event);
       A.add_read_event(cstr_event);
     } catch (const cl::Error& e) {
+      delete dep_events;
       check_opencl_error("copy (OpenCL)->(OpenCL)", e);
+    } catch (...) {
+      delete dep_events;
+      throw;
     }
   }
 
diff --git a/stan/math/prim/fun/vec_concat.hpp b/stan/math/prim/fun/vec_concat.hpp
@@ -37,7 +37,9 @@ inline void append_vectors(VecInOut& x) {}
 template <typename VecInOut, typename VecIn, typename... VecArgs>
 inline void append_vectors(VecInOut& x, const VecIn& y,
                            const VecArgs&... args) {
-  x.insert(x.end(), y.begin(), y.end());
+  for (auto& yy : y) {
+    x.push_back(yy);
+  }
   append_vectors(x, args...);
 }
 }  // namespace internal
@@ -53,7 +55,7 @@ inline void append_vectors(VecInOut& x, const VecIn& y,
  */
 template <typename Vec, typename... Args>
 inline auto vec_concat(const Vec& v1, const Args&... args) {
-  std::vector<value_type_t<Vec>> vec;
+  Vec vec;
   vec.reserve(internal::sum_vector_sizes(v1, args...));
   internal::append_vectors(vec, v1, args...);
   return vec;
diff --git a/test/unit/math/opencl/rev/normal_lccdf_test.cpp b/test/unit/math/opencl/rev/normal_lccdf_test.cpp
@@ -5,139 +5,35 @@
 #include <test/unit/math/opencl/util.hpp>
 #include <vector>
 
-TEST(ProbDistributionsNormalLccdf, error_checking) {
-  int N = 3;
-
-  Eigen::VectorXd y(N);
-  y << 0.3, 0.8, 1.0;
-  Eigen::VectorXd y_size(N - 1);
-  y_size << 0.3, 0.8;
-  Eigen::VectorXd y_value(N);
-  y_value << 0.3, NAN, 0.5;
-
-  Eigen::VectorXd mu(N);
-  mu << 0.3, 0.8, 1.0;
-  Eigen::VectorXd mu_size(N - 1);
-  mu_size << 0.3, 0.8;
-  Eigen::VectorXd mu_value(N);
-  mu_value << 0.3, -INFINITY, 0.5;
-
-  Eigen::VectorXd sigma(N);
-  sigma << 0.3, 0.8, 1.0;
-  Eigen::VectorXd sigma_size(N - 1);
-  sigma_size << 0.3, 0.8;
-  Eigen::VectorXd sigma_value(N);
-  sigma_value << 0.3, 0, 0.5;
-
-  stan::math::matrix_cl<double> y_cl(y);
-  stan::math::matrix_cl<double> y_size_cl(y_size);
-  stan::math::matrix_cl<double> y_value_cl(y_value);
-  stan::math::matrix_cl<double> mu_cl(mu);
-  stan::math::matrix_cl<double> mu_size_cl(mu_size);
-  stan::math::matrix_cl<double> mu_value_cl(mu_value);
-  stan::math::matrix_cl<double> sigma_cl(sigma);
-  stan::math::matrix_cl<double> sigma_size_cl(sigma_size);
-  stan::math::matrix_cl<double> sigma_value_cl(sigma_value);
-
-  EXPECT_NO_THROW(stan::math::normal_lccdf(y_cl, mu_cl, sigma_cl));
-
-  EXPECT_THROW(stan::math::normal_lccdf(y_size_cl, mu_cl, sigma_cl),
-               std::invalid_argument);
-  EXPECT_THROW(stan::math::normal_lccdf(y_cl, mu_size_cl, sigma_cl),
-               std::invalid_argument);
-  EXPECT_THROW(stan::math::normal_lccdf(y_cl, mu_cl, sigma_size_cl),
-               std::invalid_argument);
-
-  EXPECT_THROW(stan::math::normal_lccdf(y_value_cl, mu_cl, sigma_cl),
-               std::domain_error);
-  EXPECT_THROW(stan::math::normal_lccdf(y_cl, mu_value_cl, sigma_cl),
-               std::domain_error);
-  EXPECT_THROW(stan::math::normal_lccdf(y_cl, mu_cl, sigma_value_cl),
-               std::domain_error);
-}
 
 auto normal_lccdf_functor
     = [](const auto& y, const auto& mu, const auto& sigma) {
         return stan::math::normal_lccdf(y, mu, sigma);
       };
 
-TEST(ProbDistributionsNormalLccdf, opencl_matches_cpu_small) {
-  int N = 3;
-  int M = 2;
 
-  Eigen::VectorXd y(N);
-  y << 0.3, 0.8, 1.0;
-  Eigen::VectorXd mu(N);
-  mu << -0.3, -0.8, 1.01;
-  Eigen::VectorXd sigma(N);
-  sigma << 0.3, 0.1, 1.0;
-
-  stan::math::test::compare_cpu_opencl_prim_rev(normal_lccdf_functor, y, mu,
-                                                sigma);
-  stan::math::test::compare_cpu_opencl_prim_rev(
-      normal_lccdf_functor, y.transpose().eval(), mu.transpose().eval(),
-      sigma.transpose().eval());
-}
-
-TEST(ProbDistributionsNormalLccdf, opencl_broadcast_y) {
-  int N = 3;
-
-  double y_scal = 12.3;
-  Eigen::VectorXd mu(N);
-  mu << 0.5, 1.2, 1.0;
-  Eigen::VectorXd sigma(N);
-  sigma << 0.3, 0.8, 1.0;
-
-  stan::math::test::test_opencl_broadcasting_prim_rev<0>(normal_lccdf_functor,
-                                                         y_scal, mu, sigma);
-  stan::math::test::test_opencl_broadcasting_prim_rev<0>(
-      normal_lccdf_functor, y_scal, mu.transpose().eval(), sigma);
-}
-
-TEST(ProbDistributionsNormalLccdf, opencl_broadcast_mu) {
-  int N = 3;
-
-  Eigen::VectorXd y(N);
-  y << 0.3, 0.8, 1.0;
-  double mu_scal = 12.3;
-  Eigen::VectorXd sigma(N);
-  sigma << 0.3, 0.8, 1.0;
-
-  stan::math::test::test_opencl_broadcasting_prim_rev<1>(normal_lccdf_functor,
-                                                         y, mu_scal, sigma);
-  stan::math::test::test_opencl_broadcasting_prim_rev<1>(
-      normal_lccdf_functor, y.transpose().eval(), mu_scal, sigma);
-}
-
-TEST(ProbDistributionsNormalLccdf, opencl_broadcast_sigma) {
-  int N = 3;
-
-  Eigen::VectorXd y(N);
-  y << 0.3, 0.8, 1.0;
-  Eigen::VectorXd mu(N);
-  mu << 0.3, 0.8, 1.0;
-  double sigma_scal = 12.3;
-
-  stan::math::test::test_opencl_broadcasting_prim_rev<2>(normal_lccdf_functor,
-                                                         y, mu, sigma_scal);
-  stan::math::test::test_opencl_broadcasting_prim_rev<2>(
-      normal_lccdf_functor, y.transpose().eval(), mu, sigma_scal);
-}
 
 TEST(ProbDistributionsNormalLccdf, opencl_matches_cpu_big) {
   int N = 153;
 
-  Eigen::Matrix<double, Eigen::Dynamic, 1> y
-      = Eigen::Array<double, Eigen::Dynamic, 1>::Random(N, 1).abs();
+std::srand(123);
+for (int i = 0; i < 10; ++i) {
   Eigen::Matrix<double, Eigen::Dynamic, 1> mu
-      = Eigen::Array<double, Eigen::Dynamic, 1>::Random(N, 1).abs();
+      = Eigen::Array<double, Eigen::Dynamic, 1>::Random(N, 1) + 1.0;
   Eigen::Matrix<double, Eigen::Dynamic, 1> sigma
       = Eigen::Array<double, Eigen::Dynamic, 1>::Random(N, 1).abs() + 0.01;
-
+  Eigen::Matrix<double, Eigen::Dynamic, 1> y = (mu.array() * sigma.array()).matrix();
+  std::cout << "Iter: " << i << " mu, sigma, y" << std::endl;
+  for (int j = 0; j < N; j++) {
+    std::cout << mu(j) << ", " << sigma(j) << ", " << y(j) << std::endl;
+  }
+  std::cout << "-----------compare_cpu_opencl_prim_rev" << std::endl;
   stan::math::test::compare_cpu_opencl_prim_rev(normal_lccdf_functor, y, mu,
                                                 sigma);
+  std::cout << "-----------compare_cpu_opencl_prim_rev transpose" << std::endl;
   stan::math::test::compare_cpu_opencl_prim_rev(
       normal_lccdf_functor, y.transpose().eval(), mu.transpose().eval(),
       sigma.transpose().eval());
 }
+}
 #endif
