0x01 前言

XNNPACK是一个由Google维护的算子库，在TensorFlowLite，ExecuTorch，ONNX RT等众多知名框架中使用。笔者最近在做mllm的xnnpack后端适配工作，因xnnpack缺少文档，在此记录。

xnnpack的test中展示了大部分xnnpack的API和使用方式，读者在碰到API使用问题的时候不妨去test文件夹下面找找答案；xnnpack遵循标准的doxygen注释，也较好的说明了函数和class的使用方法。

在xnn中使用的是静态图构建的方法，在开始构建静态图之前，需要初始化xnn：

xnn_initialize(nullptr /* allocator */)

使用如下API可以构建出一张subgraph，接下来的所有操作都在更改这张子图，

xnn_subgraph_t subgraph_ = nullptr;
auto status = xnn_create_subgraph(external_nums, 0, &subgraph_);

0x02 定义 Tensor

定义未经过量化的Tensor使用的API是：

uint32_t uuid;
status = xnn_define_tensor_value(
                /*subgraph*/..., /*dtype*/...,
                dims.size(), dims.data(),
                /*data=*/...,
                /*external id*/XNN_INVALID_VALUE_ID, /*flag*/0, /*id*/&uuid);

这里需要特殊解释的是external_id、flag和uuid三个值：

  1. external_id 是对于EXternal Inputs和Outputs才需要设置的，对于xnnpack内部管理的Tensor，不需要设置这个值，给出默认的XNN_INVALID_VALUE_ID就行。external_id的作用是让xnnpack可以从runtime中传入的external_values中索引到需要的Tensor值。详细解释如下：

在xnnpack中，每个Tensor都会有一个uuid，对于xnnpack自己管理的Tensor，uuid在定义的时候会由xnnpack自己生成。还记得在xnn_create_subgraph创建的时候需要传入external_nums吗？这里的external_nums就是用户侧预留的uuid。比如external_nums是3的时候，xnn_define_tensor_value就会从4开始计数给新创建的Tensor。而前3个Tensor，即external Tensor的uuid(external_id)就是1，2，3。

  1. flag

flag是用来标识这个Tensor是不是External Inputs，Outputs或者是其他类型的Tensor。比如inputs tensor的flag是flags = XNN_VALUE_FLAG_EXTERNAL_INPUT;, outputs 是 flags = XNN_VALUE_FLAG_EXTERNAL_OUTPUT;

  1. uuid

是每个Tensor的全局索引标识

题外话：

在mllm中，Tensor的define过程如下：

   void defineXpTensor(XnnpackBackend *xpb, Tensor *t, XpTensorType ttype) {
        if (t->uuid() != XNN_INVALID_VALUE_ID) return;

        auto xp_dtype = XnnpackBackend::mllmDType2XnnDType(t->dtype());

        xnn_status status;
        std::vector<size_t> dims;
        for (auto d : t->shape()) dims.push_back(d);

        uint32_t flags;
        uint32_t external_id = XNN_INVALID_VALUE_ID;

        switch (ttype) {
        case XpTensorType::Normal:
            flags = XNN_INVALID_VALUE_ID;
            break;
        case XpTensorType::ExternalInput:
            flags = XNN_VALUE_FLAG_EXTERNAL_INPUT;
            external_id = xpb->getNewEXternalId();
            break;
        case XpTensorType::ExternalOutput:
            flags = XNN_VALUE_FLAG_EXTERNAL_OUTPUT;
            external_id = xpb->getNewEXternalId();
            break;
        }

        switch (xp_dtype) {
        case xnn_datatype_fp32: {
            status = xnn_define_tensor_value(
                xpb->getXnnSubgraph(), xp_dtype,
                dims.size(), dims.data(),
                /*data=*/nullptr,
                external_id, flags, &t->uuid());
        }
        default:
            break;
        }

        switch (ttype) {
        case XpTensorType::Normal:
            break;
        case XpTensorType::ExternalInput:
        case XpTensorType::ExternalOutput:
            xpb->registerExternalValue(t->uuid(), xnn_external_value{.id = t->uuid(), .data = t->rawHostPtr()});
            xpb->registerUuidTensor(t->uuid(), t);
            break;
        }

        if (status != xnn_status_success) {
            Log::error("xnnpack backend defineXpTensor Error");
            exit(-1);
        }
    }

0x03 定义Op

这里笔者以Fully Connect Op为例，因为这个算子包含了Weight，可以更好的展示Weight Tensor是怎么定义的。

weight的定义不同于0x02中所述，因为weight有固定的数据，在mllm中是这样调用的：

status = xnn_define_tensor_value(
                xpb->getXnnSubgraph(), xp_dtype,
                dims.size(), dims.data(),
                /*data=*/t->rawHostPtr(),
                XNN_INVALID_VALUE_ID, 0, &t->uuid());

可以发现，我们直接传入了一个data ptr，这个data ptr指向的就是weight的数据。

Op的定义过程很简单，对于FP32的Fully Connect Op：

    auto status = xnn_define_fully_connected(
        xpb->getXnnSubgraph(),
        std::numeric_limits<float>::min(),
        std::numeric_limits<float>::max(),
        inputs[0]->uuid(),
        weight_params_.uuid(),
        bias_ ? bias_params_.uuid() : XNN_INVALID_VALUE_ID,
        outputs[0]->uuid(),
        0);

0x04 执行

执行的过程有以下几步：

  1. create model

通过之前定义好的subgraph来prealloc external inputs和outputs的Memory，xnnpack中提供了xnn_tensor_get_size和xnn_allocate_zero_simd_memory来帮助读者实现这一点。

  1. create runtime

读者需要事先创建好pthread线程池。

xnn_create_runtime_v4(model_.get(), nullptr, nullptr, threadpool_, flags, &runtime_);

  1. reshape runtime

xnn_reshape_runtime(runtime_);

  1. setup runtime

在setup的时候需要传入external_values。

xnn_setup_runtime_v2(runtime_, external_values_.size(), external_values_.data());

  1. invoke

xnn_invoke_runtime(runtime_)

题外话：

在mllm中，所有的执行过程被放在XpDispatch中：

ErrorCode XpDispatch::execute(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) {
    auto xnnbk = (XnnpackBackend *)(this->backend());

    // recreate runtime
    auto m_rt = xnnbk->recreateModelRuntime(thread_count_);

    // create Model
    m_rt->createModel(xnnbk->getXnnSubgraph());

    // create runtime
    m_rt->createRuntime(0);

    // reshape
    m_rt->reshapeRuntime();

    // setup
    m_rt->setupRuntime();

    // run
    if (!m_rt->invoke()) {
        return ErrorCode::NO_EXECUTION;
    }

    // update all output's ptr
    xnnbk->assignPtrToTensor();

    return ErrorCode::MLLM_NO_ERROR;
}