batch_identity.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNELS_IDENTITY_HPP
#define STAN_MATH_OPENCL_KERNELS_IDENTITY_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* batch_identity_kernel_code = STRINGIFY(
    // \endcond
 
    __kernel void batch_identity(__global double* A, unsigned int batch_rows,
                                 unsigned int size) {
      // The ID of the matrix in the batch the thread is assigned to
      int batch_id = get_global_id(0);
      // The row and column of the matrix in the batch
      int batch_row = get_global_id(1);
      int batch_col = get_global_id(2);
      int index = batch_id * batch_rows * batch_rows + batch_col * batch_rows
                  + batch_row;
      // Check for potential overflows of A.
      if (index < size) {
        if (batch_row == batch_col) {
          A[index] = 1.0;
        } else {
          A[index] = 0.0;
        }
      }
    }
    // \cond
);
// \endcond
 
const kernel_cl<out_buffer, int, int> batch_identity(
    "batch_identity", {indexing_helpers, batch_identity_kernel_code});
 
}  // namespace opencl_kernels
}  // namespace math
}  // namespace stan
#endif
#endif