operation_cl.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_OPERATION_CL_HPP
#define STAN_MATH_OPENCL_KERNEL_GENERATOR_OPERATION_CL_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err/check_nonnegative.hpp>
#include <stan/math/opencl/kernel_generator/assignment_ops.hpp>
#include <stan/math/opencl/kernel_generator/type_str.hpp>
#include <stan/math/opencl/kernel_generator/name_generator.hpp>
#include <stan/math/opencl/matrix_cl_view.hpp>
#include <stan/math/opencl/matrix_cl.hpp>
#include <stan/math/opencl/kernel_cl.hpp>
#include <CL/opencl.hpp>
#include <algorithm>
#include <string>
#include <utility>
#include <tuple>
#include <map>
#include <array>
#include <numeric>
#include <vector>
 
namespace stan {
namespace math {
 
struct kernel_parts {
  std::string includes;  // any function definitions - as if they were included
                         // at the start of kernel source
  std::string declarations;    // declarations of any local variables
  std::string initialization;  // the code for initializations done by all
                               // threads, even if they have no work
  std::string
      body_prefix;   // the code that should be run at the start of the kernel
                     // body (before the code for arguments of an operation)
  std::string body;  // the body of the kernel - code executing operations
  std::string body_suffix;  // the code that should be run at the end of the
                            // kernel body
  std::string
      reduction_1d;  // the code for reductions within work group by all
                     // threads, even if they have no work. Run once per column.
  std::string
      reduction_2d;  // the code for reductions within work group by all
                     // threads, even if they have no work. Run only once.
  std::string args;  // kernel arguments
 
  kernel_parts operator+(const kernel_parts& other) {
    return {includes + other.includes,
            declarations += other.declarations,
            initialization + other.initialization,
            body_prefix + other.body_prefix,
            body + other.body,
            body_suffix + other.body_suffix,
            reduction_1d + other.reduction_1d,
            reduction_2d + other.reduction_2d,
            args + other.args};
  }
 
  kernel_parts& operator+=(const kernel_parts& other) {
    includes += other.includes;
    declarations += other.declarations;
    initialization += other.initialization;
    body_prefix += other.body_prefix;
    body += other.body;
    body_suffix += other.body_suffix;
    reduction_1d += other.reduction_1d;
    reduction_2d += other.reduction_2d;
    args += other.args;
    return *this;
  }
};
 
inline std::ostream& operator<<(std::ostream& os, kernel_parts& parts) {
  os << "args:" << std::endl;
  os << parts.args.substr(0, parts.args.size() - 2) << std::endl;
  os << "Decl:" << std::endl;
  os << parts.declarations << std::endl;
  os << "Init:" << std::endl;
  os << parts.initialization << std::endl;
  os << "body:" << std::endl;
  os << parts.body << std::endl;
  os << "body_suffix:" << std::endl;
  os << parts.body_suffix << std::endl;
  os << "reduction_1d:" << std::endl;
  os << parts.reduction_1d << std::endl;
  os << "reduction_2d:" << std::endl;
  os << parts.reduction_2d << std::endl;
  return os;
}
 
template <typename Derived, typename Scalar, typename... Args>
class operation_cl : public operation_cl_base {
  static_assert(
      conjunction<std::is_base_of<operation_cl_base,
                                  std::remove_reference_t<Args>>...>::value,
      "operation_cl: all arguments to operation must be operations!");
 
 protected:
  std::tuple<Args...> arguments_;
  mutable std::string var_name_;  // name of the variable that holds result of
                                  // this operation in the kernel
 
 public:
  inline Derived& derived() { return *static_cast<Derived*>(this); }
 
  inline const Derived& derived() const {
    return *static_cast<const Derived*>(this);
  }
 
  using Deriv = Derived;
  using ArgsTuple = std::tuple<Args...>;
  static const bool require_specific_local_size;
  // number of arguments this operation has
  static constexpr int N = sizeof...(Args);
  using view_transitivity = std::tuple<std::is_same<Args, void>...>;
  // value representing a not yet determined size
  static constexpr int dynamic = -1;
 
  template <size_t N>
  const auto& get_arg() const {
    return std::get<N>(arguments_);
  }
 
  explicit operation_cl(Args&&... arguments)
      : arguments_(std::forward<Args>(arguments)...) {}
 
  matrix_cl<Scalar> eval() const {
    int rows = derived().rows();
    int cols = derived().cols();
    const char* function = "operation_cl.eval()";
    if (rows < 0) {
      invalid_argument(function, "Number of rows of expression", rows,
                       " must be nonnegative, but is ",
                       " (broadcasted expressions can not be evaluated)");
    }
    if (cols < 0) {
      invalid_argument(function, "Number of columns of expression", cols,
                       " must be nonnegative, but is ",
                       " (broadcasted expressions can not be evaluated)");
    }
    matrix_cl<Scalar> res(rows, cols, derived().view());
    if (res.size() > 0) {
      this->evaluate_into(res);
    }
    return res;
  }
 
  template <typename T_lhs>
  inline void evaluate_into(T_lhs& lhs) const;
 
  template <typename T_lhs>
  inline std::string get_kernel_source_for_evaluating_into(
      const T_lhs& lhs) const;
 
