Add arena_matrix specialization for sparse matrices

stan/math/rev/core/arena_matrix.hpp

@@ -4,19 +4,15 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
 #include <stan/math/rev/core/chainablestack.hpp>
+#include <stan/math/rev/core/var_value_fwd_declare.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
 
 namespace stan {
 namespace math {
 
-/**
- * Equivalent to `Eigen::Matrix`, except that the data is stored on AD stack.
- * That makes these objects triviali destructible and usable in `vari`s.
- *
- * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`
- * ...)
- */
 template <typename MatrixType>
-class arena_matrix : public Eigen::Map<MatrixType> {
+class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>>
+    : public Eigen::Map<MatrixType> {
  public:
   using Scalar = value_type_t<MatrixType>;
   using Base = Eigen::Map<MatrixType>;
@@ -139,6 +135,264 @@ class arena_matrix : public Eigen::Map<MatrixType> {
   }
 };
 
+template <typename MatrixType>
+class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
+    : public Eigen::Map<MatrixType> {
+ public:
+  using Scalar = value_type_t<MatrixType>;
+  using Base = Eigen::Map<MatrixType>;
+  using PlainObject = std::decay_t<MatrixType>;
+  using StorageIndex = typename PlainObject::StorageIndex;
+
+  /**
+   * Default constructor.
+   */
+  arena_matrix() : Base::Map(0, 0, 0, nullptr, nullptr, nullptr) {}
+  /**
+   * Constructor for CDC formatted data. Assumes compressed and does not own
+   * memory.
+   *
+   * Constructs a read-write Map to a sparse matrix of size rows x cols,
+   * containing nnz non-zero coefficients, stored as a sparse format as defined
+   * by the pointers outerIndexPtr, innerIndexPtr, and valuePtr. If the optional
+   * parameter innerNonZerosPtr is the null pointer, then a standard compressed
+   * format is assumed.
+   *
+   * @param rows Number of rows
+   * @param cols Number of columns
+   * @param nnz Number of nonzero items in matrix
+   * @param outerIndexPtr A pointer to the array of outer indices.
+   * @param innerIndexPtr A pointer to the array of inner indices.
+   * @param valuePtr A pointer to the array of values.
+   * @param innerNonZerosPtr A pointer to the array of the number of non zeros
+   * of the inner vectors.
+   *
+   */
+  arena_matrix(Eigen::Index rows, Eigen::Index cols, Eigen::Index nnz,
+               StorageIndex* outerIndexPtr, StorageIndex* innerIndexPtr,
+               Scalar* valuePtr, StorageIndex* innerNonZerosPtr = nullptr)
+      : Base::Map(rows, cols, nnz, outerIndexPtr, innerIndexPtr, valuePtr,
+                  innerNonZerosPtr) {}
+
+ private:
+  template <typename T, typename Integral>
+  auto* copy_vector(const T* ptr, Integral size) {
+    T* ret = ChainableStack::instance_->memalloc_.alloc_array<T>(size);
+    std::copy_n(ptr, size, ret);
+    return ret;
+  }
+
+ public:
+  /**
+   * Constructs `arena_matrix` from an `Eigen::SparseMatrix`.
+   * @param other Eigen Sparse Matrix class
+   */
+  template <typename T, require_same_t<T, PlainObject>* = nullptr>
+  arena_matrix(T&& other)  // NOLINT
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
+                  copy_vector(other.innerIndexPtr(), other.nonZeros()),
+                  copy_vector(other.valuePtr(), other.nonZeros())) {}
+
+  /**
+   * Constructs `arena_matrix` from an Eigen expression
+   * @param other An expression
+   */
+  template <typename S, require_convertible_t<S&, PlainObject>* = nullptr,
+            require_not_same_t<S, PlainObject>* = nullptr>
+  arena_matrix(S&& other)  // NOLINT
+      : arena_matrix(PlainObject(std::forward<S>(other))) {}
+
+  /**
+   * Constructs `arena_matrix` from an expression. This makes an assumption that
+   * any other `Eigen::Map` also contains memory allocated in the arena.
+   * @param other expression
+   */
+  arena_matrix(const Base& other)  // NOLINT
+      : Base(other) {}
+
+  /**
+   * Const copy constructor. No actual copy is performed
+   * @note Since the memory for the arena matrix sits in Stan's memory arena all
+   * copies/moves of arena matrices are shallow moves
+   * @param other matrix to copy from
+   */
+  arena_matrix(const arena_matrix<MatrixType>& other)
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  other.outerIndexPtr(), other.innerIndexPtr(),
+                  other.valuePtr()) {}
+  /**
+   * Move constructor.
+   * @note Since the memory for the arena matrix sits in Stan's memory arena all
+   * copies/moves of arena matrices are shallow moves
+   * @param other matrix to copy from
+   */
+  arena_matrix(arena_matrix<MatrixType>&& other)
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  other.outerIndexPtr(), other.innerIndexPtr(),
+                  other.valuePtr()) {}
+  /**
+   * Copy constructor. No actual copy is performed
+   * @note Since the memory for the arena matrix sits in Stan's memory arena all
+   * copies/moves of arena matrices are shallow moves
+   * @param other matrix to copy from
+   */
+  arena_matrix(arena_matrix<MatrixType>& other)
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  other.outerIndexPtr(), other.innerIndexPtr(),
+                  other.valuePtr()) {}
+
+  // without this using, compiler prefers combination of implicit construction
+  // and copy assignment to the inherited operator when assigned an expression
+  using Base::operator=;
+
+  /**
+   * Copy assignment operator.
+   * @tparam ArenaMatrix An `arena_matrix` type
+   * @param other matrix to copy from
+   * @return `*this`
+   */
+  template <typename ArenaMatrix,
+            require_same_t<ArenaMatrix, arena_matrix<MatrixType>>* = nullptr>
+  arena_matrix& operator=(ArenaMatrix&& other) {
+    // placement new changes what data map points to - there is no allocation
+    new (this) Base(other.rows(), other.cols(), other.nonZeros(),
+                    const_cast<StorageIndex*>(other.outerIndexPtr()),
+                    const_cast<StorageIndex*>(other.innerIndexPtr()),
+                    const_cast<Scalar*>(other.valuePtr()));
+    return *this;
+  }
+
+  /**
+   * Assignment operator for assigning an expression.
+   * @tparam Expr An expression that an `arena_matrix<MatrixType>` can be
+   * constructed from
+   * @param expr expression to evaluate into this
+   * @return `*this`
+   */
+  template <typename Expr,
+            require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  arena_matrix& operator=(Expr&& expr) {
+    *this = arena_matrix(std::forward<Expr>(expr));
+    return *this;
+  }
+
+ private:
+  /**
+   * inplace operations functor for `Sparse (.*)= Sparse`.
+   * @note This assumes that each matrix is of the same size and sparsity
+   * pattern.
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam Expr Type derived from `Eigen::SparseMatrixBase`
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename Expr,
+            require_convertible_t<Expr&, MatrixType>* = nullptr>
+  void inplace_ops_impl(F&& f, Expr&& other) {
+    auto&& x = to_ref(other);
+    for (int k = 0; k < (*this).outerSize(); ++k) {
+      typename Base::InnerIterator it(*this, k);
+      typename std::decay_t<Expr>::InnerIterator iz(x, k);
+      for (; static_cast<bool>(it) && static_cast<bool>(iz); ++it, ++iz) {
+        f(it.valueRef(), iz.value());
+      }
+    }
+  }
+
+  /**
+   * inplace operations functor for `Sparse (.*)= Dense`.
+   * @note This assumes the user intends to perform the inplace operation for
+   * the nonzero parts of `this`
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam Expr Type derived from `Eigen::DenseBase`
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename Expr,
+            require_not_convertible_t<Expr&, MatrixType>* = nullptr,
+            require_eigen_dense_base_t<Expr>* = nullptr>
+  void inplace_ops_impl(F&& f, Expr&& other) {
+    auto&& x = to_ref(other);
+    for (int k = 0; k < (*this).outerSize(); ++k) {
+      typename Base::InnerIterator it(*this, k);
+      for (; static_cast<bool>(it); ++it) {
+        f(it.valueRef(), x(it.row(), it.col()));
+      }
+    }
+  }
+
+  /**
+   * inplace operations functor for `Sparse (.*)= Scalar`.
+   * @note This assumes the user intends to perform the inplace operation for
+   * the nonzero parts of `this`
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam T A scalar type
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename T,
+            require_convertible_t<T&, Scalar>* = nullptr>
+  void inplace_ops_impl(F&& f, T&& other) {
+    for (int k = 0; k < (*this).outerSize(); ++k) {
+      typename Base::InnerIterator it(*this, k);
+      for (; static_cast<bool>(it); ++it) {
+        f(it.valueRef(), other);
+      }
+    }
+  }
+
+ public:
+  /**
+   * Inplace operator `+=`
+   * @note Caution! Inplace operators assume that either
+   *  1. The right hand side sparse matrix has the same sparcity pattern
+   *  2. You only intend to add a scalar or dense matrix coefficients to the
+   * nonzero values of `this` This is intended to be used within the reverse
+   * pass for accumulation of the adjoint and is built as such. Any other use
+   * case should be be sure the above assumptions are satisfied.
+   * @tparam T A type derived from `Eigen::SparseMatrixBase` or
+   * `Eigen::DenseMatrixBase` or a `Scalar`
+   * @param other value to be added inplace to the matrix.
+   */
+  template <typename T>
+  arena_matrix& operator+=(T&& other) {
+    inplace_ops_impl(
+        [](auto&& x, auto&& y) {
+          x += y;
+          return;
+        },
+        std::forward<T>(other));
+    return *this;
+  }
+
+  /**
+   * Inplace operator `+=`
+   * @note Caution!! Inplace operators assume that either
+   *  1. The right hand side sparse matrix has the same sparcity pattern
+   *  2. You only intend to add a scalar or dense matrix coefficients to the
+   * nonzero values of `this` This is intended to be used within the reverse
+   * pass for accumulation of the adjoint and is built as such. Any other use
+   * case should be be sure the above assumptions are satisfied.
+   * @tparam T A type derived from `Eigen::SparseMatrixBase` or
+   * `Eigen::DenseMatrixBase` or a `Scalar`
+   * @param other value to be added inplace to the matrix.
+   */
+  template <typename T>
+  arena_matrix& operator-=(T&& other) {
+    inplace_ops_impl(
+        [](auto&& x, auto&& y) {
+          x -= y;
+          return;
+        },
+        std::forward<T>(other));
+    return *this;
+  }
+};
+
 }  // namespace math
 }  // namespace stan
 

