as_column_vector_or_scalar.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_AS_COLUMN_VECTOR_OR_SCALAR_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_AS_COLUMN_VECTOR_OR_SCALAR_HPP
#ifdef STAN_OPENCL
 
#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_lhs.hpp>
#include <stan/math/opencl/kernel_generator/as_operation_cl.hpp>
#include <stan/math/opencl/kernel_generator/constant.hpp>
#include <stan/math/prim/err/check_vector.hpp>
#include <stan/math/prim/meta.hpp>
#include <algorithm>
#include <map>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace stan {
namespace math {
 
template <typename T>
class as_column_vector_or_scalar_
    : public operation_cl_lhs<as_column_vector_or_scalar_<T>,
                              typename std::remove_reference_t<T>::Scalar, T> {
 public:
  using Scalar = typename std::remove_reference_t<T>::Scalar;
  using base = operation_cl_lhs<as_column_vector_or_scalar_<T>, Scalar, T>;
  using base::var_name_;
  using base::operator=;
 
  explicit as_column_vector_or_scalar_(T&& a) : base(std::forward<T>(a)) {
    check_vector("as_column_vector_or_scalar", "a", a);
  }
 
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return as_column_vector_or_scalar_<
        std::remove_reference_t<decltype(arg_copy)>>{std::move(arg_copy)};
  }
 
  inline kernel_parts get_kernel_parts(
      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, bool view_handled) const {
    kernel_parts res{};
    if (generated.count(this) == 0) {
      this->var_name_ = name_gen.generate();
      generated[this] = "";
      std::string row_index_name_arg = row_index_name;
      std::string col_index_name_arg = col_index_name;
      modify_argument_indices(row_index_name_arg, col_index_name_arg);
      std::unordered_map<const void*, const char*> generated2;
      res = this->template get_arg<0>().get_kernel_parts(
          generated2, generated_all, name_gen, row_index_name_arg,
          col_index_name_arg, view_handled);
      kernel_parts my
          = this->generate(row_index_name, col_index_name, view_handled,
                           this->template get_arg<0>().var_name_);
      if (generated_all.count(this) == 0) {
        generated_all[this] = "";
      } else {
        my.args = "";
      }
      res += my;
      res.body = res.body_prefix + res.body;
      res.body_prefix = "";
    }
    return res;
  }
 
  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_arg) const {
    kernel_parts res;
    res.args = "int " + var_name_ + "_transpose, ";
    res.body
        = type_str<Scalar>() + " " + var_name_ + " = " + var_name_arg + ";\n";
    return res;
  }
 
  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 {
    if (generated.count(this) == 0) {
      generated[this] = "";
      this->var_name_ = name_gen.generate();
    }
    std::string row_index_name_arg = row_index_name;
    std::string col_index_name_arg = col_index_name;
    modify_argument_indices(row_index_name_arg, col_index_name_arg);
    std::unordered_map<const void*, const char*> generated2;
    kernel_parts res = this->template get_arg<0>().get_kernel_parts_lhs(
        generated2, generated_all, name_gen, row_index_name_arg,
        col_index_name_arg);
    res += this->derived().generate_lhs(row_index_name, col_index_name,
                                        this->template get_arg<0>().var_name_);
    if (generated_all.count(this) == 0) {
      generated_all[this] = "";
    } else {
      res.args = "";
    }
    return res;
  }
 
  inline kernel_parts generate_lhs(const std::string& i, const std::string& j,
                                   const std::string& var_name_arg) const {
    kernel_parts res;
    res.args = "int " + var_name_ + "_transpose, ";
    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 {
    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 (generated_all.count(this) == 0) {
        generated_all[this] = "";
        kernel.setArg(arg_num++,
                      static_cast<int>(this->template get_arg<0>().rows()
                                       < this->template get_arg<0>().cols()));
      }
    }
  }
 
  inline void modify_argument_indices(std::string& row_index_name,
                                      std::string& col_index_name) const {
    std::string row_index_name2 = "(" + var_name_ + "_transpose ? "
                                  + col_index_name + " : " + row_index_name
                                  + ")";
    col_index_name = "(" + var_name_ + "_transpose ? " + row_index_name + " : "
                     + col_index_name + ")";
    row_index_name = std::move(row_index_name2);
  }
 
  inline int rows() const {
    return std::max(this->template get_arg<0>().rows(),
                    this->template get_arg<0>().cols());
  }
 
  inline int cols() const {
    return std::min(this->template get_arg<0>().rows(),
                    this->template get_arg<0>().cols());
  }
 
  inline void set_view(int bottom_diagonal, int top_diagonal,
                       int bottom_zero_diagonal, int top_zero_diagonal) const {
    auto& arg = this->template get_arg<0>();
    if (arg.rows() >= arg.cols()) {
      arg.set_view(bottom_diagonal, top_diagonal, bottom_zero_diagonal,
                   top_zero_diagonal);
    } else {
      arg.set_view(top_diagonal, bottom_diagonal, top_zero_diagonal,
                   bottom_zero_diagonal);
    }
  }
 
  inline std::pair<int, int> extreme_diagonals() const {
    auto& arg = this->template get_arg<0>();
    std::pair<int, int> arg_diags = arg.extreme_diagonals();
    if (arg.rows() >= arg.cols()) {
      return arg_diags;
    } else {
      return {-arg_diags.second, -arg_diags.first};
    }
  }
 
  inline void check_assign_dimensions(int rows, int cols) const {
    // use a dummy expression with same number of rows and cols to simplify the
    // check
    check_vector("as_column_vector_or_scalar_::check_assign_dimensions()",
                 "expression assigned to as_column_vector_or_scalar",
                 constant(0, rows, cols));
    auto& arg = this->template get_arg<0>();
    int arg_rows = arg.rows();
    int arg_cols = arg.cols();
    if (arg_rows >= arg_cols) {
      check_size_match("block_.check_assign_dimensions", "Rows of ",
                       "check_assign_dimensions argument", arg_rows, "rows of ",
                       "expression", rows);
      check_size_match("block_.check_assign_dimensions", "Columns of ",
                       "check_assign_dimensions argument", arg_cols,
                       "columns of ", "expression", cols);
    } else {
      check_size_match("block_.check_assign_dimensions", "Columns of ",
                       "check_assign_dimensions argument", arg_cols, "rows of ",
                       "expression", rows);
      check_size_match("block_.check_assign_dimensions", "Rows of ",
                       "check_assign_dimensions argument", arg_rows,
                       "columns of ", "expression", cols);
    }
  }
};
 
template <typename T,
          require_all_kernel_expressions_and_none_scalar_t<T>* = nullptr>
inline auto as_column_vector_or_scalar(T&& a) {
  auto&& a_operation = as_operation_cl(std::forward<T>(a)).deep_copy();
  return as_column_vector_or_scalar_<
      std::remove_reference_t<decltype(a_operation)>>(std::move(a_operation));
}
}  // namespace math
}  // namespace stan
 
#endif
#endif