pybind11 + pytorch + cmake で pytorch(Python) + C++ 連携のメモ

## 背景

pytorch 1.8 から fft(CUDA 対応) やら linalg(beta version) ついたりして, pytorch を数値演算フレームワークとして使うでもいい感じなので, python + C++ アプリに pytorch を pybind11 で組み込みたい.

Python 世界だけだと, 大規模データや並列処理がやはりどうしても苦手になってしまう. また Multiprocessing とか Pipe だと無駄にメモリを消費したり扱えるデータサイズに制限があったりでやっぱり大規模には向かない.

必要なところだけ C++ で効率的にやりつつも, 残りは pytorch つかって効率的にコーディングしたい.

Ubuntu(Linux)を想定します.

## conda installed pytorch

python 環境との親和性を考えて, libtorch ではなく conda や pip で入る通常の pytorch を利用します.

conda などで python 環境を作っているとします.

## TorchConfig.cmake の探索

以下のように python site パッケージで pytorch のインストールしてあるパスを取得して指定します.

cmake で `Torch_DIR` で 

```bash
TORCH_INSTALLED_DIR=`python -c "import site; print (site.getsitepackages()[0])"`/torch/share/cmake

cmake -DTorch_DIR=${TORCH_INSTALLED_DIR} ...
```

cmake では, `Xyz_DIR` とすると, `XyzConfig.cmake` or `xyz-config.cmake` を探すようになっています.

`CMAKE_MODULE_PATH` で指定もできますが, 他のカスタム cmake ファイルの探索とかぶってうまくいかないときもありますから, `Torch_DIR` がよいでしょう.

`Torch_DIR` は関連する cmake オプションより手前に定義する必要がありますので, 環境変数で処理するのがいいかもしれません.
(`CMAKE_MODULE_PATH` は優先順位高いのか cmake 内で他のオプションより先に解釈されるっぽい)

```bash
Torch_DIR=${TORCH_INSTALLED_DIR} cmake ....
```

## C++ compiler ABI

conda とかで入る pytorch は互換性のためか new GCC CXX11 ABI off(昔の ABI)でコンパイルされています. コンパイラも gcc を使っているようです.

```
# When we build libtorch with the old GCC ABI, dependent libraries must too.
if("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
  set(TORCH_CXX_FLAGS "-D_GLIBCXX_USE_CXX11_ABI=0")
endif()
```

TorchConfig.cmake では, gcc の場合は自動で `_GLIBCXX_USE_CXX11_ABI=0` を define しますが, clang などの場合はこのフラグを明示的に `target_compile_options` で追加する必要があります.

定義しないと, コンパイルはいけますが実行時にクラッシュします.

### C++ STL

また, STL は gnustl なので, libc++ を使いたい場合は pytorch をソースコードからビルドする必要があるでしょう.

## リンク

TorchConfig.cmake では, `torch_python` とのリンクが抜けていますので追加します.

lib のパスは

`${TORCH_INSTALL_PREFIX}/lib` で得られます.

## 例

最終的に CMake スクリプトは以下のようになります.

```
find_package(Torch REQUIRED)

# path to pybind11
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/pybind11 pybind11_build)

add_library(torch_testmod SHARED
    src/torch_library.cc)

target_include_directories(torch_testmod PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/pybind11/include)

target_compile_options(torch_testmod PRIVATE "-D_GLIBCXX_USE_CXX11_ABI=0")

target_link_libraries(torch_testmod "${TORCH_LIBRARIES}"
    torch_python
    pybind11::module pybind11::lto pybind11::windows_extras)
set_property(TARGET torch_testmod PROPERTY CXX_STANDARD 14)

# TORCH_INSTALL_PREFIX will be available in find_package(Torch)
target_link_directories(torch_testmod PRIVATE ${TORCH_INSTALL_PREFIX}/lib)

pybind11_extension(torch_testmod)
pybind11_strip(torch_testmod)

```

C++ コード例. torch のヘッダのどこかに pybind11 インクルードがあるようで, pybind11 ヘッダの明示的なインクルードは不要

```c++
#include <torch/extension.h>

#include <iostream>

torch::Tensor d_sigmoid(torch::Tensor z) {
  auto s = torch::sigmoid(z);
  return (1 - s) * s;
}

PYBIND11_MODULE(torch_testmod, m) {
  m.def("sigmoid", &d_sigmoid, "Sigmoid");
}
```

cmake bootstrap は以下のようになります.

```bash

# Path to TorchConfig cmake files.
TORCH_INSTALLED_DIR=`python -c "import site; print (site.getsitepackages()[0])"`/torch/share/cmake

echo "TORCH_DIR = " ${TORCH_INSTALLED_DIR}

mkdir build
cd build

Torch_DIR=${TORCH_INSTALLED_DIR} \
CXX=clang++-11 CC=clang-11 cmake -DCMAKE_BUILD_TYPE=Debug \
  -DCMAKE_VERBOSE_MAKEFILE=1 \
  ..
```

ビルドしたあとの, サンプル .py は以下のようになります.

```py
>>> import torch
>>> import torch_testmod
>>> a = torch.Tensor([1, 2, 3])
>>> torch_testmod.sigmoid(a)
tensor([0.1966, 0.1050, 0.0452])

```

Voila! :tada:

## 問題点

libtorch および pybind11 がヘッダたくさん, テンプレートたくさんでコンパイルが激遅です.  :cry:

上記の数行の C++ コードでも Threadripper 1950X で 10 秒くらいかかります.

torch と関連するコードはなるべく変更なしにして, ccache でコンパイルキャッシュ作って速くするなどの対応が求められるでしょう.