select.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_SELECT_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_SELECT_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/fun/select.hpp>
#include <stan/math/opencl/matrix_cl_view.hpp>
#include <stan/math/opencl/kernel_generator/type_str.hpp>
#include <stan/math/opencl/kernel_generator/name_generator.hpp>
#include <stan/math/opencl/kernel_generator/operation_cl.hpp>
#include <stan/math/opencl/kernel_generator/as_operation_cl.hpp>
#include <stan/math/opencl/kernel_generator/common_return_scalar.hpp>
#include <algorithm>
#include <set>
#include <string>
#include <type_traits>
#include <utility>
 
namespace stan {
namespace math {
 
template <typename T_condition, typename T_then, typename T_else>
class select_ : public operation_cl<select_<T_condition, T_then, T_else>,
                                    common_scalar_t<T_then, T_else>,
                                    T_condition, T_then, T_else> {
 public:
  using Scalar = common_scalar_t<T_then, T_else>;
  using base = operation_cl<select_<T_condition, T_then, T_else>, Scalar,
                            T_condition, T_then, T_else>;
  using base::var_name_;
 
  select_(T_condition&& condition, T_then&& then, T_else&& els)  // NOLINT
      : base(std::forward<T_condition>(condition), std::forward<T_then>(then),
             std::forward<T_else>(els)) {
    if (condition.rows() != base::dynamic && then.rows() != base::dynamic) {
      check_size_match("select", "Rows of ", "condition", condition.rows(),
                       "rows of ", "then", then.rows());
    }
    if (condition.cols() != base::dynamic && then.cols() != base::dynamic) {
      check_size_match("select", "Columns of ", "condition", condition.cols(),
                       "columns of ", "then", then.cols());
    }
 
    if (condition.rows() != base::dynamic && els.rows() != base::dynamic) {
      check_size_match("select", "Rows of ", "condition", condition.rows(),
                       "rows of ", "else", els.rows());
    }
    if (condition.cols() != base::dynamic && els.cols() != base::dynamic) {
      check_size_match("select", "Columns of ", "condition", condition.cols(),
                       "columns of ", "else", els.cols());
    }
  }
 
  inline auto deep_copy() const {
    auto&& condition_copy = this->template get_arg<0>().deep_copy();
    auto&& then_copy = this->template get_arg<1>().deep_copy();
    auto&& else_copy = this->template get_arg<2>().deep_copy();
    return select_<std::remove_reference_t<decltype(condition_copy)>,
                   std::remove_reference_t<decltype(then_copy)>,
                   std::remove_reference_t<decltype(else_copy)>>(
        std::move(condition_copy), std::move(then_copy), std::move(else_copy));
  }
 
  inline kernel_parts generate(const std::string& row_index_name,
                               const std::string& col_index_name,
                               const bool view_handled,
                               const std::string& var_name_condition,
                               const std::string& var_name_then,
                               const std::string& var_name_else) const {
    kernel_parts res{};
    res.body = type_str<Scalar>() + " " + var_name_ + " = " + var_name_condition
               + " ? " + var_name_then + " : " + var_name_else + ";\n";
    return res;
  }
 
  inline std::pair<int, int> extreme_diagonals() const {
    using std::max;
    using std::min;
    std::pair<int, int> condition_diags
        = this->template get_arg<0>().extreme_diagonals();
    std::pair<int, int> then_diags
        = this->template get_arg<1>().extreme_diagonals();
    std::pair<int, int> else_diags
        = this->template get_arg<2>().extreme_diagonals();
    // Where the condition is 0 we get else's values. Otherwise we get the more
    // extreme of then's and else's.
    return {max(min(then_diags.first, else_diags.first),
                min(condition_diags.first, else_diags.first)),
            min(max(then_diags.second, else_diags.second),
                max(condition_diags.second, else_diags.second))};
  }
};
 
template <
    typename T_condition, typename T_then, typename T_else,
    require_all_kernel_expressions_t<T_condition, T_then, T_else>* = nullptr,
    require_any_not_arithmetic_t<T_condition, T_then, T_else>* = nullptr>
inline select_<as_operation_cl_t<T_condition>, as_operation_cl_t<T_then>,
               as_operation_cl_t<T_else>>
select(T_condition&& condition, T_then&& then, T_else&& els) {  // NOLINT
  return {as_operation_cl(std::forward<T_condition>(condition)),
          as_operation_cl(std::forward<T_then>(then)),
          as_operation_cl(std::forward<T_else>(els))};
}
 
}  // namespace math
}  // namespace stan
#endif
#endif