1. Cloud TPU Driver API
ソースコード解析
作成：2019.11.16-17
＠Vengineer

2. ブログ (2007年～) : Vengineerの戯言
http://blogs.yahoo.co.jp/verification_engineer
SlideShare :
https://www.slideshare.net/ssuser479fa3
Twitter (2009年～) :
＠Vengineer
ソースコード解析職人

3. CQ出版社：雑誌インターフェース 2017年8月号～2019年8月号に
「TensorFlow XLA および Lite」
に関することを寄稿しました
2017年8月号 2017年9月号 2018年2月号 2018年8月号 2019年1月号
XLA AOT XLA AOT XLA JIT Lite & XLA Lite
2019年8月号
Lite
edge TPU

4. 今日のお話は、
Cloud TPUなので、
久しぶりの XLA です

5. https://cloud.google.com/tpu/docs/system-architecture
Cloud TPU
右図のように、
Google Compute Engine VMとHostsは
gRPCで接続
HostsとCloud TPUは、PCIe x16 x 2で接続
Computional Graphを gRPC 経由で受け取り、
TensorFlow Server内でXLA JITにて、TPU 用の
バイナリコードを生成して、PCIe経由でCloud
TPU内の命令コードバッファにダウンロードする

6. https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xl
a/python/tpu_driver
TPU Driver API
This repository contains the TPU driver API and network
(gRPC) transport implementation for high-performance
access to TPU hardware.
Release TPU Driver API into OSS (2019.11.13)

7. https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/
xla/python/tpu_driver/tpu_driver.h
TpuDriver
A TpuDriver manages a set of operations scheduled to run on a TPU system.
By default, two independently scheduled operations may execute in any
order. Ordering can be imposed in one of two ways:
1. Users can specify event dependencies via the `wait_for` argument.
2. Operations using buffer or program handles implicitly wait for the
handles to become ready before executing.

8. https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/
xla/python/tpu_driver/tpu_driver.h
TpuDriverは、何をするの？
● QuerySystemInfo
○ システム情報
● Allocate・AllocateTuple・Deallocate
○ メモリ関連
● TransferToDevice・TransferFromDevice・TransferFromDeviceToDevice
○ データの移動 (Host ⇔ Device)
● CompileProgram・LoadProgram・UnloadProgram・ExecuteProgram
○ プログラム関連

9. https://cloud.google.com/tpu/docs/system-architecture
● QuerySystemInfo
● Allocate
● AllocateTuple
● Deallocate
● TransferToDevice
● TransferFromDevice
● TransferFromDeviceToDevice
● CompileProgram
● LoadProgram
● UnloadProgram
● ExecuteProgram
TpuDriver

10. https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xl
a/python/tpu_driver
2つのTpuDriver
● GrpcTpuDriver
● RecordingTpuDriver (各APIのトレースするもの?)
(デバッグ用：TpuDriverをトレースするときに使う?)
下記の2つのマクロ(REGISTER_TPU_DRIVER)にて登録済み
REGISTER_TPU_DRIVER("grpc://",
[](const TpuDriverConfig& config)
REGISTER_TPU_DRIVER("record://", RegisterRecordingTpuDriver);

11. https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/
xla/python/tpu_driver/grpc_tpu_driver.cc
GrpcTpuDriver
下位階層(GrpcTpuStream)のラッパーファイルって感じ

12. QuerySystemInfo
void GrpcTpuDriver::QuerySystemInfo(SystemInfo* system_info) {
auto stub = CreateTpuDriverStub(config_);
::grpc::ClientContext ctx;
ctx.set_fail_fast(false);
ctx.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(10));
QuerySystemInfoRequest req;
QuerySystemInfoResponse resp;
::grpc::Status status = stub->QuerySystemInfo(&ctx, req, &resp);
*system_info = resp.system_info();
}
引数をそのまま下位階層に渡すだけ

13. メモリ関連
● Allocate メモリ割り当て
● AllocateTuple タプルの割り当て
● Deallocate メモリ開放

