Add Eigen Dense and Sparse vari_value types

stan/math/opencl/copy.hpp

@@ -11,6 +11,7 @@
 #include <stan/math/opencl/matrix_cl.hpp>
 #include <stan/math/opencl/matrix_cl_view.hpp>
 #include <stan/math/opencl/opencl_context.hpp>
+#include <stan/math/opencl/value_type.hpp>
 #include <stan/math/opencl/kernels/copy.hpp>
 #include <stan/math/opencl/kernels/pack.hpp>
 #include <stan/math/opencl/kernels/unpack.hpp>
@@ -101,6 +102,23 @@ inline Eigen::Matrix<T, R, C> from_matrix_cl(const matrix_cl<T>& src) {
   return dst;
 }
 
+/** \ingroup opencl
+ * Copies result of a kernel generator expression to the
+ * destination Eigen matrix.
+ *
+ * @tparam R rows type of the destination
+ * @tparam C cols type of the destination
+ * @tparam T type of expression
+ * @param src source expression
+ * @return Eigen matrix with a copy of the data in the source matrix
+ */
+template <int R = Eigen::Dynamic, int C = Eigen::Dynamic, typename T,
+          require_all_kernel_expressions_t<T>* = nullptr,
+          require_not_matrix_cl_t<T>* = nullptr>
+inline Eigen::Matrix<value_type_t<T>, R, C> from_matrix_cl(const T& src) {
+  return from_matrix_cl<R, C>(src.eval());
+}
+
 /** \ingroup opencl
  * Packs the flat triangular matrix on the OpenCL device and
  * copies it to the std::vector.

stan/math/opencl/kernel_generator.hpp

@@ -114,6 +114,7 @@
 
 #include <stan/math/opencl/kernel_generator/load.hpp>
 #include <stan/math/opencl/kernel_generator/scalar.hpp>
+#include <stan/math/opencl/kernel_generator/constant.hpp>
 #include <stan/math/opencl/kernel_generator/append.hpp>
 #include <stan/math/opencl/kernel_generator/binary_operation.hpp>
 #include <stan/math/opencl/kernel_generator/unary_function_cl.hpp>

stan/math/opencl/kernel_generator/constant.hpp

@@ -0,0 +1,132 @@
+#ifndef STAN_MATH_OPENCL_KERNEL_GENERATOR_CONSTANT_HPP
+#define STAN_MATH_OPENCL_KERNEL_GENERATOR_CONSTANT_HPP
+#ifdef STAN_OPENCL
+
+#include <stan/math/opencl/matrix_cl.hpp>
+#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.hpp>
+#include <limits>
+#include <set>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+
+namespace stan {
+namespace math {
+
+/** \addtogroup opencl_kernel_generator
+ *  @{
+ */
+
+/**
+ * Represents a matrix of single repeated value in kernel generator expressions.
+ * @tparam T type of the scalar
+ */
+template <typename T>
+class constant_ : public operation_cl<constant_<T>, T> {
+  T a_;
+  int rows_;
+  int cols_;
+
+ public:
+  static_assert(std::is_arithmetic<T>::value,
+                "class scalar_<T>: std::is_arithmetic<T> must be true!");
+  using Scalar = T;
+  using base = operation_cl<constant_<T>, T>;
+  using base::var_name_;
+
+  /**
+   * Constructor
+   * @param a scalar value
+   * @param rows number of rows of the matrix
+   * @param cols number of columns of the matrix
+   */
+  explicit constant_(const T a, int rows, int cols)
+      : a_(a), rows_(rows), cols_(cols) {}
+
+  /**
+   * Creates a deep copy of this expression.
+   * @return copy of \c *this
+   */
+  inline constant_<T> deep_copy() const {
+    return constant_<T>(a_, rows_, cols_);
+  }
+
+  /**
+   * Generates kernel code for this expression.
+   * @param row_index_name row index variable name
+   * @param col_index_name column index variable name
+   * @param view_handled whether whether caller already handled matrix view
+   * @return part of kernel with code for this expression
+   */
+  inline kernel_parts generate(const std::string& row_index_name,
+                               const std::string& col_index_name,
+                               const bool view_handled) const {
+    kernel_parts res{};
+    res.args = type_str<Scalar>() + " " + var_name_ + ", ";
+    return res;
+  }
+
+  /**
+   * Sets kernel arguments for this and nested expressions.
+   * @param[in,out] generated set of expressions that already set their kernel
+   * arguments
+   * @param kernel kernel to set arguments on
+   * @param[in,out] arg_num consecutive number of the first argument to set.
+   * This is incremented for each argument set by this function.
+   */
+  inline void set_args(std::set<const operation_cl_base*>& generated,
+                       cl::Kernel& kernel, int& arg_num) const {
+    kernel.setArg(arg_num++, a_);
+  }
+
+  /**
+   * Number of rows of a matrix that would be the result of evaluating this
+   * expression.
+   * @return number of rows
+   */
+  inline int rows() const { return rows_; }
+
+  /**
+   * Number of columns of a matrix that would be the result of evaluating this
+   * expression.
+   * @return number of columns
+   */
+  inline int cols() const { return cols_; }
+
+  /**
+   * Determine indices of extreme sub- and superdiagonals written.
+   * @return pair of indices - bottom and top diagonal
+   */
+  inline std::pair<int, int> extreme_diagonals() const {
+    return {std::numeric_limits<int>::min(), std::numeric_limits<int>::max()};
+  }
+};
+
+/**
+ * Matrix of repeated values in kernel generator expressions.
+ *
+ * In most cases scalars should be directly used instead of this. This is,
+ * however, useful for initializing some expression to specific value if that
+ * expresssion could also be plain `matrix_cl`.
+ *
+ * @tparam T type of argument
+ * @param a input argument
+ * @param rows number of rows
+ * @param cols number of columns
+ * @return Block of given expression
+ */
+template <typename T, typename = require_arithmetic_t<T>>
+inline auto constant(const T a, int rows, int cols) {
+  return constant_<T>(a, rows, cols);
+}
+
+/** @}*/
+}  // namespace math
+}  // namespace stan
+
+#endif
+#endif

