(ソースコードメモ)PyTorchでのCPU側並列処理

#はじめに
PyTorchのCPU側の並列処理は、ATen/Parallelで主に行う。[CPUの並列処理の概要も文書](https://pytorch.org/docs/stable/notes/cpu_threading_torchscript_inference.html)に記載されている。現状の並列処理設定を確認するには、以下のスクリプトで確認できる。([PyTorch 1.2から以下のスクリプトで確認できる。](https://github.com/pytorch/pytorch/commit/481b6d026848b6d94edc0d84f32fb59c50590625))

```python
import torch
print(torch.__config__.parallel_info()) 
```
Google Colabで実行すると、以下のように表示される。この場合、OpenMP設定になっており、環境変数`OMP_NUM_THREADS`を設定すれば、CPUの負荷を制限できる。

```shell

ATen/Parallel:
	at::get_num_threads() : 1
	at::get_num_interop_threads() : 1
OpenMP 201511 (a.k.a. OpenMP 4.5)
	omp_get_max_threads() : 1
Intel(R) Math Kernel Library Version 2019.0.4 Product Build 20190411 for Intel(R) 64 architecture applications
	mkl_get_max_threads() : 1
Intel(R) MKL-DNN v0.20.5 (Git Hash 0125f28c61c1f822fd48570b4c1066f96fcb9b2e)
std::thread::hardware_concurrency() : 2
Environment variables:
	OMP_NUM_THREADS : [not set]
	MKL_NUM_THREADS : [not set]
ATen parallel backend: OpenMP
```

現在は、[OpenMPのみからATen/Parallelをベースとした`parallel_for/parallel_reduce`への書き換えが進行中](http://web.mit.edu/~ezyang/Public/pytorch-internals.pdf#page=43)である。ここでは、PyTorchの新規処理とともに、かつての[OpenMP](https://gcc.gnu.org/onlinedocs/libgomp/index.html)を使った並列処理について説明する。なお、PyTorch 1.3時点での[デフォルトは、OpenMPである](https://github.com/pytorch/pytorch/blob/v1.3.0/CMakeLists.txt#L162)。
PyTorchの処理は、データ処理演算と、データロード(DataLoader)に分かれる。データ処理演算で使われるATen/Parallelは、Pythonより下の演算処理であるため、一つのプロセスが数百%となる。そして、データローダは、num_workersで指定した数を、別プロセスとして起動している。

#PyTorch独自関数について
[`at::parallel_for`関数や`at::parallel_reduce`関数](https://github.com/pytorch/pytorch/blob/v1.3.0/aten/src/ATen/Parallel.h#L36)で並列処理が出来る。なお、この関数のバックエンドは、OpenMPだけでなく、NativeおよびThread Building Blocks(TBB)も提供されている。ただし、コンパイル時に設定するので、実行時に変更出来ない。[通常はOpenMP、モバイル系はNativeを使う設定](https://github.com/pytorch/pytorch/blob/v1.3.0/caffe2/CMakeLists.txt#L908)になる。
##at::parallel_forの利用例
parallel_forは、様々な関数で用いられている。例えば、[Adaptive Max Pooling](https://github.com/pytorch/pytorch/blob/v1.3.0/aten/src/ATen/native/AdaptiveMaxPooling2d.cpp#L39)では、以下のように定義されている。該当関数の引数は、begin, end, grain_size, fの4つであった。それぞれ、0, sizeD, 0, [&]が設定されている。ちなみに、[[&]は、ラムダ式のキャプチャ記法](https://cpprefjp.github.io/lang/cpp11/lambda_expressions.html)(Lambda capture)である。これにより、関数の前で設定された変数へのアクセスが内部で可能になる。

```c++

  at::parallel_for(0, sizeD, 0, [&](int64_t start, int64_t end) {
    for (auto d = start; d < end; d++)
    {
      /* loop over output */
      int64_t oh, ow;
      for(oh = 0; oh < osizeH; oh++)
      {
        int istartH = start_index(oh, osizeH, isizeH);
        int iendH   = end_index(oh, osizeH, isizeH);
        int kH = iendH - istartH;

        for(ow = 0; ow < osizeW; ow++)
        {
          int istartW = start_index(ow, osizeW, isizeW);
          int iendW   = end_index(ow, osizeW, isizeW);
          int kW = iendW - istartW;

          /* local pointers */
          scalar_t *ip = input_p   + d*istrideD + istartH*istrideH + istartW*istrideW;
          scalar_t *op = output_p  + d*osizeH*osizeW + oh*osizeW + ow;
          int64_t *indp = ind_p   + d*osizeH*osizeW + oh*osizeW + ow;

          /* compute local max: */
          int64_t maxindex = -1;
          scalar_t maxval = -std::numeric_limits<float>::max();
          int ih, iw;
          for(ih = 0; ih < kH; ih++)
          {
            for(iw = 0; iw < kW; iw++)
            {
              scalar_t val = *(ip + ih*istrideH + iw*istrideW);
              if ((val > maxval) || std::isnan(val))
              {
                maxval = val;
                maxindex = (ih+istartH)*isizeW + (iw+istartW);
              }
            }
          }

          /* set output to local max */
          *op = maxval;

          /* store location of max */
          *indp = maxindex;
        }
      }
    }
  });
```

# OpenMPについて
OpenMPは、[fork-joinモデル](https://www.openmp.org/wp-content/uploads/omp-hands-on-SC08.pdf#page=14)により、並列化している。[コンパイラディレクティブ](https://en.wikipedia.org/wiki/Directive_(programming))である。以下で示す`#pragma omp`で処理をコンパイラに指示しバイナリを生成する。gccの場合、実体の処理はgcc同梱のlibgompでおこなっている。
さて、OpenMPは、PyTorchでは、ATenの中のat::parallel_forのバックエンドや、従前のTHTensor直で用いている。

##OpenMPのフラグをオンにする
gccコンパイル時に、[`-fopenmp`で宣言を行う。](https://gcc.gnu.org/onlinedocs/libgomp/Enabling-OpenMP.html#Enabling-OpenMP)このフラグにより、`#pragma omp`の宣言が可能になる。

##OpenMPの使い方
ここでは、

- OpenMPの設定
- OpenMPの直接呼出し(コーディング)

について説明する。

###環境変数の設定
[デフォルトでは、コア数(もしくはハイパースレッドを考慮してx2)が指定される。しかし、必要なら、`OMP_NUM_THREADS`で設定することが出来る。](https://gcc.gnu.org/onlinedocs/libgomp/OMP_005fNUM_005fTHREADS.html#OMP_005fNUM_005fTHREADS)

###ソースコードでの記述
`#pragma omp`以下で定義できるパラメータは、libgompのマニュアルの[10 The libgomp ABI](https://gcc.gnu.org/onlinedocs/libgomp/The-libgomp-ABI.html#The-libgomp-ABI)に記載されている。また、OpenMPのマニュアルの[parallel構造体](https://www.openmp.org/spec-html/5.0/openmpse14.html)にも仕様として記載されている。

#### parallel
以下のように[並列処理を記述する](https://gcc.gnu.org/onlinedocs/libgomp/Implementing-PARALLEL-construct.html#Implementing-PARALLEL-construct)。

```c++

 #pragma omp parallel
  {
    body;
  }
```
#### for
[parallel for]
(https://gcc.gnu.org/onlinedocs/libgomp/Implementing-FOR-construct.html#Implementing-FOR-construct)による記載もできる。forで示されたループごとに、複数の計算資源に渡す。以下のコードの場合、`ub - lb`個に分けられ処理が行われる。

```c++

#pragma omp parallel for
  for (i = lb; i <= ub; i++)
    body;
```
#### private
変数の共有可否の指定を[private](https://gcc.gnu.org/onlinedocs/libgomp/Implementing-PRIVATE-clause.html#Implementing-PRIVATE-clause)等で行う。ここでは、変数kがスレッド毎に別々に利用されることを示す。
この為、このループの外側で定義された変数(k)と、`#pragma omp`内の変数は関係しない。

```c++

#pragma omp parallel for private(k)
    for (k = 0; k < r_->size(0)*r_->size(1); k++) {
    {
      scalar_t* ptr_output = output_data + k*nOutputCols*nOutputRows;
      int64_t l;
      for (l = 0; l < nOutputRows*nOutputCols; l++)
        ptr_output[l] = 0.0;
    }
  }
```
##### 参考資料
- [(nag-j)5 データ共有属性の指定について](https://www.nag-j.co.jp/openMP/openMPDataSharing.html)


#### critical
[同一名称(ここだとblasgemm)のクリティカルセクションに入るのを防ぐ宣言](https://gcc.gnu.org/onlinedocs/libgomp/Implementing-CRITICAL-construct.html#Implementing-CRITICAL-construct)

```c++
#pragma omp critical(blasgemm)
   body:
```
#####参考資料
- [(nag-j)14 critical指示構文](https://www.nag-j.co.jp/openMP/openMPCriticalSection.html)

### OpenMPの利用例
[at::parallel_forのバックエンド](https://github.com/pytorch/pytorch/blob/v1.3.0/aten/src/ATen/ParallelOpenMP.h#L16)で使われている。なお、[divup](https://github.com/pytorch/pytorch/blob/v1.3.0/aten/src/ATen/Parallel.h#L15)は、ATen/Parallel.hで定義されている。このコードでわかるように、スレッド数に設定した範囲で処理を分配して実行する。

```c++

template <class F>
inline void parallel_for(
    const int64_t begin,
    const int64_t end,
    const int64_t grain_size,
    const F& f) {
  TORCH_CHECK(grain_size >= 0);
  if (begin >= end) {
    return;
  }
#ifdef _OPENMP
  std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
  std::exception_ptr eptr;
#pragma omp parallel if (!omp_in_parallel() && ((end - begin) >= grain_size))
  {
    int64_t num_threads = omp_get_num_threads();
    int64_t tid = omp_get_thread_num();
    int64_t chunk_size = divup((end - begin), num_threads);
    int64_t begin_tid = begin + tid * chunk_size;
    if (begin_tid < end) {
      try {
        f(begin_tid, std::min(end, chunk_size + begin_tid));
      } catch (...) {
        if (!err_flag.test_and_set()) {
          eptr = std::current_exception();
        }
      }
    }
  }
  if (eptr) {
    std::rethrow_exception(eptr);
  }
#else
  f(begin, end);
#endif
}
```



#参考資料
## PyTorch
- [Native ATen/Parallel backend #20087](https://github.com/pytorch/pytorch/pull/20087)
- [Intra-operator parallelism settings in PyTorch #19001](https://github.com/pytorch/pytorch/issues/19001)

## OpenMP
- [OpenMP Specifications](https://www.openmp.org/specifications/)
   - [OpenMP Application Programming Interface 5.0](https://www.openmp.org/spec-html/5.0/openmp.html)
- [gcc OpenMP](https://gcc.gnu.org/wiki/openmp)
   - [GNU libgomp](https://gcc.gnu.org/onlinedocs/libgomp/index.html) gompは、the GNU Offloading and Multi Processing Runtime Libraryの略語
      - libgompにオフロードされて、並列に処理されている。
      - 文書
         - [If undefined, dynamic scheduling and a chunk size of 1 is used.](https://gcc.gnu.org/onlinedocs/libgomp/OMP_005fSCHEDULE.html#OMP_005fSCHEDULE) 
      - ソースコード
         - [struct gomp_task_icv](https://github.com/gcc-mirror/gcc/blob/gcc_5_3_0_release/libgomp/libgomp.h#L228) (OpenMP 4.0制御用構造体)
         - [struct gomp_task_icv gomp_global_icv](https://github.com/gcc-mirror/gcc/blob/gcc_5_3_0_release/libgomp/env.c#L57)
         - [gomp_loop_init](https://github.com/gcc-mirror/gcc/blob/gcc_5_3_0_release/libgomp/loop.c#L36)(OpenMPのループでの処理関数等)
      - (テストコード)[libgomptestsuite](https://github.com/gcc-mirror/gcc/tree/gcc_5_3_0_release/libgomp/testsuite)
   - [GNU Compiler Collection (GCC) Internals](https://gcc.gnu.org/onlinedocs/gccint/index.html)
      - [11.7.7 OpenMP](https://gcc.gnu.org/onlinedocs/gccint/OpenMP.html#OpenMP) 
   - ソースコード
      - [#pragma omp定義箇所](https://github.com/gcc-mirror/gcc/blob/gcc_5_3_0_release/gcc/c-family/c-pragma.h#L27)

- [OpenMP Loop Scheduling](https://iadac.github.io/events/adac6/ciorba.pdf#page=5)
- [(StackOverflow)Using GDB with OpenMP](https://stackoverflow.com/questions/18768674/using-gdb-with-openmp)

## C++
- [Lambda expressions (since C++11)](https://en.cppreference.com/w/cpp/language/lambda)
   - [構文](https://cpprefjp.github.io/lang/cpp11/lambda_expressions.html#syntax)
   - [キャプチャ](https://cpprefjp.github.io/lang/cpp11/lambda_expressions.html#capture)