14. Allocate：その1
std::unique_ptr<BufferHandle> Allocate(
int32_t core_id, MemoryRegion region, int64_t num_bytes,
absl::Span<Event* const> wait_for) override {
return streams_[core_id]->Allocate(core_id, region, num_bytes, wait_for);
}
引数をそのまま下位階層に渡すだけ

15. Allocate：その2
std::unique_ptr<BufferHandle> Allocate(
int32_t core_id, MemoryRegion region, const xla::ShapeProto& shape,
absl::Span<Event* const> wait_for) override {
return streams_[core_id]->Allocate(core_id, region, shape, wait_for);
}
引数をそのまま下位階層に渡すだけ

16. AllocateTuple
std::unique_ptr<BufferHandle> AllocateTuple(
int32_t core_id, MemoryRegion region,
absl::Span<BufferHandle* const> children,
absl::Span<Event* const> wait_for) override {
return streams_[core_id]->AllocateTuple(core_id, region, children,
wait_for);
}
引数をそのまま下位階層に渡すだけ

17. Deallocate
std::unique_ptr<Event> Deallocate(
std::unique_ptr<BufferHandle> handle,
absl::Span<Event* const> wait_for) override {
auto* stream = static_cast<GrpcBufferHandle*>(handle.get())->stream();
return stream->Deallocate(std::move(handle), wait_for);
}
引数をそのまま下位階層に渡すだけ

18. プログラム関連
● CompileProgram コンパイル
● LoadProgram ロード
● UnloadProgram アンロード
● ExecuteProgram 実行

19. CompileProgram
std::unique_ptr<CompiledProgramHandle> CompileProgram(
const xla::HloProto& source, int32_t num_replicas,
absl::Span<Event* const> wait_for) override {
// Always compile using the first/default core's stream.
return streams_[0]->CompileProgram(source, num_replicas, wait_for);
}
HLOのコードをそのまま下位階層に渡すだけ。
戻り値は、コンパイルしたデータのハンドルのポインタ

20. LoadProgram
std::unique_ptr<LoadedProgramHandle> LoadProgram(
int32_t core_id, const CompiledProgramHandle* handle,
absl::Span<Event* const> wait_for) override {
return streams_[core_id]->LoadProgram(core_id, handle, wait_for);
}
引数をそのまま下位階層に渡すだけ

21. UnloadProgram
std::unique_ptr<Event> UnloadProgram(
std::unique_ptr<LoadedProgramHandle> handle,
absl::Span<Event* const> wait_for) override {
auto* stream =
static_cast<const
GrpcLoadedProgramHandle*>(handle.get())->stream();
return stream->UnloadProgram(std::move(handle), wait_for);
}
引数をそのまま下位階層に渡すだけ

22. ExecuteProgram
std::unique_ptr<Event> ExecuteProgram(
LoadedProgramHandle* program, absl::Span<BufferHandle* const> inputs,
absl::Span<BufferHandle* const> outputs,
const xla::DeviceAssignmentProto& device_assignment,
absl::Span<Event* const> wait_for) override {
auto* stream =
static_cast<const GrpcLoadedProgramHandle*>(program)->stream();
return stream->ExecuteProgram(program, inputs, outputs, device_assignment, wait_for);
}
引数をそのまま下位階層に渡すだけ
programは、stream情報を持っている

23. データの移動
● TransferToDevice
● TransferFromDevice
● TransferFromDeviceToDevice

24. TransferToDevice
std::unique_ptr<Event> TransferToDevice(
const void* src, BufferHandle* dst,
absl::Span<Event* const> wait_for) override {
auto* stream = static_cast<GrpcBufferHandle*>(dst)->stream();
return stream->TransferToDevice(src, dst, wait_for);
}
引数をそのまま下位階層に渡すだけ
dst(デバイス側のメモリ)は、stream情報を持っている

25. TransferFromDevice
std::unique_ptr<Event> TransferFromDevice(
const BufferHandle* src, void* dst,
absl::Span<Event* const> wait_for) override {
auto* stream = static_cast<const GrpcBufferHandle*>(src)->stream();
return stream->TransferFromDevice(src, dst, wait_for);
}
引数をそのまま下位階層に渡すだけ
src(デバイス側のメモリ)は、stream情報を持っている