stan/math/opencl/kernel_generator/is_kernel_expression.hpp

@@ -7,16 +7,21 @@
 #include <type_traits>
 
 namespace stan {
-namespace math {
 
 /** \addtogroup opencl_kernel_generator
  *  @{
  */
+
 /**
- * Non-templated base of \c operation is needed for easy checking if something
- * is a subclass of \c operation.
+ * Non-templated base of `operation_cl` is needed for easy checking if something
+ * is a subclass of `operation_cl`.
  */
 class operation_cl_base {};
+/**
+ * Non-templated base of `operation_cl_lhs` is needed for easy checking if
+ * something is a subclass of `operation_cl_lhs`.
+ */
+class operation_cl_lhs_base {};
 
 /**
  * Determines whether a type is non-scalar type that is a valid kernel generator
@@ -26,7 +31,6 @@ template <typename T, typename = void>
 struct is_kernel_expression_and_not_scalar
     : bool_constant<std::is_base_of<operation_cl_base,
                                     std::remove_reference_t<T>>::value> {};
-
 template <typename T>
 struct is_kernel_expression_and_not_scalar<T, require_matrix_cl_t<T>>
     : std::true_type {};
@@ -57,8 +61,21 @@ using require_all_kernel_expressions_and_none_scalar_t
 template <typename... Types>
 using require_all_kernel_expressions_t
     = require_all_t<is_kernel_expression<Types>...>;
+
+/**
+ * Determines whether a type is an assignable kernel generator
+ * expression.
+ */
+template <typename T, typename = void>
+struct is_kernel_expression_lhs
+    : bool_constant<std::is_base_of<operation_cl_lhs_base,
+                                    std::remove_reference_t<T>>::value> {};
+template <typename T>
+struct is_kernel_expression_lhs<T, require_matrix_cl_t<T>> : std::true_type {};
+
 /** @}*/
-}  // namespace math
+STAN_ADD_REQUIRE_UNARY(kernel_expression_lhs, is_kernel_expression_lhs,
+                       opencl_kernel_generator);
 }  // namespace stan
 
 #endif

stan/math/opencl/kernel_generator/operation_cl.hpp

@@ -314,6 +314,8 @@ class operation_cl : public operation_cl_base {
    * @return number of rows
    */
   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()...});
@@ -326,6 +328,8 @@ class operation_cl : public operation_cl_base {
    * @return number of columns
    */
   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()...});
@@ -352,6 +356,9 @@ class operation_cl : public operation_cl_base {
    * @return pair of indices - bottom and top diagonal
    */
   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()...);

stan/math/opencl/kernel_generator/operation_cl_lhs.hpp

