TensorFlow XLAのコード解析をしました。 この資料は、TensorFlow XLAのAOT部分に関するものです。 I analyzed the code of TensorFlow XLA. This document pertains to JIT part of TensorFlow XLA.

1. @Vengineer
2017/3/20
TensorFlow XLA
コード解析 : AOT編
チラ見版

2. 勉強会主催 :
Xilinx Zynq MPSoC (2016/02/20)
Altera SDK for OpenCL (2016/06/10)
Xilinx SDSoC (2017/01/28)
PYNQ祭り (2017/03/04)
FPGAディープラーニング実践懇親会 (2017/05/20)
ブログ : Vengineerの戯言
http://blogs.yahoo.co.jp/verification_engineer
Twitter : ＠Vengineer
書籍 : SystemVerilogスタートアップ
http://www.cqpub.co.jp/hanbai/books/36/36191.htm
自己紹介

3. TensorFlow XLAとは
https://www.tensorflow.org/performance/xla/
XLA(Accelerated Linear Algebra)は、TensorFlow計算を最適化
する線形代数のドメイン固有のコンパイラです。 結果として、サー
バーおよびモバイルプラットフォームでの速度、メモリ使用率、移植性
が向上します。 当初、ほとんどのユーザーはXLAの大きなメリットは
見られませんが、JIT(Just-In-Time)コンパイルや
AOT(Ahead-Of-Time)コンパイルを使用してXLAを使用することで
実験を開始できます。 新しいハードウェアアクセラレータをターゲット
とする開発者は、XLAを試すことを特にお勧めします。
原文(英語)をそのまま、Google翻訳にお願いしました。

4. ブログにも書きました
TensorFlow XLAの衝撃
2017年2月20日
http://blogs.yahoo.co.jp/verification_engineer/71016304.html

5. 簡単にまとめると
TensorFlow XLAでは、次の2つをサポートした
1)、JIT (Just-In-Time) コンパイル
ただし、単一マシンのみで、GPUは1つ
2)、AOT (Ahead-Of-Time) コンパイル
CPUのみ
x86/x86-64/ARM/AARCH64/PowerPC

6. この資料は、
TensorFlow XLAの
AOTに関するコードを解析したものをまと
めたです。
ご利用は、自己責任でお願いします。

7. Using AOT compilation
https://www.tensorflow.org/performance/xla/tfcompile
・tfcompileって、何？
・tfcompileは、何をする？
・tfcompileを使うには！
現時点（TensorFlow r1.0) では、AOTのターゲットは
CPU(x86/x86-64/ARM/ARM64/PowerPC)のみサポート。

8. tfcompileって、何？
・TensorFlowグラフを実行可能コードにコンパイルす
るためのツール
・バイナリサイズおよびランタイムオーバーヘッドを減ら
す
・利用例：推論グラフをモバイルデバイス用の実行可
能コードにコンパイル

9. ランタイムが無くなる
TensorFlowグラフは通常、TensorFlowランタイムに
よって実行されます。 これにより、グラフ内の各ノード
の実行ではランタイムオーバヘッドを招く。 グラフ自体
に加えて、TensorFlowランタイム用のコードが必要で
あるため、バイナリサイズが大きくなる。
tfcompileによって生成される実行可能コードは、
TensorFlowランタイムを使用せず、計算で実際に使
用されるカーネルにのみ依存する。

10. tfcompileは、何をする？
tfcompileは、TensorFlowサブグラフからそのサブグ
ラフを実装する関数を生成する。
Feedは関数の入力引数、Fetchは関数の出力引数と
なる。
すべてのPalceholdersとVariablesは、関数の入力引
数としてFeedとして指定する必要がある。
tfcompileによって生成されたファイルは、関数のオブ
ジェクトファイルとして利用できる。

11. tfcompileを使うには！
ステップ1：コンパイルするサブグラフを構成する
ステップ2：tf_libraryビルドマクロを使用してサブグラ
フをコンパイルする
ステップ3：サブグラフを呼び出すコードを書く
ステップ4：最終的なバイナリを作成する

12. tfcompile
バイナリでは提供されていない
ので、ソースコードからビルドす
る必要がある

13. ディレクトリ構成
compilerディレクトリがTensorFlow XLA
・aot
・jit
・tests
・tf2xla
・xla
AOT関連は、主に、aotディレクトリ内にある

