More Related Content Similar to TensorFlow XLAは、 中で何をやっているのか? (20) More from Mr. Vengineer (20) TensorFlow XLAは、 中で何をやっているのか?1. TensorFlow User Group ハード部 #2
2017/4/21
TensorFlow XLAは、
中で何をやっているのか?
TensorFlow r1.0(r1.1)
で公開されたXLAの
ソースコードを追ってみ
ました
@Vengineer
2. 勉強会主催 :
Xilinx Zynq MPSoC (2016/02/20)
Altera SDK for OpenCL (2016/06/10)
Xilinx SDSoC (2017/01/28)
PYNQ祭り (2017/03/04)
ブログ : Vengineerの戯言
http://blogs.yahoo.co.jp/verification_engineer
Twitter : @Vengineer
書籍 : SystemVerilogスタートアップ
http://www.cqpub.co.jp/hanbai/books/36/36191.htm
自己紹介
8. TensorFlow XLAって何?
Recap of TensorFlow DEV SUMMIT 2017で
発表された「XLAコンパイラ」
足立昌彦さん(株式会社カブク)
資料と解説を見てちょうだい
詳しくは、「TensorFlow XLAの情報と発表」
http://blogs.yahoo.co.jp/verification_engineer/71068236.html
19. RewriteGraphForExecution
core/graph/subgraph.cc
Feedノードを追加 (_Recv : .Attr("client_terminated", true))
if (!fed_outputs.empty()) {
FeedInputs( g, device_info, fed_outputs, &name_index );
}
Fetchノードを追加 (_Send : .Attr("client_terminated", true))
std::vector<Node*> fetch_nodes;
if (!fetch_outputs.empty()) {
FetchOutputs( g, device_info, fetch_outputs,
&name_index, &fetch_nodes );
}
20. SimplePlacer::Run
core/common_runtime/simple_placer.cc
1. First add all of the nodes.
2. Enumerate the constraint edges,
and use them to update the disjoint node set.
3. For each node, assign a device based on the constraints in the
disjoint node set.
4. Perform a second pass assignment for those nodes explicitly
skipped during the first pass.
23. デバイスを gpu にすると
def test_gpu(self):
with tf.Session() as sess:
x = tf.placeholder(tf.float32, [2], name="x")
with tf.device("gpu"):
y = x * 2
result = sess.run(y, {x: [1.5, 0.5]})
34. gpu を XLA_CPU に変更
def testXLA_JIT(self):
with tf.Session() as sess:
x = tf.placeholder(tf.float32, [2], name="x")
with tf.device("device:XLA_CPU:0"):
y = x * 2
result = sess.run(y, {x: [1.5, 0.5]})
46. TensorFlow XLA : JITでは!
同じデバイス内で実行できるSubgraph単位の
ノードをギュギュッと1つにまとめて、
_XlaLaunch Op
内で実行する
_XlaLaunchは、
TensorFlow XLA専用のOpとして実装
47. _XlaLaunch Opで実装は?
・Register the new Op in a C++ file
・Implement the Op in C++
compiler/kernels/xla_local_launch_op.h
compiler/kernels/xla_local_launch_op.cc
48. _XlaLaunch Op の登録
REGISTER_OP("_XlaLaunch")
.Input("constants: Tconstants")
.Attr("Tconstants: list(type) >= 0")
.Input("args: Targs")
.Attr("Targs: list(type) >= 0")
.Output("results: Tresults")
.Attr("Tresults: list(type) >= 0")
.Attr("function: func")
.Doc("XLA Launch Op. For use by the XLA JIT only.");
50. TensorFlow w/XLA: TensorFlow, Compiled! Expressiveness with performance
https://autodiff-workshop.github.io/slides/JeffDean.pdf
Computeの処理
ここに LLVM を使っている
compiler->Compile executable->Run
55. TensorFlow w/XLA: TensorFlow, Compiled! Expressiveness with performance
https://autodiff-workshop.github.io/slides/JeffDean.pdf
BuildExecutable
BuildHloModule
backend->compiler()->Compile
CpuExecutable
56. Service::BuildExecutable
xla/service/service.cc
グラフからXLA HLOへ変換
for (const VersionedComputationHandle& versioned_handle : versioned_handles) {
auto module = computation_tracker_.BuildHloModule(
Versioned_handle, true));
modules.push_back(std::move(module));
}
….
XLA HLOからLLVM IR => Executableに変換
std::vector<std::unique_ptr<Executable>> executables =
backend->compiler()->Compile(
std::move(modules), std::move(module_configs),
hlo_dumper, std::move(executors)));
58. TensorFlow w/XLA: TensorFlow, Compiled! Expressiveness with performance
https://autodiff-workshop.github.io/slides/JeffDean.pdf
executable->Runを実行
executable->RunCpuExecutable
60. ExecuteOnStream
xla/service/cpu/cpu_executable.cc
se::Stream* stream = run_options->stream();
メモリの割当て
DeviceMemoryAllocator* memory_allocator = run_options->allocator();
std::vector<se::DeviceMemoryBase> buffers(assignment_->Allocations().size());
AllocateBuffers(
memory_allocator, stream->parent()->device_ordinal(), &buffers);
関数の実行
ExecuteComputeFunction(run_options, arguments, buffers,
hlo_execution_profile));
70. 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( …., …., …. );
72. グラフの初期化
aot/compile.cc : InitGraph
グラフ定義とグラフを生成
std::unique_ptr<Graph> g(new Graph(flib)); グラフ
GraphDef copy_def(graph_def); グラフ定義
AddDefaultAttrsToGraphDef(©_def, *g->op_registry(), 0);
グラフ定義(GraphDef)からグラフに変換
ConvertGraphDefToGraph(GraphConstructorOptions(), copy_def, g.get());
Feed/Fetchをノード(_Arg/_Retval)としてグラフに追加
RewriteAndPruneGraph(g.get(), config, flags));
75. グラフのコンパイル
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);
79. ConvertGraphToXla
aot/compile.cc
XlaCompilerのCompileGraph関数を実行
std::unique_ptr<FunctionLibraryRuntime> flib_run(NewFunctionLibraryRuntime(
compiler.device_mgr(), Env::Default(), compiler.device(),
graph->versions().producer(), flib_def, OptimizerOptions()));
XlaCompiler::CompilationResult result;
compiler.CompileGraph("tfcompile", std::move(graph),
flib_run.get(), xla_args, false, &result);
グラフのコンパイル結果を XLA Computation を取り出す
*computation = std::move(result.computation);
85. CompileXla
aot/compile.cc : CompileXla
xla::LocalClient* client;
xla::LocalClient::AheadOfTimeComputationInstance instance;
instance.computation = &computation;
instance.argument_layouts = std::move(arg_layouts);
instance.result_layout = &pshape->result();
xla::StatusOr<std::vector<std::unique_ptr<xla::AotCompilationResult>>>
aot_or = client->CompileAheadOfTime({instance}, aot_opts);
88. BuildHloModule
xla/service/computation_tracker.cc
for (auto versioned_handle : post_order) {
UserComputation* computation =
ResolveInternal(versioned_handle.handle).ValueOrDie();
std::unique_ptr<HloComputation> hlo_computation =
computation->BuildHloComputation(
versioned_handle.version, resolver, include_unused_parameters));
hlo_computations[versioned_handle] = hlo_computation.get();
if (computation == entry_computation) {
module->AddEntryComputation(std::move(hlo_computation));
} else {
module->AddEmbeddedComputation(std::move(hlo_computation));
}
}