Merge pull request #2971 from stan-dev/fix/sparse-vari-impl

Add arena_matrix specialization for sparse matrices

Author: andrjohns

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,294 @@ 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>
+  inline T* copy_vector(const T* ptr, Integral size) {
+    if (size == 0)
+      return nullptr;
+    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()),
+          copy_vector(
+              other.innerNonZeroPtr(),
+              other.innerNonZeroPtr() == nullptr ? 0 : other.innerSize())) {}
+
+  /**
+   * 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(), other.innernonZeroPtr()) {}
+  /**
+   * 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(), other.innerNonZeroPtr()) {}
+  /**
+   * 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(), other.innerNonZeroPtr()) {}
+
+  // 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()),
+                    const_cast<StorageIndex*>(other.innerNonZeroPtr()));
+    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,
+            require_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  inline void inplace_ops_impl(F&& f, Expr&& other) {
+    auto&& x = to_ref(other);
+    auto* val_ptr = (*this).valuePtr();
+    auto* x_val_ptr = x.valuePtr();
+    const auto non_zeros = (*this).nonZeros();
+    for (Eigen::Index i = 0; i < non_zeros; ++i) {
+      f(val_ptr[i], x_val_ptr[i]);
+    }
+  }
+
+  /**
+   * 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,
+            require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  inline 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>
+  inline 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>
+  inline void inplace_ops_impl(F&& f, T&& other) {
+    auto* val_ptr = (*this).valuePtr();
+    const auto non_zeros = (*this).nonZeros();
+    for (Eigen::Index i = 0; i < non_zeros; ++i) {
+      f(val_ptr[i], other);
+    }
+  }
+
+ public:
+  /**
+   * Inplace operator `+=`
+   * @note Caution! Inplace operators assume that either
+   *  1. The right hand side sparse matrix has the same sparsity 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 that 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>
+  inline 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 sparsity 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 that 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>
+  inline 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 trivially 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