rowwise_reduction.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_ROWWISE_REDUCTION_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_ROWWISE_REDUCTION_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/opencl/matrix_cl_view.hpp>
#include <stan/math/opencl/kernel_generator/as_operation_cl.hpp>
#include <stan/math/opencl/kernel_generator/broadcast.hpp>
#include <stan/math/opencl/kernel_generator/binary_operation.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/type_str.hpp>
#include <map>
#include <string>
#include <type_traits>
#include <utility>
 
namespace stan {
namespace math {
namespace internal {
 
template <typename Arg>
struct matvec_mul_opt {
  // in general the optimization is not possible
  enum { is_possible = 0 };
 
  static matrix_cl_view view(const Arg&) { return matrix_cl_view::Entire; }
 
  static kernel_parts get_kernel_parts(
      const Arg& a, 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) {
    return {};
  }
};
 
template <typename Mat, typename VecT>
struct matvec_mul_opt<elt_multiply_<Mat, broadcast_<VecT, true, false>>> {
  // if the argument of rowwise reduction is multiplication with a broadcast
  // vector we can do the optimization
  enum { is_possible = 1 };
  using Arg = elt_multiply_<Mat, broadcast_<VecT, true, false>>;
 
  static matrix_cl_view view(const Arg& a) {
    return a.template get_arg<1>().template get_arg<0>().view();
  }
 
  static kernel_parts get_kernel_parts(
      const Arg& mul, 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) {
    kernel_parts res{};
    if (generated.count(&mul) == 0) {
      mul.var_name_ = name_gen.generate();
      generated[&mul] = "";
 
      const auto& matrix = mul.template get_arg<0>();
      const auto& broadcast = mul.template get_arg<1>();
      res = matrix.get_kernel_parts(generated, generated_all, name_gen,
                                    row_index_name, col_index_name, true);
      if (generated.count(&broadcast) == 0) {
        broadcast.var_name_ = name_gen.generate();
        generated[&broadcast] = "";
 
        const auto& vec = broadcast.template get_arg<0>();
        std::string row_index_name_bc = row_index_name;
        std::string col_index_name_bc = col_index_name;
        broadcast.modify_argument_indices(row_index_name_bc, col_index_name_bc);
        res += vec.get_kernel_parts(generated, generated_all, name_gen,
                                    row_index_name_bc, col_index_name_bc, true);
        res += broadcast.generate(row_index_name, col_index_name, true,
                                  vec.var_name_);
      }
      res += mul.generate(row_index_name, col_index_name, true,
                          matrix.var_name_, broadcast.var_name_);
    }
    return res;
  }
};
 
}  // namespace internal
 
template <typename Derived, typename T, typename operation, bool PassZero>
class rowwise_reduction
    : public operation_cl<Derived, typename std::remove_reference_t<T>::Scalar,
                          T> {
 public:
  using T_no_ref = std::remove_reference_t<T>;
  using Scalar = typename T_no_ref::Scalar;
  using base = operation_cl<Derived, Scalar, T>;
  using base::var_name_;
 
 protected:
  std::string init_;
 
 public:
  using base::rows;
  explicit rowwise_reduction(T&& a, const std::string& init)
      : base(std::forward<T>(a)), init_(init) {}
 
  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::unordered_map<const void*, const char*> generated2;
      if (PassZero && internal::matvec_mul_opt<T_no_ref>::is_possible) {
        res = internal::matvec_mul_opt<T_no_ref>::get_kernel_parts(
            this->template get_arg<0>(), generated2, generated_all, name_gen,
            row_index_name, var_name_ + "_j");
      } else {
        res = this->template get_arg<0>().get_kernel_parts(
            generated2, generated_all, name_gen, row_index_name,
            var_name_ + "_j", view_handled || PassZero);
      }
      kernel_parts my_part
          = generate(row_index_name, col_index_name, view_handled,
                     this->template get_arg<0>().var_name_);
      res += my_part;
      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.body_prefix
        = type_str<Scalar>() + " " + var_name_ + " = " + init_ + ";\n";
    if (PassZero) {
      res.body_prefix += "int " + var_name_ + "_start = contains_nonzero("
                         + var_name_ + "_view, LOWER) ? 0 : " + row_index_name
                         + ";\n";
      if (internal::matvec_mul_opt<T_no_ref>::is_possible) {
        res.body_prefix += "int " + var_name_ + "_end_temp = contains_nonzero("
                           + var_name_ + "_view, UPPER) ? " + var_name_
                           + "_cols : min(" + var_name_ + "_cols, "
                           + row_index_name + " + 1);\n";
        res.body_prefix += "int " + var_name_ + "_end = contains_nonzero("
                           + var_name_ + "_vec_view, UPPER) ? " + var_name_
                           + "_end_temp : min(1, " + var_name_
                           + "_end_temp);\n";
      } else {
        res.body_prefix += "int " + var_name_ + "_end = contains_nonzero("
                           + var_name_ + "_view, UPPER) ? " + var_name_
                           + "_cols : min(" + var_name_ + "_cols, "
                           + row_index_name + " + 1);\n";
      }
      res.body_prefix += "for(int " + var_name_ + "_j = " + var_name_
                         + "_start; " + var_name_ + "_j < " + var_name_
                         + "_end; " + var_name_ + "_j++){\n";
    } else {
      res.body_prefix += "for(int " + var_name_ + "_j = 0; " + var_name_
                         + "_j < " + var_name_ + "_cols; " + var_name_
                         + "_j++){\n";
    }
    res.body += var_name_ + " = " + operation::generate(var_name_, var_name_arg)
                + ";\n}\n";
    res.args = "int " + var_name_ + "_view, int " + var_name_ + "_cols, ";
    if (PassZero && internal::matvec_mul_opt<T_no_ref>::is_possible) {
      res.args += "int " + var_name_ + "_vec_view, ";
    }
    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);
      kernel.setArg(arg_num++, this->template get_arg<0>().view());
      kernel.setArg(arg_num++, this->template get_arg<0>().cols());
      if (PassZero && internal::matvec_mul_opt<T>::is_possible) {
        kernel.setArg(arg_num++, internal::matvec_mul_opt<T_no_ref>::view(
                                     this->template get_arg<0>()));
      }
    }
  }
 
  inline int cols() const { return 1; }
 
  inline std::pair<int, int> extreme_diagonals() const {
    return {-rows() + 1, cols() - 1};
  }
};
 
