Future Request: Using arena_t in the model class

Wanted to make a note before I forgot. Once we can pass around expressions in Stan math we'll also be able to pass around data and variables that we know are already allocated by adding arena_matrix types to the model class and it's functions

I'll use the model below that gives this c++ code

data {
  int N;
  int M;
  vector[N] y;
  matrix[N, M] x;
}
parameters {
  vector[M] betas;
  real alpha;
  real<lower=0> sigma;
}
model {
  target += normal_id_glm_lpdf(y | x, alpha, betas, sigma);
}

So we have this new meta type in stan math now called arena_t<EigenType> that makes an eigen matrix (or vector) whose dynamic memory is fully allocated by our stack allocator. This is really nice because, since we know all the memory was created by our allocator, we can guarantee it will exist until we manually clean up that memory. So we can do a shallow copy (copy a pointer) vs a deep copy (allocate a whole new matrix for the reverse pass)

for example for the function below

template <typename VarMatrix, typename ArithMatrix>
inline auto some_func(const VarMatrix& a, const ArithMatrix& b) {
  // if arena_t is input, do shallow copies (copy pointer) 
  //   else make a deep copy
  arena_t<VarMat> arena_a(a);
  arena_t<ArithMatrix> arena_b(b);
  arena_t<Eigen::MatrixXd> ret(a.val() + b);
  // add a callback for computing the reverse pass
  reverse_pass_callback(
    [ret, arena_a, arena_b]() mutable {
      arena_a.adj() += ret.adj() + arena_b; 
  });
  return ret;
}

In the above, if the input a or b is an arena_t<matrix> then when we make a copy for the reverse pass we can do a shallow copy instead of a deep copy for either/or.

There's two ideas we can use this with

1. Inside of log_prob we can read directly from the reader for types defined in parameters and instead of assigning them to an eigen matrix we assign them to an arena_t<matrix>. For instance the below

    stan::io::reader<local_scalar_t__> in__(params_r__, params_i__);
    local_scalar_t__ DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());
    (void)DUMMY_VAR__;  // suppress unused var warning

    try {
     // HERE
      Eigen::Matrix<local_scalar_t__, -1, 1> betas;
      betas = Eigen::Matrix<local_scalar_t__, -1, 1>(M);
      stan::math::fill(betas, DUMMY_VAR__);

      current_statement__ = 1;
      betas = in__.vector(M);

Becomes

    stan::io::reader<local_scalar_t__> in__(params_r__, params_i__);
    local_scalar_t__ DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());
    (void)DUMMY_VAR__;  // suppress unused var warning

    try {
      arena_t<Eigen::Matrix<local_scalar_t__, -1, 1>> betas = in__.vector(M);

What's nice here is that later when we call add or any function in stan math that uses arena_t we can just do a shallow copy instead of a deep copy for those var matrices

2. We can do this with data if we are a little clever

Right now for the data in the example model we have

  Eigen::Matrix<double, -1, 1> y;
  Eigen::Matrix<double, -1, -1> x;
  // ...
  // in model constructor
  y = Eigen::Matrix<double, -1, 1>(N);

But if

A. We make these arena_t like

  // in class
  arena_t<Eigen::Matrix<double, -1, 1>> y;
  arena_t<Eigen::Matrix<double, -1, -1>> x;

B. Give the model it's own stack allocator

    stan::math::stack_alloc memalloc_;

Then when we make x and y we can do

      y = arena_t<Eigen::Matrix<double, -1, 1>>(N, memalloc_);

y will now also be allocated in a stack allocator (and same for x). So when we call some_func(beta, y) in log_prob instead of having to make any deep copies we can just have two shallow copies. This will dramatically cut out the amount of memory we have to allocate, essentially down to temporaries and function returns.

(1) can be done as soon as we start returning expressions. We can't yet do (2) in the math library but if the compiler team thinks (2) is possible then we can add the stuff we need in Stan math.

The issue names no specific files, tests, or entry points. It describes a future design involving generated model classes, Stan Math arena_t support, expression returns, and a model-owned stack allocator; first establish whether the prerequisite compiler and math-library work is possible, then define implementation scope and validation criteria.