Tensorflow Lite Internal

BuiltinOpResolver::BuiltinOpResolver() {
    AddBuiltin(BuiltinOperator_ABS, Register_ABS());
    AddBuiltin(BuiltinOperator_RELU, Register_RELU(), 1, 2);
    // ...
    AddBuiltin(BuiltinOperator_MAX_POOL_2D, Register_MAX_POOL_2D(), 1, 3);
    AddBuiltin(BuiltinOperator_L2_POOL_2D, Register_L2_POOL_2D());
    AddBuiltin(BuiltinOperator_CONV_2D, Register_CONV_2D(), 1, 5);
    // ....
    AddBuiltin(BuiltinOperator_FULLY_CONNECTED, Register_FULLY_CONNECTED(), 1,
               9);
    // ...
}

void MutableOpResolver::AddBuiltin(tflite::BuiltinOperator op,
                                   const TfLiteRegistration* registration,
                                   int version) {
    TfLiteRegistration new_registration = *registration;
    new_registration.custom_name = nullptr;
    new_registration.builtin_code = op;
    new_registration.version = version;
    auto op_key = std::make_pair(op, version);
    builtins_[op_key] = new_registration;
}

TfLiteRegistration* Register_ABS() {
    static TfLiteRegistration r = {
        // init
        nullptr,
        // free
        nullptr,
        // prepare
        elementwise::GenericPrepare<elementwise::IsNumericSupportedType,
                                    elementwise::kAbsName>,
        // invoke
        elementwise::AbsEval};
    return &r;
}

TfLiteStatus AbsEval(TfLiteContext* context, TfLiteNode* node) {
    return EvalNumeric(context, node, std::abs);
}

inline TfLiteStatus EvalNumeric(TfLiteContext* context, TfLiteNode* node,
                                float float_func(float)) {
    return EvalImpl<float>(context, node, float_func, kTfLiteFloat32);
}

template <typename T>
inline TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node,
                             T func(T), TfLiteType expected_type) {
    const TfLiteTensor* input = GetInput(context, node, 0);
    TfLiteTensor* output = GetOutput(context, node, 0);
    TF_LITE_ENSURE_TYPES_EQ(context, input->type, expected_type);
    const int64_t num_elements = NumElements(input);
    const T* in_data = GetTensorData<T>(input);
    T* out_data = GetTensorData<T>(output);
    for (int64_t i = 0; i < num_elements; ++i) {
        // !!! out_data[i] = std::abs(in_data[i])
        out_data[i] = func(in_data[i]);
    }
    return kTfLiteOk;
}

TfLiteStatus InterpreterBuilder::operator()(
    std::unique_ptr<Interpreter>* interpreter, int num_threads) {
    // !!! 查找 resolver 中记录的 registration
    BuildLocalIndexToRegistrationMapping();

    auto* subgraphs = model_->subgraphs();
    auto* buffers = model_->buffers();

    if (subgraphs->size() > 1) {
        (*interpreter)->AddSubgraphs(subgraphs->size() - 1);
    }

    for (int subgraph_index = 0; subgraph_index < subgraphs->size();
         ++subgraph_index) {
        const tflite::SubGraph* subgraph = (*subgraphs)[subgraph_index];
        tflite::Subgraph* modified_subgraph =
            (*interpreter)->subgraph(subgraph_index);
        auto operators = subgraph->operators();
        auto tensors = subgraph->tensors();

        modified_subgraph->AddTensors(tensors->size());

        modified_subgraph->SetInputs(
            FlatBufferIntArrayToVector(subgraph->inputs()));
        modified_subgraph->SetOutputs(
            FlatBufferIntArrayToVector(subgraph->outputs()));

        // !!! 根据 operator 生成 node 及 execution_plan
        ParseNodes(operators, modified_subgraph);
        ParseTensors(buffers, tensors, modified_subgraph);

        std::vector<int> variables;
        for (int i = 0; i < modified_subgraph->tensors_size(); ++i) {
            auto* tensor = modified_subgraph->tensor(i);
            if (tensor->is_variable) {
                variables.push_back(i);
            }
        }
        modified_subgraph->SetVariables(std::move(variables));
    }

    // !!! 是否使用 flex delegate 支持 tensorflow op (而不是 tflite builtin op)
    // https://www.tensorflow.org/lite/guide/ops_select?hl=zh-cn
    ApplyDelegates(interpreter->get(), num_threads)
}

Interpreter::Invoke:
    primary_subgraph().Invoke()
        for (int execution_plan_index = 0; execution_plan_index < execution_plan_.size(); execution_plan_index++):
            int node_index = execution_plan_[execution_plan_index];
            TfLiteNode& node = nodes_and_registration_[node_index].first;
            const TfLiteRegistration& registration = nodes_and_registration_[node_index].second;
            OpInvoke(registration, &node)
                // !!! 调用到 resolver 中注册的 registration 中的 .invoke, 例如 AbsEval
                registration.invoke(node)

