OneFlow ソースコード分析: Eager モードでの Tensor ストレージ管理

著者 | 鄭建華

さまざまな Tensor タイプのストレージ管理方法

Lazy Tensor のストレージは、Runtime や Actor などのオブジェクトによって管理されます。静的グラフがコンパイルされた後、必要なオブジェクトの数とストレージ容量が決定されます. ランタイムなどは、初期化時にストレージを割り当て、終了時にリソースを再利用します.

Eager モードでは、Global Tensor は Local Tensor の分散カプセル化と見なすことができ、EagerGlobalTensorImpl のローカルデータは

EagerLocalTensorImpl オブジェクト。EagerLocalTensorImpl を調べることで、熱心なモードでのテンソルのストレージ管理を理解できます。

参考までにサンプルコードは以下の通りです。

import numpy as np
 import oneflow as flow
 
 a = np.random.randn(1, 4)
 flow.tensor(a, device=flow.device("cpu"), dtype=flow.float)

テンソルは関連するクラスの関係を保存します

EagerLocalTensorImpl のストレージ関連のクラス関係は次のとおりです。

サンプルコードの実行プロセスをたどって、ダイアグラム内のオブジェクトがいつ、どのように構築され、誰がストレージを保持し、どのように割り当ておよび解放されるかを確認します。

仮想マシン命令を使用して Tensor 用のストレージを割り当てる

テンソルコンストラクターは、function::_legacy_tensor_ctor によって、Python C API を介して PyTensorObject_init として登録されます。

署名に従って転送します。

サンプルコードは TensorWithDataFunctor に対応しています

MakeLocalTensorFromData を呼び出してテンソルを構築し、この関数で Functional::Empty と EmptyFunctor を呼び出してストレージを割り当てます。関連する属性を EmptyFunctor の attrs に格納し、OpInterpUtil::Dispatch を呼び出して、vm 命令の実行準備プロセス中にストレージを割り当てます。

EmptyFunctor によって返されるテンソルは、ストレージスペースのみを持ち、データを持たないオブジェクトです。データは後で CopyLocalTensorFromUntypedArray によってコピーされます

終了。

3.1 ストレージ関連オブジェクトの構築

これはイーガーモードのローカルテンソルであるため、OpInterpUtil::Dispatch は実行のために NaiveInterpret に転送されます。コード例では、この関数の入力パラメーターは次のとおりです。

入力は空の配列です
出力には要素が 1 つしかなく、null ポインターです

出力のテンソルポインターはすべて空であるため、one::TensorStorage メンバー変数が null ポインターである EagerLocalTensorImpl オブジェクトを作成する必要があります。

output_eager_blob_objects の要素が初期化されていないため、tensor_impl->InitEagerBlobObject が呼び出されます

初期化します。tensor_storage_ はまだ空であるため、プロセスは次のことを行います。

vm::TensorStorage オブジェクトを作成する
EagerBlobObject オブジェクトを作成する
set_eager_blob_object
- UpdateTensorStorage
  - one::TensorStorage オブジェクトを作成する
  - テンソルストレージリリースのコールバック関数を設定する

上記のオブジェクトの作成は、関連する情報のみを記録し、テンソルのストレージ割り当てを含みません。

one::TensorStorage に登録されたコールバック関数はメンバー変数 releaser_hook_ に割り当てられ、この関数は仮想マシン命令を通じてテンソルを解放することに注意してください。

3.2 命令実行中のテンソルストレージの割り当て

テンソルストレージを割り当てるプロセスは次のとおりです。

vm::命令::計算
vm::InstructionPolicy::ComputeIf
vm::OpCallInstructionPolicy::Compute
OpCallInstructionUtil::計算
メモリアロケータを取得する
OpCallInstructionUtil::AllocateOutputBlobsMemory
blob_object->TryAllocateBlobBlobBodyMemory
アロケータ -> 割り当て

EagerBlobObject::TryAllocateBlobBodyMemory では、アロケーターによって割り当てられたストレージアドレスが dptr に割り当てられ、ストレージアドレス dptr と Free 関数がスマートポインターを構築し、それを vm::TensorStorage の blob_dptr_ 変数に割り当てます。

仮想マシン命令による Tensor ストレージのリリース