14. TensorFlowは、
Bazelを使ってビルドしていますので、
Bazel : https://bazel.build/

15. まずは、BUILDファイル
aot/BUILD
cc_binary(
name = "tfcompile",
visibility = ["//visibility:public"],
deps = [":tfcompile_main"],
)
tfcompile_main

16. tfcompile_main
aot/BUILD
cc_library(
name = "tfcompile_main",
srcs = ["tfcompile_main.cc"],
visibility = ["//visibility:public"],
deps = [
":tfcompile_lib",
":tfcompile_proto",
…..
],
)
tfcompile_main.cc

17. tfcompile_main.cc
int main(int argc, char** argv) {
各種処理フラグの設定
tensorflow::tfcompile::MainFlags flags;
flags.target_triple = "x86_64-pc-linux";
flags.out_object = "out.o";
flags.out_header = "out.h";
std::vector<tensorflow::Flag> flag_list;
AppendMainFlags(&flag_list, &flags);
xla::legacy_flags::AppendCompilerFunctorFlags(&flag_list);
xla::legacy_flags::AppendCpuCompilerFlags(&flag_list);
xla::legacy_flags::AppendCpuRuntimeFlags(&flag_list);

18. tfcompile_main.cc
引数の処理
tensorflow::string usage = tensorflow::tfcompile::kUsageHeader;
usage += tensorflow::Flags::Usage(argv[0], flag_list);
bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list);
tensorflow::port::InitMain(usage.c_str(), &argc, &argv);
tensorflow::tfcompile::Main(flags);
return 0;
}

19. tfcompile::Main
aot/tfcompile_main.cc
コンフィグファイルとグラフファイルの読み込み
ReadProtoFile("config", flags.config, &config);
ReadProtoFile("graph", flags.graph, &graph_def);
グラフの初期化
InitGraph(graph_def, config, flags, &flib, &graph);
グラフのコンパイル
CompileGraph(std::move(graph), flags, &flib, &compile_result);
ファイル(オブジェクト、ヘッダ)の書き出し
WriteStringToFile( …., …., …. );

20. グラフ
情報
コンフィグ
情報
グラフ情報
をHLO(最適化)に変換
HLOをLLVMで
CPU実行コードに変換
オブジェクトファイルへ
の出力

21. ProtoFileの読み込み
aot/tfcompile_main.cc
if (StringPiece(fname).ends_with(".pbtxt")) {
core/platform/env.cc
return ReadTextProto(Env::Default(), fname, proto);
} else {
core/platform/env.cc
return ReadBinaryProto(Env::Default(), fname, proto);
}

22. グラフの初期化
aot/compile.cc : InitGraph
新しいグラフを生成
std::unique_ptr<Graph> g(new Graph(flib));
GraphDef copy_def(graph_def);
AddDefaultAttrsToGraphDef(&copy_def, *g->op_registry(), 0);
グラフ定義(GraphDef)からグラフに変換
ConvertGraphDefToGraph(GraphConstructorOptions(), copy_def, g.get());
Feed/Fetchをノード(_Arg/Retval)としてグラフに追加
RewriteAndPruneGraph(g.get(), config, flags));

23. グラフのコンパイル
aot/compile.cc : CompileGraph
TensorFlowグラフをXLA(HLO)フォーマットに変換
ConvertGraphToXla(client, std::move(graph), flib,
&computation, &compile_result->has_context_arg);
コンパイルオプションの設定
xla::cpu::CpuAotCompilationOptions aot_opts(
flags.target_triple, flags.target_cpu, flags.target_features,
flags.entry_point,
xla::cpu::CpuAotCompilationOptions::RelocationModel::BigPic);
XLA(HLO)をコンパイル
return CompileXla(client, computation, aot_opts, compile_result);

24. ファイルの書き出し
aot/tfcompile_main.cc:Main
オブジェクトファイルの書き出し
const std::vector<char>& obj = compile_result.aot->object_file_data();
WriteStringToFile(env, flags.out_object, StringPiece(obj.data(), obj.size()));
CPPクラス名の解析
ParseCppClass(flags.cpp_class, &header_opts.class_name,
&header_opts.namespaces));
ヘッダの生成
GenerateHeader(header_opts, config, compile_result, &header));
ファイルの書き出し
WriteStringToFile(env, flags.out_header, header));

25. ありがとうございました