binary_operation.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_BINARY_OPERATOR_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_BINARY_OPERATOR_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/matrix_cl_view.hpp>
#include <stan/math/opencl/err/check_opencl.hpp>
#include <stan/math/prim/meta.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/scalar.hpp>
#include <stan/math/opencl/kernel_generator/as_operation_cl.hpp>
#include <stan/math/opencl/kernel_generator/common_return_scalar.hpp>
#include <algorithm>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace stan {
namespace math {
 
template <typename Derived, typename T_res, typename T_a, typename T_b>
class binary_operation : public operation_cl<Derived, T_res, T_a, T_b> {
 public:
  using Scalar = T_res;
  using base = operation_cl<Derived, Scalar, T_a, T_b>;
  using base::var_name_;
 
 protected:
  std::string op_;
  using base::arguments_;
 
 public:
  binary_operation(T_a&& a, T_b&& b, const std::string& op)  // NOLINT
      : base(std::forward<T_a>(a), std::forward<T_b>(b)), op_(op) {
    const std::string function = "binary_operator" + op;
    if (a.rows() != base::dynamic && b.rows() != base::dynamic) {
      check_size_match(function.c_str(), "Rows of ", "a", a.rows(), "rows of ",
                       "b", b.rows());
    }
    if (a.cols() != base::dynamic && b.cols() != base::dynamic) {
      check_size_match(function.c_str(), "Columns of ", "a", a.cols(),
                       "columns of ", "b", b.cols());
    }
  }
 
  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_a,
                               const std::string& var_name_b) const {
    kernel_parts res{};
    res.body = type_str<Scalar>() + " " + var_name_ + " = " + var_name_a + " "
               + op_ + " " + var_name_b + ";\n";
    return res;
  }
};
 
#define COMMA ,
 
#define ADD_BINARY_OPERATION(class_name, function_name, scalar_type_expr,     \
                             operation)                                       \
  template <typename T_a, typename T_b>                                       \
  class class_name : public binary_operation<class_name<T_a, T_b>,            \
                                             scalar_type_expr, T_a, T_b> {    \
    using base                                                                \
        = binary_operation<class_name<T_a, T_b>, scalar_type_expr, T_a, T_b>; \
    using base::arguments_;                                                   \
                                                                              \
   public:                                                                    \
    using base::rows;                                                         \
    using base::cols;                                                         \
    class_name(T_a&& a, T_b&& b) /* NOLINT */                                 \
        : base(std::forward<T_a>(a), std::forward<T_b>(b), operation) {}      \
    inline auto deep_copy() const {                                           \
      auto&& a_copy = this->template get_arg<0>().deep_copy();                \
      auto&& b_copy = this->template get_arg<1>().deep_copy();                \
      return class_name<std::remove_reference_t<decltype(a_copy)>,            \
                        std::remove_reference_t<decltype(b_copy)>>(           \
          std::move(a_copy), std::move(b_copy));                              \
    }                                                                         \
  };                                                                          \
                                                                              \
  template <typename T_a, typename T_b,                                       \
            require_all_kernel_expressions_t<T_a, T_b>* = nullptr,            \
            require_any_not_arithmetic_t<T_a, T_b>* = nullptr>                \
  inline class_name<as_operation_cl_t<T_a>, as_operation_cl_t<T_b>>           \
  function_name(T_a&& a, T_b&& b) { /* NOLINT */                              \
    return {as_operation_cl(std::forward<T_a>(a)),                            \
            as_operation_cl(std::forward<T_b>(b))};                           \
  }
 
#define ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(                                \
    class_name, function_name, scalar_type_expr, operation, ...)              \
  template <typename T_a, typename T_b>                                       \
  class class_name : public binary_operation<class_name<T_a, T_b>,            \
                                             scalar_type_expr, T_a, T_b> {    \
    using base                                                                \
        = binary_operation<class_name<T_a, T_b>, scalar_type_expr, T_a, T_b>; \
    using base::arguments_;                                                   \
                                                                              \
   public:                                                                    \
    using base::rows;                                                         \
    using base::cols;                                                         \
    class_name(T_a&& a, T_b&& b) /* NOLINT */                                 \
        : base(std::forward<T_a>(a), std::forward<T_b>(b), operation) {}      \
    inline auto deep_copy() const {                                           \
      auto&& a_copy = this->template get_arg<0>().deep_copy();                \
      auto&& b_copy = this->template get_arg<1>().deep_copy();                \
      return class_name<std::remove_reference_t<decltype(a_copy)>,            \
                        std::remove_reference_t<decltype(b_copy)>>(           \
          std::move(a_copy), std::move(b_copy));                              \
    }                                                                         \
    __VA_ARGS__                                                               \
  };                                                                          \
                                                                              \
  template <typename T_a, typename T_b,                                       \
            require_all_kernel_expressions_t<T_a, T_b>* = nullptr,            \
            require_any_not_arithmetic_t<T_a, T_b>* = nullptr>                \
  inline class_name<as_operation_cl_t<T_a>, as_operation_cl_t<T_b>>           \
  function_name(T_a&& a, T_b&& b) { /* NOLINT */                              \
    return {as_operation_cl(std::forward<T_a>(a)),                            \
            as_operation_cl(std::forward<T_b>(b))};                           \
  }
 
