Future Request: Using arena_t in the model class

Aperta
  • Qualcosa non va?
#736 21 commenti 2 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
15/100
Tipo di issue
Funzionalità
Chiarezza
Da chiarire
Stato di attività
Ferma
Stack tecnologico
cpp
Ambito
compilers

Direzione di ricerca

L’issue non indica file, test o punti di ingresso specifici. Descrive una progettazione futura che coinvolge classi di modelli generate, il supporto di Stan Math per arena_t, restituzioni di espressioni e un allocatore dello stack di proprietà del modello; occorre prima stabilire se è possibile svolgere il lavoro preliminare necessario sul compilatore e sulla libreria matematica, quindi definire l’ambito dell’implementazione e i criteri di validazione.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

feature performance

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

  1. 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.

Lingua principale
OCaml
Stelle
162
Fork
60
Merge medio
2g 22h
PR unite (30g)
14

Preparare l'ambiente

  • Nessun Dockerfile né file Docker Compose
  • Ha un modello di pull request
  • Nessuna guida per i contributori

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di stan-dev/stanc3

Tutte le issue di stan-dev/stanc3

Issue simili

Altre issue su Compilers

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.