import tensorflow as tf
from tensorflow.keras import layers, Model

input=layers.Input(shape=(1, ))
output=layers.Lambda(lambda x: tf.sin(x)+x+tf.sin(x*2))(input)
model=Model(inputs=input,outputs=output)
model.summary()
print(model.predict([2.]))
model.save("/tmp/sin")

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model("/tmp/sin")
tflite = converter.convert()
with open ("/tmp/sin.tflite","wb") as f:
    f.write(tflite)
    print("size of sin.tflite:", len(tflite))

#include <cstdio>

#include "tensorflow/lite/builtin_ops.h"
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#include "tensorflow/lite/optional_debug_tools.h"
#include "tensorflow/lite/util.h"

using namespace tflite;

class MyDelegate {
  public:
    static bool SupportedOp(const TfLiteRegistration* registration) {
        return registration->builtin_code == kTfLiteBuiltinSin;
    }

    bool Invoke(TfLiteContext* context, TfLiteNode* node) {
        // !!! 这里 invoke 的实现是有问题的, 它假设了 subset 里只有一个 sin, 实
        // 际上多个连续的 sin 例如 tf.sin(tf.sin(x)) 会生成多个 sin node, 而对一
        // 个 subset 只会生成一个 delegate node...
        // 正确的做法是从 TfLiteDelegateParams 中获得所有的 node 信息但依次执行,
        // 例如 flex::Eval 的做法:
        //
        // TfLiteStatus DelegateKernel::Eval(TfLiteContext* context, TfLiteNode* node) {
        // // Execute the TensorFlow Ops sequentially.
        // for (auto& node_data : op_data_->nodes) {
        //   auto status = ExecuteFlexOp(context, buffer_map, node_data.get());
        //   TF_LITE_ENSURE_OK(context, ConvertStatus(context, status));
        // }
        //

        printf("sunway : invoke\n");
        const TfLiteTensor* input = GetInput(context, node, 0);
        TfLiteTensor* output = GetOutput(context, node, 0);

        float* input_data = input->data.f;
        float* output_data = output->data.f;

        const int M = input->dims->size;
        const int N = input->dims->data[0];
        for (int i = 0; i < M; i++) {
            for (int j = 0; j < N; ++j) {
                output_data[i * N + j] = sin(input_data[i * N + j]);
            }
        }

        return kTfLiteOk;
    }
};

TfLiteRegistration GetMyDelegateNodeRegistration() {
    TfLiteRegistration kernel_registration;
    kernel_registration.builtin_code = kTfLiteBuiltinDelegate;
    kernel_registration.custom_name = "MyDelegate";
    kernel_registration.free = NULL;
    kernel_registration.init = [](TfLiteContext* context, const char* buffer,
                                  size_t) -> void* {
        const TfLiteDelegateParams* delegate_params =
                reinterpret_cast<const TfLiteDelegateParams*>(buffer);
        MyDelegate* my_delegate = new MyDelegate;
        return my_delegate;
    };
    kernel_registration.invoke = [](TfLiteContext* context,
                                    TfLiteNode* node) -> TfLiteStatus {
        printf("kerenl_invoke\n");
        MyDelegate* kernel = reinterpret_cast<MyDelegate*>(node->user_data);
        kernel->Invoke(context, node);
        return kTfLiteOk;
    };
    kernel_registration.prepare = NULL;

    return kernel_registration;
}

TfLiteStatus MyDelegatePrepare(TfLiteContext* context,
                               TfLiteDelegate* delegate) {
    std::vector<int> supported_nodes;
    TfLiteIntArray* plan;
    TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));
    TfLiteNode* node;
    TfLiteRegistration* registration;
    for (int i = 0; i < plan->size; i++) {
        int node_index = plan->data[i];
        TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
            context, node_index, &node, &registration));
        if (MyDelegate::SupportedOp(registration)) {
            supported_nodes.push_back(node_index);
        }
    }
    TfLiteRegistration my_delegate_kernel_registration =
            GetMyDelegateNodeRegistration();

    TfLiteIntArray* supported_nodes_int_array =
            ::tflite::ConvertVectorToTfLiteIntArray(supported_nodes);
    // !!! replace node with new delegate registration
    auto status = context->ReplaceNodeSubsetsWithDelegateKernels(
        context, my_delegate_kernel_registration, supported_nodes_int_array,
        delegate);
    TfLiteIntArrayFree(supported_nodes_int_array);
    return status;
}

