Notes on PyTorch Internals III

原版英文リンク:Edward Z.Yang's PyTorch internals:Inside 245-5 D
Mechanics
Code Flow
PyTorchソースウェアハウスには多くのファイルが含まれており、詳細はCONTRIBETINGを参照してください.最も重要な4つのフォルダとソースモジュールは次のとおりです.
torch/:PyTorchモジュールライブラリ.PyTorchが開発する最も一般的な機能を含み、importによってインポートされる事前定義モジュール.PyTorchのPythonフロントエンド(frontend).
torch/csrc/:PyTorchフロントエンドモジュールのC++コードは、C++コードとPythonのバインドを実現します.さらに、自動導出エンジン(autograd/)、JITコンパイラ(jit/)、およびPyTorchのC++フロントエンド(api/)が含まれる.
aten:「A Tensor Library」の略で、テンソル関連操作の実現には、自動導出機能は含まれていない.src/ディレクトリには、古いcバージョン実装(TH/,THC/,THNN/,THCUNN/)と、C++ベース実装(ATen/)の2つの実装が含まれています.異なるデバイス上のテンソル操作を実現する方法は、それぞれ異なるフォルダ(ATen/cpu,ATen/cuda,ATen/sparse,...)にある.
c 10:Caffe 2とATenの混合略語(caffe ten)、PyTorchのコア抽象と基礎機能の実現、テンソルの具体的な記憶と実現方式を含む、サーバとモバイル端末への配備をサポートする.PyTorchは、ATen/coreのベースコアインプリメンテーションをc 10/coreに移植しています.PyTorchのC++バックエンド(backend)
コアモジュールは、上位レベルの論理実装をサポートします.PyTorchの加算関数torch.addを例にとると、モジュール間呼び出しの流れは以下の通りである.

Python関数をC関数呼び出しに変換し、PythonパラメータをC++パラメータとして解析し、torch/csrc/中関数で実現する.例えば、PyTorchのadd関数を例にとると(以下のコードが自動的に生成される):

// actual binding
static PyMethodDef torch_functions[] ={
...
{"add", (PyCFunction)THPVariable_add, METH_VARARGS | METH_VARKEYWORDS | METH_STATIC, NULL}
...
}

// auto-generated codes, needed to build PyTorch to generate it
static PyObject* torch._C.VariableFunctions.add.THPVariable_add(
          PyObject* self_, PyObject* args, PyObject* kwargs){
static PythonArgParser parser(...);

ParsedArgs<4> parsed_args;
auto r = parser.parse(args, kwargs, parsed_args);
...

if(r.isNone(3)){
    return wrap(dispatch_add(r.tensor(0), r.tensor(1), r.scalar(2)));
}else{
    return wrap(dispatch_add(r.tensor(0), r.tensor(1), r.scalar(2), r.tensor(3)));
}
...
}

torch_functionsPython関数とCバージョン関数名の対応関係を定義し、このマッピングテーブルを介して対応するCバージョン関数を問い合わせる.PythonArgParserクラスはPyTorchパラメータのCバージョン解析を実現し、dispatch_addスケジューリング下位Cバージョンのaddによって実現する.計算結果はwrapでPyObjectオブジェクトに再パッケージされ、Pythonレイヤに返されます.
2.変数タイプのスケジューリング前のステップのdispathc_add関数スケジューリングは、実際にself.add(tensor, scalar)関数、すなわちテンソル自体の実装バージョンを呼び出し、このバージョンは、次の関数によって実装され、異なる変数タイプの実装バージョンを呼び出す.

// inline functions defined on the 'type'
inline Tensor Tensor::add(const Tensor& other, Scalar alpha) const {
    return type().add(*this, other, alpha);
}

関数type()は、変数の特定のタイプを決定し、実際のタイプの実装を虚関数addによってスケジューリングする.これらの処理はaten/src/Tenモジュールで完了した.
3.設定タイプとレイアウトのスケジューリングtype()は、実際には変数とデバイスタイプのスケジューリングを同時に完了し、TypeDefault、GPUFloatTypeなどのデータ型とデバイスタイプを含む説明を返す.各タイプについて、PyTorchはbuild後に具体的に次のような実装コードを生成します.

Tensor TypeDefault:add(const Tensor& self, const Tensor& other, Scalar alpha) const {
    const OptionalDeviceGuard device_guard(device_of(self))   # device type checking
    return at::native::add(self, other, alpha)    # modern c++ impl.
}

add関数は、異なるタイプの変数およびデバイスタイプの最下位実装と同じであるため、TypeDefaultを介して統一的にカプセル化することができる.計算操作が異なる場合は、GPUFloatType::add(...)と同様に、実装を拡張して対応するバージョンを呼び出す必要があります.
4.コアコードの呼び出しステップ3のコードは、より下位レベルのat::native::add(self, other, alpha)実装をカプセル化する.これらのインプリメンテーションは、aten/src/Tenのモジュールに依存し、C++バージョン(native/)または古いcバージョン(TH/,THC/,THNN/,THCUNN/)によってインプリメンテーションされる.
Kernels
PyTorchは、aten/src/Tenモジュールによってサポートされるコアコンピューティングオペレータを開発するためのツールと仕様を提供しています.完全なカスタムコア操作のサンプルコードを示します.

