Future Request: Using arena_t in the model class

オープン
  • 何か問題がありますか？
#736 コメント 21 件 リアクション 2 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
5/5
見積もり時間
1週間以上
初心者へのやさしさ
15/100
issue の種類
機能追加
明瞭さ
説明が足りない
活発さ
停滞
技術スタック
cpp
領域
compilers

調査の方向性

この issue では、具体的なファイル、テスト、エントリーポイントは指定されていません。生成されたモデルクラス、Stan Math の arena_t サポート、式の返却、モデル所有のスタックアロケータを含む将来的な設計について説明しています。まず、前提となるコンパイラおよび数学ライブラリの作業が可能かどうかを確認し、その後、実装範囲と検証基準を定義します。

索引モデルが issue の本文から書いたものです。

説明

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.

主要言語
OCaml
スター
162
フォーク
60
平均マージ
2日 22時間
マージ済み PR（30日）
14

環境構築

  • Dockerfile・Docker Compose ファイルなし
  • プルリクエストのテンプレートあり
  • コントリビューションガイドなし

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

stan-dev/stanc3 のほかの issue

stan-dev/stanc3 の issue をすべて見る

似ている issue

Compilers の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。