diag_inv.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNELS_DIAGONAL_INVERSE_LOWER_TRI_HPP
#define STAN_MATH_OPENCL_KERNELS_DIAGONAL_INVERSE_LOWER_TRI_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/kernel_cl.hpp>
#include <stan/math/opencl/buffer_types.hpp>
#include <string>
 
namespace stan {
namespace math {
namespace opencl_kernels {
// \cond
static constexpr const char* diag_inv_kernel_code = STRINGIFY(
    // \endcond
    __kernel void diag_inv(__global double* A, __global double* tmp_inv,
                           int rows) {
      int index = get_local_id(0);
      int group = get_group_id(0);
      int block_size = get_local_size(0);
      int A_offset = group * block_size;
      // offset inside the matrix with batched identities
      int tmp_offset = group * block_size * block_size + index * block_size;
 
      // The following code is the sequential version of forward
      // substitution with the identity matrix as RHS. Only the innermost loops
      // are parallelized. The rows are processed sequentially. This loop
      // process all the rows:
      for (int k = 0; k < block_size; k++) {
        double diag_ele = A(A_offset + k, A_offset + k);
 
        // Each element under the diagonal of the RHS is divided by diag_ele.
        // Each thread in a thread block does 1 division.
        // Threads that are assigned elements above the diagonal
        // skip this division.
        if (index <= k) {
          tmp_inv[tmp_offset + k] /= diag_ele;
        }
        barrier(CLK_LOCAL_MEM_FENCE);
        // Each thread updates one column in the RHS matrix
        // (ignores values above the diagonal).
        for (int i = max(k + 1, index); i < block_size; i++) {  // NOLINT
          double factor = A(A_offset + i, A_offset + k);
          tmp_inv[tmp_offset + i] -= tmp_inv[tmp_offset + k] * factor;
        }
        barrier(CLK_LOCAL_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];
      }
    }
    // \cond
);
// \endcond
 
const kernel_cl<in_out_buffer, in_out_buffer, int> diag_inv(
    "diag_inv", {indexing_helpers, diag_inv_kernel_code},
    {{"THREAD_BLOCK_SIZE", 32}});
 
}  // namespace opencl_kernels
}  // namespace math
}  // namespace stan
#endif
#endif