TensorFlow LiteのGPUコードを眺めてみた

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

先週にアップした、TensorFlow LiteのGPUのソースコードが公開されたよの続きとして、ソースコードを眺めてみた。

GPUでサポートしているのは、16ビットと32ビットの以下のOpのみ。

　・ADD v1
　・AVERAGE_POOL_2D v1
　・CONCATENATION v1
　・CONV_2D v1
　・DEPTHWISE_CONV_2D v1-2
　・FULLY_CONNECTED v1
　・LOGISTIC v1
　・LSTM v2 (Basic LSTM only)
　・MAX_POOL_2D v1
　・MUL v1
　・PAD v1
　・PRELU v1
　・RELU v1
　・RELU6 v1
　・RESHAPE v1
　・RESIZE_BILINEAR v1
　・SOFTMAX v1
　・STRIDED_SLICE v1
　・SUB v1
　・TRANSPOSE_CONV v1

OpenGL用のカーネルコードは、ここですね。

各カーネルでは、GenerateCodeなるメソッドが用意されていて、
コンパイラのCompile メソッドから呼び出されます。

OpenGL用のDelgateは、ここで、Metal用のDelgateは、ここにあります。

OpenGL用のDelgateは、このヘッダファイル内で宣言されている

TfLiteGpuDelegateCreate

TfLiteDelegate* TfLiteGpuDelegateCreate(
    const TfLiteGpuDelegateOptions* options) {
  TFLITE_LOG_PROD_ONCE(tflite::TFLITE_LOG_INFO,
                       "Created TensorFlow Lite delegate for GPU.");
  auto* gpu_delegate = new tflite::gpu::gl::Delegate(options);
  return gpu_delegate ? gpu_delegate->tflite_delegate() : nullptr;
}

gl_dedegate.hヘッダファイルには、実装が見えないようになっていますね。

TfLiteGpuDelegateCreateを呼び出して生成します。こんな感じで。。。

    // NEW: Prepare GPU delegate.
    auto* delegate = TfLiteGpuDelegateCreate(/*options=*/nullptr);
    if (interpreter->ModifyGraphWithDelegate(delegate) != kTfLiteOk) return;

切り替えは、簡単ですね。

参考資料(TensorFlow r1.5ですが)
TensorFlow Lite (r1.5) & Android 8.1 Neural Network API