Vengineerの妄想(準備期間)

人生は短いけど、長いです。人生を楽しみましょう!

TFRTのソースコードを覗いてみた (その1)

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそすべては、SystemC v0.9公開から始まった 

今回は、ドキュメント関連です。

TFRTのドキュメントがかなりあります。とりあえず、下記の3つを読めばよさそう。

下図(URL埋め込みによる引用)は、TFRTのアーキテクチャを示すものです。

User Script => TensorFlow Python Runtime で TensorFlow Graph を作り、TF Core Runtimeで実行する感じ。TF Core Runtime から Worker としての Runtime が実装される。Worker の Runtime は、Host Runtime になり、この下に Accelerators (GPUやTPU)が付くことになる。

https://github.com/tensorflow/runtime/raw/master/documents/img/tfrt-arch.svg?sanitize=true

Host Runtime は、下図(URL埋め込みによる引用)のようになっている。

https://github.com/tensorflow/runtime/raw/master/documents/img/host-runtime.svg?sanitize=true

Host Runtime の中に、Accelerators用の Device Runtime がある。Device Runtimeは同じタイプのAcceleratorである必要はない。Acceleratorで実行できない Kernel は Host Kernels で実行する。

Host Runtimeの図の左側に、Lowering とあるが、この部分は具体的に下図(URL埋め込みによる引用)のようになっている。

GraphをLowering後、Host Program にて、MLIR => BEF 変換を行い、メモリ内に書き込んだ BEFをHost Runtime 内の BEFExecutor が実行する。

https://github.com/tensorflow/runtime/raw/master/documents/img/graph-execution.svg?sanitize=true

Acceleratorsの場合も同様に、BEFExecutor にて、Accelerator 用の Kernel が実行されることになる。

BEFとBEFExecutor は、TensorFlow XLAでの HLO => 各デバイスの実行プログラムへの変換および実行と同じ感じ。

Eager Modeでは各Op毎に実行するが、BEFExecutorではBEF内のコードを実行するので実行効率がいいのだろうか?

Host Runtimeの現状では、

All the aforementioned framework components have been built. There is also support for tensors and tensor operations, including MatMul, Relu, etc. The system design has been validated with simple inference models such as MNIST and ResNet.

にあるように、簡単な推論用のOpのみサポートしている模様。

明日のサポートしているOpを見ていきます。