@@ -24,7 +24,8 @@ namespace math {
  * @tparam Args types of arguments to this operation
  */
 template <typename Derived, typename Scalar, typename... Args>
-class operation_cl_lhs : public operation_cl<Derived, Scalar, Args...> {
+class operation_cl_lhs : public operation_cl<Derived, Scalar, Args...>,
+                         public operation_cl_lhs_base {
  protected:
   using base = operation_cl<Derived, Scalar, Args...>;
   static constexpr int N = sizeof...(Args);

stan/math/opencl/kernel_generator/scalar.hpp

@@ -27,7 +27,7 @@ namespace math {
  */
 template <typename T>
 class scalar_ : public operation_cl<scalar_<T>, T> {
- private:
+ protected:
   T a_;
 
  public:

stan/math/opencl/rev/copy.hpp

@@ -16,6 +16,76 @@
 namespace stan {
 namespace math {
 
+namespace internal {
+template <typename T_arg_adj, require_eigen_t<T_arg_adj>* = nullptr>
+class op_copy_to_cl_vari final
+    : public vari_value<matrix_cl<value_type_t<T_arg_adj>>> {
+  T_arg_adj arg_adj_;
+
+ public:
+  template <typename T_arg_val, require_eigen_t<T_arg_val>* = nullptr,
+            require_vt_same<T_arg_val, T_arg_adj>* = nullptr>
+  explicit op_copy_to_cl_vari(const T_arg_val& val, T_arg_adj adj)
+      : vari_value<matrix_cl<value_type_t<T_arg_adj>>>(to_matrix_cl(val)),
+        arg_adj_(adj) {}
+
+  virtual void chain() {
+    arg_adj_ += from_matrix_cl<std::decay_t<T_arg_adj>::RowsAtCompileTime,
+                               std::decay_t<T_arg_adj>::ColsAtCompileTime>(
+        this->adj_);
+  }
+};
+}  // namespace internal
+
+/** \ingroup opencl
+ * Copies the source var containing Eigen matrices to destination var that has
+ * data stored on the OpenCL device.
+ *
+ * @tparam T type of the Eigen matrix
+ * @param a source Eigen matrix
+ * @return var with a copy of the data on the OpenCL device
+ */
+template <typename T>
+inline var_value<matrix_cl<value_type_t<T>>> to_matrix_cl(
+    const var_value<T>& a) {
+  return new internal::op_copy_to_cl_vari<decltype(a.vi_->adj_)>(a.val(),
+                                                                 a.vi_->adj_);
+}
+
+namespace internal {
+template <typename T, int Rows, int Cols,
+          require_all_kernel_expressions_t<T>* = nullptr>
+class op_copy_from_cl_vari final
+    : public vari_value<Eigen::Matrix<value_type_t<T>, Rows, Cols>> {
+  vari_value<T>& a_;
+
+ public:
+  explicit op_copy_from_cl_vari(vari_value<T>& a)
+      : vari_value<Eigen::Matrix<value_type_t<T>, Rows, Cols>>(
+            from_matrix_cl<Rows, Cols>(a.val_)),
+        a_(a) {}
+
+  virtual void chain() { a_.adj_ = a_.adj_ + to_matrix_cl(this->adj_); }
+};
+}  // namespace internal
+
+/** \ingroup opencl
+ * Copies the source var that has data stored on the OpenCL device to
+ * destination var containing Eigen matrices.
+ *
+ * @tparam Rows number of compile time rows of the destination matrix
+ * @tparam Rows number of compile time columns of the destination matrix
+ * @tparam T type of the matrix or expression on the OpenCL device
+ * @param a source matrix_cl or expression
+ * @return var with a copy of the data on the host
+ */
+template <int Rows = Eigen::Dynamic, int Cols = Eigen::Dynamic, typename T,
+          require_all_kernel_expressions_t<T>* = nullptr>
+inline var_value<Eigen::Matrix<value_type_t<T>, Rows, Cols>> from_matrix_cl(
+    const var_value<T>& a) {
+  return new internal::op_copy_from_cl_vari<T, Rows, Cols>(*a.vi_);
+}
+
 /** \ingroup opencl
  * Copies the source Eigen matrix of vars to
  * the destination matrix that is stored
@@ -26,19 +96,17 @@ namespace math {
  * @param src source Eigen matrix
  * @return matrix_cl with a copy of the data in the source matrix
  */
-template <typename T, int R, int C, require_var_t<T>* = nullptr>
-inline matrix_cl<T> to_matrix_cl(const Eigen::Matrix<T, R, C>& src) try {
-  matrix_cl<T> dst(src.rows(), src.cols());
-  if (src.size() == 0) {
-    return dst;
-  }
-  dst.val() = to_matrix_cl(src.val().eval());
-  dst.adj() = to_matrix_cl(src.adj().eval());
-  return dst;
-} catch (const cl::Error& e) {
-  check_opencl_error("copy var Eigen->(OpenCL)", e);
-  matrix_cl<T> dst(src.rows(), src.cols());
-  return dst;
+template <typename T, require_eigen_vt<is_var, T>* = nullptr>
+inline var_value<matrix_cl<value_type_t<value_type_t<T>>>> to_matrix_cl(
+    const T& src) {
+  // the matrix can go out of scope before chain() is called. So we store a map
+  // to the data
+  var* src_array
+      = ChainableStack::instance_->memalloc_.alloc_array<var>(src.size());
+  Eigen::Map<plain_type_t<T>> src_stacked(src_array, src.rows(), src.cols());
+  src_stacked = src;
+  return new internal::op_copy_to_cl_vari<decltype(src_stacked.adj())>(
+      src_stacked.val(), src_stacked.adj());
 }
 
 /** \ingroup opencl