optional_broadcast.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_OPTIONALBROADCAST_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_OPTIONALBROADCAST_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/matrix_cl_view.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/as_operation_cl.hpp>
#include <limits>
#include <string>
#include <type_traits>
#include <map>
#include <utility>
 
namespace stan {
namespace math {
 
template <typename T, bool Colwise, bool Rowwise>
class optional_broadcast_
    : public operation_cl<optional_broadcast_<T, Colwise, Rowwise>,
                          typename std::remove_reference_t<T>::Scalar, T> {
 public:
  using Scalar = typename std::remove_reference_t<T>::Scalar;
  using base
      = operation_cl<optional_broadcast_<T, Colwise, Rowwise>, Scalar, T>;
  using base::var_name_;
 
  explicit optional_broadcast_(T&& a) : base(std::forward<T>(a)) {}
 
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return optional_broadcast_<std::remove_reference_t<decltype(arg_copy)>,
                               Colwise, Rowwise>{std::move(arg_copy)};
  }
 
  inline kernel_parts generate(const std::string& row_idx_name,
                               const std::string& col_idx_name,
                               const bool view_handled,
                               const std::string& var_name_arg) const {
    kernel_parts res;
    res.body
        += type_str<Scalar>() + " " + var_name_ + " = " + var_name_arg + ";\n";
    if (Colwise) {
      res.args += "int " + var_name_ + "is_multirow, ";
    }
    if (Rowwise) {
      res.args += "int " + var_name_ + "is_multicol, ";
    }
    return res;
  }
 
  inline void modify_argument_indices(std::string& row_idx_name,
                                      std::string& col_idx_name) const {
    if (Colwise) {
      row_idx_name = "(" + row_idx_name + " * " + var_name_ + "is_multirow)";
    }
    if (Rowwise) {
      col_idx_name = "(" + col_idx_name + " * " + var_name_ + "is_multicol)";
    }
  }
 
  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 {
    if (generated.count(this) == 0) {
      generated[this] = "";
      std::unordered_map<const void*, const char*> generated2;
      this->template get_arg<0>().set_args(generated2, generated_all, kernel,
                                           arg_num);
      if (Colwise) {
        kernel.setArg(arg_num++, static_cast<int>(
                                     this->template get_arg<0>().rows() != 1));
      }
      if (Rowwise) {
        kernel.setArg(arg_num++, static_cast<int>(
                                     this->template get_arg<0>().cols() != 1));
      }
    }
  }
 
  inline int rows() const {
    return Colwise && this->template get_arg<0>().rows() == 1
               ? base::dynamic
               : this->template get_arg<0>().rows();
  }
 
  inline int cols() const {
    return Rowwise && this->template get_arg<0>().cols() == 1
               ? base::dynamic
               : this->template get_arg<0>().cols();
  }
 
  inline matrix_cl_view view() const {
    matrix_cl_view view = this->template get_arg<0>().view();
    if (Colwise && this->template get_arg<0>().rows() == 1) {
      view = either(view, matrix_cl_view::Lower);
    }
    if (Rowwise && this->template get_arg<0>().cols() == 1) {
      view = either(view, matrix_cl_view::Upper);
    }
    return view;
  }
 
  inline int bottom_diagonal() const {
    if (Colwise && this->template get_arg<0>().rows() == 1) {
      return std::numeric_limits<int>::min();
    } else {
      return this->template get_arg<0>().bottom_diagonal();
    }
  }
 
  inline int top_diagonal() const {
    if (Rowwise && this->template get_arg<0>().cols() == 1) {
      return std::numeric_limits<int>::max();
    } else {
      return this->template get_arg<0>().top_diagonal();
    }
  }
};
 
template <bool Colwise, bool Rowwise, typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline optional_broadcast_<as_operation_cl_t<T>, Colwise, Rowwise>
optional_broadcast(T&& a) {
  auto&& a_operation = as_operation_cl(std::forward<T>(a)).deep_copy();
  return optional_broadcast_<as_operation_cl_t<T>, Colwise, Rowwise>(
      std::move(a_operation));
}
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto rowwise_optional_broadcast(T&& a) {
  return optional_broadcast<false, true>(std::forward<T>(a));
}
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto colwise_optional_broadcast(T&& a) {
  return optional_broadcast<true, false>(std::forward<T>(a));
}
 
}  // namespace math
}  // namespace stan
 
#endif
#endif