Fix race conditions caused by missing fences in OpenCL kernels for cholesky_decompose, tridiagonalization and diag_inv

stan/math/opencl/cholesky_decompose.hpp

@@ -52,7 +52,8 @@ inline void cholesky_decompose(matrix_cl<T>& A) {
   if (A.rows() <= opencl_context.tuning_opts().cholesky_min_L11_size) {
     try {
       opencl_kernels::cholesky_decompose(cl::NDRange(A.rows()),
-                                         cl::NDRange(A.rows()), A, A.rows());
+                                         cl::NDRange(A.rows()), A, A.rows(),
+                                         cl::Local(sizeof(double) * A.rows()));
     } catch (const cl::Error& e) {
       check_opencl_error("cholesky_decompose", e);
     }

stan/math/opencl/kernels/cholesky_decompose.hpp

@@ -24,33 +24,40 @@ static constexpr const char* cholesky_decompose_kernel_code = STRINGIFY(
      * @param[in, out] A The input matrix and the result of the cholesky
      *  decomposition
      * @param rows The number of rows for A and B.
+     * @param row_j Local memory buffer of rows doubles that holds row j.
      * @note Code is a <code>const char*</code> held in
      * <code>cholesky_decompose_kernel_code.</code>
      *  Used in math/opencl/cholesky_decompose.hpp.
      *  This kernel uses the helper macros available in helpers.cl.
      *
      */
-    __kernel void cholesky_decompose(__global double* A, int rows) {
+    __kernel void cholesky_decompose(__global double* A, int rows,
+                                     __local double* row_j) {
       const int local_index = get_local_id(0);
       // The following code is the sequential version of the inplace
       // cholesky decomposition. Only the innermost loops are parallelized. The
       // rows are processed sequentially. This loop process all the rows:
       for (int j = 0; j < rows; j++) {
-        if (local_index == 0) {
+        if (local_index == j) {
           double sum = 0.0;
           for (int k = 0; k < j; k++) {
-            sum = sum + A(j, k) * A(j, k);
+            const double a = A(j, k);
+            row_j[k] = a;
+            sum += a * a;
           }
           A(j, j) = sqrt(A(j, j) - sum);
+          row_j[j] = A(j, j);
         }
+        // Each work item accesses only its own row of A and reads row j from
+        // local memory, so the barriers need only a local memory fence.
         barrier(CLK_LOCAL_MEM_FENCE);
         if (local_index < j) {
           A(local_index, j) = 0.0;
         } else if (local_index > j) {
           double sum = 0.0;
           for (int k = 0; k < j; k++)
-            sum = sum + A(local_index, k) * A(j, k);
-          A(local_index, j) = (A(local_index, j) - sum) / A(j, j);
+            sum += A(local_index, k) * row_j[k];
+          A(local_index, j) = (A(local_index, j) - sum) / row_j[j];
         }
         barrier(CLK_LOCAL_MEM_FENCE);
       }
@@ -63,7 +70,7 @@ static constexpr const char* cholesky_decompose_kernel_code = STRINGIFY(
  * See the docs for \link kernels/cholesky_decompose.hpp cholesky_decompose()
  * \endlink
  */
-const kernel_cl<in_out_buffer, int> cholesky_decompose(
+const kernel_cl<in_out_buffer, int, cl::LocalSpaceArg> cholesky_decompose(
     "cholesky_decompose", {indexing_helpers, cholesky_decompose_kernel_code});
 
 }  // namespace opencl_kernels

stan/math/opencl/kernels/diag_inv.hpp

@@ -72,6 +72,9 @@ static constexpr const char* diag_inv_kernel_code = STRINGIFY(
         }
         barrier(CLK_LOCAL_MEM_FENCE);
       }
+      // The copy below overwrites entries of A that other work items read
+      // above, so a global memory fence separates the two.
+      barrier(CLK_GLOBAL_MEM_FENCE);
       for (int j = 0; j < block_size; j++) {
         // Each thread copies one column.
         A(A_offset + j, A_offset + index) = tmp_inv[tmp_offset + j];

stan/math/opencl/kernels/tridiagonalization.hpp

@@ -55,7 +55,9 @@ static constexpr const char* tridiagonalization_householder_kernel_code
           // calculate column norm between threads
           __local double q_local[LOCAL_SIZE_];
           q_local[lid] = q;
-          barrier(CLK_LOCAL_MEM_FENCE);
+          // Work items below read and write entries of P that other work
+          // items wrote above, so this barrier also fences global memory.
+          barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
           for (int step = lsize / REDUCTION_STEP_SIZE; step > 0;
                step /= REDUCTION_STEP_SIZE) {
             if (lid < step) {

test/unit/math/opencl/cholesky_decompose_test.cpp

@@ -29,6 +29,20 @@ TEST(MathMatrixOpenCL, cholesky_decompose_cpu_vs_cl_small) {
   EXPECT_MATRIX_NEAR(m1, m1_res, 1e-8);
 }
 
+TEST(MathMatrixOpenCL, cholesky_decompose_cpu_vs_cl_one_work_group) {
+  int size = 200;
+  stan::math::matrix_d m = stan::math::matrix_d::Random(size, size);
+  stan::math::matrix_d m_pos_def
+      = m * m.transpose() + size * Eigen::MatrixXd::Identity(size, size);
+
+  stan::math::matrix_cl<double> m_cl(m_pos_def);
+  stan::math::matrix_d m_res = stan::math::cholesky_decompose(m_pos_def);
+
+  stan::math::opencl::cholesky_decompose(m_cl);
+
+  EXPECT_MATRIX_NEAR(m_res, stan::math::from_matrix_cl(m_cl), 1e-8);
+}
+
 namespace {
 inline void cholesky_decompose_test(int size) {
   stan::math::matrix_d m1 = stan::math::matrix_d::Random(size, size);