PyTorchの関数本体ってどこで定義されてるの？逆伝搬は？~ATenの世界へのみちしるべ~

# はじめに
本稿はPyTorch中級者以上向けのニッチな内容ですが，**どうしても計算方法の定義が見たい**という方にはおすすめです．内容は主に３つです．

- PythonからC++で書かれた**ATenライブラリ**への橋渡し
- (conv2dを例に挙げ，)**ATen内で順伝搬関数**が扱われているかを見る
- 逆伝搬関数がいかに順伝搬関数と**紐づいている**かを見る


# 隠された実装
PyTorchの[公式ドキュメント](https://pytorch.org/docs/stable/index.html)を見ながらプログラミングをし，時には[Source]ボタンを押してgithubの実装を確認したりしますよね．でも[Source]ボタンがないときありますよね？例えばtorch.nn.Conv2dには[Source]ボタンがありますが，torch.nn.functional.conv2dにはありません．　　

・・・なんで同じConv2dなのに実装が見れないの？　　

と思いつつ，PyTorchの[github](https://github.com/pytorch/pytorch)を確認してみます．torch.nn.Conv2dはtorch.nn.functional.conv2d を利用しているので，後者の定義があるはずの torch/nn/functional.py を覗いてみると以下の定義がありました．

```python
conv2d = _add_docstr(torch.conv2d, r"""
# 略
"""
)
```
はぁ・・・なんとなく[Source]ボタンがない理由がわかりました．見ても意味がないからです．ちなみにやっていることは conv2d = torch.conv2d とほぼ一緒です．
じゃあtorch.conv2dの定義を見ればいいのかと思い探します．
  
で，どこにtorch.conv2dの定義あるの・・・？
## Pythonにおける実装を探そう
実は**PyTorchの本質的機能の多くはPython内で完結していません**．そう簡単にtorch.conv2dの定義を探すことは容易ではありません．と言ってもPythonのSyntaxを破ることはできないのでtorch.conv2dも一応Pythonの中で定義されています． torch/\_\_init\_\_.py ファイル内に以下で定義されています．

```python
for name in dir(_C._VariableFunctions): 
    if name.startswith('__'): 
        continue 
    globals()[name] = getattr(_C._VariableFunctions, name) 
    __all__.append(name) 
```
コードを見せられても意味が良くわかりませんね．conv2dという文字列も見当たらないし．少しだけ説明すると，\_Cとはtorch/\_C.soファイルのことであり，C++などの低級言語からコンパイルされたモジュールです．（コンパイル前のgithub上では\_C.soは確認できないので注意してください．）その\_Cの\_VariableFunctionsの中にconv2dが含まれており，for文内部でそれ以外の関数も含めてtorch.~として使用できるようにしています．

要約すると，ここを見ていても仕方ないです．なぜならconv2dの本体はPythonではなくC++で実装されているからです．**コンパイルされる前のソース**がどこにあるのかを見つけるのが建設的というものでしょう．
## ATenの世界
ATenとは**A Ten**sor libraryのことであり，githubだと[ここ](https://github.com/pytorch/pytorch/tree/master/aten/src/ATen)にあります．ATenは基本的にC++とCUDAによって実装されており，**PyTorchの心臓部**とも言える重要な実装が行われています．torch.conv2dのような関数はPyTorch内部でnative functionと呼ばれ，ATen/native内で実装されています．ATen/native内を探検するみちしるべとして[native_functions.yaml](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml)を見ることが有効です．このyamlファイルに記載されている関数は\_C.\_VariableFunctionsに格納されますし，ATen/native内のいずれかのファイルに実装されているからです．実際に検索してみるとconv2dがこのyamlファイルの中に記載されています．

```yaml
- func: conv2d(Tensor input, Tensor weight, Tensor? bias=None, int[2] stride=1, int[2] padding=0, int[2] dilation=1, int groups=1) -> Tensor
```
ということはATen/nativeのどこかに実装されてるはずです．それらしい名前のファイルをあたってみると，ATen/native/Convolution.cpp内にconv2dの定義を見つけることができます．

```c++
at::Tensor conv2d(
    const Tensor& input, const Tensor& weight, const Tensor& bias, 
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
  return at::convolution(input, weight, bias, stride, padding, dilation,                    
                         false, {{0, 0}}, groups); 
}
```
ここまででPython〜ATen（コンパイル前のソース）への橋渡しができました．
## ATenの中で何が行われているのか？
じゃあconv2dで呼び出しているat::convolutionはどうなっているの？という疑問は当然湧いてきます．**もう少し深入り**してみましょう．コードを辿っていって，GPUの並列計算を起動する命令（CUDA Launch）かcuDNNの関数が見れたら嬉しいですね．
at::convolutionは同じファイル内で定義されており，at::\_convolutionを呼び出しています．at::\_convolutionでは詳細な畳み込み手法の選択が行われているようです．もしGPUを使用していたらcuDNNを使用していると思うので，at::cudnn\_convolutionをさらに見て行くこととしましょう．少し苦労しますが，ATen/native/cudnn/Conv.cppに見つけることが出来ます．芋づる式にcudnn\_convolution→cudnn\_convolution\_forward→raw\_cudnn\_convolution\_forward\_out→raw\_cudnn\_convolution\_forward\_out\_32bitと辿ることができます．最終的には以下のcuDNN関数までコードを辿ることができます．

```c++

   AlgoIterator<cudnnConvolutionFwdAlgoPerf_t>(args, benchmark).try_all(
     [&](const cudnnConvolutionFwdAlgoPerf_t &fwdAlgPerf){ 
       Tensor workspace = allocate_workspace(fwdAlgPerf.memory, input); 
  
       // update convDesc mathType since cudnn 7.4+ now requires both algo + mathType to figure out 
       // whether to use Tensor core kernels or not 
       // See Note [behavior of cudnnFind and cudnnGet] 
       AT_CUDNN_CHECK(cudnnSetConvolutionMathType(args.cdesc.mut_desc(), fwdAlgPerf.mathType)); 
  
       Constant one(dataType, 1); 
       Constant zero(dataType, 0); 
  
       AT_CUDNN_CHECK(cudnnConvolutionForward( 
         args.handle, 
         &one, args.idesc.desc(), input.data_ptr(), 
         args.wdesc.desc(), weight.data_ptr(), 
         args.cdesc.desc(), fwdAlgPerf.algo, workspace.data_ptr(), fwdAlgPerf.memory, 
         &zero, args.odesc.desc(), output.data_ptr())); 
       } 
   ); 
```
上のコードにおいてcudnnConvolutionForwardがメインで動いている関数です．この関数は非オープンソースのcuDNNの関数で，**遡ることのできる限界**なのでここまでとしておきましょう．
## 逆伝搬の設定
ATenまでコードを遡ってみた訳ですが，一つ**大事なことを忘れ**ています．

conv2dの逆伝搬ってどう設定してるの・・・？自動微分って事前に順伝搬と逆伝搬の関数を決めておくんだよね？でもconv2dに紐づいた逆伝搬関数をまだみてないぞ？

そう，まだ逆伝搬の定義と**順伝搬との紐付け**が一切出てきませんでした．PyTorchはいかにしてconv2dの逆伝搬を行うのでしょうか．実はtoolsというところに秘密が隠されています．toolsはPyTorchのBuildを行う上で重要なフォルダであり，[tools/autograd](https://github.com/pytorch/pytorch/blob/master/tools/autograd/)では自動微分に関して重要な関数が定義されています．ここには順伝搬と逆伝搬の関数の紐付けを決める[derivatives.yaml](https://github.com/pytorch/pytorch/blob/master/tools/autograd/derivatives.yaml)があります．このyamlファイルを調べてみるとconv2dについては記載がありませんが，代わりに以下の設定が見つかります．

```yaml
- name: cudnn_convolution(Tensor self, Tensor weight, int[] padding, int[] stride, int[] dilation, int groups, bool benchmark, bool deterministic) -> Tensor
  self, weight: "grad.defined() ? cudnn_convolution_backward(self, grad, weight, padding, stride, dilation, groups, benchmark, deterministic, grad_input_mask) : std::tuple<Tensor, Tensor>()"
```
どうやらconv2dの実装を遡っていく中で出てきたcudnn_convolutionについて逆伝搬が設定されているようです．この逆伝搬関数cudnn_convolution_backwardはATen/native/cudnn/Conv.cpp内で定義されています．ちなみにこの下には以下の逆伝搬の設定が書いてあります．

```yaml
- name: cudnn_convolution_backward(Tensor self, Tensor grad_output, Tensor weight, int[] padding, int[] stride, int[] dilation, int groups, bool benchmark, bool deterministic, bool[2] output_mask) -> (Tensor, Tensor)
  grad_output, self, weight: _convolution_double_backward(grads[0], grads[1], Tensor(), grad_output, weight, self, stride, padding, dilation, false, std::vector<int64_t>(padding.size(), 0), groups, benchmark, deterministic, true, grad_input_mask)
```
**逆伝搬関数（１次微分）のさらなる逆伝搬関数を設定することで２次微分**を得ることができるみたいですね．
## 別題：線形方程式の例
最後にconv2d以外の例としてtorch.solveのATen内での実装を探してみましょう．torch.solveは$A,b$を入力として線形方程式$Ax=b$を解き$x$を返します．[native_functions.yaml](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml)をまず見てみると，ちゃんと以下の記述があります．

```yaml
- func: solve(Tensor self, Tensor A) -> (Tensor solution, Tensor LU) 
  use_c10_dispatcher: full 
  variants: function, method 
```
ちょっとわかりにくいですが，ATen/native/BatchLinearAlgebra.cppにsolveの定義があります．どうやらat::\_solver\_helperという関数を呼び出しています．この関数を見つければいいのですが・・・見つかりません？もう一度yamlを見てみると以下の記述が見つかります．

```yaml
- func: _solve_helper(Tensor self, Tensor A) -> (Tensor, Tensor)
  use_c10_dispatcher: full 
  variants: function 
  dispatch: 
    CPU: _solve_helper_cpu 
    CUDA: _solve_helper_cuda 
```
実はこの記述により**CPUとGPUの場合の関数のスイッチング**を行えます．at::\_solver\_helperという関数の実体は存在せず，CPUならば\_solver\_helper\_cpuに，GPUならば\_solver\_helper\_cudaに引き継がれます．CPUの場合を辿っていくと，ATen/native/BatchLinearAlgebra.cppに\_solver\_helper\_cpuがあります．さらにapply\_solveという関数が呼ばれ，lapackSolveというLAPACKというライブラリを用いた線形方程式のソルバーに引き継がれていきます．
では次に逆伝搬を調べてみましょう．[derivatives.yaml](https://github.com/pytorch/pytorch/blob/master/tools/autograd/derivatives.yaml)には以下のように逆伝搬を設定してありました．

```yaml
- name: solve(Tensor self, Tensor A) -> (Tensor solution, Tensor LU) 
  self: solve_backward_self(grad, self, A) 
  A: solve_backward_A(grad, self, A, solution) 
```
$A,b$の２種類に対する逆伝搬を別々に書くことができるようです．ではsolve\_backward\_selfとsolve\_backward\_Aの実装を探してみる・・と・・・？いくら探してもATen/native内に見つけることができません．実はこれらの関数は[tools/autograd/templates/Functions.cpp](https://github.com/pytorch/pytorch/blob/master/tools/autograd/templates/Functions.cpp)内で以下のように実装されています．
> pytorch >= 1.7 では [torch/csrc/autograd/FunctionsManual.cpp](https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/FunctionsManual.cpp)に移動したようです．

```c++
Tensor solve_backward_self(const Tensor & grad, const Tensor & self, const Tensor & A) {
  return std::get<0>(at::solve(grad, A.transpose(-2, -1))); 
} 

Tensor solve_backward_A(const Tensor & grad, const Tensor & self, const Tensor & A, const Tensor & solution) { 
  Tensor grad_self = solve_backward_self(grad, self, A); 
  if (self.ndimension() == 2 && A.ndimension() == 2) { 
    return -at::mm(grad_self, solution.transpose(-2, -1)); 
  } 
  return -at::matmul(grad_self, solution.transpose(-2, -1)); 
} 
```
どうやら**逆伝搬の実装自体はATenだけではなくtools/autograd/templates内でも行う**ようです．棲み分けとしては，ATenでは複雑な逆伝搬を，templates内では容易な逆伝搬を行っているように見受けられます．なお，線形方程式の逆伝搬がなぜこのような式になっているかは[ここ](https://qiita.com/takeuchiseijin/items/c3f93d28e3ca465cbdd9)に解説してあります．
# 終わり
以上でPyTorchの内部を見てみる記事は終わりです．conv2dとsolveの内容で大概のパターンに対処できると思いますので，実装を確認したい処理があったらATenの中を探してみるのも良いでしょう．