26. TransferFromDeviceToDevice
std::unique_ptr<Event> TransferFromDeviceToDevice(
const BufferHandle* src, BufferHandle* dst,
absl::Span<Event* const> wait_for) override {
auto* stream = static_cast<const GrpcBufferHandle*>(src)->stream();
return stream->TransferFromDeviceToDevice(src, dst, wait_for);
}
引数をそのまま下位階層に渡すだけ
src(デバイス側のメモリ)は、stream情報を持っている

27. GrpcTpuStream
GrpcTpuStream::GrpcTpuStream(int32_t id, GrpcTpuDriver* driver,
std::unique_ptr<grpc::CloudTpuDriver::Stub> stub)
: id_(id),
driver_(driver),
stub_(std::move(stub)),
stream_(stub_->StreamExecute(&ctx_)),
writer_thread_(&GrpcTpuStream::StreamWriterFn, this),
reader_thread_(&GrpcTpuStream::StreamReaderFn, this) {}
2つのスレッド( StreamWriterFn, StreamReaderFn )を起動
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/
xla/python/tpu_driver/grpc_tpu_driver.cc

28. GrpcTpuStream
● QuerySystemInfo
○ システム情報
● Allocate・AllocateTuple・Deallocate
○ メモリ関連
● TransferToDevice・TransferFromDevice・TransferFromDeviceToDevice
○ データの移動 (Host ⇔ Device)
● CompileProgram・LoadProgram・UnloadProgram・ExecuteProgram
○ プログラム関連
各関数では、前処理を行って、AddWriteRequest(std::move(req)) にて Request
をキューに入れているだけ

29. StreamWriterFn
for (const auto& r : reqs) {
TraceMe activity(absl::StrCat("GrpcTpuStream::Send "));
::grpc::WriteOptions opts;
opts.set_no_compression().clear_buffer_hint();
stream_->Write(r, opts);
}
キューに入っている Request をまとめて、stream_->Write を実行している
sreamは、StreamExecute の戻り値
stream_(stub_->StreamExecute(&ctx_)), <= 後で説明します

