adds tests for the compound assignments

SteveBronder · SteveBronder · commit bfcd2223b6e2 · 2023-09-13T12:27:22.000-04:00
diff --git a/stan/math/opencl/kernel_generator/as_operation_cl.hpp b/stan/math/opencl/kernel_generator/as_operation_cl.hpp
@@ -24,7 +24,7 @@ namespace math {
  * @param a an operation
  * @return operation
  */
-template <assignment_ops_cl AssignOp = assignment_ops_cl::equals, typename T_operation,
+template <assign_op_cl AssignOp = assign_op_cl::equals, typename T_operation,
           typename = std::enable_if_t<std::is_base_of<
               operation_cl_base, std::remove_reference_t<T_operation>>::value>>
 inline T_operation&& as_operation_cl(T_operation&& a) {
@@ -38,7 +38,7 @@ inline T_operation&& as_operation_cl(T_operation&& a) {
  * @param a scalar
  * @return \c scalar_ wrapping the input
  */
-template <assignment_ops_cl AssignOp = assignment_ops_cl::equals, typename T_scalar, typename = require_arithmetic_t<T_scalar>,
+template <assign_op_cl AssignOp = assign_op_cl::equals, typename T_scalar, typename = require_arithmetic_t<T_scalar>,
           require_not_same_t<T_scalar, bool>* = nullptr>
 inline scalar_<T_scalar> as_operation_cl(const T_scalar a) {
   return scalar_<T_scalar>(a);
@@ -51,7 +51,7 @@ inline scalar_<T_scalar> as_operation_cl(const T_scalar a) {
  * @param a scalar
  * @return \c scalar_<char> wrapping the input
  */
-template <assignment_ops_cl AssignOp = assignment_ops_cl::equals>
+template <assign_op_cl AssignOp = assign_op_cl::equals>
 inline scalar_<char> as_operation_cl(const bool a) { return scalar_<char>(a); }
 
 /**
@@ -61,7 +61,7 @@ inline scalar_<char> as_operation_cl(const bool a) { return scalar_<char>(a); }
  * @param a \c matrix_cl
  * @return \c load_ wrapping the input
  */
-template <assignment_ops_cl AssignOp = assignment_ops_cl::equals, typename T_matrix_cl,
+template <assign_op_cl AssignOp = assign_op_cl::equals, typename T_matrix_cl,
           typename = require_any_t<is_matrix_cl<T_matrix_cl>,
                                    is_arena_matrix_cl<T_matrix_cl>>>
 inline load_<T_matrix_cl, AssignOp> as_operation_cl(T_matrix_cl&& a) {
@@ -76,7 +76,7 @@ inline load_<T_matrix_cl, AssignOp> as_operation_cl(T_matrix_cl&& a) {
  * rvalue reference, the reference is removed, so that a variable of this type
  * actually stores the value.
  */
-template <typename T, assignment_ops_cl AssignOp = assignment_ops_cl::equals>
+template <typename T, assign_op_cl AssignOp = assign_op_cl::equals>
 using as_operation_cl_t = std::conditional_t<
     std::is_lvalue_reference<T>::value,
     decltype(as_operation_cl<AssignOp>(std::declval<T>())),
diff --git a/stan/math/opencl/kernel_generator/assignment_ops.hpp b/stan/math/opencl/kernel_generator/assignment_ops.hpp
@@ -9,60 +9,56 @@ namespace math {
 /**
  * Ops that decide the type of assignment for LHS operations
  */
-enum class assignment_ops_cl {equals, plus_equals, minus_equals, divide_equals};
+enum class assign_op_cl {equals, plus_equals, minus_equals, divide_equals, multiply_equals};
 
+namespace internal {
 /**
  * @param value A static constexpr const char* member for printing assignment ops
  */
-template <assignment_ops_cl assign_op>
-struct assignment_op_str;
+template <assign_op_cl assign_op>
+struct assignment_op_str_impl;
 
 template <>
-struct assignment_op_str<assignment_ops_cl::equals> {
+struct assignment_op_str_impl<assign_op_cl::equals> {
   static constexpr const char* value = " = ";
 };
 
 template <>
-struct assignment_op_str<assignment_ops_cl::plus_equals> {
+struct assignment_op_str_impl<assign_op_cl::plus_equals> {
   static constexpr const char* value = " += ";
 };
 
 template <>
-struct assignment_op_str<assignment_ops_cl::minus_equals> {
-  static constexpr const char* value = " *= ";
+struct assignment_op_str_impl<assign_op_cl::minus_equals> {
+  static constexpr const char* value = " -= ";
 };
 
 template <>
-struct assignment_op_str<assignment_ops_cl::divide_equals> {
+struct assignment_op_str_impl<assign_op_cl::divide_equals> {
   static constexpr const char* value = " /= ";
 };
 
+template <>
+struct assignment_op_str_impl<assign_op_cl::multiply_equals> {
+  static constexpr const char* value = " *= ";
+};
 
-namespace internal {
 template <typename, typename = void>
-struct has_assignment_op_str : std::false_type {};
+struct assignment_op_str : assignment_op_str_impl<assign_op_cl::equals> {};
 
 template <typename T>
-struct has_assignment_op_str<T, void_t<decltype(T::assignment_op)>> : std::true_type {};
+struct assignment_op_str<T, void_t<decltype(T::assignment_op)>> : assignment_op_str_impl<T::assignment_op> {};
 
 }  // namespace internal
 
-/**
- * @tparam T A type that does not have an `assignment_op` static constexpr member type 
- * @return A constexpr const char* equal to `" = "`
- */
-template <typename T, std::enable_if_t<!internal::has_assignment_op_str<std::decay_t<T>>::value>* = nullptr>
-inline constexpr const char* assignment_op() noexcept {
-  return " = ";
-}
 
 /**
  * @tparam T A type that has an `assignment_op` static constexpr member type
  * @return The types assignment op as a constexpr const char*
  */
-template <typename T, std::enable_if_t<internal::has_assignment_op_str<T>::value>* = nullptr>
+template <typename T>
 inline constexpr const char* assignment_op() noexcept {
-  return assignment_op_str<std::decay_t<T>::assignment_op>::value;
+  return internal::assignment_op_str<std::decay_t<T>>::value;
 }
 
 }
diff --git a/stan/math/opencl/kernel_generator/load.hpp b/stan/math/opencl/kernel_generator/load.hpp
@@ -27,7 +27,7 @@ namespace math {
  * @tparam T \c matrix_cl
  * @tparam AssignOp tells higher level operations whether the final operation should be an assignment or a type of compound assignment.
  */
-template <typename T, assignment_ops_cl AssignOp = assignment_ops_cl::equals>
+template <typename T, assign_op_cl AssignOp = assign_op_cl::equals>
 class load_
     : public operation_cl_lhs<load_<T, AssignOp>,
                               typename std::remove_reference_t<T>::type> {
@@ -36,7 +36,7 @@ class load_
 
  public:
 
-  static constexpr assignment_ops_cl assignment_op = AssignOp;
+  static constexpr assign_op_cl assignment_op = AssignOp;
   using Scalar = typename std::remove_reference_t<T>::type;
   using base = operation_cl<load_<T, AssignOp>, Scalar>;
   using base::var_name_;
diff --git a/stan/math/opencl/kernel_generator/multi_result_kernel.hpp b/stan/math/opencl/kernel_generator/multi_result_kernel.hpp
@@ -323,19 +323,19 @@ class results_cl {
     });
   }
 
-  /**
+    /**
    * Incrementing \c results_ object by \c expressions_cl object
    * executes the kernel that evaluates expressions and increments results by
    * those expressions.
    * @tparam T_expressions types of expressions
    * @param exprs expressions
    */
-  template <typename... T_expressions,
+  template <assign_op_cl AssignOp = assign_op_cl::plus_equals, typename... T_expressions,
             typename = std::enable_if_t<sizeof...(T_results)
                                         == sizeof...(T_expressions)>>
-  void operator+=(const expressions_cl<T_expressions...>& exprs) {
+  void compound_assignment_impl(const expressions_cl<T_expressions...>& exprs) {
     index_apply<sizeof...(T_expressions)>([this, &exprs](auto... Is) {
-      auto tmp = std::tuple_cat(make_assignment_pair<assignment_ops_cl::plus_equals>(
+      auto tmp = std::tuple_cat(make_assignment_pair<AssignOp>(
           std::get<Is>(results_), std::get<Is>(exprs.expressions_))...);
       index_apply<std::tuple_size<decltype(tmp)>::value>(
           [this, &tmp](auto... Is2) {
@@ -345,6 +345,20 @@ class results_cl {
     });
   }
 
+  /**
+   * Incrementing \c results_ object by \c expressions_cl object
+   * executes the kernel that evaluates expressions and increments results by
+   * those expressions.
+   * @tparam T_expressions types of expressions
+   * @param exprs expressions
+   */
+  template <typename... T_expressions,
+            typename = std::enable_if_t<sizeof...(T_results)
+                                        == sizeof...(T_expressions)>>
+  void operator+=(const expressions_cl<T_expressions...>& exprs) {
+    compound_assignment_impl<assign_op_cl::plus_equals>(exprs);
+  }
+
     /**
    * Incrementing \c results_ object by \c expressions_cl object
    * executes the kernel that evaluates expressions and increments results by
@@ -356,15 +370,35 @@ class results_cl {
             typename = std::enable_if_t<sizeof...(T_results)
                                         == sizeof...(T_expressions)>>
   void operator-=(const expressions_cl<T_expressions...>& exprs) {
-    index_apply<sizeof...(T_expressions)>([this, &exprs](auto... Is) {
-      auto tmp = std::tuple_cat(make_assignment_pair<assignment_ops_cl::minus_equals>(
-          std::get<Is>(results_), std::get<Is>(exprs.expressions_))...);
-      index_apply<std::tuple_size<decltype(tmp)>::value>(
-          [this, &tmp](auto... Is2) {
-            assignment_impl(std::make_tuple(std::make_pair(
-                std::get<Is2>(tmp).first, std::get<Is2>(tmp).second)...));
-          });
-    });
+    compound_assignment_impl<assign_op_cl::minus_equals>(exprs);
+  }
+
+    /**
+   * Incrementing \c results_ object by \c expressions_cl object
+   * executes the kernel that evaluates expressions and increments results by
+   * those expressions.
+   * @tparam T_expressions types of expressions
+   * @param exprs expressions
+   */
+  template <typename... T_expressions,
+            typename = std::enable_if_t<sizeof...(T_results)
+                                        == sizeof...(T_expressions)>>
+  void operator/=(const expressions_cl<T_expressions...>& exprs) {
+    compound_assignment_impl<assign_op_cl::divide_equals>(exprs);
+  }
+
+    /**
+   * Incrementing \c results_ object by \c expressions_cl object
+   * executes the kernel that evaluates expressions and increments results by
+   * those expressions.
+   * @tparam T_expressions types of expressions
+   * @param exprs expressions
+   */
+  template <typename... T_expressions,
+            typename = std::enable_if_t<sizeof...(T_results)
+                                        == sizeof...(T_expressions)>>
+  void operator*=(const expressions_cl<T_expressions...>& exprs) {
+    compound_assignment_impl<assign_op_cl::multiply_equals>(exprs);
   }
 
   /**
@@ -551,7 +585,7 @@ class results_cl {
    * @param expression expression
    * @return a tuple of pair of result and expression
    */
-  template <assignment_ops_cl AssignOp = assignment_ops_cl::equals, typename T_result, typename T_expression,
+  template <assign_op_cl AssignOp = assign_op_cl::equals, typename T_result, typename T_expression,
             require_all_not_t<is_without_output<T_expression>,
                               conjunction<internal::is_scalar_check<T_result>,
                                           std::is_arithmetic<std::decay_t<
@@ -570,7 +604,7 @@ class results_cl {
    * @param expression expression
    * @return a tuple of pair of result and expression
    */
-  template <assignment_ops_cl AssignOp = assignment_ops_cl::equals, typename T_result, typename T_expression,
+  template <assign_op_cl AssignOp = assign_op_cl::equals, typename T_result, typename T_expression,
             require_t<is_without_output<T_expression>>* = nullptr>
   static auto make_assignment_pair(T_result&& result,
                                    T_expression&& expression) {
@@ -584,7 +618,7 @@ class results_cl {
    * @param pass bool scalar
    * @return an empty tuple
    */
-  template <assignment_ops_cl AssignOp = assignment_ops_cl::equals, typename T_check, typename T_pass,
+  template <assign_op_cl AssignOp = assign_op_cl::equals, typename T_check, typename T_pass,
             require_t<internal::is_scalar_check<T_check>>* = nullptr,
             require_integral_t<T_pass>* = nullptr>
   static std::tuple<> make_assignment_pair(T_check&& result, T_pass&& pass) {
diff --git a/stan/math/opencl/rev/adjoint_results.hpp b/stan/math/opencl/rev/adjoint_results.hpp
@@ -41,7 +41,7 @@ class adjoint_results_cl : protected results_cl<T_results...> {
     index_apply<sizeof...(T_expressions)>([&](auto... Is) {
       auto scalars = std::tuple_cat(select_scalar_assignments(
           std::get<Is>(this->results_), std::get<Is>(exprs.expressions_))...);
-      auto nonscalars_tmp = std::tuple_cat(select_nonscalar_plusequals(
+      auto nonscalars_tmp = std::tuple_cat(select_nonscalar_assignments<assign_op_cl::plus_equals>(
           std::get<Is>(this->results_), std::get<Is>(exprs.expressions_))...);
 
       index_apply<std::tuple_size<decltype(nonscalars_tmp)>::value>(
@@ -55,7 +55,7 @@ class adjoint_results_cl : protected results_cl<T_results...> {
                   // evaluate all expressions
                   this->assignment_impl(std::tuple_cat(
                       nonscalars,
-                      this->template make_assignment_pair<assignment_ops_cl::plus_equals>(
+                      this->template make_assignment_pair<assign_op_cl::plus_equals>(
                           std::get<2>(std::get<Is_scal>(scalars)),
                           sum_2d(std::get<1>(std::get<Is_scal>(scalars))))...));
 
@@ -109,12 +109,12 @@ class adjoint_results_cl : protected results_cl<T_results...> {
    * @return pair of result and expression or empty tuple (if the result is
    * check or the expression is `calc_if<false,T>`.
    */
-  template <typename T_result, typename T_expression,
+  template <assign_op_cl AssignOp, typename T_result, typename T_expression,
             require_not_stan_scalar_t<T_result>* = nullptr,
             require_st_var<T_result>* = nullptr>
-  auto select_nonscalar_plusequals(T_result&& result,
+  auto select_nonscalar_assignments(T_result&& result,
                                     T_expression&& expression) {
-    return results_cl<T_results...>::template make_assignment_pair<assignment_ops_cl::plus_equals>(
+    return results_cl<T_results...>::template make_assignment_pair<AssignOp>(
         result.adj(), std::forward<T_expression>(expression));
   }
   /**
@@ -126,11 +126,11 @@ class adjoint_results_cl : protected results_cl<T_results...> {
    * @param expression expression
    * @return empty tuple
    */
-  template <
+  template <assign_op_cl AssignOp,
       typename T_result, typename T_expression,
       std::enable_if_t<is_stan_scalar<T_result>::value
                        || !is_var<scalar_type_t<T_result>>::value>* = nullptr>
-  auto select_nonscalar_plusequals(T_result&& result,
+  auto select_nonscalar_assignments(T_result&& result,
                                     T_expression&& expression) {
     return std::make_tuple();
   }
diff --git a/test/unit/math/opencl/kernel_generator/assignment_ops_test.cpp b/test/unit/math/opencl/kernel_generator/assignment_ops_test.cpp
@@ -8,20 +8,79 @@
 #include <gtest/gtest.h>
 #include <string>
 
-TEST(KernelGenerator, assign_ops) {
+TEST(KernelGenerator, plus_equals) {
   using stan::math::matrix_cl;
   using stan::math::var_value;
   using stan::math::var;
   using stan::math::to_matrix_cl;
+  using stan::math::from_matrix_cl;
   Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 10);
   Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 10);
   Eigen::MatrixXd C = Eigen::MatrixXd::Random(10, 10);
-  C += A + B;
   matrix_cl<double> A_cl = to_matrix_cl(A);
   matrix_cl<double> B_cl = to_matrix_cl(B);
   matrix_cl<double> C_cl = to_matrix_cl(C);
+  C += A + B;
   results(C_cl) += expressions(A_cl + B_cl);
+  Eigen::MatrixXd C_cl_host = from_matrix_cl(C_cl);
+  EXPECT_MATRIX_EQ(C_cl_host, C)
+}
+
+TEST(KernelGenerator, minus_equals) {
+  using stan::math::matrix_cl;
+  using stan::math::var_value;
+  using stan::math::var;
+  using stan::math::to_matrix_cl;
+  using stan::math::from_matrix_cl;
+  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 10);
+  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 10);
+  Eigen::MatrixXd C = Eigen::MatrixXd::Random(10, 10);
+  matrix_cl<double> A_cl = to_matrix_cl(A);
+  matrix_cl<double> B_cl = to_matrix_cl(B);
+  matrix_cl<double> C_cl = to_matrix_cl(C);
+  C -= A + B;
+  results(C_cl) -= expressions(A_cl + B_cl);
+  Eigen::MatrixXd C_cl_host = from_matrix_cl(C_cl);
+  EXPECT_MATRIX_EQ(C_cl_host, C)
+}
+
+TEST(KernelGenerator, divide_equals) {
+  using stan::math::matrix_cl;
+  using stan::math::var_value;
+  using stan::math::var;
+  using stan::math::to_matrix_cl;
+  using stan::math::from_matrix_cl;
+  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 10);
+  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 10);
+  Eigen::MatrixXd C = Eigen::MatrixXd::Random(10, 10);
+  matrix_cl<double> A_cl = to_matrix_cl(A);
+  matrix_cl<double> B_cl = to_matrix_cl(B);
+  matrix_cl<double> C_cl = to_matrix_cl(C);
+  C.array() /= A.array() + B.array();
+  results(C_cl) /= expressions(A_cl + B_cl);
+  Eigen::MatrixXd C_cl_host = from_matrix_cl(C_cl);
+  EXPECT_MATRIX_EQ(C_cl_host, C)
+}
+
+TEST(KernelGenerator, times_equals) {
+  using stan::math::matrix_cl;
+  using stan::math::var_value;
+  using stan::math::var;
+  using stan::math::to_matrix_cl;
+  using stan::math::from_matrix_cl;
+
+  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 10);
+  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 10);
+  Eigen::MatrixXd C = Eigen::MatrixXd::Random(10, 10);
+  matrix_cl<double> A_cl = to_matrix_cl(A);
+  matrix_cl<double> B_cl = to_matrix_cl(B);
+  matrix_cl<double> C_cl = to_matrix_cl(C);
+  C.array() *= A.array() + B.array();
+  results(C_cl) *= expressions(A_cl + B_cl);
+  Eigen::MatrixXd C_cl_host = from_matrix_cl(C_cl);
+  EXPECT_MATRIX_EQ(C_cl_host, C)
 }
 
 
+
 #endif
