check_cl.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_CHECK_CL_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_CHECK_CL_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/opencl/copy.hpp>
#include <stan/math/opencl/value_type.hpp>
#include <stan/math/opencl/kernel_generator/as_operation_cl.hpp>
#include <stan/math/opencl/kernel_generator/operation_cl_lhs.hpp>
#include <stan/math/opencl/kernel_generator/constant.hpp>
#include <stan/math/opencl/kernel_generator/scalar.hpp>
#include <map>
 
namespace stan {
namespace math {
 
template <typename T>
class check_cl_ : public operation_cl_lhs<check_cl_<T>, bool> {
 public:
  using Scalar = bool;
  using base = operation_cl_lhs<check_cl_<T>, bool>;
  using base::var_name_;
  using base::operator=;
 
 protected:
  // buffer[0] is indicator whether exception needs to be raised
  // buffer[1,2] are problematic indices
  matrix_cl<int> buffer_;
  matrix_cl<value_type_t<T>> value_;
 
 public:
  T arg_;
  const char* function_;
  const char* err_variable_;
  const char* must_be_;
 
  check_cl_(const char* function, const char* err_variable, T&& y,
            const char* must_be)
      : buffer_(constant(0, 3, 1)),
        value_(1, 1),
        arg_(std::forward<T>(y)),
        function_(function),
        err_variable_(err_variable),
        must_be_(must_be) {
    buffer_.view(matrix_cl_view::Entire);
  }
 
  // this operation can not be used on the right hand side of assignment
  inline kernel_parts get_kernel_parts(...) = delete;
  inline void extreme_diagonals(...) = delete;
 
  inline kernel_parts get_kernel_parts_lhs(
      std::unordered_map<const void*, const char*>& generated,
      std::unordered_map<const void*, const char*>& generated_all,
      name_generator& name_gen, const std::string& row_index_name,
      const std::string& col_index_name) const {
    kernel_parts res;
    this->var_name_ = name_gen.generate();
    generated[this] = "";
    res = arg_.get_kernel_parts(generated, generated_all, name_gen,
                                row_index_name, col_index_name, false);
 
    res.args += "__global int* " + var_name_ + "_buffer, __global "
                + type_str<value_type_t<T>>() + "* " + var_name_ + "_value, ";
    res.body += "bool " + var_name_;
    res.body_suffix += "if(!" + var_name_ +
            " && atomic_xchg(" + var_name_ + "_buffer, 1) == 0){\n"
          + var_name_ + "_buffer[1] = " + row_index_name + ";\n"
          + var_name_ + "_buffer[2] = " + col_index_name + ";\n"
          + var_name_ + "_value[0] = " + arg_.var_name_ + ";\n"
          "}";
    return res;
  }
 
  inline void set_args(
      std::unordered_map<const void*, const char*>& generated,
      std::unordered_map<const void*, const char*>& generated_all,
      cl::Kernel& kernel, int& arg_num) const {
    generated[this] = "";
    arg_.set_args(generated, generated_all, kernel, arg_num);
    kernel.setArg(arg_num++, buffer_.buffer());
    kernel.setArg(arg_num++, value_.buffer());
  }
 
  inline void check_assign_dimensions(int rows, int cols) const {
    if (arg_.rows() != base::dynamic) {
      check_size_match("check_cl_.check_assign_dimensions", "Rows of ",
                       err_variable_, arg_.rows(), "rows of ",
                       "assigned expression", rows);
    }
    if (arg_.cols() != base::dynamic) {
      check_size_match("check_cl_.check_assign_dimensions", "Columns of ",
                       err_variable_, arg_.cols(), "columns of ",
                       "assigned expression", cols);
    }
  }
 
  inline void get_clear_read_write_events(
      std::vector<cl::Event>& events) const {
    arg_.get_write_events(events);
    events.insert(events.end(), buffer_.read_events().begin(),
                  buffer_.read_events().end());
    events.insert(events.end(), buffer_.write_events().begin(),
                  buffer_.write_events().end());
  }
 
  inline void add_write_event(cl::Event& e) const {
    e.wait();
    Eigen::VectorXi res = from_matrix_cl<Eigen::VectorXi>(buffer_);
    if (res.coeff(0)) {
      double value = from_matrix_cl<scalar_type_t<T>>(value_);
      std::stringstream s;
      s << function_ << ": " << err_variable_ << "[" << res[1] << ", " << res[2]
        << "] = " << value << ", but it must be " << must_be_ << "!";
      throw std::domain_error(s.str());
    }
  }
 
  inline int rows() const { return arg_.rows(); }
 
  inline int cols() const { return arg_.cols(); }
 
  void operator=(bool condition);  // implemented in multi_result_kernel.hpp
};
 
namespace internal {
template <typename T>
struct is_scalar_check_impl : std::false_type {};
 
template <typename T>
struct is_scalar_check_impl<check_cl_<scalar_<T>>> : std::true_type {};
 
template <typename T>
using is_scalar_check = is_scalar_check_impl<std::decay_t<T>>;
}  // namespace internal
 
template <typename T, typename = require_all_kernel_expressions_t<T>>
inline auto check_cl(const char* function, const char* var_name, T&& y,
                     const char* must_be) {
  return check_cl_<as_operation_cl_t<T>>(
      function, var_name, as_operation_cl(std::forward<T>(y)), must_be);
}
 
}  // namespace math
}  // namespace stan
 
#endif
#endif  // CHECK_CL_HPP