Google Edge TPUのlibedgetpuのソースコードを調べてみた

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

Runtimeの中の executable.fbs の中を覗いてみたら、なんか、見つけた。

• 割り込みは、4つ。
• 画像をどのように分割するか？

まだ、いろいろあるっぽい。

また、tflite のディレクトリには、Custom Opのコードもあるね。

Google Edge TPUでTensorFlow Liteを使った時に何をやっているのかを妄想してみる from Mr. Vengineer

www.slideshare.net

このスライドの続編として、ソースコードの解析する必要があるかな？と思ったので、調べてみた。下記の2つのツイートのスレッドにて、ぶつぶつしました。

Google Coral Edge TPU の TensorFlow Lite のコードは、
5か月前に公開されているので、ちょっと眺めています。

このファイルでは、Google Edge TPU Compilerが吐き出したモデルからCustomOpを抽出し、Opの内容をチェックしている。https://t.co/y3S9vRVw2q

— Vengineer＠ (@Vengineer) 2020年11月7日

EdgeTPU Compilerにて、TensorFlow Liteモデルを変換したものは、EdgeTPUで実行できるOpsを1つのCustomOpにぎゅぎゅっと詰め込んでいます。その後のOpsはCPUで実行される。CustomOpの中のUserDataの中には、EdgeTPU内の命令などが入っている。どんなものが入っているもソースコードに書かれていたよ。 https://t.co/2fJ7L2oLc2

— Vengineer＠ (@Vengineer) 2020年11月7日

 APIとして、tflite_plugin_create_delegate にて、Delegateのインスタンスを生成しますが、そのコードはこんな感じ。内で、CreateEdgeTpuDelegateForCustomOpを返しています。

EDGETPU_EXPORT TfLiteDelegate* tflite_plugin_create_delegate(
                               char** options_keys, char** options_values, size_t num_options,
                               ErrorHandler error_handler) {
    DeviceOptions options;
    for (size_t i = 0; i < num_options; ++i)
        options[options_keys[i]] = options_values[i];

    auto context = GetEdgeTpuContext(options);
    if (!context) return nullptr;
    return platforms::darwinn::tflite::CreateEdgeTpuDelegateForCustomOp(context);
}

生成される EdgeTpuDelegateForCustomOpは、ここ で定義されています。

class EdgeTpuDelegateForCustomOp : public TfLiteDelegate {
    public:
    EdgeTpuDelegateForCustomOp(std::shared_ptr<edgetpu::EdgeTpuContext> context)
: TfLiteDelegate(TfLiteDelegateCreate()), context_(context) {
    this->data_ = context.get();
    this->Prepare = PrepareImpl;
    this->flags = kTfLiteDelegateFlagsAllowDynamicTensors;
    }

    private:
        std::shared_ptr<edgetpu::EdgeTpuContext> context_;
};

この中の this->Prepare に登録される PrepareImpl は、こちら。

edgetpu::kCustomOp のものを edgetpu_nodes に push_back しています。

この edgetpu_nodes に登録されたものが、Edge TPUで実行されます。

TfLiteStatus PrepareImpl(TfLiteContext* context, TfLiteDelegate* delegate) {
    context->SetExternalContext(
    context, kTfLiteEdgeTpuContext,
    static_cast<edgetpu::EdgeTpuContext*>(delegate->data_));

    TfLiteIntArray* plan;
    TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));

    std::vector<int> edgetpu_nodes;
    for (int node_index : TfLiteIntArrayView(plan)) {
        TfLiteNode* node;
        TfLiteRegistration* registration;
        TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
            context, node_index, &node, &registration));

        if (registration->custom_name &&
            std::strcmp(registration->custom_name, edgetpu::kCustomOp) == 0) {
            edgetpu_nodes.push_back(node_index);
        }
    }

    TfLiteRegistration registration = *edgetpu::RegisterCustomOp();
    registration.init = DelegateInit;
    registration.custom_name = kDelegateName;
    registration.version = kDelegateVersion;

    for (int node_index : edgetpu_nodes) {
        TfLiteIntArray* nodes = ConvertVectorToTfLiteIntArray({node_index});
        context->ReplaceNodeSubsetsWithDelegateKernels(
        context, registration, nodes, delegate);
        TfLiteIntArrayFree(nodes);
    }

    return kTfLiteOk;
}

上記の DelegateInit では、

void* DelegateInit(TfLiteContext* context, const char* buffer, size_t length) {
    const TfLiteDelegateParams* params =
    reinterpret_cast<const TfLiteDelegateParams*>(buffer);
    CHECK(params);

    TfLiteIntArray* nodes = params->nodes_to_replace;
    CHECK_EQ(nodes->size, 1);
    const int node_index = nodes->data[0];

    TfLiteNode* node;
    TfLiteRegistration* registration;
    CHECK(context->GetNodeAndRegistration(context, node_index, &node,
&registration) == kTfLiteOk);

    return edgetpu::RegisterCustomOp()->init(
                              context, static_cast<const char*>(node->custom_initial_data),
                              node->custom_initial_data_size);
}

RegistorCustomOp()->init は、ここ で以下のように定義されています。

この中の、CustomOpInit が対応するメソッドです。

TfLiteRegistration* RegisterCustomOp() {
    static TfLiteRegistration registration = {
        platforms::darwinn::tflite::CustomOpInit,
        platforms::darwinn::tflite::CustomOpFreeDirect, 
        platforms::darwinn::tflite::CustomOpPrepareDirect,
       platforms::darwinn::tflite::CustomOpInvoke,
    };
return &registration;
}

CustomOpInitは、下記のように、CustomOpUserDataDirectを呼んでいます。

void* CustomOpInit(TfLiteContext* context, const char* buffer, size_t length) {
    // Create new operator-specific user data.
   // Note this data is different from interpreter-specific data recorded in
    // context->GetExternalContext, which is probably not set yet when
    // this function is called.
    return new CustomOpUserDataDirect(reinterpret_cast<const uint8_t*>(buffer),
length);
}

CustomOpUserDataDirectでは、raw_model_data_にDeserializeCustomOpDataの戻り値を設定しています。このDeserializeCustomOpDataの中でCustonOpUserDataの内部解析をしています。

CustomOpUserDataDirect::CustomOpUserDataDirect(const uint8_t* buffer,
size_t length)
: raw_model_data_(DeserializeCustomOpData(buffer, length)) {}

これが、Edge TPU Compilerが吐き出したものを解析する部分っぽいですね。

// CustomOpData struct will be serialized as a Flexbuffer map with the
// following keys.
// "1" ---> integer (version)
// "2" ---> string (DEPRECATED; chip_name).
// "3" ---> string (DEPRECATED; serialized parameter-caching executable)
// "4" ---> string (serialized executable)

static const char kKeyVersion = "1";
// DEPRECATED (Don't reuse key for something else).
// static const char kKeyChipName = "2";
static const char kKeyParameterCachingExecutable = "3";
static const char kKeyExecutable = "4";
static const char kExecutionPreference[] = "5";

} // namespace