TfLiteDelegate* CreateMyDelegate() {
    TfLiteDelegate* delegate = new TfLiteDelegate;

    delegate->data_ = nullptr;
    delegate->flags = kTfLiteDelegateFlagsNone;
    delegate->Prepare = &MyDelegatePrepare;
    delegate->CopyFromBufferHandle = NULL;
    delegate->CopyToBufferHandle = NULL;
    delegate->FreeBufferHandle = NULL;

    return delegate;
}

int main(int argc, char* argv[]) {
    const char* filename = "/tmp/sin.tflite";

    std::unique_ptr<tflite::FlatBufferModel> model =
            tflite::FlatBufferModel::BuildFromFile(filename);

    tflite::ops::builtin::BuiltinOpResolver resolver;
    InterpreterBuilder builder(*model, resolver);
    std::unique_ptr<Interpreter> interpreter;
    builder(&interpreter);

    auto* my_delegate = CreateMyDelegate();
    // !!! ModifyGraphWithDelegate
    if (interpreter->ModifyGraphWithDelegate(my_delegate) != kTfLiteOk) {
        printf("sunway:ModifyGraphWithDelegate failed\n");
    }

    interpreter->AllocateTensors();
    float* input = interpreter->typed_input_tensor<float>(0);
    input[0] = 2.0;
    interpreter->Invoke();
    float* output = interpreter->typed_output_tensor<float>(0);
    printf("%f\n", output[0]);

    return 0;
}

TfLiteStatus Subgraph::ModifyGraphWithDelegate(TfLiteDelegate* delegate):
    TfLiteStatus status = delegate->Prepare(&context_, delegate);
    // !!! Prepare 如前面的 MyDelegatePrepare, 负责扫描所有 node, 找到能处理的
    // node, 然后调用 ReplaceNodeSubsetsWithDelegateKernels 来替换 node 和
    // execution plan

digraph G {
    SIN1 [color=red]
    SIN2 [color=red]
    MULTIPLY [color=green]
    ADD1 [color=green]
    ADD2 [color=green]
    X->SIN1->ADD1->ADD2
    X->ADD1
    X->MULTIPLY
    2->MULTIPLY->SIN2->ADD2->Y
}

TfLiteStatus Subgraph::ReplaceNodeSubsetsWithDelegateKernels(
    TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace,
    TfLiteDelegate* delegate) {
    registration.builtin_code = BuiltinOperator_DELEGATE;

    // Analyze the graph to find all independent node_subsets that are either
    // fully not\-this\-delegate or this\-delegate computation.
    InterpreterInfo info(this);
    std::vector<NodeSubset> node_subsets;
    // !!! 分为多个 subset
    PartitionGraphIntoIndependentNodeSubsets(&info, nodes_to_replace,
                                             &node_subsets);

    execution_plan_.clear();

    // !!! 按顺序遍历 subset, 因为各个 subset 之间已经是 topo 有序的
    for (auto& node_subset : node_subsets) {
        switch (node_subset.type) {
            // !!! node 不属于 delegate, 不需要替换, 直接追加到 execution_plan_
            case NodeSubset::kTfNonPartition:
                for (auto it = node_subset.nodes.begin(); it != node_subset.nodes.end();
                     ++it) {
                    execution_plan_.push_back(*it);
                }
                break;
            // !!! 对于 delegate, 需要生成一个新的 node 代替subset 中的 node(可
            // 能为多个 node, delegate 的 TfLiteDelegateParams 能拿到哪些 node
            // 被替换), 然后追加到 execution\_plan\_, 新的 node 会使用delegate 的
            // registration
            case NodeSubset::kTfPartition: {
                int node_index;

                TfLiteDelegateParams* params =
                        CreateDelegateParams(delegate, node_subset);
                AddNodeWithParameters(
                    node_subset.input_tensors, node_subset.output_tensors, {}, nullptr,
                    0, params, &registration, &node_index)

                // Initialize the output tensors's delegate-related fields.
                for (int tensor_index : node_subset.output_tensors) {
                    TfLiteTensor* tensor = &tensors_[tensor_index];
                    TF_LITE_ENSURE(&context_, tensor->delegate == nullptr ||
                                   tensor->delegate == delegate);
                    tensor->delegate = delegate;
                }

                // Associate the node with the delegate.
                TfLiteNode* node = &nodes_and_registration_[node_index].first;
                node->delegate = delegate;
            } break;
        }
    }
}

while (changed):
    subset = new array()
    changed = false
    is_delegate = nil

    for node in nodes:
        skip if node.input is not ready
        skip if node.visited
        if is_delegate is nil:
            is_delegate = node.is_delegate
        skip if node.is_delegate != is_delegate
        subset.push(node)
        node.visited = true
        node.output.ready = true
    if subset is not empty:
        subsets.push(subset)
        changed = true