struct sum_op {
  inline static std::string generate(const std::string& a,
                                     const std::string& b) {
    return a + " + " + b;
  }
};
 
template <typename T>
class rowwise_sum_
    : public rowwise_reduction<rowwise_sum_<T>, T, sum_op, true> {
  using base = rowwise_reduction<rowwise_sum_<T>, T, sum_op, true>;
  using base::arguments_;
 
 public:
  explicit rowwise_sum_(T&& a) : base(std::forward<T>(a), "0") {}
 
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return rowwise_sum_<std::remove_reference_t<decltype(arg_copy)>>(
        std::move(arg_copy));
  }
};
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto rowwise_sum(T&& a) {
  auto&& arg_copy = as_operation_cl(std::forward<T>(a)).deep_copy();
  return rowwise_sum_<std::remove_reference_t<decltype(arg_copy)>>(
      std::move(arg_copy));
}
 
struct prod_op {
  inline static std::string generate(const std::string& a,
                                     const std::string& b) {
    return a + " * " + b;
  }
};
 
template <typename T>
class rowwise_prod_
    : public rowwise_reduction<rowwise_prod_<T>, T, prod_op, false> {
  using base = rowwise_reduction<rowwise_prod_<T>, T, prod_op, false>;
  using base::arguments_;
 
 public:
  explicit rowwise_prod_(T&& a) : base(std::forward<T>(a), "1") {}
 
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return rowwise_prod_<std::remove_reference_t<decltype(arg_copy)>>(
        std::move(arg_copy));
  }
};
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto rowwise_prod(T&& a) {
  auto&& arg_copy = as_operation_cl(std::forward<T>(a)).deep_copy();
  return rowwise_prod_<std::remove_reference_t<decltype(arg_copy)>>(
      std::move(arg_copy));
}
 
template <typename T>
struct max_op {
  inline static std::string generate(const std::string& a,
                                     const std::string& b) {
    if (std::is_floating_point<T>()) {
      return "fmax(" + a + ", " + b + ")";
    }
    return "max(" + a + ", " + b + ")";
  }
 
  inline static std::string init() {
    if (std::is_floating_point<T>()) {
      return "-INFINITY";
    }
    return "INT_MIN";
  }
};
 
template <typename T>
class rowwise_max_
    : public rowwise_reduction<
          rowwise_max_<T>, T,
          max_op<typename std::remove_reference_t<T>::Scalar>, false> {
  using op = max_op<typename std::remove_reference_t<T>::Scalar>;
  using base = rowwise_reduction<rowwise_max_<T>, T, op, false>;
  using base::arguments_;
 
 public:
  explicit rowwise_max_(T&& a) : base(std::forward<T>(a), op::init()) {}
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return rowwise_max_<std::remove_reference_t<decltype(arg_copy)>>(
        std::move(arg_copy));
  }
};
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto rowwise_max(T&& a) {
  auto&& arg_copy = as_operation_cl(std::forward<T>(a)).deep_copy();
  return rowwise_max_<std::remove_reference_t<decltype(arg_copy)>>(
      std::move(arg_copy));
}
template <typename T>
struct min_op {
  inline static std::string generate(const std::string& a,
                                     const std::string& b) {
    if (std::is_floating_point<T>()) {
      return "fmin(" + a + ", " + b + ")";
    }
    return "min(" + a + ", " + b + ")";
  }
 
  inline static std::string init() {
    if (std::is_floating_point<T>()) {
      return "INFINITY";
    }
    return "INT_MAX";
  }
};
 
template <typename T>
class rowwise_min_
    : public rowwise_reduction<
          rowwise_min_<T>, T,
          min_op<typename std::remove_reference_t<T>::Scalar>, false> {
  using op = min_op<typename std::remove_reference_t<T>::Scalar>;
  using base = rowwise_reduction<rowwise_min_<T>, T, op, false>;
  using base::arguments_;
 
 public:
  explicit rowwise_min_(T&& a) : base(std::forward<T>(a), op::init()) {}
  inline auto deep_copy() const {
    auto&& arg_copy = this->template get_arg<0>().deep_copy();
    return rowwise_min_<std::remove_reference_t<decltype(arg_copy)>>(
        std::move(arg_copy));
  }
};
 
template <typename T,
          typename = require_all_kernel_expressions_and_none_scalar_t<T>>
inline auto rowwise_min(T&& a) {
  auto&& arg_copy = as_operation_cl(std::forward<T>(a)).deep_copy();
  return rowwise_min_<std::remove_reference_t<decltype(arg_copy)>>(
      std::move(arg_copy));
}
}  // namespace math
}  // namespace stan
 
#endif
#endif