broadcast.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_BROADCAST_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_BROADCAST_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 <set>
#include <utility>
 
namespace stan {
namespace math {
 
template <typename T, bool Colwise, bool Rowwise>
class broadcast_
    : public operation_cl<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<broadcast_<T, Colwise, Rowwise>, Scalar, T>;
  using base::var_name_;
 
  explicit broadcast_(T&& a) : base(std::forward<T>(a)) {
    const char* function = "broadcast";
    if (Colwise) {
      check_size_match(function, "Rows of ", "a", a.rows(), "", "", 1);
    }
    if (Rowwise) {
      check_size_match(function, "Columns of ", "a", a.cols(), "", "", 1);
    }
  }
 
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return broadcast_<std::remove_reference_t<decltype(arg_copy)>, Colwise,
                      Rowwise>{std::move(arg_copy)};
  }
 
  inline void modify_argument_indices(std::string& row_index_name,
                                      std::string& col_index_name) const {
    if (Colwise) {
      row_index_name = "0";
    }
    if (Rowwise) {
      col_index_name = "0";
    }
  }
 
  inline int rows() const {
    return Colwise ? base::dynamic : this->template get_arg<0>().rows();
  }
 
  inline int cols() const {
    return Rowwise ? base::dynamic : this->template get_arg<0>().cols();
  }
 
  inline std::pair<int, int> extreme_diagonals() const {
    std::pair<int, int> arg_diags
        = this->template get_arg<0>().extreme_diagonals();
    return {Colwise ? std::numeric_limits<int>::min() : arg_diags.first,
            Rowwise ? std::numeric_limits<int>::max() : arg_diags.second};
  }
};
 
template <bool Colwise, bool Rowwise, typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto broadcast(T&& a) {
  auto&& a_operation = as_operation_cl(std::forward<T>(a)).deep_copy();
  return broadcast_<std::remove_reference_t<decltype(a_operation)>, Colwise,
                    Rowwise>(std::move(a_operation));
}
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto rowwise_broadcast(T&& a) {
  return broadcast<false, true>(std::forward<T>(a));
}
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto colwise_broadcast(T&& a) {
  return broadcast<true, false>(std::forward<T>(a));
}
}  // namespace math
}  // namespace stan
#endif
#endif