前のセクション 3.1 で述べたように、EagerLocalTensorImpl は、EagerBlobObject を初期化し、one::TensorStorage を作成するときに、コールバック関数を設定してテンソルを解放します。コールバック関数は、変数 releaser_hook_ に格納されます。

、このコールバック関数は one::TensorStorage が破棄されたときに呼び出されます。この情報をまとめると、one::TensorStorage は破棄されるときに次の操作を実行します。

vm::InstructionList instruction_list;
 InstructionsBuilder instructions_builder(&instruction_list);
 
 // JUST(Build(&instructions_builder));
 if (eager_blob_object->producer_stream().has_value()) {
   JUST(instructions_builder->ReleaseTensor(eager_blob_object));
 }
 
 JUST(vm::Run(instructions_builder.mut_instruction_list()));

InstructionsBuilder::ReleaseTensor では、他のストリームが最近eager_blob_object を使用した場合、それらは SoftSyncStreamBetween を通じて同期されます。このようにして、ストレージ依存の問題は解決されます。

通常、ストレージは tensor の producer_stream を通じて解放され、このオブジェクトに従って対応する vm::Stream オブジェクトが取得され、それに応じて命令命令 (eager_blob_object および vm_stream を含む) が構築されます。コードは FastReleaseTensorInstructionPolicy であり、その Compute メソッドは特定のストレージ解放ロジックを実行します。プロセスは次のとおりです。

ReleaseTensorInstructionPolicy::Release()
eager_blob_object->DeallocateBlobDataPtr()
tensor_storage_->Release()
tensor_storage_->_Release()
blob_dptr_.reset()
- スマートポインタがリセットされ、ストレージの割り当て時に指定された Free メソッドが呼び出されます

形状変更などのシナリオのストレージ管理

In scenario such as reshape, slice, and transpose, the parameters of the EagerLocalTensorImpl constructor called include input tensor_storage, so the tensor_storage_ variable of this tensor is not empty. InitEagerBlobObject が実行されると、形状やストライドなどの情報を提供するために EagerBlobObject のみが作成されます。 ; One::TensorStorage は再度作成されませんが、入力のストレージは再利用されます。

2 つの TensorStorage タイプをマージできますか?

one::TensorStorage によって保存されたコールバック関数が、破棄されたときに vm::TensorStorage 内のストレージを解放するためにトリガーされるのはなぜですか?

one::TensorStorage にはもう 1 つのリリーサーしかありません。これら 2 つのストレージタイプをマージできますか?

現在の設計では、2 つのタイプをマージすることはできません。one::TensorStorage::releaser_hook_ は EagerBlobObject のスマートポインターを保持するため、EagerBlobObject は vm::TensorStorage のスマートポインターも保持します。2 つのストレージタイプが 1 つにマージされると、循環参照が発生し、オブジェクトを破棄できなくなり、メモリリークが発生します。

したがって、vm::TensorStorage は、複数のテンソル間で共有できる単純なストレージです。EagerBlobObject には、ストレージと、形状、ストライド、および data_type などの一意のオブジェクト情報の両方が含まれます。one::TensorStorage は、循環参照を回避するために導入され、ストレージの解放を担当します。

付録

GDB ブレークポイントの例

break oneflow::one::MakeLocalTensorFromData
 break oneflow::one::NaiveInterpret
 break oneflow::vm::VirtualMachineEngine::DispatchInstruction
 break oneflow::vm::OpCallInstructionUtil::Compute
 break oneflow::vm::OpCallInstructionUtil::AllocateOutputBlobsMemory
 break oneflow::vm::EagerBlobObject::TryAllocateBlobBodyMemory
 break oneflow::vm::ReleaseTensorInstructionPolicy::Release
 break oneflow/core/eager/eager_blob_object.cpp:107

参考文献

OneFlow（https://github.com/Oneflow-Inc/oneflow/tree/b51cb72430619f6088e47bbb8b8226f37299573a）
OneFlow ソースコード分析: Tensor 型システムと Local Tensor

他のみんなが見ている

ウェルカムスター、OneFlow を試す: github.com/Oneflow-Inc/oneflow/ http://github.com/Oneflow-Inc/oneflow/

OneFlow ソース コード分析: Eager モードでの Tensor ストレージ管理

おすすめ

OneFlow ソースコード分析: Eager モードでの Tensor ストレージ管理