Tensor my_op(Tensor& result, const Tensor& self, const Tensor& other){
    // error checking
    TORCH_CHECK(result.is_cpu() && self.is_cpu() && other.is_cpu());
    TORCH_CHECK(self.dim() == 1);
    TORCH_CHECK(self.sizes() == other.sizes());

    // output allocation
    result.resize_(self.sizes());
    
    // data type (dtype) dispatch
    AT_DISPATCH_FORALL_TYPES(
        self.scalar_type(), "my_op", [&]{
            my_op_cpu(result, self, other), 
        }
    );
}

template
void my_op_cpu(Tensor& result, const Tensor& self, const Tensor& other){
    // data access
    auto result_accessor = result.accessor();
    auto self_accessor = self.accessor();
    auto other_accessor = other.accessor();
  
    // parallelization
    parallel_for(0, self.size(0), 0, [&](int64_t start, int64_t end){
        ... self_accessor[i] ...
    });
}

次のセクションがあります.
メタデータ登録:PyTorchによって提供されるメタデータ要件は、Pythonを生成するバインディングコード(前節で説明したPythonとCコード間の変換とパラメータ解析)を自動化するために使用されます.各定義されたコア・オペレーションには、次のメタデータ・モードが必要です.

    - func: func_name(ArgType arg0, ArgType arg1, ...) -> Return
    variants: function, method
    dispatch:
        CPU: func_cpu
        CUDA: func_cuda 

ここで、func_name:定義されたコア計算アクション関数の名前.ArgType:パラメータタイプ、Tensor, Tensor[], int, int[], float, Scalarなどです.variants:PyTorchがPythonバージョン関数を自動的に生成する名前がテンソルメソッド(function)かネーミングスペースの関数(method)かを制御する2つのタイプを含む.t.foo()バリエーションを使用する場合、at::foo()パラメータを含める必要があります.Pythonバージョンの関数名が自動的に生成されると、このmethodパラメータはパラメータリストから削除されます.たとえば、selfの関数宣言.selfに設定すると、where(BoolTensor cond, Tensor self, Tensor other)の関数名が自動的に生成されます.methodに設定すると、self.where(cond, other)の関数名が自動的に生成されます.デフォルトでは、ATenはnative関数に対してfunction方式名のみを生成し、テンソルに関連するコアオペレータ(e.g,at::where(cond, self, other)など)に対してfunction方式名を使用することができる.add, sub:異なる設定タイプに対してスケジューリング可能な実際の関数名を指定します.異なる設定タイプに対して、異なるバージョンを生成する関数名を指定できます.
より詳細な仕様要求はaten/src/Ten/native/READMEを参照することができる.md.カスタムコア計算関数のメタデータは、上記の要件に従って作成し、native_に追加する必要があります.functions.yamlファイルに登録します.PyTorchは登録した関数に対してメタデータ記述の要求に従って、自動的にPythonのバインディングを生成します.
上記のカスタムコア操作関数は、次のように記述されます.

-func: my_op(Tensor& result, const Tensor& self, const Tensor& other) -> Tensor
variants: function, method
dispath:
    CPU: my_op_cpu
    CUDA: my_op_cuda

逆勾配計算をサポートする必要があるコア計算操作については、同様の方法で導出操作関数のメタデータを提供する必要がある.具体的にはderivatives.yaml
エラー検出(Error Checking)エラーチェックは、コアコードを記述する際に非常に重要です.PyTorchは2つのエラーチェックのツールを提供して開発者を便利にします:low levelの方式はmethodマクロを提供しました;High level方式は、dispatchをTORCH_CHECKにカプセル化し、Tensorなどの検出関数を提供する.
出力メモリ割り当て(Output Allocation)は、結果を出力する前にメモリを予め割り当てる必要がある.PyTorchは、プリアサイメント出力、その場出力、コピー出力などの出力結果をサポートします.実装中、その場出力とコピー出力は、予め割り当てられた出力の単純なパッケージにすぎない.たとえば

// pre-allocate storage for 'result' outside
Tensor& abs_out(Tensor& result, const Tensor& self){
    result.resize_(self.sizes());
    // ... the real impl.
}

// a new allocated operation
Tensor& abs(const Tensor& self){
    Tensor result = at::empty({0}, self.options());
    abs_out(result, self);
    return result
}

// in-place operation
Tensor& abs_(const Tensor& self){
    return abs_out(self, self);
}

データ型スケジューリング(Dtype Dispatch)は、TensorArgマクロによってデータ型スケジューリングを定義する.このマクロは、変数の現在のタイプを特化し、実際に一致するタイプの実装をスケジューリングするテンプレート関数です.
データアクセス(Data Access)PyTorchは、テンソルに対する3つの異なるアクセス方式をサポートします.カプセル化されたアクセス方式は、元のデータポインタにアクセスするよりも便利であり、最下位のcheckDimまたはレイアウトを自動的に処理することができる.AT_DISPATCH_ALL_TYPESアクセステンソルの特定の場所のデータをサポートします.strideサポート規則方式ポーリングアクセス;CPUのシーケンス化には、TensorAccesor等のシーケンス化記述アクセスが提供される.
Notes on PyTorch Internalsシリーズ記事
Notes on PyTorch Internals I Notes on PyTorch Internals II Notes on PyTorch Internals III