  template <typename T_result>
  kernel_parts get_whole_kernel_parts(
      std::unordered_map<const void*, const char*>& generated,
      std::unordered_map<const void*, const char*>& generated_all,
      name_generator& ng, const std::string& row_index_name,
      const std::string& col_index_name, const T_result& result) const {
    kernel_parts parts = derived().get_kernel_parts(
        generated, generated_all, ng, row_index_name, col_index_name, false);
    kernel_parts out_parts = result.get_kernel_parts_lhs(
        generated, generated_all, ng, row_index_name, col_index_name);
    out_parts.body += assignment_op<T_result>() + derived().var_name_ + ";\n";
    parts += out_parts;
    return parts;
  }
 
  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;
      derived().modify_argument_indices(row_index_name_arg, col_index_name_arg);
      std::array<kernel_parts, N> args_parts = index_apply<N>([&](auto... Is) {
        std::unordered_map<const void*, const char*> generated2;
        return std::array<kernel_parts, N>{this->get_arg<Is>().get_kernel_parts(
            &Derived::modify_argument_indices
                    == &operation_cl::modify_argument_indices
                ? generated
                : generated2,
            generated_all, name_gen, row_index_name_arg, col_index_name_arg,
            view_handled
                && std::tuple_element_t<
                    Is, typename Deriv::view_transitivity>::value)...};
      });
      res = std::accumulate(args_parts.begin(), args_parts.end(),
                            kernel_parts{});
      kernel_parts my_part = index_apply<N>([&](auto... Is) {
        return this->derived().generate(row_index_name, col_index_name,
                                        view_handled,
                                        this->get_arg<Is>().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 {
    var_name_ = var_name_arg;
    return {};
  }
 
  inline void modify_argument_indices(std::string& row_index_name,
                                      std::string& col_index_name) const {}
 
  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] = "";
      // parameter pack expansion returns a comma-separated list of values,
      // which can not be used as an expression. We work around that by using
      // comma operator to get a list of ints, which we use to construct an
      // initializer_list from. Cast to voids avoids warnings about unused
      // expression.
      index_apply<N>([&](auto... Is) {
        std::unordered_map<const void*, const char*> generated2;
        static_cast<void>(std::initializer_list<int>{
            (this->get_arg<Is>().set_args(
                 &Derived::modify_argument_indices
                         == &operation_cl::modify_argument_indices
                     ? generated
                     : generated2,
                 generated_all, kernel, arg_num),
             0)...});
      });
    }
  }
 
  inline void add_read_event(cl::Event& e) const {
    index_apply<N>([&](auto... Is) {
      static_cast<void>(std::initializer_list<int>{
          (this->get_arg<Is>().add_read_event(e), 0)...});
    });
  }
 
  inline void get_write_events(std::vector<cl::Event>& events) const {
    index_apply<N>([&](auto... Is) {
      static_cast<void>(std::initializer_list<int>{
          (this->template get_arg<Is>().get_write_events(events), 0)...});
    });
  }
 
  inline int rows() const {
    static_assert(
        N > 0, "default rows does not work on expressions with no arguments!");
    return index_apply<N>([&](auto... Is) {
      // assuming all non-dynamic sizes match
      return std::max({this->get_arg<Is>().rows()...});
    });
  }
 
  inline int cols() const {
    static_assert(
        N > 0, "default cols does not work on expressions with no arguments!");
    return index_apply<N>([&](auto... Is) {
      // assuming all non-dynamic sizes match
      return std::max({this->get_arg<Is>().cols()...});
    });
  }
 
  inline int size() const { return derived().rows() * derived().cols(); }
 
  inline int thread_rows() const { return derived().rows(); }
 
  inline int thread_cols() const { return derived().cols(); }
 
  inline std::pair<int, int> extreme_diagonals() const {
    static_assert(N > 0,
                  "default extreme_diagonals does not work on expressions with "
                  "no arguments!");
    return index_apply<N>([&](auto... Is) {
      auto arg_diags
          = std::make_tuple(this->get_arg<Is>().extreme_diagonals()...);
      int bottom = std::min(
          std::initializer_list<int>({std::get<Is>(arg_diags).first...}));
      int top = std::max(
          std::initializer_list<int>({std::get<Is>(arg_diags).second...}));
      return std::make_pair(bottom, top);
    });
  }
 
  inline matrix_cl_view view() const {
    std::pair<int, int> diagonals = derived().extreme_diagonals();
    matrix_cl_view view;
    if (diagonals.first < 0) {
      view = matrix_cl_view::Lower;
    } else {
      view = matrix_cl_view::Diagonal;
    }
    if (diagonals.second > 0) {
      view = either(view, matrix_cl_view::Upper);
    }
    return view;
  }
 
  inline void get_unique_matrix_accesses(
      std::vector<int>& uids, std::unordered_map<const void*, int>& id_map,
      int& next_id) const {
    index_apply<N>([&](auto... Is) {
      static_cast<void>(std::initializer_list<int>{(
          this->get_arg<Is>().get_unique_matrix_accesses(uids, id_map, next_id),
          0)...});
    });
  }
};
 
template <typename Derived, typename Scalar, typename... Args>
const bool operation_cl<Derived, Scalar, Args...>::require_specific_local_size
    = std::max({false,
                std::decay_t<Args>::Deriv::require_specific_local_size...});
}  // namespace math
}  // namespace stan
 
#endif
#endif