30. StreamReaderFn
while (stream_->Read(&resp)) {
for (const StreamResponse::Entry entry : resp.entry()) {
if (entry.has_transfer_from()) {
memcpy(it->second.dst, entry.transfer_from().data().data(),
it->second.num_bytes);
stream_->Read にて、Responseを受け取り、処理をしている
デバイスからのデータ移動があるときに、memcpy を実行している
sreamは、StreamExecute の戻り値
stream_(stub_->StreamExecute(&ctx_)), <= 後で説明します

31. https://github.com/tensorflow/tensorflow/blob/de014ee62647d05e7ecb6039
29682832b6d82624/tensorflow/compiler/xla/python/tpu_driver/tpu_service.
proto#L159-L170
CloudTpuDriver
以下のサービスを提供：
service CloudTpuDriver {
rpc Open(OpenRequest) returns (OpenResponse);
rpc Close(CloseRequest) returns (CloseResponse);
rpc QuerySystemInfo(QuerySystemInfoRequest) returns
(QuerySystemInfoResponse);
rpc StreamExecute(stream StreamRequest) returns (stream
StreamResponse);
}

32. StreamRequest
AllocateRequest alloc = 1;
AllocateTupleRequest alloc_tuple = 2;
DeallocateRequest dealloc = 3;
TransferToDeviceRequest transfer_to = 4;
TransferFromDeviceRequest transfer_from = 5;
TransferFromDeviceToDeviceRequest transfer_from_to = 10;
CompileRequest compile = 6;
LoadProgramRequest load = 7;
UnloadProgramRequest unload = 8;
ExecuteRequest execute = 9;
TpuDriverに対応するRequest

33. StreamResponse
TransferFromDeviceResponse transfer_from = 3;
CompileResponse compile = 4;
TpuDriverに対応するRequest
CompileResponseは、TpuDriverにはないが、CompileRequestの戻り値

34. TPU Client
tensorflow.compiler.xla.python.tpu_driver.client.tpu_client

35. TpuBackend
class TpuBackend(xla_client.Backend):
"""XLA backend implemented using the Tpu driver API."""
def __init__(self, client):
"""Creates a new TpuBackend.
Args:
client: A _tpu_client.TpuClient object.
"""
super(TpuBackend, self).__init__('tpu')
self.client = client
https://github.com/tensorflow/tensorflow/blob/de014ee62647d05e7ecb6039
29682832b6d82624/tensorflow/compiler/xla/python/tpu_driver/client/tpu_cli
ent.py

36. TpuBackend::create
@staticmethod
def create(worker=None):
elif worker == 'local':
if TpuBackend._local_backend is None:
logging.info('Starting the local TPU driver.')
TpuBackend._local_backend = TpuBackend(
_tpu_client.TpuClient.Get('local://'))
return TpuBackend._local_backend
else:
return TpuBackend(_tpu_client.TpuClient.Get(worker))
worker には、’local’ か それ以外 (workerは設定されていないとダメ)

37. PyTpuClient::Get
StatusOr<std::shared_ptr<PyTpuClient>> PyTpuClient::Get(
const std::string& worker) {
tpu_driver::TpuDriverConfig driver_config;
driver_config.worker = worker;
auto client_status = tpu_driver::TpuDriverRegistry::Open(driver_config);
auto client = client_status.ConsumeValueOrDie();
std::vector<std::shared_ptr<Device>> devices;
devices.reserve(num_cores);
return std::make_shared<PyTpuClient>("tpu", std::move(client),
std::move(devices), /*host_id=*/0);
}
tpu_driver::TpuDriverRegistry::Open を呼んでいる

38. Backend
@six.add_metaclass(abc.ABCMeta)
class Backend(object):
"""Abstract base class for XLA backends."""
def __init__(self, platform):
"""Creates a new Backend.
Args:
platform: A string naming the platform; for example 'gpu'.
"""
self.platform = platform
https://github.com/tensorflow/tensorflow/blob/4406ec35b6056ea7a1314979
292407f1b1dd6409/tensorflow/compiler/xla/python/xla_client.py

39. TpuBackend::Compile
def compile(self, c_computation, compile_options):
options = _xla.ExecutableBuildOptions()
options.num_replicas = compile_options.num_replicas
if compile_options.result_layout:
options.result_layout = compile_options.result_layout
options.debug_options.xla_cpu_fast_math_honor_infs = True
options.debug_options.xla_cpu_fast_math_honor_nans = True
options.debug_options.xla_cpu_fast_math_honor_division = True
options.debug_options.xla_cpu_fast_math_honor_functions = True
options.debug_options.xla_gpu_enable_fast_min_max = False
return _tpu_client.TpuExecutable.Compile(c_computation,
compile_options.argument_layouts,
options, self.client,
compile_options.device_assignment)

40. TpuExecutable
py::class_<PyTpuExecutable>(m, "TpuExecutable")
.def_static("Compile", &PyTpuExecutable::Compile,
py::call_guard<py::gil_scoped_release>())
.def("local_devices", &PyTpuExecutable::local_devices)
.def("DeviceOrdinals", …..
.def("SizeOfGeneratedCodeInBytes",
&PyTpuExecutable::SizeOfGeneratedCodeInBytes)
.def("Delete", &PyTpuExecutable::Delete)
.def("Execute", &PyTpuExecutable::Execute,
py::call_guard<py::gil_scoped_release>(), py::arg("arguments"))
.def("ExecutePerReplica", &PyTpuExecutable::ExecutePerReplica,
py::call_guard<py::gil_scoped_release>(), py::arg("arguments"));
https://github.com/tensorflow/tensorflow/blob/tensorflow/compiler/xla/pytho
n/tpu_driver/client/tpu_client_extension.cc#L159

41. PyTpuExecutable::Compile
std::unique_ptr<tpu_driver::CompiledProgramHandle> compiled_program =
client->driver()->CompileProgram(hlo_proto, options.num_replicas(), {});
for (int replica = 0; replica < options.num_replicas(); ++replica) {
loaded_programs[replica] = client->driver()->LoadProgram(
device_ordinal, compiled_program.get(), {});
}
コンパイル後、replicaの分、プログラムをロード
https://github.com/tensorflow/tensorflow/blob/tensorflow/compiler/xla/pytho
n/tpu_driver/client/tpu_client.cc

42. PyTpuExecutable::Execute
ExecuteResult result =
ExecuteHelper(absl::MakeSpan(&all_core_arguments, 1),
argument_handles, /*replica=*/0, RunId());
Status status = WaitForExecuteEvent(result.on_execute_finished.get());
return std::move(result.buffer);
プログラムを実行し、イベント待ち、結果のバッファを返す

43. PyTpuExecutable::ExecuteHelper
std::unique_ptr<tpu_driver::Event> on_execute_finished =
client_->driver()->ExecuteProgram(
executables_[replica].get(), inputs,
{output_buffer->DeviceBuffer()->handle.get()}, device_assignment,
{ready_to_execute});
return {std::move(output_buffer), std::move(on_execute_finished)};
入力データと出力データのハンドル、デバイスの割り当てを渡して、プログラムを実
行。出力データのバッファのポインタを返すが、この時点では出力データは書き込め
られているとは限らない

44. https://www.slideshare.net/ssuser479fa3/tensorflow-local-python-xla-client
local Python XLA client
2018年2月にはソースコード解析していた

45. https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/
python/xla_client.py
LocalBackend
localClient backend
・CPU
・GPU
XLA backend implemented using
the in-process xla::LocalClient API
CPU or GPU の場合XLA client

46. xla_client : An XLA client in Python
tensorflow/compiler/xla/python/xla_client.py
テストコード
tensorflow/compiler/xla/python/xla_client_test.py
https://github.com/tensorflow/tensorflow/tree/4406ec35b6056ea7a131497929
2407f1b1dd6409/tensorflow/compiler/xla/python

47. JAX: Autograd and XLA
https://github.com/google/jax/tree/master

48. JAX: Autograd and XLA
JAX is Autograd and XLA, brought together for high-performance
machine learning research.
What’s new is that JAX uses XLA to compile and run your
NumPy programs on GPUs and TPUs.
● Cloud TPU support
● Multi-GPU and multi-TPU support

49. TpuBackendは、jaxで使われていた
from . import tpu_client
def _get_tpu_driver_backend(platform):
del platform
backend_target = FLAGS.jax_backend_target
if backend_target is None:
raise ValueError('When using TPU Driver as the backend, you must specify
'
'--jax_backend_target=<hostname>:8470.')
return tpu_client.TpuBackend.create(worker=backend_target)

50. Add TPU Driver to jaxlib #1673
Add TPU Driver: a low-level TPU API as a JAX backend for Cloud TPU.
https://github.com/tensorflow/tensorflow/commit/4406ec35b6056ea7a131497
9292407f1b1dd6409

51. LocalBackend
localClient backend
・CPU
・GPU
https://github.com/google/jax/blob/f4aa5150e87cd90debe86707f5958204a8d
9d968/jax/lib/xla_bridge.py#L98-L117
XLA backend implemented using
the in-process xla::LocalClient API
CPU or GPU の場合XLA client

52. TpuBackend
TPU の場合
https://cloud.google.com/tpu/docs/system-architecture
ここに書いてあるように、
Google Compute Engine VM上の
TpuBackend が
TensorFlow Clientに対応する？
XLA client
CloudTpuDriver

53. まとめ
Cloud TPU Driver API は、
Google / JAX で TPU を利用するために
OSS 化されたのかな？
と、思っています

54. あたしは、
ディープラーニング職人 ではありません
コンピュータエンジニア です
ありがとうございました
＠Vengineer
ソースコード解析職人