stan/math/rev/core/var_value_fwd_declare.hpp

@@ -6,6 +6,16 @@ namespace math {
 // forward declaration of var
 template <typename T, typename = void>
 class var_value;
+
+/**
+ * Equivalent to `Eigen::Matrix`, except that the data is stored on AD stack.
+ * That makes these objects triviali destructible and usable in `vari`s.
+ *
+ * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`
+ * ...)
+ */
+template <typename MatrixType, typename = void>
+class arena_matrix;
 }  // namespace math
 }  // namespace stan
 #endif

stan/math/rev/core/vari.hpp

@@ -807,11 +807,11 @@ class vari_value<T, require_all_t<is_plain_type<T>, is_eigen_dense_base<T>>>
  *
  */
 template <typename T>
-class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
-                                                      chainable_alloc {
+class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
  public:
   using PlainObject = plain_type_t<T>;  // Base type of Eigen class
   using value_type = PlainObject;       // vari's adj_ and val_ member type
+  using InnerIterator = typename arena_matrix<PlainObject>::InnerIterator;
   /**
    * Rows at compile time
    */
@@ -825,12 +825,12 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    * The adjoint of this variable, which is the partial derivative
    * of this variable with respect to the root variable.
    */
-  PlainObject adj_;
+  arena_matrix<PlainObject> adj_;
 
   /**
    * The value of this variable.
    */
-  const PlainObject val_;
+  arena_matrix<PlainObject> val_;
 
   /**
    * Construct a variable implementation from a value. The
@@ -847,8 +847,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    * @param x Value of the constructed variable.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  explicit vari_value(S&& x)
-      : chainable_alloc(), adj_(x), val_(std::forward<S>(x)) {
+  explicit vari_value(S&& x) : adj_(x), val_(std::forward<S>(x)) {
     this->set_zero_adjoint();
     ChainableStack::instance_->var_stack_.push_back(this);
   }
@@ -870,8 +869,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    * that its `chain()` method is not called.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  vari_value(S&& x, bool stacked)
-      : chainable_alloc(), adj_(x), val_(std::forward<S>(x)) {
+  vari_value(S&& x, bool stacked) : adj_(x), val_(std::forward<S>(x)) {
     this->set_zero_adjoint();
     if (stacked) {
       ChainableStack::instance_->var_stack_.push_back(this);
@@ -921,11 +919,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    * result with respect to itself to be 1.
    */
   inline void init_dependent() {
-    for (int k = 0; k < adj_.outerSize(); ++k) {
-      for (typename PlainObject::InnerIterator it(adj_, k); it; ++it) {
-        it.valueRef() = 1.0;
-      }
-    }
+    std::fill(adj_.valuePtr(), adj_.valuePtr() + adj_.nonZeros(), 1.0);
   }
 
   /**
@@ -934,11 +928,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    * example in a Jacobian calculation).
    */
   inline void set_zero_adjoint() noexcept final {
-    for (int k = 0; k < adj_.outerSize(); ++k) {
-      for (typename PlainObject::InnerIterator it(adj_, k); it; ++it) {
-        it.valueRef() = 0.0;
-      }
-    }
+    std::fill(adj_.valuePtr(), adj_.valuePtr() + adj_.nonZeros(), 0.0);
   }
 
   /**

stan/math/rev/meta/arena_type.hpp

@@ -6,13 +6,9 @@
 #include <stan/math/prim/meta/plain_type.hpp>
 #include <stan/math/rev/core/arena_allocator.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
+#include <stan/math/rev/core/var_value_fwd_declare.hpp>
 
 namespace stan {
-namespace math {
-// forward declaration
-template <typename MatrixType>
-class arena_matrix;
-}  // namespace math
 
 namespace internal {
 template <typename T, typename = void, typename = void>