ADD_BINARY_OPERATION(addition_operator_, operator+,
                     common_scalar_t<T_a COMMA T_b>, "+");
ADD_BINARY_OPERATION(addition_, add, common_scalar_t<T_a COMMA T_b>, "+");
ADD_BINARY_OPERATION(subtraction_operator_, operator-,
                     common_scalar_t<T_a COMMA T_b>, "-");
ADD_BINARY_OPERATION(subtraction_, subtract, common_scalar_t<T_a COMMA T_b>,
                     "-");
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    elt_multiply_, elt_multiply, common_scalar_t<T_a COMMA T_b>, "*",
    using view_transitivity = std::tuple<std::true_type, std::true_type>;
    inline std::pair<int, int> extreme_diagonals() const {
      std::pair<int, int> diags0
          = this->template get_arg<0>().extreme_diagonals();
      std::pair<int, int> diags1
          = this->template get_arg<1>().extreme_diagonals();
      return {std::max(diags0.first, diags1.first),
              std::min(diags0.second, diags1.second)};
    });
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    elt_divide_, elt_divide, common_scalar_t<T_a COMMA T_b>, "/",
    inline std::pair<int, int> extreme_diagonals() const {
      return {-rows() + 1, cols() - 1};
    });
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    elt_modulo_, operator%, common_scalar_t<T_a COMMA T_b>, "%",
    static_assert(
        std::is_integral<scalar_type_t<T_a>>::value&&
            std::is_integral<scalar_type_t<T_b>>::value,
        "both operands to operator% must have integral scalar types!");
    inline std::pair<int, int> extreme_diagonals() const {
      return {-rows() + 1, cols() - 1};
    });
 
ADD_BINARY_OPERATION(less_than_, operator<, bool, "<");
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    less_than_or_equal_, operator<=, bool,
    "<=", inline std::pair<int, int> extreme_diagonals() const {
      return {-rows() + 1, cols() - 1};
    });
ADD_BINARY_OPERATION(greater_than_, operator>, bool, ">");
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    greater_than_or_equal_, operator>=, bool,
    ">=", inline std::pair<int, int> extreme_diagonals() const {
      return {-rows() + 1, cols() - 1};
    });
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    equals_, operator==, bool,
    "==", inline std::pair<int, int> extreme_diagonals() const {
      return {-rows() + 1, cols() - 1};
    });
ADD_BINARY_OPERATION(not_equals_, operator!=, bool, "!=");
 
ADD_BINARY_OPERATION(logical_or_, operator||, bool, "||");
ADD_BINARY_OPERATION_WITH_CUSTOM_CODE(
    logical_and_, operator&&, bool, "&&",
    using view_transitivity = std::tuple<std::true_type, std::true_type>;
    inline std::pair<int, int> extreme_diagonals() const {
      std::pair<int, int> diags0
          = this->template get_arg<0>().extreme_diagonals();
      std::pair<int, int> diags1
          = this->template get_arg<1>().extreme_diagonals();
      return {std::max(diags0.first, diags1.first),
              std::min(diags0.second, diags1.second)};
    });
 
template <typename T_a, typename T_b, typename = require_arithmetic_t<T_a>,
          typename = require_all_kernel_expressions_t<T_b>>
inline elt_multiply_<as_operation_cl_t<T_a>, as_operation_cl_t<T_b>> operator*(
    T_a a, T_b&& b) {  // NOLINT
  return {as_operation_cl(a), as_operation_cl(std::forward<T_b>(b))};
}
 
template <typename T_a, typename T_b,
          typename = require_all_kernel_expressions_t<T_a>,
          typename = require_arithmetic_t<T_b>>
inline elt_multiply_<as_operation_cl_t<T_a>, as_operation_cl_t<T_b>> operator*(
    T_a&& a, const T_b b) {  // NOLINT
  return {as_operation_cl(std::forward<T_a>(a)), as_operation_cl(b)};
}
 
#undef COMMA
#undef ADD_BINARY_OPERATION
#undef ADD_BINARY_OPERATION_WITH_CUSTOM_CODE
}  // namespace math
}  // namespace stan
#endif
#endif